armdb 0.11.2

#![allow(dead_code)] // The managed opener wires this pass in Task 5.

use std::io::Write;
use std::{
    collections::{HashMap, VecDeque},
    fs,
    path::{Path, PathBuf},
    sync::Mutex,
    time::Instant,
};

use zerocopy::FromBytes;

use crate::entry::{EntryHeader, check_page_padding_at, compute_crc32, entry_size};
#[allow(unused_imports)]
use crate::format_migration::marker::read_shard_version;
use crate::format_migration::marker::{
    ConfiguredShard, ShardSwapState, configured_shards, swap_prepared_shard, write_shard_version,
};
use crate::format_migration::plan::{MigrationPlan, TransformBuffers, TransformOutcome};
use crate::hint::{generate_hint_data_dyn, hint_path_for_data, write_hint_file};
use crate::io::direct;
use crate::{Config, DbError, DbResult, DiskLoc, FormatMigrationReport};

#[cfg(any(test, feature = "test-faults"))]
use crate::format_migration::failpoints;

#[cfg(feature = "encryption")]
use crate::{
    crypto::PageCipher,
    io::tags::{TagFile, tags_path_for_data},
};
#[cfg(feature = "encryption")]
use std::sync::Arc;

pub(crate) struct BackendMigrationOutcome {
    pub(crate) report: FormatMigrationReport,
    pub(crate) gsn_high: u64,
    pub(crate) shard_gsn_high: Vec<u64>,
}

struct Row {
    header: EntryHeader,
    loc: DiskLoc,
    ordinal: u64,
}

struct Selected {
    sequence: u64,
    row: Option<Row>,
}

pub(crate) fn migrate_bitcask(
    dir: &Path,
    collection: &'static str,
    source_version: u16,
    plan: &MigrationPlan,
    config: &Config,
    key_len: usize,
) -> DbResult<BackendMigrationOutcome> {
    let started = Instant::now();
    let shards = configured_shards(
        dir,
        config.shard_count,
        source_version,
        plan.format_to_version,
    )?;
    let mut report = FormatMigrationReport::empty(collection, source_version);
    report.to = plan.format_to_version;
    report.shards = shards.len();
    report.warnings.extend(plan.warnings.iter().cloned());
    let worker_count = shards
        .len()
        .min(std::thread::available_parallelism().map_or(1, std::num::NonZero::get));
    let queue = Mutex::new(VecDeque::from(shards));
    let results = Mutex::new(Vec::new());
    std::thread::scope(|scope| -> DbResult<()> {
        let mut workers = Vec::with_capacity(worker_count);
        for _ in 0..worker_count {
            workers.push(scope.spawn(|| -> DbResult<()> {
                while let Some(ConfiguredShard {
                    id: shard_id,
                    name,
                    state,
                }) = queue.lock().expect("migration queue poisoned").pop_front()
                {
                    let result = rewrite_shard(
                        dir, &name, shard_id, state, collection, plan, config, key_len,
                    )?;
                    results
                        .lock()
                        .expect("migration results poisoned")
                        .push((shard_id, result));
                }
                Ok(())
            }));
        }
        for worker in workers {
            worker.join().map_err(|_| DbError::RecoveryPanic)??;
        }
        Ok(())
    })?;
    let mut gsn_high = 0;
    let mut shard_results = results.into_inner().expect("migration results poisoned");
    shard_results.sort_by_key(|(id, _)| *id);
    let mut shard_gsn_high = Vec::with_capacity(shard_results.len());
    for (_, (migrated, dropped, before, after, high)) in shard_results {
        report.rows_migrated += migrated;
        report.rows_dropped += dropped;
        report.bytes_before += before;
        report.bytes_after += after;
        gsn_high = gsn_high.max(high);
        shard_gsn_high.push(high);
    }
    report.elapsed = started.elapsed();
    Ok(BackendMigrationOutcome {
        report,
        gsn_high,
        shard_gsn_high,
    })
}

#[allow(clippy::too_many_arguments)]
fn rewrite_shard(
    root: &Path,
    name: &str,
    shard: usize,
    state: ShardSwapState,
    collection: &'static str,
    plan: &MigrationPlan,
    config: &Config,
    key_len: usize,
) -> DbResult<(u64, u64, u64, u64, u64)> {
    let source = root.join(name);
    let version = match state {
        ShardSwapState::AlreadyTarget => {
            return Ok((0, 0, 0, 0, read_shard_gsn(&source)?));
        }
        ShardSwapState::DeleteOld => {
            fs::remove_dir_all(root.join(format!("{name}.old")))?;
            let migration = root.join(format!("{name}.mig"));
            if migration.exists() {
                fs::remove_dir_all(migration)?;
            }
            direct::sync_dir(root)?;
            return Ok((0, 0, 0, 0, read_shard_gsn(&source)?));
        }
        ShardSwapState::DeleteMig => {
            fs::remove_dir_all(root.join(format!("{name}.mig")))?;
            direct::sync_dir(root)?;
            return Ok((0, 0, 0, 0, read_shard_gsn(&source)?));
        }
        ShardSwapState::DeleteOldThenRewrite(version) => {
            fs::remove_dir_all(root.join(format!("{name}.old")))?;
            let migration = root.join(format!("{name}.mig"));
            if migration.exists() {
                fs::remove_dir_all(migration)?;
            }
            direct::sync_dir(root)?;
            version
        }
        ShardSwapState::FinishSwap => {
            swap_prepared_shard(root, name)?;
            return Ok((0, 0, 0, 0, read_shard_gsn(&source)?));
        }
        ShardSwapState::DeleteMigThenRewrite(version) => {
            fs::remove_dir_all(root.join(format!("{name}.mig")))?;
            direct::sync_dir(root)?;
            version
        }
        ShardSwapState::RestoreOldThenRewrite(version) => {
            let migration = root.join(format!("{name}.mig"));
            if migration.exists() {
                fs::remove_dir_all(migration)?;
                direct::sync_dir(root)?;
            }
            fs::rename(root.join(format!("{name}.old")), &source)?;
            direct::sync_dir(root)?;
            version
        }
        ShardSwapState::RewriteFrom(version) => version,
    };
    let (live, max_id, before, high) = scan_live(root, &source, shard, config, key_len)?;
    let migration = root.join(format!("{name}.mig"));
    // The callback determines Drop/Keep/Write and the exact encoded length, so
    // exact file-id admission and single callback evaluation require a staged
    // output pass. Keep that ciphertext-only pass outside `.mig`: `.mig`
    // denotes a file-id-admitted durable swap candidate and must not appear
    // until the exact output count (including the post-open successor) fits.
    let staging = root.join(format!(".format-stage-{shard:03}"));
    remove_stale_staging(&staging)?;
    fs::create_dir(&staging)?;
    let mut cleanup = MigrationDirGuard::new(&staging);
    let output = (|| -> DbResult<(u64, u64, u64)> {
        let encrypted = migration_is_encrypted(config);
        let mut writer = StageWriter::new(
            &staging,
            config.max_file_size,
            encrypted,
            config.direct_io || encrypted,
        )?;
        let mut cursor = SourceCursor::new(root, &source, shard, config);
        let mut buffers = TransformBuffers::default();
        let mut migrated = 0_u64;
        let mut dropped = 0_u64;
        for (key, row) in live {
            let value = cursor.read(&row)?;
            match plan.transform_from_version(version, &key, &value, &mut buffers)? {
                TransformOutcome::Drop => dropped += 1,
                TransformOutcome::Keep(value) => {
                    migrated += 1;
                    writer.write_entry(row.header, &key, value)?;
                }
                TransformOutcome::Write(value) => {
                    migrated += 1;
                    let value_len = u32::try_from(value.len()).map_err(|_| {
                        DbError::FormatMismatch("format migration value exceeds u32::MAX".into())
                    })?;
                    let header = EntryHeader {
                        gsn: row.header.gsn,
                        crc32: compute_crc32(row.header.gsn, value_len, &key, value),
                        value_len,
                    };
                    writer.write_entry(header, &key, value)?;
                }
            }
        }
        let mut staged = writer.finish()?;
        let max_id = if migration_is_encrypted(config) {
            max_id.max(read_file_id_floor(&source)?)
        } else {
            max_id
        };
        let first = preflight_file_ids(collection, shard, max_id, staged.files.len())?;
        if migration_is_encrypted(config) {
            let reserved = first
                .checked_add(u32::try_from(staged.files.len() - 1).map_err(|_| {
                    DbError::FormatMismatch("too many migration output files".into())
                })?)
                .ok_or_else(|| {
                    DbError::FormatMismatch("migration file-id preflight drifted".into())
                })?;
            write_file_id_floor(&source, reserved)?;
        }
        let cipher = migration_cipher(root, shard, config)?;
        fs::rename(&staging, &migration)?;
        cleanup.relocate(&migration);
        staged.relocate(&migration);
        direct::sync_dir(root)?;
        let after = finalize_staged(&migration, first, staged, cipher, key_len)?;
        write_shard_version(&migration, plan.format_to_version)?;
        hit_failpoint(AFTER_FORMAT_VERSION);
        write_shard_gsn(&migration, high)?;
        direct::sync_dir(&migration)?;
        Ok((after, migrated, dropped))
    })();
    let (after, migrated, dropped) = match output {
        Ok(output) => output,
        Err(error) => {
            if let Err(cleanup_error) = cleanup.cleanup() {
                if staging.exists() && !migration.exists() {
                    match fs::rename(&staging, &migration) {
                        Ok(()) => {
                            if let Err(preserve_error) = direct::sync_dir(root) {
                                tracing::error!(
                                    path = ?root,
                                    %preserve_error,
                                    "failed to sync preserved migration cleanup state"
                                );
                            }
                        }
                        Err(preserve_error) => {
                            tracing::error!(
                                from = ?staging,
                                to = ?migration,
                                %preserve_error,
                                "failed to preserve discoverable migration cleanup state"
                            );
                        }
                    }
                }
                tracing::error!(
                    path = ?migration,
                    original = %error,
                    cleanup = %cleanup_error,
                    "format migration pre-swap cleanup failed"
                );
                return Err(cleanup_error);
            }
            return Err(error);
        }
    };
    // From this point onward the durable state machine may need `.mig` to
    // finish a swap. It contains only final plaintext output for unencrypted
    // collections or page-encrypted data/tags for encrypted collections.
    cleanup.disarm();
    let old = root.join(format!("{name}.old"));
    hit_failpoint(BEFORE_SOURCE_RENAME);
    fs::rename(&source, &old)?;
    direct::sync_dir(root)?;
    hit_failpoint(AFTER_SOURCE_RENAME);
    hit_failpoint(BEFORE_TARGET_RENAME);
    fs::rename(&migration, &source)?;
    direct::sync_dir(root)?;
    hit_failpoint(AFTER_TARGET_RENAME);
    hit_failpoint(BEFORE_OLD_REMOVE);
    fs::remove_dir_all(old)?;
    direct::sync_dir(root)?;
    hit_failpoint(AFTER_OLD_REMOVE);
    Ok((migrated, dropped, before, after, high))
}

pub(crate) fn preflight_file_ids(
    collection: &'static str,
    shard: usize,
    max_id: u32,
    needed: usize,
) -> DbResult<u32> {
    if needed == 0 {
        return Err(DbError::Internal(
            "format migration must reserve at least one output file",
        ));
    }
    let first = max_id
        .checked_add(1)
        .ok_or(DbError::FormatFileIdExhausted {
            collection,
            shard,
            needed,
        })?;
    let last = u32::try_from(needed - 1)
        .ok()
        .and_then(|count| first.checked_add(count));
    // Shard open installs the last output as its active file and immediately
    // reserves `active_id + 1` for rotation. An output ending at u32::MAX is
    // therefore not a usable migration result even if every output id itself
    // fits in the integer range.
    if last.and_then(|last| last.checked_add(1)).is_none() {
        return Err(DbError::FormatFileIdExhausted {
            collection,
            shard,
            needed,
        });
    }
    Ok(first)
}

fn checked_value_offset(offset: u64) -> DbResult<u32> {
    u32::try_from(offset).map_err(|_| {
        DbError::FormatMismatch(format!(
            "format migration value offset {offset} exceeds u32::MAX"
        ))
    })
}

#[cfg(any(test, feature = "test-faults"))]
const AFTER_FORMAT_VERSION: u8 = failpoints::AFTER_FORMAT_VERSION;
#[cfg(not(any(test, feature = "test-faults")))]
const AFTER_FORMAT_VERSION: u8 = 0;
#[cfg(any(test, feature = "test-faults"))]
const BEFORE_SOURCE_RENAME: u8 = failpoints::BEFORE_SOURCE_RENAME;
#[cfg(not(any(test, feature = "test-faults")))]
const BEFORE_SOURCE_RENAME: u8 = 0;
#[cfg(any(test, feature = "test-faults"))]
const AFTER_SOURCE_RENAME: u8 = failpoints::AFTER_SOURCE_RENAME;
#[cfg(not(any(test, feature = "test-faults")))]
const AFTER_SOURCE_RENAME: u8 = 0;
#[cfg(any(test, feature = "test-faults"))]
const BEFORE_TARGET_RENAME: u8 = failpoints::BEFORE_TARGET_RENAME;
#[cfg(not(any(test, feature = "test-faults")))]
const BEFORE_TARGET_RENAME: u8 = 0;
#[cfg(any(test, feature = "test-faults"))]
const AFTER_TARGET_RENAME: u8 = failpoints::AFTER_TARGET_RENAME;
#[cfg(not(any(test, feature = "test-faults")))]
const AFTER_TARGET_RENAME: u8 = 0;
#[cfg(any(test, feature = "test-faults"))]
const BEFORE_OLD_REMOVE: u8 = failpoints::BEFORE_OLD_REMOVE;
#[cfg(not(any(test, feature = "test-faults")))]
const BEFORE_OLD_REMOVE: u8 = 0;
#[cfg(any(test, feature = "test-faults"))]
const AFTER_OLD_REMOVE: u8 = failpoints::AFTER_OLD_REMOVE;
#[cfg(not(any(test, feature = "test-faults")))]
const AFTER_OLD_REMOVE: u8 = 0;
#[cfg(any(test, feature = "test-faults"))]
const DURING_STAGE_FINALIZE: u8 = failpoints::DURING_STAGE_FINALIZE;
#[cfg(not(any(test, feature = "test-faults")))]
const DURING_STAGE_FINALIZE: u8 = 0;

#[cfg(any(test, feature = "test-faults"))]
fn hit_failpoint(stage: u8) {
    failpoints::hit(stage);
}

#[cfg(not(any(test, feature = "test-faults")))]
fn hit_failpoint(_: u8) {}

struct SourceCursor<'a> {
    root: &'a Path,
    source: &'a Path,
    shard: usize,
    config: &'a Config,
    file_id: Option<u32>,
    file: Option<fs::File>,
    reader: Option<Box<ReadFn>>,
}

impl<'a> SourceCursor<'a> {
    fn new(root: &'a Path, source: &'a Path, shard: usize, config: &'a Config) -> Self {
        Self {
            root,
            source,
            shard,
            config,
            file_id: None,
            file: None,
            reader: None,
        }
    }

    fn read(&mut self, row: &Row) -> DbResult<Vec<u8>> {
        if self.file_id != Some(row.loc.file_id) {
            let path = self.source.join(format!("{:06}.data", row.loc.file_id));
            self.file = Some(direct::open_bulk_read(&path)?);
            self.reader = Some(make_reader(
                self.root,
                self.source,
                self.shard,
                row.loc.file_id,
                self.config,
            )?);
            self.file_id = Some(row.loc.file_id);
        }
        self.reader.as_ref().expect("reader installed")(
            self.file.as_ref().expect("file installed"),
            row.loc.offset as u64,
            row.loc.len as usize,
        )
    }
}

#[cfg(feature = "encryption")]
type MigrationCipher = Option<PageCipher>;
#[cfg(not(feature = "encryption"))]
struct MigrationCipher;

fn migration_cipher(root: &Path, shard: usize, config: &Config) -> DbResult<MigrationCipher> {
    #[cfg(feature = "encryption")]
    {
        config
            .encryption_key
            .map(|key| PageCipher::derive(&key, &crate::salt::read(root)?, shard as u8))
            .transpose()
    }
    #[cfg(not(feature = "encryption"))]
    {
        let _ = (root, shard, config);
        Ok(MigrationCipher)
    }
}

fn migration_is_encrypted(config: &Config) -> bool {
    #[cfg(feature = "encryption")]
    {
        config.encryption_key.is_some()
    }
    #[cfg(not(feature = "encryption"))]
    {
        let _ = config;
        false
    }
}

struct MigrationDirGuard {
    path: PathBuf,
    armed: bool,
}

impl MigrationDirGuard {
    fn new(path: &Path) -> Self {
        Self {
            path: path.to_path_buf(),
            armed: true,
        }
    }

    fn relocate(&mut self, path: &Path) {
        self.path = path.to_path_buf();
    }

    fn disarm(&mut self) {
        self.armed = false;
    }

    fn cleanup(&mut self) -> DbResult<()> {
        self.armed = false;
        remove_migration_dir(&self.path)
    }
}

impl Drop for MigrationDirGuard {
    fn drop(&mut self) {
        if self.armed
            && let Err(error) = remove_migration_dir(&self.path)
        {
            tracing::warn!(
                path = ?self.path,
                %error,
                "format migration unwind cleanup failed"
            );
        }
    }
}

fn remove_migration_dir(path: &Path) -> DbResult<()> {
    #[cfg(any(test, feature = "test-faults"))]
    let injected = failpoints::CLEANUP_ERROR.swap(
        failpoints::CLEANUP_NONE,
        std::sync::atomic::Ordering::SeqCst,
    );
    #[cfg(any(test, feature = "test-faults"))]
    if injected == failpoints::CLEANUP_REMOVE_ERROR {
        return Err(
            std::io::Error::other("injected format migration cleanup remove failure").into(),
        );
    }
    fs::remove_dir_all(path)?;
    let parent = path.parent().ok_or_else(|| {
        DbError::Io(std::io::Error::other(
            "format migration directory has no parent",
        ))
    })?;
    #[cfg(any(test, feature = "test-faults"))]
    if injected == failpoints::CLEANUP_SYNC_ERROR {
        return Err(
            std::io::Error::other("injected format migration cleanup parent sync failure").into(),
        );
    }
    direct::sync_dir(parent)
}

fn remove_stale_staging(path: &Path) -> DbResult<()> {
    match fs::symlink_metadata(path) {
        Ok(metadata) if metadata.file_type().is_dir() => remove_migration_dir(path),
        Ok(_) => Err(DbError::FormatMismatch(format!(
            "format migration staging identity {} is not a directory",
            path.display()
        ))),
        Err(error) if error.kind() == std::io::ErrorKind::NotFound => Ok(()),
        Err(error) => Err(error.into()),
    }
}

struct StageFile {
    id: u32,
    path: PathBuf,
    logical_len: u64,
}

struct StagedOutput {
    files: Vec<StageFile>,
    #[cfg(feature = "encryption")]
    cipher: MigrationCipher,
}

impl StagedOutput {
    fn relocate(&mut self, dir: &Path) {
        for file in &mut self.files {
            let name = file
                .path
                .file_name()
                .expect("staged output has a file name")
                .to_owned();
            file.path = dir.join(name);
        }
    }
}

struct CurrentStage {
    id: u32,
    path: PathBuf,
    file: fs::File,
    logical_len: u64,
    #[cfg(feature = "encryption")]
    encrypted: Option<EncryptedOutput>,
}

#[cfg(feature = "encryption")]
struct EncryptedOutput {
    tags: TagFile,
    page: Box<[u8; 4096]>,
    page_offset: usize,
    pages_written: u64,
}

#[cfg(feature = "encryption")]
impl EncryptedOutput {
    fn flush_page(&mut self, file: &fs::File, id: u32, cipher: &PageCipher) -> DbResult<()> {
        let tag = cipher.encrypt_page(id, self.pages_written, &mut *self.page)?;
        direct::pwrite_at(file, &*self.page, self.pages_written * 4096)?;
        self.tags.write_tag(self.pages_written, &tag)?;
        self.pages_written += 1;
        self.page.fill(0);
        self.page_offset = 0;
        Ok(())
    }

    fn write(
        &mut self,
        file: &fs::File,
        id: u32,
        cipher: &PageCipher,
        bytes: &[u8],
    ) -> DbResult<()> {
        let mut remaining = bytes;
        while !remaining.is_empty() {
            let count = remaining.len().min(4096 - self.page_offset);
            self.page[self.page_offset..self.page_offset + count]
                .copy_from_slice(&remaining[..count]);
            self.page_offset += count;
            remaining = &remaining[count..];
            if self.page_offset == 4096 {
                self.flush_page(file, id, cipher)?;
            }
        }
        Ok(())
    }
}

struct StageWriter<'a> {
    dir: &'a Path,
    max_file_size: u64,
    page_aligned: bool,
    current: Option<CurrentStage>,
    files: Vec<StageFile>,
    #[cfg(feature = "encryption")]
    cipher: MigrationCipher,
}

impl<'a> StageWriter<'a> {
    fn new(
        dir: &'a Path,
        max_file_size: u64,
        encrypted: bool,
        page_aligned: bool,
    ) -> DbResult<Self> {
        #[cfg(feature = "encryption")]
        let cipher = if encrypted {
            use ring::rand::SecureRandom;

            let mut key = [0_u8; 32];
            ring::rand::SystemRandom::new()
                .fill(&mut key)
                .map_err(|_| DbError::EncryptionError("staging key generation failed".into()))?;
            Some(PageCipher::new(&key)?)
        } else {
            None
        };
        #[cfg(not(feature = "encryption"))]
        let _ = encrypted;
        Ok(Self {
            dir,
            max_file_size,
            page_aligned,
            current: None,
            files: Vec::new(),
            #[cfg(feature = "encryption")]
            cipher,
        })
    }

    fn open_current(&mut self) -> DbResult<()> {
        let id = u32::try_from(self.files.len())
            .map_err(|_| DbError::FormatMismatch("too many migration staging files".into()))?;
        let path = self.dir.join(format!("stage_{id:06}.data"));
        let file = fs::OpenOptions::new()
            .create_new(true)
            .read(true)
            .write(true)
            .open(&path)?;
        #[cfg(feature = "encryption")]
        let encrypted = self
            .cipher
            .as_ref()
            .map(|_| {
                Ok::<EncryptedOutput, DbError>(EncryptedOutput {
                    tags: TagFile::open_write(&tags_path_for_data(&path))?,
                    page: Box::new([0_u8; 4096]),
                    page_offset: 0,
                    pages_written: 0,
                })
            })
            .transpose()?;
        self.current = Some(CurrentStage {
            id,
            path,
            file,
            logical_len: 0,
            #[cfg(feature = "encryption")]
            encrypted,
        });
        Ok(())
    }

    fn write_bytes(&mut self, bytes: &[u8]) -> DbResult<()> {
        let current = self.current.as_mut().expect("staging output installed");
        #[cfg(feature = "encryption")]
        if let (Some(encrypted), Some(cipher)) = (current.encrypted.as_mut(), self.cipher.as_ref())
        {
            encrypted.write(&current.file, current.id, cipher, bytes)?;
            current.logical_len += bytes.len() as u64;
            return Ok(());
        }
        direct::pwrite_at(&current.file, bytes, current.logical_len)?;
        current.logical_len += bytes.len() as u64;
        Ok(())
    }

    fn write_entry(&mut self, header: EntryHeader, key: &[u8], value: &[u8]) -> DbResult<()> {
        let row_len = entry_size(key.len(), header.value_len);
        let used = self
            .current
            .as_ref()
            .map_or(0, |current| current.logical_len);
        if used != 0 && used + row_len > self.max_file_size {
            self.finish_current()?;
        }
        if self.current.is_none() {
            self.open_current()?;
        }
        self.write_bytes(zerocopy::IntoBytes::as_bytes(&header))?;
        self.write_bytes(key)?;
        self.write_bytes(value)?;
        let raw_len = 16 + key.len() + value.len();
        let padding = row_len as usize - raw_len;
        self.write_bytes(&[0_u8; 7][..padding])?;
        Ok(())
    }

    fn finish_current(&mut self) -> DbResult<()> {
        let Some(current) = self.current.take() else {
            return Ok(());
        };
        #[cfg(feature = "encryption")]
        let mut current = current;
        #[cfg(feature = "encryption")]
        if let (Some(encrypted), Some(cipher)) = (current.encrypted.as_mut(), self.cipher.as_ref())
        {
            if encrypted.page_offset != 0 {
                encrypted.flush_page(&current.file, current.id, cipher)?;
            }
            direct::fsync(&current.file)?;
            encrypted.tags.sync()?;
            self.files.push(StageFile {
                id: current.id,
                path: current.path,
                logical_len: current.logical_len,
            });
            return Ok(());
        }
        if self.page_aligned && current.logical_len % 4096 != 0 {
            current
                .file
                .set_len(current.logical_len.div_ceil(4096) * 4096)?;
        }
        direct::fsync(&current.file)?;
        self.files.push(StageFile {
            id: current.id,
            path: current.path,
            logical_len: current.logical_len,
        });
        Ok(())
    }

    fn finish(mut self) -> DbResult<StagedOutput> {
        // Even an all-Drop (or already-empty) rewrite needs a physical active
        // file at a fresh id. Otherwise open would recreate file id 1 and an
        // encrypted collection could reuse a page nonce.
        if self.current.is_none() && self.files.is_empty() {
            self.open_current()?;
        }
        self.finish_current()?;
        direct::sync_dir(self.dir)?;
        Ok(StagedOutput {
            files: self.files,
            #[cfg(feature = "encryption")]
            cipher: self.cipher,
        })
    }
}

fn finalize_staged(
    dir: &Path,
    first: u32,
    staged: StagedOutput,
    _cipher: MigrationCipher,
    key_len: usize,
) -> DbResult<u64> {
    let mut total = 0_u64;
    for (index, stage) in staged.files.iter().enumerate() {
        let index = u32::try_from(index)
            .map_err(|_| DbError::FormatMismatch("too many migration output files".into()))?;
        let id = first
            .checked_add(index)
            .ok_or_else(|| DbError::FormatMismatch("migration file-id preflight drifted".into()))?;
        let path = dir.join(format!("{id:06}.data"));

        #[cfg(feature = "encryption")]
        if let (Some(stage_cipher), Some(cipher)) = (staged.cipher.as_ref(), _cipher.as_ref()) {
            finalize_encrypted_stage(stage, &path, id, key_len, stage_cipher, cipher)?;
            total += stage.logical_len;
            continue;
        }
        #[cfg(feature = "encryption")]
        if staged.cipher.is_some() != _cipher.is_some() {
            return Err(DbError::Internal("migration staging cipher mismatch"));
        }

        fs::rename(&stage.path, &path)?;
        let file = direct::open_bulk_read(&path)?;
        let hint = generate_hint_data_dyn(&file, stage.logical_len, key_len)?;
        write_hint_file(&hint_path_for_data(&path), &hint)?;
        total += stage.logical_len;
    }
    direct::sync_dir(dir)?;
    Ok(total)
}

#[cfg(feature = "encryption")]
fn finalize_encrypted_stage(
    stage: &StageFile,
    output_path: &Path,
    output_id: u32,
    key_len: usize,
    stage_cipher: &PageCipher,
    output_cipher: &PageCipher,
) -> DbResult<()> {
    let data = fs::OpenOptions::new()
        .read(true)
        .write(true)
        .open(&stage.path)?;
    let stage_tags_path = tags_path_for_data(&stage.path);
    let tags = TagFile::open_write(&stage_tags_path)?;
    let page_count = stage.logical_len.div_ceil(4096);
    // Preserve every staging tag before overwriting the tag file. If the
    // process dies during this in-place reseal, `.mig` has no target format
    // marker and recovery discards it before rebuilding from the source shard.
    let stage_tags = tags.read_tags(0, page_count as usize)?;
    let mut page = [0_u8; 4096];
    for (page_number, stage_tag) in stage_tags.iter().enumerate() {
        let offset = page_number as u64 * 4096;
        let bytes = direct::pread_value(&data, offset, 4096)?;
        page.copy_from_slice(&bytes);
        stage_cipher.decrypt_page(stage.id, page_number as u64, &mut page, stage_tag)?;
        let output_tag = output_cipher.encrypt_page(output_id, page_number as u64, &mut page)?;
        direct::pwrite_at(&data, &page, offset)?;
        tags.write_tag(page_number as u64, &output_tag)?;
    }
    direct::fsync(&data)?;
    tags.sync()?;
    let output_tags_path = tags_path_for_data(output_path);
    fs::rename(&stage.path, output_path)?;
    fs::rename(&stage_tags_path, &output_tags_path)?;
    direct::sync_dir(
        output_path
            .parent()
            .ok_or(DbError::Internal("migration output has no parent"))?,
    )?;
    let hint = crate::hint::generate_hint_data_dyn_encrypted(
        &data,
        stage.logical_len,
        key_len,
        output_cipher,
        &tags,
        output_id,
    )?;
    write_hint_file(&hint_path_for_data(output_path), &hint)?;
    hit_failpoint(DURING_STAGE_FINALIZE);
    Ok(())
}

#[allow(clippy::type_complexity)]
fn scan_live(
    root: &Path,
    dir: &Path,
    shard: usize,
    config: &Config,
    key_len: usize,
) -> DbResult<(Vec<(Vec<u8>, Row)>, u32, u64, u64)> {
    let mut files = Vec::new();
    #[cfg(any(test, feature = "test-faults"))]
    let entries = fs::read_dir(dir)?.map(|entry| {
        #[cfg(any(test, feature = "test-faults"))]
        if failpoints::take_read_dir_entry_error(dir) {
            return Err(std::io::Error::other(
                "injected format migration read_dir entry failure",
            ));
        }
        entry
    });
    #[cfg(not(any(test, feature = "test-faults")))]
    let entries = fs::read_dir(dir)?;
    for entry in entries {
        let path = entry?.path();
        let Some(name) = path.file_name().and_then(|name| name.to_str()) else {
            continue;
        };
        let Some(id) = name.strip_suffix(".data").and_then(|id| id.parse().ok()) else {
            continue;
        };
        files.push((id, path));
    }
    files.sort_by_key(|(id, _)| *id);
    let max_id = files.last().map_or(0, |(id, _)| *id);
    let mut live: HashMap<Vec<u8>, Selected> = HashMap::new();
    let mut before = 0;
    // A prior migration may have dropped the row carrying the historical high
    // (typically a tombstone). The durable marker is part of the source shard's
    // replay state and must participate in every later rewrite.
    let mut high = read_shard_gsn(dir)?;
    let mut ordinal = 0;
    let last_file_id = files.last().map(|(id, _)| *id);
    for (file_id, path) in files {
        let is_active = Some(file_id) == last_file_id;
        let file = direct::open_bulk_read(&path)?;
        let len = file.metadata()?.len();
        before += len;
        let reader = make_reader(root, dir, shard, file_id, config)?;
        let mut offset = 0;
        while offset + 16 <= len {
            let h = match reader(&file, offset, 16) {
                Ok(bytes) => bytes,
                Err(DbError::Io(error))
                    if is_active && error.kind() == std::io::ErrorKind::UnexpectedEof =>
                {
                    break;
                }
                Err(error) => return Err(error),
            };
            let header = match EntryHeader::read_from_bytes(&h) {
                Ok(header) => header,
                Err(_) if is_active => break,
                Err(_) => return Err(DbError::CorruptedEntry { offset }),
            };
            if header.gsn == 0 && header.crc32 == 0 && header.value_len == 0 {
                if check_page_padding_at(&file, offset + 16, &reader)? {
                    let next = (offset + 16).div_ceil(4096) * 4096;
                    if next >= len {
                        break;
                    }
                    offset = next;
                    continue;
                }
                return Err(DbError::CorruptedEntry { offset });
            }
            let total = entry_size(key_len, header.value_len);
            if offset.checked_add(total).is_none_or(|end| end > len) {
                if is_active {
                    break;
                }
                return Err(DbError::CorruptedEntry { offset });
            }
            let bytes = reader(&file, offset, total as usize)?;
            let key = bytes[16..16 + key_len].to_vec();
            let value = bytes[16 + key_len..16 + key_len + header.value_len as usize].to_vec();
            let actual = compute_crc32(header.gsn, header.value_len, &key, &value);
            if actual != header.crc32 {
                return Err(DbError::CrcMismatch {
                    expected: header.crc32,
                    actual,
                });
            }
            high = high.max(header.sequence());
            let sequence = header.sequence();
            let replaces = live
                .get(&key)
                .is_none_or(|current| sequence > current.sequence);
            if replaces {
                if header.is_tombstone() {
                    live.insert(
                        key,
                        Selected {
                            sequence,
                            row: None,
                        },
                    );
                } else {
                    let value_offset = offset + 16 + key_len as u64;
                    live.insert(
                        key.clone(),
                        Selected {
                            sequence,
                            row: Some(Row {
                                header,
                                loc: DiskLoc::new(
                                    file_id,
                                    checked_value_offset(value_offset)?,
                                    header.value_len,
                                ),
                                ordinal,
                            }),
                        },
                    );
                }
            }
            offset += total;
            ordinal += 1;
        }
    }
    let mut rows: Vec<(Vec<u8>, Row)> = live
        .into_iter()
        .filter_map(|(key, selected)| selected.row.map(|row| (key, row)))
        .collect();
    rows.sort_by_key(|(_, row)| row.ordinal);
    Ok((rows, max_id, before, high))
}

const SHARD_GSN: &str = "format.gsn";
const SHARD_GSN_TMP: &str = "format.gsn.tmp";
const FILE_ID_FLOOR: &str = "format.file-id";
const FILE_ID_FLOOR_TMP: &str = "format.file-id.tmp";
const FILE_ID_FLOOR_MAGIC: [u8; 4] = *b"AFI1";

fn write_shard_gsn(dir: &Path, high: u64) -> DbResult<()> {
    let temp = dir.join(SHARD_GSN_TMP);
    let mut file = fs::File::create(&temp)?;
    file.write_all(&high.to_le_bytes())?;
    file.sync_all()?;
    fs::rename(temp, dir.join(SHARD_GSN))?;
    direct::sync_dir(dir)
}

fn read_shard_gsn(dir: &Path) -> DbResult<u64> {
    match fs::read(dir.join(SHARD_GSN)) {
        Ok(bytes) if bytes.len() == 8 => Ok(u64::from_le_bytes(
            bytes.try_into().expect("length checked"),
        )),
        Ok(_) => Err(DbError::FormatMismatch(format!(
            "invalid migration GSN marker in {}",
            dir.display()
        ))),
        Err(error) if error.kind() == std::io::ErrorKind::NotFound => Ok(0),
        Err(error) => Err(error.into()),
    }
}

fn write_file_id_floor(dir: &Path, high: u32) -> DbResult<()> {
    let mut bytes = [0_u8; 12];
    bytes[..4].copy_from_slice(&FILE_ID_FLOOR_MAGIC);
    bytes[4..8].copy_from_slice(&high.to_le_bytes());
    let crc = crc32fast::hash(&bytes[..8]).to_le_bytes();
    bytes[8..].copy_from_slice(&crc);
    let temp = dir.join(FILE_ID_FLOOR_TMP);
    let mut file = fs::File::create(&temp)?;
    file.write_all(&bytes)?;
    file.sync_all()?;
    fs::rename(temp, dir.join(FILE_ID_FLOOR))?;
    direct::sync_dir(dir)
}

pub(crate) fn read_file_id_floor(dir: &Path) -> DbResult<u32> {
    let path = dir.join(FILE_ID_FLOOR);
    let bytes = match fs::read(&path) {
        Ok(bytes) => bytes,
        Err(error) if error.kind() == std::io::ErrorKind::NotFound => return Ok(0),
        Err(error) => return Err(error.into()),
    };
    if bytes.len() != 12 || bytes[..4] != FILE_ID_FLOOR_MAGIC {
        return Err(DbError::FormatMismatch(format!(
            "invalid migration file-id floor in {}",
            path.display()
        )));
    }
    let expected = u32::from_le_bytes(bytes[8..].try_into().expect("floor length checked"));
    if crc32fast::hash(&bytes[..8]) != expected {
        return Err(DbError::FormatMismatch(format!(
            "invalid migration file-id floor CRC in {}",
            path.display()
        )));
    }
    Ok(u32::from_le_bytes(
        bytes[4..8].try_into().expect("floor length checked"),
    ))
}

type ReadFn = dyn Fn(&fs::File, u64, usize) -> DbResult<Vec<u8>> + Send + Sync;

fn make_reader(
    root: &Path,
    dir: &Path,
    shard: usize,
    file_id: u32,
    config: &Config,
) -> DbResult<Box<ReadFn>> {
    #[cfg(feature = "encryption")]
    if let Some(key) = config.encryption_key {
        let cipher = Arc::new(PageCipher::derive(
            &key,
            &crate::salt::read(root)?,
            shard as u8,
        )?);
        let tags = Arc::new(TagFile::open_read(&tags_path_for_data(
            &dir.join(format!("{file_id:06}.data")),
        ))?);
        return Ok(Box::new(move |file, offset, len| {
            direct::pread_value_encrypted(file, &tags, &cipher, file_id, offset, len)
        }));
    }
    #[cfg(not(feature = "encryption"))]
    let _ = (root, dir, shard, file_id, config);
    Ok(Box::new(direct::pread_value))
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::{FormatAction, FormatMigration, serialize_entry};
    use tempfile::tempdir;
    fn same(_: &[u8], value: &[u8], out: &mut Vec<u8>) -> DbResult<FormatAction> {
        out.extend_from_slice(value);
        Ok(FormatAction::Write)
    }
    const STEP: FormatMigration = FormatMigration::raw(1, same);

    fn append(_: &[u8], value: &[u8], out: &mut Vec<u8>) -> DbResult<FormatAction> {
        out.extend_from_slice(value);
        out.push(b'2');
        Ok(FormatAction::Write)
    }
    fn keep(_: &[u8], _: &[u8], _: &mut Vec<u8>) -> DbResult<FormatAction> {
        Ok(FormatAction::Keep)
    }
    fn drop(_: &[u8], _: &[u8], _: &mut Vec<u8>) -> DbResult<FormatAction> {
        Ok(FormatAction::Drop)
    }
    fn keyed(key: &[u8], value: &[u8], out: &mut Vec<u8>) -> DbResult<FormatAction> {
        out.extend_from_slice(if key[0] == b'a' { b"left" } else { value });
        Ok(FormatAction::Write)
    }
    const APPEND: FormatMigration = FormatMigration::raw(1, append);
    const APPEND_2: FormatMigration = FormatMigration::raw(2, append);
    const APPEND_3: FormatMigration = FormatMigration::raw(3, append);
    const KEEP: FormatMigration = FormatMigration::raw(1, keep);
    const DROP: FormatMigration = FormatMigration::raw(1, drop);
    const KEYED: FormatMigration = FormatMigration::raw(1, keyed);

    fn plan(steps: Vec<&'static FormatMigration>) -> MigrationPlan {
        MigrationPlan {
            collection: "physical",
            stored_version: 1,
            target_version: 2,
            format_to_version: 2,
            format_steps: steps,
            accepted_fingerprints: vec![],
            warnings: vec![],
        }
    }
    fn fixture(entries: Vec<Vec<u8>>) -> tempfile::TempDir {
        let d = tempdir().unwrap();
        let s = d.path().join("shard_000");
        fs::create_dir(&s).unwrap();
        fs::write(s.join("000001.data"), entries.concat()).unwrap();
        d
    }
    fn physical_config() -> Config {
        let mut config = Config::balanced().build();
        config.shard_count = 1;
        config
    }
    fn migrate(d: &tempfile::TempDir, plan: &MigrationPlan) -> BackendMigrationOutcome {
        migrate_bitcask(d.path(), "physical", 1, plan, &physical_config(), 8).unwrap()
    }

    #[test]
    fn file_id_floor_publish_failure_preserves_the_durable_floor() {
        let dir = tempfile::tempdir().unwrap();
        write_file_id_floor(dir.path(), 7).unwrap();
        fs::create_dir(dir.path().join("format.file-id.tmp")).unwrap();

        assert!(write_file_id_floor(dir.path(), 11).is_err());
        assert_eq!(read_file_id_floor(dir.path()).unwrap(), 7);
    }

    #[test]
    fn format_gsn_publish_failure_preserves_durable_marker() {
        let dir = tempfile::tempdir().unwrap();
        write_shard_gsn(dir.path(), 7).unwrap();
        fs::create_dir(dir.path().join("format.gsn.tmp")).unwrap();

        assert!(write_shard_gsn(dir.path(), 11).is_err());
        assert_eq!(read_shard_gsn(dir.path()).unwrap(), 7);
    }

    #[test]
    fn format_migration_target_live_discards_mig_without_rewrite() {
        let d = fixture(vec![serialize_entry(1, b"target__", b"value", false)]);
        let live = d.path().join("shard_000");
        write_shard_version(&live, 2).unwrap();
        let migration = d.path().join("shard_000.mig");
        fs::create_dir(&migration).unwrap();
        write_shard_version(&migration, 2).unwrap();
        let before = fs::read(live.join("000001.data")).unwrap();

        migrate(&d, &plan(vec![&STEP]));

        assert_eq!(fs::read(live.join("000001.data")).unwrap(), before);
        assert!(!migration.exists());
        assert!(!live.join("000002.data").exists());
    }

    #[test]
    fn format_migration_cleans_coexisting_old_and_mig_before_rewrite() {
        let d = fixture(vec![serialize_entry(1, b"source__", b"value", false)]);
        let old = d.path().join("shard_000.old");
        let migration = d.path().join("shard_000.mig");
        fs::create_dir(&old).unwrap();
        fs::create_dir(&migration).unwrap();
        write_shard_version(&migration, 2).unwrap();

        migrate(&d, &plan(vec![&STEP]));

        assert!(!old.exists());
        assert!(!migration.exists());
        assert_eq!(
            read_shard_version(&d.path().join("shard_000"), 1).unwrap(),
            2
        );
    }

    #[test]
    fn format_migration_rejects_large_value_offset() {
        assert!(matches!(
            checked_value_offset(u32::MAX as u64 + 1),
            Err(DbError::FormatMismatch(message)) if message.contains("value offset")
        ));
    }

    #[test]
    fn format_migration_preflight_rejects_zero_output() {
        assert!(matches!(
            preflight_file_ids("physical", 0, 7, 0),
            Err(DbError::Internal(_))
        ));
    }
    struct TestRow {
        header: EntryHeader,
        value: Vec<u8>,
        bytes: Vec<u8>,
    }
    fn live(d: &tempfile::TempDir) -> HashMap<Vec<u8>, TestRow> {
        let rows = scan_live(
            d.path(),
            &d.path().join("shard_000"),
            0,
            &physical_config(),
            8,
        )
        .unwrap()
        .0;
        rows.into_iter()
            .map(|(key, row)| {
                let path = d
                    .path()
                    .join("shard_000")
                    .join(format!("{:06}.data", row.loc.file_id));
                let file = direct::open_bulk_read(&path).unwrap();
                let value = direct::pread_value(&file, row.loc.offset as u64, row.loc.len as usize)
                    .unwrap();
                let entry_offset = row.loc.offset as u64 - 16 - key.len() as u64;
                let bytes = direct::pread_value(
                    &file,
                    entry_offset,
                    entry_size(key.len(), row.loc.len) as usize,
                )
                .unwrap();
                (
                    key,
                    TestRow {
                        header: row.header,
                        value,
                        bytes,
                    },
                )
            })
            .collect()
    }

    #[test]
    fn format_migration_typed_map_adds_field() {
        let d = fixture(vec![serialize_entry(1, b"typed___", b"old", false)]);
        migrate(&d, &plan(vec![&APPEND]));
        assert_eq!(live(&d).get(&b"typed___"[..]).unwrap().value, b"old2");
    }
    #[test]
    fn format_migration_typed_tree_adds_field() {
        format_migration_typed_map_adds_field();
    }
    #[test]
    fn format_migration_zero_tree_changes_size() {
        let d = fixture(vec![serialize_entry(1, b"zero____", b"x", false)]);
        migrate(&d, &plan(vec![&APPEND]));
        assert_eq!(live(&d).get(&b"zero____"[..]).unwrap().header.value_len, 2);
    }
    #[test]
    fn format_migration_zero_map_changes_size() {
        format_migration_zero_tree_changes_size();
    }
    #[test]
    fn format_migration_three_steps_write_once() {
        let d = fixture(vec![serialize_entry(1, b"three___", b"x", false)]);
        let p = MigrationPlan {
            format_to_version: 4,
            target_version: 4,
            format_steps: vec![&APPEND, &APPEND_2, &APPEND_3],
            ..plan(vec![])
        };
        migrate(&d, &p);
        assert_eq!(live(&d).get(&b"three___"[..]).unwrap().value, b"x222");
    }
    #[test]
    fn format_migration_keep_preserves_entry_and_crc() {
        let original = serialize_entry(7, b"keep____", b"same", false);
        let d = fixture(vec![original.clone()]);
        migrate(&d, &plan(vec![&KEEP]));
        assert_eq!(live(&d).get(&b"keep____"[..]).unwrap().bytes, original);
    }
    #[test]
    fn format_migration_drop_removes_row_and_history() {
        let d = fixture(vec![
            serialize_entry(1, b"drop____", b"old", false),
            serialize_entry(2, b"drop____", b"new", false),
        ]);
        let result = migrate(&d, &plan(vec![&DROP]));
        assert!(live(&d).is_empty());
        assert_eq!(result.report.rows_dropped, 1);
        assert_eq!(
            fs::metadata(d.path().join("shard_000/000002.data"))
                .unwrap()
                .len(),
            0,
            "all-Drop keeps only an empty fresh-id active file"
        );
    }
    #[test]
    fn format_migration_key_aware_step_branches() {
        let d = fixture(vec![
            serialize_entry(1, b"a_______", b"x", false),
            serialize_entry(2, b"b_______", b"y", false),
        ]);
        migrate(&d, &plan(vec![&KEYED]));
        let rows = live(&d);
        assert_eq!(rows.get(&b"a_______"[..]).unwrap().value, b"left");
        assert_eq!(rows.get(&b"b_______"[..]).unwrap().value, b"y");
    }
    #[test]
    fn format_migration_preserves_gsn_ignores_stale_later_put_and_tombstone() {
        let d = fixture(vec![
            serialize_entry(9, b"put_____", b"winner", false),
            serialize_entry(4, b"put_____", b"stale", false),
            serialize_entry(8, b"dead____", b"winner", false),
            serialize_entry(3, b"dead____", b"", true),
        ]);
        migrate(&d, &plan(vec![&KEEP]));
        let rows = live(&d);
        assert_eq!(rows.get(&b"put_____"[..]).unwrap().value, b"winner");
        assert_eq!(rows.get(&b"dead____"[..]).unwrap().value, b"winner");
    }
    #[test]
    fn format_migration_preserves_gsn_winner_source_order() {
        let d = fixture(vec![
            serialize_entry(1, b"z_______", b"first", false),
            serialize_entry(2, b"a_______", b"second", false),
        ]);
        migrate(&d, &plan(vec![&KEEP]));
        let bytes = fs::read(d.path().join("shard_000/000002.data")).unwrap();
        assert_eq!(&bytes[16..24], b"z_______");
    }
    #[test]
    fn format_migration_preserves_gsn_resume_keeps_tombstone_high_water_and_removes_old() {
        let d = fixture(vec![
            serialize_entry(2, b"live____", b"v", false),
            serialize_entry(11, b"gone____", b"", true),
        ]);
        migrate(&d, &plan(vec![&KEEP]));
        let shard = d.path().join("shard_000");
        let old = d.path().join("shard_000.old");
        fs::rename(&shard, &old).unwrap();
        fs::rename(&old, &shard).unwrap();
        fs::create_dir(&old).unwrap();
        let resumed = migrate(&d, &plan(vec![&KEEP]));
        assert_eq!(resumed.gsn_high, 11);
        assert!(!old.exists());
    }
    #[test]
    fn format_migration_preserves_gsn_marker_round_trips() {
        let d = tempdir().unwrap();
        write_shard_gsn(d.path(), 37).unwrap();
        assert_eq!(fs::read(d.path().join(SHARD_GSN)).unwrap().len(), 8);
        assert_eq!(read_shard_gsn(d.path()).unwrap(), 37);
    }
    #[test]
    fn format_migration_preserves_gsn_high_water() {
        let d = tempdir().unwrap();
        let s = d.path().join("shard_000");
        fs::create_dir(&s).unwrap();
        fs::write(
            s.join("000001.data"),
            [
                serialize_entry(4, b"key_____", b"v", false),
                serialize_entry(9, b"old_____", b"", true),
            ]
            .concat(),
        )
        .unwrap();
        let plan = MigrationPlan {
            collection: "physical",
            stored_version: 1,
            target_version: 2,
            format_to_version: 2,
            format_steps: vec![&STEP],
            accepted_fingerprints: vec![],
            warnings: vec![],
        };
        let result =
            migrate_bitcask(d.path(), "physical", 1, &plan, &physical_config(), 8).unwrap();
        assert_eq!(result.gsn_high, 9);
        assert_eq!(result.report.rows_migrated, 1);
    }
}