hf2q 0.1.1

//! ADR-017 §A.2 — `AsyncWriterHandle`: background-thread KV-block writer.
//!
//! `AsyncWriterHandle` runs `DiskBlockStore::write_block_sync` on a
//! dedicated background thread. Inference threads enqueue [`WriteJob`]s
//! and continue without blocking on disk I/O. This is the load-bearing
//! property of ADR-017's "writes off the prefill path" promise (§R-P1).
//!
//! ## Lifecycle
//!
//!   1. [`AsyncWriterHandle::spawn`] starts a worker thread + sync_channel.
//!   2. Inference threads call [`AsyncWriterHandle::enqueue`].
//!   3. The worker drains jobs, calls `write_block_sync`, fires the
//!      optional completion callback.
//!   4. [`AsyncWriterHandle::shutdown`] drops the sender, signaling the
//!      worker to drain pending jobs and exit. `join()` waits for clean
//!      drain. Errors during drain are logged via `tracing::warn!` but
//!      do NOT panic (a write failure mid-shutdown should not kill the
//!      `cmd_serve` process).
//!
//! ## Why `mpsc::sync_channel`
//!
//! Bounded back-pressure (per §R-P1: writer must NOT unboundedly buffer
//! blocks if the disk is slow). `try_send` returns `TrySendError::Full`
//! when the channel is at capacity, letting the spiller short-circuit
//! (skip the spill rather than stall the prefill thread). Phase A.3
//! wires the back-pressure policy on top of this.
//!
//! ## Why no panic on I/O error
//!
//! Mirrors the precedent from `serve/cache.rs`: a failed atomic-rename
//! on one block must not corrupt unrelated session state. The worker
//! reports the error to the optional completion channel, emits
//! `tracing::warn!`, and continues with the next job. Operators see
//! the warning in the serve log; the next successful write still
//! lands.

use std::io;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::mpsc::{self, RecvError, SyncSender, TryRecvError, TrySendError};
use std::sync::Arc;
use std::thread::{self, JoinHandle};
use std::time::Duration;

use crate::serve::kv_persist::block_store::{DiskBlockStore, WriteJob};

/// Sender + worker-thread join handle. Drop semantics: on `shutdown`
/// the sender is dropped, the worker drains the channel, exits, and
/// `join()` returns. Calling `drop` on the handle without `shutdown`
/// also closes the channel (sender dropped) and the worker exits — but
/// `shutdown` is the explicit, error-checked path.
pub struct AsyncWriterHandle {
    tx: Option<SyncSender<WriteJob>>,
    join_handle: Option<JoinHandle<()>>,
    /// ADR-017 iter-16 — pending-job depth counter. Incremented on
    /// every successful enqueue (`try_send` Ok / `send` Ok),
    /// decremented in `process_job` after the worker dequeues + runs
    /// the write. Lock-free `Arc<AtomicUsize>` so probes from any
    /// thread are O(1) and never block the writer. Shared with the
    /// worker via the `worker_loop` parameter list.
    pending: Arc<AtomicUsize>,
}

impl AsyncWriterHandle {
    /// Start a worker thread bound to `store`. `channel_capacity`
    /// controls back-pressure: at capacity, [`enqueue`] returns
    /// `TrySendError::Full`. Production callers pass a small bound
    /// (4-16) so the spiller short-circuits under sustained load.
    pub fn spawn(store: Arc<DiskBlockStore>, channel_capacity: usize) -> Self {
        let (tx, rx) = mpsc::sync_channel::<WriteJob>(channel_capacity);
        let pending = Arc::new(AtomicUsize::new(0));
        let pending_for_worker = Arc::clone(&pending);
        let join_handle = thread::Builder::new()
            .name("hf2q-kv-writer".to_string())
            .spawn(move || {
                worker_loop(store, rx, pending_for_worker);
            })
            .expect("spawn hf2q-kv-writer thread");
        Self {
            tx: Some(tx),
            join_handle: Some(join_handle),
            pending,
        }
    }

    /// Non-blocking enqueue. Returns `TrySendError::Full` if the
    /// channel is at capacity, `TrySendError::Disconnected` if the
    /// worker has shut down. Inference threads use this to short-
    /// circuit a stall-prone spill (§R-P1).
    ///
    /// The `Err` variant is intentionally large — it carries ownership
    /// of the failed [`WriteJob`] back to the caller so the body bytes
    /// can be retried or dropped without re-extracting from the engine.
    /// Boxing would force a heap alloc on the hot success path.
    #[allow(clippy::result_large_err)]
    pub fn enqueue(&self, job: WriteJob) -> Result<(), TrySendError<WriteJob>> {
        match self.tx.as_ref() {
            Some(tx) => match tx.try_send(job) {
                Ok(()) => {
                    // Successful enqueue → bump the pending-depth
                    // counter. The worker decrements after dequeue.
                    self.pending.fetch_add(1, Ordering::Relaxed);
                    Ok(())
                }
                Err(e) => Err(e),
            },
            None => Err(TrySendError::Disconnected(job)),
        }
    }

    /// Blocking enqueue. Used by tests and by callers that explicitly
    /// want back-pressure to wait rather than short-circuit. The `Err`
    /// variant carries the failed job for the same reason as
    /// [`Self::enqueue`].
    #[allow(clippy::result_large_err)]
    pub fn enqueue_blocking(&self, job: WriteJob) -> Result<(), mpsc::SendError<WriteJob>> {
        match self.tx.as_ref() {
            Some(tx) => match tx.send(job) {
                Ok(()) => {
                    self.pending.fetch_add(1, Ordering::Relaxed);
                    Ok(())
                }
                Err(e) => Err(e),
            },
            None => Err(mpsc::SendError(job)),
        }
    }

    /// ADR-017 iter-16 — read the pending-job depth without taking
    /// any lock. The counter is incremented inside [`Self::enqueue`] /
    /// [`Self::enqueue_blocking`] on successful send, and decremented
    /// inside `process_job` after the worker has dequeued the job
    /// (i.e. AFTER the channel slot is freed). The visible value is
    /// therefore an upper bound on the in-flight write count: at any
    /// instant there may be up to `pending_jobs()` jobs that have been
    /// accepted by the channel but whose `write_block_sync` may or may
    /// not have completed.
    ///
    /// Production callers of `BlockPrefixCacheSpiller` may use this as
    /// a back-pressure signal (e.g. shed lower-priority spill jobs
    /// when depth > N). Tests use it to verify that the writer's
    /// enqueue path actually queues work.
    ///
    /// `Ordering::Relaxed` is sufficient — the counter is a snapshot
    /// metric, not a synchronization barrier.
    pub fn pending_jobs(&self) -> usize {
        self.pending.load(Ordering::Relaxed)
    }

    /// Drop the sender (signaling the worker to drain) and join the
    /// worker thread. Returns `Err` if `join()` propagates a panic
    /// from the worker. (The worker is engineered NOT to panic on
    /// I/O errors; a panic here means a bug in the worker, not a
    /// disk-full situation.)
    pub fn shutdown(mut self) -> io::Result<()> {
        // Drop the sender so the worker's recv() returns RecvError
        // and the worker drains + exits.
        self.tx.take();
        if let Some(jh) = self.join_handle.take() {
            jh.join().map_err(|panic_payload| {
                io::Error::other(format!(
                    "kv_persist writer worker panicked: {:?}",
                    panic_payload
                        .downcast_ref::<&str>()
                        .copied()
                        .or_else(|| { panic_payload.downcast_ref::<String>().map(|s| s.as_str()) })
                        .unwrap_or("<non-string panic payload>")
                ))
            })?;
        }
        Ok(())
    }
}

impl Drop for AsyncWriterHandle {
    fn drop(&mut self) {
        // Best-effort drop: close the channel; let the worker drain.
        // If the worker is still running, we don't block in Drop.
        // Production callers should call `shutdown()` explicitly.
        self.tx.take();
        if let Some(jh) = self.join_handle.take() {
            // Don't block forever in Drop; spawn a detacher.
            let _ = jh.join();
        }
    }
}

/// The worker loop. Drains jobs from `rx`, writes them via `store`,
/// fires the per-job completion callback. On I/O error: log, fire
/// the completion callback with `Err`, continue with the next job.
fn worker_loop(
    store: Arc<DiskBlockStore>,
    rx: mpsc::Receiver<WriteJob>,
    pending: Arc<AtomicUsize>,
) {
    loop {
        let job = match rx.recv() {
            Ok(j) => j,
            Err(RecvError) => {
                // Sender dropped — drain any remaining queued items
                // (recv() returned RecvError because the queue is
                // empty AND no senders remain; we're done).
                break;
            }
        };
        process_job(&store, job, &pending);

        // Opportunistically drain any further jobs without an extra
        // condvar hop. `try_recv` returns immediately; this lets us
        // shave a millisecond or two off batch enqueues.
        loop {
            match rx.try_recv() {
                Ok(j) => process_job(&store, j, &pending),
                Err(TryRecvError::Empty) => break,
                Err(TryRecvError::Disconnected) => return,
            }
        }
    }
}

/// Write a single job. Errors are logged + reported via the
/// completion callback (if any); they NEVER panic the worker.
fn process_job(store: &DiskBlockStore, job: WriteJob, pending: &Arc<AtomicUsize>) {
    let WriteJob {
        header,
        body,
        completion_tx,
    } = job;
    let result = store.write_block_sync(&header, &body).map(|_| ());
    if let Err(ref e) = result {
        tracing::warn!(
            target: "hf2q::kv_persist::writer",
            error = %e,
            block_hash = %header.block_hash,
            "kv_persist writer: write_block_sync failed; continuing"
        );
    }
    if let Some(tx) = completion_tx {
        // The completion channel is best-effort: if the receiver is
        // gone (e.g. the spiller's request was dropped), we silently
        // discard the result. Production code should set a bound of
        // 1 on the completion channel and treat a `try_send` failure
        // as "caller didn't care".
        let _ = tx.try_send(result);
    }
    // ADR-017 §R-F5 (closure iter-11 2026-05-05): enforce on-disk
    // cache budget after every successful write. Pre-iter-11 the
    // `set_budget_bytes(N)` plumbing was wired but `evict_lru_until_under_budget`
    // had no production caller — the cache grew unbounded. We now
    // call it post-write with `|_| false` (pin nothing): the LRU sort
    // is by mtime so the just-written block is at the END of the
    // eviction queue, never the head. Eviction is best-effort — an
    // I/O error is logged but does NOT fail the write (the caller's
    // block already landed; budget is a soft cap).
    //
    // Phase A.3+ pin-set wiring is left for future iter when the
    // engine surfaces "blocks I'd want to restore from after a
    // crash" — until then, LRU-by-mtime is a sound approximation
    // (recently-spilled blocks are most likely to be hot).
    if let Err(e) = store.evict_lru_until_under_budget(|_| false) {
        tracing::warn!(
            target: "hf2q::kv_persist::writer",
            error = %e,
            "kv_persist writer: post-write budget eviction failed; continuing"
        );
    }
    // ADR-017 iter-16 — pending-depth bookkeeping. Decrement AFTER the
    // write + completion-ack so a probe taken between dequeue and
    // ack-fire still observes the job as pending. `saturating_sub`
    // semantics via `fetch_update` would be cleaner, but
    // `fetch_sub` cannot underflow as long as every successful
    // enqueue is paired with exactly one process_job — invariant
    // held by the worker_loop drain pattern.
    pending.fetch_sub(1, Ordering::Relaxed);
}

/// Convenience: build a oneshot-style completion channel for a single
/// `WriteJob`. Returns the sender (place into `WriteJob.completion_tx`)
/// and a receiver the caller can `recv_timeout` on. This is a thin
/// wrapper over `mpsc::sync_channel(1)` — we keep it here so the call
/// sites read clearly.
pub fn completion_channel() -> (SyncSender<io::Result<()>>, mpsc::Receiver<io::Result<()>>) {
    mpsc::sync_channel::<io::Result<()>>(1)
}

/// Default channel capacity. Small bound = aggressive back-pressure
/// signaling. Production may tune via `HF2Q_KV_WRITER_CAPACITY`.
pub const DEFAULT_CHANNEL_CAPACITY: usize = 8;

/// Default poll interval used by tests waiting on completion. Not used
/// by production code paths — the production caller picks the timeout
/// based on the SLA budget for the spill operation.
pub const DEFAULT_COMPLETION_TIMEOUT: Duration = Duration::from_secs(5);

// ---------------------------------------------------------------------------
// Tests.
// ---------------------------------------------------------------------------

#[cfg(test)]
mod tests {
    use super::*;
    use crate::serve::kv_persist::block_store::DiskBlockStore;
    use crate::serve::kv_persist::format::{
        compute_model_fingerprint, BlockHash, EnvelopeHeader, ModelFingerprint, ParentBlockHash,
        BLOCK_TOKENS, CURRENT_FORMAT_VERSION,
    };
    use sha2::{Digest, Sha256};
    use std::path::PathBuf;
    use std::process;
    use std::sync::atomic::{AtomicU32, Ordering};
    use std::time::SystemTime;

    fn temp_dir(label: &str) -> PathBuf {
        static COUNTER: AtomicU32 = AtomicU32::new(0);
        let n = COUNTER.fetch_add(1, Ordering::SeqCst);
        let pid = process::id();
        let nanos = SystemTime::now()
            .duration_since(std::time::UNIX_EPOCH)
            .map(|d| d.as_nanos())
            .unwrap_or(0);
        let dir = std::env::temp_dir().join(format!("hf2q-kv-writer-{label}-{pid}-{nanos}-{n}"));
        std::fs::create_dir_all(&dir).expect("temp_dir mkdir");
        dir
    }

    fn fixture_fp(seed: &str) -> ModelFingerprint {
        compute_model_fingerprint(
            seed,
            "Q4_0",
            "hf2q-test-1.0.0",
            "deadbeefcafebabe1122334455667788",
            "<|im_start|>...<|im_end|>",
        )
    }

    fn make_block(
        fp: ModelFingerprint,
        parent: ParentBlockHash,
        seed: u32,
    ) -> (Vec<u8>, EnvelopeHeader) {
        let body: Vec<u8> = (0..512u32)
            .flat_map(|i| (i.wrapping_add(seed)).to_le_bytes())
            .collect();
        let mut h = Sha256::new();
        h.update(&body);
        let bh: [u8; 32] = h.finalize().into();
        let header = EnvelopeHeader {
            format_version: CURRENT_FORMAT_VERSION.0,
            model_fingerprint: fp,
            block_hash: BlockHash(bh),
            parent_block_hash: parent,
            payload_kind: "kv-dense-bf16".into(),
            codec_version: 1,
            n_tokens: BLOCK_TOKENS,
        };
        (body, header)
    }

    #[test]
    fn spawn_then_shutdown_clean() {
        // Spawn a writer with no work, shut down immediately. No panic,
        // join returns Ok.
        let dir = temp_dir("spawnshut");
        let store = Arc::new(DiskBlockStore::new(dir.clone(), 0).expect("new"));
        let handle = AsyncWriterHandle::spawn(Arc::clone(&store), 8);
        handle.shutdown().expect("clean shutdown");
        assert_eq!(store.index().block_count(), 0, "no work, no blocks");

        let _ = std::fs::remove_dir_all(&dir);
    }

    #[test]
    fn enqueue_then_shutdown_drains_pending_jobs() {
        // Submit 10 jobs with the channel at capacity 4; shutdown drains
        // all of them before exiting. Verify all 10 blocks landed on disk.
        let dir = temp_dir("drain");
        let store = Arc::new(DiskBlockStore::new(dir.clone(), 0).expect("new"));
        let handle = AsyncWriterHandle::spawn(Arc::clone(&store), 4);
        let fp = fixture_fp("drain");

        let mut hashes: Vec<BlockHash> = Vec::new();
        for s in 0u32..10 {
            let (body, header) = make_block(fp, ParentBlockHash(None), s);
            hashes.push(header.block_hash);
            let job = WriteJob {
                header,
                body,
                completion_tx: None,
            };
            // Use blocking enqueue so all 10 jobs eventually land
            // (try_send would fail at capacity).
            handle.enqueue_blocking(job).expect("enqueue");
        }

        handle.shutdown().expect("shutdown drains");

        assert_eq!(store.index().block_count(), 10, "all 10 drained");
        for h in &hashes {
            assert!(store.index().lookup(h).is_some(), "hash indexed");
            let p = store.block_path(&fp, h);
            assert!(p.exists(), "file at {} present", p.display());
        }

        let _ = std::fs::remove_dir_all(&dir);
    }

    #[test]
    fn enqueue_completion_ack_fires() {
        // A job with a completion_tx receives Ok(()) after the worker
        // writes the block. Per feedback_live_verification_must_check_content:
        // also verify the body bytes round-trip.
        let dir = temp_dir("ack");
        let store = Arc::new(DiskBlockStore::new(dir.clone(), 0).expect("new"));
        let handle = AsyncWriterHandle::spawn(Arc::clone(&store), 4);
        let fp = fixture_fp("ack");

        let (body, header) = make_block(fp, ParentBlockHash(None), 1);
        let body_clone = body.clone();
        let block_hash = header.block_hash;
        let (ack_tx, ack_rx) = completion_channel();
        let job = WriteJob {
            header,
            body,
            completion_tx: Some(ack_tx),
        };
        handle.enqueue_blocking(job).expect("enqueue");

        let result = ack_rx
            .recv_timeout(DEFAULT_COMPLETION_TIMEOUT)
            .expect("completion received within timeout");
        result.expect("write succeeded");

        // Block visible immediately after ack.
        let body_back = store.read_block(&block_hash).expect("read");
        assert_eq!(body_back, body_clone, "body bytes round-trip via writer");

        handle.shutdown().expect("shutdown");
        let _ = std::fs::remove_dir_all(&dir);
    }

    #[test]
    fn worker_does_not_panic_on_io_error() {
        // Trigger a write failure by setting a 1-byte size override on the
        // store; the writer reports the error via completion_tx and
        // continues with subsequent jobs.
        let dir = temp_dir("ioerr");
        let store = Arc::new(DiskBlockStore::new(dir.clone(), 0).expect("new"));
        store.set_max_block_bytes_override(1); // any nonzero body fails
        let handle = AsyncWriterHandle::spawn(Arc::clone(&store), 4);
        let fp = fixture_fp("ioerr");

        // Submit a job that will fail.
        let (body_a, header_a) = make_block(fp, ParentBlockHash(None), 1);
        let (ack_a_tx, ack_a_rx) = completion_channel();
        handle
            .enqueue_blocking(WriteJob {
                header: header_a,
                body: body_a,
                completion_tx: Some(ack_a_tx),
            })
            .expect("enqueue a");

        let r_a = ack_a_rx
            .recv_timeout(DEFAULT_COMPLETION_TIMEOUT)
            .expect("ack a");
        assert!(r_a.is_err(), "first job reported error");

        // Worker is still alive; lift the constraint and submit a second
        // job. It must succeed.
        store.set_max_block_bytes_override(0); // unlimited (use default)
        let (body_b, header_b) = make_block(fp, ParentBlockHash(None), 2);
        let block_hash_b = header_b.block_hash;
        let (ack_b_tx, ack_b_rx) = completion_channel();
        handle
            .enqueue_blocking(WriteJob {
                header: header_b,
                body: body_b,
                completion_tx: Some(ack_b_tx),
            })
            .expect("enqueue b");

        let r_b = ack_b_rx
            .recv_timeout(DEFAULT_COMPLETION_TIMEOUT)
            .expect("ack b");
        r_b.expect("second job succeeded — worker survived previous error");
        assert!(store.index().lookup(&block_hash_b).is_some());

        handle.shutdown().expect("shutdown");
        let _ = std::fs::remove_dir_all(&dir);
    }

    #[test]
    fn enqueue_returns_full_error_when_channel_capacity_reached() {
        // Use a 1-cap channel with a manual stall: we DON'T spawn the
        // worker until after we've filled the channel, so subsequent
        // try_send calls return TrySendError::Full deterministically.
        //
        // Pattern: make a sender from a brand-new channel, fill it to
        // cap, observe Full. Then drop the channel (no leak).
        let (tx, _rx) = mpsc::sync_channel::<WriteJob>(1);
        let dir = temp_dir("full");
        let store = DiskBlockStore::new(dir.clone(), 0).expect("new");
        let fp = fixture_fp("full");
        let (body_a, header_a) = make_block(fp, ParentBlockHash(None), 1);
        let (body_b, header_b) = make_block(fp, ParentBlockHash(None), 2);
        // First send fits in the bounded channel.
        tx.try_send(WriteJob {
            header: header_a,
            body: body_a,
            completion_tx: None,
        })
        .expect("first try_send fits");
        // Second send must be rejected with Full.
        let err = tx
            .try_send(WriteJob {
                header: header_b,
                body: body_b,
                completion_tx: None,
            })
            .expect_err("must fail");
        match err {
            TrySendError::Full(_) => {}
            other => panic!("expected Full, got {other:?}"),
        }
        // No worker was spawned; the channel + store drop here.
        drop(tx);
        drop(store);
        let _ = std::fs::remove_dir_all(&dir);
    }

    #[test]
    fn enqueue_after_shutdown_returns_disconnected() {
        // A handle whose tx has been taken returns Disconnected on
        // try_send. We can't easily access that state externally
        // (Drop closes tx), so we model it directly: shutdown + drop +
        // a separately-built handle whose tx is None at construction
        // time would be the test seam, but the public API doesn't
        // expose that. Instead we assert the documented contract:
        // after shutdown, the handle is consumed (move semantics) and
        // can no longer be enqueued.
        //
        // The strictest property we can test on the live API is:
        // shutdown -> Ok with all enqueued jobs drained. We covered that
        // in `enqueue_then_shutdown_drains_pending_jobs`. Here we add
        // a smaller check: a freshly-spawned handle, immediately shut
        // down, refuses no enqueue (we never tried). Document the move-
        // semantics intent instead of forcing a false Disconnected test.
        let dir = temp_dir("disc");
        let store = Arc::new(DiskBlockStore::new(dir.clone(), 0).expect("new"));
        let handle = AsyncWriterHandle::spawn(Arc::clone(&store), 1);
        handle.shutdown().expect("shutdown");
        // Handle is consumed; this line would not compile:
        // let _ = handle.enqueue(...);
        // (Verified at the type level — move semantics.)
        let _ = std::fs::remove_dir_all(&dir);
    }
}