haematite_cli/serve/

embedding.rs

//! EW1: the embedding driver.
//!
//! One thread reads pending pages through held immutable snapshots and has
//! the worker embed their texts, then admits one durable batch commit.
//! Foreground admission excludes later embedding commits while a write waits.
//! At most one already-admitted batch, sized by `embed_page_records`, finishes
//! one full flush before that writer; model, diff and text reads hold no writer.
//!
//! What waits is a deduplicated queue of fields, woken by a write that
//! touches a declared collection and by a `vector.declare`, and filled at
//! start with every field declared on every live branch. A page that did not
//! end its walk puts its field at the back. Failed pages also go to the back,
//! with a per-field delay that doubles to the restart ceiling. Only a page
//! commit clears that count; a successful worker probe does not.
//!
//! The worker is started by the thread, off the serving path. EWR1: a worker
//! that died or timed out is started again by the thread, at once the first
//! time, then after a wait that doubles up to a ceiling (`Restart`), for as
//! long as the service runs; the field it was making goes back to wait, and
//! its walk carries on from its last commit. The wait is on the driver's own
//! signal, never with the storage lock held. Until a worker is ready a text
//! search waits for one up to `embed_ready_ms`, then is refused as
//! `embed_worker_starting` or `embed_worker_restarting`. Shutdown tells the
//! thread and joins it: a page under way ends its worker call, within the
//! call bound, and its commit.

use std::collections::{BTreeMap, BTreeSet, VecDeque};
use std::num::NonZeroUsize;
use std::sync::{Arc, Condvar, Mutex, MutexGuard, PoisonError};
use std::thread::JoinHandle;
use std::time::{Duration, Instant};

#[path = "embedding_admission.rs"]
mod admission;
pub(in crate::serve) use admission::{Admission, Guard};

use haematite::SyncCounts;
use haematite::branch::{Timestamp, current_timestamp};
use serde::Serialize;
use serde_json::{Value, json};

use super::embed::job::{self, Bounds, Embed};
use super::embed::question::Ask;
use super::embed::settings::Restart;
use super::embed::{Embedded, Embedder, Settings};
use super::protocol::{Selection, Version};
use super::storage::embed_job::Reach;
use super::storage::embed_pending::PendingReader;
use super::storage::embedded::Position;
use super::storage::{Failure, Storage};

/// What the driver needs of a worker: documents for the walk, a question
/// for a search.
pub trait Worker: Embed + Ask + Send + Sync {}

impl<W: Embed + Ask + Send + Sync> Worker for W {}

/// One page reads `embed_page_records` records, has them embedded in worker
/// calls of `embed_request_items` texts, and commits once (EW1-fix).
#[must_use]
pub fn bounds(settings: &Settings) -> Bounds {
    let input_bytes = usize::try_from(settings.limits.input_bytes)
        .ok()
        .and_then(NonZeroUsize::new)
        .unwrap_or(NonZeroUsize::MIN);
    Bounds {
        page_records: settings.page_records,
        batch: settings.request_items,
        input_bytes,
        call: settings.document_call(),
    }
}

/// HANG1: the most stopping the driver waits for its threads: a page under
/// way (each of its worker calls, `embed_call_ms` and with windows on that
/// for each window a text may need, ⌈`page_records` / batch⌉ of them), its
/// commit's wait for a writer (`embed_writer_wait_ms`), and
/// the question worker's start (`embed_ready_ms`). `None` past the clock's
/// range, when stopping waits for nothing.
fn stop_within(bounds: Bounds, writer_wait: Duration, restart: Restart) -> Option<Duration> {
    let calls = bounds.page_records.get().div_ceil(bounds.batch.get());
    bounds
        .call
        .checked_mul(u32::try_from(calls).ok()?)?
        .checked_add(writer_wait)?
        .checked_add(restart.ready)
}

/// One declared field of one branch.
pub type Field = (String, String, String);

/// The worker's state, as the status names it.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)]
#[serde(rename_all = "snake_case")]
pub enum WorkerState {
    EmbedWorkerStarting,
    Ready,
    /// It died or timed out, and is to be started again (EWR1).
    Restarting,
}

/// A failure kept for the status, with the field and the time.
#[derive(Debug, Clone, Serialize)]
pub struct Failed {
    pub branch: String,
    pub collection: String,
    pub field: String,
    pub at: Timestamp,
    pub failure: Failure,
}

/// Where the walk's thread is.
#[derive(Debug, Default, Clone, Copy, PartialEq, Eq)]
enum Phase {
    /// Waiting for a field, or for the worker's start.
    #[default]
    Waiting,
    /// Making a page of a field it took.
    Stepping,
    /// STOP2: a page's commit has started and not yet been answered.
    Committing,
    /// HANG1: the thread has left `run` and says nothing more.
    Finished,
}

impl Phase {
    /// A page is under way: being made, or committed.
    const fn stepping(self) -> bool {
        matches!(self, Self::Stepping | Self::Committing)
    }
}

#[derive(Debug, Default)]
struct Waiting {
    order: VecDeque<Field>,
    names: BTreeSet<Field>,
}

impl Waiting {
    fn insert(&mut self, field: Field) -> bool {
        if !self.names.insert(field.clone()) {
            return false;
        }
        self.order.push_back(field);
        true
    }

    fn requeue(&mut self, field: Field) {
        self.order.retain(|queued| queued != &field);
        self.names.insert(field.clone());
        self.order.push_back(field);
    }

    fn contains(&self, field: &Field) -> bool {
        self.names.contains(field)
    }

    #[cfg(test)]
    fn is_empty(&self) -> bool {
        self.order.is_empty()
    }

    fn pop_ready(
        &mut self,
        failures: &BTreeMap<Field, PageFailure>,
        now: Instant,
    ) -> Option<Field> {
        for _ in 0..self.order.len() {
            let field = self.order.pop_front()?;
            if failures
                .get(&field)
                .and_then(|failed| failed.due)
                .is_some_and(|due| due > now)
            {
                self.order.push_back(field);
            } else {
                self.names.remove(&field);
                return Some(field);
            }
        }
        None
    }
}

#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
struct PagePosition {
    committed: Option<Position>,
    toward: Version,
    end_after: Option<String>,
}

#[derive(Debug, Serialize)]
struct PageFailure {
    position: Option<PagePosition>,
    consecutive_failures: u32,
    failure: Failure,
    at: Timestamp,
    next_attempt_due: Option<Timestamp>,
    #[serde(skip)]
    due: Option<Instant>,
}

enum Scheduled {
    Stop,
    Start,
    Page(Field),
    Wait(Option<Duration>),
}

#[derive(Debug, Default)]
struct State {
    waiting: Waiting,
    failures: BTreeMap<Field, PageFailure>,
    /// Every field the driver knows is declared, for the wake on a write.
    known: BTreeSet<Field>,
    worker: Option<WorkerState>,
    stopping: bool,
    phase: Phase,
    /// STOP2: `stop` gave up on the thread, which starts no commit after.
    left: bool,
    died: Option<String>,
    pages: u64,
    /// Vectors made, by field, since the service started.
    made: Vec<(Field, u64)>,
    /// EMB1: records given no vector, by field, since the service started,
    /// because the model reads only part of their text.
    longer_than_model: Vec<(Field, u64)>,
    /// EMB-E1: records passed with no vector, by field, since the service
    /// started, because their text is empty or only whitespace.
    embed_input_empty: Vec<(Field, u64)>,
    /// EMB-E2: records given windows in place of one vector, by field.
    windowed: Vec<(Field, u64)>,
    /// EMB-E2: records given no vector because their whole cover needs more
    /// windows than the cap, by field.
    windows_over_cap: Vec<(Field, u64)>,
    /// Walks that reached their end, by field.
    ended: BTreeSet<Field>,
    last_failure: Option<Failed>,
    durable: SyncCounts,
    /// Times the worker was started since the service began (EWR1).
    starts: u64,
    /// Deaths since the worker last answered.
    deaths: u32,
    /// When the next start is due, by the clock and as a time.
    next_start: Option<(Instant, Timestamp)>,
    /// Each wait before a start, in order.
    #[cfg(test)]
    waits: Vec<Duration>,
    /// Searches waiting for a worker now.
    #[cfg(test)]
    askers: usize,
}

impl State {
    fn failed_page(
        &mut self,
        field: &Field,
        position: Option<PagePosition>,
        failure: Failure,
        restart: Restart,
        now: Instant,
    ) -> Result<(), String> {
        let consecutive = self
            .failures
            .get(field)
            .map_or(1, |failed| failed.consecutive_failures.saturating_add(1));
        let wait = restart.wait(consecutive.saturating_add(1));
        let due = now
            .checked_add(wait)
            .ok_or("embed_page_retry_clock: retry deadline cannot be represented")?;
        let nanos = u64::try_from(wait.as_nanos())
            .map_err(|error| format!("embed_page_retry_clock: {error}"))?;
        let at = current_timestamp();
        self.failures.insert(
            field.clone(),
            PageFailure {
                position,
                consecutive_failures: consecutive,
                failure,
                at,
                next_attempt_due: Some(at.saturating_add(nanos)),
                due: Some(due),
            },
        );
        self.waiting.requeue(field.clone());
        Ok(())
    }

    fn committed_page(&mut self, field: &Field) {
        if let Some(failed) = self.failures.get_mut(field) {
            failed.consecutive_failures = 0;
            failed.next_attempt_due = None;
            failed.due = None;
        }
    }

    fn next_page(&mut self, now: Instant) -> Option<Field> {
        self.waiting.pop_ready(&self.failures, now)
    }

    fn scheduled(&mut self, now: Instant) -> Scheduled {
        if self.stopping {
            return Scheduled::Stop;
        }
        if self.worker == Some(WorkerState::Restarting) {
            return match self
                .next_start
                .and_then(|(due, _at)| due.checked_duration_since(now))
            {
                Some(left) if !left.is_zero() => Scheduled::Wait(Some(left)),
                _ => Scheduled::Start,
            };
        }
        if let Some(field) = self.next_page(now) {
            self.phase = Phase::Stepping;
            return Scheduled::Page(field);
        }
        Scheduled::Wait(self.page_wait(now))
    }

    fn page_wait(&self, now: Instant) -> Option<Duration> {
        self.waiting
            .names
            .iter()
            .filter_map(|field| self.failures.get(field).and_then(|failed| failed.due))
            .filter_map(|due| due.checked_duration_since(now))
            .min()
    }

    /// The worker died or timed out: it is to be started again after the
    /// wait its deaths in a row call for.
    fn died(&mut self, restart: Restart, failed: Failed) {
        self.died_at(restart, failed, Instant::now());
    }

    fn died_at(&mut self, restart: Restart, failed: Failed, now: Instant) {
        self.worker = Some(WorkerState::Restarting);
        self.deaths = self.deaths.saturating_add(1);
        let wait = restart.wait(self.deaths);
        let nanos = u64::try_from(wait.as_nanos()).unwrap_or(u64::MAX);
        let due = now.checked_add(wait).unwrap_or(now);
        self.next_start = Some((due, failed.at.saturating_add(nanos)));
        self.last_failure = Some(failed);
        #[cfg(test)]
        self.waits.push(wait);
    }
}

#[derive(Debug, Default)]
struct Shared {
    state: Mutex<State>,
    changed: Condvar,
}

impl Shared {
    fn lock(&self) -> MutexGuard<'_, State> {
        self.state.lock().unwrap_or_else(PoisonError::into_inner)
    }

    fn wait<'a>(&self, state: MutexGuard<'a, State>) -> MutexGuard<'a, State> {
        self.changed
            .wait(state)
            .unwrap_or_else(PoisonError::into_inner)
    }

    /// As [`Self::wait`], for at most `most`; the lock is let go meanwhile.
    fn wait_for<'a>(&self, state: MutexGuard<'a, State>, most: Duration) -> MutexGuard<'a, State> {
        self.changed.wait_timeout(state, most).map_or_else(
            |poisoned| poisoned.into_inner().0,
            |(state, _timeout)| state,
        )
    }
}

/// SL1: the worker that answers questions alone, apart from the walk's, so
/// a search never waits behind a walk batch: one `Embedder` keeps one worker
/// behind one lock for a whole batch (embedder.rs `embed`). SL2: it is
/// started when the service starts, on a thread of its own, so the first
/// search does not pay the model's load; a search that comes first waits for
/// it up to `embed_ready_ms`. After a death the next question starts it
/// again at once; a death after that is followed by the same doubling waits
/// as the walk's worker (EWR1), and a question before the start is due is
/// refused as `embed_worker_restarting`. It holds its own copy of the model,
/// so with it the service's workers may use twice `embed_memory_bytes`.
struct Questions {
    worker: Arc<dyn Worker>,
    restart: Restart,
    state: Mutex<Asked>,
    changed: Condvar,
}

/// The question worker's state, as the status names it.
#[derive(Debug, Default)]
struct Asked {
    /// The start begun with the service is under way (SL2).
    starting: bool,
    ready: bool,
    starts: u64,
    deaths: u32,
    next_start: Option<(Instant, Timestamp)>,
    last_failure: Option<Failed>,
}

impl Questions {
    fn lock(&self) -> MutexGuard<'_, Asked> {
        self.state.lock().unwrap_or_else(PoisonError::into_inner)
    }

    /// Begin the start on a thread of its own; the caller does not wait. A
    /// thread that cannot be made leaves the start to the first question,
    /// and says why in the status.
    fn begin(questions: &Arc<Self>) -> Option<JoinHandle<()>> {
        questions.lock().starting = true;
        let starting = Arc::clone(questions);
        let begun = std::thread::Builder::new()
            .name("haematite-question".to_owned())
            .spawn(move || {
                let started = starting.worker.question("ready");
                starting.lock().starting = false;
                starting.answered(true, &started);
            });
        match begun {
            Ok(thread) => Some(thread),
            Err(error) => {
                let mut asked = questions.lock();
                asked.starting = false;
                asked.last_failure = Some(Failed {
                    branch: String::new(),
                    collection: String::new(),
                    field: String::new(),
                    at: current_timestamp(),
                    failure: Failure::new("embed_thread_unstarted", error.to_string()),
                });
                drop(asked);
                None
            }
        }
    }

    /// Keep what a call of the worker came to: ready, or a death and when
    /// the next start is due; `starting` counts it as a start.
    fn answered(&self, starting: bool, answer: &Result<Embedded, Failure>) {
        let mut asked = self.lock();
        if starting {
            asked.starts = asked.starts.saturating_add(1);
        }
        match answer {
            Ok(_embedded) => {
                asked.ready = true;
                asked.deaths = 0;
                asked.next_start = None;
            }
            Err(failure) if failure.code.starts_with("embed_worker") => {
                asked.ready = false;
                asked.deaths = asked.deaths.saturating_add(1);
                let wait = self.restart.wait(asked.deaths);
                let at = current_timestamp();
                let nanos = u64::try_from(wait.as_nanos()).unwrap_or(u64::MAX);
                let due = Instant::now()
                    .checked_add(wait)
                    .unwrap_or_else(Instant::now);
                asked.next_start = Some((due, at.saturating_add(nanos)));
                asked.last_failure = Some(Failed {
                    branch: String::new(),
                    collection: String::new(),
                    field: String::new(),
                    at,
                    failure: failure.clone(),
                });
            }
            // A refused question is the call's, not the worker's: the worker
            // stays ready, and the refusal is kept as the last failure, as a
            // page refusal is.
            Err(failure) => {
                asked.last_failure = Some(Failed {
                    branch: String::new(),
                    collection: String::new(),
                    field: String::new(),
                    at: current_timestamp(),
                    failure: failure.clone(),
                });
            }
        }
        drop(asked);
        self.changed.notify_all();
    }

    fn status(&self) -> Value {
        let asked = self.lock();
        let state = match (asked.starting, asked.ready, asked.deaths) {
            (true, _, _) => "starting",
            (false, true, _) => "ready",
            (false, false, 0) => "not_started",
            (false, false, _) => "restarting",
        };
        let answer = json!({
            "state": state,
            "starts": asked.starts,
            "next_start_due": asked.next_start.map(|(_due, at)| at),
            "last_failure": asked.last_failure,
        });
        drop(asked);
        answer
    }
}

impl Ask for Questions {
    /// The question, asked of this worker alone with no lock held; while
    /// the start begun with the service is under way, after it, waiting up
    /// to `embed_ready_ms` on this worker's own signal.
    fn question(&self, text: &str) -> Result<Embedded, Failure> {
        let until = Instant::now().checked_add(self.restart.ready);
        let mut asked = self.lock();
        while asked.starting {
            let left = until.and_then(|until| until.checked_duration_since(Instant::now()));
            let Some(left) = left.filter(|left| !left.is_zero()) else {
                return Err(Failure::new(
                    "embed_worker_starting",
                    "the question worker has not answered yet; text search waits for it",
                ));
            };
            asked = self.changed.wait_timeout(asked, left).map_or_else(
                |poisoned| poisoned.into_inner().0,
                |(asked, _timeout)| asked,
            );
        }
        let due = asked.next_start.map(|(due, _at)| due);
        let starting = !asked.ready;
        drop(asked);
        if due.is_some_and(|due| due > Instant::now()) {
            return Err(Failure::new(
                "embed_worker_restarting",
                "the question worker died and is being started again; store.stats says when",
            ));
        }
        let answer = self.worker.question(text);
        self.answered(starting, &answer);
        answer
    }
}

/// The embedding thread and what it shares with the service.
pub struct Driver {
    admission: Arc<Admission>,
    shared: Arc<Shared>,
    worker: Arc<dyn Worker>,
    /// SL1: the question worker, when questions are kept apart from the walk.
    questions: Option<Arc<Questions>>,
    /// SL2: the thread starting it, begun with the service.
    question_start: Mutex<Option<JoinHandle<()>>>,
    restart: Restart,
    thread: Mutex<Option<JoinHandle<()>>>,
    /// HANG1: the most `stop` waits for the threads ([`stop_within`]).
    stop_within: Option<Duration>,
}

impl std::fmt::Debug for Driver {
    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        formatter.debug_struct("Driver").finish_non_exhaustive()
    }
}

/// The service's embedder over its settings.
#[must_use]
pub fn embedder(settings: Settings) -> Arc<dyn Worker> {
    Arc::new(Embedder::new(Some(settings)))
}

impl Driver {
    /// Start the thread over `storage` with `worker`, restarted as
    /// [`Restart::default`] says: every field declared on every live branch
    /// waits.
    ///
    /// # Errors
    /// The declared fields unread, or the thread unstarted, by name.
    pub fn start(
        storage: Arc<Mutex<Storage>>,
        worker: Arc<dyn Worker>,
        bounds: Bounds,
        writer_wait: Duration,
    ) -> Result<Self, Failure> {
        Self::start_with(storage, worker, bounds, writer_wait, Restart::default())
    }

    /// As [`Self::start`], with the worker started again as `restart` says.
    ///
    /// # Errors
    /// The declared fields unread, or the thread unstarted, by name.
    pub fn start_with(
        storage: Arc<Mutex<Storage>>,
        worker: Arc<dyn Worker>,
        bounds: Bounds,
        writer_wait: Duration,
        restart: Restart,
    ) -> Result<Self, Failure> {
        Self::start_apart(storage, worker, None, bounds, writer_wait, restart)
    }

    /// As [`Self::start_with`], with questions asked of `questions` when it
    /// is given, never of the walk's `worker` (SL1). The question worker
    /// never commits, so it never takes the admission (JOIN1).
    ///
    /// # Errors
    /// The declared fields unread, or the thread unstarted, by name.
    pub fn start_apart(
        storage: Arc<Mutex<Storage>>,
        worker: Arc<dyn Worker>,
        questions: Option<Arc<dyn Worker>>,
        bounds: Bounds,
        writer_wait: Duration,
        restart: Restart,
    ) -> Result<Self, Failure> {
        let admission = Arc::new(Admission::new(writer_wait)?);
        let (fields, pending_reader) = {
            let writer = storage.lock().map_err(|_| {
                Failure::new("writer_unavailable", "writer state requires recovery")
            })?;
            (declared(&writer)?, writer.pending_reader())
        };
        let shared = Arc::new(Shared::default());
        let mut state = shared.lock();
        state.known.extend(fields.iter().cloned());
        for field in fields {
            state.waiting.insert(field);
        }
        state.worker = Some(WorkerState::EmbedWorkerStarting);
        drop(state);
        let thread_shared = Arc::clone(&shared);
        let thread_worker = Arc::clone(&worker);
        let thread_admission = Arc::clone(&admission);
        let thread = std::thread::Builder::new()
            .name("haematite-embed".to_owned())
            .spawn(move || {
                let ended = run(
                    &thread_shared,
                    &storage,
                    &thread_admission,
                    &pending_reader,
                    thread_worker.as_ref(),
                    bounds,
                    restart,
                );
                let mut state = thread_shared.lock();
                state.phase = Phase::Finished;
                if !state.stopping {
                    state.died = Some(ended.err().unwrap_or_else(|| {
                        "embed_thread_ended: it stopped without being told".to_owned()
                    }));
                }
                drop(state);
                thread_shared.changed.notify_all();
            })
            .map_err(|error| Failure::new("embed_thread_unstarted", error.to_string()))?;
        let questions = questions.map(|worker| {
            Arc::new(Questions {
                worker,
                restart,
                state: Mutex::new(Asked::default()),
                changed: Condvar::new(),
            })
        });
        // SL2: begun now, off the serving path; nothing waits for it here.
        let question_start = questions.as_ref().and_then(Questions::begin);
        Ok(Self {
            admission,
            shared,
            worker,
            questions,
            question_start: Mutex::new(question_start),
            restart,
            thread: Mutex::new(Some(thread)),
            stop_within: stop_within(bounds, writer_wait, restart),
        })
    }

    pub(in crate::serve) fn admission(&self) -> &Admission {
        &self.admission
    }

    /// A write on `branch` touched `collection`: each field declared there
    /// waits, once.
    pub fn wrote(&self, branch: &str, collection: &str) {
        let mut guard = self.shared.lock();
        let state = &mut *guard;
        let mut added = false;
        for field in state
            .known
            .iter()
            .filter(|(known, at, _field)| known == branch && at == collection)
        {
            added |= state.waiting.insert(field.clone());
        }
        drop(guard);
        if added {
            self.shared.changed.notify_all();
        }
    }

    /// `field` of `collection` on `branch` was declared: it is known, and
    /// its walk waits.
    pub fn declared(&self, branch: &str, collection: &str, field: &str) {
        let key = (branch.to_owned(), collection.to_owned(), field.to_owned());
        let mut state = self.shared.lock();
        state.known.insert(key.clone());
        state.ended.remove(&key);
        state.waiting.insert(key);
        drop(state);
        self.shared.changed.notify_all();
    }

    /// The worker, for a question, once it is ready: a search that finds
    /// none ready waits for one up to `embed_ready_ms`, on the driver's
    /// signal, with no storage lock held.
    ///
    /// # Errors
    /// `embed_worker_starting` before the worker first answered;
    /// `embed_worker_restarting` while one that died is started again.
    pub fn asker(&self) -> Result<&dyn Ask, Failure> {
        // SL1: a question worker of its own answers whatever the walk does.
        if let Some(questions) = &self.questions {
            return Ok(questions.as_ref());
        }
        let until = Instant::now().checked_add(self.restart.ready);
        let mut state = self.shared.lock();
        #[cfg(test)]
        {
            state.askers = state.askers.saturating_add(1);
            self.shared.changed.notify_all();
        }
        let worker = loop {
            let left = until.and_then(|until| until.checked_duration_since(Instant::now()));
            match (state.worker, left) {
                (Some(WorkerState::Ready), _) => break WorkerState::Ready,
                (_, Some(left)) if !left.is_zero() && !state.stopping && state.died.is_none() => {
                    state = self.shared.wait_for(state, left);
                }
                (worker, _) => break worker.unwrap_or(WorkerState::EmbedWorkerStarting),
            }
        };
        #[cfg(test)]
        {
            state.askers = state.askers.saturating_sub(1);
        }
        drop(state);
        match worker {
            WorkerState::Ready => Ok(self.worker.as_ref()),
            WorkerState::Restarting => Err(Failure::new(
                "embed_worker_restarting",
                "the embedding worker died and is being started again; store.stats says when",
            )),
            WorkerState::EmbedWorkerStarting => Err(Failure::new(
                "embed_worker_starting",
                "the embedding worker has not answered yet; text search waits for it",
            )),
        }
    }

    /// What `vectors.fields` says of a field this driver serves: the
    /// worker's state while it is not ready, else whether its walk is open.
    #[must_use]
    pub fn field_status(&self, branch: &str, collection: &str, field: &str) -> &'static str {
        let worker = self.shared.lock().worker;
        match worker {
            Some(WorkerState::Ready) if self.walk_open(branch, collection, field) => "embedding",
            Some(WorkerState::Ready) => "embedded",
            Some(WorkerState::Restarting) => "embed_worker_restarting",
            Some(WorkerState::EmbedWorkerStarting) | None => "embed_worker_starting",
        }
    }

    /// EMB1: the records of a field this driver has given no vector since
    /// the service started, because the model reads only part of their text.
    #[must_use]
    pub fn longer_than_model(&self, branch: &str, collection: &str, field: &str) -> u64 {
        let key = (branch.to_owned(), collection.to_owned(), field.to_owned());
        let state = self.shared.lock();
        let count = counted(&state.longer_than_model, &key);
        drop(state);
        count
    }

    /// EMB-E1: the records of a field this driver has passed with no vector
    /// since the service started, because their text is empty or only
    /// whitespace (`embed_input_empty`).
    #[must_use]
    pub fn embed_input_empty(&self, branch: &str, collection: &str, field: &str) -> u64 {
        let key = (branch.to_owned(), collection.to_owned(), field.to_owned());
        let state = self.shared.lock();
        let count = counted(&state.embed_input_empty, &key);
        drop(state);
        count
    }

    /// EMB-E2: the records of a field this driver has given windows, and
    /// those given none for needing more than the cap, since it started.
    #[must_use]
    pub fn windows(&self, branch: &str, collection: &str, field: &str) -> (u64, u64) {
        let key = (branch.to_owned(), collection.to_owned(), field.to_owned());
        let state = self.shared.lock();
        let counts = (
            counted(&state.windowed, &key),
            counted(&state.windows_over_cap, &key),
        );
        drop(state);
        counts
    }

    #[must_use]
    pub fn field_failure(&self, branch: &str, collection: &str, field: &str) -> Option<Value> {
        let key = (branch.to_owned(), collection.to_owned(), field.to_owned());
        let state = self.shared.lock();
        state.failures.get(&key).map(|failed| json!(failed))
    }

    /// Whether the walk of a field is still open as far as this driver
    /// knows: it waits, or a page of it is being made.
    #[must_use]
    pub fn walk_open(&self, branch: &str, collection: &str, field: &str) -> bool {
        let key = (branch.to_owned(), collection.to_owned(), field.to_owned());
        let state = self.shared.lock();
        let open = state.waiting.contains(&key) || !state.ended.contains(&key);
        drop(state);
        open
    }

    /// The driver's state for the service's status answer.
    #[must_use]
    pub fn status(&self) -> Value {
        // A handle `stop` has taken says nothing; the thread's own flag does.
        let handle = self
            .thread
            .lock()
            .unwrap_or_else(PoisonError::into_inner)
            .as_ref()
            .map(JoinHandle::is_finished);
        let question = self
            .questions
            .as_ref()
            .map_or_else(|| json!("off"), |questions| questions.status());
        let state = self.shared.lock();
        let finished = handle.unwrap_or(state.phase == Phase::Finished);
        let thread = match (&state.died, finished, state.stopping) {
            (Some(_), _, _) | (None, true, false) => "died",
            (None, true, true) => "stopped",
            (None, false, _) if state.phase.stepping() => "embedding",
            (None, false, _) => "waiting",
        };
        let fields: Vec<Value> = state
            .known
            .iter()
            .map(|key| {
                let made = counted(&state.made, key);
                let walk = if state.ended.contains(key) && !state.waiting.contains(key) {
                    "done"
                } else {
                    "open"
                };
                json!({
                    "branch": key.0, "collection": key.1, "field": key.2,
                    "made_since_start": made, "walk": walk,
                    "longer_than_model": counted(&state.longer_than_model, key),
                    "embed_input_empty": counted(&state.embed_input_empty, key),
                    "windowed": counted(&state.windowed, key),
                    "windows_over_cap": counted(&state.windows_over_cap, key),
                    "failure": state.failures.get(key),
                })
            })
            .collect();
        let answer = json!({
            "thread": thread,
            "died": state.died,
            "worker": state.worker,
            // SL1: both workers by name; the walk's is `worker` above.
            "workers": {"walk": state.worker, "question": question},
            "pages": state.pages,
            "fields": fields,
            "last_failure": state.last_failure,
            "starts": state.starts,
            "next_start_due": state.next_start.map(|(_due, at)| at),
            "durable": {
                "record_installs": state.durable.record_installs,
                "barrier_device_flushes": state.durable.barrier_device_flushes,
                "ref_file_full_flushes": state.durable.ref_file_full_flushes,
                "directory_full_flushes": state.durable.directory_full_flushes,
            },
        });
        drop(state);
        answer
    }
}

impl Driver {
    /// HANG1: tell the threads to stop and wait for them on their own
    /// signals, for at most [`stop_within`] in all: a page under way ends its
    /// worker calls and its commit, and no new one starts. A thread still
    /// running then is named and left, never waited on for ever. Each
    /// thread is waited on once; a later call finds nothing to wait for.
    ///
    /// # Errors
    /// `embed_thread_unended`, naming each thread that did not end in time.
    pub fn stop(&self) -> Result<(), Failure> {
        let until = self
            .stop_within
            .and_then(|within| Instant::now().checked_add(within));
        let left = || until.and_then(|until| until.checked_duration_since(Instant::now()));
        let mut state = self.shared.lock();
        state.stopping = true;
        self.shared.changed.notify_all();
        while state.phase != Phase::Finished {
            let Some(wait) = left().filter(|wait| !wait.is_zero()) else {
                break;
            };
            state = self.shared.wait_for(state, wait);
        }
        let walk_ended = state.phase == Phase::Finished;
        // STOP2: from here the thread, if it is left, starts no commit; one
        // already started finishes, and is said.
        state.left = !walk_ended;
        let under_way = state.phase == Phase::Committing;
        drop(state);
        let question_ended = self.questions.as_ref().is_none_or(|questions| {
            let mut asked = questions.lock();
            while asked.starting {
                let Some(wait) = left().filter(|wait| !wait.is_zero()) else {
                    break;
                };
                asked = questions.changed.wait_timeout(asked, wait).map_or_else(
                    |poisoned| poisoned.into_inner().0,
                    |(asked, _timeout)| asked,
                );
            }
            !asked.starting
        });
        let mut unended = Vec::new();
        for (name, slot, ended) in [
            ("haematite-embed", &self.thread, walk_ended),
            ("haematite-question", &self.question_start, question_ended),
        ] {
            let thread = slot.lock().unwrap_or_else(PoisonError::into_inner).take();
            match thread {
                Some(thread) if ended => drop(thread.join()),
                Some(_left) => unended.push(name),
                None => {}
            }
        }
        if unended.is_empty() {
            return Ok(());
        }
        let commit = match (walk_ended, under_way) {
            (true, _) => "",
            (false, true) => "; its commit already under way finishes",
            (false, false) => "; it starts no commit",
        };
        Err(Failure::new(
            "embed_thread_unended",
            format!(
                "{} did not end within {:?} of being told to stop, and is left running{commit}",
                unended.join(" and "),
                self.stop_within.unwrap_or_default()
            ),
        ))
    }
}

impl Drop for Driver {
    /// [`Driver::stop`]; a thread it leaves is named on the service's
    /// standard error, the one place a drop can say it.
    fn drop(&mut self) {
        if let Err(failure) = self.stop() {
            eprintln!("haematite embedding: {failure}");
        }
    }
}

/// Every field declared on every live branch.
fn declared(storage: &Storage) -> Result<Vec<Field>, Failure> {
    let mut fields = Vec::new();
    for branch in storage.branch_names() {
        let at = storage.capture(&Selection {
            branch: branch.clone(),
            version: None,
        })?;
        for (collection, field) in storage.embedded_fields(&at)? {
            fields.push((branch.clone(), collection, field));
        }
    }
    Ok(fields)
}

/// The thread: start the worker, then take a waiting field and make one page
/// of it, until told to stop. A worker that died is started again once its
/// start is due, before any page.
fn run(
    shared: &Shared,
    storage: &Mutex<Storage>,
    admission: &Admission,
    pending_reader: &PendingReader,
    worker: &dyn Worker,
    bounds: Bounds,
    restart: Restart,
) -> Result<(), String> {
    start_worker(shared, worker, restart);
    loop {
        let mut state = shared.lock();
        let field = loop {
            match state.scheduled(Instant::now()) {
                Scheduled::Stop => return Ok(()),
                Scheduled::Start => break None,
                Scheduled::Page(field) => break Some(field),
                Scheduled::Wait(Some(left)) => state = shared.wait_for(state, left),
                Scheduled::Wait(None) => state = shared.wait(state),
            }
        };
        drop(state);
        // JOIN1: a restart takes neither the storage lock nor the
        // admission; only a page's commit goes through admission.background.
        match field {
            None => start_worker(shared, worker, restart),
            Some(field) => page(
                shared,
                (storage, admission),
                pending_reader,
                worker,
                bounds,
                (restart, None),
                field,
            )?,
        }
    }
}

/// Start the worker, off the serving path and with no lock held: one
/// question, whose answer only says the worker is ready.
fn start_worker(shared: &Shared, worker: &dyn Worker, restart: Restart) {
    let started = worker.question("ready");
    let mut state = shared.lock();
    state.starts = state.starts.saturating_add(1);
    match started {
        Ok(_ready) => {
            state.worker = Some(WorkerState::Ready);
            state.deaths = 0;
            state.next_start = None;
        }
        Err(failure) => state.died(
            restart,
            Failed {
                branch: String::new(),
                collection: String::new(),
                field: String::new(),
                at: current_timestamp(),
                failure,
            },
        ),
    }
    drop(state);
    shared.changed.notify_all();
}

/// A snapshot page and model call hold no writer. Admission lets one
/// foreground writer finish before the next batch commit can begin; the
/// storage and the admission its commit goes through come together.
fn page(
    shared: &Shared,
    (storage, admission): (&Mutex<Storage>, &Admission),
    pending_reader: &PendingReader,
    worker: &dyn Worker,
    bounds: Bounds,
    (restart, attempted): (Restart, Option<Instant>),
    field: Field,
) -> Result<(), String> {
    let (branch, collection, name) = (&field.0, &field.1, &field.2);
    let read = pending_reader.read_page(
        branch,
        collection,
        name,
        bounds.page_records,
        bounds.input_bytes,
    );
    let mut position = None;
    let outcome = match read {
        Err(failure) => Err(failure),
        Ok(None) => Ok(None),
        Ok(Some(read)) => {
            let pending = read.pending;
            position = Some(PagePosition {
                committed: read.position,
                toward: pending.toward,
                end_after: match &pending.reach {
                    Reach::Upto(after) => Some(after.clone()),
                    Reach::End => None,
                },
            });
            match job::make(
                &pending,
                collection,
                name,
                &read.declaration,
                worker,
                bounds,
            ) {
                Err(failure) => Err(failure),
                Ok(made) => {
                    let mut writer = admission
                        .background(storage)
                        .map_err(|failure| failure.to_string())?;
                    // STOP2: the last word before a commit. Under the lock
                    // `stop` gives up under, so a commit either started
                    // before it gave up, and finishes, or never starts: the
                    // page is not written and is made again on the next start.
                    let mut state = shared.lock();
                    if state.left {
                        state.phase = Phase::Waiting;
                        drop(state);
                        drop(writer);
                        return Ok(());
                    }
                    state.phase = Phase::Committing;
                    drop(state);
                    let before = writer.sync_counts();
                    let committed = writer.embed_commit(branch, &made.page, bounds.page_records);
                    let durable = writer.sync_counts().since(before);
                    drop(writer);
                    committed.map(|_| Some((made, durable)))
                }
            }
        }
    };
    finish_page(shared, field, position, outcome, restart, attempted)
}

fn finish_page(
    shared: &Shared,
    field: Field,
    position: Option<PagePosition>,
    outcome: Result<Option<(job::MadePage, SyncCounts)>, Failure>,
    restart: Restart,
    attempted: Option<Instant>,
) -> Result<(), String> {
    let mut state = shared.lock();
    state.phase = Phase::Waiting;
    let now = attempted.unwrap_or_else(Instant::now);
    match outcome {
        // Nothing pending: the walk is at the head.
        Ok(None) => {
            state.ended.insert(field);
        }
        Ok(Some((made, durable))) => {
            state.committed_page(&field);
            state.worker = Some(WorkerState::Ready);
            state.pages = state.pages.saturating_add(1);
            state.durable = plus(state.durable, durable);
            count(&mut state.made, &field, made.page.made.len());
            // EMB-E2 (i): a text past the cap is one the model reads part of.
            let over = made.windows_over_cap.len();
            let longer = made.longer_than_model.len().saturating_add(over);
            count(&mut state.longer_than_model, &field, longer);
            count(&mut state.embed_input_empty, &field, made.empty_input.len());
            count(&mut state.windowed, &field, made.page.windowed.len());
            count(&mut state.windows_over_cap, &field, over);
            // Asked again until the walk says it is at the head.
            if made.page.reach != Reach::End {
                state.ended.remove(&field);
            }
            state.waiting.insert(field);
        }
        // The worker died or timed out: it is started again, and the field
        // waits to carry on from its last commit (EWR1).
        Err(failure) if failure.code.starts_with("embed_worker") => {
            state.failed_page(&field, position, failure.clone(), restart, now)?;
            state.died_at(
                restart,
                Failed {
                    branch: field.0,
                    collection: field.1,
                    field: field.2,
                    at: current_timestamp(),
                    failure,
                },
                now,
            );
        }
        // A page refusal keeps its position and yields to every other field.
        Err(failure) => {
            state.failed_page(&field, position, failure.clone(), restart, now)?;
            state.last_failure = Some(Failed {
                branch: field.0,
                collection: field.1,
                field: field.2,
                at: current_timestamp(),
                failure,
            });
        }
    }
    drop(state);
    shared.changed.notify_all();
    Ok(())
}

/// `counts`' number for `field`; 0 for a field it does not name.
fn counted(counts: &[(Field, u64)], field: &Field) -> u64 {
    counts
        .iter()
        .find(|(known, _)| known == field)
        .map_or(0, |(_, number)| *number)
}

/// Add `more` to `counts`' number for `field`, naming the field first if
/// `counts` does not yet.
fn count(counts: &mut Vec<(Field, u64)>, field: &Field, more: usize) {
    let more = u64::try_from(more).unwrap_or(u64::MAX);
    match counts.iter_mut().find(|(known, _)| known == field) {
        Some((_, number)) => *number = number.saturating_add(more),
        None => counts.push((field.clone(), more)),
    }
}

const fn plus(left: SyncCounts, right: SyncCounts) -> SyncCounts {
    SyncCounts {
        record_installs: left.record_installs.saturating_add(right.record_installs),
        barrier_device_flushes: left
            .barrier_device_flushes
            .saturating_add(right.barrier_device_flushes),
        ref_file_full_flushes: left
            .ref_file_full_flushes
            .saturating_add(right.ref_file_full_flushes),
        directory_full_flushes: left
            .directory_full_flushes
            .saturating_add(right.directory_full_flushes),
    }
}

#[cfg(test)]
impl Driver {
    pub(in crate::serve) fn for_steps(
        storage: &Mutex<Storage>,
        worker: Arc<dyn Worker>,
        bounds: Bounds,
        restart: Restart,
    ) -> Result<Self, Failure> {
        let writer = storage
            .lock()
            .map_err(|_| Failure::new("writer_unavailable", "writer state requires recovery"))?;
        let fields = declared(&writer)?;
        drop(writer);
        let shared = Arc::new(Shared::default());
        {
            let mut state = shared.lock();
            state.known.extend(fields.iter().cloned());
            for field in fields {
                state.waiting.insert(field);
            }
        }
        start_worker(&shared, worker.as_ref(), restart);
        Ok(Self {
            admission: Arc::new(Admission::new(Duration::from_secs(1))?),
            shared,
            worker,
            questions: None,
            question_start: Mutex::new(None),
            restart,
            thread: Mutex::new(None),
            stop_within: stop_within(bounds, Duration::from_secs(1), restart),
        })
    }

    pub(in crate::serve) fn step_at(
        &self,
        storage: &Mutex<Storage>,
        bounds: Bounds,
        now: Instant,
    ) -> Result<bool, String> {
        let mut state = self.shared.lock();
        let scheduled = state.scheduled(now);
        drop(state);
        let field = match scheduled {
            Scheduled::Page(field) => Some(field),
            Scheduled::Start => None,
            Scheduled::Wait(_) | Scheduled::Stop => return Ok(false),
        };
        if let Some(field) = field {
            let reader = storage
                .lock()
                .map_err(|error| error.to_string())?
                .pending_reader();
            page(
                &self.shared,
                (storage, &self.admission),
                &reader,
                self.worker.as_ref(),
                bounds,
                (self.restart, Some(now)),
                field,
            )?;
        } else {
            start_worker(&self.shared, self.worker.as_ref(), self.restart);
        }
        Ok(true)
    }

    #[must_use]
    pub(in crate::serve) fn retry_due(
        &self,
        branch: &str,
        collection: &str,
        field: &str,
    ) -> Option<Instant> {
        self.shared
            .lock()
            .failures
            .get(&(branch.to_owned(), collection.to_owned(), field.to_owned()))
            .and_then(|failed| failed.due)
    }

    /// Wait until the worker answered or failed, nothing waits, and no page
    /// is being made, or the thread ended.
    pub fn wait_idle(&self) {
        let mut state = self.shared.lock();
        while (state.worker == Some(WorkerState::EmbedWorkerStarting)
            || state.worker == Some(WorkerState::Restarting)
            || !state.waiting.is_empty()
            || state.phase.stepping())
            && state.died.is_none()
        {
            state = self.shared.wait(state);
        }
        drop(state);
    }

    #[must_use]
    pub fn pages(&self) -> u64 {
        self.shared.lock().pages
    }

    #[must_use]
    pub fn durable(&self) -> SyncCounts {
        self.shared.lock().durable
    }

    #[must_use]
    pub fn last_failure(&self) -> Option<Failed> {
        self.shared.lock().last_failure.clone()
    }

    /// Each wait before a start of the worker, in order.
    #[must_use]
    pub fn waits(&self) -> Vec<Duration> {
        self.shared.lock().waits.clone()
    }

    /// Wait until the walk's thread has left `run` (HANG1).
    pub fn wait_walk_ended(&self) {
        let mut state = self.shared.lock();
        while state.phase != Phase::Finished {
            state = self.shared.wait(state);
        }
        drop(state);
    }

    /// Wait until the question worker's start begun with the service has
    /// answered or failed (SL2).
    pub fn wait_question_started(&self) {
        let Some(questions) = &self.questions else {
            return;
        };
        let mut asked = questions.lock();
        while asked.starting {
            asked = questions
                .changed
                .wait(asked)
                .unwrap_or_else(PoisonError::into_inner);
        }
        drop(asked);
    }

    /// Wait until `count` searches wait for a worker.
    pub fn wait_askers(&self, count: usize) {
        let mut state = self.shared.lock();
        while state.askers < count {
            state = self.shared.wait(state);
        }
        drop(state);
    }
}

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