Capped Hashmap

.vscode/launch.json

@@ -0,0 +1,26 @@
+{
+  // Use IntelliSense to learn about possible attributes.
+  // Hover to view descriptions of existing attributes.
+  // For more information, visit: https://go.microsoft.com/fwlink/?linkid=830387
+  "version": "0.2.0",
+  "configurations": [
+    {
+      "type": "lldb",
+      "request": "launch",
+      "name": "Debug zkbtc",
+      "cargo": {
+        "args": [
+          "build",
+          "--bin=zkbtc",
+          "--package=zkbitcoin"
+        ],
+        "filter": {
+          "name": "zkbtc",
+          "kind": "bin"
+        }
+      },
+      "args": [],
+      "cwd": "${workspaceFolder}"
+    },
+  ]
+}

src/capped_hashmap.rs

@@ -19,11 +19,11 @@ where
 {
     pub fn new(capacity: usize) -> Self {
         Self {
+            capacity,
             // We add a little bit margin to the capacity, because we always insert first and then check
             // if collection is full. So we don't want the hashmap to double it's capacity because of that.
-            capacity: capacity + 1,
-            inner: HashMap::with_capacity(capacity),
-            last_items: VecDeque::with_capacity(capacity),
+            inner: HashMap::with_capacity(capacity + 1),
+            last_items: VecDeque::with_capacity(capacity + 1),
         }
     }
 
@@ -33,11 +33,11 @@ where
         let mut ret = None;
 
         // replacing a value should not push any new items to last_items
-        if let None = self.inner.insert(k, v) {
+        if self.inner.insert(k, v).is_none() {
             self.last_items.push_front(k);
         }
 
-        if self.last_items.len() > self.capacity - 1 {
+        if self.last_items.len() > self.capacity {
             // remove the oldest item. We an safely unwrap because we know the last_items is not empty at this point
             let key = self.last_items.pop_back().unwrap();
             assert!(self.remove(&key).is_some());
@@ -54,12 +54,11 @@ where
             return None;
         };
 
-        self.last_items = self
+        self
             .last_items
             .iter()
-            .filter(|key| *key != k)
-            .map(|key| *key)
-            .collect::<VecDeque<_>>();
+            .position(|key| key == k)
+            .and_then(|pos| self.last_items.remove(pos));
 
         Some(v)
     }
@@ -98,6 +97,7 @@ mod tests {
 
     #[test]
     fn test_insert_should_return_removed_key() {
+        // The real capacity will be 14. Read here for how this is calculated https://stackoverflow.com/a/76114888/512783
         let mut col: CappedHashMap<u8, u8> = CappedHashMap::new(10);
 
         for i in 0..10 {
@@ -107,6 +107,9 @@ mod tests {
         for i in 10..30 {
             // the nth oldest key will be removed
             let key_removed = col.insert(i, i);
+            // our hashmap and vecqueue should never grow i.e. capacity doesn't change
+            assert_eq!(col.last_items.capacity(), 11);
+
             assert!(key_removed.is_some());
             assert_eq!(key_removed.unwrap(), i - 10);
             assert_eq!(col.size(), 10);

src/committee/node.rs

@@ -17,10 +17,8 @@ use rand::thread_rng;
 use serde::{Deserialize, Serialize};
 
 use crate::{
-    bob_request::{BobRequest, SmartContract},
-    capped_hashmap::CappedHashMap,
-    frost,
-    mpc_sign_tx::get_digest_to_hash,
+    bob_request::{BobRequest, SmartContract}, capped_hashmap::CappedHashMap,
+    constants::MAX_SIGNING_TASK, frost, mpc_sign_tx::get_digest_to_hash,
 };
 
 //
@@ -35,7 +33,7 @@ pub struct NodeState {
     /// The public key stuff they need.
     pub pubkey_package: frost::PublicKeyPackage,
 
-    // TODO: ensure that this cannot grow like crazy? prune old tasks?
+    /// The current pending signing tasks
     pub signing_tasks: RwLock<CappedHashMap<Txid, LocalSigningTask>>,
 }
 
@@ -234,7 +232,7 @@ pub async fn run_server(
     let ctx = NodeState {
         key_package,
         pubkey_package,
-        signing_tasks: RwLock::new(CappedHashMap::new(100)),
+        signing_tasks: RwLock::new(CappedHashMap::new(MAX_SIGNING_TASK)),
     };
 
     let server = Server::builder()

src/constants.rs

@@ -43,3 +43,5 @@ pub const KEEPALIVE_WAIT_SECONDS: u64 = 5;
 
 /// The total number of fibonacci backoff retries before considering an MPC node offline
 pub const KEEPALIVE_MAX_RETRIES: u8 = 10;
+
+pub const MAX_SIGNING_TASK: usize = 100;

src/lib.rs

@@ -3,6 +3,7 @@
 use anyhow::Context;
 use secp256k1::hashes::Hash;
 
+pub mod capped_hashmap;
 pub mod committee;
 pub mod constants;
 pub mod frost;
@@ -21,7 +22,6 @@ pub mod bob_request;
 /// 3. The zkBitcoin committee produce a collaborative schnorr signature to unlock the funds for Bob.
 pub mod mpc_sign_tx;
 
-pub mod capped_hashmap;
 //
 // Helpers
 //