pallet-scheduler: Ensure we request a preimage

frame/scheduler/src/lib.rs

@@ -777,6 +777,8 @@ impl<T: Config> Pallet<T> {
 	) -> Result<TaskAddress<T::BlockNumber>, DispatchError> {
 		let when = Self::resolve_time(when)?;
 
+		let lookup_hash = call.lookup_hash();
+
 		// sanitize maybe_periodic
 		let maybe_periodic = maybe_periodic
 			.filter(|p| p.1 > 1 && !p.0.is_zero())
@@ -790,7 +792,14 @@ impl<T: Config> Pallet<T> {
 			origin,
 			_phantom: PhantomData,
 		};
-		Self::place_task(when, task).map_err(|x| x.0)
+		let res = Self::place_task(when, task).map_err(|x| x.0)?;
+
+		if let Some(hash) = lookup_hash {
+			// Request the call to be made available.
+			T::Preimages::request(&hash);
+		}
+
+		Ok(res)
 	}
 
 	fn do_cancel(
@@ -862,6 +871,8 @@ impl<T: Config> Pallet<T> {
 
 		let when = Self::resolve_time(when)?;
 
+		let lookup_hash = call.lookup_hash();
+
 		// sanitize maybe_periodic
 		let maybe_periodic = maybe_periodic
 			.filter(|p| p.1 > 1 && !p.0.is_zero())
@@ -876,7 +887,14 @@ impl<T: Config> Pallet<T> {
 			origin,
 			_phantom: Default::default(),
 		};
-		Self::place_task(when, task).map_err(|x| x.0)
+		let res = Self::place_task(when, task).map_err(|x| x.0)?;
+
+		if let Some(hash) = lookup_hash {
+			// Request the call to be made available.
+			T::Preimages::request(&hash);
+		}
+
+		Ok(res)
 	}
 
 	fn do_cancel_named(origin: Option<T::PalletsOrigin>, id: TaskName) -> DispatchResult {
@@ -1027,6 +1045,7 @@ impl<T: Config> Pallet<T> {
 		} else {
 			Agenda::<T>::remove(when);
 		}
+
 		postponed == 0
 	}
 

frame/scheduler/src/tests.rs

@@ -58,9 +58,10 @@ fn scheduling_with_preimages_works() {
 			RuntimeCall::Logger(LoggerCall::log { i: 42, weight: Weight::from_ref_time(10) });
 		let hash = <Test as frame_system::Config>::Hashing::hash_of(&call);
 		let len = call.using_encoded(|x| x.len()) as u32;
-		let hashed = Preimage::pick(hash, len);
-		assert_ok!(Preimage::note_preimage(RuntimeOrigin::signed(0), call.encode()));
+		// Important to use here `Bounded::Lookup` to ensure that we request the hash.
+		let hashed = Bounded::Lookup { hash, len };
 		assert_ok!(Scheduler::do_schedule(DispatchTime::At(4), None, 127, root(), hashed));
+		assert_ok!(Preimage::note_preimage(RuntimeOrigin::signed(0), call.encode()));
 		assert!(Preimage::is_requested(&hash));
 		run_to_block(3);
 		assert!(logger::log().is_empty());
@@ -479,8 +480,9 @@ fn scheduler_handles_periodic_unavailable_preimage() {
 		let call = RuntimeCall::Logger(LoggerCall::log { i: 42, weight: (max_weight / 3) * 2 });
 		let hash = <Test as frame_system::Config>::Hashing::hash_of(&call);
 		let len = call.using_encoded(|x| x.len()) as u32;
-		let bound = Preimage::pick(hash, len);
-		assert_ok!(Preimage::note(call.encode().into()));
+		let bound = Bounded::Lookup { hash, len };
+		// The preimage isn't requested yet.
+		assert!(!Preimage::is_requested(&hash));
 
 		assert_ok!(Scheduler::do_schedule(
 			DispatchTime::At(4),
@@ -489,19 +491,23 @@ fn scheduler_handles_periodic_unavailable_preimage() {
 			root(),
 			bound.clone(),
 		));
+
+		// The preimage is requested.
+		assert!(Preimage::is_requested(&hash));
+
+		// Note the preimage.
+		assert_ok!(Preimage::note_preimage(RuntimeOrigin::signed(1), call.encode()));
+
 		// Executes 1 times till block 4.
 		run_to_block(4);
 		assert_eq!(logger::log().len(), 1);
 
-		// Unnote the preimage.
+		// Unnote the preimage
 		Preimage::unnote(&hash);
 
 		// Does not ever execute again.
 		run_to_block(100);
 		assert_eq!(logger::log().len(), 1);
-
-		// The preimage is not requested anymore.
-		assert!(!Preimage::is_requested(&hash));
 	});
 }
 
@@ -1129,7 +1135,8 @@ fn postponed_named_task_cannot_be_rescheduled() {
 			RuntimeCall::Logger(LoggerCall::log { i: 42, weight: Weight::from_ref_time(1000) });
 		let hash = <Test as frame_system::Config>::Hashing::hash_of(&call);
 		let len = call.using_encoded(|x| x.len()) as u32;
-		let hashed = Preimage::pick(hash, len);
+		// Important to use here `Bounded::Lookup` to ensure that we request the hash.
+		let hashed = Bounded::Lookup { hash, len };
 		let name: [u8; 32] = hash.as_ref().try_into().unwrap();
 
 		let address = Scheduler::do_schedule_named(

frame/support/src/traits/preimages.rs

@@ -26,6 +26,9 @@ use sp_std::borrow::Cow;
 pub type Hash = H256;
 pub type BoundedInline = crate::BoundedVec<u8, ConstU32<128>>;
 
+/// The maximum we expect a single legacy hash lookup to be.
+const MAX_LEGACY_LEN: u32 = 1_000_000;
+
 #[derive(
 	Encode, Decode, MaxEncodedLen, Clone, Eq, PartialEq, scale_info::TypeInfo, RuntimeDebug,
 )]
@@ -67,20 +70,25 @@ impl<T> Bounded<T> {
 	/// Returns the hash of the preimage.
 	///
 	/// The hash is re-calculated every time if the preimage is inlined.
-	pub fn hash(&self) -> H256 {
+	pub fn hash(&self) -> Hash {
 		use Bounded::*;
 		match self {
-			Legacy { hash, .. } => *hash,
+			Lookup { hash, .. } | Legacy { hash, .. } => *hash,
 			Inline(x) => blake2_256(x.as_ref()).into(),
-			Lookup { hash, .. } => *hash,
 		}
 	}
-}
 
-// The maximum we expect a single legacy hash lookup to be.
-const MAX_LEGACY_LEN: u32 = 1_000_000;
+	/// Returns the hash to lookup the preimage.
+	///
+	/// If this is a `Bounded::Inline`, `None` is returned as no lookup is required.
+	pub fn lookup_hash(&self) -> Option<Hash> {
+		use Bounded::*;
+		match self {
+			Lookup { hash, .. } | Legacy { hash, .. } => Some(*hash),
+			Inline(_) => None,
+		}
+	}
 
-impl<T> Bounded<T> {
 	/// Returns the length of the preimage or `None` if the length is unknown.
 	pub fn len(&self) -> Option<u32> {
 		match self {
@@ -168,8 +176,11 @@ pub trait QueryPreimage {
 		}
 	}
 
-	/// Create a `Bounded` instance based on the `hash` and `len` of the encoded value. This may not
-	/// be `peek`-able or `realize`-able.
+	/// Create a `Bounded` instance based on the `hash` and `len` of the encoded value.
+	///
+	/// It also directly requests the given `hash` using [`Self::request`].
+	///
+	/// This may not be `peek`-able or `realize`-able.
 	fn pick<T>(hash: Hash, len: u32) -> Bounded<T> {
 		Self::request(&hash);
 		Bounded::Lookup { hash, len }
@@ -228,10 +239,12 @@ pub trait StorePreimage: QueryPreimage {
 		Self::unrequest(hash)
 	}
 
-	/// Convert an otherwise unbounded or large value into a type ready for placing in storage. The
-	/// result is a type whose `MaxEncodedLen` is 131 bytes.
+	/// Convert an otherwise unbounded or large value into a type ready for placing in storage.
+	///
+	/// The result is a type whose `MaxEncodedLen` is 131 bytes.
 	///
 	/// NOTE: Once this API is used, you should use either `drop` or `realize`.
+	/// The value is also noted using [`Self::note`].
 	fn bound<T: Encode>(t: T) -> Result<Bounded<T>, DispatchError> {
 		let data = t.encode();
 		let len = data.len() as u32;