From 44d73dac58a508c5975425de4dfb305b498ac109 Mon Sep 17 00:00:00 2001 From: Lewis Date: Sun, 31 May 2026 13:24:03 +0300 Subject: [PATCH] feat(store): rebuild & rewrite missing/corrupt MST blocks Lewis: May this revision serve well! --- .../src/blockstore/compaction.rs | 105 ++++++++++++++++ .../src/blockstore/group_commit.rs | 78 ++++++++++++ .../tranquil-store/src/blockstore/repair.rs | 113 ++++++++++++++++++ 3 files changed, 296 insertions(+) create mode 100644 crates/tranquil-store/src/blockstore/repair.rs diff --git a/crates/tranquil-store/src/blockstore/compaction.rs b/crates/tranquil-store/src/blockstore/compaction.rs index 1a31de9..61ba353 100644 --- a/crates/tranquil-store/src/blockstore/compaction.rs +++ b/crates/tranquil-store/src/blockstore/compaction.rs @@ -186,6 +186,111 @@ fn purge_phantom_file( }) } +#[allow(clippy::too_many_arguments)] +pub(super) fn repair_blocks_on_writer_thread( + manager: &DataFileManager, + index: &BlockIndex, + blocks: &[(CidBytes, Vec)], + current_epoch: CommitEpoch, + file_ids: &FileIdAllocator, + hint_positions: &super::group_commit::ShardHintPositions, + epoch: &super::types::EpochCounter, +) -> Result { + if blocks.is_empty() { + return Ok(0); + } + + let new_file_id = file_ids.allocate(); + let new_handle = manager.open_for_append(new_file_id)?; + let mut writer = DataFileWriter::new(manager.io(), new_handle.fd(), new_file_id)?; + + let hint_path = hint_file_path(manager.data_dir(), new_file_id); + let hint_fd = manager.io().open(&hint_path, OpenOptions::read_write())?; + let mut hint_writer = HintFileWriter::new(manager.io(), hint_fd); + + let mut relocations: Vec<(CidBytes, BlockLocation)> = Vec::with_capacity(blocks.len()); + let write_result = blocks.iter().try_for_each(|(cid, data)| { + let refcount = index.get(cid).map(|e| e.refcount.raw()).unwrap_or(1).max(1); + let loc = writer.append_block(cid, data)?; + hint_writer.append_relocate(cid, &loc, refcount)?; + relocations.push((*cid, loc)); + Ok::<_, CompactionError>(()) + }); + + let record_count = u32::try_from(relocations.len()).unwrap_or(u32::MAX); + let writer_position = writer.position(); + let finalize = write_result + .and_then(|()| writer.sync().map_err(CompactionError::from)) + .and_then(|()| { + hint_writer + .append_commit_marker( + current_epoch.raw(), + record_count, + new_file_id, + writer_position, + ) + .map_err(CompactionError::from) + }) + .and_then(|()| hint_writer.sync().map_err(CompactionError::from)) + .and_then(|()| { + manager + .io() + .sync_dir(manager.data_dir()) + .map_err(CompactionError::from) + }) + .and_then(|()| manager.io().barrier().map_err(CompactionError::from)); + + let final_hint_offset = hint_writer.position(); + let _ = manager.io().close(hint_fd); + + if let Err(e) = + finalize.and_then(|()| verify_repaired_blocks(manager, new_file_id, &relocations)) + { + manager.delete_data_file(new_file_id).ok(); + manager + .io() + .delete(&hint_file_path(manager.data_dir(), new_file_id)) + .ok(); + return Err(e); + } + + hint_positions.record_extra(new_file_id, final_hint_offset); + index.apply_compaction(&relocations, &[]); + index + .write_checkpoint(epoch.current(), hint_positions) + .map_err(CompactionError::Io)?; + + tracing::info!( + dest = %new_file_id, + repaired = relocations.len(), + "structural repair complete" + ); + + Ok(relocations.len() as u64) +} + +fn verify_repaired_blocks( + manager: &DataFileManager, + file_id: DataFileId, + relocations: &[(CidBytes, BlockLocation)], +) -> Result<(), CompactionError> { + let handle = manager.open_for_read(file_id)?; + let file_size = manager.io().file_size(handle.fd())?; + relocations.iter().try_for_each(|(cid, loc)| { + match super::data_file::decode_block_record( + manager.io(), + handle.fd(), + loc.offset, + file_size, + ) { + Ok(Some(ReadBlockRecord::Valid { cid_bytes, .. })) if cid_bytes == *cid => Ok(()), + _ => Err(CompactionError::Io(io::Error::other( + "repaired block failed read-back verification", + ))), + } + }) +} + #[allow(clippy::too_many_arguments)] fn stream_compact( manager: &DataFileManager, diff --git a/crates/tranquil-store/src/blockstore/group_commit.rs b/crates/tranquil-store/src/blockstore/group_commit.rs index 69ea0c6..c728512 100644 --- a/crates/tranquil-store/src/blockstore/group_commit.rs +++ b/crates/tranquil-store/src/blockstore/group_commit.rs @@ -175,6 +175,9 @@ type PutResponse = tokio::sync::oneshot::Sender, Commi type ApplyResponse = tokio::sync::oneshot::Sender>; type CompactResponse = tokio::sync::oneshot::Sender>; type RepairResponse = tokio::sync::oneshot::Sender>; +type RepairBlocksResponse = tokio::sync::oneshot::Sender>; +type RepairBlockSet = Vec<(CidBytes, Vec)>; +type DeferredRepairBlock = (RepairBlockSet, RepairBlocksResponse); type QuiesceResponse = tokio::sync::oneshot::Sender; type QuiesceResume = tokio::sync::oneshot::Receiver<()>; @@ -197,6 +200,10 @@ pub enum CommitRequest { leaked_cids: Vec<(CidBytes, RefCount)>, response: RepairResponse, }, + RepairBlocks { + blocks: RepairBlockSet, + response: RepairBlocksResponse, + }, Quiesce { response: QuiesceResponse, resume: QuiesceResume, @@ -668,6 +675,10 @@ enum ClassifyResult { leaked_cids: Vec<(CidBytes, RefCount)>, response: RepairResponse, }, + RepairBlocks { + blocks: RepairBlockSet, + response: RepairBlocksResponse, + }, Quiesce { response: QuiesceResponse, resume: QuiesceResume, @@ -704,6 +715,9 @@ fn classify_request(req: CommitRequest) -> ClassifyResult { leaked_cids, response, }, + CommitRequest::RepairBlocks { blocks, response } => { + ClassifyResult::RepairBlocks { blocks, response } + } CommitRequest::Quiesce { response, resume } => ClassifyResult::Quiesce { response, resume }, CommitRequest::Shutdown => ClassifyResult::Shutdown, } @@ -720,6 +734,7 @@ struct DrainResult { shutdown: bool, deferred_compacts: Vec<(DataFileId, u64, CompactResponse)>, deferred_repairs: Vec<(Vec<(CidBytes, RefCount)>, RepairResponse)>, + deferred_repair_blocks: Vec, deferred_quiesces: Vec<(QuiesceResponse, QuiesceResume)>, } @@ -735,6 +750,7 @@ fn drain_batch( shutdown: true, deferred_compacts: Vec::new(), deferred_repairs: Vec::new(), + deferred_repair_blocks: Vec::new(), deferred_quiesces: Vec::new(), }; } @@ -748,6 +764,7 @@ fn drain_batch( shutdown: false, deferred_compacts: vec![(file_id, grace_period_ms, response)], deferred_repairs: Vec::new(), + deferred_repair_blocks: Vec::new(), deferred_quiesces: Vec::new(), }; } @@ -760,6 +777,17 @@ fn drain_batch( shutdown: false, deferred_compacts: Vec::new(), deferred_repairs: vec![(leaked_cids, response)], + deferred_repair_blocks: Vec::new(), + deferred_quiesces: Vec::new(), + }; + } + ClassifyResult::RepairBlocks { blocks, response } => { + return DrainResult { + entries: Vec::new(), + shutdown: false, + deferred_compacts: Vec::new(), + deferred_repairs: Vec::new(), + deferred_repair_blocks: vec![(blocks, response)], deferred_quiesces: Vec::new(), }; } @@ -769,6 +797,7 @@ fn drain_batch( shutdown: false, deferred_compacts: Vec::new(), deferred_repairs: Vec::new(), + deferred_repair_blocks: Vec::new(), deferred_quiesces: vec![(response, resume)], }; } @@ -781,6 +810,7 @@ fn drain_batch( shutdown: false, deferred_compacts: Vec::new(), deferred_repairs: Vec::new(), + deferred_repair_blocks: Vec::new(), deferred_quiesces: Vec::new(), }; @@ -803,6 +833,10 @@ fn drain_batch( r.deferred_repairs.push((leaked_cids, response)); false } + ClassifyResult::RepairBlocks { blocks, response } => { + r.deferred_repair_blocks.push((blocks, response)); + false + } ClassifyResult::Quiesce { response, resume } => { r.deferred_quiesces.push((response, resume)); false @@ -946,6 +980,19 @@ fn commit_loop( let _ = response.send(Ok(repaired)); continue; } + Ok(CommitRequest::RepairBlocks { blocks, response }) => { + let result = compaction::repair_blocks_on_writer_thread( + manager, + index, + &blocks, + epoch.current(), + &ctx.file_ids, + &ctx.hint_positions, + epoch, + ); + let _ = response.send(result); + continue; + } Ok(CommitRequest::Quiesce { response, resume }) => { handle_quiesce(manager, state, epoch, response, resume); continue; @@ -1018,6 +1065,22 @@ fn commit_loop( let _ = response.send(Ok(repaired)); }); + drain + .deferred_repair_blocks + .into_iter() + .for_each(|(blocks, response)| { + let result = compaction::repair_blocks_on_writer_thread( + manager, + index, + &blocks, + epoch.current(), + &ctx.file_ids, + &ctx.hint_positions, + epoch, + ); + let _ = response.send(result); + }); + maybe_checkpoint( index, epoch, @@ -1062,6 +1125,7 @@ fn drain_and_process_remaining( let mut entries: Vec = Vec::new(); let mut compacts: Vec<(DataFileId, u64, CompactResponse)> = Vec::new(); let mut repairs: Vec<(Vec<(CidBytes, RefCount)>, RepairResponse)> = Vec::new(); + let mut repair_blocks: Vec = Vec::new(); std::iter::from_fn(|| receiver.try_recv().ok()).for_each(|req| match classify_request(req) { ClassifyResult::Batch(entry) => entries.push(entry), @@ -1074,6 +1138,7 @@ fn drain_and_process_remaining( leaked_cids, response, } => repairs.push((leaked_cids, response)), + ClassifyResult::RepairBlocks { blocks, response } => repair_blocks.push((blocks, response)), ClassifyResult::Shutdown | ClassifyResult::Quiesce { .. } => {} }); @@ -1111,6 +1176,19 @@ fn drain_and_process_remaining( let _ = response.send(Ok(repaired)); }); + repair_blocks.into_iter().for_each(|(blocks, response)| { + let result = compaction::repair_blocks_on_writer_thread( + manager, + index, + &blocks, + epoch.current(), + &ctx.file_ids, + &ctx.hint_positions, + epoch, + ); + let _ = response.send(result); + }); + shutdown_checkpoint(index, epoch, &ctx.hint_positions); } diff --git a/crates/tranquil-store/src/blockstore/repair.rs b/crates/tranquil-store/src/blockstore/repair.rs new file mode 100644 index 0000000..19fc63c --- /dev/null +++ b/crates/tranquil-store/src/blockstore/repair.rs @@ -0,0 +1,113 @@ +use std::sync::Arc; + +use cid::Cid; +use jacquard_repo::error::RepoError; +use jacquard_repo::mst::Mst; +use jacquard_repo::storage::MemoryBlockStore; + +use crate::clock::Clock; +use crate::io::StorageIO; + +use super::store::{TranquilBlockStore, cid_to_bytes}; +use super::types::CidBytes; + +const REPAIR_FILE_BYTE_BUDGET: usize = 64 * 1024 * 1024; + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct RepairOutcome { + pub nodes_total: usize, + pub nodes_repaired: u64, +} + +fn rebuild_err(context: &str, e: impl std::fmt::Display) -> RepoError { + RepoError::storage(std::io::Error::other(format!("{context}: {e}"))) +} + +async fn rebuild_node_blocks( + entries: Vec<(String, Cid)>, + expected_root: Cid, +) -> Result)>, RepoError> { + let scratch = Arc::new(MemoryBlockStore::new()); + let mut mst = Mst::new(scratch); + for (key, cid) in &entries { + mst.add_mut(key.as_str(), *cid) + .await + .map_err(|e| rebuild_err("mst rebuild add", e))?; + } + + let (root, blocks) = mst + .collect_blocks() + .await + .map_err(|e| rebuild_err("mst collect_blocks", e))?; + + if root != expected_root { + return Err(RepoError::storage(std::io::Error::new( + std::io::ErrorKind::InvalidData, + format!( + "rebuilt MST root {root} does not match expected root {expected_root}, refusing to repair" + ), + ))); + } + + blocks + .into_iter() + .map(|(cid, bytes)| Ok((cid_to_bytes(&cid)?, bytes.to_vec()))) + .collect() +} + +fn batch_by_bytes(blocks: Vec<(CidBytes, Vec)>) -> Vec)>> { + blocks + .into_iter() + .fold( + (Vec::)>>::new(), 0usize), + |(mut batches, current_bytes), item| { + let item_len = item.1.len(); + match batches.last_mut() { + Some(last) if current_bytes + item_len <= REPAIR_FILE_BYTE_BUDGET => { + last.push(item); + (batches, current_bytes + item_len) + } + _ => { + batches.push(vec![item]); + (batches, item_len) + } + } + }, + ) + .0 +} + +pub async fn rebuild_and_repair_mst( + store: &TranquilBlockStore, + entries: &[(String, Cid)], + expected_root: Cid, +) -> Result +where + S: StorageIO + Send + Sync + 'static, + C: Clock, +{ + let store = store.clone(); + let entries = entries.to_vec(); + tokio::task::spawn_blocking(move || -> Result { + let handle = tokio::runtime::Handle::current(); + let node_blocks = handle.block_on(rebuild_node_blocks(entries, expected_root))?; + let nodes_total = node_blocks.len(); + + let to_repair: Vec<(CidBytes, Vec)> = node_blocks + .into_iter() + .filter(|(cid, _)| !matches!(store.get_block_sync(cid), Ok(Some(_)))) + .collect(); + + let nodes_repaired = batch_by_bytes(to_repair) + .into_iter() + .try_fold(0u64, |acc, batch| store.repair_blocks(batch).map(|n| acc + n)) + .map_err(|e| rebuild_err("repair_blocks", e))?; + + Ok(RepairOutcome { + nodes_total, + nodes_repaired, + }) + }) + .await + .map_err(RepoError::task_failed)? +}