use std::error::Error;
use std::num::{NonZeroU32, NonZeroU64, NonZeroUsize};
use std::sync::{Arc, Condvar, Mutex, MutexGuard, PoisonError};
use std::time::{Duration, Instant};
use haem_frames::embedding::{Distance, input_refusal};
use haematite::SyncCounts;
use haematite::db::DatabaseConfig;
use haematite::store::NodeCacheBudget;
use serde_json::Value;
use serde_json::value::RawValue;
use super::Worker;
use crate::serve::embed::job::{Bounds, Embed};
use crate::serve::embed::question::Ask;
use crate::serve::embed::settings::{PAGE_RECORDS_MOST, Restart, writer_wait};
use crate::serve::embed::{Answer, Embedded, Longer, ModelId, Options, WindowVector};
use crate::serve::protocol::{
Change, CollectionWrite, DocumentWrite, Metric, Mutation, Operation, Selection, StoreRead,
VectorDeclaration, VectorSearchText, Version,
};
use crate::serve::query::Setting;
use crate::serve::service::Service;
use crate::serve::storage::embedded::{Declaration, ModelName, VectorKind};
use crate::serve::storage::{Failure, Storage};
type TestResult = Result<(), Box<dyn Error>>;
const FIELD: &str = "body_vec";
const MANIFEST: &str = "0000000000000000000000000000000000000000000000000000000000005eed";
#[derive(Default)]
struct Calls {
made: usize,
holds: Vec<usize>,
held: Option<usize>,
fail_on: Option<usize>,
asked: usize,
question_holds: Vec<usize>,
questions_failing: usize,
serial: bool,
in_batch: bool,
sent: Vec<Vec<String>>,
}
#[derive(Default)]
struct StandIn {
calls: Mutex<Calls>,
changed: Condvar,
}
impl StandIn {
fn lock(&self) -> MutexGuard<'_, Calls> {
self.calls.lock().unwrap_or_else(PoisonError::into_inner)
}
fn hold(&self, call: usize) {
self.lock().holds.push(call);
}
fn fail(&self, call: usize) {
self.lock().fail_on = Some(call);
}
fn hold_question(&self, question: usize) {
self.lock().question_holds.push(question);
}
fn serial() -> Self {
let stand_in = Self::default();
stand_in.lock().serial = true;
stand_in
}
fn fail_questions(&self, count: usize) {
self.lock().questions_failing = count;
}
fn wait_held(&self) {
let mut calls = self.lock();
while calls.held.is_none() {
calls = self
.changed
.wait(calls)
.unwrap_or_else(PoisonError::into_inner);
}
drop(calls);
}
fn let_go(&self) {
let mut calls = self.lock();
if let Some(call) = calls.held.take() {
calls.holds.retain(|held| *held != call);
calls.question_holds.retain(|held| *held != call);
}
drop(calls);
self.changed.notify_all();
}
fn made(texts: &[String]) -> Embedded {
Embedded {
model: ModelId {
name: "stand-in".to_owned(),
revision: "1".to_owned(),
manifest_sha256: MANIFEST.to_owned(),
dimensions: 2,
distance: Distance::Cosine,
normalised: false,
},
answers: texts
.iter()
.map(|text| {
if text.starts_with("long") {
Answer::Longer(Longer {
tokens_total: 600,
tokens_limit: 512,
})
} else if text.starts_with("overcap") {
Answer::OverCap(Longer {
tokens_total: 60_000,
tokens_limit: 512,
})
} else if text.starts_with("windows") {
let length = u64::try_from(text.len()).unwrap_or(u64::MAX);
let middle = length / 2;
Answer::Windows(
[(0, middle + 1), (middle, length)]
.into_iter()
.zip([vec![1.0, 0.0], vec![0.0, 1.0]])
.map(|((start, end), vector)| WindowVector { start, end, vector })
.collect(),
)
} else if text.contains('a') {
Answer::Vector(vec![1.0, 0.0])
} else {
Answer::Vector(vec![0.0, 1.0])
}
})
.collect(),
}
}
}
impl Embed for StandIn {
fn documents(&self, texts: &[String]) -> Result<Embedded, Failure> {
let mut calls = self.lock();
calls.made = calls.made.saturating_add(1);
calls.sent.push(texts.to_vec());
calls.in_batch = true;
let call = calls.made;
if calls.holds.contains(&call) {
calls.held = Some(call);
self.changed.notify_all();
while calls.holds.contains(&call) {
calls = self
.changed
.wait(calls)
.unwrap_or_else(PoisonError::into_inner);
}
}
let fail = calls.fail_on == Some(call);
calls.in_batch = false;
drop(calls);
if fail {
return Err(Failure::new("embed_worker_exited", "scripted to die"));
}
if let Some(index) = texts
.iter()
.position(|text| input_refusal(text, u64::MAX).is_some())
{
return Err(Failure::new(
"embed_input_empty",
format!("text {index} cannot be embedded"),
));
}
Ok(Self::made(texts))
}
}
impl Ask for StandIn {
fn question(&self, text: &str) -> Result<Embedded, Failure> {
let mut calls = self.lock();
if calls.serial && calls.in_batch {
return Err(Failure::new("embed_busy", "a question behind a walk batch"));
}
calls.asked = calls.asked.saturating_add(1);
let question = calls.asked;
if calls.question_holds.contains(&question) {
calls.held = Some(question);
self.changed.notify_all();
while calls.question_holds.contains(&question) {
calls = self
.changed
.wait(calls)
.unwrap_or_else(PoisonError::into_inner);
}
}
let fail = calls.questions_failing > 0;
calls.questions_failing = calls.questions_failing.saturating_sub(1);
drop(calls);
if fail {
return Err(Failure::new("embed_worker_timeout", "scripted to time out"));
}
Ok(Self::made(&[text.to_owned()]))
}
}
fn bounds_of(page: usize, batch: usize) -> Result<Bounds, Box<dyn Error>> {
Ok(Bounds {
page_records: NonZeroUsize::new(page).ok_or("not zero")?,
batch: NonZeroUsize::new(batch).ok_or("not zero")?,
input_bytes: NonZeroUsize::new(4096).ok_or("not zero")?,
call: Duration::from_secs(1),
})
}
fn bounds() -> Result<Bounds, Box<dyn Error>> {
bounds_of(2, 2)
}
fn service_with(records: &[(&str, &str)]) -> Result<(tempfile::TempDir, Service), Box<dyn Error>> {
let directory = tempfile::tempdir()?;
let storage = Storage::create(
DatabaseConfig {
data_dir: directory.path().join("database"),
shard_count: 2,
distributed: None,
executor_threads: Some(1),
node_cache_budget: Some(NodeCacheBudget::Unlimited),
retention: haematite::Retention::reclaim(0),
},
NodeCacheBudget::Unlimited,
fixture_long_wait_interval(),
)?;
let mut service = Service::new(storage, NonZeroU32::new(100).ok_or("not zero")?)?;
service.embed_search_bound = Some(NonZeroU64::new(1_000).ok_or("not zero")?);
mutate(
&service,
Change::Collection(CollectionWrite {
collection: "notes".to_owned(),
schema: None,
}),
)?;
for (id, body) in records {
put(&service, id, body)?;
}
Ok((directory, service))
}
fn start(service: &mut Service, stand_in: &Arc<StandIn>) -> TestResult {
start_with(service, stand_in, bounds()?)
}
fn start_with(service: &mut Service, stand_in: &Arc<StandIn>, bounds: Bounds) -> TestResult {
let worker: Arc<dyn Worker> = Arc::<StandIn>::clone(stand_in);
service.start_embedding(worker, bounds, Duration::from_secs(1))?;
Ok(())
}
fn start_apart(service: &mut Service, walk: &Arc<StandIn>, questions: &Arc<StandIn>) -> TestResult {
let walk: Arc<dyn Worker> = Arc::<StandIn>::clone(walk);
let questions: Arc<dyn Worker> = Arc::<StandIn>::clone(questions);
let wait = Duration::from_secs(1);
service.start_embedding_apart(walk, Some(questions), bounds()?, wait, Restart::default())?;
service
.embedding()
.ok_or("no thread")?
.wait_question_started();
Ok(())
}
fn start_restarting(
service: &mut Service,
stand_in: &Arc<StandIn>,
restart: Restart,
) -> TestResult {
let worker: Arc<dyn Worker> = Arc::<StandIn>::clone(stand_in);
service.start_embedding_with(worker, bounds()?, Duration::from_secs(1), restart)?;
Ok(())
}
fn selection() -> Selection {
Selection {
branch: "main".to_owned(),
version: None,
}
}
fn head(service: &Service) -> Result<Version, Box<dyn Error>> {
let answer = service.dispatch(Operation::StoreStats(StoreRead {
selection: selection(),
}))?;
Ok(serde_json::from_value(
answer["selection"]["version"].clone(),
)?)
}
fn mutate(service: &Service, change: Change) -> TestResult {
let expected = head(service)?;
service.dispatch(Operation::Mutate(Mutation {
request_id: format!("embed-{}-{}", expected.generation, expected.sequence),
branch: "main".to_owned(),
expected,
change,
}))?;
Ok(())
}
fn put(service: &Service, id: &str, body: &str) -> TestResult {
mutate(
service,
Change::Put(DocumentWrite {
collection: "notes".to_owned(),
document_id: id.to_owned(),
document: RawValue::from_string(format!(r#"{{"body":"{body}"}}"#))?,
}),
)
}
fn declare(service: &Service) -> TestResult {
mutate(
service,
Change::VectorDeclare(VectorDeclaration {
collection: "notes".to_owned(),
field: FIELD.to_owned(),
declaration: Declaration {
kind: VectorKind::Dense,
source: "body".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,
},
}),
)
}
fn search(service: &Service, text: &str) -> Result<Value, Box<dyn Error>> {
Ok(
service.dispatch(Operation::VectorSearchText(VectorSearchText {
selection: selection(),
collection: "notes".to_owned(),
field: FIELD.to_owned(),
text: text.to_owned(),
k: NonZeroU32::new(1).ok_or("not zero")?,
}))?,
)
}
fn field_status(service: &Service) -> Result<Value, Box<dyn Error>> {
let driver = service.embedding().ok_or("no embedding thread")?;
Ok(driver.status()["fields"][0].clone())
}
fn wait_idle(service: &Service) -> TestResult {
service
.embedding()
.ok_or("no embedding thread")?
.wait_idle();
Ok(())
}
const FIVE: [(&str, &str); 5] = [
("a", "alpha"),
("b", "zulu"),
("c", "echo"),
("d", "kilo"),
("e", "mike"),
];
#[test]
fn declare_walks_every_existing_record_in_batches_and_ends() -> TestResult {
let (_directory, mut service) = service_with(&FIVE)?;
let stand_in = Arc::new(StandIn::default());
start(&mut service, &stand_in)?;
declare(&service)?;
wait_idle(&service)?;
let status = field_status(&service)?;
assert_eq!(status["made_since_start"], 5);
assert_eq!(status["walk"], "done");
assert_eq!(stand_in.lock().made, 3);
assert_eq!(service.embedding().ok_or("no thread")?.pages(), 3);
Ok(())
}
#[test]
fn a_later_write_is_embedded_without_a_second_declare() -> TestResult {
let (_directory, mut service) = service_with(&FIVE[..2])?;
let stand_in = Arc::new(StandIn::default());
start(&mut service, &stand_in)?;
declare(&service)?;
wait_idle(&service)?;
put(&service, "f", "papa")?;
wait_idle(&service)?;
let status = field_status(&service)?;
assert_eq!(status["made_since_start"], 3);
assert_eq!(status["walk"], "done");
let answer = search(&service, "papa")?;
assert_eq!(answer["coverage"]["state"], "complete");
Ok(())
}
#[test]
fn search_text_answers_the_nearest_record_and_names_coverage() -> TestResult {
let (_directory, mut service) = service_with(&FIVE)?;
let stand_in = Arc::new(StandIn::default());
start(&mut service, &stand_in)?;
declare(&service)?;
wait_idle(&service)?;
let answer = search(&service, "banana")?;
assert_eq!(answer["hits"][0]["document_id"], "a");
assert!(answer["hits"][0].get("span").is_none(), "{answer}");
assert_eq!(answer["coverage"]["state"], "complete");
assert_eq!(answer["coverage"]["walk"], "done");
Ok(())
}
#[test]
fn search_during_an_open_walk_says_it_is_open() -> TestResult {
let (_directory, mut service) = service_with(&FIVE)?;
let stand_in = Arc::new(StandIn::default());
stand_in.hold(2);
start(&mut service, &stand_in)?;
declare(&service)?;
stand_in.wait_held();
let answer = search(&service, "banana")?;
assert_eq!(answer["coverage"]["walk"], "open");
assert_eq!(answer["coverage"]["state"], "partial");
stand_in.let_go();
wait_idle(&service)?;
Ok(())
}
#[test]
fn a_worker_killed_mid_walk_is_started_again_and_the_walk_ends() -> TestResult {
let (_directory, mut service) = service_with(&FIVE)?;
let stand_in = Arc::new(StandIn::default());
stand_in.fail(2);
start(&mut service, &stand_in)?;
declare(&service)?;
wait_idle(&service)?;
let driver = service.embedding().ok_or("no thread")?;
let failed = driver.last_failure().ok_or("no failure kept")?;
assert_eq!(failed.failure.code, "embed_worker_exited");
let status = field_status(&service)?;
assert_eq!(status["made_since_start"], 5);
assert_eq!(status["walk"], "done");
assert_eq!(driver.pages(), 3);
assert_eq!(stand_in.lock().made, 4);
assert_eq!(driver.status()["worker"], "ready");
assert_eq!(driver.status()["starts"], 2);
Ok(())
}
#[test]
fn a_search_while_the_worker_is_down_is_answered_after_the_restart() -> TestResult {
let (_directory, mut service) = service_with(&FIVE)?;
let stand_in = Arc::new(StandIn::default());
stand_in.fail(2);
stand_in.hold_question(2);
let restart = Restart {
ready: Duration::from_secs(60),
..Restart::default()
};
start_restarting(&mut service, &stand_in, restart)?;
declare(&service)?;
stand_in.wait_held();
let driver = service.embedding().ok_or("no thread")?;
assert_eq!(driver.status()["worker"], "restarting");
let answer = std::thread::scope(|scope| {
let asking = scope.spawn(|| search(&service, "banana").map_err(|error| error.to_string()));
driver.wait_askers(1);
stand_in.let_go();
asking
.join()
.map_err(|_panic| "the search panicked".to_owned())
})??;
assert!(answer["hits"].is_array(), "{answer}");
wait_idle(&service)?;
assert_eq!(field_status(&service)?["made_since_start"], 5);
Ok(())
}
#[test]
fn repeated_deaths_wait_longer_each_time_up_to_the_ceiling() -> TestResult {
let (_directory, mut service) = service_with(&FIVE[..1])?;
let stand_in = Arc::new(StandIn::default());
stand_in.fail_questions(4);
let restart = Restart {
first: Duration::from_millis(1),
most: Duration::from_millis(3),
ready: Duration::ZERO,
};
start_restarting(&mut service, &stand_in, restart)?;
wait_idle(&service)?;
let driver = service.embedding().ok_or("no thread")?;
let waits: Vec<u128> = driver.waits().iter().map(Duration::as_millis).collect();
assert_eq!(waits, [0, 1, 2, 3]);
let status = driver.status();
assert_eq!(status["starts"], 5);
assert_eq!(status["worker"], "ready");
assert_eq!(status["next_start_due"], Value::Null);
assert_eq!(
status["last_failure"]["failure"]["code"],
"embed_worker_timeout"
);
Ok(())
}
#[test]
fn the_restart_wait_doubles_from_the_second_start_to_its_ceiling() {
let restart = Restart {
first: Duration::from_secs(1),
most: Duration::from_secs(5),
ready: Duration::ZERO,
};
let waits: Vec<u64> = (1..=6)
.map(|deaths| restart.wait(deaths).as_secs())
.collect();
assert_eq!(waits, [0, 1, 2, 4, 5, 5]);
assert_eq!(restart.wait(u32::MAX), Duration::from_secs(5));
}
#[test]
fn stats_count_the_starts_and_keep_the_last_failure() -> TestResult {
let (_directory, mut service) = service_with(&FIVE)?;
let stand_in = Arc::new(StandIn::default());
stand_in.fail(2);
stand_in.hold_question(2);
start(&mut service, &stand_in)?;
declare(&service)?;
stand_in.wait_held();
let driver = service.embedding().ok_or("no thread")?;
let held = driver.status();
assert_eq!(held["starts"], 1);
assert!(held["next_start_due"].is_u64(), "{held}");
assert_eq!(
held["last_failure"]["failure"]["code"],
"embed_worker_exited"
);
stand_in.let_go();
wait_idle(&service)?;
let status = driver.status();
assert_eq!(status["starts"], 2);
assert_eq!(status["next_start_due"], Value::Null);
assert_eq!(
status["last_failure"]["failure"]["code"],
"embed_worker_exited"
);
Ok(())
}
#[test]
fn vectors_fields_says_a_served_field_is_embedded() -> TestResult {
let (_directory, mut service) = service_with(&FIVE)?;
let stand_in = Arc::new(StandIn::default());
start(&mut service, &stand_in)?;
declare(&service)?;
wait_idle(&service)?;
let answer = service.dispatch(Operation::VectorFields(StoreRead {
selection: selection(),
}))?;
assert_eq!(answer["embedded_fields"][0]["status"], "embedded");
Ok(())
}
#[test]
fn the_storage_lock_is_free_during_a_worker_call() -> TestResult {
let (_directory, mut service) = service_with(&FIVE)?;
let stand_in = Arc::new(StandIn::default());
stand_in.hold(1);
start(&mut service, &stand_in)?;
declare(&service)?;
stand_in.wait_held();
put(&service, "f", "papa")?;
stand_in.let_go();
wait_idle(&service)?;
assert_eq!(field_status(&service)?["walk"], "done");
Ok(())
}
#[test]
fn one_durable_operation_per_batch() -> TestResult {
let (_directory, mut service) = service_with(&FIVE)?;
let stand_in = Arc::new(StandIn::default());
start(&mut service, &stand_in)?;
declare(&service)?;
wait_idle(&service)?;
let driver = service.embedding().ok_or("no thread")?;
let (before, pages) = (driver.durable(), driver.pages());
put(&service, "f", "papa")?;
wait_idle(&service)?;
let driver = service.embedding().ok_or("no thread")?;
assert_eq!(driver.pages(), pages + 1);
assert_eq!(
driver.durable().since(before),
SyncCounts {
record_installs: 1,
barrier_device_flushes: 1,
ref_file_full_flushes: 1,
directory_full_flushes: 1,
}
);
Ok(())
}
#[test]
fn shutdown_during_a_batch_joins_when_the_call_ends() -> TestResult {
let (directory, mut service) = service_with(&FIVE)?;
let stand_in = Arc::new(StandIn::default());
stand_in.hold(1);
start(&mut service, &stand_in)?;
declare(&service)?;
stand_in.wait_held();
let dropping = std::thread::spawn(move || drop(service));
stand_in.let_go();
dropping.join().map_err(|_| "the drop panicked")?;
assert_eq!(stand_in.lock().made, 1);
drop(directory);
Ok(())
}
const fn each(count: u64) -> SyncCounts {
SyncCounts {
record_installs: count,
barrier_device_flushes: count,
ref_file_full_flushes: count,
directory_full_flushes: count,
}
}
#[test]
fn a_page_of_six_in_batches_of_two_is_three_calls_and_one_commit() -> TestResult {
let (_directory, mut service) = service_with(&FIVE[..1])?;
let stand_in = Arc::new(StandIn::default());
stand_in.hold(2);
start_with(&mut service, &stand_in, bounds_of(6, 2)?)?;
declare(&service)?;
wait_idle(&service)?;
put(&service, "f", "foxtrot")?;
stand_in.wait_held();
let driver = service.embedding().ok_or("no thread")?;
let (before, pages) = (driver.durable(), driver.pages());
for (id, body) in [("g", "golf"), ("h", "hotel"), ("i", "india")] {
put(&service, id, body)?;
}
for (id, body) in [("j", "juliet"), ("k", "kilo"), ("l", "lima")] {
put(&service, id, body)?;
}
stand_in.let_go();
wait_idle(&service)?;
let driver = service.embedding().ok_or("no thread")?;
assert_eq!(driver.pages(), pages + 2);
assert_eq!(driver.durable().since(before), each(2));
assert_eq!(stand_in.lock().made, 5);
assert_eq!(field_status(&service)?["made_since_start"], 8);
Ok(())
}
#[test]
fn a_failure_in_batch_two_of_three_commits_nothing() -> TestResult {
let (_directory, mut service) = service_with(&FIVE)?;
let stand_in = Arc::new(StandIn::default());
stand_in.fail(2);
stand_in.hold_question(2);
start_with(&mut service, &stand_in, bounds_of(5, 2)?)?;
declare(&service)?;
stand_in.wait_held();
let driver = service.embedding().ok_or("no thread")?;
let failed = driver.last_failure().ok_or("no failure kept")?;
assert_eq!(failed.failure.code, "embed_worker_exited");
assert_eq!(driver.pages(), 0);
assert_eq!(driver.durable(), each(0));
assert_eq!(stand_in.lock().made, 2);
assert_eq!(field_status(&service)?["made_since_start"], 0);
let refused = search(&service, "banana").err().ok_or("search answered")?;
assert!(
refused.to_string().contains("embed_worker_restarting"),
"{refused}"
);
stand_in.let_go();
wait_idle(&service)?;
assert_eq!(driver.pages(), 1);
assert_eq!(stand_in.lock().made, 5);
assert_eq!(field_status(&service)?["made_since_start"], 5);
Ok(())
}
#[test]
fn a_write_between_read_and_commit_leaves_that_record_out() -> TestResult {
let (_directory, mut service) = service_with(&[("a", "zulu"), ("b", "kilo")])?;
let stand_in = Arc::new(StandIn::default());
stand_in.hold(1);
stand_in.hold(2);
start_with(&mut service, &stand_in, bounds_of(4, 2)?)?;
declare(&service)?;
stand_in.wait_held();
put(&service, "a", "alpha")?;
stand_in.let_go();
stand_in.wait_held();
let answer = search(&service, "banana")?;
assert_eq!(answer["compared"], 1);
assert_eq!(answer["stale"], 0);
assert_eq!(answer["hits"][0]["document_id"], "b");
stand_in.let_go();
wait_idle(&service)?;
let answer = search(&service, "banana")?;
assert_eq!(answer["hits"][0]["document_id"], "a");
Ok(())
}
fn options(page_records: Setting<u64>) -> Options {
Options {
worker: Setting::Present("/stand-in/worker".into()),
model_dir: Setting::Present("/stand-in/model".into()),
model_sha256: Setting::Present("abc".to_owned()),
memory_bytes: Setting::Present(1 << 30),
threads: Setting::Present(1),
batch_items: Setting::Present(4),
input_bytes: Setting::Present(16_384),
tokens: Setting::Present(512),
frame_bytes: Setting::Present(1 << 20),
ready_ms: Setting::Present(1_000),
call_ms: Setting::Present(1_000),
waiting: Setting::Present(4),
request_items: Setting::Present(8),
page_records,
writer_wait_ms: Setting::Present(1_000),
windows: Setting::Absent,
max_windows: Setting::Absent,
overlap_tokens: Setting::Absent,
}
}
#[test]
fn bounds_page_by_embed_page_records_in_calls_of_embed_request_items() -> TestResult {
let settings = options(Setting::Present(256))
.resolve()?
.ok_or("no embed settings")?;
let bounds = super::bounds(&settings);
assert_eq!(bounds.page_records.get(), 256);
assert_eq!(bounds.batch.get(), 8);
Ok(())
}
#[test]
fn embed_page_records_is_refused_at_zero_and_past_its_ceiling() {
let most = u64::try_from(PAGE_RECORDS_MOST).unwrap_or(u64::MAX);
for refused in [
Setting::Absent,
Setting::Present(0),
Setting::Present(most + 1),
] {
let answer = options(refused).resolve();
assert!(
answer
.as_ref()
.is_err_and(|error| error.contains("embed_page_records")),
"{answer:?}"
);
}
let taken = options(Setting::Present(most)).resolve();
assert!(
taken.as_ref().is_ok_and(|settings| settings
.as_ref()
.is_some_and(|settings| settings.page_records.get() == PAGE_RECORDS_MOST)),
"{taken:?}"
);
}
#[test]
fn foreground_is_admitted_after_a_held_embedding_read() -> TestResult {
let storage = Mutex::new(0);
let admission = super::Admission::new(Duration::from_secs(1))?;
let held = admission.background(&storage)?;
std::thread::scope(|scope| -> TestResult {
let writer = scope.spawn(|| -> Result<(), Failure> {
let mut writer = admission.foreground(&storage, false)?;
*writer = 1;
drop(writer);
Ok(())
});
admission.wait_foreground()?;
drop(held);
writer.join().map_err(|_| "writer thread panicked")??;
assert_eq!(*storage.lock().map_err(|_| "storage poisoned")?, 1);
Ok(())
})
}
#[test]
fn embedding_resumes_after_the_foreground_writer() -> TestResult {
let storage = Mutex::new(0);
let admission = super::Admission::new(Duration::from_secs(1))?;
let foreground = admission.foreground(&storage, false)?;
std::thread::scope(|scope| -> TestResult {
let walk = scope.spawn(|| -> Result<(), Failure> {
let mut walk = admission.background(&storage)?;
*walk += 1;
drop(walk);
Ok(())
});
drop(foreground);
walk.join().map_err(|_| "walk thread panicked")??;
assert_eq!(*storage.lock().map_err(|_| "storage poisoned")?, 1);
Ok(())
})
}
#[test]
fn foreground_against_foreground_still_refuses_busy() -> TestResult {
let storage = Mutex::new(0);
let admission = super::Admission::new(Duration::from_secs(1))?;
let first = admission.foreground(&storage, false)?;
let second = admission.foreground(&storage, false);
assert!(second.as_ref().is_err_and(|failure| failure.code == "busy"));
drop(first);
assert!(admission.foreground(&storage, false).is_ok());
Ok(())
}
#[test]
fn embedding_writer_wait_defaults_and_rejects_zero() -> TestResult {
assert_eq!(
writer_wait(Setting::Absent)?,
Duration::from_millis(crate::serve::embed::settings::WRITER_WAIT_DEFAULT_MS,)
);
assert!(
writer_wait(Setting::Present(0))
.as_ref()
.is_err_and(|failure| failure.contains("embed_writer_wait_ms"))
);
assert_eq!(
writer_wait(Setting::Present(17))?,
Duration::from_millis(17)
);
Ok(())
}
#[test]
fn a_waiting_foreground_writer_prevents_the_next_embedding_commit() -> TestResult {
use std::sync::atomic::{AtomicUsize, Ordering};
let storage = Mutex::new(0);
let admission = super::Admission::new(Duration::from_secs(1))?;
let commits = AtomicUsize::new(0);
let first = admission.background(&storage)?;
std::thread::scope(|scope| -> TestResult {
let writer = scope.spawn(|| -> Result<(), Failure> {
let mut writer = admission.foreground(&storage, false)?;
assert_eq!(commits.load(Ordering::SeqCst), 0);
*writer = 1;
drop(writer);
Ok(())
});
admission.wait_foreground()?;
let next = scope.spawn(|| -> Result<(), Failure> {
let storage = admission.background(&storage)?;
assert_eq!(*storage, 1);
drop(storage);
commits.fetch_add(1, Ordering::SeqCst);
Ok(())
});
admission.wait_background()?;
drop(first);
writer.join().map_err(|_| "writer thread panicked")??;
next.join().map_err(|_| "embedding thread panicked")??;
assert_eq!(commits.load(Ordering::SeqCst), 1);
Ok(())
})
}
#[test]
fn embedding_starts_from_configuration_without_writer_wait() -> TestResult {
let (directory, mut service) = service_with(&[("a", "alpha")])?;
let config = serde_json::json!({
"data_dir": directory.path().join("database"),
"socket": directory.path().join("haematite.sock"),
"branch_cache_bytes_per_shard": 1_048_576,
"max_frame": 65_536,
"max_clients": 1,
"max_page": 8,
"timeout_ms": 1000,
"embed_worker": "/stand-in/worker",
"embed_model_dir": "/stand-in/model",
"embed_model_sha256": MANIFEST,
"embed_memory_bytes": 1 << 30,
"embed_threads": 1,
"embed_batch_items": 2,
"embed_input_bytes": 4096,
"embed_tokens": 512,
"embed_frame_bytes": 1 << 20,
"embed_ready_ms": 1000,
"embed_call_ms": 1000,
"embed_waiting": 4,
"embed_request_items": 2,
"embed_page_records": 2
});
assert!(config.get("embed_writer_wait_ms").is_none());
let config: crate::serve::cli::Config = serde_json::from_value(config)?;
assert!(!config.embed_writer_wait_ms.is_set());
let settings = config
.embed_settings()?
.ok_or("embedding settings absent")?;
assert_eq!(
settings.writer_wait,
Duration::from_millis(crate::serve::embed::settings::WRITER_WAIT_DEFAULT_MS,)
);
let stand_in = Arc::new(StandIn::default());
let worker: Arc<dyn Worker> = Arc::<StandIn>::clone(&stand_in);
service.start_embedding(worker, super::bounds(&settings), settings.writer_wait)?;
declare(&service)?;
wait_idle(&service)?;
assert_eq!(field_status(&service)?["made_since_start"], 1);
assert_eq!(field_status(&service)?["walk"], "done");
assert_eq!(stand_in.lock().made, 1);
Ok(())
}
#[test]
fn a_search_is_answered_while_a_walk_batch_is_held() -> TestResult {
let (_directory, mut service) = service_with(&FIVE)?;
let walk = Arc::new(StandIn::serial());
let questions = Arc::new(StandIn::default());
walk.hold(2);
start_apart(&mut service, &walk, &questions)?;
declare(&service)?;
walk.wait_held();
let answer = search(&service, "banana")?;
assert_eq!(answer["hits"][0]["document_id"], "a");
assert_eq!(answer["coverage"]["walk"], "open");
assert_eq!(walk.lock().asked, 1);
assert_eq!(questions.lock().asked, 2);
walk.let_go();
wait_idle(&service)?;
Ok(())
}
#[test]
fn the_question_worker_is_started_again_after_a_death() -> TestResult {
let (_directory, mut service) = service_with(&FIVE)?;
let walk = Arc::new(StandIn::default());
let questions = Arc::new(StandIn::default());
questions.fail_questions(1);
start_apart(&mut service, &walk, &questions)?;
declare(&service)?;
wait_idle(&service)?;
let status = service.embedding().ok_or("no thread")?.status();
assert_eq!(status["workers"]["question"]["state"], "restarting");
let answer = search(&service, "banana")?;
assert_eq!(answer["hits"][0]["document_id"], "a");
let status = service.embedding().ok_or("no thread")?.status();
let question = &status["workers"]["question"];
assert_eq!(question["state"], "ready");
assert_eq!(question["starts"], 2);
assert_eq!(
question["last_failure"]["failure"]["code"],
"embed_worker_timeout"
);
Ok(())
}
#[test]
fn stats_name_both_workers() -> TestResult {
let (_directory, mut service) = service_with(&FIVE)?;
let walk = Arc::new(StandIn::default());
let questions = Arc::new(StandIn::default());
start_apart(&mut service, &walk, &questions)?;
declare(&service)?;
wait_idle(&service)?;
search(&service, "banana")?;
let status = service.embedding().ok_or("no thread")?.status();
assert_eq!(status["workers"]["walk"], "ready");
assert_eq!(status["workers"]["question"]["state"], "ready");
assert_eq!(status["workers"]["question"]["starts"], 1);
let (_other, mut alone) = service_with(&FIVE[..1])?;
start(&mut alone, &Arc::new(StandIn::default()))?;
wait_idle(&alone)?;
let status = alone.embedding().ok_or("no thread")?.status();
assert_eq!(status["workers"]["question"], "off");
Ok(())
}
#[test]
fn the_question_worker_starts_with_the_service() -> TestResult {
let (_directory, mut service) = service_with(&FIVE)?;
let walk = Arc::new(StandIn::default());
let questions = Arc::new(StandIn::default());
start_apart(&mut service, &walk, &questions)?;
assert_eq!(questions.lock().asked, 1);
declare(&service)?;
wait_idle(&service)?;
search(&service, "banana")?;
assert_eq!(questions.lock().asked, 2);
let status = service.embedding().ok_or("no thread")?.status();
assert_eq!(status["workers"]["question"]["starts"], 1);
Ok(())
}
#[test]
fn stats_show_the_question_worker_starting_then_ready() -> TestResult {
let (_directory, mut service) = service_with(&FIVE[..1])?;
let walk: Arc<dyn Worker> = Arc::new(StandIn::default());
let questions = Arc::new(StandIn::default());
questions.hold_question(1);
let asked: Arc<dyn Worker> = Arc::<StandIn>::clone(&questions);
let wait = Duration::from_secs(1);
service.start_embedding_apart(walk, Some(asked), bounds()?, wait, Restart::default())?;
questions.wait_held();
let driver = service.embedding().ok_or("no thread")?;
assert_eq!(driver.status()["workers"]["question"]["state"], "starting");
questions.let_go();
driver.wait_question_started();
assert_eq!(driver.status()["workers"]["question"]["state"], "ready");
Ok(())
}
#[test]
fn a_thread_that_will_not_end_is_named_within_the_stop_bound() -> TestResult {
let (_directory, mut service) = service_with(&FIVE)?;
let stand_in = Arc::new(StandIn::default());
stand_in.hold(1);
let mut bounds = bounds()?;
bounds.call = Duration::from_millis(50);
start_with(&mut service, &stand_in, bounds)?;
declare(&service)?;
stand_in.wait_held();
let driver = service.embedding().ok_or("no thread")?;
let refused = driver.stop().err().ok_or("a held thread was waited for")?;
assert_eq!(refused.code, "embed_thread_unended");
assert!(
refused.message.contains("haematite-embed did not end"),
"{}",
refused.message
);
assert!(refused.message.contains("1.05s"), "{}", refused.message);
stand_in.let_go();
driver.wait_walk_ended();
assert!(driver.stop().is_ok());
Ok(())
}
#[test]
fn a_held_call_let_go_after_stop_gave_up_makes_no_commit() -> TestResult {
let (_directory, mut service) = service_with(&FIVE)?;
let stand_in = Arc::new(StandIn::default());
stand_in.hold(1);
let mut bounds = bounds()?;
bounds.call = Duration::from_millis(50);
start_with(&mut service, &stand_in, bounds)?;
declare(&service)?;
stand_in.wait_held();
let driver = service.embedding().ok_or("no thread")?;
let refused = driver.stop().err().ok_or("a held thread was waited for")?;
assert!(
refused.message.contains("it starts no commit"),
"{}",
refused.message
);
stand_in.let_go();
driver.wait_walk_ended();
assert_eq!(driver.pages(), 0);
assert_eq!(driver.durable(), each(0));
assert_eq!(field_status(&service)?["made_since_start"], 0);
Ok(())
}
#[test]
fn a_longer_text_is_counted_and_the_rest_of_its_batch_made() -> TestResult {
let records = [("a", "alpha"), ("b", "long text"), ("c", "echo")];
let (_directory, mut service) = service_with(&records)?;
let stand_in = Arc::new(StandIn::default());
start_with(&mut service, &stand_in, bounds_of(3, 3)?)?;
declare(&service)?;
wait_idle(&service)?;
let status = field_status(&service)?;
assert_eq!(status["made_since_start"], 2);
assert_eq!(status["longer_than_model"], 1);
assert_eq!(status["walk"], "done");
assert_eq!(stand_in.lock().sent, [["alpha", "long text", "echo"]]);
assert_eq!(service.embedding().ok_or("no thread")?.pages(), 1);
let answer = service.dispatch(Operation::VectorFields(StoreRead {
selection: selection(),
}))?;
assert_eq!(answer["embedded_fields"][0]["longer_than_model"], 1);
assert_eq!(answer["embedded_fields"][0]["status"], "embedded");
Ok(())
}
#[test]
fn a_longer_text_commits_once_and_is_not_sent_again() -> TestResult {
let (_directory, mut service) = service_with(&FIVE)?;
let stand_in = Arc::new(StandIn::default());
start(&mut service, &stand_in)?;
declare(&service)?;
wait_idle(&service)?;
let driver = service.embedding().ok_or("no thread")?;
let (before, pages) = (driver.durable(), driver.pages());
put(&service, "f", "long text")?;
wait_idle(&service)?;
let driver = service.embedding().ok_or("no thread")?;
assert_eq!(driver.pages(), pages + 1);
assert_eq!(driver.durable().since(before), each(1));
put(&service, "g", "golf")?;
wait_idle(&service)?;
let sent = stand_in.lock().sent.clone();
let last_two = sent
.get(sent.len().saturating_sub(2)..)
.ok_or("fewer than two calls")?;
assert_eq!(last_two, [["long text"], ["golf"]]);
let status = field_status(&service)?;
assert_eq!(status["made_since_start"], 6);
assert_eq!(status["longer_than_model"], 1);
Ok(())
}
#[test]
fn a_walk_of_only_longer_texts_makes_nothing_and_ends() -> TestResult {
let (_directory, mut service) = service_with(&[("a", "long a"), ("b", "long b")])?;
let stand_in = Arc::new(StandIn::default());
start(&mut service, &stand_in)?;
declare(&service)?;
wait_idle(&service)?;
let status = field_status(&service)?;
assert_eq!(status["made_since_start"], 0);
assert_eq!(status["longer_than_model"], 2);
assert_eq!(status["walk"], "done");
assert_eq!(stand_in.lock().made, 1);
Ok(())
}
fn blanks_are_passed_and_the_rest_made(blank: &str) -> TestResult {
let records = [
("a", blank),
("b", "bravo"),
("c", blank),
("d", "delta"),
("e", blank),
];
let (_directory, mut service) = service_with(&records)?;
declare(&service)?;
let stand_in = Arc::new(StandIn::default());
let worker: Arc<dyn Worker> = Arc::<StandIn>::clone(&stand_in);
let bounds = bounds_of(5, 5)?;
service.start_embedding_steps(worker, bounds, Restart::default())?;
let mut steps = 0;
while service.embedding_step_at(bounds, Instant::now())? {
steps += 1;
assert!(steps < 10, "the walk did not settle in 10 steps");
}
let status = field_status(&service)?;
assert_eq!(status["made_since_start"], 2, "{status}");
assert_eq!(status["embed_input_empty"], 3);
assert_eq!(status["walk"], "done");
assert_eq!(stand_in.lock().sent, [["bravo", "delta"]]);
let store = service.dispatch(Operation::StoreStats(StoreRead {
selection: selection(),
}))?;
assert_eq!(store["embedding"]["fields"][0]["embed_input_empty"], 3);
let fields = service.dispatch(Operation::VectorFields(StoreRead {
selection: selection(),
}))?;
assert_eq!(fields["embedded_fields"][0]["embed_input_empty"], 3);
Ok(())
}
#[test]
fn windows_and_over_cap_are_counted_per_field() -> TestResult {
let records = [
("a", "windows of text"),
("b", "bravo"),
("c", "overcap text"),
];
let (_directory, mut service) = service_with(&records)?;
declare(&service)?;
let worker: Arc<dyn Worker> = Arc::new(StandIn::default());
let bounds = bounds_of(5, 5)?;
service.start_embedding_steps(worker, bounds, Restart::default())?;
let mut steps = 0;
while service.embedding_step_at(bounds, Instant::now())? {
steps += 1;
assert!(steps < 10, "the walk did not settle in 10 steps");
}
let status = field_status(&service)?;
assert_eq!(status["made_since_start"], 1, "{status}");
assert_eq!(status["windowed"], 1);
assert_eq!(status["windows_over_cap"], 1);
assert_eq!(status["longer_than_model"], 1);
let store = service.dispatch(Operation::StoreStats(StoreRead {
selection: selection(),
}))?;
assert_eq!(store["embedding"]["fields"][0]["windowed"], 1);
let fields = service.dispatch(Operation::VectorFields(StoreRead {
selection: selection(),
}))?;
assert_eq!(fields["embedded_fields"][0]["windowed"], 1);
assert_eq!(fields["embedded_fields"][0]["windows_over_cap"], 1);
Ok(())
}
#[test]
fn a_windowed_hit_answers_its_nearest_window_span() -> TestResult {
let (_directory, mut service) = service_with(&[("a", "windows of text")])?;
declare(&service)?;
let worker: Arc<dyn Worker> = Arc::new(StandIn::default());
let bounds = bounds_of(5, 5)?;
service.start_embedding_steps(worker, bounds, Restart::default())?;
let mut steps = 0;
while service.embedding_step_at(bounds, Instant::now())? {
steps += 1;
assert!(steps < 10, "the walk did not settle in 10 steps");
}
for (question, start, end) in [("banana", 0, 8), ("zulu", 7, 15)] {
let answer = search(&service, question)?;
let hit = &answer["hits"][0];
assert_eq!(hit["document_id"], "a", "{answer}");
assert_eq!(hit["span"]["start"], start, "{answer}");
assert_eq!(hit["span"]["end"], end, "{answer}");
}
Ok(())
}
#[test]
fn an_empty_text_first_between_and_last_is_passed_and_the_rest_made() -> TestResult {
blanks_are_passed_and_the_rest_made("")
}
#[test]
fn a_whitespace_text_first_between_and_last_is_passed_and_the_rest_made() -> TestResult {
blanks_are_passed_and_the_rest_made(r" \t\n ")
}
#[path = "embedding_retry_tests.rs"]
mod retry;
const fn fixture_long_wait_interval() -> std::time::Duration {
std::time::Duration::from_secs(1)
}