change: Stake data source improvements (#500)

node/node/Cargo.toml

@@ -107,6 +107,7 @@ main-chain-follower-mock = { workspace = true, features = [
 ] }
 tokio = { workspace = true }
 cli-commands = { workspace = true }
+sp-stake-distribution = { workspace = true, features = ["std"] }
 
 [build-dependencies]
 substrate-build-script-utils = { workspace = true }

node/node/src/main_chain_follower.rs

@@ -3,16 +3,18 @@ use db_sync_follower::{
 	block::BlockDataSourceImpl, candidates::CandidatesDataSourceImpl,
 	mc_hash::McHashDataSourceImpl, metrics::McFollowerMetrics,
 	native_token::NativeTokenManagementDataSourceImpl, sidechain_rpc::SidechainRpcDataSourceImpl,
+	stake_distribution::StakeDistributionDataSourceImpl,
 };
 use main_chain_follower_mock::{
 	block::BlockDataSourceMock, candidate::AuthoritySelectionDataSourceMock,
 	mc_hash::McHashDataSourceMock, native_token::NativeTokenDataSourceMock,
-	sidechain_rpc::SidechainRpcDataSourceMock,
+	sidechain_rpc::SidechainRpcDataSourceMock, stake_distribution::StakeDistributionDataSourceMock,
 };
 use pallet_sidechain_rpc::SidechainRpcDataSource;
 use sc_service::error::Error as ServiceError;
 use sidechain_mc_hash::McHashDataSource;
 use sp_native_token_management::NativeTokenManagementDataSource;
+use sp_stake_distribution::StakeDistributionDataSource;
 use std::{error::Error, sync::Arc};
 
 #[derive(Clone)]
@@ -21,6 +23,7 @@ pub struct DataSources {
 	pub authority_selection: Arc<dyn AuthoritySelectionDataSource + Send + Sync>,
 	pub native_token: Arc<dyn NativeTokenManagementDataSource + Send + Sync>,
 	pub sidechain_rpc: Arc<dyn SidechainRpcDataSource + Send + Sync>,
+	pub stake_distribution: Arc<dyn StakeDistributionDataSource + Send + Sync>,
 }
 
 pub(crate) async fn create_cached_main_chain_follower_data_sources(
@@ -57,10 +60,12 @@ pub fn create_mock_data_sources(
 		mc_hash: Arc::new(McHashDataSourceMock::new(block)),
 		authority_selection: Arc::new(AuthoritySelectionDataSourceMock::new_from_env()?),
 		native_token: Arc::new(NativeTokenDataSourceMock::new()),
+		stake_distribution: Arc::new(StakeDistributionDataSourceMock::new()),
 	})
 }
 
 pub const CANDIDATES_FOR_EPOCH_CACHE_SIZE: usize = 64;
+pub const STAKE_CACHE_SIZE: usize = 100;
 
 pub async fn create_cached_data_sources(
 	metrics_opt: Option<McFollowerMetrics>,
@@ -80,8 +85,13 @@ pub async fn create_cached_data_sources(
 				.cached(CANDIDATES_FOR_EPOCH_CACHE_SIZE)?,
 		),
 		native_token: Arc::new(NativeTokenManagementDataSourceImpl::new_from_env(
+			pool.clone(),
+			metrics_opt.clone(),
+		)?),
+		stake_distribution: Arc::new(StakeDistributionDataSourceImpl::new(
 			pool,
 			metrics_opt,
-		)?),
+			STAKE_CACHE_SIZE,
+		)),
 	})
 }

toolkit/mainchain-follower/db-sync-follower/src/stake_distribution/mod.rs

@@ -30,7 +30,7 @@ impl StakeDistributionDataSource for StakeDistributionDataSourceImpl {
 		pool_hash: MainchainKeyHash,
 	) -> Result<PoolDelegation, Box<dyn std::error::Error + Send + Sync>> {
 		Ok(self
-			.get_stake_pool_delegation_distribution_for_pools(epoch, vec![pool_hash])
+			.get_stake_pool_delegation_distribution_for_pools(epoch, &[pool_hash])
 			.await?
 			.0
 			.entry(pool_hash)
@@ -41,15 +41,15 @@ impl StakeDistributionDataSource for StakeDistributionDataSourceImpl {
 	async fn get_stake_pool_delegation_distribution_for_pools(
 		&self,
 		epoch: McEpochNumber,
-		pool_hashes: Vec<MainchainKeyHash>,
+		pool_hashes: &[MainchainKeyHash],
 	) -> Result<StakeDistribution, Box<dyn std::error::Error + Send + Sync>> {
 		let mut pool_hashes_to_query = Vec::<[u8; 28]>::new();
 		let mut stake_distribution = BTreeMap::<MainchainKeyHash, PoolDelegation>::new();
 
 		for pool_hash in pool_hashes {
 			match self.cache.get_distribution_for_pool(epoch, pool_hash) {
 				Some(pool_delegation) => {
-					stake_distribution.insert(pool_hash, pool_delegation);
+					stake_distribution.insert(pool_hash.clone(), pool_delegation);
 				},
 				None => pool_hashes_to_query.push(pool_hash.0),
 			}
@@ -121,10 +121,10 @@ impl Cache {
 	fn get_distribution_for_pool(
 		&self,
 		epoch: McEpochNumber,
-		pool_hash: MainchainKeyHash,
+		pool_hash: &MainchainKeyHash,
 	) -> Option<PoolDelegation> {
 		if let Ok(mut cache) = self.distribution_per_pool_cache.lock() {
-			cache.get(&(epoch, pool_hash)).map(|e| e.clone())
+			cache.get(&(epoch, *pool_hash)).map(|e| e.clone())
 		} else {
 			None
 		}

toolkit/mainchain-follower/db-sync-follower/src/stake_distribution/tests.rs

@@ -42,7 +42,7 @@ async fn stake_pool_delegation_distribution_for_pools_works(pool: PgPool) {
 	let distribution = make_source(pool)
 		.get_stake_pool_delegation_distribution_for_pools(
 			epoch,
-			vec![stake_pool_key_hash_1(), stake_pool_key_hash_2()],
+			&[stake_pool_key_hash_1(), stake_pool_key_hash_2()],
 		)
 		.await
 		.unwrap()

toolkit/mainchain-follower/mock/Cargo.toml

@@ -17,6 +17,7 @@ sidechain-domain = { workspace = true }
 sp-timestamp = { workspace = true }
 sp-native-token-management = { workspace = true, optional = true }
 sidechain-mc-hash = { workspace = true, optional = true }
+sp-stake-distribution = { workspace = true, optional = true }
 pallet-sidechain-rpc = { workspace = true, optional = true }
 authority-selection-inherents = { workspace = true, optional = true }
 
@@ -27,10 +28,11 @@ std = [
 	"sidechain-domain/std",
 	"rand/std",
 	"sp-native-token-management?/std",
+	"sp-stake-distribution?/std"
 ]
 block-source = []
 candidate-source = ["authority-selection-inherents"]
 native-token = ["sp-native-token-management"]
 mc-hash = ["sidechain-mc-hash"]
 sidechain-rpc = ["pallet-sidechain-rpc"]
-stake-distribution = []
+stake-distribution = ["sp-stake-distribution"]

toolkit/mainchain-follower/mock/src/lib.rs

@@ -11,6 +11,8 @@ pub mod mc_hash;
 pub mod native_token;
 #[cfg(feature = "sidechain-rpc")]
 pub mod sidechain_rpc;
+#[cfg(feature = "stake-distribution")]
+pub mod stake_distribution;
 
 #[allow(unused)]
 pub(crate) struct UnimplementedMocks;

toolkit/mainchain-follower/mock/src/stake_distribution.rs

@@ -0,0 +1,29 @@
+use sidechain_domain::*;
+use sp_stake_distribution::StakeDistributionDataSource;
+
+pub struct StakeDistributionDataSourceMock;
+
+impl StakeDistributionDataSourceMock {
+	pub fn new() -> Self {
+		Self
+	}
+}
+
+#[async_trait::async_trait]
+impl StakeDistributionDataSource for StakeDistributionDataSourceMock {
+	async fn get_stake_pool_delegation_distribution_for_pool(
+		&self,
+		_epoch: McEpochNumber,
+		_pool_hash: MainchainKeyHash,
+	) -> Result<PoolDelegation, Box<dyn std::error::Error + Send + Sync>> {
+		Ok(PoolDelegation::default())
+	}
+
+	async fn get_stake_pool_delegation_distribution_for_pools(
+		&self,
+		_epoch: McEpochNumber,
+		_pool_hashes: &[MainchainKeyHash],
+	) -> Result<StakeDistribution, Box<dyn std::error::Error + Send + Sync>> {
+		Ok(StakeDistribution::default())
+	}
+}

toolkit/primitives/domain/src/lib.rs

@@ -11,8 +11,10 @@ extern crate core;
 extern crate num_derive;
 
 pub use alloc::collections::btree_map::BTreeMap;
+#[cfg(feature = "std")]
+use alloc::format;
 pub use alloc::vec::Vec;
-use alloc::{format, str::FromStr, string::String, string::ToString, vec};
+use alloc::{str::FromStr, string::String, string::ToString, vec};
 use byte_string_derive::byte_string;
 use core::{
 	fmt::{Display, Formatter},

toolkit/primitives/domain/src/mainchain_epoch.rs

@@ -17,6 +17,7 @@ pub struct MainchainEpochConfig {
 	pub slot_duration_millis: Duration,
 }
 
+#[cfg(feature = "std")]
 fn default_slot_duration() -> Duration {
 	Duration::from_millis(1000)
 }

toolkit/primitives/native-token-management/src/lib.rs

@@ -1,14 +1,15 @@
 #![cfg_attr(not(feature = "std"), no_std)]
 
-use core::str::FromStr;
+#[cfg(feature = "std")]
+use {core::str::FromStr, sp_runtime::traits::Block as BlockT};
 
 #[cfg(feature = "std")]
 pub use inherent_provider::*;
 
 use parity_scale_codec::{Decode, Encode, MaxEncodedLen};
 use sidechain_domain::*;
 use sp_inherents::*;
-use sp_runtime::{scale_info::TypeInfo, traits::Block as BlockT};
+use sp_runtime::scale_info::TypeInfo;
 
 #[cfg(test)]
 mod tests;

toolkit/primitives/session-manager/src/pallet_session_support.rs

@@ -1,5 +1,3 @@
-#![cfg_attr(not(feature = "std"), no_std)]
-
 use core::marker::PhantomData;
 use derive_new::new;
 use frame_system::pallet_prelude::BlockNumberFor;

toolkit/primitives/session-validator-management/src/lib.rs

@@ -1,5 +1,6 @@
 #![cfg_attr(not(feature = "std"), no_std)]
 
+#[cfg(feature = "std")]
 use core::str::FromStr;
 
 use scale_info::TypeInfo;

toolkit/primitives/stake-distribution/src/lib.rs

@@ -15,6 +15,6 @@ pub trait StakeDistributionDataSource {
 	async fn get_stake_pool_delegation_distribution_for_pools(
 		&self,
 		epoch: McEpochNumber,
-		pool_hashes: Vec<MainchainKeyHash>,
+		pool_hashes: &[MainchainKeyHash],
 	) -> Result<StakeDistribution, Box<dyn std::error::Error + Send + Sync>>;
 }