batched walking

examples/disk-read-file/main.rs

···
 
     let reader = CarReader::new(reader).await?;
 
-    let redb_store = repo_stream::disk_redb::RedbStore::new(tmpfile)?;
+    let redb_store = repo_stream::disk_redb::RedbStore::new(tmpfile).await?;
 
     let root = reader
         .header()

src/disk_drive.rs

···
 use std::collections::VecDeque;
 use std::error::Error;
 
-use crate::disk_walk::{Step, Trip, Walker};
+use crate::disk_walk::{Trip, Walker};
 use crate::mst::Commit;
 
 use ipld_core::cid::Cid;
···
 
 /// Errors that can happen while consuming and emitting blocks and records
 #[derive(Debug, thiserror::Error)]
-pub enum DriveError {
+pub enum DriveError<E: Error> {
     #[error("Failed to initialize CarReader: {0}")]
     CarReader(#[from] iroh_car::Error),
     #[error("Car block stream error: {0}")]
···
     MissingBlock(Cid),
     #[error("Failed to walk the mst tree: {0}")]
     Tripped(#[from] Trip),
+    #[error("whatever: {0}")]
+    WalkingProblem(#[from] WalkError),
+    #[error("whatever: {0}")]
+    Boooooo(String),
+    #[error("processing error: {0}")]
+    ProcessingError(E),
 }
 
+/// Limited subset of errors that can happen while walking
+#[derive(Debug, thiserror::Error)]
+pub enum WalkError {
+    #[error("The MST block {0} could not be found")]
+    MissingBlock(Cid),
+    #[error("Failed to walk the mst tree: {0}")]
+    Tripped(#[from] Trip),
+}
+
+/// Storage backend for caching large-repo blocks
+///
+/// Since
 pub trait BlockStore<MPB: Serialize + DeserializeOwned> {
-    fn put_batch(&self, blocks: Vec<(Cid, MPB)>) -> impl std::future::Future<Output = ()> + Send; // unwraps for now
-    fn get(&self, key: Cid) -> Option<MPB>;
+    fn put_batch(&self, blocks: Vec<(Cid, MPB)>) -> impl Future<Output = ()> + Send; // unwraps for now
+    fn walk_batch(
+        &self,
+        walker: Walker,
+        n: usize,
+    ) -> impl Future<Output = Result<(Walker, Vec<(String, MPB)>), String>>; // boo string error for now because
 }
 
 type CarBlock<E> = Result<(Cid, Vec<u8>), E>;
···
 impl<SE, S, T, BS, P, PE> Vehicle<SE, S, T, BS, P, PE>
 where
     SE: Error + 'static,
-    S: Stream<Item = CarBlock<SE>> + Unpin,
-    T: Clone + Serialize + DeserializeOwned,
-    BS: BlockStore<Vec<u8>>,
-    P: Fn(&[u8]) -> Result<T, PE>,
+    S: Stream<Item = CarBlock<SE>> + Unpin + Send,
+    T: Clone + Serialize + DeserializeOwned + Send,
+    BS: BlockStore<Vec<u8>> + Send,
+    P: Fn(&[u8]) -> Result<T, PE> + Send,
     PE: Error,
 {
     /// Set up the stream
···
         block_stream: S,
         block_store: BS,
         process: P,
-    ) -> Result<(Commit, Self), DriveError> {
+    ) -> Result<(Commit, Self), DriveError<PE>> {
         let mut commit = None;
 
         log::warn!("init: load blocks");
···
         Ok((commit, me))
     }
 
-    async fn load_chunk(&mut self, n: usize) -> Result<(), DriveError> {
-        self.out_cache.reserve(n);
-        for _ in 0..n {
-            let item = match self.walker.step(&mut self.block_store, &self.process)? {
-                Step::Step { rkey, data } => (rkey, data),
-                Step::Finish => break,
-                Step::Rest(cid) => return Err(DriveError::MissingBlock(cid)),
-            };
-            self.out_cache.push_back(item);
-        }
+    async fn load_chunk(&mut self, n: usize) -> Result<(), DriveError<PE>> {
+        let walker = std::mem::take(&mut self.walker);
+        let (walker, batch) = self
+            .block_store
+            .walk_batch(walker, n)
+            .await
+            .map_err(DriveError::Boooooo)?;
+        self.walker = walker;
+
+        let processed = batch
+            .into_iter()
+            .map(|(k, raw)| (self.process)(&raw).map(|t| (k, t)))
+            .collect::<Result<Vec<_>, _>>()
+            .map_err(DriveError::ProcessingError)?;
+
+        self.out_cache.extend(processed);
         Ok(())
     }
 
···
     /// (but non-zero), even if it's not the last chunk.
     ///
     /// an empty vec will be returned to signal the end.
-    pub async fn next_chunk(&mut self, n: usize) -> Result<Vec<(String, T)>, DriveError> {
+    pub async fn next_chunk(&mut self, n: usize) -> Result<Vec<(String, T)>, DriveError<PE>> {
         if self.out_cache.is_empty() {
             self.load_chunk(n).await?;
         }
···
     }
 
     /// Manually step through the record outputs
-    pub async fn next_record(&mut self) -> Result<Option<(String, T)>, DriveError> {
+    pub async fn next_record(&mut self) -> Result<Option<(String, T)>, DriveError<PE>> {
         if self.out_cache.is_empty() {
             self.load_chunk(64).await?; // TODO
         }
···
     }
 
     /// Convert to a futures::stream of record outputs
-    pub fn stream(self) -> impl Stream<Item = Result<(String, T), DriveError>> {
+    pub fn stream(self) -> impl Stream<Item = Result<(String, T), DriveError<PE>>> {
         futures::stream::try_unfold(self, |mut this| async move {
             let maybe_record = this.next_record().await?;
             Ok(maybe_record.map(|b| (b, this)))

src/disk_redb.rs

···
 use crate::disk_drive::BlockStore;
+use crate::disk_walk::{Need, Walker};
 use ipld_core::cid::Cid;
 use redb::{Database, Durability, Error, ReadableDatabase, TableDefinition};
 use std::path::Path;
···
 }
 
 impl RedbStore {
-    pub fn new(path: impl AsRef<Path>) -> Result<Self, Error> {
+    pub async fn new(path: impl AsRef<Path> + 'static + Send) -> Result<Self, Error> {
         log::warn!("redb new");
-        let db = Database::create(path)?;
+        let db = tokio::task::spawn_blocking(|| Database::create(path))
+            .await
+            .unwrap()?;
         log::warn!("db created");
         Ok(Self { db: db.into() })
     }
···
         .unwrap();
     }
 
-    fn get(&self, c: Cid) -> Option<Vec<u8>> {
-        let key_bytes = c.to_bytes();
-        let tx = self.db.begin_read().unwrap();
-        let table = tx.open_table(TABLE).unwrap();
-        let t = table.get(&*key_bytes).unwrap()?.value().to_vec();
-        Some(t)
+    async fn walk_batch(
+        &self,
+        mut walker: Walker,
+        n: usize,
+    ) -> Result<(Walker, Vec<(String, Vec<u8>)>), String> {
+        let db = self.db.clone();
+        tokio::task::spawn_blocking(move || {
+            let tx = db.begin_read().unwrap();
+            let table = tx.open_table(TABLE).unwrap();
+
+            let mut out = Vec::with_capacity(n);
+            loop {
+                let Some(need) = walker.next_needed() else {
+                    break;
+                };
+                let cid = need.cid();
+                let Some(res) = table.get(&*cid.to_bytes()).unwrap() else {
+                    return Err(format!("missing block: {cid:?}"));
+                };
+                let block = res.value();
+
+                match need {
+                    Need::Node(_) => walker
+                        .handle_node(block)
+                        .map_err(|e| format!("failed to handle mst node: {e}"))?,
+                    Need::Record { rkey, .. } => {
+                        out.push((rkey, block.to_vec()));
+                        if out.len() >= n {
+                            break;
+                        }
+                    }
+                }
+            }
+            Ok((walker, out))
+        })
+        .await
+        .unwrap() // tokio join
     }
 }

src/disk_walk.rs

···
 //! Depth-first MST traversal
 
-use crate::disk_drive::BlockStore;
 use crate::mst::Node;
+use std::convert::Infallible;
 
 use ipld_core::cid::Cid;
 use serde::{Serialize, de::DeserializeOwned};
-use std::error::Error;
 
 /// Errors that can happen while walking
 #[derive(Debug, thiserror::Error)]
···
     #[error("empty mst nodes are not allowed")]
     NodeEmpty,
     #[error("Failed to decode commit block: {0}")]
-    BadCommit(Box<dyn std::error::Error>),
+    BadCommit(serde_ipld_dagcbor::DecodeError<Infallible>),
     #[error("Action node error: {0}")]
     RkeyError(#[from] RkeyError),
     #[error("Process failed: {0}")]
···
 }
 
 #[derive(Debug, Clone, PartialEq)]
-enum Need {
+pub enum Need {
     Node(Cid),
     Record { rkey: String, cid: Cid },
 }
 
+impl Need {
+    pub fn cid(&self) -> Cid {
+        match self {
+            Need::Node(cid) => *cid,
+            Need::Record { cid, .. } => *cid,
+        }
+    }
+}
+
 fn push_from_node(stack: &mut Vec<Need>, node: &Node) -> Result<(), RkeyError> {
     let mut entries = Vec::with_capacity(node.entries.len());
 
···
 /// Traverser of an atproto MST
 ///
 /// Walks the tree from left-to-right in depth-first order
-#[derive(Debug)]
+///
+/// (turning into more of a navigator)
+/// it doesn't quite feel like the divisions of responsibility are right around
+/// here yet.
+#[derive(Debug, Default)]
 pub struct Walker {
     stack: Vec<Need>,
     prev: String,
···
         }
     }
 
-    /// Advance through nodes until we find a record or can't go further
-    pub fn step<T: Clone + Serialize + DeserializeOwned, E: Error>(
-        &mut self,
-        block_store: &mut impl BlockStore<Vec<u8>>,
-        process: impl Fn(&[u8]) -> Result<T, E>,
-    ) -> Result<Step<T>, Trip> {
-        loop {
-            let Some(mut need) = self.stack.last() else {
-                log::trace!("tried to walk but we're actually done.");
-                return Ok(Step::Finish);
-            };
+    pub fn next_needed(&mut self) -> Option<Need> {
+        self.stack.pop()
+        // TODO:
+        // let need = self.stack.pop()?;
+        // if let Need::Record { ref rkey, .. } = need {
+        //     // rkeys *must* be in order or else the tree is invalid (or
+        //     // we have a bug)
+        //     if *rkey <= self.prev {
+        //         return Err(Trip::RkeyOutOfOrder);
+        //     }
+        //     self.prev = rkey.clone();
+        // }
+        // Some(need)
+    }
 
-            match &mut need {
-                Need::Node(cid) => {
-                    log::trace!("need node {cid:?}");
-                    let Some(block) = block_store.get(*cid) else {
-                        log::trace!("node not found, resting");
-                        return Ok(Step::Rest(*cid));
-                    };
-
-                    let node = serde_ipld_dagcbor::from_slice::<Node>(&block)
-                        .map_err(|e| Trip::BadCommit(e.into()))?;
-
-                    // found node, make sure we remember
-                    self.stack.pop();
-
-                    // queue up work on the found node next
-                    push_from_node(&mut self.stack, &node)?;
-                }
-                Need::Record { rkey, cid } => {
-                    log::trace!("need record {cid:?}");
-                    let Some(block) = block_store.get(*cid) else {
-                        log::trace!("record block not found, resting");
-                        return Ok(Step::Rest(*cid));
-                    };
-                    let rkey = rkey.clone();
-
-                    let data = process(&block).map_err(|e| Trip::ProcessFailed(e.to_string()));
-
-                    // found node, make sure we remember
-                    self.stack.pop();
-
-                    log::trace!("emitting a block as a step. depth={}", self.stack.len());
-
-                    let data = data.map_err(|e| Trip::ProcessFailed(e.to_string()))?;
-
-                    // rkeys *must* be in order or else the tree is invalid (or
-                    // we have a bug)
-                    if rkey <= self.prev {
-                        return Err(Trip::RkeyOutOfOrder);
-                    }
-                    self.prev = rkey.clone();
-
-                    return Ok(Step::Step { rkey, data });
-                }
-            }
-        }
+    /// hacky: this must be called after next_needed if it was a node
+    pub fn handle_node(&mut self, block: &[u8]) -> Result<(), Trip> {
+        let node = serde_ipld_dagcbor::from_slice::<Node>(block).map_err(Trip::BadCommit)?;
+        push_from_node(&mut self.stack, &node)?;
+        Ok(())
     }
 }