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 nanoseconds) for cached entries in the get_utxos cache.
+    get_utxos_cache_expiration: 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 nanoseconds) for cached entries in the get_utxos cache.
+    get_utxos_cache_expiration: 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: 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,7 +112,7 @@ 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)]
 pub struct GetUtxosResponse {
@@ -121,8 +121,8 @@ pub struct GetUtxosResponse {
     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,
         }
     }
 }
@@ -1297,7 +1297,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 {
@@ -1340,3 +1352,96 @@ impl From<u64> for Timestamp {
         Self(timestamp)
     }
 }
+
+#[derive(Clone, Eq, PartialEq, Ord, PartialOrd, Debug, CandidType, Deserialize, Serialize)]
+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.
+///
+/// On the other hand, lookups are not affected by expiration because they are ready only.
+/// If an old entry is found, it is returned regardless of its timestamp.
+#[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 expired = self.values.split_off(&pivot);
+            std::mem::swap(&mut self.values, &mut expired);
+            expired.keys().for_each(|key| {
+                self.keys.remove(key.inner.as_ref().unwrap());
+            });
+        }
+    }
+    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 nanoseconds) for cached entries in
+    /// the get_utxos cache.
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub get_utxos_cache_expiration: 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 nanoseconds) for cached entries in
+    /// the get_utxos cache.
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub get_utxos_cache_expiration: 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: u64,
         req: GetUtxosRequest,
         source: CallSource,
         runtime: &R,
@@ -165,7 +166,18 @@ 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| {
+                crate::state::mutate_state(|s| {
+                    s.get_utxos_cache.prune(now);
+                    s.get_utxos_cache.insert(req, res.clone(), now)
+                })
+            })
+        }
     }
 
     // Record start time of method execution for metrics
@@ -176,7 +188,7 @@ pub async fn get_utxos<R: CanisterRuntime>(
         filter: Some(UtxoFilter::MinConfirmations(min_confirmations)),
     };
 
-    let mut response = bitcoin_get_utxos(request.clone(), source, runtime).await?;
+    let mut response = bitcoin_get_utxos(start_time, request.clone(), source, runtime).await?;
 
     let mut utxos = std::mem::take(&mut response.utxos);
     let mut num_pages: usize = 1;
@@ -187,7 +199,7 @@ 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(runtime.time(), paged_request, source, runtime).await?;
         utxos.append(&mut response.utxos);
         num_pages += 1;
     }

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

@@ -13,6 +13,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 const BUCKETS_MS: [u64; 8] = [500, 1_000, 2_000, 4_000, 8_000, 16_000, 32_000, u64::MAX];
@@ -267,7 +269,9 @@ pub fn encode_metrics(
             "Number of get_utxos calls the minter issued, labeled by source.",
         )?
         .value(&[("source", "client")], GET_UTXOS_CLIENT_CALLS.get() as f64)?
-        .value(&[("source", "minter")], GET_UTXOS_MINTER_CALLS.get() as f64)?;
+        .value(&[("source", "minter")], GET_UTXOS_MINTER_CALLS.get() as f64)?
+        .value(&[("cache", "hits")], GET_UTXOS_CACHE_HITS.get() as f64)?
+        .value(&[("cache", "misses")], GET_UTXOS_CACHE_MISSES.get() as f64)?;
 
     metrics.encode_gauge(
         "ckbtc_minter_btc_balance",

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,
         }: 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 {
+            self.get_utxos_cache
+                .set_expiration(Duration::from_nanos(expiration));
+        }
     }
 
     #[allow(deprecated)]
@@ -464,6 +473,7 @@ impl CkBtcMinterState {
             btc_checker_principal,
             kyt_principal: _,
             kyt_fee,
+            get_utxos_cache_expiration,
         }: 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 {
+            self.get_utxos_cache
+                .set_expiration(Duration::from_nanos(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_nanos(
+                args.get_utxos_cache_expiration.unwrap_or_default(),
+            )),
         }
     }
 }

rs/bitcoin/ckbtc/minter/src/test_fixtures.rs

@@ -30,6 +30,7 @@ pub fn init_args() -> InitArgs {
         check_fee: None,
         kyt_principal: None,
         kyt_fee: None,
+        get_utxos_cache_expiration: None,
     }
 }
 
@@ -348,6 +349,7 @@ pub mod arbitrary {
                 kyt_fee: option::of(any::<u64>()),
                 btc_checker_principal: option::of(canister_id()),
                 kyt_principal: option::of(canister_id()),
+                get_utxos_cache_expiration: option::of(any::<u64>()),
             })
         }
 
@@ -361,6 +363,7 @@ pub mod arbitrary {
                 kyt_fee: option::of(any::<u64>()),
                 btc_checker_principal: option::of(canister_id()),
                 kyt_principal: option::of(canister_id()),
+                get_utxos_cache_expiration: option::of(any::<u64>()),
             })
         }