fix(tranquil-store): checkpoint-hint race & missing dedup hints

Lewis: May this revision serve well! <lu5a@proton.me>

crates/tranquil-store/Cargo.toml

@@ -52,6 +52,7 @@ k256 = { workspace = true }
 rand = { workspace = true }
 tikv-jemallocator = "0.6"
 tracing-subscriber = { workspace = true, features = ["env-filter"] }
+libc = "0.2"
 
 [[bench]]
 name = "blockstore"

crates/tranquil-store/src/blockstore/compaction.rs

@@ -93,8 +93,7 @@ pub(super) fn compact_on_writer_thread<S: StorageIO>(
             Err(e)
         }
         Ok((new_size, live_count, dead_count)) => {
-            let positions = hint_positions.snapshot();
-            if let Err(e) = index.write_checkpoint(epoch.current(), &positions) {
+            if let Err(e) = index.write_checkpoint(epoch.current(), hint_positions) {
                 tracing::warn!(error = %e, "pre-delete checkpoint failed during compaction");
             }
 

crates/tranquil-store/src/blockstore/group_commit.rs

@@ -813,8 +813,7 @@ fn maybe_checkpoint(
     if !elapsed && !threshold {
         return;
     }
-    let positions = hint_positions.snapshot();
-    match index.write_checkpoint(epoch.current(), &positions) {
+    match index.write_checkpoint(epoch.current(), hint_positions) {
         Ok(()) => {
             *last_checkpoint = std::time::Instant::now();
             *writes_since_checkpoint = 0;
@@ -831,8 +830,7 @@ fn shutdown_checkpoint(
     epoch: &EpochCounter,
     hint_positions: &ShardHintPositions,
 ) {
-    let positions = hint_positions.snapshot();
-    match index.write_checkpoint(epoch.current(), &positions) {
+    match index.write_checkpoint(epoch.current(), hint_positions) {
         Ok(()) => tracing::debug!("shutdown checkpoint written"),
         Err(e) => tracing::warn!(error = %e, "shutdown checkpoint failed"),
     }
@@ -924,8 +922,6 @@ fn commit_loop<S: StorageIO>(
             if let Ok((ref dedup, _)) = result {
                 writes_since_checkpoint =
                     writes_since_checkpoint.saturating_add(dedup.len() as u64);
-                ctx.hint_positions
-                    .update(ctx.shard_id, state.file_id, state.hint_position);
             }
 
             dispatch_responses(drain.entries, result.map(|(dedup, _proof)| dedup));
@@ -1024,8 +1020,6 @@ fn drain_and_process_remaining<S: StorageIO>(
 
         if let Ok((ref _dedup, ref proof)) = result {
             run_post_sync_hook(post_sync_hook, proof);
-            ctx.hint_positions
-                .update(ctx.shard_id, state.file_id, state.hint_position);
         }
 
         dispatch_responses(entries, result.map(|(dedup, _proof)| dedup));
@@ -1102,6 +1096,7 @@ fn process_batch<S: StorageIO>(
             let location = match dedup.get(cid_bytes) {
                 Some(&loc) => {
                     dedup_hits = dedup_hits.saturating_add(1);
+                    hint_writer.append_hint(cid_bytes, &loc)?;
                     loc
                 }
                 None => {
@@ -1194,7 +1189,19 @@ fn process_batch<S: StorageIO>(
     };
     let t = std::time::Instant::now();
     index
-        .batch_put(&index_entries, &all_decrements, cursor, current_epoch, now)
+        .batch_put_and_advance_position(
+            &index_entries,
+            &all_decrements,
+            cursor,
+            current_epoch,
+            now,
+            super::hash_index::PositionUpdate {
+                hint_positions: &ctx.hint_positions,
+                shard_id: ctx.shard_id,
+                file_id: state.file_id,
+                offset: state.hint_position,
+            },
+        )
         .map_err(CommitError::from)?;
     let index_nanos = t.elapsed().as_nanos() as u64;
 

crates/tranquil-store/src/blockstore/hash_index.rs

@@ -5,11 +5,19 @@ use std::path::{Path, PathBuf};
 use parking_lot::RwLock;
 
 use super::data_file::CID_SIZE;
+use super::group_commit::ShardHintPositions;
 use super::types::{
     BlockLength, BlockLocation, BlockOffset, CidBytes, CollectionResult, CommitEpoch, DataFileId,
-    HintOffset, IndexEntry, LivenessInfo, RefCount, WallClockMs, WriteCursor,
+    HintOffset, IndexEntry, LivenessInfo, RefCount, ShardId, WallClockMs, WriteCursor,
 };
 
+pub struct PositionUpdate<'a> {
+    pub hint_positions: &'a ShardHintPositions,
+    pub shard_id: ShardId,
+    pub file_id: DataFileId,
+    pub offset: HintOffset,
+}
+
 const EMPTY_CID: [u8; CID_SIZE] = [0u8; CID_SIZE];
 
 fn is_empty(cid: &[u8; CID_SIZE]) -> bool {
@@ -1197,6 +1205,30 @@ impl BlockIndex {
         cursor: WriteCursor,
         epoch: CommitEpoch,
         now: WallClockMs,
+    ) -> Result<(), BlockIndexError> {
+        self.batch_put_inner(entries, decrements, cursor, epoch, now, None)
+    }
+
+    pub fn batch_put_and_advance_position(
+        &self,
+        entries: &[([u8; CID_SIZE], BlockLocation)],
+        decrements: &[[u8; CID_SIZE]],
+        cursor: WriteCursor,
+        epoch: CommitEpoch,
+        now: WallClockMs,
+        position_update: PositionUpdate<'_>,
+    ) -> Result<(), BlockIndexError> {
+        self.batch_put_inner(entries, decrements, cursor, epoch, now, Some(position_update))
+    }
+
+    fn batch_put_inner(
+        &self,
+        entries: &[([u8; CID_SIZE], BlockLocation)],
+        decrements: &[[u8; CID_SIZE]],
+        cursor: WriteCursor,
+        epoch: CommitEpoch,
+        now: WallClockMs,
+        position_update: Option<PositionUpdate<'_>>,
     ) -> Result<(), BlockIndexError> {
         let mut table = self.table.write();
 
@@ -1217,6 +1249,12 @@ impl BlockIndex {
         });
 
         table.set_write_cursor(cursor);
+
+        if let Some(pos) = position_update {
+            pos.hint_positions
+                .update(pos.shard_id, pos.file_id, pos.offset);
+        }
+
         Ok(())
     }
 
@@ -1418,6 +1456,17 @@ impl BlockIndex {
     }
 
     pub fn write_checkpoint(
+        &self,
+        epoch: CommitEpoch,
+        hint_positions: &ShardHintPositions,
+    ) -> io::Result<()> {
+        let _guard = self.checkpoint_lock.lock();
+        let table = self.table.read();
+        let positions = hint_positions.snapshot();
+        write_checkpoint_ab(&table, &self.index_dir, epoch, &positions)
+    }
+
+    pub fn write_checkpoint_with_positions(
         &self,
         epoch: CommitEpoch,
         positions: &CheckpointPositions,

crates/tranquil-store/tests/checkpoint_race.rs

@@ -0,0 +1,223 @@
+mod common;
+
+use std::io;
+use std::sync::Arc;
+use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
+
+use tranquil_store::blockstore::{
+    BlockStoreConfig, BlocksSynced, CidBytes, GroupCommitConfig, TranquilBlockStore,
+};
+use tranquil_store::PostBlockstoreHook;
+
+struct SlowHook;
+
+impl PostBlockstoreHook for SlowHook {
+    fn on_blocks_synced(&self, _proof: &BlocksSynced) -> io::Result<()> {
+        std::thread::sleep(std::time::Duration::from_millis(1));
+        Ok(())
+    }
+}
+
+fn refcount(store: &TranquilBlockStore, cid: &CidBytes) -> Option<u32> {
+    store.block_index().get(cid).map(|e| e.refcount.raw())
+}
+
+fn race_config(dir: &std::path::Path) -> BlockStoreConfig {
+    BlockStoreConfig {
+        data_dir: dir.join("data"),
+        index_dir: dir.join("index"),
+        max_file_size: 256 * 1024,
+        group_commit: GroupCommitConfig {
+            checkpoint_interval_ms: 10,
+            checkpoint_write_threshold: 20,
+            ..GroupCommitConfig::default()
+        },
+        shard_count: 4,
+    }
+}
+
+fn cid_for(shard: u8, seq: u32) -> CidBytes {
+    let mut cid = [0u8; 36];
+    cid[0] = 0x01;
+    cid[1] = 0x71;
+    cid[2] = 0x12;
+    cid[3] = 0x20;
+    cid[4] = shard;
+    cid[8..12].copy_from_slice(&seq.to_le_bytes());
+    (12..36).for_each(|i| cid[i] = (seq as u8).wrapping_add(i as u8));
+    cid
+}
+
+fn write_phase(base: &std::path::Path, use_hook: bool) -> Vec<CidBytes> {
+    let config = race_config(base);
+    let hook: Option<Arc<dyn PostBlockstoreHook>> = use_hook.then(|| Arc::new(SlowHook) as _);
+    let store = Arc::new(TranquilBlockStore::open_with_hook(config, hook).unwrap());
+
+    let running = Arc::new(AtomicBool::new(true));
+    let total_cycles = Arc::new(AtomicU64::new(0));
+
+    let writers: Vec<_> = (0..4u8)
+        .map(|shard| {
+            let store = Arc::clone(&store);
+            let running = Arc::clone(&running);
+            let total_cycles = Arc::clone(&total_cycles);
+            std::thread::spawn(move || {
+                let mut targets = Vec::new();
+                let mut seq = 0u32;
+                while running.load(Ordering::Relaxed) {
+                    let cid = cid_for(shard, seq);
+                    store
+                        .put_blocks_blocking(vec![(cid, vec![shard; 60])])
+                        .unwrap();
+                    store
+                        .put_blocks_blocking(vec![(cid, vec![shard; 60])])
+                        .unwrap();
+                    store
+                        .apply_commit_blocking(vec![], vec![cid])
+                        .unwrap();
+                    targets.push(cid);
+                    seq += 1;
+                    total_cycles.fetch_add(1, Ordering::Relaxed);
+                }
+                targets
+            })
+        })
+        .collect();
+
+    while total_cycles.load(Ordering::Relaxed) < 500 {
+        std::thread::yield_now();
+    }
+
+    running.store(false, Ordering::Relaxed);
+
+    let all_targets: Vec<CidBytes> = writers
+        .into_iter()
+        .flat_map(|w| w.join().unwrap())
+        .collect();
+
+    all_targets.iter().for_each(|cid| {
+        assert_eq!(refcount(&store, cid), Some(1), "pre-crash sanity");
+    });
+
+    let store = Arc::try_unwrap(store).ok().unwrap();
+    std::mem::forget(store);
+
+    all_targets
+}
+
+fn verify_phase(base: &std::path::Path, targets: &[CidBytes]) -> usize {
+    let config = race_config(base);
+    let store = TranquilBlockStore::open(config).unwrap();
+    let bad = targets
+        .iter()
+        .filter(|cid| refcount(&store, cid) != Some(1))
+        .count();
+    drop(store);
+    bad
+}
+
+#[test]
+fn crash_recovery_preserves_refcounts() {
+    common::with_runtime(|| {
+        let mut corrupted = 0u32;
+        let total = 20u32;
+
+        (0..total).for_each(|_| {
+            let dir = tempfile::TempDir::new().unwrap();
+            let exe = std::env::current_exe().unwrap();
+            let dir_str = dir.path().to_str().unwrap();
+
+            let output = std::process::Command::new(&exe)
+                .arg("--exact")
+                .arg("__crash_write_phase")
+                .env("CRASH_TEST_DIR", dir_str)
+                .env("CRASH_TEST_HOOK", "0")
+                .output()
+                .unwrap();
+
+            assert!(output.status.success() || output.status.code() == Some(0));
+
+            let target_bytes = std::fs::read(dir.path().join("targets.bin")).unwrap();
+            let targets: Vec<CidBytes> = target_bytes
+                .chunks_exact(36)
+                .map(|chunk| {
+                    let mut cid = [0u8; 36];
+                    cid.copy_from_slice(chunk);
+                    cid
+                })
+                .collect();
+
+            if verify_phase(dir.path(), &targets) > 0 {
+                corrupted += 1;
+            }
+        });
+
+        assert_eq!(
+            corrupted, 0,
+            "{corrupted}/{total} iterations had refcount corruption after crash recovery"
+        );
+    });
+}
+
+#[test]
+fn crash_with_slow_hook_preserves_refcounts() {
+    common::with_runtime(|| {
+        let mut corrupted = 0u32;
+        let total = 20u32;
+
+        (0..total).for_each(|_| {
+            let dir = tempfile::TempDir::new().unwrap();
+            let exe = std::env::current_exe().unwrap();
+            let dir_str = dir.path().to_str().unwrap();
+
+            let output = std::process::Command::new(&exe)
+                .arg("--exact")
+                .arg("__crash_write_phase")
+                .env("CRASH_TEST_DIR", dir_str)
+                .env("CRASH_TEST_HOOK", "1")
+                .output()
+                .unwrap();
+
+            assert!(output.status.success() || output.status.code() == Some(0));
+
+            let target_bytes = std::fs::read(dir.path().join("targets.bin")).unwrap();
+            let targets: Vec<CidBytes> = target_bytes
+                .chunks_exact(36)
+                .map(|chunk| {
+                    let mut cid = [0u8; 36];
+                    cid.copy_from_slice(chunk);
+                    cid
+                })
+                .collect();
+
+            if verify_phase(dir.path(), &targets) > 0 {
+                corrupted += 1;
+            }
+        });
+
+        assert_eq!(
+            corrupted, 0,
+            "{corrupted}/{total} iterations had refcount corruption after crash with slow hook"
+        );
+    });
+}
+
+#[test]
+fn __crash_write_phase() {
+    let dir = match std::env::var("CRASH_TEST_DIR") {
+        Ok(d) => d,
+        Err(_) => return,
+    };
+    let use_hook = std::env::var("CRASH_TEST_HOOK").map(|v| v == "1").unwrap_or(false);
+    let base = std::path::Path::new(&dir);
+
+    let rt = tokio::runtime::Runtime::new().unwrap();
+    let _guard = rt.enter();
+
+    let targets = write_phase(base, use_hook);
+
+    let target_bytes: Vec<u8> = targets.iter().flat_map(|cid| cid.iter().copied()).collect();
+    std::fs::write(base.join("targets.bin"), &target_bytes).unwrap();
+
+    unsafe { libc::_exit(0) }
+}

crates/tranquil-store/tests/sim_blockstore.rs

@@ -112,7 +112,7 @@ impl SimHarness {
             HintOffset::new(entries.len() as u64 * HINT_RECORD_SIZE as u64),
         );
         index
-            .write_checkpoint(CommitEpoch::zero(), &positions)
+            .write_checkpoint_with_positions(CommitEpoch::zero(), &positions)
             .unwrap();
     }