candystore 1.2.2

use std::sync::{Arc, atomic::Ordering};

use crate::{
    crash_point,
    data_file::DataFile,
    index_file::{EntryPointer, RowWriteGuard},
    internal::{
        EntryType, FILE_OFFSET_ALIGNMENT, HashCoord, KVRef, KeyNamespace, MAX_DATA_FILES,
        ROW_WIDTH, aligned_data_entry_size, aligned_tombstone_entry_waste, invalid_data_error,
        is_corrupt_entry_error,
    },
    types::{Error, MAX_USER_KEY_SIZE, MAX_USER_VALUE_SIZE, Result},
};

use super::CandyStore;

#[derive(Clone, Copy)]
enum RebuildMode {
    TailFrom(u64),
    FullFile,
}

impl CandyStore {
    /// How many bytes of replayed data between progressive checkpoints.
    const REBUILD_CHECKPOINT_INTERVAL_BYTES: u64 = 256 * 1024;

    pub(super) fn recover_index(&self) -> Result<()> {
        let ordered_files = self.inner.ordered_data_files();
        let Some(last_file) = ordered_files.last().cloned() else {
            return Ok(());
        };

        for data_file in &ordered_files {
            data_file.sync_to_current()?;
        }

        if self.has_truncated_index_entries() {
            self.persist_rebuild_checkpoint(ordered_files[0].file_ordinal, 0, 0)?;
            // Purge only after the rewind is durable so an interrupted replay cannot rewind again.
            for data_file in &ordered_files {
                self.purge_file_range(data_file.file_idx, data_file.used_bytes(), u64::MAX)?;
            }
        }

        let (commit_file_ordinal, commit_offset) = self.inner.index_file.checkpoint_cursor();

        let start_idx = ordered_files
            .iter()
            .position(|data_file| data_file.file_ordinal >= commit_file_ordinal)
            .unwrap_or(ordered_files.len() - 1);
        let start_file = &ordered_files[start_idx];
        let rebuild_mode = if start_file.file_ordinal == commit_file_ordinal {
            self.validated_commit_offset(start_file, commit_offset)?
        } else {
            RebuildMode::TailFrom(0)
        };
        let rewind_ordinal = (start_idx != 0
            || !matches!(
                rebuild_mode,
                RebuildMode::TailFrom(0) | RebuildMode::FullFile
            ))
        .then_some(ordered_files[0].file_ordinal);

        // Recompute the runtime-only delta from the persisted replay cursor.
        self.inner
            .uncommitted_entries_delta
            .store(0, Ordering::Relaxed);
        let mut match_scratch = Vec::new();
        let mut bytes_since_checkpoint = 0u64;
        let mut pending_committed_delta = 0i64;

        let mut final_cursor = (last_file.file_ordinal, last_file.used_bytes());
        for (idx, data_file) in ordered_files.iter().enumerate().skip(start_idx) {
            let file_mode = if idx == start_idx {
                rebuild_mode
            } else {
                RebuildMode::TailFrom(0)
            };
            let Some(durable_extent) = self.rebuild_file_from(
                data_file,
                file_mode,
                rewind_ordinal,
                &mut bytes_since_checkpoint,
                &mut pending_committed_delta,
                &mut match_scratch,
            )?
            else {
                return self.recover_index();
            };
            final_cursor = (data_file.file_ordinal, durable_extent);
        }

        // The header count and the rows mmap are flushed independently, so after
        // a crash they can disagree; the rebuilt rows are the only source of truth.
        self.reset_recovery_stats_from_rows()?;
        self.persist_rebuild_checkpoint(final_cursor.0, final_cursor.1, 0)?;
        debug_assert_eq!(
            self.inner.uncommitted_entries_delta.load(Ordering::Relaxed),
            0
        );

        Ok(())
    }

    fn has_truncated_index_entries(&self) -> bool {
        let files = self.inner.data_files.read();
        let rows = self.inner.index_file.rows_table();
        for row_idx in 0..self.inner.index_file.num_rows() {
            let row = rows.row(row_idx);
            let split_level = row.split_level.load(Ordering::Acquire);
            if split_level == 0 {
                continue;
            }
            for col in 0..ROW_WIDTH {
                let ptr = row.pointers[col];
                if row.signatures[col] != HashCoord::INVALID_SIG
                    && ptr.is_valid()
                    && row.entry_belongs_to_row(col, row_idx, split_level)
                    && files
                        .get(&ptr.file_idx())
                        .is_some_and(|file| ptr.file_offset() >= file.used_bytes())
                {
                    return true;
                }
            }
        }
        false
    }

    fn rebuild_file_from(
        &self,
        data_file: &Arc<DataFile>,
        rebuild_mode: RebuildMode,
        rewind_ordinal: Option<u64>,
        bytes_since_checkpoint: &mut u64,
        pending_committed_delta: &mut i64,
        match_scratch: &mut Vec<u8>,
    ) -> Result<Option<u64>> {
        let start_offset = match rebuild_mode {
            RebuildMode::TailFrom(offset) => offset,
            RebuildMode::FullFile => 0,
        };

        // Pre-purge any index entries that point past the file's durable
        // extent. This handles the case where the data file was truncated
        // (e.g. disk-full or corruption) and ensures the replay loop won't
        // encounter stale pointers when comparing existing entries.
        let pre_rebuild_tail_upper_bound = data_file.recovery_tail_upper_bound();
        let pre_purge_extent = pre_rebuild_tail_upper_bound.next_multiple_of(FILE_OFFSET_ALIGNMENT);
        self.apply_recovery_delta(
            self.purge_file_range(data_file.file_idx, pre_purge_extent, u64::MAX)?,
            pending_committed_delta,
        );

        if matches!(rebuild_mode, RebuildMode::FullFile) {
            // The saved checkpoint within this file is no longer trustworthy.
            // Remove every pointer into it and rebuild its contribution from 0.
            self.apply_recovery_delta(
                self.purge_file_range(data_file.file_idx, 0, u64::MAX)?,
                pending_committed_delta,
            );
        }

        let mut offset = start_offset;
        let mut read_buf = Vec::new();
        let mut buf_file_offset = 0u64;
        let mut last_durable_offset = start_offset;
        while let Some((kv, entry_offset, next_offset)) =
            data_file.read_next_entry_ref(offset, &mut read_buf, &mut buf_file_offset)?
        {
            let expected_offset = offset.next_multiple_of(FILE_OFFSET_ALIGNMENT);
            if entry_offset != expected_offset
                && self.rewind_for_damaged_range(
                    data_file,
                    expected_offset,
                    entry_offset,
                    rewind_ordinal,
                    pending_committed_delta,
                )?
            {
                return Ok(None);
            }
            let entry_bytes = next_offset - offset;
            offset = next_offset;

            let Some(ns) = KeyNamespace::from_u8(kv.ns) else {
                return Err(invalid_data_error("unknown key namespace in data file"));
            };

            self.apply_recovery_delta(
                self.recover_entry(data_file, ns, kv, entry_offset, match_scratch)?,
                pending_committed_delta,
            );
            self.inner
                .stats
                .num_rebuilt_entries
                .fetch_add(1, Ordering::Relaxed);
            last_durable_offset = next_offset;
            crash_point("rebuild_entry");

            *bytes_since_checkpoint += entry_bytes;
            if *bytes_since_checkpoint >= Self::REBUILD_CHECKPOINT_INTERVAL_BYTES {
                self.persist_rebuild_checkpoint(
                    data_file.file_ordinal,
                    offset,
                    *pending_committed_delta,
                )?;
                *pending_committed_delta = 0;
                *bytes_since_checkpoint = 0;
            }
        }

        let durable_extent = last_durable_offset.next_multiple_of(FILE_OFFSET_ALIGNMENT);

        if durable_extent < pre_rebuild_tail_upper_bound {
            if self.rewind_for_damaged_range(
                data_file,
                durable_extent,
                u64::MAX,
                rewind_ordinal,
                pending_committed_delta,
            )? {
                return Ok(None);
            }
            self.inner.stats.num_rebuild_purged_bytes.fetch_add(
                pre_rebuild_tail_upper_bound - durable_extent,
                Ordering::Relaxed,
            );
            data_file.truncate_to_offset(durable_extent)?;
        }

        self.apply_recovery_delta(
            self.purge_file_range(data_file.file_idx, durable_extent, u64::MAX)?,
            pending_committed_delta,
        );
        Ok(Some(durable_extent))
    }

    /// Purges index pointers into a damaged range; returns true if replay must restart from the oldest file.
    fn rewind_for_damaged_range(
        &self,
        data_file: &DataFile,
        start: u64,
        end: u64,
        rewind_ordinal: Option<u64>,
        pending_committed_delta: &mut i64,
    ) -> Result<bool> {
        let rewind = rewind_ordinal
            .filter(|_| self.has_index_entries_in_range(data_file.file_idx, start, end));
        if let Some(ordinal) = rewind {
            self.persist_rebuild_checkpoint(ordinal, 0, *pending_committed_delta)?;
            *pending_committed_delta = 0;
        }
        self.apply_recovery_delta(
            self.purge_file_range(data_file.file_idx, start, end)?,
            pending_committed_delta,
        );
        Ok(rewind.is_some())
    }

    fn has_index_entries_in_range(&self, file_idx: u16, start: u64, end: u64) -> bool {
        let rows = self.inner.index_file.rows_table();
        (0..self.inner.index_file.num_rows()).any(|row_idx| {
            let row = rows.row(row_idx);
            row.split_level.load(Ordering::Acquire) != 0
                && (0..ROW_WIDTH).any(|col| {
                    let ptr = row.pointers[col];
                    row.signatures[col] != HashCoord::INVALID_SIG
                        && ptr.is_valid()
                        && ptr.file_idx() == file_idx
                        && (start..end).contains(&ptr.file_offset())
                })
        })
    }

    fn validated_commit_offset(
        &self,
        active_file: &Arc<DataFile>,
        checkpoint_offset: u64,
    ) -> Result<RebuildMode> {
        if checkpoint_offset == 0 {
            return Ok(RebuildMode::TailFrom(0));
        }

        let used_bytes = active_file.used_bytes();
        if checkpoint_offset > used_bytes {
            return Ok(RebuildMode::FullFile);
        }
        if checkpoint_offset == used_bytes {
            return Ok(RebuildMode::TailFrom(checkpoint_offset));
        }

        let mut probe_buf = Vec::new();
        let mut probe_file_offset = 0u64;
        match active_file.read_next_entry_ref(
            checkpoint_offset,
            &mut probe_buf,
            &mut probe_file_offset,
        )? {
            Some((_, entry_offset, _)) if entry_offset == checkpoint_offset => {
                Ok(RebuildMode::TailFrom(checkpoint_offset))
            }
            _ => Ok(RebuildMode::FullFile),
        }
    }

    fn persist_rebuild_checkpoint(&self, ordinal: u64, offset: u64, delta: i64) -> Result<()> {
        let resume_offset = offset.next_multiple_of(FILE_OFFSET_ALIGNMENT);

        self.inner
            .index_file
            .sync_rows(self.inner.index_file.rows_table_mut())?;
        self.inner.fold_checkpointed_num_entries(delta);
        self.inner
            .persist_checkpoint_cursor(ordinal, resume_offset)?;

        self.inner.index_file.sync_header()
    }

    /// Rebuilds entry counts and per-file waste from the final live index.
    fn reset_recovery_stats_from_rows(&self) -> Result<()> {
        let files = self.inner.data_files.read();
        let row_table = self.inner.index_file.rows_table();
        let num_rows = self.inner.index_file.num_rows();
        let mut live = 0u64;
        let mut live_entries = vec![Vec::new(); MAX_DATA_FILES as usize];

        for row_idx in 0..num_rows {
            let row = row_table.row(row_idx);
            let sl = row.split_level.load(Ordering::Acquire);
            if sl == 0 {
                continue;
            }
            for col in 0..ROW_WIDTH {
                let ptr = row.pointers[col];
                if row.signatures[col] != HashCoord::INVALID_SIG
                    && ptr.is_valid()
                    && row.entry_belongs_to_row(col, row_idx, sl)
                    && files.contains_key(&ptr.file_idx())
                {
                    live += 1;
                    live_entries[ptr.file_idx() as usize].push(ptr);
                }
            }
        }

        for (file_idx, entries) in live_entries.into_iter().enumerate() {
            let file_idx = file_idx as u16;
            let waste = match files.get(&file_idx) {
                Some(file) => file
                    .used_bytes()
                    .saturating_sub(file.indexed_live_bytes(&entries)?),
                None => 0,
            };
            self.inner
                .index_file
                .set_file_waste(file_idx, waste.min(u32::MAX as u64) as u32);
        }

        self.inner
            .index_file
            .header_ref()
            .committed_num_entries
            .store(live, Ordering::Relaxed);
        self.inner
            .uncommitted_entries_delta
            .store(0, Ordering::Relaxed);
        Ok(())
    }

    /// Removes index entries pointing into `[start, end)` of `file_idx`.
    fn purge_file_range(&self, file_idx: u16, start: u64, end: u64) -> Result<i64> {
        let row_table = self.inner.index_file.rows_table();
        let num_rows = self.inner.index_file.num_rows();
        let mut removed = 0i64;

        for row_idx in 0..num_rows {
            let mut row = row_table.row_mut(row_idx);
            if row.split_level.load(Ordering::Acquire) == 0 {
                continue;
            }
            for col in 0..ROW_WIDTH {
                if row.signatures[col] == HashCoord::INVALID_SIG {
                    continue;
                }
                let ptr = row.pointers[col];
                if !ptr.is_valid() {
                    continue;
                }
                if ptr.file_idx() == file_idx && (start..end).contains(&ptr.file_offset()) {
                    row.remove(col);
                    removed += 1;
                }
            }
        }
        Ok(-removed)
    }

    fn recover_entry(
        &self,
        data_file: &Arc<DataFile>,
        ns: KeyNamespace,
        kv: KVRef<'_>,
        entry_offset: u64,
        match_scratch: &mut Vec<u8>,
    ) -> Result<i64> {
        match kv.entry_type {
            EntryType::Insert | EntryType::Update => {
                self.recover_data_entry(data_file, ns, kv, entry_offset, match_scratch)
            }
            EntryType::Tombstone => self.recover_tombstone_entry(data_file, ns, kv, match_scratch),
            _ => Ok(0),
        }
    }

    fn apply_recovery_delta(&self, delta: i64, pending_committed_delta: &mut i64) {
        if delta == 0 {
            return;
        }

        self.inner.add_uncommitted_num_entries(delta);
        *pending_committed_delta += delta;
    }

    /// Fix index pointers for a data/update entry and return its live-entry delta.
    fn recover_data_entry(
        &self,
        data_file: &Arc<DataFile>,
        ns: KeyNamespace,
        kv: KVRef<'_>,
        entry_offset: u64,
        match_scratch: &mut Vec<u8>,
    ) -> Result<i64> {
        let key = kv.key();
        let val = kv.value();
        self.validate_recovered_data_entry(key, val)?;
        let entry_len = 4 + 4 + key.len() + val.len() + 2;
        let aligned_len = entry_len.next_multiple_of(FILE_OFFSET_ALIGNMENT as usize);
        let hc = HashCoord::new(ns, key, self.inner.config.hash_key);
        let ptr = EntryPointer::new(
            data_file.file_idx,
            entry_offset,
            aligned_len,
            hc.masked_row_selector(),
        );

        self.inner._mut_op(ns, key, &[], |hc, mut row, key, _| {
            let (existing, mut delta) =
                self.find_recovery_match(&mut row, hc, key, match_scratch)?;
            if let Some((col, entry)) = existing {
                if entry == ptr {
                    // Already points at this entry — nothing to fix.
                    return Ok(delta);
                }
                if entry.file_idx() == data_file.file_idx && entry.file_offset() > ptr.file_offset()
                {
                    // A newer active-file entry already exists — skip.
                    return Ok(delta);
                }
                // Older pointer — replace with this newer one.
                row.replace_pointer(col, ptr);
                return Ok(delta);
            }
            // Key not in index — insert it.
            if let Some(col) = row.find_free_slot() {
                row.insert(col, hc.sig, ptr);
                delta += 1;
                Ok(delta)
            } else {
                Err(Error::SplitRow(row.split_level.load(Ordering::Relaxed)))
            }
        })
    }

    /// Finds `key`'s slot in `row`, unlinking corrupt candidates; returns the live-entry delta of those unlinks.
    fn find_recovery_match(
        &self,
        row: &mut RowWriteGuard<'_>,
        hc: HashCoord,
        key: &[u8],
        match_scratch: &mut Vec<u8>,
    ) -> Result<(Option<(usize, EntryPointer)>, i64)> {
        let files = self.inner.data_files.read();
        let mut corrupt_cols = Vec::new();
        let mut existing = None;
        for (col, entry) in row.iter_matches(hc) {
            let file = files
                .get(&entry.file_idx())
                .ok_or(Error::MissingDataFile(entry.file_idx()))?;
            match file.read_kv_into(entry.file_offset(), entry.size_hint(), match_scratch) {
                Ok(kv) if kv.key() == key => {
                    existing = Some((col, entry));
                    break;
                }
                Ok(_) => {}
                Err(err) if is_corrupt_entry_error(&err) => corrupt_cols.push(col),
                Err(err) => return Err(err),
            }
        }
        let delta = -(corrupt_cols.len() as i64);
        for col in corrupt_cols {
            self.inner.unlink_corrupt_entry(row, col);
        }
        Ok((existing, delta))
    }

    /// Fix index pointers for a tombstone entry and return its live-entry delta.
    fn recover_tombstone_entry(
        &self,
        _data_file: &Arc<DataFile>,
        ns: KeyNamespace,
        kv: KVRef<'_>,
        match_scratch: &mut Vec<u8>,
    ) -> Result<i64> {
        let key = kv.key();
        self.validate_recovered_tombstone_entry(key)?;

        self.inner._mut_op(ns, key, &[], |hc, mut row, key, _| {
            let (existing, mut delta) =
                self.find_recovery_match(&mut row, hc, key, match_scratch)?;
            if let Some((col, _)) = existing {
                row.remove(col);
                delta -= 1;
            }
            Ok(delta)
        })
    }

    fn validate_recovered_data_entry(&self, key: &[u8], val: &[u8]) -> Result<()> {
        let entry_size = aligned_data_entry_size(key.len(), val.len()) as usize;
        if key.len() > MAX_USER_KEY_SIZE
            || val.len() > MAX_USER_VALUE_SIZE
            || entry_size > self.inner.config.max_data_file_size as usize
        {
            return Err(invalid_data_error(
                "recovered data entry exceeds configured limits",
            ));
        }
        Ok(())
    }

    fn validate_recovered_tombstone_entry(&self, key: &[u8]) -> Result<()> {
        let entry_size = aligned_tombstone_entry_waste(key.len()) as usize;
        if key.len() > MAX_USER_KEY_SIZE
            || entry_size > self.inner.config.max_data_file_size as usize
        {
            return Err(invalid_data_error(
                "recovered tombstone entry exceeds configured limits",
            ));
        }
        Ok(())
    }
}