Scheduler: remove empty agenda on cancel (#12989)

* Scheduler: remove empty agenda on cancel

* use iter any

* fix benches

* remove trailing None

* Add CleanupAgendas migration

Signed-off-by: Oliver Tale-Yazdi <oliver.tale-yazdi@parity.io>

* fix ci

Signed-off-by: Oliver Tale-Yazdi <oliver.tale-yazdi@parity.io>

* Count non-empty agendas in migration

Signed-off-by: Oliver Tale-Yazdi <oliver.tale-yazdi@parity.io>

Signed-off-by: Oliver Tale-Yazdi <oliver.tale-yazdi@parity.io>
Co-authored-by: Oliver Tale-Yazdi <oliver.tale-yazdi@parity.io>

substrate/frame/scheduler/src/benchmarking.rs

@@ -244,13 +244,17 @@ benchmarks! {
 	}: _<SystemOrigin<T>>(schedule_origin, when, 0)
 	verify {
 		ensure!(
-			Lookup::<T>::get(u32_to_name(0)).is_none(),
-			"didn't remove from lookup"
+			s == 1 || Lookup::<T>::get(u32_to_name(0)).is_none(),
+			"didn't remove from lookup if more than 1 task scheduled for `when`"
 		);
 		// Removed schedule is NONE
 		ensure!(
-			Agenda::<T>::get(when)[0].is_none(),
-			"didn't remove from schedule"
+			s == 1 || Agenda::<T>::get(when)[0].is_none(),
+			"didn't remove from schedule if more than 1 task scheduled for `when`"
+		);
+		ensure!(
+			s > 1 || Agenda::<T>::get(when).len() == 0,
+			"remove from schedule if only 1 task scheduled for `when`"
 		);
 	}
 
@@ -280,13 +284,17 @@ benchmarks! {
 	}: _(RawOrigin::Root, u32_to_name(0))
 	verify {
 		ensure!(
-			Lookup::<T>::get(u32_to_name(0)).is_none(),
-			"didn't remove from lookup"
+			s == 1 || Lookup::<T>::get(u32_to_name(0)).is_none(),
+			"didn't remove from lookup if more than 1 task scheduled for `when`"
 		);
 		// Removed schedule is NONE
 		ensure!(
-			Agenda::<T>::get(when)[0].is_none(),
-			"didn't remove from schedule"
+			s == 1 || Agenda::<T>::get(when)[0].is_none(),
+			"didn't remove from schedule if more than 1 task scheduled for `when`"
+		);
+		ensure!(
+			s > 1 || Agenda::<T>::get(when).len() == 0,
+			"remove from schedule if only 1 task scheduled for `when`"
 		);
 	}
 

substrate/frame/scheduler/src/lib.rs

@@ -752,6 +752,22 @@ impl<T: Config> Pallet<T> {
 		Ok(index)
 	}
 
+	/// Remove trailing `None` items of an agenda at `when`. If all items are `None` remove the
+	/// agenda record entirely.
+	fn cleanup_agenda(when: T::BlockNumber) {
+		let mut agenda = Agenda::<T>::get(when);
+		match agenda.iter().rposition(|i| i.is_some()) {
+			Some(i) if agenda.len() > i + 1 => {
+				agenda.truncate(i + 1);
+				Agenda::<T>::insert(when, agenda);
+			},
+			Some(_) => {},
+			None => {
+				Agenda::<T>::remove(when);
+			},
+		}
+	}
+
 	fn do_schedule(
 		when: DispatchTime<T::BlockNumber>,
 		maybe_periodic: Option<schedule::Period<T::BlockNumber>>,
@@ -802,6 +818,7 @@ impl<T: Config> Pallet<T> {
 			if let Some(id) = s.maybe_id {
 				Lookup::<T>::remove(id);
 			}
+			Self::cleanup_agenda(when);
 			Self::deposit_event(Event::Canceled { when, index });
 			Ok(())
 		} else {
@@ -824,6 +841,7 @@ impl<T: Config> Pallet<T> {
 			ensure!(!matches!(task, Some(Scheduled { maybe_id: Some(_), .. })), Error::<T>::Named);
 			task.take().ok_or(Error::<T>::NotFound)
 		})?;
+		Self::cleanup_agenda(when);
 		Self::deposit_event(Event::Canceled { when, index });
 
 		Self::place_task(new_time, task).map_err(|x| x.0)
@@ -880,6 +898,7 @@ impl<T: Config> Pallet<T> {
 					}
 					Ok(())
 				})?;
+				Self::cleanup_agenda(when);
 				Self::deposit_event(Event::Canceled { when, index });
 				Ok(())
 			} else {
@@ -905,6 +924,7 @@ impl<T: Config> Pallet<T> {
 			let task = agenda.get_mut(index as usize).ok_or(Error::<T>::NotFound)?;
 			task.take().ok_or(Error::<T>::NotFound)
 		})?;
+		Self::cleanup_agenda(when);
 		Self::deposit_event(Event::Canceled { when, index });
 		Self::place_task(new_time, task).map_err(|x| x.0)
 	}

substrate/frame/scheduler/src/migration.rs

@@ -198,6 +198,119 @@ pub mod v3 {
 	}
 }
 
+mod v4 {
+	use super::*;
+	use frame_support::pallet_prelude::*;
+
+	/// This migration cleans up empty agendas of the V4 scheduler.
+	///
+	/// This should be run on a scheduler that does not have
+	/// <https://github.com/paritytech/substrate/pull/12989> since it piles up `None`-only agendas. This does not modify the pallet version.
+	pub struct CleanupAgendas<T>(sp_std::marker::PhantomData<T>);
+
+	impl<T: Config> OnRuntimeUpgrade for CleanupAgendas<T> {
+		#[cfg(feature = "try-runtime")]
+		fn pre_upgrade() -> Result<Vec<u8>, &'static str> {
+			assert_eq!(
+				StorageVersion::get::<Pallet<T>>(),
+				4,
+				"Can only cleanup agendas of the V4 scheduler"
+			);
+
+			let agendas = Agenda::<T>::iter_keys().count();
+			let non_empty_agendas =
+				Agenda::<T>::iter_values().filter(|a| a.iter().any(|s| s.is_some())).count();
+			log::info!(
+				target: TARGET,
+				"There are {} total and {} non-empty agendas",
+				agendas,
+				non_empty_agendas
+			);
+
+			Ok((agendas as u32, non_empty_agendas as u32).encode())
+		}
+
+		fn on_runtime_upgrade() -> Weight {
+			let version = StorageVersion::get::<Pallet<T>>();
+			if version != 4 {
+				log::warn!(target: TARGET, "Skipping CleanupAgendas migration since it was run on the wrong version: {:?} != 4", version);
+				return T::DbWeight::get().reads(1)
+			}
+
+			let keys = Agenda::<T>::iter_keys().collect::<Vec<_>>();
+			let mut writes = 0;
+			for k in &keys {
+				let mut schedules = Agenda::<T>::get(k);
+				let all_schedules = schedules.len();
+				let suffix_none_schedules =
+					schedules.iter().rev().take_while(|s| s.is_none()).count();
+
+				match all_schedules.checked_sub(suffix_none_schedules) {
+					Some(0) => {
+						log::info!(
+							"Deleting None-only agenda {:?} with {} entries",
+							k,
+							all_schedules
+						);
+						Agenda::<T>::remove(k);
+						writes.saturating_inc();
+					},
+					Some(ne) if ne > 0 => {
+						log::info!(
+							"Removing {} schedules of {} from agenda {:?}, now {:?}",
+							suffix_none_schedules,
+							all_schedules,
+							ne,
+							k
+						);
+						schedules.truncate(ne);
+						Agenda::<T>::insert(k, schedules);
+						writes.saturating_inc();
+					},
+					Some(_) => {
+						frame_support::defensive!(
+							// Bad but let's not panic.
+							"Cannot have more None suffix schedules that schedules in total"
+						);
+					},
+					None => {
+						log::info!("Agenda {:?} does not have any None suffix schedules", k);
+					},
+				}
+			}
+
+			// We don't modify the pallet version.
+
+			T::DbWeight::get().reads_writes(1 + keys.len().saturating_mul(2) as u64, writes)
+		}
+
+		#[cfg(feature = "try-runtime")]
+		fn post_upgrade(state: Vec<u8>) -> Result<(), &'static str> {
+			assert_eq!(StorageVersion::get::<Pallet<T>>(), 4, "Version must not change");
+
+			let (old_agendas, non_empty_agendas): (u32, u32) =
+				Decode::decode(&mut state.as_ref()).expect("Must decode pre_upgrade state");
+			let new_agendas = Agenda::<T>::iter_keys().count() as u32;
+
+			match old_agendas.checked_sub(new_agendas) {
+				Some(0) => log::warn!(
+					target: TARGET,
+					"Did not clean up any agendas. v4::CleanupAgendas can be removed."
+				),
+				Some(n) =>
+					log::info!(target: TARGET, "Cleaned up {} agendas, now {}", n, new_agendas),
+				None => unreachable!(
+					"Number of agendas cannot increase, old {} new {}",
+					old_agendas, new_agendas
+				),
+			}
+			assert_eq!(new_agendas, non_empty_agendas, "Expected to keep all non-empty agendas");
+
+			Ok(())
+		}
+	}
+}
+
 #[cfg(test)]
 #[cfg(feature = "try-runtime")]
 mod test {
@@ -396,6 +509,64 @@ mod test {
 		});
 	}
 
+	#[test]
+	fn cleanup_agendas_works() {
+		use sp_core::bounded_vec;
+		new_test_ext().execute_with(|| {
+			StorageVersion::new(4).put::<Scheduler>();
+
+			let call = RuntimeCall::System(frame_system::Call::remark { remark: vec![] });
+			let bounded_call = Preimage::bound(call).unwrap();
+			let some = Some(ScheduledOf::<Test> {
+				maybe_id: None,
+				priority: 1,
+				call: bounded_call,
+				maybe_periodic: None,
+				origin: root(),
+				_phantom: Default::default(),
+			});
+
+			// Put some empty, and some non-empty agendas in there.
+			let test_data: Vec<(
+				BoundedVec<Option<ScheduledOf<Test>>, <Test as Config>::MaxScheduledPerBlock>,
+				Option<
+					BoundedVec<Option<ScheduledOf<Test>>, <Test as Config>::MaxScheduledPerBlock>,
+				>,
+			)> = vec![
+				(bounded_vec![some.clone()], Some(bounded_vec![some.clone()])),
+				(bounded_vec![None, some.clone()], Some(bounded_vec![None, some.clone()])),
+				(bounded_vec![None, some.clone(), None], Some(bounded_vec![None, some.clone()])),
+				(bounded_vec![some.clone(), None, None], Some(bounded_vec![some.clone()])),
+				(bounded_vec![None, None], None),
+				(bounded_vec![None, None, None], None),
+				(bounded_vec![], None),
+			];
+
+			// Insert all the agendas.
+			for (i, test) in test_data.iter().enumerate() {
+				Agenda::<Test>::insert(i as u64, test.0.clone());
+			}
+
+			// Run the migration.
+			let data = v4::CleanupAgendas::<Test>::pre_upgrade().unwrap();
+			let _w = v4::CleanupAgendas::<Test>::on_runtime_upgrade();
+			v4::CleanupAgendas::<Test>::post_upgrade(data).unwrap();
+
+			// Check that the post-state is correct.
+			for (i, test) in test_data.iter().enumerate() {
+				match test.1.clone() {
+					None => assert!(
+						!Agenda::<Test>::contains_key(i as u64),
+						"Agenda {} should be removed",
+						i
+					),
+					Some(new) =>
+						assert_eq!(Agenda::<Test>::get(i as u64), new, "Agenda wrong {}", i),
+				}
+			}
+		});
+	}
+
 	fn signed(i: u64) -> OriginCaller {
 		system::RawOrigin::Signed(i).into()
 	}

substrate/frame/scheduler/src/tests.rs

@@ -1291,40 +1291,6 @@ fn scheduler_v3_anon_reschedule_works() {
 	});
 }
 
-/// Cancelling a call and then scheduling a second call for the same
-/// block results in different addresses.
-#[test]
-fn scheduler_v3_anon_schedule_does_not_resuse_addr() {
-	use frame_support::traits::schedule::v3::Anon;
-	new_test_ext().execute_with(|| {
-		let call =
-			RuntimeCall::Logger(LoggerCall::log { i: 42, weight: Weight::from_ref_time(10) });
-
-		// Schedule both calls.
-		let addr_1 = <Scheduler as Anon<_, _, _>>::schedule(
-			DispatchTime::At(4),
-			None,
-			127,
-			root(),
-			Preimage::bound(call.clone()).unwrap(),
-		)
-		.unwrap();
-		// Cancel the call.
-		assert_ok!(<Scheduler as Anon<_, _, _>>::cancel(addr_1));
-		let addr_2 = <Scheduler as Anon<_, _, _>>::schedule(
-			DispatchTime::At(4),
-			None,
-			127,
-			root(),
-			Preimage::bound(call).unwrap(),
-		)
-		.unwrap();
-
-		// Should not re-use the address.
-		assert!(addr_1 != addr_2);
-	});
-}
-
 #[test]
 fn scheduler_v3_anon_next_schedule_time_works() {
 	use frame_support::traits::schedule::v3::Anon;
@@ -1531,45 +1497,6 @@ fn scheduler_v3_anon_reschedule_fills_holes() {
 	});
 }
 
-/// Re-scheduling into the same block produces a different address
-/// if there is still space in the agenda.
-#[test]
-fn scheduler_v3_anon_reschedule_does_not_resuse_addr_if_agenda_not_full() {
-	use frame_support::traits::schedule::v3::Anon;
-	let max: u32 = <Test as Config>::MaxScheduledPerBlock::get();
-	assert!(max > 1, "This test only makes sense for MaxScheduledPerBlock > 1");
-
-	new_test_ext().execute_with(|| {
-		let call =
-			RuntimeCall::Logger(LoggerCall::log { i: 42, weight: Weight::from_ref_time(10) });
-
-		// Schedule both calls.
-		let addr_1 = <Scheduler as Anon<_, _, _>>::schedule(
-			DispatchTime::At(4),
-			None,
-			127,
-			root(),
-			Preimage::bound(call.clone()).unwrap(),
-		)
-		.unwrap();
-		// Cancel the call.
-		assert_ok!(<Scheduler as Anon<_, _, _>>::cancel(addr_1));
-		let addr_2 = <Scheduler as Anon<_, _, _>>::schedule(
-			DispatchTime::At(5),
-			None,
-			127,
-			root(),
-			Preimage::bound(call).unwrap(),
-		)
-		.unwrap();
-		// Re-schedule `call` to block 4.
-		let addr_3 = <Scheduler as Anon<_, _, _>>::reschedule(addr_2, DispatchTime::At(4)).unwrap();
-
-		// Should not re-use the address.
-		assert!(addr_1 != addr_3);
-	});
-}
-
 /// The scheduler can be used as `v3::Named` trait.
 #[test]
 fn scheduler_v3_named_basic_works() {
@@ -1767,3 +1694,149 @@ fn scheduler_v3_named_next_schedule_time_works() {
 		);
 	});
 }
+
+#[test]
+fn cancel_last_task_removes_agenda() {
+	new_test_ext().execute_with(|| {
+		let when = 4;
+		let call =
+			RuntimeCall::Logger(LoggerCall::log { i: 42, weight: Weight::from_ref_time(10) });
+		let address = Scheduler::do_schedule(
+			DispatchTime::At(when),
+			None,
+			127,
+			root(),
+			Preimage::bound(call.clone()).unwrap(),
+		)
+		.unwrap();
+		let address2 = Scheduler::do_schedule(
+			DispatchTime::At(when),
+			None,
+			127,
+			root(),
+			Preimage::bound(call).unwrap(),
+		)
+		.unwrap();
+		// two tasks at agenda.
+		assert!(Agenda::<Test>::get(when).len() == 2);
+		assert_ok!(Scheduler::do_cancel(None, address));
+		// still two tasks at agenda, `None` and `Some`.
+		assert!(Agenda::<Test>::get(when).len() == 2);
+		// cancel last task from `when` agenda.
+		assert_ok!(Scheduler::do_cancel(None, address2));
+		// if all tasks `None`, agenda fully removed.
+		assert!(Agenda::<Test>::get(when).len() == 0);
+	});
+}
+
+#[test]
+fn cancel_named_last_task_removes_agenda() {
+	new_test_ext().execute_with(|| {
+		let when = 4;
+		let call =
+			RuntimeCall::Logger(LoggerCall::log { i: 42, weight: Weight::from_ref_time(10) });
+		Scheduler::do_schedule_named(
+			[1u8; 32],
+			DispatchTime::At(when),
+			None,
+			127,
+			root(),
+			Preimage::bound(call.clone()).unwrap(),
+		)
+		.unwrap();
+		Scheduler::do_schedule_named(
+			[2u8; 32],
+			DispatchTime::At(when),
+			None,
+			127,
+			root(),
+			Preimage::bound(call).unwrap(),
+		)
+		.unwrap();
+		// two tasks at agenda.
+		assert!(Agenda::<Test>::get(when).len() == 2);
+		assert_ok!(Scheduler::do_cancel_named(None, [2u8; 32]));
+		// removes trailing `None` and leaves one task.
+		assert!(Agenda::<Test>::get(when).len() == 1);
+		// cancel last task from `when` agenda.
+		assert_ok!(Scheduler::do_cancel_named(None, [1u8; 32]));
+		// if all tasks `None`, agenda fully removed.
+		assert!(Agenda::<Test>::get(when).len() == 0);
+	});
+}
+
+#[test]
+fn reschedule_last_task_removes_agenda() {
+	new_test_ext().execute_with(|| {
+		let when = 4;
+		let call =
+			RuntimeCall::Logger(LoggerCall::log { i: 42, weight: Weight::from_ref_time(10) });
+		let address = Scheduler::do_schedule(
+			DispatchTime::At(when),
+			None,
+			127,
+			root(),
+			Preimage::bound(call.clone()).unwrap(),
+		)
+		.unwrap();
+		let address2 = Scheduler::do_schedule(
+			DispatchTime::At(when),
+			None,
+			127,
+			root(),
+			Preimage::bound(call).unwrap(),
+		)
+		.unwrap();
+		// two tasks at agenda.
+		assert!(Agenda::<Test>::get(when).len() == 2);
+		assert_ok!(Scheduler::do_cancel(None, address));
+		// still two tasks at agenda, `None` and `Some`.
+		assert!(Agenda::<Test>::get(when).len() == 2);
+		// reschedule last task from `when` agenda.
+		assert_eq!(
+			Scheduler::do_reschedule(address2, DispatchTime::At(when + 1)).unwrap(),
+			(when + 1, 0)
+		);
+		// if all tasks `None`, agenda fully removed.
+		assert!(Agenda::<Test>::get(when).len() == 0);
+	});
+}
+
+#[test]
+fn reschedule_named_last_task_removes_agenda() {
+	new_test_ext().execute_with(|| {
+		let when = 4;
+		let call =
+			RuntimeCall::Logger(LoggerCall::log { i: 42, weight: Weight::from_ref_time(10) });
+		Scheduler::do_schedule_named(
+			[1u8; 32],
+			DispatchTime::At(when),
+			None,
+			127,
+			root(),
+			Preimage::bound(call.clone()).unwrap(),
+		)
+		.unwrap();
+		Scheduler::do_schedule_named(
+			[2u8; 32],
+			DispatchTime::At(when),
+			None,
+			127,
+			root(),
+			Preimage::bound(call).unwrap(),
+		)
+		.unwrap();
+		// two tasks at agenda.
+		assert!(Agenda::<Test>::get(when).len() == 2);
+		assert_ok!(Scheduler::do_cancel_named(None, [1u8; 32]));
+		// still two tasks at agenda, `None` and `Some`.
+		assert!(Agenda::<Test>::get(when).len() == 2);
+		// reschedule last task from `when` agenda.
+		assert_eq!(
+			Scheduler::do_reschedule_named([2u8; 32], DispatchTime::At(when + 1)).unwrap(),
+			(when + 1, 0)
+		);
+		// if all tasks `None`, agenda fully removed.
+		assert!(Agenda::<Test>::get(when).len() == 0);
+	});
+}