Add retry mechanics to pallet-scheduler

cumulus/parachains/runtimes/collectives/collectives-westend/src/weights/pallet_scheduler.rs

@@ -203,4 +203,28 @@ impl<T: frame_system::Config> pallet_scheduler::WeightInfo for WeightInfo<T> {
 			.saturating_add(T::DbWeight::get().reads(2))
 			.saturating_add(T::DbWeight::get().writes(2))
 	}
+	/// Dummy impl
+	fn check_retry(s: u32, ) -> Weight {
+		Weight::from_parts(2_963_000, 0)
+			.saturating_add(Weight::from_parts(0, 159279))
+			.saturating_add(Weight::from_parts(909_557, 0).saturating_mul(s.into()))
+			.saturating_add(T::DbWeight::get().reads(1))
+			.saturating_add(T::DbWeight::get().writes(1))
+	}
+	/// Dummy impl
+	fn set_retry(s: u32, ) -> Weight {
+		Weight::from_parts(2_963_000, 0)
+			.saturating_add(Weight::from_parts(0, 159279))
+			.saturating_add(Weight::from_parts(909_557, 0).saturating_mul(s.into()))
+			.saturating_add(T::DbWeight::get().reads(1))
+			.saturating_add(T::DbWeight::get().writes(1))
+	}
+	/// Dummy impl
+	fn set_retry_named(s: u32, ) -> Weight {
+		Weight::from_parts(2_963_000, 0)
+			.saturating_add(Weight::from_parts(0, 159279))
+			.saturating_add(Weight::from_parts(909_557, 0).saturating_mul(s.into()))
+			.saturating_add(T::DbWeight::get().reads(1))
+			.saturating_add(T::DbWeight::get().writes(1))
+	}
 }

polkadot/runtime/rococo/src/weights/pallet_scheduler.rs

@@ -201,4 +201,28 @@ impl<T: frame_system::Config> pallet_scheduler::WeightInfo for WeightInfo<T> {
 			.saturating_add(T::DbWeight::get().reads(2))
 			.saturating_add(T::DbWeight::get().writes(2))
 	}
+	/// Dummy impl
+	fn check_retry(s: u32, ) -> Weight {
+		Weight::from_parts(2_963_000, 0)
+			.saturating_add(Weight::from_parts(0, 159279))
+			.saturating_add(Weight::from_parts(909_557, 0).saturating_mul(s.into()))
+			.saturating_add(T::DbWeight::get().reads(1))
+			.saturating_add(T::DbWeight::get().writes(1))
+	}
+	/// Dummy impl
+	fn set_retry(s: u32, ) -> Weight {
+		Weight::from_parts(2_963_000, 0)
+			.saturating_add(Weight::from_parts(0, 159279))
+			.saturating_add(Weight::from_parts(909_557, 0).saturating_mul(s.into()))
+			.saturating_add(T::DbWeight::get().reads(1))
+			.saturating_add(T::DbWeight::get().writes(1))
+	}
+	/// Dummy impl
+	fn set_retry_named(s: u32, ) -> Weight {
+		Weight::from_parts(2_963_000, 0)
+			.saturating_add(Weight::from_parts(0, 159279))
+			.saturating_add(Weight::from_parts(909_557, 0).saturating_mul(s.into()))
+			.saturating_add(T::DbWeight::get().reads(1))
+			.saturating_add(T::DbWeight::get().writes(1))
+	}
 }

polkadot/runtime/westend/src/weights/pallet_scheduler.rs

@@ -204,4 +204,28 @@ impl<T: frame_system::Config> pallet_scheduler::WeightInfo for WeightInfo<T> {
 			.saturating_add(T::DbWeight::get().reads(2))
 			.saturating_add(T::DbWeight::get().writes(2))
 	}
+	/// Dummy impl
+	fn check_retry(s: u32, ) -> Weight {
+		Weight::from_parts(2_963_000, 0)
+			.saturating_add(Weight::from_parts(0, 159279))
+			.saturating_add(Weight::from_parts(909_557, 0).saturating_mul(s.into()))
+			.saturating_add(T::DbWeight::get().reads(1))
+			.saturating_add(T::DbWeight::get().writes(1))
+	}
+	/// Dummy impl
+	fn set_retry(s: u32, ) -> Weight {
+		Weight::from_parts(2_963_000, 0)
+			.saturating_add(Weight::from_parts(0, 159279))
+			.saturating_add(Weight::from_parts(909_557, 0).saturating_mul(s.into()))
+			.saturating_add(T::DbWeight::get().reads(1))
+			.saturating_add(T::DbWeight::get().writes(1))
+	}
+	/// Dummy impl
+	fn set_retry_named(s: u32, ) -> Weight {
+		Weight::from_parts(2_963_000, 0)
+			.saturating_add(Weight::from_parts(0, 159279))
+			.saturating_add(Weight::from_parts(909_557, 0).saturating_mul(s.into()))
+			.saturating_add(T::DbWeight::get().reads(1))
+			.saturating_add(T::DbWeight::get().writes(1))
+	}
 }

substrate/frame/scheduler/src/benchmarking.rs

@@ -306,5 +306,59 @@ benchmarks! {
 		);
 	}
 
+	check_retry {
+		let s in 1 .. T::MaxScheduledPerBlock::get();
+		let when = BLOCK_NUMBER.into();
+
+		fill_schedule::<T>(when, s)?;
+		let name = u32_to_name(s - 1);
+		let address = Lookup::<T>::get(name).unwrap();
+		let period: BlockNumberFor<T> = 1u32.into();
+		let root: <T as Config>::PalletsOrigin = frame_system::RawOrigin::Root.into();
+		Retries::<T>::insert(address, RetryConfig { total_retries: 10, remaining: 10, period });
+		let (mut when, index) = address;
+		let task = Agenda::<T>::get(when)[index as usize].clone().unwrap();
+	}: {
+		Scheduler::<T>::check_retry(when, when, index, task);
+	} verify {
+		when = when + BlockNumberFor::<T>::one();
+		assert_eq!(
+			Retries::<T>::get((when, 0)),
+			Some(RetryConfig { total_retries: 10, remaining: 9, period })
+		);
+	}
+
+	set_retry {
+		let s in 1 .. T::MaxScheduledPerBlock::get();
+		let when = BLOCK_NUMBER.into();
+
+		fill_schedule::<T>(when, s)?;
+		let name = u32_to_name(s - 1);
+		let address = Lookup::<T>::get(name).unwrap();
+		let (when, index) = address;
+	}: _(RawOrigin::Root, when, index, 10, BlockNumberFor::<T>::one())
+	verify {
+		assert_eq!(
+			Retries::<T>::get((when, index)),
+			Some(RetryConfig { total_retries: 10, remaining: 10, period: BlockNumberFor::<T>::one() })
+		);
+	}
+
+	set_retry_named {
+		let s in 1 .. T::MaxScheduledPerBlock::get();
+		let when = BLOCK_NUMBER.into();
+
+		fill_schedule::<T>(when, s)?;
+		let name = u32_to_name(s - 1);
+		let address = Lookup::<T>::get(name).unwrap();
+		let (when, index) = address;
+	}: _(RawOrigin::Root, name, 10, BlockNumberFor::<T>::one())
+	verify {
+		assert_eq!(
+			Retries::<T>::get((when, index)),
+			Some(RetryConfig { total_retries: 10, remaining: 10, period: BlockNumberFor::<T>::one() })
+		);
+	}
+
 	impl_benchmark_test_suite!(Scheduler, crate::mock::new_test_ext(), crate::mock::Test);
 }

substrate/frame/scheduler/src/lib.rs

@@ -122,6 +122,18 @@ pub type CallOrHashOf<T> =
 pub type BoundedCallOf<T> =
 	Bounded<<T as Config>::RuntimeCall, <T as frame_system::Config>::Hashing>;
 
+/// The configuration of the retry mechanism for a given task along with its current state.
+#[cfg_attr(any(feature = "std", test), derive(PartialEq, Eq))]
+#[derive(Clone, RuntimeDebug, Encode, Decode, MaxEncodedLen, TypeInfo)]
+pub struct RetryConfig<Period> {
+	/// Initial amount of retries allowed.
+	total_retries: u8,
+	/// Amount of retries left.
+	remaining: u8,
+	/// Period of time between retry attempts.
+	period: Period,
+}
+
 #[cfg_attr(any(feature = "std", test), derive(PartialEq, Eq))]
 #[derive(Clone, RuntimeDebug, Encode, Decode)]
 struct ScheduledV1<Call, BlockNumber> {
@@ -273,6 +285,15 @@ pub mod pallet {
 		ValueQuery,
 	>;
 
+	#[pallet::storage]
+	pub type Retries<T: Config> = StorageMap<
+		_,
+		Blake2_128Concat,
+		TaskAddress<BlockNumberFor<T>>,
+		RetryConfig<BlockNumberFor<T>>,
+		OptionQuery,
+	>;
+
 	/// Lookup from a name to the block number and index of the task.
 	///
 	/// For v3 -> v4 the previously unbounded identities are Blake2-256 hashed to form the v4
@@ -299,6 +320,8 @@ pub mod pallet {
 		CallUnavailable { task: TaskAddress<BlockNumberFor<T>>, id: Option<TaskName> },
 		/// The given task was unable to be renewed since the agenda is full at that block.
 		PeriodicFailed { task: TaskAddress<BlockNumberFor<T>>, id: Option<TaskName> },
+		/// The given task was unable to be retried since the agenda is full at that block.
+		RetryFailed { task: TaskAddress<BlockNumberFor<T>>, id: Option<TaskName> },
 		/// The given task can never be executed since it is overweight.
 		PermanentlyOverweight { task: TaskAddress<BlockNumberFor<T>>, id: Option<TaskName> },
 	}
@@ -440,6 +463,86 @@ pub mod pallet {
 			)?;
 			Ok(())
 		}
+
+		/// Set a retry configuration for a task so that, in case its scheduled run fails, it will
+		/// be retried after `period` blocks, for a total amount of `retries` retries or until it
+		/// succeeds.
+		///
+		/// If a retried task runs successfully before running out of retries, its remaining retry
+		/// counter will be reset to the initial value. If a retried task runs out of retries, it
+		/// will be removed from the schedule.
+		///
+		/// Tasks which need to be scheduled for a retry are still subject to weight metering and
+		/// agenda space, same as a regular task. Periodic tasks will have their periodic schedule
+		/// put on hold while the task is retrying.
+		#[pallet::call_index(6)]
+		#[pallet::weight(<T as Config>::WeightInfo::set_retry(T::MaxScheduledPerBlock::get()))]
+		pub fn set_retry(
+			origin: OriginFor<T>,
+			when: BlockNumberFor<T>,
+			index: u32,
+			retries: u8,
+			period: BlockNumberFor<T>,
+		) -> DispatchResult {
+			T::ScheduleOrigin::ensure_origin(origin.clone())?;
+			let origin = <T as Config>::RuntimeOrigin::from(origin);
+			let agenda = Agenda::<T>::get(when);
+			match agenda.get(index as usize) {
+				Some(Some(task)) =>
+					if matches!(
+						T::OriginPrivilegeCmp::cmp_privilege(origin.caller(), &task.origin),
+						Some(Ordering::Less) | None
+					) {
+						return Err(BadOrigin.into())
+					},
+				_ => return Err(Error::<T>::NotFound.into()),
+			}
+			Retries::<T>::insert(
+				(when, index),
+				RetryConfig { total_retries: retries, remaining: retries, period },
+			);
+			Ok(())
+		}
+
+		/// Set a retry configuration for a named task so that, in case its scheduled run fails, it
+		/// will be retried after `period` blocks, for a total amount of `retries` retries or until
+		/// it succeeds.
+		///
+		/// If a retried task runs successfully before running out of retries, its remaining retry
+		/// counter will be reset to the initial value. If a retried task runs out of retries, it
+		/// will be removed from the schedule.
+		///
+		/// Tasks which need to be scheduled for a retry are still subject to weight metering and
+		/// agenda space, same as a regular task. Periodic tasks will have their periodic schedule
+		/// put on hold while the task is retrying.
+		#[pallet::call_index(7)]
+		#[pallet::weight(<T as Config>::WeightInfo::set_retry_named(T::MaxScheduledPerBlock::get()))]
+		pub fn set_retry_named(
+			origin: OriginFor<T>,
+			id: TaskName,
+			retries: u8,
+			period: BlockNumberFor<T>,
+		) -> DispatchResult {
+			T::ScheduleOrigin::ensure_origin(origin.clone())?;
+			let origin = <T as Config>::RuntimeOrigin::from(origin);
+			let (when, agenda_index) = Lookup::<T>::get(&id).ok_or(Error::<T>::NotFound)?;
+			let agenda = Agenda::<T>::get(when);
+			match agenda.get(agenda_index as usize) {
+				Some(Some(task)) =>
+					if matches!(
+						T::OriginPrivilegeCmp::cmp_privilege(origin.caller(), &task.origin),
+						Some(Ordering::Less) | None
+					) {
+						return Err(BadOrigin.into())
+					},
+				_ => return Err(Error::<T>::NotFound.into()),
+			}
+			Retries::<T>::insert(
+				(when, agenda_index),
+				RetryConfig { total_retries: retries, remaining: retries, period },
+			);
+			Ok(())
+		}
 	}
 }
 
@@ -1089,7 +1192,7 @@ impl<T: Config> Pallet<T> {
 		when: BlockNumberFor<T>,
 		agenda_index: u32,
 		is_first: bool,
-		mut task: ScheduledOf<T>,
+		task: ScheduledOf<T>,
 	) -> Result<(), (ServiceTaskError, Option<ScheduledOf<T>>)> {
 		if let Some(ref id) = task.maybe_id {
 			Lookup::<T>::remove(id);
@@ -1124,11 +1227,24 @@ impl<T: Config> Pallet<T> {
 			},
 			Err(()) => Err((Overweight, Some(task))),
 			Ok(result) => {
+				let failed = result.is_err();
 				Self::deposit_event(Event::Dispatched {
 					task: (when, agenda_index),
 					id: task.maybe_id,
 					result,
 				});
+
+				let mut task = if failed {
+					let _ = weight
+						.try_consume(T::WeightInfo::check_retry(T::MaxScheduledPerBlock::get()));
+					match Self::check_retry(now, when, agenda_index, task) {
+						Some(task) => task,
+						None => return Ok(()),
+					}
+				} else {
+					task
+				};
+
 				if let &Some((period, count)) = &task.maybe_periodic {
 					if count > 1 {
 						task.maybe_periodic = Some((period, count - 1));
@@ -1137,18 +1253,31 @@ impl<T: Config> Pallet<T> {
 					}
 					let wake = now.saturating_add(period);
 					match Self::place_task(wake, task) {
-						Ok(_) => {},
+						Ok(new_address) => {
+							if let Some(RetryConfig { total_retries, remaining, period }) =
+								Retries::<T>::take((when, agenda_index))
+							{
+								// Reset the counter if the task ran successfully.
+								let remaining = if failed { remaining } else { total_retries };
+								Retries::<T>::insert(
+									new_address,
+									RetryConfig { total_retries, remaining, period },
+								);
+							}
+						},
 						Err((_, task)) => {
 							// TODO: Leave task in storage somewhere for it to be rescheduled
 							// manually.
 							T::Preimages::drop(&task.call);
+							let _ = Retries::<T>::take((when, agenda_index));
 							Self::deposit_event(Event::PeriodicFailed {
 								task: (when, agenda_index),
 								id: task.maybe_id,
 							});
 						},
 					}
 				} else {
+					let _ = Retries::<T>::take((when, agenda_index));
 					T::Preimages::drop(&task.call);
 				}
 				Ok(())
@@ -1192,6 +1321,52 @@ impl<T: Config> Pallet<T> {
 		let _ = weight.try_consume(call_weight);
 		Ok(result)
 	}
+
+	/// Check if a task has a retry configutation in place and, if so, try to reschedule it.
+	///
+	/// If the task was successfully rescheduled for later, this function will return `None`.
+	/// Otherwise, if the task was not rescheduled, either because there was no retry configuration
+	/// in place, there were no more retry attempts left, or the agenda was full, this function
+	/// will return the task so it can be handled somewhere else.
+	fn check_retry(
+		now: BlockNumberFor<T>,
+		when: BlockNumberFor<T>,
+		agenda_index: u32,
+		task: ScheduledOf<T>,
+	) -> Option<ScheduledOf<T>> {
+		// Check if we have a retry configuration set.
+		match Retries::<T>::take((when, agenda_index)) {
+			Some(RetryConfig { total_retries, mut remaining, period }) => {
+				remaining = match remaining.checked_sub(1) {
+					Some(n) => n,
+					None => return Some(task),
+				};
+				let wake = now.saturating_add(period);
+				match Self::place_task(wake, task) {
+					Ok(address) => {
+						// Reinsert the retry config to the new address of the task after it was
+						// placed.
+						Retries::<T>::insert(
+							address,
+							RetryConfig { total_retries, remaining, period },
+						);
+						None
+					},
+					Err((_, task)) => {
+						// TODO: Leave task in storage somewhere for it to be
+						// rescheduled manually.
+						T::Preimages::drop(&task.call);
+						Self::deposit_event(Event::RetryFailed {
+							task: (when, agenda_index),
+							id: task.maybe_id,
+						});
+						Some(task)
+					},
+				}
+			},
+			None => Some(task),
+		}
+	}
 }
 
 impl<T: Config> schedule::v2::Anon<BlockNumberFor<T>, <T as Config>::RuntimeCall, T::PalletsOrigin>