feat: Add retain function

Cargo.toml

@@ -16,7 +16,7 @@ categories = ["asynchronous", "data-structures"]
 
 [dependencies]
 cid = "0.11.0"
-dashmap = "5.5.3"
+dashmap = "6.0.1"
 lru = { version = "0.12.2", optional = true }
 multihash = "0.19.1"
 thiserror = "1.0.40"
@@ -31,11 +31,12 @@ tokio = { version = "1.29.0", features = ["macros", "rt"], optional = true }
 
 [target.'cfg(target_arch = "wasm32")'.dependencies]
 js-sys = { version = "0.3.68", optional = true }
-rexie = { version = "0.5.0", optional = true }
+rexie = { version = "0.6.1", optional = true }
+idb = { version = "0.6.2", optional = true }
 wasm-bindgen = { version = "0.2.91", optional = true }
 
 [dev-dependencies]
-rstest = "0.18.2"
+rstest = "0.22.0"
 tokio = { version = "1.29.0", features = ["macros", "rt"] }
 tempfile = "3.10"
 

src/in_memory_blockstore.rs

@@ -44,6 +44,13 @@ impl<const MAX_MULTIHASH_SIZE: usize> InMemoryBlockstore<MAX_MULTIHASH_SIZE> {
     fn remove_cid(&self, cid: &CidGeneric<MAX_MULTIHASH_SIZE>) {
         self.map.remove(cid);
     }
+
+    fn retain<F>(&self, mut predicate: F)
+    where
+        F: FnMut(&[u8]) -> bool,
+    {
+        self.map.retain(|cid, _block| predicate(&cid.to_bytes()));
+    }
 }
 
 impl<const MAX_MULTIHASH_SIZE: usize> Blockstore for InMemoryBlockstore<MAX_MULTIHASH_SIZE> {
@@ -67,6 +74,14 @@ impl<const MAX_MULTIHASH_SIZE: usize> Blockstore for InMemoryBlockstore<MAX_MULT
         let cid = convert_cid(cid)?;
         Ok(self.contains_cid(&cid))
     }
+
+    async fn retain<F>(&self, predicate: F) -> Result<()>
+    where
+        F: FnMut(&[u8]) -> bool,
+    {
+        self.retain(predicate);
+        Ok(())
+    }
 }
 
 impl<const MAX_MULTIHASH_SIZE: usize> Default for InMemoryBlockstore<MAX_MULTIHASH_SIZE> {

src/indexed_db_blockstore.rs

@@ -1,5 +1,5 @@
 use cid::CidGeneric;
-use js_sys::Uint8Array;
+use js_sys::{ArrayBuffer, Uint8Array};
 use rexie::{KeyRange, ObjectStore, Rexie, Store, TransactionMode};
 use wasm_bindgen::{JsCast, JsValue};
 
@@ -9,6 +9,7 @@ use crate::{Blockstore, Error, Result};
 const DB_VERSION: u32 = 1;
 
 const BLOCK_STORE: &str = "BLOCKSTORE.BLOCKS";
+const RETAIN_BATCH_SIZE: u32 = 1024;
 
 /// A [`Blockstore`] implementation backed by an [IndexedDB] database.
 ///
@@ -50,22 +51,20 @@ impl Blockstore for IndexedDbBlockstore {
             .db
             .transaction(&[BLOCK_STORE], TransactionMode::ReadOnly)?;
         let blocks = tx.store(BLOCK_STORE)?;
-        let block = blocks.get(&cid).await?;
-
-        if block.is_undefined() {
-            Ok(None)
-        } else {
-            let arr = block.dyn_ref::<Uint8Array>().ok_or_else(|| {
-                Error::StoredDataError(format!(
-                    "expected 'Uint8Array', got '{}'",
-                    block
-                        .js_typeof()
-                        .as_string()
-                        .expect("typeof must be a string")
-                ))
-            })?;
-            Ok(Some(arr.to_vec()))
-        }
+        let Some(block) = blocks.get(cid.into()).await? else {
+            return Ok(None);
+        };
+
+        let arr = block.dyn_ref::<Uint8Array>().ok_or_else(|| {
+            Error::StoredDataError(format!(
+                "expected 'Uint8Array', got '{}'",
+                block
+                    .js_typeof()
+                    .as_string()
+                    .expect("typeof must be a string")
+            ))
+        })?;
+        Ok(Some(arr.to_vec()))
     }
 
     async fn put_keyed<const S: usize>(&self, cid: &CidGeneric<S>, data: &[u8]) -> Result<()> {
@@ -91,7 +90,7 @@ impl Blockstore for IndexedDbBlockstore {
             .transaction(&[BLOCK_STORE], TransactionMode::ReadWrite)?;
         let blocks = tx.store(BLOCK_STORE)?;
 
-        blocks.delete(&cid).await?;
+        blocks.delete(cid.into()).await?;
 
         Ok(())
     }
@@ -106,6 +105,44 @@ impl Blockstore for IndexedDbBlockstore {
 
         has_key(&blocks, &cid).await
     }
+
+    async fn retain<F>(&self, predicate: F) -> Result<()>
+    where
+        F: Fn(&[u8]) -> bool + 'static,
+    {
+        let tx = self
+            .db
+            .transaction(&[BLOCK_STORE], TransactionMode::ReadWrite)?;
+        let blocks = tx.store(BLOCK_STORE)?;
+        let mut last_key = None;
+        loop {
+            let keys = blocks
+                .get_all_keys(
+                    last_key
+                        .map(|key| KeyRange::lower_bound(&key, Some(true)))
+                        .transpose()?,
+                    Some(RETAIN_BATCH_SIZE),
+                )
+                .await?;
+            last_key = keys.last().cloned();
+            let count = keys.len();
+
+            for key in keys {
+                let cid = Uint8Array::new(&ArrayBuffer::from(key).into());
+                // TODO: can this copy be elided?
+                if !predicate(cid.to_vec().as_ref()) {
+                    blocks.delete(cid.into()).await?;
+                }
+            }
+
+            if count < RETAIN_BATCH_SIZE as usize {
+                break;
+            }
+        }
+
+        tx.commit().await?;
+        Ok(())
+    }
 }
 
 impl From<rexie::Error> for Error {
@@ -114,9 +151,15 @@ impl From<rexie::Error> for Error {
     }
 }
 
+impl From<idb::Error> for Error {
+    fn from(value: idb::Error) -> Self {
+        Error::FatalDatabaseError(value.to_string())
+    }
+}
+
 async fn has_key(store: &Store, key: &JsValue) -> Result<bool> {
     let key_range = KeyRange::only(key)?;
-    let count = store.count(Some(&key_range)).await?;
+    let count = store.count(Some(key_range)).await?;
     Ok(count > 0)
 }