add GovernedMap cached DataSource implementation

Author: jankun4

demo/node/src/data_sources.rs

@@ -2,7 +2,7 @@ use authority_selection_inherents::authority_selection_inputs::AuthoritySelectio
 use pallet_sidechain_rpc::SidechainRpcDataSource;
 use partner_chains_db_sync_data_sources::{
 	block::BlockDataSourceImpl, candidates::CandidatesDataSourceImpl,
-	governed_map::GovernedMapDataSourceImpl, mc_hash::McHashDataSourceImpl,
+	governed_map::GovernedMapDataSourceCachedImpl, mc_hash::McHashDataSourceImpl,
 	metrics::McFollowerMetrics, native_token::NativeTokenManagementDataSourceImpl,
 	sidechain_rpc::SidechainRpcDataSourceImpl, stake_distribution::StakeDistributionDataSourceImpl,
 };
@@ -69,7 +69,6 @@ pub fn create_mock_data_sources()
 
 pub const CANDIDATES_FOR_EPOCH_CACHE_SIZE: usize = 64;
 pub const STAKE_CACHE_SIZE: usize = 100;
-
 pub async fn create_cached_db_sync_data_sources(
 	metrics_opt: Option<McFollowerMetrics>,
 ) -> Result<DataSources, Box<dyn Error + Send + Sync + 'static>> {
@@ -96,6 +95,9 @@ pub async fn create_cached_db_sync_data_sources(
 			metrics_opt.clone(),
 			STAKE_CACHE_SIZE,
 		)),
-		governed_map: Arc::new(GovernedMapDataSourceImpl::new(pool, metrics_opt)),
+		governed_map: Arc::new(GovernedMapDataSourceCachedImpl::new_from_env(
+			pool,
+			metrics_opt.clone(),
+		)?),
 	})
 }

toolkit/data-sources/db-sync/src/db_model.rs

@@ -114,6 +114,14 @@ pub(crate) struct DatumOutput {
 	pub datum: DbDatum,
 }
 
+#[derive(Debug, Clone, sqlx::FromRow, PartialEq)]
+pub(crate) struct DatumChangeOutput {
+	pub datum: DbDatum,
+	pub block_no: BlockNumber,
+	pub block_index: TxIndexInBlock,
+	pub action: String,
+}
+
 #[derive(Debug, Clone, PartialEq)]
 pub(crate) struct NativeTokenAmount(pub u128);
 impl From<NativeTokenAmount> for sidechain_domain::NativeTokenAmount {
@@ -367,6 +375,53 @@ pub(crate) async fn get_token_utxo_for_epoch(
 		.await?)
 }
 
+#[cfg(feature = "governed-map")]
+pub(crate) async fn get_changes(
+	pool: &Pool<Postgres>,
+	address: &Address,
+	last_block: Option<BlockNumber>,
+	new_block: BlockNumber,
+	asset: Asset,
+) -> Result<Vec<DatumChangeOutput>, SqlxError> {
+	let query = "
+		(SELECT
+			datum.value as datum, origin_block.block_no as block_no, origin_tx.block_index as block_index, 'upsert' as action
+		FROM tx_out
+		INNER JOIN tx origin_tx			ON tx_out.tx_id = origin_tx.id
+		INNER JOIN block origin_block	ON origin_tx.block_id = origin_block.id
+		INNER JOIN datum				ON tx_out.data_hash = datum.hash
+		INNER JOIN ma_tx_out			ON tx_out.id = ma_tx_out.tx_out_id
+		INNER JOIN multi_asset			ON multi_asset.id = ma_tx_out.ident
+		WHERE
+			tx_out.address = $1 AND ($2 IS NULL OR origin_block.block_no > $2) AND origin_block.block_no <= $3
+			AND multi_asset.policy = $4
+			AND multi_asset.name = $5)
+		UNION
+		(SELECT
+			datum.value as datum, consuming_block.block_no as block_no, consuming_tx.block_index as block_index, 'remove' as action
+		FROM tx_out
+		LEFT JOIN tx_in consuming_tx_in	ON tx_out.tx_id = consuming_tx_in.tx_out_id AND tx_out.index = consuming_tx_in.tx_out_index
+		LEFT JOIN tx consuming_tx		ON consuming_tx_in.tx_in_id = consuming_tx.id
+		LEFT JOIN block consuming_block	ON consuming_tx.block_id = consuming_block.id
+		INNER JOIN datum				ON tx_out.data_hash = datum.hash
+		INNER JOIN ma_tx_out			ON tx_out.id = ma_tx_out.tx_out_id
+		INNER JOIN multi_asset			ON multi_asset.id = ma_tx_out.ident
+		WHERE
+			tx_out.address = $1
+			AND (consuming_tx_in.id IS NOT NULL AND ($2 IS NULL OR consuming_block.block_no > $2) AND consuming_block.block_no <= $3)
+			AND multi_asset.policy = $4
+			AND multi_asset.name = $5)
+		ORDER BY block_no ASC";
+	Ok(sqlx::query_as::<_, DatumChangeOutput>(query)
+		.bind(&address.0)
+		.bind(last_block)
+		.bind(new_block)
+		.bind(&asset.policy_id.0)
+		.bind(&asset.asset_name.0)
+		.fetch_all(pool)
+		.await?)
+}
+
 #[cfg(feature = "governed-map")]
 pub(crate) async fn get_datums_at_address_with_token(
 	pool: &Pool<Postgres>,

toolkit/data-sources/db-sync/src/governed_map/mod.rs

@@ -1,14 +1,18 @@
-use crate::DataSourceError::ExpectedDataNotFound;
+use crate::DataSourceError::{self, ExpectedDataNotFound};
 use crate::Result;
+use crate::db_model::Block;
 use crate::{metrics::McFollowerMetrics, observed_async_trait};
+use db_sync_sqlx::{BlockNumber, TxIndexInBlock};
 use derive_new::new;
 use log::warn;
 use partner_chains_plutus_data::governed_map::GovernedMapDatum;
 use sidechain_domain::byte_string::ByteString;
 use sidechain_domain::*;
 use sp_governed_map::{GovernedMapDataSource, MainChainScriptsV1};
 use sqlx::PgPool;
+use std::cmp::{max, min};
 use std::collections::HashMap;
+use std::sync::{Arc, Mutex};
 
 #[cfg(test)]
 pub mod tests;
@@ -87,3 +91,209 @@ impl GovernedMapDataSourceImpl {
 		Ok(mappings)
 	}
 }
+
+#[derive(new)]
+pub struct GovernedMapDataSourceCachedImpl {
+	pub pool: PgPool,
+	pub metrics_opt: Option<McFollowerMetrics>,
+	security_parameter: u32,
+	cache_size: u16,
+	#[new(default)]
+	cache: Arc<Mutex<Cache>>,
+}
+
+observed_async_trait!(
+	impl GovernedMapDataSource for GovernedMapDataSourceCachedImpl {
+		async fn get_mapping_changes(
+			&self,
+			since_mc_block: Option<McBlockHash>,
+			up_to_mc_block: McBlockHash,
+			scripts: MainChainScriptsV1,
+		) -> std::result::Result<
+			Vec<(String, Option<ByteString>)>,
+			Box<dyn std::error::Error + Send + Sync>,
+		> {
+			let since_block_number = match since_mc_block {
+				Some(hash) =>
+					Some(crate::db_model::get_block_by_hash(&self.pool, hash.clone())
+						.await?
+						.ok_or_else(|| Box::new(ExpectedDataNotFound(format!("Block hash: {hash}"))) as Box<dyn std::error::Error + Send + Sync>)?
+						.block_no),
+				None => None,
+			};
+
+			let Some(up_to_block) =
+				crate::db_model::get_block_by_hash(&self.pool, up_to_mc_block.clone()).await?
+			else {
+				return Err(Box::new(ExpectedDataNotFound(format!("Block hash: {up_to_mc_block}"))));
+			};
+
+			if let Some(cached_changes) = self.get_changes_from_cache(since_block_number, up_to_block.block_no).await? {
+				return Ok(cached_changes);
+			}
+
+			let latest_block = get_latest_block(&self.pool).await?;
+
+			let latest_stable_block = BlockNumber(latest_block.block_no.0.saturating_sub(self.security_parameter));
+
+			let since_block_plus = BlockNumber(since_block_number.unwrap_or(BlockNumber(0)).0 + self.cache_size as u32);
+
+			let max_search_block = min(latest_stable_block, max(up_to_block.block_no, since_block_plus));
+
+			let changes = self.get_changes_in_range_to_cache(since_block_number, max_search_block, scripts).await?;
+
+			if let Ok(mut cache) = self.cache.lock() {
+				cache.add_changes(changes.clone());
+			}
+
+			Ok(deduplicate_changes(filter_changes_in_range(changes, since_block_number, up_to_block.block_no)))
+		}
+	}
+);
+
+impl GovernedMapDataSourceCachedImpl {
+	pub fn new_from_env(
+		pool: PgPool,
+		metrics_opt: Option<McFollowerMetrics>,
+	) -> std::result::Result<Self, &'static str> {
+		let security_parameter: u32 = std::env::var("CARDANO_SECURITY_PARAMETER")
+			.ok()
+			.and_then(|s| s.parse().ok())
+			.ok_or("Couldn't read env variable CARDANO_SECURITY_PARAMETER as u32")?;
+		Ok(Self {
+			pool,
+			metrics_opt,
+			security_parameter,
+			cache_size: 1000,
+			cache: Default::default(),
+		})
+	}
+
+	async fn get_changes_from_cache(
+		&self,
+		since_block: Option<BlockNumber>,
+		up_to_block: BlockNumber,
+	) -> Result<Option<Vec<(String, Option<ByteString>)>>> {
+		if let Ok(cache) = self.cache.lock() {
+			if let Some(changes) = cache.get_changes_in_range(since_block, up_to_block) {
+				return Ok(Some(deduplicate_changes(changes)));
+			}
+		};
+		return Ok(None);
+	}
+
+	async fn get_changes_in_range_to_cache(
+		&self,
+		since_block: Option<BlockNumber>,
+		up_to_block: BlockNumber,
+		scripts: MainChainScriptsV1,
+	) -> Result<Vec<(String, Option<ByteString>, BlockNumber, TxIndexInBlock)>> {
+		let latest_block_number = {
+			let cache = self.cache.lock().expect("Failed to lock cache");
+
+			cache.newest_block_number.or(since_block)
+		};
+		let changes = crate::db_model::get_changes(
+			&self.pool,
+			&scripts.validator_address.into(),
+			latest_block_number,
+			up_to_block,
+			scripts.asset.into(),
+		)
+		.await?;
+
+		let mut result = Vec::new();
+
+		for change in changes {
+			match GovernedMapDatum::try_from(change.datum.0) {
+				Ok(GovernedMapDatum { key, value }) => match change.action.as_str() {
+					"remove" => result.push((key, None, change.block_no, change.block_index)),
+					"upsert" => {
+						result.push((key, Some(value), change.block_no, change.block_index))
+					},
+					_ => warn!("Unknown action: {}", change.action),
+				},
+				Err(err) => warn!("Failed decoding map entry: {err}"),
+			}
+		}
+		Ok(result)
+	}
+}
+
+#[derive(Default)]
+pub(crate) struct Cache {
+	newest_block_number: Option<BlockNumber>,
+	changes: Vec<(String, Option<ByteString>, BlockNumber, TxIndexInBlock)>,
+}
+
+fn deduplicate_changes(
+	mut changes: Vec<(String, Option<ByteString>, BlockNumber, TxIndexInBlock)>,
+) -> Vec<(String, Option<ByteString>)> {
+	let mut result = HashMap::new();
+
+	changes.sort_by_key(|(_, _, block_number, block_index)| (*block_number, *block_index));
+
+	for (key, value, _, _) in changes {
+		if result.contains_key(&key) {
+			result.remove(&key);
+		}
+		result.insert(key, value);
+	}
+
+	result.into_iter().collect()
+}
+
+fn filter_changes_in_range(
+	changes: Vec<(String, Option<ByteString>, BlockNumber, TxIndexInBlock)>,
+	since_block: Option<BlockNumber>,
+	up_to_block: BlockNumber,
+) -> Vec<(String, Option<ByteString>, BlockNumber, TxIndexInBlock)> {
+	changes
+		.iter()
+		.filter(|(_, _, block_number, _)| {
+			block_number.0 <= up_to_block.0
+				&& since_block.map(|b| block_number.0 > b.0).unwrap_or(true)
+		})
+		.map(|x| x.clone())
+		.collect()
+}
+
+async fn get_latest_block(pool: &PgPool) -> Result<Block> {
+	crate::db_model::get_latest_block_info(pool).await?.ok_or(
+		DataSourceError::ExpectedDataNotFound(
+			"The latest block not found when querying for native token transfers".to_string(),
+		),
+	)
+}
+
+impl Cache {
+	pub fn get_changes_in_range(
+		&self,
+		since_block: Option<BlockNumber>,
+		up_to_block: BlockNumber,
+	) -> Option<Vec<(String, Option<ByteString>, BlockNumber, TxIndexInBlock)>> {
+		let Some(newest_block_number) = self.newest_block_number else {
+			return None;
+		};
+
+		if newest_block_number.0 < up_to_block.0 {
+			return None;
+		}
+
+		Some(filter_changes_in_range(self.changes.clone(), since_block, up_to_block))
+	}
+
+	pub fn add_changes(
+		&mut self,
+		changes: Vec<(String, Option<ByteString>, BlockNumber, TxIndexInBlock)>,
+	) {
+		self.changes.extend(changes);
+		self.newest_block_number = Some(
+			self.changes
+				.iter()
+				.max_by_key(|(_, _, block_number, _)| block_number.0)
+				.map(|(_, _, block_number, _)| *block_number)
+				.unwrap_or(BlockNumber(0)),
+		);
+	}
+}