refactor(redb-store): Optimization for small file import in redb store

iroh-bytes/src/store/file/import_flat_store.rs

@@ -210,7 +210,8 @@ impl ActorState {
         }
 
         let txn = db.begin_write()?;
-        let mut tables = Tables::new(&txn)?;
+        let mut delete_after_commit = Default::default();
+        let mut tables = Tables::new(&txn, &mut delete_after_commit)?;
         for (hash, entry) in index {
             if tables.blobs.get(hash)?.is_some() {
                 tracing::info!("hash {} already exists in the db", hash.to_hex());
@@ -232,15 +233,15 @@ impl ActorState {
                     None
                 };
                 if let Err(cause) =
-                    std::fs::rename(data_path, self.path_options.owned_data_path(&hash))
+                    std::fs::rename(data_path, self.options.path.owned_data_path(&hash))
                 {
                     tracing::error!("failed to move data file: {}", cause);
                     continue;
                 }
                 if let Some(outboard_path) = outboard_path {
                     if let Err(cause) = copy_outboard(
                         &outboard_path,
-                        &self.path_options.owned_outboard_path(&hash),
+                        &self.options.path.owned_outboard_path(&hash),
                     ) {
                         tracing::error!("failed to move outboard file: {}", cause);
                         continue;
@@ -281,7 +282,7 @@ impl ActorState {
                 if let Some(outboard_path) = outboard_path {
                     if let Err(cause) = copy_outboard(
                         &outboard_path,
-                        &self.path_options.owned_outboard_path(&hash),
+                        &self.options.path.owned_outboard_path(&hash),
                     ) {
                         tracing::error!("failed to move outboard file: {}", cause);
                         continue;
@@ -321,15 +322,15 @@ impl ActorState {
                     None
                 };
                 if let Err(cause) =
-                    std::fs::rename(data_path, self.path_options.owned_data_path(&hash))
+                    std::fs::rename(data_path, self.options.path.owned_data_path(&hash))
                 {
                     tracing::error!("failed to move data file: {}", cause);
                     continue;
                 }
                 if let Some(outboard_path) = outboard_path {
                     if let Err(cause) = copy_outboard(
                         &outboard_path,
-                        &self.path_options.owned_outboard_path(&hash),
+                        &self.options.path.owned_outboard_path(&hash),
                     ) {
                         tracing::error!("failed to move outboard file: {}", cause);
                         continue;

iroh-bytes/src/store/file/tables.rs

@@ -1,9 +1,11 @@
 //! Table definitions and accessors for the redb database.
+use std::collections::BTreeSet;
+
 use redb::{ReadableTable, TableDefinition, TableError};
 
 use iroh_base::hash::{Hash, HashAndFormat};
 
-use super::EntryState;
+use super::{EntryState, PathOptions};
 use crate::util::Tag;
 
 pub(super) const BLOBS_TABLE: TableDefinition<Hash, EntryState> = TableDefinition::new("blobs-0");
@@ -33,15 +35,27 @@ pub(super) struct Tables<'a, 'b> {
     pub tags: redb::Table<'a, 'b, Tag, HashAndFormat>,
     pub inline_data: redb::Table<'a, 'b, Hash, &'static [u8]>,
     pub inline_outboard: redb::Table<'a, 'b, Hash, &'static [u8]>,
+    pub delete_after_commit: &'b mut DeleteSet,
+}
+
+#[derive(Debug, PartialEq, Eq, PartialOrd, Ord)]
+pub(super) enum BaoFilePart {
+    Outboard,
+    Data,
+    Sizes,
 }
 
 impl<'db, 'txn> Tables<'db, 'txn> {
-    pub fn new(tx: &'txn redb::WriteTransaction<'db>) -> std::result::Result<Self, TableError> {
+    pub fn new(
+        tx: &'txn redb::WriteTransaction<'db>,
+        delete_after_commit: &'txn mut DeleteSet,
+    ) -> std::result::Result<Self, TableError> {
         Ok(Self {
             blobs: tx.open_table(BLOBS_TABLE)?,
             tags: tx.open_table(TAGS_TABLE)?,
             inline_data: tx.open_table(INLINE_DATA_TABLE)?,
             inline_outboard: tx.open_table(INLINE_OUTBOARD_TABLE)?,
+            delete_after_commit,
         })
     }
 }
@@ -95,3 +109,54 @@ impl ReadableTables for ReadOnlyTables<'_> {
         &self.inline_outboard
     }
 }
+
+/// Helper to keep track of files to delete after a transaction is committed.
+#[derive(Debug, Default)]
+pub(super) struct DeleteSet(BTreeSet<(Hash, BaoFilePart)>);
+
+impl DeleteSet {
+    /// Mark a file as to be deleted after the transaction is committed.
+    pub fn insert(&mut self, hash: Hash, parts: impl IntoIterator<Item = BaoFilePart>) {
+        for part in parts {
+            self.0.insert((hash, part));
+        }
+    }
+
+    /// Mark a file as to be kept after the transaction is committed.
+    ///
+    /// This will cancel any previous delete for the same file in the same transaction.
+    pub fn remove(&mut self, hash: Hash, parts: impl IntoIterator<Item = BaoFilePart>) {
+        for part in parts {
+            self.0.remove(&(hash, part));
+        }
+    }
+
+    /// Get the inner set of files to delete.
+    pub fn into_inner(self) -> BTreeSet<(Hash, BaoFilePart)> {
+        self.0
+    }
+
+    /// Apply the delete set and clear it.
+    ///
+    /// This will delete all files marked for deletion and then clear the set.
+    /// Errors will just be logged.
+    pub fn apply_and_clear(&mut self, options: &PathOptions) {
+        for (hash, to_delete) in &self.0 {
+            tracing::info!("deleting {:?}", to_delete);
+            let path = match to_delete {
+                BaoFilePart::Data => options.owned_data_path(hash),
+                BaoFilePart::Outboard => options.owned_outboard_path(hash),
+                BaoFilePart::Sizes => options.owned_sizes_path(hash),
+            };
+            if let Err(cause) = std::fs::remove_file(&path) {
+                tracing::warn!(
+                    "failed to delete {:?} {}: {}",
+                    to_delete,
+                    path.display(),
+                    cause
+                );
+            }
+        }
+        self.0.clear();
+    }
+}

iroh-bytes/src/store/file/test_support.rs

@@ -139,7 +139,7 @@ impl StoreInner {
             .send_async(ActorMessage::Blobs { filter, tx })
             .await?;
         let blobs = rx.await?;
-        let res = blobs
+        let res = blobs?
             .into_iter()
             .map(|r| {
                 r.map(|(hash, _)| hash)
@@ -163,9 +163,9 @@ impl ActorState {
         tables: &impl ReadableTables,
         hash: Hash,
     ) -> ActorResult<Option<EntryData>> {
-        let data_path = self.path_options.owned_data_path(&hash);
-        let outboard_path = self.path_options.owned_outboard_path(&hash);
-        let sizes_path = self.path_options.owned_sizes_path(&hash);
+        let data_path = self.options.path.owned_data_path(&hash);
+        let outboard_path = self.options.path.owned_outboard_path(&hash);
+        let sizes_path = self.options.path.owned_sizes_path(&hash);
         let entry = match tables.blobs().get(hash)? {
             Some(guard) => match guard.value() {
                 EntryState::Complete {
@@ -243,9 +243,9 @@ impl ActorState {
         hash: Hash,
         entry: Option<EntryData>,
     ) -> ActorResult<()> {
-        let data_path = self.path_options.owned_data_path(&hash);
-        let outboard_path = self.path_options.owned_outboard_path(&hash);
-        let sizes_path = self.path_options.owned_sizes_path(&hash);
+        let data_path = self.options.path.owned_data_path(&hash);
+        let outboard_path = self.options.path.owned_outboard_path(&hash);
+        let sizes_path = self.options.path.owned_sizes_path(&hash);
         // tabula rasa
         std::fs::remove_file(&outboard_path).ok();
         std::fs::remove_file(&data_path).ok();
@@ -260,15 +260,15 @@ impl ActorState {
         let entry = match entry {
             EntryData::Complete { data, outboard } => {
                 let data_size = data.len() as u64;
-                let data_location = if data_size > self.inline_options.max_data_inlined {
+                let data_location = if data_size > self.options.inline.max_data_inlined {
                     std::fs::write(data_path, &data)?;
                     DataLocation::Owned(data_size)
                 } else {
                     tables.inline_data.insert(hash, data.as_slice())?;
                     DataLocation::Inline(())
                 };
                 let outboard_size = outboard.len() as u64;
-                let outboard_location = if outboard_size > self.inline_options.max_outboard_inlined
+                let outboard_location = if outboard_size > self.options.inline.max_outboard_inlined
                 {
                     std::fs::write(outboard_path, &outboard)?;
                     OutboardLocation::Owned

iroh-bytes/src/store/file/tests.rs

@@ -51,6 +51,7 @@ async fn create_test_db() -> (tempfile::TempDir, Store) {
     let db_path = testdir.path().join("db.redb");
     let options = Options {
         path: PathOptions::new(testdir.path()),
+        batch: Default::default(),
         inline: Default::default(),
     };
     let db = Store::new(db_path, options).await.unwrap();
@@ -803,6 +804,7 @@ async fn actor_store_smoke() {
     let db_path = testdir.path().join("test.redb");
     let options = Options {
         path: PathOptions::new(testdir.path()),
+        batch: Default::default(),
         inline: Default::default(),
     };
     let db = Store::new(db_path, options).await.unwrap();

iroh-bytes/src/store/file/util.rs

@@ -2,6 +2,7 @@ use std::{
     fs::OpenOptions,
     io::{self, Write},
     path::Path,
+    time::{Duration, Instant},
 };
 
 /// A reader that calls a callback with the number of bytes read after each read.
@@ -63,3 +64,120 @@ pub fn read_and_remove(path: &Path) -> io::Result<Vec<u8>> {
     std::fs::remove_file(path)?;
     Ok(data)
 }
+
+/// A wrapper for a flume receiver that allows peeking at the next message.
+#[derive(Debug)]
+pub(super) struct PeekableFlumeReceiver<T> {
+    msg: Option<T>,
+    recv: flume::Receiver<T>,
+}
+
+#[allow(dead_code)]
+impl<T> PeekableFlumeReceiver<T> {
+    pub fn new(recv: flume::Receiver<T>) -> Self {
+        Self { msg: None, recv }
+    }
+
+    /// Peek at the next message.
+    ///
+    /// Will block if there are no messages.
+    /// Returns None only if there are no more messages (sender is dropped).
+    pub fn peek(&mut self) -> Option<&T> {
+        if self.msg.is_none() {
+            self.msg = self.recv.recv().ok();
+        }
+        self.msg.as_ref()
+    }
+
+    /// Receive the next message.
+    ///
+    /// Will block if there are no messages.
+    /// Returns None only if there are no more messages (sender is dropped).
+    pub fn recv(&mut self) -> Option<T> {
+        if let Some(msg) = self.msg.take() {
+            return Some(msg);
+        }
+        self.recv.recv().ok()
+    }
+
+    /// Try to peek at the next message.
+    ///
+    /// Will not block.
+    /// Returns None if reading would block, or if there are no more messages (sender is dropped).
+    pub fn try_peek(&mut self) -> Option<&T> {
+        if self.msg.is_none() {
+            self.msg = self.recv.try_recv().ok();
+        }
+        self.msg.as_ref()
+    }
+
+    /// Try to receive the next message.
+    ///
+    /// Will not block.
+    /// Returns None if reading would block, or if there are no more messages (sender is dropped).
+    pub fn try_recv(&mut self) -> Option<T> {
+        if let Some(msg) = self.msg.take() {
+            return Some(msg);
+        }
+        self.recv.try_recv().ok()
+    }
+
+    pub fn recv_timeout(&mut self, timeout: std::time::Duration) -> Option<T> {
+        if let Some(msg) = self.msg.take() {
+            return Some(msg);
+        }
+        self.recv.recv_timeout(timeout).ok()
+    }
+
+    /// Create an iterator that pulls messages from the receiver for at most
+    /// `count` messages or `max_duration` time.
+    pub fn batch_iter(&mut self, count: usize, max_duration: Duration) -> BatchIter<T> {
+        BatchIter::new(self, count, max_duration)
+    }
+
+    /// Push back a message. This will only work if there is room for it.
+    /// Otherwise, it will fail and return the message.
+    pub fn push_back(&mut self, msg: T) -> std::result::Result<(), T> {
+        if self.msg.is_none() {
+            self.msg = Some(msg);
+            Ok(())
+        } else {
+            Err(msg)
+        }
+    }
+}
+
+pub(super) struct BatchIter<'a, T> {
+    recv: &'a mut PeekableFlumeReceiver<T>,
+    start: Instant,
+    remaining: usize,
+    max_duration: Duration,
+}
+
+impl<'a, T> BatchIter<'a, T> {
+    fn new(recv: &'a mut PeekableFlumeReceiver<T>, count: usize, max_duration: Duration) -> Self {
+        Self {
+            recv,
+            start: Instant::now(),
+            remaining: count,
+            max_duration,
+        }
+    }
+}
+
+impl<'a, T> Iterator for BatchIter<'a, T> {
+    type Item = T;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        if self.remaining == 0 {
+            return None;
+        }
+        let elapsed = self.start.elapsed();
+        if elapsed >= self.max_duration {
+            return None;
+        }
+        let remaining_time = self.max_duration - elapsed;
+        self.remaining -= 1;
+        self.recv.recv_timeout(remaining_time)
+    }
+}