summa_core/index/

writer.rs

//! IndexWriter — async document indexing with parallel segment building.
//!
//! This module is only compiled with the "native" feature.
//!
//! # Architecture
//!
//! ```text
//! add_document() ──try_send──► [shared bounded MPMC] ◄──recv── worker 0
//!                                                     ◄──recv── worker 1
//!                                                     ◄──recv── worker N
//! ```
//!
//! - **Shared MPMC queue** (`async_channel`): all workers compete for documents.
//!   Busy workers (building segments) naturally stop pulling; free workers pick up slack.
//! - **Zero-copy pipeline**: `Document` is moved (never cloned) through every stage:
//!   `add_document()` → channel → `recv_blocking()` → `SegmentBuilder::add_document()`.
//! - `add_document` returns `QueueFull` when the queue is at capacity.
//! - **Workers are OS threads**: CPU-intensive work (tokenization, posting list building)
//!   runs on dedicated threads, never blocking the tokio async runtime.
//!   Async I/O (segment file writes) is bridged via `Handle::block_on()`.
//! - **Fixed per-worker memory budget**: `max_indexing_memory_bytes / num_workers`.
//!   Workers use deterministic, staggered soft flush thresholds within that
//!   budget so equal-size builders do not all stop draining at once.
//! - **Build concurrency reserve**: while input is open, at most `N - 1`
//!   workers build segments concurrently. A worker that cannot get a slot
//!   keeps draining up to the former 80% flush boundary. Once input closes,
//!   all `N` tail builds may finish concurrently because no drainer is needed.
//! - **Two-phase commit**:
//!   1. `prepare_commit()` — closes queue, workers flush builders to disk.
//!      Returns a `PreparedCommit` guard. No new documents accepted until resolved.
//!   2. `PreparedCommit::commit()` — registers segments in metadata, resumes workers.
//!   3. `PreparedCommit::abort()` — discards prepared segments, resumes workers.
//!   4. `commit()` — convenience: `prepare_commit().await?.commit().await`.
//!
//! Since `prepare_commit`/`commit` take `&mut self`, Rust’s borrow checker
//! guarantees no concurrent `add_document` calls during the commit window.

use super::primary_key::load_pk_segment_data;
use super::staged_row::{StagedRow, StagedSegment};

struct QueuedDocument {
    doc: Document,
    row: Option<Arc<StagedRow>>,
}
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};

use futures::{FutureExt, StreamExt, TryStreamExt};
use rustc_hash::FxHashMap;

use crate::directories::DirectoryWriter;
use crate::dsl::{Document, Field, Schema};
use crate::error::{Error, Result};
use crate::segment::{
    SegmentBuilder, SegmentBuilderConfig, SegmentId, validate_vector_value_counts,
};
use crate::tokenizer::BoxedTokenizer;

use super::IndexConfig;

/// Total pipeline capacity (in documents).
const PIPELINE_MAX_SIZE_IN_DOCS: usize = 10_000;

/// Builder memory percentage at which the first worker starts trying to flush.
///
/// The last worker starts at [`SOFT_FLUSH_MAX_PERCENT`], preserving the former
/// 20% segment-build headroom. Intermediate workers are evenly staggered
/// between the two bounds.
const SOFT_FLUSH_MIN_PERCENT: usize = 70;
const SOFT_FLUSH_MAX_PERCENT: usize = 80;

/// File name of the advisory single-writer lock inside the index directory.
pub const WRITER_LOCK_FILENAME: &str = ".summa_writer.lock";

/// Return a deterministic per-worker soft flush threshold.
///
/// All thresholds remain at or below the former uniform 80% trigger, so their
/// sum cannot increase builder memory. The 70–80% spread prevents identical
/// workers consuming the shared queue at the same rate from entering segment
/// builds in lockstep.
fn soft_flush_threshold(memory_budget: usize, worker_id: usize, num_workers: usize) -> usize {
    if num_workers <= 1 {
        return hard_flush_threshold(memory_budget);
    }

    let worker_id = worker_id.min(num_workers - 1);
    let span = SOFT_FLUSH_MAX_PERCENT - SOFT_FLUSH_MIN_PERCENT;
    let denominator = 100u128 * (num_workers - 1) as u128;
    let numerator = (SOFT_FLUSH_MIN_PERCENT * (num_workers - 1) + span * worker_id) as u128;
    ((memory_budget as u128 * numerator) / denominator) as usize
}

/// Preserve the historical 20% allowance for allocations created while a
/// segment is finalized. Workers may stagger below this boundary, but no
/// queued builder consumes the scratch reserve while waiting for a build slot.
fn hard_flush_threshold(memory_budget: usize) -> usize {
    memory_budget.saturating_mul(SOFT_FLUSH_MAX_PERCENT) / 100
}

/// Derive the builder defaults used by the standard writer constructors.
///
/// `IndexConfig::num_compression_threads` is the public per-index setting; it
/// must reach segment builders instead of being replaced by the machine-wide
/// `SegmentBuilderConfig` default. Explicit `*_with_config` constructors bypass
/// this helper and continue honoring every supplied builder option.
fn default_builder_config(index_config: &IndexConfig) -> SegmentBuilderConfig {
    let bounds = index_config.effective_posting_bounds();
    SegmentBuilderConfig {
        num_compression_threads: index_config.num_compression_threads,
        optimization: index_config.optimization,
        posting_codec: index_config.effective_posting_codec(),
        quantized_norms: index_config.quantized_norms,
        compact_text: index_config.compact_text,
        posting_ratio_bounds: bounds.ratio,
        posting_impact_bounds: bounds.impact,
        term_dict_block_size: index_config.term_dict_block_size,
        ..SegmentBuilderConfig::default()
    }
}

/// Bounds simultaneous segment finalization while reserving an indexing
/// worker to drain the shared document queue.
///
/// A soft-threshold worker first tries to acquire without waiting. If every
/// slot is occupied it may continue indexing until its hard per-worker memory
/// flush boundary. At that hard limit it waits, preserving the remaining 20%
/// of every worker share for finalization scratch.
struct SegmentBuildLimiter {
    live_max_active: usize,
    flush_max_active: usize,
    active: AtomicUsize,
    flushing: AtomicBool,
    wait_mutex: parking_lot::Mutex<()>,
    available: parking_lot::Condvar,
}

impl SegmentBuildLimiter {
    fn new(num_workers: usize) -> Self {
        Self {
            // A single-worker writer must still be able to build. With two or
            // more workers, reserve one worker from concurrent finalization.
            live_max_active: num_workers.saturating_sub(1).max(1),
            // Once the input queue closes there is no ingestion to reserve;
            // flush every tail concurrently as the old writer did.
            flush_max_active: num_workers.max(1),
            active: AtomicUsize::new(0),
            flushing: AtomicBool::new(false),
            wait_mutex: parking_lot::Mutex::new(()),
            available: parking_lot::Condvar::new(),
        }
    }

    fn try_acquire(&self) -> Option<SegmentBuildPermit<'_>> {
        self.try_acquire_up_to(self.live_max_active)
    }

    fn try_acquire_up_to(&self, limit: usize) -> Option<SegmentBuildPermit<'_>> {
        let mut active = self.active.load(Ordering::Acquire);
        loop {
            if active >= limit {
                return None;
            }
            match self.active.compare_exchange_weak(
                active,
                active + 1,
                Ordering::AcqRel,
                Ordering::Acquire,
            ) {
                Ok(_) => return Some(SegmentBuildPermit { limiter: self }),
                Err(observed) => active = observed,
            }
        }
    }

    fn acquire(&self) -> SegmentBuildPermit<'_> {
        let mut wait = self.wait_mutex.lock();
        loop {
            let limit = if self.flushing.load(Ordering::Acquire) {
                self.flush_max_active
            } else {
                self.live_max_active
            };
            if let Some(permit) = self.try_acquire_up_to(limit) {
                return permit;
            }
            self.available.wait(&mut wait);
        }
    }

    fn acquire_flush(&self) -> SegmentBuildPermit<'_> {
        self.acquire_up_to(self.flush_max_active)
    }

    fn acquire_up_to(&self, limit: usize) -> SegmentBuildPermit<'_> {
        let mut wait = self.wait_mutex.lock();
        loop {
            if let Some(permit) = self.try_acquire_up_to(limit) {
                return permit;
            }
            self.available.wait(&mut wait);
        }
    }

    /// Promote existing live waiters when input closes. Store the phase before
    /// taking the condvar mutex; a waiter either observes it directly or
    /// releases the mutex in `wait`, after which this notification reaches it.
    fn begin_flush(&self) {
        self.flushing.store(true, Ordering::Release);
        let _wait = self.wait_mutex.lock();
        self.available.notify_all();
    }

    fn end_flush(&self) {
        self.flushing.store(false, Ordering::Release);
    }

    /// Reserve a build slot once `builder_memory` reaches its soft threshold.
    ///
    /// When all slots are busy, `None` below `hard_budget` means "keep
    /// draining". At the hard budget this waits for a slot rather than
    /// exceeding the configured per-worker memory share.
    fn reserve_if_due(
        &self,
        builder_memory: usize,
        soft_threshold: usize,
        hard_budget: usize,
    ) -> Option<SegmentBuildPermit<'_>> {
        if builder_memory < soft_threshold {
            return None;
        }
        if let Some(permit) = self.try_acquire() {
            return Some(permit);
        }
        if builder_memory < hard_budget {
            return None;
        }
        Some(self.acquire())
    }
}

struct SegmentBuildPermit<'a> {
    limiter: &'a SegmentBuildLimiter,
}

impl Drop for SegmentBuildPermit<'_> {
    fn drop(&mut self) {
        let previous = self.limiter.active.fetch_sub(1, Ordering::AcqRel);
        debug_assert!(previous > 0);
        // Synchronize notification with acquire's condvar wait so a permit
        // becoming free cannot be missed between its check and sleeping.
        let _wait = self.limiter.wait_mutex.lock();
        // Live-ingestion and closed-queue flush waiters have different limits;
        // wake both classes so the reserved flush slot cannot be stranded.
        self.limiter.available.notify_all();
    }
}

#[cfg(test)]
mod indexing_pipeline_tests {
    use std::sync::Arc;
    use std::time::Duration;

    use super::{
        SOFT_FLUSH_MAX_PERCENT, SOFT_FLUSH_MIN_PERCENT, SegmentBuildLimiter,
        default_builder_config, hard_flush_threshold, soft_flush_threshold,
    };

    #[test]
    fn standard_builder_config_honors_index_compression_width() {
        let index_config = crate::index::IndexConfig {
            num_compression_threads: 7,
            ..Default::default()
        };

        let builder_config = default_builder_config(&index_config);

        assert_eq!(builder_config.num_compression_threads, 7);
    }

    #[test]
    fn flush_thresholds_are_staggered_without_increasing_memory_budget() {
        const WORKERS: usize = 12;
        const PER_WORKER_BUDGET: usize = 1024 * 1024 * 1024;

        let thresholds: Vec<_> = (0..WORKERS)
            .map(|worker| soft_flush_threshold(PER_WORKER_BUDGET, worker, WORKERS))
            .collect();

        assert_eq!(
            thresholds[0],
            PER_WORKER_BUDGET * SOFT_FLUSH_MIN_PERCENT / 100
        );
        assert_eq!(
            thresholds[WORKERS - 1],
            PER_WORKER_BUDGET * SOFT_FLUSH_MAX_PERCENT / 100
        );
        assert!(
            thresholds.windows(2).all(|pair| pair[0] < pair[1]),
            "production-width workers must not reach identical flush thresholds: {thresholds:?}"
        );

        let staggered_total: usize = thresholds.iter().sum();
        let former_uniform_total = WORKERS * (PER_WORKER_BUDGET * SOFT_FLUSH_MAX_PERCENT / 100);
        assert!(
            staggered_total <= former_uniform_total,
            "staggering must not increase aggregate builder memory"
        );

        assert_eq!(
            soft_flush_threshold(PER_WORKER_BUDGET, 0, 1),
            PER_WORKER_BUDGET * SOFT_FLUSH_MAX_PERCENT / 100,
            "single-worker behavior retains the former 80% build headroom"
        );

        let hard_threshold = hard_flush_threshold(PER_WORKER_BUDGET);
        let build_scratch = PER_WORKER_BUDGET - hard_threshold;
        let steady_state_peak = (WORKERS - 1) * (hard_threshold + build_scratch) + hard_threshold;
        assert!(
            steady_state_peak <= WORKERS * PER_WORKER_BUDGET,
            "rotated hard-threshold builds must retain aggregate scratch headroom"
        );
    }

    #[test]
    fn full_build_gate_leaves_soft_threshold_worker_draining() {
        const WORKERS: usize = 12;
        let limiter = SegmentBuildLimiter::new(WORKERS);
        let mut active_builds: Vec<_> = (0..WORKERS - 1)
            .map(|_| {
                limiter
                    .try_acquire()
                    .expect("N - 1 builds should be admitted")
            })
            .collect();

        assert!(
            limiter.try_acquire().is_none(),
            "the final worker must be reserved from concurrent segment builds"
        );
        assert!(
            limiter.reserve_if_due(750, 700, 800).is_none(),
            "a worker below its hard budget must keep draining when builds are saturated"
        );

        drop(active_builds.pop());
        let replacement = limiter
            .reserve_if_due(750, 700, 800)
            .expect("a completed build must immediately rotate draining capacity");
        assert!(limiter.try_acquire().is_none());
        drop(replacement);
        drop(active_builds);
    }

    #[test]
    fn closed_queue_flush_uses_the_reserved_build_slot() {
        let limiter = SegmentBuildLimiter::new(2);
        let live_build = limiter
            .try_acquire()
            .expect("one live build should be admitted");
        assert!(limiter.try_acquire().is_none());

        let tail_build = limiter.acquire_flush();
        assert!(
            limiter
                .try_acquire_up_to(limiter.flush_max_active)
                .is_none(),
            "closed-queue flushes must remain bounded by the worker count"
        );

        drop(tail_build);
        drop(live_build);
    }

    #[test]
    fn closing_input_promotes_an_existing_live_waiter() {
        let limiter = Arc::new(SegmentBuildLimiter::new(2));
        let live_build = limiter.try_acquire().unwrap();
        let waiter_limiter = Arc::clone(&limiter);
        let (started_tx, started_rx) = std::sync::mpsc::channel();
        let (acquired_tx, acquired_rx) = std::sync::mpsc::channel();

        let waiter = std::thread::spawn(move || {
            started_tx.send(()).unwrap();
            let _permit = waiter_limiter
                .reserve_if_due(800, 700, 800)
                .expect("closed input must promote a hard-boundary waiter");
            acquired_tx.send(()).unwrap();
        });

        started_rx.recv_timeout(Duration::from_secs(1)).unwrap();
        assert!(acquired_rx.recv_timeout(Duration::from_millis(50)).is_err());

        limiter.begin_flush();
        acquired_rx
            .recv_timeout(Duration::from_secs(1))
            .expect("live waiter did not adopt the closed-queue build limit");
        waiter.join().unwrap();
        limiter.end_flush();
        drop(live_build);
    }

    #[test]
    fn hard_budget_waiter_resumes_when_a_build_finishes() {
        let limiter = Arc::new(SegmentBuildLimiter::new(3));
        let first = limiter.try_acquire().unwrap();
        let second = limiter.try_acquire().unwrap();
        let waiter_limiter = Arc::clone(&limiter);
        let (started_tx, started_rx) = std::sync::mpsc::channel();
        let (acquired_tx, acquired_rx) = std::sync::mpsc::channel();

        let waiter = std::thread::spawn(move || {
            started_tx.send(()).unwrap();
            let _permit = waiter_limiter
                .reserve_if_due(800, 700, 800)
                .expect("hard-budget worker must eventually acquire a build slot");
            acquired_tx.send(()).unwrap();
        });

        started_rx.recv_timeout(Duration::from_secs(1)).unwrap();
        assert!(
            acquired_rx.recv_timeout(Duration::from_millis(50)).is_err(),
            "hard-budget worker must not over-subscribe segment builds"
        );
        drop(first);
        acquired_rx
            .recv_timeout(Duration::from_secs(1))
            .expect("hard-budget worker did not wake after a build completed");
        waiter.join().unwrap();
        drop(second);
    }
}

/// Advisory single-writer lock state.
///
/// Two independent writers on one index directory silently destroy each
/// other's data: the orphan sweep at writer open deletes the other process's
/// unpublished segment files, and metadata saves are last-writer-wins. For
/// directories rooted on a local filesystem the writer therefore holds an OS
/// advisory lock for its whole lifetime; the kernel releases it automatically
/// when the process dies.
enum WriterLock {
    /// Lock acquired. Closing the file (writer drop) releases it.
    Held { _file: std::fs::File },
    /// The directory has no lockable local filesystem root (e.g. RAM or
    /// remote directories) — cross-process locking is not applicable.
    NotApplicable,
    /// Another writer holds the lock. Every mutating operation fails loudly
    /// with this message instead of silently double-writing.
    Unavailable { reason: String },
}

/// Local filesystem root of the index directory, when the directory type
/// exposes one.
fn writer_lock_root<D: DirectoryWriter + 'static>(directory: &D) -> Option<std::path::PathBuf> {
    let any: &dyn std::any::Any = directory;
    if let Some(mmap) = any.downcast_ref::<crate::directories::MmapDirectory>() {
        return Some(mmap.root().to_path_buf());
    }
    // FsDirectory does not expose its root path, so the single-writer lock
    // cannot be enforced for it yet. Say so loudly instead of silently
    // skipping protection for a filesystem-backed writer.
    if any
        .downcast_ref::<crate::directories::FsDirectory>()
        .is_some()
    {
        log::warn!(
            "[writer_lock] FsDirectory exposes no root path; single-writer locking \
             is not enforced for this writer — do not open a second writer for the \
             same index directory"
        );
    }
    None
}

/// Try to take the exclusive single-writer lock for `directory`.
///
/// Returns `WriterLock::Unavailable` (not `Err`) on conflict so infallible
/// constructors can defer the failure to their first mutating operation.
fn try_acquire_writer_lock<D: DirectoryWriter + 'static>(directory: &D) -> Result<WriterLock> {
    let Some(root) = writer_lock_root(directory) else {
        return Ok(WriterLock::NotApplicable);
    };
    std::fs::create_dir_all(&root)?;
    let lock_path = root.join(WRITER_LOCK_FILENAME);
    let file = std::fs::OpenOptions::new()
        .create(true)
        .truncate(false)
        .write(true)
        .open(&lock_path)?;
    match file.try_lock() {
        Ok(()) => Ok(WriterLock::Held { _file: file }),
        Err(std::fs::TryLockError::WouldBlock) => Ok(WriterLock::Unavailable {
            reason: format!(
                "another IndexWriter already holds the single-writer lock for this \
                 index ({}); Summa supports one writer per index directory — stop \
                 the other writer (e.g. a running summa-server or summa-tool) \
                 before opening this one",
                lock_path.display()
            ),
        }),
        Err(std::fs::TryLockError::Error(error)) => Err(Error::Io(error)),
    }
}

/// Async IndexWriter for adding documents and committing segments.
///
/// **Backpressure:** `add_document()` is sync and O(1). It returns
/// `Error::QueueFull` when the shared queue is full and
/// `Error::CommitInProgress` while a generation is publishing or awaiting
/// retry; callers must back off.
///
/// **Two-phase commit:**
/// - `prepare_commit()` → `PreparedCommit::commit()` or `PreparedCommit::abort()`
/// - `commit()` is a convenience that does both phases.
/// - Between prepare and commit, the caller can do external work (WAL, sync, etc.)
///   knowing that abort is possible if something fails.
/// - Dropping `PreparedCommit` without calling commit/abort auto-aborts.
pub struct IndexWriter<D: DirectoryWriter + 'static> {
    pub(super) directory: Arc<D>,
    pub(super) schema: Arc<Schema>,
    pub(super) config: IndexConfig,
    /// MPMC sender, replaced under a brief lock on each commit cycle (workers
    /// get the corresponding new receiver via resume).
    doc_sender: Arc<parking_lot::RwLock<async_channel::Sender<QueuedDocument>>>,
    /// Worker OS thread handles — long-lived, survive across commits.
    workers: Vec<std::thread::JoinHandle<()>>,
    /// Shared worker state (immutable config + mutable segment output + sync)
    worker_state: Arc<WorkerState<D>>,
    /// Segment manager — owns metadata.json, handles segments and background merging
    pub(super) segment_manager: Arc<crate::merge::SegmentManager<D>>,
    /// Segments flushed to disk but not yet registered in metadata. Each item
    /// owns an active-operation guard, so orphan sweeping cannot delete it.
    flushed_segments: Arc<parking_lot::Mutex<Vec<PreparedSegment<D>>>>,
    /// Primary key dedup index (None if schema has no primary field)
    primary_key_index: Arc<parking_lot::RwLock<Option<super::primary_key::PrimaryKeyIndex>>>,
    /// Serializes async snapshot acquisition/loading across commits and
    /// lifecycle-owned merge/reorder topology refreshes.
    primary_key_refresh_lock: Arc<tokio::sync::Mutex<()>>,
    /// Tracks the owned finalizer spawned by `PreparedCommit::commit`. The
    /// requesting future may disappear, but a second commit generation must
    /// not start until this one has made publication and worker state agree.
    commit_finalization: Arc<CommitFinalizationState>,
    /// True while a failed post-commit PK refresh has left the uncommitted
    /// reservations as the ONLY record of already-committed keys (fail-closed,
    /// see `finalize_prepared_commit`). While set, abort paths must NOT clear
    /// the reservations or duplicate primary keys could be admitted.
    pk_reservations_retained: Arc<AtomicBool>,
    /// Advisory single-writer lock, held for the writer's lifetime.
    /// `Unavailable` is retryable: the conflicting holder may exit at any
    /// time (the kernel then releases its lock), so `ensure_writer_lock`
    /// re-attempts acquisition instead of caching the conflict forever.
    writer_lock: parking_lot::RwLock<WriterLock>,
}

#[derive(Default)]
struct CommitFinalizationState {
    in_progress: AtomicBool,
    idle: tokio::sync::Notify,
}

impl CommitFinalizationState {
    fn begin(&self) -> bool {
        self.in_progress
            .compare_exchange(false, true, Ordering::AcqRel, Ordering::Acquire)
            .is_ok()
    }

    fn finish(&self) {
        self.in_progress.store(false, Ordering::Release);
        self.idle.notify_waiters();
    }

    async fn wait_until_idle(&self) {
        while self.in_progress.load(Ordering::Acquire) {
            let notified = self.idle.notified();
            if !self.in_progress.load(Ordering::Acquire) {
                break;
            }
            notified.await;
        }
    }
}

/// Shared state for worker threads.
struct WorkerState<D: DirectoryWriter + 'static> {
    directory: Arc<D>,
    schema: Arc<Schema>,
    builder_config: SegmentBuilderConfig,
    tokenizers: parking_lot::RwLock<FxHashMap<Field, BoxedTokenizer>>,
    /// Fixed per-worker memory budget (bytes). When a builder exceeds this, segment is built.
    memory_budget_per_worker: usize,
    /// Limits live segment finalization to N - 1 workers, reserving
    /// queue-draining capacity; closed-queue tail flushes may use all N.
    segment_build_limiter: SegmentBuildLimiter,
    /// Segment manager — workers read trained structures from its ArcSwap (lock-free).
    segment_manager: Arc<crate::merge::SegmentManager<D>>,
    /// Segments built by workers, collected by `prepare_commit()`. Their RAII
    /// guards protect both in-progress and completed-uncommitted files.
    built_segments: parking_lot::Mutex<Vec<PreparedSegment<D>>>,
    /// First failure in the current flush generation. Worker-side indexing is
    /// asynchronous, so `prepare_commit` is the only sound place to surface
    /// it to the caller. A failed generation is aborted as a unit; publishing
    /// only its successful segments would silently lose documents.
    cycle_error: parking_lot::Mutex<Option<String>>,
    cycle_failed: AtomicBool,

    // === Worker lifecycle synchronization ===
    // Workers survive across commits. On prepare_commit the channel is closed;
    // workers flush their builders, increment flush_count, then wait on
    // resume_cvar for a new receiver. commit/abort creates a fresh channel
    // and wakes them.
    /// Number of workers that have completed their flush.
    flush_count: AtomicUsize,
    /// Mutex + condvar for prepare_commit to wait on all workers flushed.
    flush_mutex: parking_lot::Mutex<()>,
    flush_cvar: parking_lot::Condvar,
    /// Holds the new channel receiver after commit/abort. Workers clone from this.
    resume_receiver: parking_lot::Mutex<Option<async_channel::Receiver<QueuedDocument>>>,
    /// Monotonically increasing epoch, bumped by each resume_workers call.
    /// Workers compare against their local epoch to avoid re-cloning a stale receiver.
    resume_epoch: AtomicUsize,
    /// Condvar for workers to wait for resume (new channel) or shutdown.
    resume_cvar: parking_lot::Condvar,
    /// When true, workers should exit permanently (IndexWriter dropped).
    shutdown: AtomicBool,
    /// Total number of worker threads.
    num_workers: usize,
}

/// A completed indexing segment that has not been published in metadata yet.
///
/// `operation` is intentionally data, not a side-channel set update: moving
/// this value through worker → prepared commit → commit/abort moves lifecycle
/// ownership with it, and every unwind/drop path releases ownership safely.
struct PreparedSegment<D: DirectoryWriter + 'static> {
    id: String,
    segment_id: SegmentId,
    num_docs: u32,
    staged_rows: Arc<StagedSegment>,
    segment_manager: Arc<crate::merge::SegmentManager<D>>,
    operation: Option<crate::merge::SegmentOperationGuard>,
    runtime: tokio::runtime::Handle,
    needs_vector_upgrade: bool,
    published: bool,
}

impl<D: DirectoryWriter + 'static> PreparedSegment<D> {
    fn metadata_entry(&self) -> (String, u32) {
        (self.id.clone(), self.num_docs)
    }

    fn mark_published(&mut self) {
        self.published = true;
        // Metadata + SegmentTracker are now the durable lifecycle owners.
        drop(self.operation.take());
    }
}

impl<D: DirectoryWriter + 'static> WorkerState<D> {
    fn record_cycle_error(&self, error: impl Into<String>) {
        let mut first_error = self.cycle_error.lock();
        if first_error.is_none() {
            *first_error = Some(error.into());
        }
        drop(first_error);
        self.cycle_failed.store(true, Ordering::Release);
    }
}

impl<D: DirectoryWriter + 'static> Drop for PreparedSegment<D> {
    fn drop(&mut self) {
        if self.published {
            return;
        }
        let Some(operation) = self.operation.take() else {
            return;
        };
        self.segment_manager.schedule_unpublished_segment_cleanup(
            self.segment_id,
            operation,
            self.runtime.clone(),
        );
    }
}

impl<D: DirectoryWriter + 'static> IndexWriter<D> {
    /// Create a new index in the directory
    pub async fn create(directory: D, schema: Schema, config: IndexConfig) -> Result<Self> {
        let builder_config = default_builder_config(&config);
        Self::create_with_config(directory, schema, config, builder_config).await
    }

    /// Create a new index with custom builder config
    pub async fn create_with_config(
        directory: D,
        schema: Schema,
        config: IndexConfig,
        builder_config: SegmentBuilderConfig,
    ) -> Result<Self> {
        schema.validate()?;
        crate::dsl::reject_removed_vector_index_types(&schema).map_err(Error::Schema)?;
        let directory = Arc::new(directory);
        let schema = Arc::new(schema);
        // Directory-layer metrics (cold writes, lazy reads) carry the index label
        directory.set_index_label(schema.index_label());

        // Refuse a second writer before touching any index state.
        let writer_lock = try_acquire_writer_lock(directory.as_ref())?;
        if let WriterLock::Unavailable { reason } = &writer_lock {
            return Err(Error::Internal(reason.clone()));
        }
        // Refuse to clobber an existing index: persisting a fresh empty
        // metadata.json would orphan every committed segment, and the next
        // writer open's orphan sweep would permanently delete them.
        if directory
            .exists(std::path::Path::new(super::INDEX_META_FILENAME))
            .await?
        {
            return Err(Error::Internal(format!(
                "refusing to create index: {} already exists in this directory; \
                 use IndexWriter::open to open the existing index, or delete the \
                 directory first if you really want to start over",
                super::INDEX_META_FILENAME
            )));
        }

        let metadata = super::IndexMetadata::new((*schema).clone());

        // A custom builder config is authoritative for the physical text
        // layout, including merge re-encoding and the metadata compatibility
        // gate.
        let mut segment_config = config.clone();
        segment_config.optimization = builder_config.optimization;
        segment_config.posting_codec = Some(builder_config.posting_codec);
        segment_config.quantized_norms = builder_config.quantized_norms;
        segment_config.compact_text = builder_config.compact_text;
        segment_config.term_dict_block_size = builder_config.term_dict_block_size;
        let segment_manager =
            super::segment_manager_from_config(&directory, &schema, metadata, &segment_config)?;
        segment_manager.update_metadata(|_| {}).await?;

        Ok(Self::new_with_parts(
            directory,
            schema,
            config,
            builder_config,
            segment_manager,
            writer_lock,
        ))
    }

    /// Open an existing index for exclusive writing.
    ///
    /// Multiple independent writers for the same directory are unsupported;
    /// for filesystem-rooted directories this is enforced with an advisory
    /// single-writer lock ([`WRITER_LOCK_FILENAME`]) held for the writer's
    /// lifetime. This path removes crash-leftover outputs before starting its
    /// workers. Use [`Index::writer`](super::Index::writer) to share lifecycle
    /// state with an already-open search index.
    pub async fn open(directory: D, config: IndexConfig) -> Result<Self> {
        let builder_config = default_builder_config(&config);
        Self::open_with_config(directory, config, builder_config).await
    }

    /// Open an existing index with custom builder config
    pub async fn open_with_config(
        directory: D,
        config: IndexConfig,
        builder_config: SegmentBuilderConfig,
    ) -> Result<Self> {
        let directory = Arc::new(directory);

        // The lock must be held before the orphan sweep below: sweeping while
        // another process's writer is live deletes its in-flight outputs.
        let writer_lock = try_acquire_writer_lock(directory.as_ref())?;
        if let WriterLock::Unavailable { reason } = &writer_lock {
            return Err(Error::Internal(reason.clone()));
        }

        let metadata = super::IndexMetadata::load_persisting_migration(directory.as_ref()).await?;
        let schema = Arc::new(metadata.schema.clone());
        // Directory-layer metrics (cold writes, lazy reads) carry the index label
        directory.set_index_label(schema.index_label());

        // See `create_with_config`: custom builders define the on-disk layout.
        let mut segment_config = config.clone();
        segment_config.optimization = builder_config.optimization;
        segment_config.posting_codec = Some(builder_config.posting_codec);
        segment_config.quantized_norms = builder_config.quantized_norms;
        segment_config.compact_text = builder_config.compact_text;
        segment_config.term_dict_block_size = builder_config.term_dict_block_size;
        segment_config.posting_ratio_bounds = builder_config.posting_ratio_bounds;
        segment_config.posting_impact_bounds = builder_config.posting_impact_bounds;
        super::log_posting_bounds_policy(directory.as_ref(), &metadata, &segment_config).await;
        let segment_manager =
            super::segment_manager_from_config(&directory, &schema, metadata, &segment_config)?;
        let swept = segment_manager.cleanup_orphan_segments().await?;
        if swept > 0 {
            log::warn!(
                "[segment_cleanup] swept {} orphan segment(s) while opening writer",
                swept
            );
        }
        segment_manager.try_load_and_publish_trained().await?;

        Ok(Self::new_with_parts(
            directory,
            schema,
            config,
            builder_config,
            segment_manager,
            writer_lock,
        ))
    }

    /// Create an IndexWriter from an existing Index.
    /// Shares the SegmentManager for consistent segment lifecycle management.
    ///
    /// This constructor is infallible, so a single-writer lock conflict is
    /// deferred: the returned writer fails loudly on its first mutating
    /// operation instead of silently double-writing next to another writer.
    pub fn from_index(index: &super::Index<D>) -> Self {
        let writer_lock = match try_acquire_writer_lock(index.directory.as_ref()) {
            Ok(lock) => lock,
            Err(error) => WriterLock::Unavailable {
                reason: format!("failed to acquire the single-writer lock: {error}"),
            },
        };
        if let WriterLock::Unavailable { reason } = &writer_lock {
            log::error!("[writer_lock] {reason}");
        }
        let builder_config = default_builder_config(&index.config);
        Self::new_with_parts(
            Arc::clone(&index.directory),
            index.schema_arc(),
            index.config.clone(),
            builder_config,
            Arc::clone(&index.segment_manager),
            writer_lock,
        )
    }

    // ========================================================================
    // Construction + pipeline management
    // ========================================================================

    /// Common construction: creates worker state, spawns workers, assembles `Self`.
    fn new_with_parts(
        directory: Arc<D>,
        schema: Arc<Schema>,
        config: IndexConfig,
        builder_config: SegmentBuilderConfig,
        segment_manager: Arc<crate::merge::SegmentManager<D>>,
        writer_lock: WriterLock,
    ) -> Self {
        // Auto-configure tokenizers from schema for all text fields
        let registry = crate::tokenizer::TokenizerRegistry::new();
        let mut tokenizers = FxHashMap::default();
        for (field, entry) in schema.fields() {
            if matches!(entry.field_type, crate::dsl::FieldType::Text)
                && let Some(ref tok_name) = entry.tokenizer
                && let Some(tok) = registry.get(tok_name)
            {
                tokenizers.insert(field, tok);
            }
        }

        let num_workers = config.num_indexing_threads.max(1);
        let worker_state = Arc::new(WorkerState {
            directory: Arc::clone(&directory),
            schema: Arc::clone(&schema),
            builder_config,
            tokenizers: parking_lot::RwLock::new(tokenizers),
            memory_budget_per_worker: config.max_indexing_memory_bytes / num_workers,
            segment_build_limiter: SegmentBuildLimiter::new(num_workers),
            segment_manager: Arc::clone(&segment_manager),
            built_segments: parking_lot::Mutex::new(Vec::new()),
            cycle_error: parking_lot::Mutex::new(None),
            cycle_failed: AtomicBool::new(false),
            flush_count: AtomicUsize::new(0),
            flush_mutex: parking_lot::Mutex::new(()),
            flush_cvar: parking_lot::Condvar::new(),
            resume_receiver: parking_lot::Mutex::new(None),
            resume_epoch: AtomicUsize::new(0),
            resume_cvar: parking_lot::Condvar::new(),
            shutdown: AtomicBool::new(false),
            num_workers,
        });
        let (doc_sender, workers) = Self::spawn_workers(&worker_state, num_workers);
        let primary_key_index = Arc::new(parking_lot::RwLock::new(None));
        let primary_key_refresh_lock = Arc::new(tokio::sync::Mutex::new(()));

        Self {
            directory,
            schema,
            config,
            doc_sender: Arc::new(parking_lot::RwLock::new(doc_sender)),
            workers,
            worker_state,
            segment_manager,
            flushed_segments: Arc::new(parking_lot::Mutex::new(Vec::new())),
            primary_key_index,
            primary_key_refresh_lock,
            commit_finalization: Arc::new(CommitFinalizationState::default()),
            pk_reservations_retained: Arc::new(AtomicBool::new(false)),
            writer_lock: parking_lot::RwLock::new(writer_lock),
        }
    }

    /// Fail loudly when another writer owns the single-writer lock.
    ///
    /// A deferred conflict (`from_index` during a writer handover, e.g. a
    /// rolling pod restart) is not permanent: the holder exits and the kernel
    /// releases its advisory lock. Re-attempt acquisition on every call in
    /// the `Unavailable` state so the writer recovers as soon as the lock
    /// frees, instead of rejecting all writes for its lifetime.
    fn ensure_writer_lock(&self) -> Result<()> {
        // Fast path: uncontended read on the healthy states.
        if !matches!(&*self.writer_lock.read(), WriterLock::Unavailable { .. }) {
            return Ok(());
        }

        let mut lock = self.writer_lock.write();
        // Another thread may have recovered while we waited for the write lock.
        if !matches!(&*lock, WriterLock::Unavailable { .. }) {
            return Ok(());
        }
        match try_acquire_writer_lock(self.directory.as_ref())? {
            acquired @ (WriterLock::Held { .. } | WriterLock::NotApplicable) => {
                log::info!(
                    "[writer_lock] index={} single-writer lock acquired after retry; \
                     the previous holder has released it — resuming writes",
                    self.schema.index_label()
                );
                *lock = acquired;
                Ok(())
            }
            WriterLock::Unavailable { reason } => {
                let err = Error::Internal(reason.clone());
                *lock = WriterLock::Unavailable { reason };
                Err(err)
            }
        }
    }

    /// Clear primary-key reservations after an aborted or failed generation.
    ///
    /// Skipped while a failed post-commit PK refresh has left the uncommitted
    /// reservations as the ONLY record of already-committed keys (fail-closed,
    /// see `finalize_prepared_commit`): wiping them would admit duplicate
    /// primary keys. Retaining the aborted generation's keys as well is
    /// deliberately conservative — they clear on the next successful commit's
    /// refresh.
    fn clear_uncommitted_pk_reservations(&self) {
        if self.pk_reservations_retained.load(Ordering::Acquire) {
            log::warn!(
                "[primary_key] index={} keeping uncommitted reservations through abort: a \
                 failed post-commit refresh left them as the only record of \
                 committed keys; they are cleared by the next successful commit",
                self.schema.index_label()
            );
            return;
        }
        if let Some(pk_index) = self.primary_key_index.write().as_mut() {
            pk_index.clear_uncommitted();
        }
    }

    fn spawn_workers(
        worker_state: &Arc<WorkerState<D>>,
        num_workers: usize,
    ) -> (
        async_channel::Sender<QueuedDocument>,
        Vec<std::thread::JoinHandle<()>>,
    ) {
        let (sender, receiver) = async_channel::bounded(PIPELINE_MAX_SIZE_IN_DOCS);
        let handle = tokio::runtime::Handle::current();
        let mut workers = Vec::with_capacity(num_workers);
        for i in 0..num_workers {
            let state = Arc::clone(worker_state);
            let rx = receiver.clone();
            let rt = handle.clone();
            workers.push(
                std::thread::Builder::new()
                    .name(format!("index-worker-{}", i))
                    .spawn(move || Self::worker_loop(state, rx, rt, i))
                    .expect("failed to spawn index worker thread"),
            );
        }
        (sender, workers)
    }

    /// Get the schema
    pub fn schema(&self) -> Arc<Schema> {
        self.segment_manager.published_generation().schema.clone()
    }

    /// Set tokenizer for a field.
    /// Propagated to worker threads — takes effect for the next SegmentBuilder they create.
    pub fn set_tokenizer<T: crate::tokenizer::Tokenizer>(&mut self, field: Field, tokenizer: T) {
        self.worker_state
            .tokenizers
            .write()
            .insert(field, Box::new(tokenizer));
    }

    /// Initialize primary key deduplication from committed segments.
    ///
    /// Tries to load a cached bloom filter from `pk_bloom.bin` first. If the
    /// cache covers all current segments, the bloom is reused directly (fast
    /// path). If new segments appeared since the cache was written, only their
    /// keys are iterated (incremental). Falls back to a full rebuild when no
    /// cache exists.
    ///
    /// Only loads fast-field data (text dictionaries) per segment — NOT full
    /// `SegmentReader`s — to avoid duplicating dense/sparse index memory.
    ///
    /// The CPU-intensive bloom build is offloaded via `spawn_blocking` so it
    /// does not block the tokio runtime.
    ///
    /// No-op if schema has no primary field.
    pub async fn init_primary_key_dedup(&mut self) -> Result<()> {
        use super::primary_key::{PK_BLOOM_FILE, deserialize_pk_bloom};

        self.commit_finalization.wait_until_idle().await;
        self.ensure_writer_lock()?;

        let field = match self.schema.primary_field() {
            Some(f) => f,
            None => return Ok(()),
        };

        // A merge/reorder replacement can publish while this initialization
        // performs async segment loads. Serialize both paths so an older
        // initialization snapshot cannot overwrite the replacement refresh
        // and keep retired source segments pinned indefinitely.
        let _refresh_guard = self.primary_key_refresh_lock.lock().await;
        {
            let callback_directory = Arc::clone(&self.directory);
            let callback_schema = Arc::clone(&self.schema);
            let callback_manager = Arc::downgrade(&self.segment_manager);
            let callback_primary_key = Arc::downgrade(&self.primary_key_index);
            let callback_refresh_lock = Arc::downgrade(&self.primary_key_refresh_lock);
            self.segment_manager.set_replacement_refresh(move || {
                let directory = Arc::clone(&callback_directory);
                let schema = Arc::clone(&callback_schema);
                let manager = callback_manager.clone();
                let primary_key = callback_primary_key.clone();
                let refresh_lock = callback_refresh_lock.clone();
                async move {
                    let (Some(manager), Some(primary_key), Some(refresh_lock)) = (
                        manager.upgrade(),
                        primary_key.upgrade(),
                        refresh_lock.upgrade(),
                    ) else {
                        return Ok(());
                    };
                    refresh_primary_key_snapshot(
                        &directory,
                        &schema,
                        &manager,
                        &primary_key,
                        &refresh_lock,
                        PrimaryKeyRefresh::Replacement,
                    )
                    .await
                }
            });
        }

        let snapshot = self.segment_manager.acquire_snapshot().await;
        let current_seg_ids: Vec<String> = snapshot.segment_ids().to_vec();

        // Try to load persisted bloom filter.
        let cached = match self
            .directory
            .open_read(std::path::Path::new(PK_BLOOM_FILE))
            .await
        {
            Ok(handle) => {
                let data = handle.read_bytes_range(0..handle.len()).await;
                match data {
                    Ok(bytes) => deserialize_pk_bloom(bytes.as_slice()),
                    Err(_) => None,
                }
            }
            Err(_) => None,
        };

        // Load lightweight fast-field data for all segments concurrently.
        let load_futures: Vec<_> = current_seg_ids
            .iter()
            .map(|seg_id_str| {
                let seg_id_str = seg_id_str.clone();
                let dir = self.directory.as_ref();
                let schema = Arc::clone(&self.schema);
                let deletion = snapshot.deletions().get(&seg_id_str).cloned();
                async move { load_pk_segment_data(dir, &seg_id_str, &schema, deletion).await }
            })
            .collect();
        let all_data = futures::stream::iter(load_futures)
            .buffer_unordered(4)
            .try_collect::<Vec<_>>()
            .await?;

        if let Some((persisted_seg_ids, bloom)) = cached {
            // Partition: old segments (covered by bloom) first, new segments at end.
            let mut pk_data = Vec::with_capacity(all_data.len());
            let mut new_data = Vec::new();
            for d in all_data {
                if persisted_seg_ids.contains(&d.segment_id) {
                    pk_data.push(d);
                } else {
                    new_data.push(d);
                }
            }
            let needs_persist = !new_data.is_empty();
            let new_start = pk_data.len();
            pk_data.extend(new_data);

            let pk_index = if new_start == pk_data.len() {
                // Fast path: all segments covered by cache.
                super::primary_key::PrimaryKeyIndex::from_persisted(field, bloom, pk_data, snapshot)
            } else {
                // Incremental: only iterate new segments' keys.
                let index_label = self.schema.index_label().to_owned();
                tokio::task::spawn_blocking(move || {
                    // Insert new segments' keys into the bloom, then construct
                    // PrimaryKeyIndex with the pre-populated bloom.
                    let mut bloom = bloom;
                    let mut added = 0usize;
                    let num_new = pk_data.len() - new_start;
                    for data in &pk_data[new_start..] {
                        if let Some(ff) = data.fast_fields.get(&field.0)
                            && let Some(dict) = ff.text_dict()
                        {
                            for key in dict.iter() {
                                bloom.insert(key.as_bytes());
                                added += 1;
                            }
                        }
                    }
                    if added > 0 {
                        log::info!(
                            "[primary_key] index={index_label} bloom: added {} keys from {} new segment(s)",
                            added,
                            num_new,
                        );
                    }
                    super::primary_key::PrimaryKeyIndex::from_persisted(
                        field, bloom, pk_data, snapshot,
                    )
                })
                .await
                .map_err(|e| Error::Internal(format!("spawn_blocking failed: {}", e)))?
            };

            if needs_persist {
                self.persist_pk_bloom(&pk_index, &current_seg_ids).await;
            }

            *self.primary_key_index.write() = Some(pk_index);
        } else {
            // No cache — full rebuild, offloaded to blocking thread.
            let pk_index = tokio::task::spawn_blocking(move || {
                super::primary_key::PrimaryKeyIndex::new(field, all_data, snapshot)
            })
            .await
            .map_err(|e| Error::Internal(format!("spawn_blocking failed: {}", e)))?;

            self.persist_pk_bloom(&pk_index, &current_seg_ids).await;
            *self.primary_key_index.write() = Some(pk_index);
        }

        // The freshly built index covers every committed segment, so any
        // reservations retained after a failed post-commit refresh are
        // superseded by committed_data.
        self.pk_reservations_retained
            .store(false, Ordering::Release);

        Ok(())
    }

    /// Persist the primary-key bloom filter to `pk_bloom.bin`.
    /// Best-effort: errors are logged but not propagated.
    async fn persist_pk_bloom(
        &self,
        pk_index: &super::primary_key::PrimaryKeyIndex,
        segment_ids: &[String],
    ) {
        use super::primary_key::PK_BLOOM_FILE;

        let writer = match self
            .directory
            .streaming_writer(std::path::Path::new(PK_BLOOM_FILE))
            .await
        {
            Ok(writer) => writer,
            Err(error) => {
                log::warn!(
                    "[primary_key] index={} failed to open bloom cache: {}",
                    self.schema.index_label(),
                    error
                );
                return;
            }
        };
        let result = crate::segment::block_in_place_if_multithread(|| {
            write_pk_bloom_stream(pk_index, segment_ids, writer)
        });
        if let Err(e) = result {
            log::warn!(
                "[primary_key] index={} failed to persist bloom cache: {}",
                self.schema.index_label(),
                e
            );
        }
    }

    /// Add a document to the indexing queue (sync, O(1)).
    ///
    /// `Document` is moved into the channel (zero-copy). Workers compete to pull it.
    /// Returns an explicit backpressure error when the queue is at capacity or
    /// a prepared commit generation is not yet resolved.
    pub fn add_document(&self, doc: Document) -> Result<()> {
        self.enqueue_document(doc, false)
    }

    fn enqueue_document(&self, doc: Document, replace: bool) -> Result<()> {
        self.ensure_writer_lock()?;
        if self.worker_state.shutdown.load(Ordering::Acquire) {
            return Err(Error::IndexClosed);
        }
        if self.commit_finalization.in_progress.load(Ordering::Acquire) {
            return Err(Error::CommitInProgress);
        }
        let sender = self.doc_sender.read().clone();
        // A publication error deliberately leaves the prepared generation and
        // its workers paused for a lossless retry. Report this as backpressure
        // instead of inserting/rolling back a PK key against a closed channel.
        if sender.is_closed() {
            return Err(Error::CommitInProgress);
        }
        // Reject unencodable documents before they enter a worker queue. A
        // worker discovers these limits only after mutating a segment builder,
        // which invalidates every sibling document in the commit generation.
        validate_vector_value_counts(&doc, &self.schema)?;
        super::content_hash::document_hash(&doc, &self.schema)?;
        let primary_key_index = self.primary_key_index.read();
        let enqueue = |doc, row| {
            sender
                .try_send(QueuedDocument { doc, row })
                .map_err(|error| match error {
                    async_channel::TrySendError::Full(_) => Error::QueueFull,
                    async_channel::TrySendError::Closed(_) => Error::CommitInProgress,
                })
        };
        if let Some(pk) = primary_key_index.as_ref() {
            pk.admit_document(doc, &self.schema, replace, |doc, row| {
                enqueue(doc, Some(row))
            })
        } else {
            enqueue(doc, None)
        }
    }

    /// Stage deletion of the latest committed or unpublished row by exact key.
    /// Commit publishes visibility atomically. Missing keys are idempotent.
    pub fn delete_primary_key(&mut self, key: &str) -> Result<()> {
        self.ensure_writer_lock()?;
        if self.worker_state.shutdown.load(Ordering::Acquire) {
            return Err(Error::IndexClosed);
        }
        if self.commit_finalization.in_progress.load(Ordering::Acquire)
            || self.doc_sender.read().is_closed()
        {
            return Err(Error::CommitInProgress);
        }
        if self.pk_reservations_retained.load(Ordering::Acquire) {
            return Err(Error::CommitInProgress);
        }
        let guard = self.primary_key_index.read();
        let pk = guard.as_ref().ok_or_else(|| {
            Error::Schema("row deletion requires initialized primary-key deduplication".into())
        })?;
        pk.delete(key)?;
        Ok(())
    }

    /// Replace a committed row (or insert a missing key). Deletion and the new
    /// document become visible together at commit. Queue rejection rolls back
    /// this call's deletion so a failed upsert cannot remove the old row.
    /// Equal configured content hashes are accepted no-ops. Stored-hash I/O
    /// precedes mutation, so cancelling that read leaves pending work unchanged.
    pub async fn upsert_document(&mut self, doc: Document) -> Result<()> {
        let field = self
            .schema
            .primary_field()
            .ok_or_else(|| Error::Schema("upserts require a primary key".into()))?;
        let key = super::primary_key::document_key(&doc, field)?;
        validate_vector_value_counts(&doc, &self.schema)?;
        let hash = super::content_hash::document_hash(&doc, &self.schema)?;
        // Validate writer admission before touching the reservation set.
        self.ensure_writer_lock()?;
        if self.worker_state.shutdown.load(Ordering::Acquire) {
            return Err(Error::IndexClosed);
        }
        if self.commit_finalization.in_progress.load(Ordering::Acquire)
            || self.doc_sender.read().is_closed()
            || self.pk_reservations_retained.load(Ordering::Acquire)
        {
            return Err(Error::CommitInProgress);
        }
        if let Some(hash) = hash {
            let target = {
                let guard = self.primary_key_index.read();
                let pk = guard.as_ref().ok_or_else(|| {
                    Error::Schema("upserts require initialized primary-key deduplication".into())
                })?;
                match pk.staged_hash_matches(key, hash) {
                    Some(true) => {
                        log::debug!(
                            "[content_hash] index={} skipped unchanged staged upsert",
                            self.schema.index_label()
                        );
                        return Ok(());
                    }
                    Some(false) => None,
                    None => pk.content_hash_target(key)?,
                }
            };
            if let Some(target) = target
                && target.matches(hash, &self.schema).await?
            {
                return Ok(());
            }
        }
        if self.primary_key_index.read().is_none() {
            return Err(Error::Schema(
                "upserts require initialized primary-key deduplication".into(),
            ));
        }
        self.enqueue_document(doc, true)
    }

    /// Add multiple documents to the indexing queue.
    ///
    /// Returns the number of documents successfully queued. Stops at the first
    /// backpressure error and returns the count queued so far.
    pub fn add_documents(&self, documents: Vec<Document>) -> Result<usize> {
        let total = documents.len();
        for (i, doc) in documents.into_iter().enumerate() {
            match self.add_document(doc) {
                Ok(()) => {}
                Err(Error::QueueFull | Error::CommitInProgress) => return Ok(i),
                Err(e) => return Err(e),
            }
        }
        Ok(total)
    }

    // ========================================================================
    // Worker loop
    // ========================================================================

    /// Worker loop — runs on a dedicated OS thread, survives across commits.
    ///
    /// Outer loop: each iteration processes one commit cycle.
    ///   Inner loop: pull documents from MPMC queue, index them, build segments
    ///   when memory budget is exceeded.
    ///   On channel close (prepare_commit): flush current builder, signal
    ///   flush_count, wait for resume with new receiver.
    ///   On shutdown (Drop): exit permanently.
    fn worker_loop(
        state: Arc<WorkerState<D>>,
        initial_receiver: async_channel::Receiver<QueuedDocument>,
        handle: tokio::runtime::Handle,
        worker_id: usize,
    ) {
        let mut receiver = initial_receiver;
        let mut my_epoch = 0usize;
        let soft_flush_threshold =
            soft_flush_threshold(state.memory_budget_per_worker, worker_id, state.num_workers);
        let hard_flush_threshold = hard_flush_threshold(state.memory_budget_per_worker);

        loop {
            // Wrap the recv+build phase in catch_unwind so a panic doesn't
            // prevent flush_count from being signaled (which would hang
            // prepare_commit forever).
            let build_result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
                let mut builder: Option<SegmentBuilder> = None;
                let mut staged_rows = Arc::new(StagedSegment::default());

                while let Ok(doc) = receiver.recv_blocking() {
                    if state.shutdown.load(Ordering::Acquire) {
                        break;
                    }
                    // Another worker already invalidated this generation.
                    // Drain the shared queue so prepare_commit can complete,
                    // but do not spend CPU/RAM building outputs that must be
                    // discarded transactionally.
                    if state.cycle_failed.load(Ordering::Acquire) {
                        continue;
                    }
                    // Initialize builder if needed
                    if builder.is_none() {
                        match SegmentBuilder::new(
                            state.segment_manager.published_generation().schema.clone(),
                            state.builder_config.clone(),
                        ) {
                            Ok(mut b) => {
                                for (field, tokenizer) in state.tokenizers.read().iter() {
                                    b.set_tokenizer(*field, tokenizer.clone_box());
                                }
                                builder = Some(b);
                            }
                            Err(e) => {
                                log::error!("Failed to create segment builder: {:?}", e);
                                state.record_cycle_error(format!(
                                    "failed to create segment builder: {e}"
                                ));
                                continue;
                            }
                        }
                    }

                    let b = builder.as_mut().unwrap();
                    if let Some(row) = &doc.row
                        && !row.attach(&staged_rows, b.num_docs())
                    {
                        continue;
                    }
                    if let Err(e) = b.add_document(doc.doc) {
                        log::error!("Failed to index document: {:?}", e);
                        state.record_cycle_error(format!("failed to index document: {e}"));
                        continue;
                    }

                    let builder_memory = b.estimated_memory_bytes();

                    if b.num_docs() & 0x3FFF == 0 {
                        log::debug!(
                            "[indexing] index={} docs={}, memory={}, budget={}",
                            state.schema.index_label(),
                            b.num_docs(),
                            crate::format_bytes(builder_memory as u64),
                            crate::format_bytes(state.memory_budget_per_worker as u64)
                        );
                    }

                    // Require minimum 100 docs before flushing to avoid tiny segments
                    const MIN_DOCS_BEFORE_FLUSH: u32 = 100;

                    if b.num_docs() >= MIN_DOCS_BEFORE_FLUSH
                        && let Some(_build_permit) = state.segment_build_limiter.reserve_if_due(
                            builder_memory,
                            soft_flush_threshold,
                            hard_flush_threshold,
                        )
                    {
                        log::info!(
                            "[indexing] index={} memory budget reached, building segment: \
                             worker={}, docs={}, memory={}, soft_budget={}, hard_budget={}",
                            state.schema.index_label(),
                            worker_id,
                            b.num_docs(),
                            crate::format_bytes(builder_memory as u64),
                            crate::format_bytes(soft_flush_threshold as u64),
                            crate::format_bytes(hard_flush_threshold as u64),
                        );
                        let full_builder = builder.take().unwrap();
                        Self::build_segment_inline(
                            &state,
                            full_builder,
                            std::mem::take(&mut staged_rows),
                            &handle,
                        );
                    }
                }

                // Channel closed — flush current builder
                if !state.cycle_failed.load(Ordering::Acquire)
                    && let Some(b) = builder.take()
                    && b.num_docs() > 0
                {
                    let _build_permit = state.segment_build_limiter.acquire_flush();
                    Self::build_segment_inline(&state, b, staged_rows, &handle);
                }
            }));

            if build_result.is_err() {
                log::error!(
                    "[worker] index={} panic during indexing cycle — documents in this cycle may be lost",
                    state.schema.index_label()
                );
                state.record_cycle_error("indexing worker panicked while building the batch");
            }

            // Signal flush completion (always, even after panic — prevents
            // prepare_commit from hanging)
            let prev = state.flush_count.fetch_add(1, Ordering::Release);
            if prev + 1 == state.num_workers {
                // Last worker — wake prepare_commit. notify_all, not
                // notify_one: a cancelled commit leaves its detached
                // spawn_blocking waiter parked on this condvar, and with a
                // single notification that dead waiter would consume the
                // only wakeup, stalling a retried prepare_commit for its
                // full deadline.
                let _lock = state.flush_mutex.lock();
                state.flush_cvar.notify_all();
            }

            // Wait for resume (new channel) or shutdown.
            // Check resume_epoch to avoid re-cloning a stale receiver from
            // a previous cycle.
            {
                let mut lock = state.resume_receiver.lock();
                loop {
                    if state.shutdown.load(Ordering::Acquire) {
                        return;
                    }
                    let current_epoch = state.resume_epoch.load(Ordering::Acquire);
                    if current_epoch > my_epoch
                        && let Some(rx) = lock.as_ref()
                    {
                        receiver = rx.clone();
                        my_epoch = current_epoch;
                        break;
                    }
                    state.resume_cvar.wait(&mut lock);
                }
            }
        }
    }

    /// Build a segment on the worker thread. Uses `Handle::block_on()` to bridge
    /// into async context for I/O (streaming writers). CPU work (rayon) stays on
    /// the worker thread / rayon pool.
    fn build_segment_inline(
        state: &WorkerState<D>,
        builder: SegmentBuilder,
        staged_rows: Arc<StagedSegment>,
        handle: &tokio::runtime::Handle,
    ) {
        let segment_id = SegmentId::new();
        let segment_hex = segment_id.to_hex();
        // Claim the ID before the first file write. The guard is moved into
        // `PreparedSegment` on success and otherwise releases automatically.
        let operation = match state
            .segment_manager
            .protect_new_segment(segment_hex.clone())
        {
            Ok(operation) => operation,
            Err(e) => {
                log::error!(
                    "[segment_build_failed] index={} segment_id={} lifecycle_error={}",
                    state.schema.index_label(),
                    segment_hex,
                    e,
                );
                state.record_cycle_error(format!(
                    "failed to claim segment {segment_hex} for building: {e}"
                ));
                return;
            }
        };
        let trained = state.segment_manager.trained_for_segment_build();
        let doc_count = builder.num_docs();
        let build_start = std::time::Instant::now();

        log::info!(
            "[segment_build] index={} segment_id={} doc_count={} ann={}",
            state.schema.index_label(),
            segment_hex,
            doc_count,
            trained.is_some()
        );

        // Construct the cleanup owner before building. It keeps lifecycle
        // ownership through async deletion on ordinary error, abort, and
        // panic unwind; crash recovery is the only path left to the sweeper.
        let mut prepared = PreparedSegment {
            id: segment_hex.clone(),
            segment_id,
            num_docs: doc_count,
            staged_rows,
            segment_manager: Arc::clone(&state.segment_manager),
            operation: Some(operation),
            runtime: handle.clone(),
            needs_vector_upgrade: trained.is_none(),
            published: false,
        };

        match handle.block_on(builder.build(
            state.directory.as_ref(),
            segment_id,
            trained.as_deref(),
        )) {
            Ok(meta) if meta.num_docs == doc_count && meta.num_docs > 0 => {
                let duration_ms = build_start.elapsed().as_millis() as u64;
                log::info!(
                    "[segment_build_done] index={} segment_id={} doc_count={} duration_ms={}",
                    state.schema.index_label(),
                    segment_hex,
                    meta.num_docs,
                    duration_ms,
                );
                prepared.num_docs = meta.num_docs;
                state.built_segments.lock().push(prepared);
            }
            Ok(meta) => {
                let error = format!(
                    "segment {segment_hex} built {} docs from a {doc_count}-document builder",
                    meta.num_docs
                );
                log::error!(
                    "[segment_build_failed] index={} {error}",
                    state.schema.index_label()
                );
                state.record_cycle_error(error);
            }
            Err(e) => {
                log::error!(
                    "[segment_build_failed] index={} segment_id={} error={:?}",
                    state.schema.index_label(),
                    segment_hex,
                    e
                );
                // `prepared` owns the lifecycle claim and schedules one
                // tracked, idempotent cleanup pass when this scope ends.
                state.record_cycle_error(format!("failed to build segment {segment_hex}: {e}"));
            }
        }
    }

    // ========================================================================
    // Public API — commit, merge, etc.
    // ========================================================================

    /// Check merge policy and spawn a background merge if needed.
    pub async fn maybe_merge(&self) {
        self.segment_manager.maybe_merge().await;
    }

    /// Drain all in-flight merge tasks.
    /// Blocking merge phases cannot be cancelled safely once started.
    pub async fn abort_merges(&self) {
        self.segment_manager.abort_merges().await;
    }

    /// Stop accepting lifecycle work, stop and join indexing workers, and
    /// discard unpublished segments. Index deletion calls this while holding
    /// the registry writer lock so in-flight requests finish first and stale
    /// writer Arcs cannot restart work afterward.
    pub async fn shutdown(&mut self) -> Result<()> {
        self.segment_manager.begin_shutdown();
        self.signal_worker_shutdown();

        // A cancelled commit request leaves its owned finalizer running. Do not
        // clear shared PK/prepared state while that task may still publish or
        // refresh it. Worker shutdown is signalled first, so a successful
        // finalizer cannot restart ingestion while deletion is waiting.
        self.commit_finalization.wait_until_idle().await;

        let workers = std::mem::take(&mut self.workers);
        let panicked = tokio::task::spawn_blocking(move || {
            workers
                .into_iter()
                .map(|worker| worker.join().is_err())
                .filter(|panicked| *panicked)
                .count()
        })
        .await
        .map_err(|error| Error::Internal(format!("failed to join index workers: {}", error)))?;
        if panicked > 0 {
            log::error!(
                "[index_shutdown] index={} {} indexing worker(s) panicked",
                self.schema.index_label(),
                panicked
            );
        }

        // No commit is possible after shutdown. Dropping these RAII values
        // releases their lifecycle ownership before directory deletion.
        self.flushed_segments.lock().clear();
        self.worker_state.built_segments.lock().clear();
        if let Some(pk_index) = self.primary_key_index.write().as_mut() {
            pk_index.clear_uncommitted();
        }
        Ok(())
    }

    /// Wait for the in-flight background merge to complete (if any).
    pub async fn wait_for_merging_thread(&self) {
        self.segment_manager.wait_for_merging_thread().await;
    }

    /// Wait for all eligible merges to complete, including cascading merges.
    pub async fn wait_for_all_merges(&self) {
        self.segment_manager.wait_for_all_merges().await;
    }

    /// Wait until an owned commit finalizer has reconciled durable metadata,
    /// primary-key state, and worker availability. Normally callers need not
    /// use this: it exists for orderly shutdown and request supervisors that
    /// want to observe completion after cancelling their original waiter.
    pub async fn wait_for_commit_finalization(&self) {
        self.commit_finalization.wait_until_idle().await;
    }

    /// Get the segment tracker for sharing with readers.
    pub fn tracker(&self) -> std::sync::Arc<crate::segment::SegmentTracker> {
        self.segment_manager.tracker()
    }

    /// Acquire a snapshot of current segments for reading.
    pub async fn acquire_snapshot(&self) -> crate::segment::SegmentSnapshot {
        self.segment_manager.acquire_snapshot().await
    }

    /// Clean up orphan segment files not registered in metadata.
    ///
    /// Requires the single-writer lock: sweeping while another process's
    /// writer is live would delete its in-flight segment outputs.
    pub async fn cleanup_orphan_segments(&self) -> Result<usize> {
        self.ensure_writer_lock()?;
        self.segment_manager.cleanup_orphan_segments().await
    }

    /// Prepare commit — signal workers to flush, wait for completion, collect segments.
    ///
    /// All documents sent via `add_document` before this call are guaranteed
    /// to be written to segment files on disk. Segments are NOT yet registered
    /// in metadata — call `PreparedCommit::commit()` for that.
    ///
    /// Workers are NOT destroyed — they flush their builders and wait for
    /// `resume_workers()` to give them a new channel.
    ///
    /// `add_document` returns `CommitInProgress` until commit/abort resumes workers.
    /// On `CommitFlushTimeout`, retry the same generation: workers may still
    /// be building and must not be resumed or discarded before they finish.
    pub async fn prepare_commit(&mut self) -> Result<PreparedCommit<'_, D>> {
        self.prepare_commit_with_timeout(std::time::Duration::from_secs(300))
            .await
    }

    pub(super) async fn prepare_commit_with_timeout(
        &mut self,
        flush_timeout: std::time::Duration,
    ) -> Result<PreparedCommit<'_, D>> {
        self.ensure_writer_lock()?;
        if self.worker_state.shutdown.load(Ordering::Acquire) {
            return Err(Error::IndexClosed);
        }
        if self.commit_finalization.in_progress.load(Ordering::Acquire) {
            return Err(Error::CommitInProgress);
        }
        // 1. Close channel → workers drain remaining docs and flush builders
        self.doc_sender.read().close();
        self.worker_state.segment_build_limiter.begin_flush();

        // Wake any workers still waiting on resume_cvar from previous cycle.
        // They'll clone the stale receiver, enter recv_blocking, get Err
        // immediately (sender already closed), flush, and signal completion.
        self.worker_state.resume_cvar.notify_all();

        // 2. Wait for all workers to complete their flush (via spawn_blocking
        //    to avoid blocking the tokio runtime)
        let state = Arc::clone(&self.worker_state);
        let index_label = self.schema.index_label().to_owned();
        let all_flushed = tokio::task::spawn_blocking(move || {
            let mut lock = state.flush_mutex.lock();
            let deadline = std::time::Instant::now() + flush_timeout;
            while state.flush_count.load(Ordering::Acquire) < state.num_workers {
                let remaining = deadline.saturating_duration_since(std::time::Instant::now());
                if remaining.is_zero() {
                    log::error!(
                        "[prepare_commit] index={index_label} timed out waiting for workers: {}/{} flushed",
                        state.flush_count.load(Ordering::Acquire),
                        state.num_workers
                    );
                    return false;
                }
                state.flush_cvar.wait_for(&mut lock, remaining);
            }
            true
        })
        .await
        .map_err(|e| Error::Internal(format!("Failed to wait for workers: {}", e)))?;

        if !all_flushed {
            // Keep this commit cycle paused. Resetting flush_count and handing
            // out a new receiver while an old worker is still building lets
            // that late worker increment the *next* cycle's counter. A later
            // prepare can then return before all of its workers flushed and
            // publish an incomplete set of segments. The caller may retry
            // prepare_commit; it will observe the same generation and collect
            // every completed output once the lagging worker finishes.
            return Err(Error::CommitFlushTimeout {
                flushed_workers: self.worker_state.flush_count.load(Ordering::Acquire),
                total_workers: self.worker_state.num_workers,
            });
        }

        let cycle_error = { self.worker_state.cycle_error.lock().take() };
        if let Some(error) = cycle_error {
            // No partial publication: some documents in this generation no
            // longer exist in a worker builder, so successful sibling outputs
            // cannot be committed without violating commit's all-prior-docs
            // guarantee. Their RAII drops retain ownership through deletion.
            self.flushed_segments.lock().clear();
            self.worker_state.built_segments.lock().clear();
            self.clear_uncommitted_pk_reservations();
            self.resume_workers();
            return Err(Error::Internal(format!(
                "indexing generation failed; no documents from this batch were committed: {error}"
            )));
        }

        // 3. Collect built segments
        let built = std::mem::take(&mut *self.worker_state.built_segments.lock());
        self.flushed_segments.lock().extend(built);

        Ok(PreparedCommit {
            writer: self,
            is_resolved: false,
        })
    }

    /// Commit (convenience): prepare_commit + commit in one call.
    ///
    /// Guarantees all prior `add_document` calls are committed.
    /// Vector training is decoupled — call `build_vector_index()` manually.
    pub async fn commit(&mut self) -> Result<bool> {
        self.prepare_commit().await?.commit().await
    }

    /// Commit admitted mutations, then physically remove deleted rows from
    /// the selected segment using a bounded scratch budget. Row addresses can
    /// change; primary keys remain stable.
    pub async fn compact_segment(
        &mut self,
        segment_id: &str,
        memory_budget: usize,
    ) -> Result<bool> {
        if memory_budget < 1024 * 1024 {
            return Err(Error::Schema(
                "compaction memory budget must be at least 1 MiB".into(),
            ));
        }
        if crate::segment::SegmentId::from_hex(segment_id).is_none() {
            return Err(Error::Document("invalid compaction segment ID".into()));
        }
        self.commit().await?;
        let changed = self
            .segment_manager
            .compact_segment(segment_id, memory_budget)
            .await?;
        self.persist_replacement_snapshot().await?;
        Ok(changed)
    }

    /// Compact every tombstoned segment in the current snapshot. Works even
    /// when the index contains only one segment.
    pub async fn compact(&mut self, memory_budget: usize) -> Result<usize> {
        if memory_budget < 1024 * 1024 {
            return Err(Error::Schema(
                "compaction memory budget must be at least 1 MiB".into(),
            ));
        }
        self.commit().await?;
        self.wait_for_merging_thread().await;
        let ids = self.segment_manager.get_segment_ids().await;
        let mut count = 0;
        for id in ids {
            count += usize::from(
                self.segment_manager
                    .compact_segment(&id, memory_budget)
                    .await?,
            );
        }
        self.persist_replacement_snapshot().await?;
        Ok(count)
    }

    /// Force merge all segments into one.
    pub async fn force_merge(&mut self) -> Result<()> {
        self.force_merge_with_snapshot_refresh(|| std::future::ready(Ok(())))
            .await
    }

    /// Force merge while refreshing an external segment consumer after the
    /// background-merge drain and every durable replacement.
    ///
    /// Segment publication refreshes the writer's primary-key topology through
    /// the manager's lifecycle-owned hook. Servers use this callback to reload
    /// their cached `IndexReader` as well.
    pub async fn force_merge_with_snapshot_refresh<F, Fut>(
        &mut self,
        refresh_external: F,
    ) -> Result<()>
    where
        F: FnMut() -> Fut,
        Fut: std::future::Future<Output = Result<()>>,
    {
        self.force_merge_with_compaction_and_snapshot_refresh(false, refresh_external)
            .await
    }

    /// Merge normally, optionally compacting the final outputs once. Defaults
    /// to retaining tombstones through `force_merge()`.
    pub async fn force_merge_with_compaction(&mut self, compact: bool) -> Result<()> {
        self.force_merge_with_compaction_and_snapshot_refresh(compact, || {
            std::future::ready(Ok(()))
        })
        .await
    }

    pub async fn force_merge_with_compaction_and_snapshot_refresh<F, Fut>(
        &mut self,
        compact: bool,
        refresh_external: F,
    ) -> Result<()>
    where
        F: FnMut() -> Fut,
        Fut: std::future::Future<Output = Result<()>>,
    {
        let budget = compact.then_some(self.config.compaction_memory_budget_bytes);
        if budget.is_some_and(|bytes| bytes < 1024 * 1024) {
            return Err(Error::Schema(
                "compaction memory budget must be at least 1 MiB".into(),
            ));
        }
        self.prepare_commit().await?.commit().await?;

        self.segment_manager
            .force_merge_with_compaction_and_snapshot_refresh(budget, refresh_external)
            .await?;

        // Segment IDs in the on-disk bloom cache need only the final
        // generation. Persisting the unchanged bloom after every hierarchy
        // level adds avoidable I/O on large primary-key indexes.
        self.persist_replacement_snapshot().await
    }

    /// Maintain each segment: reorder text with Recursive Graph Bisection,
    /// compact ANN runs, and consolidate sparse nomination runs.
    ///
    /// Sparse maintenance preserves forward values and limits consolidation
    /// work to its configured budget.
    pub async fn reorder(&mut self) -> Result<()> {
        self.reorder_with_snapshot_refresh(|| std::future::ready(Ok(())))
            .await
    }

    /// Reorder behind a shared writer without holding its lock during
    /// maintenance admission or segment rewriting. The retained Arc keeps
    /// the single-writer file lock alive until the operation finishes.
    pub async fn reorder_with_shared_writer<F, Fut>(
        writer: &Arc<tokio::sync::RwLock<Self>>,
        refresh_external: F,
    ) -> Result<()>
    where
        F: FnMut() -> Fut,
        Fut: std::future::Future<Output = Result<()>>,
    {
        let segment_manager = {
            let mut writer = writer.write().await;
            writer.commit().await?;
            Arc::clone(&writer.segment_manager)
        };
        segment_manager
            .reorder_segments_with_snapshot_refresh(refresh_external)
            .await?;
        writer.read().await.persist_replacement_snapshot().await
    }

    /// Reorder while refreshing an external reader after each durable segment
    /// replacement, so retired sources are released during a long pass.
    pub async fn reorder_with_snapshot_refresh<F, Fut>(&mut self, refresh_external: F) -> Result<()>
    where
        F: FnMut() -> Fut,
        Fut: std::future::Future<Output = Result<()>>,
    {
        self.prepare_commit().await?.commit().await?;

        self.segment_manager
            .reorder_segments_with_snapshot_refresh(refresh_external)
            .await?;
        self.persist_replacement_snapshot().await
    }

    /// Persist the final topology after a bounded series of replacements.
    async fn persist_replacement_snapshot(&self) -> Result<()> {
        refresh_primary_key_snapshot(
            &self.directory,
            &self.schema,
            &self.segment_manager,
            &self.primary_key_index,
            &self.primary_key_refresh_lock,
            PrimaryKeyRefresh::FinalReplacement,
        )
        .await
    }

    /// Get the segment manager (for background optimizer access).
    pub fn segment_manager(&self) -> &Arc<crate::merge::SegmentManager<D>> {
        &self.segment_manager
    }

    /// Resume workers with a fresh channel. Called after commit or abort.
    ///
    /// Workers are already alive — just give them a new channel and wake them.
    /// If the tokio runtime has shut down (e.g., program exit), this is a no-op.
    fn resume_workers(&mut self) {
        Self::resume_workers_shared(&self.worker_state, &self.doc_sender);
    }

    fn resume_workers_shared(
        worker_state: &Arc<WorkerState<D>>,
        doc_sender: &Arc<parking_lot::RwLock<async_channel::Sender<QueuedDocument>>>,
    ) {
        if worker_state.shutdown.load(Ordering::Acquire) {
            return;
        }
        if tokio::runtime::Handle::try_current().is_err() {
            // Runtime is gone — signal permanent shutdown so workers don't
            // hang forever on resume_cvar.
            worker_state.shutdown.store(true, Ordering::Release);
            worker_state.resume_cvar.notify_all();
            return;
        }

        // Reset flush count for next cycle
        worker_state.segment_build_limiter.end_flush();
        worker_state.flush_count.store(0, Ordering::Release);
        *worker_state.cycle_error.lock() = None;
        worker_state.cycle_failed.store(false, Ordering::Release);

        // Create new channel
        let (sender, receiver) = async_channel::bounded(PIPELINE_MAX_SIZE_IN_DOCS);
        *doc_sender.write() = sender;

        // Set new receiver, bump epoch, and wake all workers
        {
            let mut lock = worker_state.resume_receiver.lock();
            *lock = Some(receiver);
        }
        worker_state.resume_epoch.fetch_add(1, Ordering::Release);
        worker_state.resume_cvar.notify_all();
    }

    fn signal_worker_shutdown(&self) {
        self.worker_state.shutdown.store(true, Ordering::Release);
        self.doc_sender.read().close();
        self.worker_state.segment_build_limiter.begin_flush();
        self.worker_state.resume_cvar.notify_all();
    }

    // Vector index methods (build_vector_index, etc.) are in vector_builder.rs
}

impl<D: DirectoryWriter + 'static> Drop for IndexWriter<D> {
    fn drop(&mut self) {
        self.signal_worker_shutdown();
        for w in std::mem::take(&mut self.workers) {
            let _ = w.join();
        }
    }
}

/// A prepared commit that can be finalized or aborted.
///
/// Two-phase commit guard. Between `prepare_commit()` and
/// `commit()`/`abort()`, segments are on disk but NOT in metadata.
/// Dropping without calling either will auto-abort (discard segments,
/// respawn workers).
pub struct PreparedCommit<'a, D: DirectoryWriter + 'static> {
    writer: &'a mut IndexWriter<D>,
    is_resolved: bool,
}

/// Returns prepared segments to the writer if an owned commit finalizer fails
/// or unwinds before it can establish that metadata owns them. Retrying commit
/// is safe even when publication actually won the race: `SegmentManager::commit`
/// is idempotent and the operation guards keep the files protected meanwhile.
struct PreparedSegmentsGuard<D: DirectoryWriter + 'static> {
    segments: Option<Vec<PreparedSegment<D>>>,
    retry_slot: Arc<parking_lot::Mutex<Vec<PreparedSegment<D>>>>,
}

impl<D: DirectoryWriter + 'static> PreparedSegmentsGuard<D> {
    fn metadata_entries(&self) -> Vec<(String, u32)> {
        self.segments
            .as_deref()
            .unwrap_or_default()
            .iter()
            .map(PreparedSegment::metadata_entry)
            .collect()
    }

    fn staged_deletions(&self) -> Vec<(String, Arc<StagedSegment>)> {
        self.segments
            .as_deref()
            .unwrap_or_default()
            .iter()
            .filter(|segment| segment.staged_rows.is_dirty())
            .map(|segment| (segment.id.clone(), Arc::clone(&segment.staged_rows)))
            .collect()
    }

    fn take_published(&mut self) -> Vec<PreparedSegment<D>> {
        self.segments.take().unwrap_or_default()
    }

    fn vector_upgrade_segment_ids(&self) -> Vec<String> {
        self.segments
            .as_deref()
            .unwrap_or_default()
            .iter()
            .filter(|segment| segment.needs_vector_upgrade)
            .map(|segment| segment.id.clone())
            .collect()
    }
}

impl<D: DirectoryWriter + 'static> Drop for PreparedSegmentsGuard<D> {
    fn drop(&mut self) {
        if let Some(segments) = self.segments.take() {
            self.retry_slot.lock().extend(segments);
        }
    }
}

/// Couples completion of the owned commit task to writer availability. The
/// default is deliberately fail-closed: a pre-publication error or panic keeps
/// workers paused so the retained prepared generation can be retried. Only the
/// normal published path arms resumption.
struct CommitFinalizationGuard<D: DirectoryWriter + 'static> {
    state: Arc<CommitFinalizationState>,
    worker_state: Arc<WorkerState<D>>,
    doc_sender: Arc<parking_lot::RwLock<async_channel::Sender<QueuedDocument>>>,
    resume_workers: bool,
}

impl<D: DirectoryWriter + 'static> CommitFinalizationGuard<D> {
    fn resume_on_drop(&mut self) {
        self.resume_workers = true;
    }
}

impl<D: DirectoryWriter + 'static> Drop for CommitFinalizationGuard<D> {
    fn drop(&mut self) {
        if self.resume_workers {
            IndexWriter::<D>::resume_workers_shared(&self.worker_state, &self.doc_sender);
        }
        self.state.finish();
    }
}

/// Everything needed to finish one prepared generation is moved into this
/// value before spawning. Its two guards therefore reconcile segment
/// ownership and writer availability even if Tokio drops the task before its
/// first poll.
struct OwnedCommitFinalization<D: DirectoryWriter + 'static> {
    directory: Arc<D>,
    schema: Arc<Schema>,
    segment_manager: Arc<crate::merge::SegmentManager<D>>,
    primary_key_index: Arc<parking_lot::RwLock<Option<super::primary_key::PrimaryKeyIndex>>>,
    primary_key_refresh_lock: Arc<tokio::sync::Mutex<()>>,
    prepared: PreparedSegmentsGuard<D>,
    finalization: Option<CommitFinalizationGuard<D>>,
    publication_observed: Arc<AtomicBool>,
    pk_reservations_retained: Arc<AtomicBool>,
}

#[derive(Clone, Copy)]
enum PrimaryKeyRefresh {
    /// A commit may introduce genuinely new keys and persists the cache.
    Commit,
    /// A merge/reorder only changes segment topology; keys are already in the
    /// monotonic bloom and the intermediate segment IDs need not be persisted.
    Replacement,
    /// Final topology refresh: still no key hashing, but persist the new set of
    /// segment IDs alongside the unchanged bloom.
    FinalReplacement,
}

async fn refresh_primary_key_snapshot<D: DirectoryWriter + 'static>(
    directory: &Arc<D>,
    schema: &Arc<Schema>,
    segment_manager: &Arc<crate::merge::SegmentManager<D>>,
    primary_key_index: &Arc<parking_lot::RwLock<Option<super::primary_key::PrimaryKeyIndex>>>,
    primary_key_refresh_lock: &Arc<tokio::sync::Mutex<()>>,
    refresh: PrimaryKeyRefresh,
) -> Result<()> {
    let _refresh_guard = primary_key_refresh_lock.lock().await;
    let snapshot = segment_manager.acquire_snapshot().await;
    let existing_ids: std::collections::HashSet<String> = {
        let guard = primary_key_index.read();
        let Some(pk_index) = guard.as_ref() else {
            return Ok(());
        };
        pk_index
            .committed_visibility()
            .filter(|(id, deletion)| *deletion == snapshot.deletions().get(*id).map(|(_, d)| d))
            .map(|(id, _)| id.to_owned())
            .collect()
    };

    let load_futures: Vec<_> = snapshot
        .segment_ids()
        .iter()
        .filter(|id| !existing_ids.contains(id.as_str()))
        .map(|seg_id_str| {
            let seg_id_str = seg_id_str.clone();
            let dir = directory.as_ref();
            let schema = Arc::clone(schema);
            let deletion = snapshot.deletions().get(&seg_id_str).cloned();
            async move { load_pk_segment_data(dir, &seg_id_str, &schema, deletion).await }
        })
        .collect();
    let mut new_data = futures::stream::iter(load_futures)
        .buffer_unordered(4)
        .try_collect::<Vec<_>>()
        .await?;
    if let Some(field) = schema.primary_field() {
        new_data = tokio::task::spawn_blocking(move || {
            for data in &mut new_data {
                data.prepare_live_keys(field);
            }
            new_data
        })
        .await
        .map_err(|error| {
            Error::Internal(format!("primary-key visibility refresh failed: {error}"))
        })?;
    }
    let seg_ids: Vec<String> = snapshot.segment_ids().to_vec();

    let persist_bloom = {
        let mut guard = primary_key_index.write();
        let Some(pk_index) = guard.as_mut() else {
            return Ok(());
        };
        match refresh {
            PrimaryKeyRefresh::Commit => pk_index.refresh_incremental(new_data, snapshot),
            PrimaryKeyRefresh::Replacement | PrimaryKeyRefresh::FinalReplacement => {
                pk_index.refresh_replacement(new_data, snapshot);
            }
        }
        matches!(
            refresh,
            PrimaryKeyRefresh::Commit | PrimaryKeyRefresh::FinalReplacement
        )
    };

    if persist_bloom {
        let writer = match directory
            .streaming_writer(std::path::Path::new(super::primary_key::PK_BLOOM_FILE))
            .await
        {
            Ok(writer) => writer,
            Err(error) => {
                log::warn!(
                    "[primary_key] index={} failed to open bloom cache: {}",
                    schema.index_label(),
                    error
                );
                return Ok(());
            }
        };
        // The outer read guard prevents replacement of the PK index while the
        // inner state lock streams its bloom. No corpus-sized Vec is created.
        let guard = primary_key_index.read();
        if let Some(pk_index) = guard.as_ref()
            && let Err(error) = crate::segment::block_in_place_if_multithread(|| {
                write_pk_bloom_stream(pk_index, &seg_ids, writer)
            })
        {
            log::warn!(
                "[primary_key] index={} failed to persist bloom cache: {}",
                schema.index_label(),
                error
            );
        }
    }
    Ok(())
}

fn write_pk_bloom_stream(
    pk_index: &super::primary_key::PrimaryKeyIndex,
    segment_ids: &[String],
    mut writer: Box<dyn crate::directories::StreamingWriter>,
) -> std::io::Result<()> {
    pk_index.write_bloom_cache(segment_ids, writer.as_mut())?;
    writer.finish()
}

async fn finalize_prepared_commit<D: DirectoryWriter + 'static>(
    mut commit: OwnedCommitFinalization<D>,
) -> Result<bool> {
    let metadata_entries = commit.prepared.metadata_entries();
    let published_segment_ids = commit.prepared.vector_upgrade_segment_ids();

    // This entire future is owned by a Tokio task. Cancelling the RPC only
    // drops its JoinHandle; it cannot split durable metadata publication from
    // PK reservations or worker resumption.
    let deletes = commit
        .primary_key_index
        .read()
        .as_ref()
        .map_or_else(Vec::new, |pk| pk.pending_deletes());
    commit
        .segment_manager
        .commit_with_deletes(
            &metadata_entries,
            deletes,
            commit.prepared.staged_deletions(),
        )
        .await?;
    commit.publication_observed.store(true, Ordering::Release);
    if let Some(pk) = commit.primary_key_index.read().as_ref() {
        pk.mark_deletes_published();
    }

    let mut published = commit.prepared.take_published();
    for segment in &mut published {
        segment.mark_published();
    }
    drop(published);
    commit
        .segment_manager
        .schedule_vector_segment_upgrades(published_segment_ids);
    // Publication is irreversible. From here onward every exit path, including
    // panic unwind, must make the writer available again while PK reservations
    // remain fail-closed until refresh succeeds.
    if let Some(finalization) = commit.finalization.as_mut() {
        finalization.resume_on_drop();
    } else {
        log::error!("owned commit finalization guard was already released after publication");
    }

    // Metadata publication is the commit point. Cache refresh is fail-closed:
    // retaining the generation's uncommitted keys may cause conservative
    // duplicate rejections, but can never admit a duplicate or turn a durable
    // commit into an API error.
    match refresh_primary_key_snapshot(
        &commit.directory,
        &commit.schema,
        &commit.segment_manager,
        &commit.primary_key_index,
        &commit.primary_key_refresh_lock,
        PrimaryKeyRefresh::Commit,
    )
    .await
    {
        // A successful refresh folded every committed key into committed_data
        // and cleared the reservations — nothing retained anymore.
        Ok(()) => commit
            .pk_reservations_retained
            .store(false, Ordering::Release),
        Err(error) => {
            // The retained reservations are now the ONLY record of the
            // published segments' keys. Abort paths must not clear them
            // (see clear_uncommitted_pk_reservations) or duplicates would
            // be admitted.
            commit
                .pk_reservations_retained
                .store(true, Ordering::Release);
            log::error!(
                "[primary_key] committed metadata but failed to refresh dedup state; \
                 retaining reservations until a later successful commit: {}",
                error,
            );
        }
    }

    // Merge scheduling is optional post-commit work and may briefly wait on
    // manager state. Reconcile worker availability first so it cannot extend
    // ingestion backpressure after metadata and PK state already agree.
    drop(commit.finalization.take());
    commit.segment_manager.maybe_merge().await;
    Ok(true)
}

impl<'a, D: DirectoryWriter + 'static> PreparedCommit<'a, D> {
    /// Finalize: register segments in metadata, evaluate merge policy, resume workers.
    ///
    /// Returns `true` if new segments were committed, `false` if nothing changed.
    pub async fn commit(mut self) -> Result<bool> {
        let segments = std::mem::take(&mut *self.writer.flushed_segments.lock());

        // Fast path: nothing to commit
        if segments.is_empty()
            && self
                .writer
                .primary_key_index
                .read()
                .as_ref()
                .is_none_or(|pk| pk.pending_deletes().is_empty())
        {
            log::debug!(
                "[commit] index={} no segments to commit, skipping",
                self.writer.schema.index_label()
            );
            self.is_resolved = true;
            self.writer.resume_workers();
            return Ok(false);
        }

        if !self.writer.commit_finalization.begin() {
            self.writer.flushed_segments.lock().extend(segments);
            // Keep the prepared generation paused. Letting `Drop` auto-abort
            // here would delete the retryable segments owned by another
            // finalization state transition.
            self.is_resolved = true;
            return Err(Error::CommitInProgress);
        }

        let publication_observed = Arc::new(AtomicBool::new(false));
        let owned = OwnedCommitFinalization {
            directory: Arc::clone(&self.writer.directory),
            schema: Arc::clone(&self.writer.schema),
            segment_manager: Arc::clone(&self.writer.segment_manager),
            primary_key_index: Arc::clone(&self.writer.primary_key_index),
            primary_key_refresh_lock: Arc::clone(&self.writer.primary_key_refresh_lock),
            prepared: PreparedSegmentsGuard {
                segments: Some(segments),
                retry_slot: Arc::clone(&self.writer.flushed_segments),
            },
            finalization: Some(CommitFinalizationGuard {
                state: Arc::clone(&self.writer.commit_finalization),
                worker_state: Arc::clone(&self.writer.worker_state),
                doc_sender: Arc::clone(&self.writer.doc_sender),
                resume_workers: false,
            }),
            publication_observed: Arc::clone(&publication_observed),
            pk_reservations_retained: Arc::clone(&self.writer.pk_reservations_retained),
        };

        // From this point the owned value, not this cancel-sensitive guard,
        // controls every segment and the paused worker generation. Resolve the
        // local guard before spawning so even a runtime-spawn panic cannot
        // auto-abort the retryable generation during unwind.
        self.is_resolved = true;
        let task_publication = Arc::clone(&publication_observed);
        let task = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
            tokio::spawn(async move {
                match std::panic::AssertUnwindSafe(finalize_prepared_commit(owned))
                    .catch_unwind()
                    .await
                {
                    Ok(result) => result,
                    Err(_) if task_publication.load(Ordering::Acquire) => {
                        log::error!(
                            "owned commit finalizer panicked after metadata publication; \
                             treating the durable generation as committed"
                        );
                        Ok(true)
                    }
                    Err(_) => Err(Error::Internal(
                        "owned commit finalizer panicked before metadata publication".into(),
                    )),
                }
            })
        }))
        .map_err(|_| Error::Internal("runtime rejected owned commit finalizer".into()))?;

        match task.await {
            Ok(result) => result,
            Err(error) if publication_observed.load(Ordering::Acquire) => {
                log::error!(
                    "owned commit finalizer terminated after metadata publication: {}; \
                     treating the durable generation as committed",
                    error,
                );
                Ok(true)
            }
            Err(error) => Err(Error::Internal(format!(
                "owned commit finalizer terminated unexpectedly: {error}"
            ))),
        }
    }

    /// Abort: discard prepared segments, delete their files asynchronously,
    /// and resume workers. Lifecycle ownership is held until deletion ends.
    pub fn abort(mut self) {
        self.is_resolved = true;
        self.writer.flushed_segments.lock().clear();
        self.writer.clear_uncommitted_pk_reservations();
        self.writer.resume_workers();
    }
}

impl<D: DirectoryWriter + 'static> Drop for PreparedCommit<'_, D> {
    fn drop(&mut self) {
        if !self.is_resolved {
            log::warn!("PreparedCommit dropped without commit/abort — auto-aborting");
            self.writer.flushed_segments.lock().clear();
            self.writer.clear_uncommitted_pk_reservations();
            self.writer.resume_workers();
        }
    }
}

#[cfg(test)]
#[path = "tests/staged_admission.rs"]
mod staged_admission;