use super::{Limits, Listener};
use crate::serve::{frame, service::Service, storage::Storage};
use haematite::{db::DatabaseConfig, store::NodeCacheBudget};
use std::io;
use std::num::{NonZeroU32, NonZeroUsize};
use std::os::unix::net::UnixStream;
use std::sync::Arc;
use std::time::Duration;
type TestResult = Result<(), Box<dyn std::error::Error>>;
fn service(path: &std::path::Path) -> Result<Arc<Service>, Box<dyn std::error::Error>> {
let storage = Storage::create(
DatabaseConfig {
data_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),
)?;
Ok(Arc::new(Service::new(storage, NonZeroU32::MIN)?))
}
fn limits() -> Limits {
Limits {
max_frame: NonZeroU32::MIN.saturating_add(65535),
max_clients: NonZeroUsize::MIN,
frame_timeout: Duration::from_secs(1),
}
}
#[test]
fn accept_errors_do_not_stop_the_next_client() -> TestResult {
let dir = tempfile::tempdir()?;
let service = service(dir.path())?;
for error in [
io::Error::from_raw_os_error(24),
io::Error::from_raw_os_error(23),
io::Error::from(io::ErrorKind::ConnectionAborted),
] {
let (_client, server) = UnixStream::pair()?;
let mut accepted = 0;
let result = Listener::accepting(
[Err(error), Ok(server)].into_iter(),
&service,
limits(),
|_, _, _, _| {
accepted += 1;
Ok(())
},
);
assert!(
result.is_ok(),
"accept error stopped the service: {result:?}"
);
assert_eq!(accepted, 1);
}
Ok(())
}
#[test]
fn spawn_failure_refuses_one_client_and_keeps_accepting() -> TestResult {
let dir = tempfile::tempdir()?;
let service = service(dir.path())?;
let (mut client, server) = UnixStream::pair()?;
let (_next, next_server) = UnixStream::pair()?;
client.set_read_timeout(Some(Duration::from_secs(2)))?;
let mut accepted = 0;
let result = Listener::accepting(
[Ok(server), Ok(next_server)].into_iter(),
&service,
limits(),
|_, _, _, _| {
accepted += 1;
if accepted == 1 {
Err(io::Error::from(io::ErrorKind::WouldBlock))
} else {
Ok(())
}
},
);
assert!(
result.is_ok(),
"spawn failure stopped the service: {result:?}"
);
assert_eq!(accepted, 2);
let answer =
frame::read(&mut client, limits().max_frame)?.ok_or("spawn refusal had no frame")?;
let answer: serde_json::Value = serde_json::from_slice(&answer)?;
assert_eq!(answer["error"]["code"], "client_spawn_failed");
Ok(())
}
#[test]
fn client_bound_sends_a_named_refusal_with_count_and_bound() -> TestResult {
let dir = tempfile::tempdir()?;
let service = service(dir.path())?;
let (_first, server) = UnixStream::pair()?;
let (mut client, excess) = UnixStream::pair()?;
client.set_read_timeout(Some(Duration::from_secs(2)))?;
let hello = serde_json::json!({"id":"overload-hello","operation":{"method":"hello","params":crate::serve::protocol::HelloAgreement::client()}});
frame::write(
&mut client,
&serde_json::to_vec(&hello)?,
limits().max_frame,
)?;
let mut held = None;
Listener::accepting(
[Ok(server), Ok(excess)].into_iter(),
&service,
limits(),
|_, _, _, permit| {
held = Some(permit);
Ok(())
},
)?;
let bytes = frame::read(&mut client, limits().max_frame)?
.ok_or("client over max_clients was dropped without a refusal frame")?;
let answer: serde_json::Value = serde_json::from_slice(&bytes)?;
assert_eq!(
answer["id"], "overload-hello",
"the client cannot match the refusal to its hello"
);
assert_eq!(answer["error"]["code"], "busy_max_clients");
let message = answer["error"]["message"]
.as_str()
.ok_or("missing message")?;
assert!(
message.contains("max_clients=1") && message.contains("active_clients=1"),
"{message}"
);
drop(held);
Ok(())
}
#[test]
fn disconnected_follower_releases_slot_without_a_branch_move() -> TestResult {
use crate::serve::protocol::{Change, CollectionWrite, Mutation, Operation, Version};
use std::sync::{
atomic::{AtomicUsize, Ordering},
mpsc,
};
let dir = tempfile::tempdir()?;
let service = service(dir.path())?;
let (mut client, server) = UnixStream::pair()?;
client.set_read_timeout(Some(Duration::from_secs(5)))?;
let active = Arc::new(AtomicUsize::new(1));
let permit = super::Permit(Arc::clone(&active));
let serving = Arc::clone(&service);
let (ended, end) = mpsc::channel();
let worker = std::thread::spawn(move || {
let result = super::connection(server, &serving, limits());
drop(permit);
ended.send(result).map_err(io::Error::other)
});
let hello = serde_json::json!({"id":"hello", "operation":{"method":"hello", "params":{
"protocol":crate::serve::protocol::PROTOCOL,
"capabilities":[crate::serve::protocol::RECEIPT_CAPABILITY],
"max_frame": limits().max_frame.get()
}}});
frame::write(
&mut client,
&serde_json::to_vec(&hello)?,
limits().max_frame,
)?;
frame::read(&mut client, limits().max_frame)?.ok_or("no hello")?;
let follow = serde_json::json!({"id":"follow","operation":{"method":"head.follow","params":{"branch":"main"}}});
frame::write(
&mut client,
&serde_json::to_vec(&follow)?,
limits().max_frame,
)?;
let bytes = frame::read(&mut client, limits().max_frame)?.ok_or("no initial head")?;
let head: serde_json::Value = serde_json::from_slice(&bytes)?;
let version: Version = serde_json::from_value(head["result"]["version"].clone())?;
drop(client);
let released = end.recv_timeout(Duration::from_secs(2));
let count = active.load(Ordering::Acquire);
service.dispatch(Operation::Mutate(Mutation {
request_id: "cleanup".into(),
branch: "main".into(),
expected: version,
change: Change::Collection(CollectionWrite {
collection: "cleanup".into(),
schema: None,
}),
}))?;
worker.join().map_err(|_| "follower panicked")??;
assert!(
released.is_ok(),
"closed follower retained its slot on an unmoving branch"
);
assert_eq!(count, 0, "closed follower still counted as active");
Ok(())
}