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};
pub trait Worker: Embed + Ask + Send + Sync {}
impl<W: Embed + Ask + Send + Sync> Worker for W {}
#[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(),
}
}
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)
}
pub type Field = (String, String, String);
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)]
#[serde(rename_all = "snake_case")]
pub enum WorkerState {
EmbedWorkerStarting,
Ready,
Restarting,
}
#[derive(Debug, Clone, Serialize)]
pub struct Failed {
pub branch: String,
pub collection: String,
pub field: String,
pub at: Timestamp,
pub failure: Failure,
}
#[derive(Debug, Default, Clone, Copy, PartialEq, Eq)]
enum Phase {
#[default]
Waiting,
Stepping,
Committing,
Finished,
}
impl Phase {
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>,
known: BTreeSet<Field>,
worker: Option<WorkerState>,
stopping: bool,
phase: Phase,
left: bool,
died: Option<String>,
pages: u64,
made: Vec<(Field, u64)>,
longer_than_model: Vec<(Field, u64)>,
embed_input_empty: Vec<(Field, u64)>,
windowed: Vec<(Field, u64)>,
windows_over_cap: Vec<(Field, u64)>,
ended: BTreeSet<Field>,
last_failure: Option<Failed>,
durable: SyncCounts,
starts: u64,
deaths: u32,
next_start: Option<(Instant, Timestamp)>,
#[cfg(test)]
waits: Vec<Duration>,
#[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()
}
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)
}
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,
)
}
}
struct Questions {
worker: Arc<dyn Worker>,
restart: Restart,
state: Mutex<Asked>,
changed: Condvar,
}
#[derive(Debug, Default)]
struct Asked {
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)
}
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
}
}
}
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(),
});
}
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 {
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
}
}
pub struct Driver {
admission: Arc<Admission>,
shared: Arc<Shared>,
worker: Arc<dyn Worker>,
questions: Option<Arc<Questions>>,
question_start: Mutex<Option<JoinHandle<()>>>,
restart: Restart,
thread: Mutex<Option<JoinHandle<()>>>,
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()
}
}
#[must_use]
pub fn embedder(settings: Settings) -> Arc<dyn Worker> {
Arc::new(Embedder::new(Some(settings)))
}
impl Driver {
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())
}
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)
}
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(),
})
});
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
}
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();
}
}
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();
}
pub fn asker(&self) -> Result<&dyn Ask, Failure> {
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",
)),
}
}
#[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",
}
}
#[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
}
#[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
}
#[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))
}
#[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
}
#[must_use]
pub fn status(&self) -> Value {
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,
"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 {
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;
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 {
fn drop(&mut self) {
if let Err(failure) = self.stop() {
eprintln!("haematite embedding: {failure}");
}
}
}
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)
}
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);
match field {
None => start_worker(shared, worker, restart),
Some(field) => page(
shared,
(storage, admission),
pending_reader,
worker,
bounds,
(restart, None),
field,
)?,
}
}
}
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();
}
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())?;
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 {
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());
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);
if made.page.reach != Reach::End {
state.ended.remove(&field);
}
state.waiting.insert(field);
}
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,
);
}
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(())
}
fn counted(counts: &[(Field, u64)], field: &Field) -> u64 {
counts
.iter()
.find(|(known, _)| known == field)
.map_or(0, |(_, number)| *number)
}
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)
}
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()
}
#[must_use]
pub fn waits(&self) -> Vec<Duration> {
self.shared.lock().waits.clone()
}
pub fn wait_walk_ended(&self) {
let mut state = self.shared.lock();
while state.phase != Phase::Finished {
state = self.shared.wait(state);
}
drop(state);
}
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);
}
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;