armdb 0.11.2

//! Background hint-file generation for retired Bitcask data files.
//!
//! Rotation retires the active data file and its hint file is a full scan of
//! that file — up to `max_file_size` bytes (256 MiB by default, 1 GiB under
//! `Config::large_values`). Running that scan inside `rotate()` holds the shard
//! lock for its whole duration, so the writer that happened to trip the rotation
//! and every other writer to that shard stall on work none of them needs: a
//! missing hint costs a one-file scan at the next open, nothing more (F8 of
//! libfunc/armour#83).
//!
//! This module moves the scan onto a worker thread owned by the collection's
//! [`Engine`](crate::engine::Engine). The thread is started lazily on the first
//! rotation that would generate a hint, so a collection that never rotates —
//! every FixedStore collection, and any Bitcask one with `hints: false` — never
//! spawns one.
//!
//! **The retired file is immutable**, and the job holds its `Arc<ImmutableFile>`
//! (hence its fd), so the scan is safe even if compaction unlinks the file
//! underneath it. What is *not* safe is writing the hint after compaction
//! retired the data file: that would leave a `{id}.hint` no `{id}.data` matches.
//! The final rename is therefore taken under the shard lock, which compaction
//! also holds while it deletes — see [`HintService::finish_job`].

use std::path::PathBuf;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::mpsc::{Receiver, SyncSender, TrySendError, sync_channel};
use std::sync::{Arc, OnceLock, Weak};
use std::thread::JoinHandle;

use crate::error::DbResult;
use crate::shard::{ImmutableFile, Shards};
use crate::sync::{self, Mutex};

/// Jobs a collection may have in flight before rotation starts dropping them.
///
/// Each queued job pins one retired file's fd and is dropped — not blocked on —
/// when the queue is full, so this bounds the fd and memory a burst of
/// rotations can pin. Losing a hint is a start-up cost at the next open, never
/// a correctness one; blocking rotation instead would reintroduce exactly the
/// stall this module exists to remove.
const QUEUE_CAPACITY: usize = 16;

/// One retired file's hint generation.
pub(crate) struct HintJob {
    pub(crate) shard_idx: usize,
    pub(crate) file_id: u32,
    /// The retired file, borrowed from the shard's `immutable` list. Holding it
    /// keeps the fd open for the scan even if compaction unlinks the path.
    pub(crate) file: Arc<ImmutableFile>,
    /// Logical end of the retired file — the rotation's pre-padding offset.
    pub(crate) file_len: u64,
    pub(crate) key_len: usize,
    /// Path of the `.data` file; the hint path is derived from it.
    pub(crate) data_path: PathBuf,
    #[cfg(feature = "encryption")]
    pub(crate) cipher: Option<Arc<crate::crypto::PageCipher>>,
    /// Test-only seam: park the generation here so a test can observe *where*
    /// it runs. See [`crate::shard::ShardInner::test_hint_gate`].
    #[cfg(test)]
    pub(crate) gate: Option<Arc<crate::shard::HintGate>>,
}

/// Per-collection hint generation service. Owned by the `Engine`, shared with
/// every shard of that collection.
pub(crate) struct HintService {
    /// Sender + thread handle, created on the first [`Self::submit`].
    worker: Mutex<Option<Worker>>,
    /// Set once by the `Engine` right after `Shards` is built; the worker needs
    /// it to take the shard lock for the final rename.
    shards: OnceLock<Weak<Shards>>,
    /// Jobs dropped because the queue was full. Read by tests; also the number
    /// worth alerting on if it ever moves in production.
    dropped: AtomicU64,
}

struct Worker {
    tx: SyncSender<HintJob>,
    handle: JoinHandle<()>,
}

impl HintService {
    pub(crate) fn new() -> Self {
        Self {
            worker: Mutex::new(None),
            shards: OnceLock::new(),
            dropped: AtomicU64::new(0),
        }
    }

    /// Give the service the shard set its jobs belong to. Called once, by the
    /// `Engine`, after `Shards` exists.
    pub(crate) fn attach_shards(&self, shards: Weak<Shards>) {
        let _ = self.shards.set(shards);
    }

    /// Number of jobs dropped because the queue was full.
    #[cfg_attr(not(test), allow(dead_code))]
    pub(crate) fn dropped(&self) -> u64 {
        self.dropped.load(Ordering::Relaxed)
    }

    /// Hand a retired file's hint generation to the worker, starting it if this
    /// is the first one. Never blocks: a full queue drops the job with a `warn!`.
    pub(crate) fn submit(&self, job: HintJob) {
        let file_id = job.file_id;
        let mut guard = sync::lock(&self.worker);
        let worker = guard.get_or_insert_with(|| self.spawn());
        match worker.tx.try_send(job) {
            Ok(()) => {}
            Err(TrySendError::Full(job)) => {
                self.dropped.fetch_add(1, Ordering::Relaxed);
                tracing::warn!(
                    file_id = job.file_id,
                    capacity = QUEUE_CAPACITY,
                    "hint queue full — skipping hint for this file; recovery \
                     will scan it instead"
                );
            }
            Err(TrySendError::Disconnected(_)) => {
                tracing::warn!(file_id, "hint worker stopped — skipping hint for this file");
            }
        }
    }

    /// The worker must **not** hold an `Arc<Self>`: the service owns the
    /// sender that keeps the worker alive, so an `Arc` back would be a cycle —
    /// the service would never drop, and its thread would outlive the
    /// collection. It gets the `Weak<Shards>` it needs and nothing else.
    fn spawn(&self) -> Worker {
        let (tx, rx) = sync_channel(QUEUE_CAPACITY);
        let shards = self.shards.get().cloned().unwrap_or_else(Weak::new);
        let handle = std::thread::Builder::new()
            .name("armdb-hint".into())
            .spawn(move || run(&shards, rx))
            .expect("spawn hint worker");
        Worker { tx, handle }
    }

    /// Drain every queued job and stop the worker. Idempotent.
    ///
    /// Dropping the sender ends the worker's `recv` loop only after it has
    /// processed what is already queued, so joining here is what makes a clean
    /// close leave every retired file's hint on disk.
    pub(crate) fn drain_and_stop(&self) {
        let worker = sync::lock(&self.worker).take();
        if let Some(Worker { tx, handle }) = worker {
            drop(tx);
            let _ = handle.join();
        }
    }
}

/// Live worker threads. A collection that is dropped must not leave one behind —
/// the worker holding an `Arc` back to its service was exactly the cycle that
/// made `Drop` unreachable, so this counter is what a test can assert on.
#[cfg(test)]
pub(crate) static LIVE_WORKERS: std::sync::atomic::AtomicUsize =
    std::sync::atomic::AtomicUsize::new(0);

fn run(shards: &Weak<Shards>, rx: Receiver<HintJob>) {
    #[cfg(test)]
    LIVE_WORKERS.fetch_add(1, Ordering::Relaxed);
    #[cfg(test)]
    let _guard = scopeguard();

    while let Ok(job) = rx.recv() {
        let file_id = job.file_id;
        if let Err(e) = process(shards, job) {
            tracing::warn!(
                file_id,
                error = %e,
                "hint generation failed — recovery will scan this file instead"
            );
        }
    }
}

/// Generate the hint outside every lock, then publish it under the shard lock.
/// Errors are the caller's to log: a failure here costs a scan at the next open,
/// which is why nothing on the write path waits for this.
fn process(shards: &Weak<Shards>, job: HintJob) -> DbResult<()> {
    #[cfg(test)]
    if let Some(gate) = &job.gate {
        gate.wait();
    }

    let data = match generate(&job) {
        Ok(data) => data,
        // The scan stopped short of the file's logical end. `rotate` treated
        // this as "skip the hint, warn" when it ran the scan itself; keep
        // that meaning here rather than reporting it as a failure.
        Err(crate::error::DbError::CorruptedEntry { offset }) => {
            tracing::warn!(
                file_id = job.file_id,
                offset,
                "hint generation stopped early — skipping hint file"
            );
            return Ok(());
        }
        Err(e) => return Err(e),
    };

    let hint_path = crate::hint::hint_path_for_data(&job.data_path);
    let tmp = crate::hint::write_hint_tmp(&hint_path, &data)?;
    finish_job(shards, &job, &tmp, &hint_path)
}

/// Publish (or discard) a generated hint under the shard lock.
///
/// Compaction deletes `{id}.data` and `{id}.hint` while holding that same
/// lock, so checking the file is still an immutable member and renaming
/// under one acquisition is what keeps a hint from outliving its data file.
fn finish_job(
    shards: &Weak<Shards>,
    job: &HintJob,
    tmp: &std::path::Path,
    hint_path: &std::path::Path,
) -> DbResult<()> {
    let Some(shards) = shards.upgrade() else {
        // The collection is gone; nothing to publish into.
        let _ = std::fs::remove_file(tmp);
        return Ok(());
    };
    let Some(shard) = shards.get(job.shard_idx) else {
        let _ = std::fs::remove_file(tmp);
        return Ok(());
    };

    // The check and the rename happen under one acquisition: compaction
    // deletes the data file while holding this same lock, so splitting them
    // would leave a window where the file dies between them and the hint is
    // published anyway.
    let inner = shard.lock();
    if !inner.immutable.iter().any(|f| f.file_id == job.file_id) {
        drop(inner);
        let _ = std::fs::remove_file(tmp);
        tracing::debug!(
            file_id = job.file_id,
            "retired file was compacted before its hint was written"
        );
        return Ok(());
    }
    crate::hint::publish_hint_tmp(tmp, hint_path)
}

impl Drop for HintService {
    fn drop(&mut self) {
        self.drain_and_stop();
    }
}

fn generate(job: &HintJob) -> DbResult<Vec<u8>> {
    #[cfg(feature = "encryption")]
    {
        if let (Some(cipher), Some(tag_file)) = (&job.cipher, &job.file.tag_file) {
            return crate::hint::generate_hint_data_dyn_encrypted(
                &job.file.file,
                job.file_len,
                job.key_len,
                cipher,
                tag_file.as_ref(),
                job.file_id,
            );
        }
    }
    crate::hint::generate_hint_data_dyn(&job.file.file, job.file_len, job.key_len)
}

/// Decrements [`LIVE_WORKERS`] when the worker's `run` returns, however it
/// returns.
#[cfg(test)]
fn scopeguard() -> impl Drop {
    struct Guard;
    impl Drop for Guard {
        fn drop(&mut self) {
            LIVE_WORKERS.fetch_sub(1, Ordering::Relaxed);
        }
    }
    Guard
}