feat(store): rebuild & rewrite missing/corrupt MST blocks

Lewis: May this revision serve well! <lu5a@proton.me>
2026-07-19 22:42:25 +00:00 · 2026-05-31 13:24:03 +03:00
parent b8cae15c12
commit 44d73dac58
3 changed files with 296 additions and 0 deletions
@@ -186,6 +186,111 @@ fn purge_phantom_file<S: StorageIO>(
    })
 }

+#[allow(clippy::too_many_arguments)]
+pub(super) fn repair_blocks_on_writer_thread<S: StorageIO>(
+    manager: &DataFileManager<S>,
+    index: &BlockIndex,
+    blocks: &[(CidBytes, Vec<u8>)],
+    current_epoch: CommitEpoch,
+    file_ids: &FileIdAllocator,
+    hint_positions: &super::group_commit::ShardHintPositions,
+    epoch: &super::types::EpochCounter,
+) -> Result<u64, CompactionError> {
+    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<S: StorageIO>(
+    manager: &DataFileManager<S>,
+    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<S: StorageIO>(
    manager: &DataFileManager<S>,
@@ -175,6 +175,9 @@ type PutResponse = tokio::sync::oneshot::Sender<Result<Vec<BlockLocation>, Commi
 type ApplyResponse = tokio::sync::oneshot::Sender<Result<(), CommitError>>;
 type CompactResponse = tokio::sync::oneshot::Sender<Result<CompactionResult, CompactionError>>;
 type RepairResponse = tokio::sync::oneshot::Sender<Result<u64, CommitError>>;
+type RepairBlocksResponse = tokio::sync::oneshot::Sender<Result<u64, CompactionError>>;
+type RepairBlockSet = Vec<(CidBytes, Vec<u8>)>;
+type DeferredRepairBlock = (RepairBlockSet, RepairBlocksResponse);
 type QuiesceResponse = tokio::sync::oneshot::Sender<BlockstoreSnapshot>;
 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<DeferredRepairBlock>,
    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<S: StorageIO, C: Clock>(
                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<S: StorageIO, C: Clock>(
                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<S: StorageIO, C: Clock>(
    let mut entries: Vec<BatchEntry> = 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<DeferredRepairBlock> = 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<S: StorageIO, C: Clock>(
            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<S: StorageIO, C: Clock>(
        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);
 }

@@ -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<Vec<(CidBytes, Vec<u8>)>, 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<u8>)>) -> Vec<Vec<(CidBytes, Vec<u8>)>> {
+    blocks
+        .into_iter()
+        .fold(
+            (Vec::<Vec<(CidBytes, Vec<u8>)>>::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<S, C>(
+    store: &TranquilBlockStore<S, C>,
+    entries: &[(String, Cid)],
+    expected_root: Cid,
+) -> Result<RepairOutcome, RepoError>
+where
+    S: StorageIO + Send + Sync + 'static,
+    C: Clock,
+{
+    let store = store.clone();
+    let entries = entries.to_vec();
+    tokio::task::spawn_blocking(move || -> Result<RepairOutcome, RepoError> {
+        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<u8>)> = 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)?
+}