From 52c22060f3954178ba1640da13326cb05daf5e30 Mon Sep 17 00:00:00 2001 From: Lewis Date: Sun, 12 Apr 2026 10:38:31 +0300 Subject: [PATCH] fix(tranquil-store): preserve refcount in hint relocate records Lewis: May this revision serve well! --- .../src/blockstore/compaction.rs | 9 +- .../src/blockstore/hash_index.rs | 33 ++- crates/tranquil-store/src/blockstore/hint.rs | 23 +- .../tests/compaction_liveness.rs | 254 ++++++++++++++++++ 4 files changed, 304 insertions(+), 15 deletions(-) create mode 100644 crates/tranquil-store/tests/compaction_liveness.rs diff --git a/crates/tranquil-store/src/blockstore/compaction.rs b/crates/tranquil-store/src/blockstore/compaction.rs index 7b2c722..0a21ea8 100644 --- a/crates/tranquil-store/src/blockstore/compaction.rs +++ b/crates/tranquil-store/src/blockstore/compaction.rs @@ -162,7 +162,13 @@ fn stream_compact( } => match index.get(&cid_bytes) { Some(e) if e.location.file_id == source_file_id && !e.refcount.is_zero() => { let loc = writer.append_block(&cid_bytes, &data)?; - hint_writer.append_relocate(&cid_bytes, loc.file_id, loc.offset, loc.length)?; + hint_writer.append_relocate( + &cid_bytes, + loc.file_id, + loc.offset, + loc.length, + e.refcount.raw(), + )?; relocations.push((cid_bytes, loc)); live_count = live_count.saturating_add(1); } @@ -187,6 +193,7 @@ fn stream_compact( loc.file_id, loc.offset, loc.length, + e.refcount.raw(), )?; relocations.push((cid_bytes, loc)); live_count = live_count.saturating_add(1); diff --git a/crates/tranquil-store/src/blockstore/hash_index.rs b/crates/tranquil-store/src/blockstore/hash_index.rs index c157624..5046d7e 100644 --- a/crates/tranquil-store/src/blockstore/hash_index.rs +++ b/crates/tranquil-store/src/blockstore/hash_index.rs @@ -378,6 +378,7 @@ impl HashTable { &mut self, cid: &[u8; CID_SIZE], new_location: BlockLocation, + refcount: RefCount, ) -> Result { if is_empty(cid) { return Ok(false); @@ -394,7 +395,9 @@ impl HashTable { let slot_cid = self.slots[idx].cid; if is_empty(&slot_cid) { - self.slots[idx] = Slot::from_location(*cid, new_location); + let mut slot = Slot::from_location(*cid, new_location); + slot.refcount = refcount; + self.slots[idx] = slot; self.count += 1; return Ok(false); } @@ -411,6 +414,7 @@ impl HashTable { let slot_dist = self.probe_distance(idx, slot_home); if slot_dist < dist { let mut displaced = Slot::from_location(*cid, new_location); + displaced.refcount = refcount; std::mem::swap(&mut self.slots[idx], &mut displaced); self.count += 1; self.relocate_displaced(displaced, idx, slot_dist); @@ -571,7 +575,7 @@ impl HashTable { removals: &[CidBytes], ) { relocations.iter().for_each(|(cid, new_loc)| { - if let Err(e) = self.relocate(cid, *new_loc) { + if let Err(e) = self.relocate(cid, *new_loc, RefCount::one()) { tracing::error!(?e, "capacity exhausted during compaction relocation"); } }); @@ -1267,18 +1271,20 @@ impl BlockIndex { pub fn batch_relocate( &self, - relocations: &[(CidBytes, BlockLocation)], + relocations: &[(CidBytes, BlockLocation, u32)], ) -> Result<(), BlockIndexError> { if relocations.is_empty() { return Ok(()); } let mut table = self.table.write(); - relocations.iter().try_for_each(|(cid, location)| { - table - .relocate(cid, *location) - .map(|_| ()) - .map_err(|_| BlockIndexError::CapacityExhausted) - }) + relocations + .iter() + .try_for_each(|(cid, location, refcount)| { + table + .relocate(cid, *location, RefCount::new(*refcount)) + .map(|_| ()) + .map_err(|_| BlockIndexError::CapacityExhausted) + }) } pub fn batch_remove(&self, cids: &[CidBytes]) { @@ -1499,15 +1505,18 @@ impl BlockIndex { file_id, offset, length, + refcount, } => { let loc = BlockLocation { file_id, offset, length, }; - table.relocate(&cid_bytes, loc).map_err(|_| { - io::Error::other("hash table capacity exhausted during rebuild") - })?; + table + .relocate(&cid_bytes, loc, RefCount::new(refcount)) + .map_err(|_| { + io::Error::other("hash table capacity exhausted during rebuild") + })?; } ReadHintRecord::Remove { cid_bytes } => { let _ = table.remove(&cid_bytes); diff --git a/crates/tranquil-store/src/blockstore/hint.rs b/crates/tranquil-store/src/blockstore/hint.rs index 4ec098a..a035763 100644 --- a/crates/tranquil-store/src/blockstore/hint.rs +++ b/crates/tranquil-store/src/blockstore/hint.rs @@ -78,6 +78,8 @@ pub(crate) fn encode_hint_record( write_hint_record(io, fd, write_offset, &record) } +const REFCOUNT_OFFSET: usize = 2; + pub(crate) fn encode_relocate_record( io: &S, fd: FileId, @@ -86,10 +88,13 @@ pub(crate) fn encode_relocate_record( file_id: DataFileId, block_offset: BlockOffset, length: BlockLength, + refcount: u32, ) -> io::Result<()> { let mut record = [0u8; HINT_RECORD_SIZE]; record[TYPE_OFFSET] = RECORD_TYPE_RELOCATE; record[VERSION_OFFSET] = HINT_FORMAT_VERSION; + let rc16 = u16::try_from(refcount).unwrap_or(u16::MAX); + record[REFCOUNT_OFFSET..REFCOUNT_OFFSET + 2].copy_from_slice(&rc16.to_le_bytes()); record[CID_OFFSET..CID_OFFSET + CID_SIZE].copy_from_slice(cid_bytes); record[FIELD_A_OFFSET..FIELD_A_OFFSET + 4].copy_from_slice(&file_id.raw().to_le_bytes()); record[FIELD_A_OFFSET + 4..FIELD_A_OFFSET + 8].copy_from_slice(&length.raw().to_le_bytes()); @@ -158,6 +163,7 @@ pub enum ReadHintRecord { file_id: DataFileId, offset: BlockOffset, length: BlockLength, + refcount: u32, }, Remove { cid_bytes: [u8; CID_SIZE], @@ -255,6 +261,15 @@ pub fn decode_hint_record( })) } RECORD_TYPE_RELOCATE => { + let rc16 = u16::from_le_bytes( + record[REFCOUNT_OFFSET..REFCOUNT_OFFSET + 2] + .try_into() + .unwrap(), + ); + let refcount = match rc16 { + 0 => 1, + n => u32::from(n), + }; let file_id = DataFileId::new(u32::from_le_bytes( record[FIELD_A_OFFSET..FIELD_A_OFFSET + 4] .try_into() @@ -278,6 +293,7 @@ pub fn decode_hint_record( file_id, offset: block_offset, length: BlockLength::new(raw_length), + refcount, })) } RECORD_TYPE_REMOVE => Ok(Some(ReadHintRecord::Remove { cid_bytes })), @@ -341,6 +357,7 @@ impl<'a, S: StorageIO> HintFileWriter<'a, S> { file_id: DataFileId, offset: BlockOffset, length: BlockLength, + refcount: u32, ) -> io::Result<()> { encode_relocate_record( self.io, @@ -350,6 +367,7 @@ impl<'a, S: StorageIO> HintFileWriter<'a, S> { file_id, offset, length, + refcount, )?; self.position = self.position.advance(HINT_RECORD_SIZE as u64); Ok(()) @@ -575,7 +593,7 @@ pub fn replay_hints_into_block_index( let mut replayed: u64 = 0; let mut put_buffer: Vec<([u8; CID_SIZE], BlockLocation)> = Vec::with_capacity(REPLAY_BATCH_SIZE); - let mut relocate_buffer: Vec<([u8; CID_SIZE], BlockLocation)> = + let mut relocate_buffer: Vec<([u8; CID_SIZE], BlockLocation, u32)> = Vec::with_capacity(REPLAY_BATCH_SIZE); let mut remove_buffer: Vec<[u8; CID_SIZE]> = Vec::with_capacity(REPLAY_BATCH_SIZE); @@ -663,13 +681,14 @@ pub fn replay_hints_into_block_index( file_id, offset, length, + refcount, } => { let loc = BlockLocation { file_id, offset, length, }; - relocate_buffer.push((cid_bytes, loc)); + relocate_buffer.push((cid_bytes, loc, refcount)); let record_end = offset.advance(BLOCK_RECORD_OVERHEAD as u64 + length.as_u64()); diff --git a/crates/tranquil-store/tests/compaction_liveness.rs b/crates/tranquil-store/tests/compaction_liveness.rs new file mode 100644 index 0000000..3f06b4b --- /dev/null +++ b/crates/tranquil-store/tests/compaction_liveness.rs @@ -0,0 +1,254 @@ +mod common; + +use std::collections::HashSet; + +use tranquil_store::blockstore::{ + BlockStoreConfig, CidBytes, GroupCommitConfig, TranquilBlockStore, +}; + +fn tiny_store_config(dir: &std::path::Path) -> BlockStoreConfig { + BlockStoreConfig { + data_dir: dir.join("data"), + index_dir: dir.join("index"), + max_file_size: 4096, + group_commit: GroupCommitConfig { + checkpoint_interval_ms: 600_000, + checkpoint_write_threshold: 1_000_000, + ..GroupCommitConfig::default() + }, + shard_count: 1, + } +} + +fn make_block(seed: u32, size: usize) -> (CidBytes, Vec) { + ( + common::test_cid(seed), + common::block_data(seed) + .into_iter() + .cycle() + .take(size) + .collect(), + ) +} + +fn verify_live_blocks(store: &TranquilBlockStore, live: &HashSet, context: &str) { + let missing: Vec = live + .iter() + .copied() + .filter(|&seed| { + store + .get_block_sync(&common::test_cid(seed)) + .unwrap() + .is_none() + }) + .collect(); + + assert!( + missing.is_empty(), + "{context}: {count} live blocks missing from store: {missing:?}", + count = missing.len(), + ); +} + +fn compact_sealed(store: &TranquilBlockStore) { + let files = store.list_data_files().unwrap(); + files + .iter() + .copied() + .take(files.len().saturating_sub(1)) + .for_each(|fid| { + let _ = store.compact_file(fid, 0); + }); +} + +fn delete_checkpoints(index_dir: &std::path::Path) { + let _ = std::fs::remove_file(index_dir.join("checkpoint_a.tqc")); + let _ = std::fs::remove_file(index_dir.join("checkpoint_b.tqc")); +} + +#[test] +fn relocate_loses_refcount_on_hint_rebuild() { + common::with_runtime(|| { + let dir = tempfile::TempDir::new().unwrap(); + + let target = common::test_cid(1); + let target_data = vec![0xABu8; 200]; + + { + let store = TranquilBlockStore::open(tiny_store_config(dir.path())).unwrap(); + + store + .put_blocks_blocking(vec![(target, target_data.clone())]) + .unwrap(); + store + .put_blocks_blocking(vec![(target, target_data.clone())]) + .unwrap(); + + let padding: Vec<_> = (100..130u32).map(|s| make_block(s, 300)).collect(); + store.put_blocks_blocking(padding).unwrap(); + + std::thread::sleep(std::time::Duration::from_millis(10)); + compact_sealed(&store); + + store.apply_commit_blocking(vec![], vec![target]).unwrap(); + + let data = store.get_block_sync(&target).unwrap(); + assert!(data.is_some(), "target should be live, refcount 2 - 1 = 1"); + } + + delete_checkpoints(&dir.path().join("index")); + + { + let store = TranquilBlockStore::open(tiny_store_config(dir.path())).unwrap(); + + let data = store.get_block_sync(&target).unwrap(); + assert!( + data.is_some(), + "BUG: target missing after hint-only rebuild. \ + RELOCATE created entry with refcount 1 instead of 2, \ + then DEC brought it to 0 instead of 1." + ); + + std::thread::sleep(std::time::Duration::from_millis(10)); + compact_sealed(&store); + + let data = store.get_block_sync(&target).unwrap(); + assert!( + data.is_some(), + "BUG: target removed by compaction after hint rebuild \ + incorrectly set refcount to 0" + ); + } + }); +} + +#[test] +fn multi_restart_with_compaction_between_put_and_dec() { + common::with_runtime(|| { + let dir = tempfile::TempDir::new().unwrap(); + + let shared = common::test_cid(42); + let shared_data = vec![0xCDu8; 200]; + + { + let store = TranquilBlockStore::open(tiny_store_config(dir.path())).unwrap(); + + store + .put_blocks_blocking(vec![(shared, shared_data.clone())]) + .unwrap(); + store + .put_blocks_blocking(vec![(shared, shared_data.clone())]) + .unwrap(); + store + .put_blocks_blocking(vec![(shared, shared_data.clone())]) + .unwrap(); + + let filler: Vec<_> = (200..230u32).map(|s| make_block(s, 300)).collect(); + store.put_blocks_blocking(filler).unwrap(); + } + + delete_checkpoints(&dir.path().join("index")); + + { + let store = TranquilBlockStore::open(tiny_store_config(dir.path())).unwrap(); + + let data = store.get_block_sync(&shared).unwrap(); + assert!( + data.is_some(), + "round 1: shared block present after rebuild" + ); + + std::thread::sleep(std::time::Duration::from_millis(10)); + compact_sealed(&store); + + store.apply_commit_blocking(vec![], vec![shared]).unwrap(); + + let data = store.get_block_sync(&shared).unwrap(); + assert!( + data.is_some(), + "round 1: shared block should survive, refcount 3 - 1 = 2" + ); + } + + delete_checkpoints(&dir.path().join("index")); + + { + let store = TranquilBlockStore::open(tiny_store_config(dir.path())).unwrap(); + + let data = store.get_block_sync(&shared).unwrap(); + assert!( + data.is_some(), + "round 2: shared block should survive hint rebuild, refcount should be 2" + ); + + store.apply_commit_blocking(vec![], vec![shared]).unwrap(); + + let data = store.get_block_sync(&shared).unwrap(); + assert!( + data.is_some(), + "round 2: shared block should survive DEC, refcount 2 - 1 = 1" + ); + + std::thread::sleep(std::time::Duration::from_millis(10)); + compact_sealed(&store); + + let data = store.get_block_sync(&shared).unwrap(); + assert!( + data.is_some(), + "BUG: shared block removed by compaction. \ + Multiple restarts with RELOCATE collapsed refcount \ + from 3 down to 1, two DECs made it 0." + ); + } + }); +} + +#[test] +fn stress_create_delete_restart_cycle_matches_bug_report() { + common::with_runtime(|| { + let dir = tempfile::TempDir::new().unwrap(); + let mut live: HashSet = HashSet::new(); + let mut rng = common::Rng::new(12345); + let mut next_seed: u32 = 0; + + (0..4).for_each(|cycle| { + { + let store = TranquilBlockStore::open(tiny_store_config(dir.path())).unwrap(); + + (0..20).for_each(|_| { + let seed_a = next_seed; + let seed_b = next_seed + 1; + next_seed += 2; + + store + .put_blocks_blocking(vec![make_block(seed_a, 150), make_block(seed_b, 150)]) + .unwrap(); + live.insert(seed_a); + live.insert(seed_b); + + if rng.next_u32() % 2 == 0 { + let victim: Option = live.iter().copied().next(); + if let Some(v) = victim { + store + .apply_commit_blocking(vec![], vec![common::test_cid(v)]) + .unwrap(); + live.remove(&v); + } + } + }); + + std::thread::sleep(std::time::Duration::from_millis(10)); + compact_sealed(&store); + + verify_live_blocks(&store, &live, &format!("cycle {cycle} before kill")); + } + + delete_checkpoints(&dir.path().join("index")); + + { + let store = TranquilBlockStore::open(tiny_store_config(dir.path())).unwrap(); + verify_live_blocks(&store, &live, &format!("cycle {cycle} after hint rebuild")); + } + }); + }); +}