use std::error::Error;
use std::num::{NonZeroU32, NonZeroUsize};
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
use std::time::{Duration, Instant};
use serde_json::Value;
use serde_json::value::RawValue;
use super::{
FIELD, FIVE, MANIFEST, StandIn, TestResult, bounds, declare, mutate, selection, service_with,
};
use crate::serve::embed::Embedded;
use crate::serve::embed::job::Embed;
use crate::serve::embed::question::Ask;
use crate::serve::embed::settings::Restart;
use crate::serve::embedding::Worker;
use crate::serve::protocol::{
Change, DocumentWrite, Metric, Operation, StoreRead, VectorDeclaration, VectorSearchText,
};
use crate::serve::service::Service;
use crate::serve::storage::Failure;
use crate::serve::storage::embedded::{Declaration, ModelName, VectorKind};
#[derive(Default)]
struct Poison {
succeeding: AtomicBool,
attempts: AtomicUsize,
}
impl Embed for Poison {
fn documents(&self, texts: &[String]) -> Result<Embedded, Failure> {
if texts.iter().any(|text| text == "alpha") {
self.attempts.fetch_add(1, Ordering::SeqCst);
if !self.succeeding.load(Ordering::SeqCst) {
return Err(Failure::new("embed_worker_exited", "poison page"));
}
}
Ok(StandIn::made(texts))
}
}
impl Ask for Poison {
fn question(&self, text: &str) -> Result<Embedded, Failure> {
Ok(StandIn::made(&[text.to_owned()]))
}
}
fn retry_service() -> Result<(tempfile::TempDir, Service, Arc<Poison>), Box<dyn Error>> {
let (directory, mut service) = service_with(&FIVE)?;
for (id, body) in FIVE {
mutate(
&service,
Change::Put(DocumentWrite {
collection: "notes".to_owned(),
document_id: id.to_owned(),
document: RawValue::from_string(format!(r#"{{"body":"{body}","safe":"banana"}}"#))?,
}),
)?;
}
declare(&service)?;
mutate(
&service,
Change::VectorDeclare(VectorDeclaration {
collection: "notes".to_owned(),
field: "z_good".to_owned(),
declaration: Declaration {
kind: VectorKind::Dense,
source: "safe".to_owned(),
model: ModelName {
name: "stand-in".to_owned(),
revision: "1".to_owned(),
manifest_sha256: MANIFEST.to_owned(),
},
dimensions: NonZeroUsize::new(2).ok_or("two is not zero")?,
distance: Metric::Cosine,
},
}),
)?;
let worker = Arc::new(Poison::default());
let erased: Arc<dyn Worker> = worker.clone();
service.start_embedding_steps(erased, bounds()?, Restart::default())?;
Ok((directory, service, worker))
}
fn failed_field(service: &Service) -> Result<Value, Box<dyn Error>> {
let stats = service.dispatch(Operation::StoreStats(StoreRead {
selection: selection(),
}))?;
stats["embedding"]["fields"]
.as_array()
.ok_or("fields missing")?
.iter()
.find(|field| field["field"] == FIELD)
.cloned()
.ok_or_else(|| "field missing".into())
}
#[test]
fn a_poison_page_yields_and_only_retries_after_its_due_step() -> TestResult {
let (_directory, service, worker) = retry_service()?;
assert!(service.embedding_step_at(bounds()?, Instant::now())?);
let driver = service.embedding().ok_or("driver missing")?;
let due = driver
.retry_due("main", "notes", FIELD)
.ok_or("retry missing")?;
let before = due
.checked_sub(Duration::from_nanos(1))
.ok_or("clock underflow")?;
for _ in 0..5 {
assert!(service.embedding_step_at(bounds()?, before)?);
}
assert!(!service.embedding_step_at(bounds()?, before)?);
assert_eq!(worker.attempts.load(Ordering::SeqCst), 1);
let stats = driver.status();
let good = stats["fields"]
.as_array()
.ok_or("fields missing")?
.iter()
.find(|field| field["field"] == "z_good")
.ok_or("good field missing")?;
assert_eq!(good["made_since_start"], 5);
assert_eq!(good["walk"], "done");
let searchable = service.dispatch(Operation::VectorSearchText(VectorSearchText {
selection: selection(),
collection: "notes".to_owned(),
field: "z_good".to_owned(),
text: "banana".to_owned(),
k: NonZeroU32::new(5).ok_or("five is not zero")?,
}))?;
assert_eq!(searchable["compared"], 5);
assert_eq!(
searchable["hits"].as_array().ok_or("hits missing")?.len(),
5
);
let failed = failed_field(&service)?;
assert_eq!(failed["made_since_start"], 0);
assert_eq!(failed["failure"]["consecutive_failures"], 1);
assert_eq!(failed["failure"]["failure"]["code"], "embed_worker_exited");
assert_eq!(failed["failure"]["failure"]["message"], "poison page");
assert!(failed["failure"]["position"]["committed"].is_null());
assert!(failed["failure"]["position"]["toward"].is_object());
assert_eq!(failed["failure"]["position"]["end_after"], "b");
assert!(failed["failure"]["next_attempt_due"].is_u64());
let fields = service.dispatch(Operation::VectorFields(StoreRead {
selection: selection(),
}))?;
let exposed = fields["embedded_fields"]
.as_array()
.ok_or("embedded fields missing")?
.iter()
.find(|field| field["field"] == FIELD)
.ok_or("failed field missing")?;
assert_eq!(exposed["failure"], failed["failure"]);
assert!(service.embedding_step_at(bounds()?, due)?);
assert_eq!(worker.attempts.load(Ordering::SeqCst), 2);
let twice = failed_field(&service)?;
assert_eq!(twice["failure"]["consecutive_failures"], 2);
assert_eq!(twice["failure"]["position"], failed["failure"]["position"]);
let at = twice["failure"]["at"]
.as_u64()
.ok_or("failure time missing")?;
let next = twice["failure"]["next_attempt_due"]
.as_u64()
.ok_or("retry time missing")?;
assert_eq!(next - at, 2_000_000_000);
let next_due = driver
.retry_due("main", "notes", FIELD)
.ok_or("retry missing")?;
worker.succeeding.store(true, Ordering::SeqCst);
assert!(service.embedding_step_at(bounds()?, next_due)?);
assert!(service.embedding_step_at(bounds()?, next_due)?);
let committed = failed_field(&service)?;
assert_eq!(committed["made_since_start"], 2);
assert_eq!(committed["failure"]["consecutive_failures"], 0);
assert!(committed["failure"]["next_attempt_due"].is_null());
assert_eq!(committed["failure"]["failure"]["message"], "poison page");
assert!(driver.retry_due("main", "notes", FIELD).is_none());
Ok(())
}
#[test]
fn the_page_retry_queue_is_fair_across_fields_collections_and_branches() {
let mut state = super::super::State::default();
let failed = ("a".to_owned(), "a".to_owned(), "a".to_owned());
let others = [
("a".to_owned(), "a".to_owned(), "z".to_owned()),
("a".to_owned(), "z".to_owned(), "a".to_owned()),
("z".to_owned(), "a".to_owned(), "a".to_owned()),
];
state.waiting.insert(failed.clone());
for field in &others {
state.waiting.insert(field.clone());
}
assert_eq!(state.next_page(Instant::now()), Some(failed.clone()));
state.waiting.requeue(failed.clone());
for field in others {
assert_eq!(state.next_page(Instant::now()), Some(field));
}
assert_eq!(state.next_page(Instant::now()), Some(failed));
}
#[test]
fn page_failure_backoff_survives_ready_and_stops_growing_at_the_ceiling() -> TestResult {
let shared = super::super::Shared::default();
let field = ("main".to_owned(), "notes".to_owned(), FIELD.to_owned());
let restart = Restart {
first: Duration::from_millis(1),
most: Duration::from_millis(5),
ready: Duration::ZERO,
};
let worker = Poison::default();
for (count, millis) in [1_u32, 2, 3, 4, 5].into_iter().zip([1_u64, 2, 4, 5, 5]) {
shared.lock().failed_page(
&field,
None,
Failure::new("embed_worker_exited", "poison page"),
restart,
Instant::now(),
)?;
super::super::start_worker(&shared, &worker, restart);
let state = shared.lock();
let failed = state.failures.get(&field).ok_or("failure missing")?;
assert_eq!(failed.consecutive_failures, count);
assert_eq!(
failed.next_attempt_due.ok_or("due missing")? - failed.at,
millis * 1_000_000,
);
drop(state);
}
shared.lock().committed_page(&field);
assert_eq!(
shared
.lock()
.failures
.get(&field)
.ok_or("failure missing")?
.consecutive_failures,
0,
);
Ok(())
}