feat(ckbtc): Add get_utxos_cache to reduce latency of update_balance calls

rs/bitcoin/ckbtc/minter/ckbtc_minter.did

@@ -208,6 +208,9 @@ type InitArgs = record {
 
     /// The canister id of the KYT canister (deprecated, use btc_checker_principal instead).
     kyt_principal: opt principal;
+
+    /// The expiration duration (in seconds) for cached entries in the get_utxos cache.
+    get_utxos_cache_expiration_seconds: opt nat64;
 };
 
 // The upgrade parameters of the minter canister.
@@ -237,6 +240,9 @@ type UpgradeArgs = record {
 
     /// The canister id of the KYT canister (deprecated, use btc_checker_principal instead).
     kyt_principal: opt principal;
+
+    /// The expiration duration (in seconds) for cached entries in the get_utxos cache.
+    get_utxos_cache_expiration_seconds: opt nat64;
 };
 
 type RetrieveBtcStatus = variant {

rs/bitcoin/ckbtc/minter/src/guard.rs

@@ -135,6 +135,7 @@ mod tests {
             check_fee: None,
             kyt_principal: None,
             kyt_fee: None,
+            get_utxos_cache_expiration_seconds: None,
         }
     }
 

rs/bitcoin/ckbtc/minter/src/lib.rs

@@ -8,7 +8,7 @@ use candid::{CandidType, Deserialize, Principal};
 use ic_btc_checker::CheckTransactionResponse;
 use ic_btc_interface::{MillisatoshiPerByte, OutPoint, Page, Satoshi, Txid, Utxo};
 use ic_canister_log::log;
-use ic_cdk::api::management_canister::bitcoin::BitcoinNetwork;
+use ic_cdk::api::management_canister::bitcoin;
 use ic_management_canister_types_private::DerivationPath;
 use icrc_ledger_types::icrc1::account::Account;
 use icrc_ledger_types::icrc1::transfer::Memo;
@@ -112,17 +112,17 @@ pub struct ECDSAPublicKey {
     pub chain_code: Vec<u8>,
 }
 
-pub type GetUtxosRequest = ic_cdk::api::management_canister::bitcoin::GetUtxosRequest;
+pub type GetUtxosRequest = bitcoin::GetUtxosRequest;
 
-#[derive(Clone, Eq, PartialEq, Debug, CandidType, Deserialize, Serialize)]
+#[derive(Clone, Eq, PartialEq, Debug)]
 pub struct GetUtxosResponse {
     pub utxos: Vec<Utxo>,
     pub tip_height: u32,
     pub next_page: Option<Page>,
 }
 
-impl From<ic_cdk::api::management_canister::bitcoin::GetUtxosResponse> for GetUtxosResponse {
-    fn from(response: ic_cdk::api::management_canister::bitcoin::GetUtxosResponse) -> Self {
+impl From<bitcoin::GetUtxosResponse> for GetUtxosResponse {
+    fn from(response: bitcoin::GetUtxosResponse) -> Self {
         Self {
             utxos: response
                 .utxos
@@ -155,12 +155,12 @@ pub enum Network {
     Regtest,
 }
 
-impl From<Network> for BitcoinNetwork {
+impl From<Network> for bitcoin::BitcoinNetwork {
     fn from(network: Network) -> Self {
         match network {
-            Network::Mainnet => BitcoinNetwork::Mainnet,
-            Network::Testnet => BitcoinNetwork::Testnet,
-            Network::Regtest => BitcoinNetwork::Regtest,
+            Network::Mainnet => bitcoin::BitcoinNetwork::Mainnet,
+            Network::Testnet => bitcoin::BitcoinNetwork::Testnet,
+            Network::Regtest => bitcoin::BitcoinNetwork::Regtest,
         }
     }
 }
@@ -1331,7 +1331,19 @@ impl CanisterRuntime for IcCanisterRuntime {
 }
 
 /// Time in nanoseconds since the epoch (1970-01-01).
-#[derive(Eq, Clone, Copy, PartialEq, Debug, Default, Serialize, CandidType, serde::Deserialize)]
+#[derive(
+    Clone,
+    Copy,
+    Eq,
+    PartialEq,
+    Ord,
+    PartialOrd,
+    Debug,
+    Default,
+    Serialize,
+    CandidType,
+    serde::Deserialize,
+)]
 pub struct Timestamp(u64);
 
 impl Timestamp {
@@ -1374,3 +1386,97 @@ impl From<u64> for Timestamp {
         Self(timestamp)
     }
 }
+
+#[derive(Clone, Eq, PartialEq, Ord, PartialOrd, Debug)]
+struct Timestamped<Inner> {
+    timestamp: Timestamp,
+    inner: Option<Inner>,
+}
+
+impl<Inner> Timestamped<Inner> {
+    fn new<T: Into<Timestamp>>(timestamp: T, inner: Inner) -> Self {
+        Self {
+            timestamp: timestamp.into(),
+            inner: Some(inner),
+        }
+    }
+}
+
+/// A cache that expires older entries upon insertion.
+///
+/// More specifically, entries are inserted with a timestamp, and
+/// then all existing entries with a timestamp less than `t - expiration` are removed before
+/// the new entry is inserted.
+///
+/// Similarly, lookups will also take an additional timestamp as argument, and only entries
+/// newer than that will be returned.
+#[derive(Clone, Eq, PartialEq, Debug)]
+pub struct CacheWithExpiration<Key, Value> {
+    expiration: Duration,
+    keys: BTreeMap<Key, Timestamp>,
+    values: BTreeMap<Timestamped<Key>, Value>,
+}
+
+impl<Key: Ord + Clone, Value: Clone> CacheWithExpiration<Key, Value> {
+    pub fn new(expiration: Duration) -> Self {
+        Self {
+            expiration,
+            keys: Default::default(),
+            values: Default::default(),
+        }
+    }
+
+    pub fn is_empty(&self) -> bool {
+        self.len() == 0
+    }
+
+    pub fn len(&self) -> usize {
+        let len = self.keys.len();
+        assert_eq!(len, self.values.len());
+        len
+    }
+
+    pub fn set_expiration(&mut self, expiration: Duration) {
+        self.expiration = expiration;
+    }
+
+    pub fn prune<T: Into<Timestamp>>(&mut self, now: T) {
+        let timestamp = now.into();
+        if let Some(expire_cutoff) = timestamp.checked_sub(self.expiration) {
+            let pivot = Timestamped {
+                timestamp: expire_cutoff,
+                inner: None,
+            };
+            let mut non_expired = self.values.split_off(&pivot);
+            self.values.keys().for_each(|key| {
+                self.keys.remove(key.inner.as_ref().unwrap());
+            });
+            std::mem::swap(&mut self.values, &mut non_expired);
+            assert_eq!(self.keys.len(), self.values.len())
+        }
+    }
+    pub fn insert<T: Into<Timestamp>>(&mut self, key: Key, value: Value, now: T) {
+        let timestamp = now.into();
+        if let Some(old_timestamp) = self.keys.insert(key.clone(), timestamp) {
+            self.values
+                .remove(&Timestamped::new(old_timestamp, key.clone()));
+        }
+        self.values.insert(Timestamped::new(timestamp, key), value);
+    }
+
+    pub fn get<T: Into<Timestamp>>(&self, key: &Key, now: T) -> Option<&Value> {
+        let now = now.into();
+        let timestamp = *self.keys.get(key)?;
+        if let Some(expire_cutoff) = now.checked_sub(self.expiration) {
+            if timestamp < expire_cutoff {
+                return None;
+            }
+        }
+        self.values.get(&Timestamped {
+            timestamp,
+            inner: Some(key.clone()),
+        })
+    }
+}
+
+pub type GetUtxosCache = CacheWithExpiration<bitcoin::GetUtxosRequest, GetUtxosResponse>;

rs/bitcoin/ckbtc/minter/src/lifecycle/init.rs

@@ -63,6 +63,11 @@ pub struct InitArgs {
     #[serde(skip_serializing_if = "Option::is_none")]
     #[deprecated(note = "use btc_checker_principal instead")]
     pub kyt_principal: Option<CanisterId>,
+
+    /// The expiration duration in seconds) for cached entries in
+    /// the get_utxos cache.
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub get_utxos_cache_expiration_seconds: Option<u64>,
 }
 
 pub fn init(args: InitArgs) {

rs/bitcoin/ckbtc/minter/src/lifecycle/upgrade.rs

@@ -44,6 +44,11 @@ pub struct UpgradeArgs {
     #[serde(skip_serializing_if = "Option::is_none")]
     #[deprecated(note = "use btc_checker_principal instead")]
     pub kyt_principal: Option<CanisterId>,
+
+    /// The expiration duration (in seconds) for cached entries in
+    /// the get_utxos cache.
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub get_utxos_cache_expiration_seconds: Option<u64>,
 }
 
 pub fn post_upgrade(upgrade_args: Option<UpgradeArgs>) {

rs/bitcoin/ckbtc/minter/src/management.rs

@@ -156,6 +156,7 @@ pub async fn get_utxos<R: CanisterRuntime>(
     runtime: &R,
 ) -> Result<GetUtxosResponse, CallError> {
     async fn bitcoin_get_utxos<R: CanisterRuntime>(
+        now: &mut u64,
         req: GetUtxosRequest,
         source: CallSource,
         runtime: &R,
@@ -165,17 +166,31 @@ pub async fn get_utxos<R: CanisterRuntime>(
             CallSource::Minter => &crate::metrics::GET_UTXOS_MINTER_CALLS,
         }
         .with(|cell| cell.set(cell.get() + 1));
-        runtime.bitcoin_get_utxos(req).await
+        if let Some(res) = crate::state::read_state(|s| s.get_utxos_cache.get(&req, *now).cloned())
+        {
+            crate::metrics::GET_UTXOS_CACHE_HITS.with(|cell| cell.set(cell.get() + 1));
+            Ok(res)
+        } else {
+            crate::metrics::GET_UTXOS_CACHE_MISSES.with(|cell| cell.set(cell.get() + 1));
+            runtime.bitcoin_get_utxos(req.clone()).await.inspect(|res| {
+                *now = runtime.time();
+                crate::state::mutate_state(|s| {
+                    s.get_utxos_cache.prune(*now);
+                    s.get_utxos_cache.insert(req, res.clone(), *now)
+                })
+            })
+        }
     }
 
     let start_time = runtime.time();
+    let mut now = start_time;
     let request = GetUtxosRequest {
         address: address.clone(),
         network: network.into(),
         filter: Some(UtxoFilter::MinConfirmations(min_confirmations)),
     };
 
-    let mut response = bitcoin_get_utxos(request.clone(), source, runtime).await?;
+    let mut response = bitcoin_get_utxos(&mut now, request.clone(), source, runtime).await?;
 
     let mut utxos = std::mem::take(&mut response.utxos);
     let mut num_pages: usize = 1;
@@ -186,12 +201,12 @@ pub async fn get_utxos<R: CanisterRuntime>(
             filter: Some(UtxoFilter::Page(page.to_vec())),
             ..request.clone()
         };
-        response = bitcoin_get_utxos(paged_request, source, runtime).await?;
+        response = bitcoin_get_utxos(&mut now, paged_request, source, runtime).await?;
         utxos.append(&mut response.utxos);
         num_pages += 1;
     }
 
-    observe_get_utxos_latency(utxos.len(), num_pages, source, start_time, runtime.time());
+    observe_get_utxos_latency(utxos.len(), num_pages, source, start_time, now);
 
     response.utxos = utxos;
 

rs/bitcoin/ckbtc/minter/src/metrics.rs

@@ -28,6 +28,8 @@ thread_local! {
     pub static UPDATE_CALL_LATENCY: RefCell<BTreeMap<NumUtxoPages,LatencyHistogram>> = RefCell::default();
     pub static GET_UTXOS_CALL_LATENCY: RefCell<BTreeMap<(NumUtxoPages, CallSource),LatencyHistogram>> = RefCell::default();
     pub static GET_UTXOS_RESULT_SIZE: RefCell<BTreeMap<CallSource,NumUtxosHistogram>> = RefCell::default();
+    pub static GET_UTXOS_CACHE_HITS : Cell<u64> = Cell::default();
+    pub static GET_UTXOS_CACHE_MISSES: Cell<u64> = Cell::default();
     pub static SIGN_WITH_ECDSA_LATENCY: RefCell<BTreeMap<MetricsResult, LatencyHistogram>> = RefCell::default();
 }
 
@@ -303,6 +305,18 @@ pub fn encode_metrics(
         .value(&[("source", "client")], GET_UTXOS_CLIENT_CALLS.get() as f64)?
         .value(&[("source", "minter")], GET_UTXOS_MINTER_CALLS.get() as f64)?;
 
+    metrics.encode_counter(
+        "ckbtc_minter_get_utxos_cache_hits",
+        GET_UTXOS_CACHE_HITS.get() as f64,
+        "Number of cache hits for get_utxos calls.",
+    )?;
+
+    metrics.encode_counter(
+        "ckbtc_minter_get_utxos_cache_misses",
+        GET_UTXOS_CACHE_MISSES.get() as f64,
+        "Number of cache misses for get_utxos calls.",
+    )?;
+
     metrics.encode_gauge(
         "ckbtc_minter_btc_balance",
         state::read_state(|s| s.get_total_btc_managed()) as f64,

rs/bitcoin/ckbtc/minter/src/state.rs

@@ -20,7 +20,7 @@ use crate::lifecycle::upgrade::UpgradeArgs;
 use crate::logs::P0;
 use crate::state::invariants::{CheckInvariants, CheckInvariantsImpl};
 use crate::updates::update_balance::SuspendedUtxo;
-use crate::{address::BitcoinAddress, ECDSAPublicKey, Network, Timestamp};
+use crate::{address::BitcoinAddress, ECDSAPublicKey, GetUtxosCache, Network, Timestamp};
 use candid::{CandidType, Deserialize, Principal};
 use ic_base_types::CanisterId;
 use ic_btc_interface::{OutPoint, Txid, Utxo};
@@ -31,6 +31,7 @@ use serde::Serialize;
 use std::collections::btree_map::Entry;
 use std::collections::btree_set;
 use std::iter::Chain;
+use std::time::Duration;
 
 /// The maximum number of finalized BTC retrieval requests that we keep in the
 /// history.
@@ -390,6 +391,9 @@ pub struct CkBtcMinterState {
 
     /// Map from burn block index to the the reimbursed request.
     pub reimbursed_transactions: BTreeMap<u64, ReimbursedDeposit>,
+
+    /// Cache of get_utxos call results
+    pub get_utxos_cache: GetUtxosCache,
 }
 
 #[derive(Clone, Eq, PartialEq, Debug, CandidType, Serialize, serde::Deserialize)]
@@ -432,6 +436,7 @@ impl CkBtcMinterState {
             btc_checker_principal,
             kyt_principal: _,
             kyt_fee,
+            get_utxos_cache_expiration_seconds,
         }: InitArgs,
     ) {
         self.btc_network = btc_network;
@@ -450,6 +455,10 @@ impl CkBtcMinterState {
         if let Some(min_confirmations) = min_confirmations {
             self.min_confirmations = min_confirmations;
         }
+        if let Some(expiration) = get_utxos_cache_expiration_seconds {
+            self.get_utxos_cache
+                .set_expiration(Duration::from_secs(expiration));
+        }
     }
 
     #[allow(deprecated)]
@@ -464,6 +473,7 @@ impl CkBtcMinterState {
             btc_checker_principal,
             kyt_principal: _,
             kyt_fee,
+            get_utxos_cache_expiration_seconds,
         }: UpgradeArgs,
     ) {
         if let Some(retrieve_btc_min_amount) = retrieve_btc_min_amount {
@@ -496,6 +506,10 @@ impl CkBtcMinterState {
         } else if let Some(kyt_fee) = kyt_fee {
             self.check_fee = kyt_fee;
         }
+        if let Some(expiration) = get_utxos_cache_expiration_seconds {
+            self.get_utxos_cache
+                .set_expiration(Duration::from_secs(expiration));
+        }
     }
 
     pub fn validate_config(&self) {
@@ -1500,6 +1514,9 @@ impl From<InitArgs> for CkBtcMinterState {
             suspended_utxos: Default::default(),
             pending_reimbursements: Default::default(),
             reimbursed_transactions: Default::default(),
+            get_utxos_cache: GetUtxosCache::new(Duration::from_secs(
+                args.get_utxos_cache_expiration_seconds.unwrap_or_default(),
+            )),
         }
     }
 }