feat(store): detect foreign&corrupt blocks on read & preserve blocks thru recovery

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

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

@@ -707,6 +707,22 @@ impl HashTable {
         })
     }
 
+    pub fn indexed_file_ends(&self) -> HashMap<DataFileId, BlockOffset> {
+        self.iter().fold(HashMap::new(), |mut ends, s| {
+            let end = s
+                .offset
+                .advance(super::data_file::BLOCK_RECORD_OVERHEAD as u64 + s.length.as_u64());
+            ends.entry(s.file_id)
+                .and_modify(|cur| {
+                    if end > *cur {
+                        *cur = end;
+                    }
+                })
+                .or_insert(end);
+            ends
+        })
+    }
+
     pub fn find_leaked_refcounts(
         &self,
         is_reachable: impl Fn(&CidBytes) -> bool,
@@ -1504,6 +1520,10 @@ impl BlockIndex {
             .liveness_by_file(current_epoch, now, grace_period_ms)
     }
 
+    pub fn indexed_file_ends(&self) -> HashMap<DataFileId, BlockOffset> {
+        self.table.read().indexed_file_ends()
+    }
+
     pub fn find_leaked_refcounts(
         &self,
         is_reachable: impl Fn(&CidBytes) -> bool,

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

@@ -6,6 +6,7 @@ pub mod hash_index;
 mod hint;
 mod manager;
 mod reader;
+mod repair;
 mod store;
 mod types;
 
@@ -25,7 +26,8 @@ pub use hint::{
     ReadHintRecord, RebuildError, decode_hint_record, hint_file_path, scan_hints_to_memory,
 };
 pub use manager::{CachedHandle, DEFAULT_MAX_FILE_SIZE, DataFileManager};
-pub use reader::{BlockStoreReader, ReadError};
+pub use reader::{BLOCK_CORRUPTION_MARKER, BlockStoreReader, ReadError};
+pub use repair::{RepairOutcome, rebuild_and_repair_mst};
 pub use store::QuiesceGuard;
 pub use store::{BlockStoreConfig, DEFAULT_SHARD_COUNT, OpenRetryPolicy, TranquilBlockStore};
 pub use types::{

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

@@ -11,6 +11,8 @@ use super::hash_index::BlockIndex;
 use super::manager::DataFileManager;
 use super::types::{BlockLocation, BlockOffset, DataFileId};
 
+pub const BLOCK_CORRUPTION_MARKER: &str = "corrupted block at";
+
 #[derive(Debug, Clone)]
 pub enum ReadError {
     Io(Arc<io::Error>),
@@ -25,7 +27,7 @@ impl std::fmt::Display for ReadError {
         match self {
             Self::Io(e) => write!(f, "io: {e}"),
             Self::Corrupted { file_id, offset } => {
-                write!(f, "corrupted block at {file_id}:{}", offset.raw())
+                write!(f, "{BLOCK_CORRUPTION_MARKER} {file_id}:{}", offset.raw())
             }
         }
     }
@@ -71,7 +73,7 @@ impl<S: StorageIO> BlockStoreReader<S> {
 
     pub fn get(&self, cid: &[u8; CID_SIZE]) -> Result<Option<Bytes>, ReadError> {
         match self.index.get(cid) {
-            Some(e) => self.read_block_at(e.location).map(Some),
+            Some(e) => self.read_block_at(e.location, cid).map(Some),
             None => Ok(None),
         }
     }
@@ -83,10 +85,10 @@ impl<S: StorageIO> BlockStoreReader<S> {
     pub fn get_many(&self, cids: &[[u8; CID_SIZE]]) -> Result<Vec<Option<Bytes>>, ReadError> {
         let mut results: Vec<Option<Bytes>> = vec![None; cids.len()];
 
-        let index_lookups: Vec<(usize, BlockLocation)> = cids
+        let index_lookups: Vec<(usize, [u8; CID_SIZE], BlockLocation)> = cids
             .iter()
             .enumerate()
-            .filter_map(|(i, cid)| self.index.get(cid).map(|entry| (i, entry.location)))
+            .filter_map(|(i, cid)| self.index.get(cid).map(|entry| (i, *cid, entry.location)))
             .collect();
         self.read_locations_into(&index_lookups, &mut results)?;
 
@@ -95,30 +97,38 @@ impl<S: StorageIO> BlockStoreReader<S> {
 
     fn read_locations_into(
         &self,
-        lookups: &[(usize, BlockLocation)],
+        lookups: &[(usize, [u8; CID_SIZE], BlockLocation)],
         results: &mut [Option<Bytes>],
     ) -> Result<(), ReadError> {
-        let mut by_file: HashMap<DataFileId, Vec<(usize, BlockLocation)>> = HashMap::new();
-        lookups.iter().for_each(|&(idx, loc)| {
-            by_file.entry(loc.file_id).or_default().push((idx, loc));
+        let mut by_file: HashMap<DataFileId, Vec<(usize, [u8; CID_SIZE], BlockLocation)>> =
+            HashMap::new();
+        lookups.iter().for_each(|&(idx, cid, loc)| {
+            by_file
+                .entry(loc.file_id)
+                .or_default()
+                .push((idx, cid, loc));
         });
 
         by_file.into_iter().try_for_each(|(file_id, mut entries)| {
             let handle = self.manager.open_for_read(file_id)?;
             let file_size = self.manager.io().file_size(handle.fd())?;
-            entries.sort_by_key(|(_, loc)| loc.offset);
-            entries.into_iter().try_for_each(|(orig_idx, loc)| {
-                let data = self.decode_and_validate(handle.fd(), file_size, loc)?;
+            entries.sort_by_key(|(_, _, loc)| loc.offset);
+            entries.into_iter().try_for_each(|(orig_idx, cid, loc)| {
+                let data = self.decode_and_validate(handle.fd(), file_size, loc, &cid)?;
                 results[orig_idx] = Some(data);
                 Ok::<_, ReadError>(())
             })
         })
     }
 
-    fn read_block_at(&self, location: BlockLocation) -> Result<Bytes, ReadError> {
+    fn read_block_at(
+        &self,
+        location: BlockLocation,
+        expected_cid: &[u8; CID_SIZE],
+    ) -> Result<Bytes, ReadError> {
         let handle = self.manager.open_for_read(location.file_id)?;
         let file_size = self.manager.io().file_size(handle.fd())?;
-        self.decode_and_validate(handle.fd(), file_size, location)
+        self.decode_and_validate(handle.fd(), file_size, location, expected_cid)
     }
 
     fn decode_and_validate(
@@ -126,38 +136,129 @@ impl<S: StorageIO> BlockStoreReader<S> {
         fd: FileId,
         file_size: u64,
         location: BlockLocation,
+        expected_cid: &[u8; CID_SIZE],
     ) -> Result<Bytes, ReadError> {
-        let attempt_once = || -> Result<Bytes, ReadError> {
-            match decode_block_record(self.manager.io(), fd, location.offset, file_size)? {
-                Some(ReadBlockRecord::Valid { data, .. })
-                    if data.len() == location.length.raw() as usize =>
+        let at_location = ReadError::Corrupted {
+            file_id: location.file_id,
+            offset: location.offset,
+        };
+        let attempt_once = || -> Result<Bytes, (ReadError, bool)> {
+            match decode_block_record(self.manager.io(), fd, location.offset, file_size) {
+                Err(e) => Err((e.into(), false)),
+                Ok(Some(ReadBlockRecord::Valid {
+                    data, cid_bytes, ..
+                })) if cid_bytes == *expected_cid
+                    && data.len() == location.length.raw() as usize =>
                 {
                     Ok(Bytes::from(data))
                 }
-                Some(ReadBlockRecord::Valid { .. }) => Err(ReadError::Corrupted {
-                    file_id: location.file_id,
-                    offset: location.offset,
-                }),
-                Some(
+                Ok(Some(ReadBlockRecord::Valid { .. })) => Err((at_location.clone(), false)),
+                Ok(Some(
                     ReadBlockRecord::Corrupted { offset } | ReadBlockRecord::Truncated { offset },
-                ) => Err(ReadError::Corrupted {
-                    file_id: location.file_id,
-                    offset,
-                }),
-                None => Err(ReadError::Corrupted {
-                    file_id: location.file_id,
-                    offset: location.offset,
-                }),
+                )) => Err((
+                    ReadError::Corrupted {
+                        file_id: location.file_id,
+                        offset,
+                    },
+                    true,
+                )),
+                Ok(None) => Err((at_location.clone(), true)),
             }
         };
         (0..READ_RETRY_ATTEMPTS.saturating_sub(1))
             .find_map(|_| match attempt_once() {
                 Ok(bytes) => Some(Ok(bytes)),
-                Err(ReadError::Corrupted { .. }) => None,
-                Err(e) => Some(Err(e)),
+                Err((_, true)) => None,
+                Err((e, false)) => Some(Err(e)),
             })
-            .unwrap_or_else(attempt_once)
+            .unwrap_or_else(|| attempt_once().map_err(|(e, _)| e))
     }
 }
 
 const READ_RETRY_ATTEMPTS: u32 = 4;
+
+#[cfg(test)]
+mod tests {
+    use super::{BlockStoreReader, ReadError};
+    use crate::blockstore::data_file::{CID_SIZE, DataFileWriter};
+    use crate::blockstore::hash_index::{BlockIndex, HashTable};
+    use crate::blockstore::manager::DataFileManager;
+    use crate::blockstore::test_cid;
+    use crate::blockstore::types::{BlockLocation, DataFileId};
+    use crate::io::StorageIO;
+    use crate::sim::SimulatedIO;
+    use std::path::{Path, PathBuf};
+    use std::sync::Arc;
+
+    const BLOCK_A: &[u8] = b"block-a-contents";
+    const BLOCK_B: &[u8] = b"block-b-contents";
+
+    fn setup() -> DataFileManager<SimulatedIO> {
+        let sim = SimulatedIO::pristine(7);
+        let dir = Path::new("/data");
+        sim.mkdir(dir).unwrap();
+        sim.sync_dir(dir).unwrap();
+        DataFileManager::new(sim, dir.to_path_buf(), 1 << 20)
+    }
+
+    fn write_two_blocks(
+        mgr: &DataFileManager<SimulatedIO>,
+    ) -> ([u8; CID_SIZE], BlockLocation, [u8; CID_SIZE], BlockLocation) {
+        let handle = mgr.open_for_append(DataFileId::new(0)).unwrap();
+        let mut writer = DataFileWriter::new(mgr.io(), handle.fd(), DataFileId::new(0)).unwrap();
+        let cid_a = test_cid(1);
+        let cid_b = test_cid(2);
+        let loc_a = writer.append_block(&cid_a, BLOCK_A).unwrap();
+        let loc_b = writer.append_block(&cid_b, BLOCK_B).unwrap();
+        writer.sync().unwrap();
+        assert_eq!(loc_a.length, loc_b.length, "blocks must be equal length");
+        (cid_a, loc_a, cid_b, loc_b)
+    }
+
+    fn index_mapping(pairs: &[([u8; CID_SIZE], BlockLocation)]) -> Arc<BlockIndex> {
+        let mut table = HashTable::with_capacity(64);
+        pairs.iter().for_each(|(cid, loc)| {
+            table.insert_or_increment(cid, *loc).unwrap();
+        });
+        Arc::new(BlockIndex::new(table, PathBuf::from("/index")))
+    }
+
+    #[test]
+    fn get_rejects_index_pointing_at_different_block() {
+        let mgr = Arc::new(setup());
+        let (cid_a, _loc_a, _cid_b, loc_b) = write_two_blocks(&mgr);
+        let index = index_mapping(&[(cid_a, loc_b)]);
+        let reader = BlockStoreReader::new(index, mgr);
+        match reader.get(&cid_a) {
+            Err(ReadError::Corrupted { .. }) => {}
+            other => {
+                panic!("expected Corrupted when CID resolves to a foreign block, got {other:?}")
+            }
+        }
+    }
+
+    #[test]
+    fn get_many_rejects_index_pointing_at_different_block() {
+        let mgr = Arc::new(setup());
+        let (cid_a, _loc_a, _cid_b, loc_b) = write_two_blocks(&mgr);
+        let index = index_mapping(&[(cid_a, loc_b)]);
+        let reader = BlockStoreReader::new(index, mgr);
+        match reader.get_many(&[cid_a]) {
+            Err(ReadError::Corrupted { .. }) => {}
+            other => panic!("expected Corrupted from get_many on foreign block, got {other:?}"),
+        }
+    }
+
+    #[test]
+    fn correct_mapping_still_resolves() {
+        let mgr = Arc::new(setup());
+        let (cid_a, loc_a, cid_b, loc_b) = write_two_blocks(&mgr);
+        let index = index_mapping(&[(cid_a, loc_a), (cid_b, loc_b)]);
+        let reader = BlockStoreReader::new(index, mgr);
+        assert_eq!(reader.get(&cid_a).unwrap().unwrap().as_ref(), BLOCK_A);
+        assert_eq!(reader.get(&cid_b).unwrap().unwrap().as_ref(), BLOCK_B);
+        let many = reader.get_many(&[cid_a, cid_b]).unwrap();
+        assert_eq!(many[0].as_ref().unwrap().as_ref(), BLOCK_A);
+        assert_eq!(many[1].as_ref().unwrap().as_ref(), BLOCK_B);
+    }
+}

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

@@ -27,7 +27,7 @@ use super::types::{
     LivenessInfo,
 };
 
-fn cid_to_bytes(cid: &Cid) -> Result<[u8; CID_SIZE], RepoError> {
+pub(crate) fn cid_to_bytes(cid: &Cid) -> Result<[u8; CID_SIZE], RepoError> {
     let raw = cid.to_bytes();
     let len = raw.len();
     raw.try_into().map_err(|_| {
@@ -64,7 +64,11 @@ fn read_error_to_repo(e: ReadError) -> RepoError {
         }
         ReadError::Corrupted { file_id, offset } => RepoError::storage(io::Error::new(
             io::ErrorKind::InvalidData,
-            format!("corrupted block at {file_id}:{}", offset.raw()),
+            format!(
+                "{} {file_id}:{}",
+                super::reader::BLOCK_CORRUPTION_MARKER,
+                offset.raw()
+            ),
         )),
     }
 }
@@ -345,10 +349,12 @@ impl<S: StorageIO + Send + Sync + 'static, C: Clock> TranquilBlockStore<S, C> {
         }
 
         let header_start = BlockOffset::new(super::data_file::BLOCK_HEADER_SIZE as u64);
+        let indexed_ends = index.indexed_file_ends();
 
         all_data_files.iter().try_for_each(|&fid| {
             let start_offset = file_cursors.get(&fid).copied().unwrap_or(header_start);
-            Self::replay_single_file(io, data_dir, index, fid, start_offset)
+            let indexed_end = indexed_ends.get(&fid).copied().unwrap_or(header_start);
+            Self::replay_single_file(io, data_dir, index, fid, start_offset, indexed_end)
         })
     }
 
@@ -358,6 +364,7 @@ impl<S: StorageIO + Send + Sync + 'static, C: Clock> TranquilBlockStore<S, C> {
         index: &BlockIndex,
         file_id: DataFileId,
         start_offset: BlockOffset,
+        indexed_end: BlockOffset,
     ) -> Result<(), RepoError> {
         let file_path = data_dir.join(format!("{file_id}.{}", super::manager::DATA_FILE_EXTENSION));
 
@@ -381,7 +388,8 @@ impl<S: StorageIO + Send + Sync + 'static, C: Clock> TranquilBlockStore<S, C> {
             })
             .is_ok();
 
-        let result = Self::scan_and_index(io, index, fd, file_id, start_offset, hint_exists);
+        let result =
+            Self::scan_and_index(io, index, fd, file_id, start_offset, indexed_end, hint_exists);
 
         let _ = io.close(fd);
 
@@ -394,6 +402,7 @@ impl<S: StorageIO + Send + Sync + 'static, C: Clock> TranquilBlockStore<S, C> {
         fd: crate::io::FileId,
         file_id: DataFileId,
         start_offset: BlockOffset,
+        indexed_end: BlockOffset,
         hint_exists: bool,
     ) -> Result<(), RepoError> {
         let file_size = io.file_size(fd).map_err(RepoError::storage)?;
@@ -453,15 +462,17 @@ impl<S: StorageIO + Send + Sync + 'static, C: Clock> TranquilBlockStore<S, C> {
             RepoError::storage(e)
         })?;
 
-        if file_size > last_valid_end.raw() {
+        let keep_end = last_valid_end.max(indexed_end);
+
+        if file_size > keep_end.raw() {
             tracing::info!(
                 file_id = %file_id,
-                truncating_from = last_valid_end.raw(),
+                truncating_from = keep_end.raw(),
                 file_size,
                 scanned_count = scanned_entries.len(),
                 "truncating partial/unacked tail"
             );
-            io.truncate(fd, last_valid_end.raw())
+            io.truncate(fd, keep_end.raw())
                 .map_err(RepoError::storage)?;
             io.sync(fd).map_err(RepoError::storage)?;
         }
@@ -475,7 +486,7 @@ impl<S: StorageIO + Send + Sync + 'static, C: Clock> TranquilBlockStore<S, C> {
             );
             let cursor = super::types::WriteCursor {
                 file_id,
-                offset: last_valid_end,
+                offset: keep_end,
             };
             index
                 .batch_put_if_absent(&scanned_entries, cursor)
@@ -599,6 +610,28 @@ impl<S: StorageIO + Send + Sync + 'static, C: Clock> TranquilBlockStore<S, C> {
             .map_err(commit_error_to_repo)
     }
 
+    pub fn repair_blocks(
+        &self,
+        blocks: Vec<(super::types::CidBytes, Vec<u8>)>,
+    ) -> Result<u64, CompactionError> {
+        if blocks.is_empty() {
+            return Ok(0);
+        }
+        let (tx, rx) = tokio::sync::oneshot::channel();
+        let sender = self
+            .writer
+            .with(|w| w.sender_round_robin().clone())
+            .map_err(|_| CompactionError::ChannelClosed)?;
+        sender
+            .send(CommitRequest::RepairBlocks {
+                blocks,
+                response: tx,
+            })
+            .map_err(|_| CompactionError::ChannelClosed)?;
+        rx.blocking_recv()
+            .map_err(|_| CompactionError::ChannelClosed)?
+    }
+
     pub fn get_block_sync(
         &self,
         cid_bytes: &[u8; CID_SIZE],
@@ -940,6 +973,7 @@ mod tests {
             &index,
             file_id,
             BlockOffset::new(BLOCK_HEADER_SIZE as u64),
+            BlockOffset::new(BLOCK_HEADER_SIZE as u64),
         );
 
         assert!(

crates/tranquil-store/tests/recovery_preserves_indexed_blocks.rs

@@ -0,0 +1,84 @@
+mod common;
+
+use std::fs;
+use std::path::Path;
+
+use common::{block_data, test_cid, with_runtime};
+use tranquil_store::blockstore::{
+    BLOCK_HEADER_SIZE, BlockStoreConfig, CID_SIZE, GroupCommitConfig, TranquilBlockStore,
+};
+
+fn config(dir: &Path) -> BlockStoreConfig {
+    BlockStoreConfig {
+        data_dir: dir.join("data"),
+        index_dir: dir.join("index"),
+        max_file_size: 1 << 20,
+        group_commit: GroupCommitConfig::default(),
+        shard_count: 1,
+    }
+}
+
+fn corrupt_nth_block_data(data_file: &Path, n: usize) {
+    let mut bytes = fs::read(data_file).expect("read data file");
+    let mut pos = BLOCK_HEADER_SIZE;
+    let mut idx = 0usize;
+    while pos + CID_SIZE + 4 <= bytes.len() {
+        let len = u32::from_le_bytes(
+            bytes[pos + CID_SIZE..pos + CID_SIZE + 4]
+                .try_into()
+                .unwrap(),
+        ) as usize;
+        let data_start = pos + CID_SIZE + 4;
+        let rec_end = data_start + len + 4;
+        if rec_end > bytes.len() {
+            break;
+        }
+        if idx == n && len > 0 {
+            bytes[data_start] ^= 0xFF;
+            fs::write(data_file, &bytes).expect("write corrupted data file");
+            return;
+        }
+        pos = rec_end;
+        idx += 1;
+    }
+    panic!("could not locate block {n} to corrupt");
+}
+
+#[test]
+fn recovery_preserves_indexed_blocks_past_a_mid_file_corruption() {
+    with_runtime(|| {
+        let dir = tempfile::TempDir::new().unwrap();
+        let seeds: Vec<u32> = (1..=10).collect();
+
+        {
+            let store = TranquilBlockStore::open(config(dir.path())).unwrap();
+            seeds.iter().for_each(|&s| {
+                store
+                    .put_blocks_blocking(vec![(test_cid(s), block_data(s))])
+                    .unwrap();
+            });
+            store
+                .repair_blocks(vec![(test_cid(999), block_data(999))])
+                .unwrap();
+            drop(store);
+        }
+
+        let data_file = dir.path().join("data").join("000001.tqb");
+        corrupt_nth_block_data(&data_file, 4);
+
+        let store = TranquilBlockStore::open(config(dir.path())).unwrap();
+
+        [1u32, 2, 3, 4].iter().for_each(|&s| {
+            assert!(
+                store.get_block_sync(&test_cid(s)).unwrap().is_some(),
+                "block {s} before the corruption was lost"
+            );
+        });
+        [6u32, 7, 8, 9, 10].iter().for_each(|&s| {
+            assert!(
+                store.get_block_sync(&test_cid(s)).unwrap().is_some(),
+                "block {s} after the corruption was lost to recovery truncation"
+            );
+        });
+    });
+}