Use simpler strategy for storage cleanup

polkadot/runtime/rococo/src/lib.rs

@@ -1307,52 +1307,6 @@ impl pallet_asset_rate::Config for Runtime {
 	type BenchmarkHelper = runtime_common::impls::benchmarks::AssetRateArguments;
 }
 
-#[frame_support::pallet]
-pub mod im_online_remover {
-	use frame_support::pallet_prelude::*;
-	use frame_system::pallet_prelude::*;
-
-	#[pallet::pallet]
-	pub struct Pallet<T>(_);
-
-	#[pallet::config]
-	pub trait Config: frame_system::Config {}
-
-	#[pallet::hooks]
-	impl<T: Config> Hooks<BlockNumberFor<T>> for Pallet<T> {
-		fn on_initialize(n: BlockNumberFor<T>) -> Weight {
-			if RemoveAtBlock::<T>::get() == None {
-				RemoveAtBlock::<T>::set(Some(n));
-			}
-			Weight::zero()
-		}
-
-		fn offchain_worker(n: BlockNumberFor<T>) {
-			const DB_PREFIX: &[u8] = b"parity/im-online-heartbeat/";
-			if let Some(remove_at) = RemoveAtBlock::<T>::get() {
-				if remove_at == n {
-					let validator_set_size =
-						pallet_session::Pallet::<crate::Runtime>::validators().len() as u32;
-					(0..validator_set_size).for_each(|idx| {
-						let key = {
-							let mut key = DB_PREFIX.to_vec();
-							key.extend(idx.encode());
-							key
-						};
-						// FIXME: `StorageLock` needed?
-						sp_runtime::offchain::storage::StorageValueRef::persistent(&key).clear();
-					});
-				}
-			}
-		}
-	}
-
-	#[pallet::storage]
-	pub(super) type RemoveAtBlock<T: Config> = StorageValue<_, BlockNumberFor<T>, OptionQuery>;
-}
-
-impl im_online_remover::Config for Runtime {}
-
 construct_runtime! {
 	pub enum Runtime
 	{
@@ -1469,8 +1423,6 @@ construct_runtime! {
 		// Pallet for sending XCM.
 		XcmPallet: pallet_xcm::{Pallet, Call, Storage, Event<T>, Origin, Config<T>} = 99,
 
-		ImOnlineRemover: im_online_remover::{Pallet, Storage} = 100,
-
 		// Pallet for migrating Identity to a parachain. To be removed post-migration.
 		IdentityMigrator: identity_migrator::{Pallet, Call, Event<T>} = 248,
 
@@ -1768,7 +1720,6 @@ mod benches {
 		[runtime_parachains::paras_inherent, ParaInherent]
 		[runtime_parachains::paras, Paras]
 		[runtime_parachains::assigner_on_demand, OnDemandAssignmentProvider]
-		[im_online_remover, ImOnlineRemover]
 		// Substrate
 		[pallet_balances, Balances]
 		[pallet_balances, NisCounterpartBalances]
@@ -1865,6 +1816,12 @@ sp_api::impl_runtime_apis! {
 
 	impl offchain_primitives::OffchainWorkerApi<Block> for Runtime {
 		fn offchain_worker(header: &<Block as BlockT>::Header) {
+			use sp_runtime::{traits::Header, DigestItem};
+
+			if header.digest().logs().iter().find(|&di| *di == DigestItem::RuntimeEnvironmentUpdated).is_some() {
+				pallet_im_online::migration::clear_offchain_storage::<Runtime>()
+			}
+
 			Executive::offchain_worker(header)
 		}
 	}

substrate/frame/im-online/Cargo.toml

@@ -23,6 +23,7 @@ frame-benchmarking = { path = "../benchmarking", default-features = false, optio
 frame-support = { path = "../support", default-features = false }
 frame-system = { path = "../system", default-features = false }
 pallet-authorship = { path = "../authorship", default-features = false }
+pallet-session = { path = "../session", default-features = false }
 sp-application-crypto = { path = "../../primitives/application-crypto", default-features = false, features = ["serde"] }
 sp-core = { path = "../../primitives/core", default-features = false, features = ["serde"] }
 sp-io = { path = "../../primitives/io", default-features = false }

substrate/frame/im-online/src/migration.rs

@@ -116,6 +116,18 @@ pub mod v1 {
 	}
 }
 
+pub fn clear_offchain_storage<T: pallet_session::Config>() {
+	let validator_set_size = pallet_session::Pallet::<T>::validators().len() as u32;
+	(0..validator_set_size).for_each(|idx| {
+		let key = {
+			let mut key = DB_PREFIX.to_vec();
+			key.extend(idx.encode());
+			key
+		};
+		sp_runtime::offchain::storage::StorageValueRef::persistent(&key).clear();
+	});
+}
+
 #[cfg(all(feature = "try-runtime", test))]
 mod test {
 	use super::*;