feat(archive): automatically select bounded parallel encoding

crates/fjall/src/archive/mod.rs

@@ -261,6 +261,11 @@ impl ArchiveStore {
         })
     }
 
+    /// Automatic encoding decisions and peak per-batch encoding resources.
+    pub fn append_stats(&self) -> dolos_flatfiles::AppendStats {
+        self.flatfiles.append_stats()
+    }
+
     /// Get a reference to the underlying database
     pub fn database(&self) -> &Database {
         &self.db
@@ -458,9 +463,22 @@ pub struct ArchiveWriter {
     batch: Mutex<OwnedWriteBatch>,
     pending_blocks: Mutex<Vec<(ChainPoint, RawBlock)>>,
     overlay: Mutex<HashMap<BlockSlot, Vec<BlockLocation>>>,
+    #[cfg(test)]
+    fail_index_commit: bool,
 }
 
 impl ArchiveWriter {
+    fn new(store: &ArchiveStore) -> Self {
+        Self {
+            batch: Mutex::new(store.db.batch()),
+            store: store.clone(),
+            pending_blocks: Mutex::new(Vec::new()),
+            overlay: Mutex::new(HashMap::new()),
+            #[cfg(test)]
+            fail_index_commit: false,
+        }
+    }
+
     fn resolve_locations(
         &self,
         overlay: &HashMap<BlockSlot, Vec<BlockLocation>>,
@@ -644,6 +662,12 @@ impl CoreArchiveWriter for ArchiveWriter {
             }
         }
 
+        #[cfg(test)]
+        if self.fail_index_commit {
+            return Err(io_err(std::io::Error::other(
+                "injected index commit failure",
+            )));
+        }
         let batch = batch.durability(Some(PersistMode::Buffer));
         batch.commit().map_err(fjall_err)?;
 
@@ -865,12 +889,7 @@ impl CoreArchiveStore for ArchiveStore {
     type ExactIter = ExactIter;
 
     fn start_writer(&self) -> Result<Self::Writer, ArchiveError> {
-        Ok(ArchiveWriter {
-            batch: Mutex::new(self.db.batch()),
-            store: self.clone(),
-            pending_blocks: Mutex::new(Vec::new()),
-            overlay: Mutex::new(HashMap::new()),
-        })
+        Ok(ArchiveWriter::new(self))
     }
 
     fn read_logs(
@@ -1244,6 +1263,32 @@ mod tests {
         std::fs::read(store.flatfiles.segment_path(segment)).unwrap()
     }
 
+    #[test]
+    fn automatic_index_failure_leaves_only_unindexed_frames_and_retry_keeps_original_locations() {
+        let store = ArchiveStore::for_tempdir(StateSchema::default()).unwrap();
+        write(&store, &[(1, body(1, 0))]);
+        let original = locations(&store, 1);
+        let second = Arc::new(vec![2; 128 << 10]);
+        let third = Arc::new(vec![3; 128 << 10]);
+        let mut writer = store.start_writer().unwrap();
+        writer.apply(&point(2), &second).unwrap();
+        writer.apply(&point(3), &third).unwrap();
+        writer.fail_index_commit = true;
+        assert!(writer.commit().is_err());
+        assert!(locations(&store, 2).is_empty());
+        assert!(locations(&store, 3).is_empty());
+        let dead_end = segment_bytes(&store, 0).len() as u64;
+        let writer = store.start_writer().unwrap();
+        writer.apply(&point(1), &body(1, 0)).unwrap();
+        writer.apply(&point(2), &second).unwrap();
+        writer.apply(&point(3), &third).unwrap();
+        writer.commit().unwrap();
+        assert_eq!(locations(&store, 1), original);
+        assert!(locations(&store, 2)[0].offset >= dead_end);
+        assert_eq!(store.get_block_by_slot(&2).unwrap().unwrap(), *second);
+        assert_eq!(store.get_block_by_slot(&3).unwrap().unwrap(), *third);
+    }
+
     #[test]
     fn a_repeated_import_keeps_the_original_frame_and_leaves_the_copy_unnamed() {
         let store = ArchiveStore::for_tempdir(StateSchema::default()).unwrap();

crates/flatfiles/Cargo.toml

@@ -4,6 +4,7 @@ version.workspace = true
 edition.workspace = true
 
 [dependencies]
+rayon.workspace = true
 tempfile = "3"
 zstd = "0.13"
 

crates/flatfiles/src/codec.rs

@@ -27,6 +27,20 @@ pub const MAX_BODY_BYTES: usize = 16 << 20;
 /// back before the decoder is pooled.
 const POOLED_BUFFER_BYTES: usize = 1 << 20;
 
+/// The most bytes the frame for a `len`-byte body can take: what an
+/// encoder reserves before compressing it, and what an import window
+/// budgets for it before it is encoded.
+pub fn frame_bound(len: usize) -> usize {
+    zstd::zstd_safe::compress_bound(len)
+}
+
+pub(crate) fn oversized(len: usize) -> io::Error {
+    io::Error::new(
+        io::ErrorKind::InvalidInput,
+        format!("a {len}-byte body exceeds the {MAX_BODY_BYTES}-byte frame limit"),
+    )
+}
+
 fn encoder_dictionary() -> &'static EncoderDictionary<'static> {
     static DICTIONARY: OnceLock<EncoderDictionary<'static>> = OnceLock::new();
     DICTIONARY.get_or_init(|| EncoderDictionary::copy(BUNDLED_DICTIONARY, COMPRESSION_LEVEL))
@@ -41,6 +55,7 @@ fn decoder_dictionary() -> &'static DecoderDictionary<'static> {
 pub struct Encoder {
     compressor: Compressor<'static>,
     frame: Vec<u8>,
+    bounded: bool,
 }
 
 impl Encoder {
@@ -52,29 +67,39 @@ impl Encoder {
         Ok(Self {
             compressor,
             frame: Vec::new(),
+            bounded: false,
+        })
+    }
+
+    pub(crate) fn for_parallel() -> io::Result<Self> {
+        Ok(Self {
+            bounded: true,
+            ..Self::new()?
         })
     }
 
     /// Compress `body` into one frame, borrowed from the encoder's buffer
     /// until the next call.
     pub fn encode(&mut self, body: &[u8]) -> io::Result<&[u8]> {
         if body.len() > MAX_BODY_BYTES {
-            return Err(io::Error::new(
-                io::ErrorKind::InvalidInput,
-                format!(
-                    "a {}-byte body exceeds the {MAX_BODY_BYTES}-byte frame limit",
-                    body.len()
-                ),
-            ));
+            return Err(oversized(body.len()));
         }
         self.frame.clear();
-        let bound = zstd::zstd_safe::compress_bound(body.len());
+        let bound = frame_bound(body.len());
         if self.frame.capacity() < bound {
-            self.frame.reserve(bound);
+            if self.bounded {
+                self.frame.reserve_exact(bound);
+            } else {
+                self.frame.reserve(bound);
+            }
         }
         let n = self.compressor.compress_to_buffer(body, &mut self.frame)?;
         Ok(&self.frame[..n])
     }
+
+    pub(crate) fn capacity(&self) -> usize {
+        self.frame.capacity()
+    }
 }
 
 /// A reusable decompression context with its frame and body buffers.

crates/flatfiles/src/lib.rs

@@ -6,14 +6,16 @@
 //! by its physical offset and length inside its segment; the archive index
 //! holds those locations and this crate decodes the frame they name. The
 //! crate knows nothing about Cardano and depends on little beyond the
-//! standard library: `zstd` for the frames and `tempfile` for throwaway
-//! stores.
+//! standard library: `zstd` for the frames, `rayon` for eligible batches'
+//! parallel encoding and `tempfile` for throwaway stores.
 
 mod codec;
 mod store;
 
-pub use codec::{BUNDLED_DICTIONARY, COMPRESSION_LEVEL, MAX_BODY_BYTES};
-pub use store::{parse_segment_filename, FlatFileStore, ResourceStats};
+pub use codec::{frame_bound, BUNDLED_DICTIONARY, COMPRESSION_LEVEL, MAX_BODY_BYTES};
+pub use store::{
+    parse_segment_filename, AppendStats, FlatFileStore, ResourceStats, ENCODE_WINDOW_BYTES,
+};
 
 /// Number of slots per segment file (one Cardano epoch).
 pub const SLOTS_PER_SEGMENT: u64 = 432_000;