use super::Service;
use crate::serve::protocol::{Change, CollectionWrite, Mutation, Operation, Selection, Version};
use crate::serve::storage::Storage;
use haematite::{db::DatabaseConfig, store::NodeCacheBudget};
use std::num::NonZeroU32;
use std::sync::mpsc;
use std::time::Duration;
type TestResult = Result<(), Box<dyn std::error::Error>>;
#[test]
fn concurrent_writers_wait_and_both_succeed() -> TestResult {
let dir = tempfile::tempdir()?;
let storage = Storage::create(
DatabaseConfig {
data_dir: dir.path().join("db"),
shard_count: 1,
distributed: None,
executor_threads: Some(1),
node_cache_budget: Some(NodeCacheBudget::Unlimited),
retention: haematite::Retention::KeepAll,
},
NodeCacheBudget::Unlimited,
Duration::from_secs(3600),
)?;
let service = Service::new(storage, NonZeroU32::MIN)?;
let initial = service
.catalogue
.capture(&Selection {
branch: "main".into(),
version: None,
})?
.version;
let request = |id: &str, collection: &str, expected| {
Operation::Mutate(Mutation {
request_id: id.into(),
branch: "main".into(),
expected,
change: Change::Collection(CollectionWrite {
collection: collection.into(),
schema: None,
}),
})
};
let first_request = request("first-write", "first", initial);
let next = Version {
generation: initial.generation,
sequence: initial.sequence + 1,
};
let second_request = request("second-write", "second", next);
std::thread::scope(|scope| -> TestResult {
let (paused, pause) = mpsc::channel();
let (release, released) = mpsc::channel();
let service = &service;
let first = scope.spawn(move || {
Storage::arm_publish_pause(paused, released);
service.dispatch_waiting(first_request, Duration::from_secs(5))
});
pause.recv_timeout(Duration::from_secs(5))?;
let (send, receive) = mpsc::channel();
let second = scope.spawn(move || {
send.send(service.dispatch_waiting(second_request, Duration::from_secs(5)))
});
let waiting = {
let (active, _) = service
.writer_available
.changed
.wait_timeout_while(
service
.writer_available
.active
.lock()
.map_err(|_| "wait state poisoned")?,
Duration::from_secs(5),
|active| !*active,
)
.map_err(|_| "wait state poisoned")?;
let waiting = *active;
drop(active);
waiting
};
release.send(())?;
let one = first.join().map_err(|_| "first writer panicked")??;
second.join().map_err(|_| "second writer panicked")??;
let two = receive.recv_timeout(Duration::from_secs(5))?;
assert!(waiting, "second writer never entered its lock wait");
assert!(two.is_ok(), "second concurrent writer was refused: {two:?}");
let one: crate::serve::storage::Receipt = serde_json::from_value(one)?;
let two: crate::serve::storage::Receipt = serde_json::from_value(two?)?;
assert_eq!(one.version, next);
assert_eq!(two.version.sequence, initial.sequence + 2);
assert_eq!(two.version.generation, initial.generation);
Ok(())
})
}
#[test]
fn writer_wait_refuses_at_request_timeout_and_names_the_wait() -> TestResult {
let dir = tempfile::tempdir()?;
let storage = Storage::create(
DatabaseConfig {
data_dir: dir.path().join("db"),
shard_count: 1,
distributed: None,
executor_threads: Some(1),
node_cache_budget: Some(NodeCacheBudget::Unlimited),
retention: haematite::Retention::KeepAll,
},
NodeCacheBudget::Unlimited,
Duration::from_secs(3600),
)?;
let service = Service::new(storage, NonZeroU32::MIN)?;
let version = service
.catalogue
.capture(&Selection {
branch: "main".into(),
version: None,
})?
.version;
let held = service.storage.lock().map_err(|_| "writer poisoned")?;
let start = std::time::Instant::now();
let refused = service
.dispatch_waiting(
Operation::Mutate(Mutation {
request_id: "waiting".into(),
branch: "main".into(),
expected: version,
change: Change::Collection(CollectionWrite {
collection: "records".into(),
schema: None,
}),
}),
Duration::from_millis(50),
)
.err()
.ok_or("held writer was admitted")?;
assert!(
start.elapsed() < Duration::from_secs(2),
"writer wait exceeded its request bound"
);
assert_eq!(refused.code, "busy");
assert!(
refused.message.contains("wait") && refused.message.contains("timeout_ms=50"),
"refusal did not name the wait: {}",
refused.message
);
drop(held);
Ok(())
}
struct ReadyWorker;
impl crate::serve::embed::job::Embed for ReadyWorker {
fn documents(
&self,
_: &[String],
) -> Result<crate::serve::embed::Embedded, crate::serve::storage::Failure> {
Err(crate::serve::storage::Failure::new(
"unexpected_document_call",
"this fixture has no embedding documents",
))
}
}
impl crate::serve::embed::question::Ask for ReadyWorker {
fn question(
&self,
_: &str,
) -> Result<crate::serve::embed::Embedded, crate::serve::storage::Failure> {
Ok(crate::serve::embed::Embedded {
model: crate::serve::embed::ModelId {
name: "fixture".into(),
revision: "1".into(),
manifest_sha256: "0".repeat(64),
dimensions: 1,
distance: haem_frames::embedding::Distance::Cosine,
normalised: false,
},
answers: vec![crate::serve::embed::Answer::Vector(vec![1.0])],
})
}
}
#[test]
fn held_embedding_commit_obeys_the_foreground_request_deadline() -> TestResult {
use std::num::NonZeroUsize;
use std::sync::Arc;
let dir = tempfile::tempdir()?;
let storage = Storage::create(
DatabaseConfig {
data_dir: dir.path().join("db"),
shard_count: 1,
distributed: None,
executor_threads: Some(1),
node_cache_budget: Some(NodeCacheBudget::Unlimited),
retention: haematite::Retention::KeepAll,
},
NodeCacheBudget::Unlimited,
Duration::from_secs(30),
)?;
let mut service = Service::new(storage, NonZeroU32::MIN)?;
service.start_embedding_steps(
Arc::new(ReadyWorker),
crate::serve::embed::job::Bounds {
page_records: NonZeroUsize::MIN,
batch: NonZeroUsize::MIN,
input_bytes: NonZeroUsize::MIN,
call: Duration::from_secs(1),
},
crate::serve::embed::settings::Restart::default(),
)?;
let selection = Selection {
branch: "main".into(),
version: None,
};
let before = service.catalogue.capture(&selection)?.version;
let held = service
.embedding()
.ok_or("missing driver")?
.admission()
.background(&service.storage)?;
let request = Operation::Mutate(Mutation {
request_id: "embedding-blocked-write".into(),
branch: "main".into(),
expected: before,
change: Change::Collection(CollectionWrite {
collection: "records".into(),
schema: None,
}),
});
let refused = service
.dispatch_waiting(request, Duration::from_millis(50))
.err()
.ok_or("held writer admitted")?;
assert_eq!(
refused.code, "busy",
"embedding wait ignored the request deadline: {refused}"
);
assert!(refused.message.contains("timeout_ms=50"), "{refused}");
assert_eq!(service.catalogue.capture(&selection)?.version, before);
drop(held);
let request = Operation::Mutate(Mutation {
request_id: "embedding-blocked-write".into(),
branch: "main".into(),
expected: before,
change: Change::Collection(CollectionWrite {
collection: "records".into(),
schema: None,
}),
});
assert!(
service
.dispatch_waiting(request, Duration::from_secs(1))
.is_ok(),
"deadline left admission reserved"
);
Ok(())
}