Nicer disk reading api and more api stuff and

benches/huge-car.rs

···
     let reader = tokio::fs::File::open(filename).await.unwrap();
     let reader = tokio::io::BufReader::new(reader);
 
-    let mb = 2_usize.pow(20);
-
-    let mut driver = match Driver::load_car(reader, |block| block.len(), 1024 * mb)
+    let mut driver = match Driver::load_car(reader, |block| block.len(), 1024)
         .await
         .unwrap()
     {
-        Driver::Lil(_, mem_driver) => mem_driver,
-        Driver::Big(_) => panic!("not doing disk for benchmark"),
+        Driver::Memory(_, mem_driver) => mem_driver,
+        Driver::Disk(_) => panic!("not doing disk for benchmark"),
     };
 
     let mut n = 0;

benches/non-huge-cars.rs

···
 }
 
 async fn drive_car(bytes: &[u8]) -> usize {
-    let mut driver = match Driver::load_car(bytes, |block| block.len(), 32 * 2_usize.pow(20))
+    let mut driver = match Driver::load_car(bytes, |block| block.len(), 32)
         .await
         .unwrap()
     {
-        Driver::Lil(_, mem_driver) => mem_driver,
-        Driver::Big(_) => panic!("not benching big cars here"),
+        Driver::Memory(_, mem_driver) => mem_driver,
+        Driver::Disk(_) => panic!("not benching big cars here"),
     };
 
     let mut n = 0;

examples/disk-read-file/main.rs

···
+/*!
+Read a CAR file by spilling to disk
+*/
+
 extern crate repo_stream;
 use clap::Parser;
-use repo_stream::{Driver, noop};
+use repo_stream::{Driver, process::noop};
 use std::path::PathBuf;
 
 #[derive(Debug, Parser)]
···
 
     // configure how much memory can be used before spilling to disk.
     // real memory usage may differ somewhat.
-    let in_mem_limit = 10 * 2_usize.pow(20);
+    let in_mem_limit = 10; // MiB
 
     // configure how much memory sqlite is allowed to use when dumping to disk
-    let db_cache_mb = 32;
+    let db_cache_mb = 32; // MiB
 
     log::info!("hello! reading the car...");
 
     // in this example we only bother handling CARs that are too big for memory
     // `noop` helper means: do no block processing, store the raw blocks
     let driver = match Driver::load_car(reader, noop, in_mem_limit).await? {
-        Driver::Lil(_, _) => panic!("try this on a bigger car"),
-        Driver::Big(big_stuff) => {
+        Driver::Memory(_, _) => panic!("try this on a bigger car"),
+        Driver::Disk(big_stuff) => {
             // we reach here if the repo was too big and needs to be spilled to
             // disk to continue
 
···
 
     log::info!("arrived! joining rx...");
 
+    // clean up the database. would be nice to do this in drop so it happens
+    // automatically, but some blocking work happens, so that's not allowed in
+    // async rust. 🤷‍♀️
     join.await?.reset_store().await?;
 
     log::info!("done. n={n} zeros={zeros}");

examples/read-file/main.rs

···
+/*!
+Read a CAR file with in-memory processing
+*/
+
 extern crate repo_stream;
 use clap::Parser;
 use repo_stream::Driver;
···
     let reader = tokio::io::BufReader::new(reader);
 
     let (commit, mut driver) =
-        match Driver::load_car(reader, |block| block.len(), 16 * 1024 * 1024).await? {
-            Driver::Lil(commit, mem_driver) => (commit, mem_driver),
-            Driver::Big(_) => panic!("can't handle big cars yet"),
+        match Driver::load_car(reader, |block| block.len(), 16 /* MiB */).await? {
+            Driver::Memory(commit, mem_driver) => (commit, mem_driver),
+            Driver::Disk(_) => panic!("this example doesn't handle big CARs"),
         };
 
     log::info!("got commit: {commit:?}");
···
         n += pairs.len();
         // log::info!("got {rkey:?}");
     }
-    log::info!("bye! {n}");
+    log::info!("bye! total records={n}");
 
     Ok(())
 }

src/disk.rs

···
         // let insert_stmt = tx.prepare("INSERT INTO blocks (key, val) VALUES (?1, ?2)")?;
         Ok(SqliteWriter { tx })
     }
-    pub fn get_reader(&'_ self) -> Result<SqliteReader<'_>, rusqlite::Error> {
+    pub fn get_reader<'conn>(&'conn self) -> Result<SqliteReader<'conn>, rusqlite::Error> {
         let select_stmt = self.conn.prepare("SELECT val FROM blocks WHERE key = ?1")?;
         Ok(SqliteReader { select_stmt })
     }

src/drive.rs

···
 }
 
 pub enum Driver<R: AsyncRead + Unpin, T: Processable> {
-    Lil(Commit, MemDriver<T>),
-    Big(BigCar<R, T>),
+    Memory(Commit, MemDriver<T>),
+    Disk(BigCar<R, T>),
 }
 
 impl<R: AsyncRead + Unpin, T: Processable> Driver<R, T> {
     pub async fn load_car(
         reader: R,
         process: fn(Vec<u8>) -> T,
-        max_size: usize,
+        max_size_mb: usize,
     ) -> Result<Driver<R, T>, DriveError> {
+        let max_size = max_size_mb * 2_usize.pow(20);
         let mut mem_blocks = HashMap::new();
 
         let mut car = CarReader::new(reader).await?;
···
             mem_size += std::mem::size_of::<Cid>() + maybe_processed.get_size();
             mem_blocks.insert(cid, maybe_processed);
             if mem_size >= max_size {
-                return Ok(Driver::Big(BigCar {
+                return Ok(Driver::Disk(BigCar {
                     car,
                     root,
                     process,
···
 
         let walker = Walker::new(commit.data);
 
-        Ok(Driver::Lil(
+        Ok(Driver::Memory(
             commit,
             MemDriver {
                 blocks: mem_blocks,
···
             commit,
             BigCarReady {
                 process: self.process,
-                store,
-                walker,
+                state: Some(BigState { store, walker }),
             },
         ))
     }
 }
 
+struct BigState {
+    store: SqliteStore,
+    walker: Walker,
+}
+
 pub struct BigCarReady<T: Clone> {
     process: fn(Vec<u8>) -> T,
-    store: SqliteStore,
-    walker: Walker,
+    state: Option<BigState>,
 }
 
 impl<T: Processable + Send + 'static> BigCarReady<T> {
-    pub async fn next_chunk(
-        mut self,
-        n: usize,
-    ) -> Result<(Self, Option<BlockChunk<T>>), DriveError> {
-        let mut out = Vec::with_capacity(n);
-        (self, out) = tokio::task::spawn_blocking(move || {
-            let store = self.store;
-            let mut reader = store.get_reader()?;
+    pub async fn next_chunk(&mut self, n: usize) -> Result<Option<BlockChunk<T>>, DriveError> {
+        let process = self.process;
+
+        // state should only *ever* be None transiently while inside here
+        let mut state = self
+            .state
+            .take()
+            .expect("BigCarReady must have Some(state)");
 
-            for _ in 0..n {
-                // walk as far as we can until we run out of blocks or find a record
-                match self.walker.disk_step(&mut reader, self.process)? {
-                    Step::Missing(cid) => return Err(DriveError::MissingBlock(cid)),
-                    Step::Finish => break,
-                    Step::Found { rkey, data } => {
-                        out.push((rkey, data));
-                        continue;
+        // the big pain here is that we don't want to leave self.state in an
+        // invalid state (None), so all the error paths have to make sure it
+        // comes out again.
+        let (state, res) = tokio::task::spawn_blocking(
+            move || -> (BigState, Result<BlockChunk<T>, DriveError>) {
+                let mut reader_res = state.store.get_reader();
+                let reader: &mut _ = match reader_res {
+                    Ok(ref mut r) => r,
+                    Err(ref mut e) => {
+                        // unfortunately we can't return the error directly because
+                        // (for some reason) it's attached to the lifetime of the
+                        // reader?
+                        // hack a mem::swap so we can get it out :/
+                        let mut e_swapped =
+                            rusqlite::Error::InvalidParameterName("this error was stolen".into());
+                        std::mem::swap(e, &mut e_swapped);
+                        // the pain: `state` *has to* outlive the reader
+                        drop(reader_res);
+                        return (state, Err(e_swapped.into()));
                     }
                 };
-            }
 
-            drop(reader); // cannot outlive store
-            self.store = store;
-            Ok::<_, DriveError>((self, out))
-        })
-        .await??;
+                let mut out = Vec::with_capacity(n);
+
+                for _ in 0..n {
+                    // walk as far as we can until we run out of blocks or find a record
+                    let step = match state.walker.disk_step(reader, process) {
+                        Ok(s) => s,
+                        Err(e) => {
+                            // the pain: `state` *has to* outlive the reader
+                            drop(reader_res);
+                            return (state, Err(e.into()));
+                        }
+                    };
+                    match step {
+                        Step::Missing(cid) => {
+                            // the pain: `state` *has to* outlive the reader
+                            drop(reader_res);
+                            return (state, Err(DriveError::MissingBlock(cid)));
+                        }
+                        Step::Finish => break,
+                        Step::Found { rkey, data } => out.push((rkey, data)),
+                    };
+                }
+
+                // `state` *has to* outlive the reader
+                drop(reader_res);
+
+                (state, Ok::<_, DriveError>(out))
+            },
+        )
+        .await?; // on tokio JoinError, we'll be left with invalid state :(
+
+        // *must* restore state before dealing with the actual result
+        self.state = Some(state);
+
+        let out = res?;
 
         if out.is_empty() {
-            Ok((self, None))
+            Ok(None)
         } else {
-            Ok((self, Some(out)))
+            Ok(Some(out))
         }
     }
 
···
         n: usize,
         tx: mpsc::Sender<Result<BlockChunk<T>, DriveError>>,
     ) -> Result<(), mpsc::error::SendError<Result<BlockChunk<T>, DriveError>>> {
-        let mut reader = match self.store.get_reader() {
+        let BigState { store, walker } = self.state.as_mut().expect("valid state");
+        let mut reader = match store.get_reader() {
             Ok(r) => r,
             Err(e) => return tx.blocking_send(Err(e.into())),
         };
···
             for _ in 0..n {
                 // walk as far as we can until we run out of blocks or find a record
 
-                let step = match self.walker.disk_step(&mut reader, self.process) {
+                let step = match walker.disk_step(&mut reader, self.process) {
                     Ok(s) => s,
                     Err(e) => return tx.blocking_send(Err(e.into())),
                 };
···
 
     pub async fn reset_store(mut self) -> Result<SqliteStore, DriveError> {
         tokio::task::spawn_blocking(move || {
-            self.store.reset()?;
-            Ok(self.store)
+            let BigState { mut store, .. } = self.state.take().expect("valid state");
+            store.reset()?;
+            Ok(store)
         })
         .await?
     }

src/lib.rs

···
-//! Fast and robust atproto CAR file processing in rust
-//!
-//! For now see the [examples](https://tangled.org/@microcosm.blue/repo-stream/tree/main/examples)
+/*!
+A robust CAR file -> MST walker for atproto
+
+Small CARs have their blocks buffered in memory. If a configurable memory limit
+is reached while reading blocks, CAR reading is suspended, and can be continued
+by providing disk storage to buffer the CAR blocks instead.
+
+A `process` function can be provided for tasks where records are transformed
+into a smaller representation, to save memory (and disk) during block reading.
+
+Once blocks are loaded, the MST is walked and emitted as chunks of pairs of
+`(rkey, processed_block)` pairs, in order (depth first, left-to-right).
+
+Some MST validations are applied
+- Keys must appear in order
+- Keys must be at the correct MST tree depth
+
+`iroh_car` additionally applies a block size limit of `2MiB`.
+
+```
+use repo_stream::{Driver, SqliteStore};
+
+# #[tokio::main]
+# async fn main() -> Result<(), Box<dyn std::error::Error>> {
+# let reader = include_bytes!("../car-samples/tiny.car").as_slice();
+let mut total_size = 0;
+let process = |rec: Vec<u8>| rec.len(); // block processing: just extract the size
+let in_mem_limit = 10; /* MiB */
+let db_cache_size = 32; /* MiB */
+
+match Driver::load_car(reader, process, in_mem_limit).await? {
+
+    // if all blocks fit within memory
+    Driver::Memory(_commit, mut driver) => {
+        while let Some(chunk) = driver.next_chunk(256).await? {
+            for (_rkey, size) in chunk {
+                total_size += size;
+            }
+        }
+    },
+
+    // if the CAR was too big for in-memory processing
+    Driver::Disk(paused) => {
+        // set up a disk store we can spill to
+        let store = SqliteStore::new("some/path.sqlite".into(), db_cache_size).await?;
+        // do the spilling, get back a (similar) driver
+        let (_commit, mut driver) = paused.finish_loading(store).await?;
+
+        while let Some(chunk) = driver.next_chunk(256).await? {
+            for (_rkey, size) in chunk {
+                total_size += size;
+            }
+        }
+
+        // clean up the disk store (drop tables etc)
+        driver.reset_store().await?;
+    }
+};
+println!("sum of size of all records: {total_size}");
+# Ok(())
+# }
+```
+
+Find more [examples in the repo](https://tangled.org/@microcosm.blue/repo-stream/tree/main/examples).
+
+*/
 
 mod mst;
 mod walk;
···
 
 pub use disk::SqliteStore;
 pub use drive::{DriveError, Driver};
-pub use process::{Processable, noop};
+pub use process::Processable;

tests/non-huge-cars.rs

···
 const MIDSIZE_CAR: &'static [u8] = include_bytes!("../car-samples/midsize.car");
 
 async fn test_car(bytes: &[u8], expected_records: usize, expected_sum: usize) {
-    let mb = 2_usize.pow(20);
-
-    let mut driver = match Driver::load_car(bytes, |block| block.len(), 10 * mb)
+    let mut driver = match Driver::load_car(bytes, |block| block.len(), 10 /* MiB */)
         .await
         .unwrap()
     {
-        Driver::Lil(_commit, mem_driver) => mem_driver,
-        Driver::Big(_) => panic!("too big"),
+        Driver::Memory(_commit, mem_driver) => mem_driver,
+        Driver::Disk(_) => panic!("too big"),
     };
 
     let mut records = 0;