Adjust session consumer part.1 (#229)

tool/state-processor/src/main.rs

@@ -18,6 +18,7 @@ mod evm;
 mod identity;
 mod indices;
 mod proxy;
+mod session;
 mod staking;
 mod system;
 mod vesting;
@@ -32,9 +33,10 @@ fn main() -> Result<()> {
 	std::env::set_var("RUST_LOG", "state_processor");
 	pretty_env_logger::init();
 
+	// <Processor<Pangolin>>::new()?.process()?;
+	<Processor<Pangolin>>::new()?.test().process()?;
 	// <Processor<Darwinia>>::new()?.process()?;
 	// <Processor<Crab>>::new()?.process()?;
-	<Processor<Pangolin>>::new()?.test().process()?;
 
 	Ok(())
 }

tool/state-processor/src/processor.rs

@@ -11,14 +11,15 @@ use std::{
 use crate::*;
 // crates.io
 use anyhow::Result;
-use fxhash::FxHashMap;
+use fxhash::{FxHashMap, FxHashSet};
 use once_cell::sync::Lazy;
 use parity_scale_codec::{Decode, Encode};
 use serde::de::DeserializeOwned;
 use serde_json::Value;
 // hack-ink
 use subspector::ChainSpec;
 
+pub type Set = FxHashSet<String>;
 pub type Map<V> = FxHashMap<String, V>;
 
 pub static NOW: Lazy<RwLock<u32>> = Lazy::new(|| RwLock::new(0));
@@ -112,6 +113,32 @@ impl<R> State<R> {
 		})
 	}
 
+	pub fn get_value<D>(&self, pallet: &[u8], item: &[u8], hash: &str, value: &mut D) -> &Self
+	where
+		D: Decode,
+	{
+		let key = full_key(pallet, item, hash);
+
+		if let Some(v) = self.map.get(&key) {
+			match decode(v) {
+				Ok(v) => *value = v,
+				Err(e) => log::error!(
+					"failed to decode `{}::{}::{hash}({v})`, due to `{e}`",
+					String::from_utf8_lossy(pallet),
+					String::from_utf8_lossy(item),
+				),
+			}
+		} else {
+			log::error!(
+				"key not found `{}::{}::{hash}`",
+				String::from_utf8_lossy(pallet),
+				String::from_utf8_lossy(item),
+			);
+		}
+
+		self
+	}
+
 	pub fn insert_raw_key_raw_value(&mut self, key: String, value: String) -> &mut Self {
 		self.map.insert(key, value);
 
@@ -127,6 +154,64 @@ impl<R> State<R> {
 		self
 	}
 
+	pub fn insert_raw_key_map(&mut self, pairs: Map<String>) -> &mut Self {
+		pairs.into_iter().for_each(|(k, v)| {
+			if self.map.contains_key(&k) {
+				log::error!("key({k}) has already existed, overriding");
+			}
+
+			self.map.insert(k, v);
+		});
+
+		self
+	}
+
+	pub fn insert_value<E>(&mut self, pallet: &[u8], item: &[u8], hash: &str, value: E) -> &mut Self
+	where
+		E: Encode,
+	{
+		self.map.insert(full_key(pallet, item, hash), encode_value(value));
+
+		self
+	}
+
+	pub fn insert_map<E, F>(&mut self, pairs: Map<E>, process_key: F) -> &mut Self
+	where
+		E: Encode,
+		F: Fn(&str) -> String,
+	{
+		pairs.into_iter().for_each(|(k, v)| {
+			self.map.insert(process_key(&k), encode_value(v));
+		});
+
+		self
+	}
+
+	pub fn contains_key(&self, key: &str) -> bool {
+		self.map.contains_key(key)
+	}
+
+	pub fn exists(&self, pallet: &[u8], item: &[u8]) -> bool {
+		self.map.keys().into_iter().any(|k| k.starts_with(&item_key(pallet, item)))
+	}
+
+	pub fn unreserve<A>(&mut self, account_id_32: A, amount: u128)
+	where
+		A: AsRef<[u8]>,
+	{
+		let account_id_32 = account_id_32.as_ref();
+		let (p, i, h) = if is_evm_address(account_id_32) {
+			(&b"System"[..], &b"Account"[..], &account_id_32[11..31])
+		} else {
+			(&b"AccountMigration"[..], &b"Accounts"[..], account_id_32)
+		};
+
+		self.mutate_value(p, i, &blake2_128_concat_to_string(h), |a: &mut AccountInfo| {
+			a.data.free += amount;
+			a.data.reserved -= amount;
+		});
+	}
+
 	pub fn take_raw_map<F>(
 		&mut self,
 		prefix: &str,
@@ -149,40 +234,19 @@ impl<R> State<R> {
 		self
 	}
 
-	pub fn insert_raw_key_map(&mut self, pairs: Map<String>) -> &mut Self {
-		pairs.into_iter().for_each(|(k, v)| {
-			if self.map.contains_key(&k) {
-				log::error!("key({k}) has already existed, overriding");
-			}
-
-			self.map.insert(k, v);
-		});
-
-		self
-	}
-
-	pub fn get_value<D>(&self, pallet: &[u8], item: &[u8], hash: &str, value: &mut D) -> &Self
+	pub fn take_keys<F>(&mut self, prefix: &str, keys: &mut Set, process_key: F) -> &mut Self
 	where
-		D: Decode,
+		F: Fn(&str, &str) -> String,
 	{
-		let key = full_key(pallet, item, hash);
+		self.map.retain(|k, _| {
+			if k.starts_with(prefix) {
+				keys.insert(process_key(k, prefix));
 
-		if let Some(v) = self.map.get(&key) {
-			match decode(v) {
-				Ok(v) => *value = v,
-				Err(e) => log::error!(
-					"failed to decode `{}::{}::{hash}({v})`, due to `{e}`",
-					String::from_utf8_lossy(pallet),
-					String::from_utf8_lossy(item),
-				),
+				false
+			} else {
+				true
 			}
-		} else {
-			log::error!(
-				"key not found `{}::{}::{hash}`",
-				String::from_utf8_lossy(pallet),
-				String::from_utf8_lossy(item),
-			);
-		}
+		});
 
 		self
 	}
@@ -219,31 +283,6 @@ impl<R> State<R> {
 		self
 	}
 
-	pub fn insert_value<E>(&mut self, pallet: &[u8], item: &[u8], hash: &str, value: E) -> &mut Self
-	where
-		E: Encode,
-	{
-		self.map.insert(full_key(pallet, item, hash), encode_value(value));
-
-		self
-	}
-
-	pub fn mutate_value<D, F>(&mut self, pallet: &[u8], item: &[u8], hash: &str, f: F) -> &mut Self
-	where
-		D: Default + Encode + Decode,
-		F: FnOnce(&mut D),
-	{
-		let mut v = D::default();
-
-		self.get_value(pallet, item, hash, &mut v);
-
-		f(&mut v);
-
-		self.insert_value(pallet, item, hash, v);
-
-		self
-	}
-
 	pub fn take_map<D, F>(
 		&mut self,
 		pallet: &[u8],
@@ -284,41 +323,20 @@ impl<R> State<R> {
 		self
 	}
 
-	pub fn insert_map<E, F>(&mut self, pairs: Map<E>, process_key: F) -> &mut Self
+	pub fn mutate_value<D, F>(&mut self, pallet: &[u8], item: &[u8], hash: &str, f: F) -> &mut Self
 	where
-		E: Encode,
-		F: Fn(&str) -> String,
+		D: Default + Encode + Decode,
+		F: FnOnce(&mut D),
 	{
-		pairs.into_iter().for_each(|(k, v)| {
-			self.map.insert(process_key(&k), encode_value(v));
-		});
-
-		self
-	}
+		let mut v = D::default();
 
-	pub fn contains_key(&self, key: &str) -> bool {
-		self.map.contains_key(key)
-	}
+		self.get_value(pallet, item, hash, &mut v);
 
-	pub fn exists(&self, pallet: &[u8], item: &[u8]) -> bool {
-		self.map.keys().into_iter().any(|k| k.starts_with(&item_key(pallet, item)))
-	}
+		f(&mut v);
 
-	pub fn unreserve<A>(&mut self, account_id_32: A, amount: u128)
-	where
-		A: AsRef<[u8]>,
-	{
-		let account_id_32 = account_id_32.as_ref();
-		let (p, i, h) = if is_evm_address(account_id_32) {
-			(&b"System"[..], &b"Account"[..], &account_id_32[11..31])
-		} else {
-			(&b"AccountMigration"[..], &b"Accounts"[..], account_id_32)
-		};
+		self.insert_value(pallet, item, hash, v);
 
-		self.mutate_value(p, i, &blake2_128_concat_to_string(h), |a: &mut AccountInfo| {
-			a.data.free += amount;
-			a.data.reserved -= amount;
-		});
+		self
 	}
 }
 

tool/state-processor/src/session/mod.rs

@@ -0,0 +1,23 @@
+// darwinia
+use crate::*;
+
+impl<S> Processor<S> {
+	pub fn process_session(&mut self) -> (Set, Set) {
+		// Session storage items.
+		// https://github.dev/darwinia-network/substrate/blob/darwinia-v0.12.5/frame/session/src/lib.rs#L532
+
+		let prefix = item_key(b"Session", b"NextKeys");
+
+		let mut s_keys = Set::default();
+		let mut p_keys = Set::default();
+
+		self.solo_state.take_keys(&prefix, &mut s_keys, |k, _| {
+			blake2_128_concat_to_string(key_to_account32(k))
+		});
+		self.para_state.take_keys(&prefix, &mut p_keys, |k, _| {
+			blake2_128_concat_to_string(key_to_account32(k))
+		});
+
+		(s_keys, p_keys)
+	}
+}

tool/state-processor/src/system/README.md

@@ -11,6 +11,7 @@
   - take para `Balances::TotalIssuance` and `Balances::Locks`
     - there should not be any locks on parachain
   - check if there are any other locks
+- decrease pallet-session account references
 - use all previous data to build the new accounts and calculate total issuances
   - merge solo and para account infos
     - how to deal with the account references? - TODO

tool/state-processor/src/system/mod.rs

@@ -24,13 +24,31 @@ where
 	pub fn process_system(&mut self) -> &mut Self {
 		// System storage items.
 		// https://github.dev/darwinia-network/substrate/blob/darwinia-v0.12.5/frame/system/src/lib.rs#L545
-		let solo_account_infos = self.process_solo_account_infos();
-		let para_account_infos = self.process_para_account_infos();
+		let mut solo_account_infos = self.process_solo_account_infos();
+		let mut para_account_infos = self.process_para_account_infos();
 		let (ring_total_issuance_storage, kton_total_issuance_storage) = self.process_balances();
+		let (solo_validators, para_collators) = self.process_session();
 		let mut accounts = Map::default();
 		let mut ring_total_issuance = u128::default();
 		let mut kton_total_issuance = u128::default();
 
+		// Skip for testnets, due to https://github.com/paritytech/substrate/issues/13172.
+		// log::info!("decrease solo pallet-session account references");
+		// solo_validators.into_iter().for_each(|k| {
+		// 	if let Some(a) = solo_account_infos.get_mut(&k) {
+		// 		a.consumers -= 1;
+		// 	}
+		// });
+
+		// Skip, due to https://github.com/paritytech/substrate/issues/13172.
+		// log::info!("decrease para pallet-session account references");
+		// para_collators.into_iter().for_each(|k| {
+		// 	if let Some(a) = para_account_infos.get_mut(&k) {
+		// 		dbg!(get_last_64(&k));
+		// 		a.consumers -= 1;
+		// 	}
+		// });
+
 		log::info!("build accounts");
 		log::info!("calculate total issuance");
 		solo_account_infos.into_iter().for_each(|(k, v)| {

tool/state-processor/src/util.rs

@@ -47,7 +47,7 @@ pub fn get_hashed_key(full_key: &str, item_key: &str) -> String {
 	full_key.trim_start_matches(item_key).into()
 }
 
-// twox128(pallet) + twox128(item) + *(item_key) -> account_id_32
+// twox128(pallet) + twox128(item) + *_concat(account_id_32) -> account_id_32
 pub fn get_last_64_key(key: &str, _: &str) -> String {
 	get_last_64(key)
 }
@@ -73,6 +73,10 @@ pub fn is_evm_address(address: &[u8]) -> bool {
 		&& address[1..31].iter().fold(address[0], |checksum, &b| checksum ^ b) == address[31]
 }
 
+pub fn key_to_account32(key: &str) -> [u8; 32] {
+	array_bytes::hex2array_unchecked(get_last_64(key))
+}
+
 pub fn build_spec(chain: &str) -> Result<()> {
 	log::info!("build {chain} spec");