introducing overflow count to handle ignored stack population when re…

…ached max size

pallets/support/src/lib.rs

@@ -66,6 +66,10 @@ pub mod pallet {
 	#[pallet::getter(fn execution_context)]
 	pub(super) type ExecutionContext<T: Config> = StorageValue<_, ExecutionIdStack, ValueQuery>;
 
+	///To handle the overflow of increasing the execution context.
+	#[pallet::storage]
+	pub(super) type OverflowCount<T: Config> = StorageValue<_, u32, ValueQuery>;
+
 	#[pallet::error]
 	pub enum Error<T> {}
 
@@ -89,8 +93,9 @@ pub mod pallet {
 	impl<T: Config> Hooks<BlockNumberFor<T>> for Pallet<T> {
 		fn on_initialize(_n: BlockNumberFor<T>) -> Weight {
 			ExecutionContext::<T>::kill();
-
-			T::DbWeight::get().reads_writes(1, 1)
+			OverflowCount::<T>::kill();
+
+			T::DbWeight::get().reads_writes(2, 2)
 		}
 	}
 
@@ -136,6 +141,7 @@ impl<T: Config> Pallet<T> {
 			//We make it fire and forget, and it should fail only in test and when if wrongly used
 			debug_assert_ne!(stack.len(), MAX_STACK_SIZE as usize, "Stack should not be full");
 			if let Err(err) = stack.try_push(execution_type(next_id)) {
+				OverflowCount::<T>::mutate(|count| *count += 1);
 				log::warn!(target: LOG_TARGET, "The max stack size of execution stack has been reached: {:?}", err);
 			}
 
@@ -152,19 +158,23 @@ impl<T: Config> Pallet<T> {
 		where
 		F: FnOnce(u32) -> ExecutionType
 	{
-		ExecutionContext::<T>::mutate(|stack| {
-			//We make it fire and forget, and it should fail only in test and when if wrongly used
-			debug_assert_ne!(stack.len(), 0, "The stack should not be empty when decreased");
-
-			if let Some(last_stack_entry) = stack.last() {
-				let expected_last_entry = expected_last_stack_entry(0);//We use a dummy 0 as id as we only compare type
-				if discriminant(last_stack_entry) == discriminant(&expected_last_entry) {
-					if stack.pop().is_none() {
-						log::warn!(target: LOG_TARGET,"The execution stack should not be empty when decreased. The stack should be populated first, or should not be decreased more than its size");
+		if OverflowCount::<T>::get() > 0 {
+			OverflowCount::<T>::mutate(|count| *count -= 1);
+		} else {
+			ExecutionContext::<T>::mutate(|stack| {
+				//We make it fire and forget, and it should fail only in test and when if wrongly used
+				debug_assert_ne!(stack.len(), 0, "The stack should not be empty when decreased");
+
+				if let Some(last_stack_entry) = stack.last() {
+					let expected_last_entry = expected_last_stack_entry(0);//We use a dummy 0 as id as we only compare type
+					if discriminant(last_stack_entry) == discriminant(&expected_last_entry) {
+						if stack.pop().is_none() {
+							log::warn!(target: LOG_TARGET,"The execution stack should not be empty when decreased. The stack should be populated first, or should not be decreased more than its size");
+						}
 					}
 				}
-			}
-		});
+			});
+		}
 	}
 
 	pub fn get_context() -> Vec<ExecutionType> {

pallets/support/src/tests/incremental_id.rs

@@ -15,6 +15,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+use frame_support::traits::Len;
 use crate::tests::mock::*;
 use crate::Event;
 
@@ -148,3 +149,54 @@ fn entry_is_removed_when_type_matched_with_last_stack_item() {
 		assert_eq!(AmmSupport::execution_context().into_inner(), vec![]);
 	});
 }
+
+#[test]
+fn removing_invalid_type_should_not_decrease_context() {
+	ExtBuilder::default().build().execute_with(|| {
+		let types = vec![ExecutionType::Omnipool,ExecutionType::Router, ExecutionType::XcmExchange, ExecutionType::Batch];
+		for i in 0..MAX_STACK_SIZE {
+			let idx = i as usize % types.len();
+			let operation_type = types[idx];
+			AmmSupport::add_to_context(operation_type);
+		}
+
+		assert_eq!(AmmSupport::execution_context().len(), 16);
+
+		AmmSupport::remove_from_context(|id|ExecutionType::DCA(0, id));
+		AmmSupport::remove_from_context(|id|ExecutionType::Xcm([1u8;32], id));
+
+		assert_eq!(AmmSupport::execution_context().len(), 16);
+	});
+}
+
+//This test is ignored because it is not possible to overflow when running tests in non-release mode
+#[ignore]
+#[test]
+fn overflow_should_be_handled_when_max_stack_size_reached() {
+	ExtBuilder::default().build().execute_with(|| {
+		for _ in 0..MAX_STACK_SIZE {
+			AmmSupport::add_to_context(ExecutionType::Batch);
+		}
+
+		AmmSupport::add_to_context(ExecutionType::Batch);
+		AmmSupport::add_to_context(ExecutionType::Batch);
+		AmmSupport::add_to_context(ExecutionType::Batch);
+
+		assert_eq!(AmmSupport::execution_context().len(), 16);
+
+		//We remove the batch 3 times to check if overflow handled
+		AmmSupport::remove_from_context(ExecutionType::Batch);
+		assert_eq!(AmmSupport::execution_context().len(), 16);
+
+		AmmSupport::remove_from_context(ExecutionType::Batch);
+		assert_eq!(AmmSupport::execution_context().len(), 16);
+
+		AmmSupport::remove_from_context(ExecutionType::Batch);
+		assert_eq!(AmmSupport::execution_context().len(), 16);
+
+		//Check if stack behaves normally after overflow
+		AmmSupport::remove_from_context(ExecutionType::Batch);
+		assert_eq!(AmmSupport::execution_context().len(), 15);
+	});
+}
+