use std::error::Error;
use std::io::{Read, Write};
use std::net::{TcpListener, TcpStream};
use std::sync::Arc;
use beamr::scheduler::ExitWatchState;
use liminal::protocol::{Frame, MessageId, ProtocolError, ProtocolVersion};
use serde_json::Value;
use crate::config::{ChannelDef, LimitsConfig, ServerConfig, ServicesConfig, WebSocketConfig};
use crate::server::connection::{
ConnectionServices, ConnectionSupervisor, LiminalConnectionServices, LoopbackClientEnd,
LoopbackDuplex, WebSocketListener,
};
use crate::server::mount::MountKind;
use crate::server::shutdown::ShutdownHandle;
use super::events::ServerEvent;
use super::sink::SelfStreams;
use super::tests::R3Fixture;
type TestResult<T = ()> = Result<T, Box<dyn Error>>;
struct Fixture {
source: R3Fixture,
services: Arc<LiminalConnectionServices>,
supervisor: ConnectionSupervisor,
sink: Arc<SelfStreams>,
}
impl Fixture {
fn new() -> TestResult<Self> {
let source = R3Fixture::new()?;
let config = ServerConfig {
listen_address: "127.0.0.1:0".parse()?,
health_listen_address: "127.0.0.1:0".parse()?,
drain_timeout_ms: 999_999,
channels: vec![ChannelDef {
name: "pressure.source".to_owned(),
schema_ref: None,
durable: false,
loaded_schema: None,
}],
routing_rules: Vec::new(),
persistence_path: None,
cluster: None,
auth: None,
services: ServicesConfig::default(),
limits: LimitsConfig::default(),
participant: Some(crate::server::participant::test_participant_config()),
websocket: None,
};
let services = Arc::new(LiminalConnectionServices::from_config_with_store(
&config,
Arc::clone(&source.store),
)?);
let shutdown = ShutdownHandle::new();
let supervisor = ConnectionSupervisor::with_fatal_shutdown(
services.clone(),
None,
None,
config.limits,
shutdown.clone(),
)?;
let sink =
crate::server::runtime::bind_self_streams(&services, &supervisor, &config, &shutdown)?
.ok_or("mount fixture has no self streams")?;
Ok(Self {
source,
services,
supervisor,
sink,
})
}
fn head(&self) -> TestResult<u64> {
Ok(self.services.self_stream_head("liminal.events")?)
}
fn rows_since(&self, start: u64) -> TestResult<Vec<Value>> {
self.sink.drain()?;
let length = self
.head()?
.checked_sub(start)
.ok_or("event head moved backwards")?;
if length == 0 {
return Ok(Vec::new());
}
let rows = liminal::durability::bridge::block_on(self.source.store.read_from(
"liminal.events:0",
start,
usize::try_from(length)?,
))??;
assert_eq!(
u64::try_from(rows.len())?,
length,
"event tail was incomplete"
);
rows.into_iter()
.map(|row| {
let envelope = liminal::durability::MessageEnvelope::deserialize(&row.payload)?;
Ok(serde_json::from_slice(&envelope.payload)?)
})
.collect()
}
}
impl Drop for Fixture {
fn drop(&mut self) {
self.supervisor.shutdown();
if let Err(error) = self.sink.drain() {
eprintln!("mount fixture self-stream drain failed: {error}");
}
}
}
enum Transport {
Tcp(TcpStream),
Loopback(LoopbackClientEnd),
WebSocket(Box<tungstenite::WebSocket<TcpStream>>),
}
struct Peer {
transport: Transport,
buffer: Vec<u8>,
listener: Option<WebSocketListener>,
}
impl Peer {
fn connect(fixture: &Fixture, mount: MountKind) -> TestResult<Self> {
let (transport, listener) = match mount {
MountKind::Tcp => {
let listener = TcpListener::bind("127.0.0.1:0")?;
let client = TcpStream::connect(listener.local_addr()?)?;
let (server, _) = listener.accept()?;
fixture.supervisor.spawn_connection(server)?;
(Transport::Tcp(client), None)
}
MountKind::Loopback => {
let (client, server) = LoopbackDuplex::bounded(256 * 1024);
fixture.supervisor.spawn_loopback_connection(server)?;
(Transport::Loopback(client), None)
}
MountKind::WebSocket => {
let listener = WebSocketListener::bind(
&WebSocketConfig {
listen_address: "127.0.0.1:0".parse()?,
path: "/mount".to_owned(),
allowed_origins: Vec::new(),
ping_interval_ms: None,
},
fixture.supervisor.clone(),
)?;
let address = listener.local_addr();
let url = format!("ws://{address}/mount");
let (socket, _) = tungstenite::client(url.as_str(), TcpStream::connect(address)?)?;
(Transport::WebSocket(Box::new(socket)), Some(listener))
}
};
let mut peer = Self {
transport,
buffer: Vec::new(),
listener,
};
peer.send(&Frame::Connect {
flags: 0,
min_version: ProtocolVersion::new(1, 0),
max_version: ProtocolVersion::new(1, 0),
auth_token: Vec::new(),
})?;
assert!(matches!(peer.read()?, Frame::ConnectAck { .. }));
Ok(peer)
}
fn send(&mut self, frame: &Frame) -> TestResult {
let mut bytes = vec![0; liminal::protocol::encoded_len(frame)?];
let length = liminal::protocol::encode(frame, &mut bytes)?;
bytes.truncate(length);
match &mut self.transport {
Transport::Tcp(socket) => socket.write_all(&bytes)?,
Transport::Loopback(client) => client.write_all(&bytes)?,
Transport::WebSocket(socket) => {
socket.send(tungstenite::Message::Binary(bytes.into()))?;
}
}
Ok(())
}
fn read(&mut self) -> TestResult<Frame> {
if let Transport::WebSocket(socket) = &mut self.transport {
let tungstenite::Message::Binary(bytes) = socket.read()? else {
return Err("WebSocket peer received a non-binary application message".into());
};
let (frame, consumed) = liminal::protocol::decode(&bytes)?;
assert_eq!(consumed, bytes.len());
return Ok(frame);
}
loop {
match liminal::protocol::decode(&self.buffer) {
Ok((frame, consumed)) => {
self.buffer.drain(..consumed);
return Ok(frame);
}
Err(
ProtocolError::IncompleteHeader { .. } | ProtocolError::TruncatedPayload { .. },
) => {
let mut bytes = [0; 8192];
let length = match &mut self.transport {
Transport::Tcp(socket) => socket.read(&mut bytes)?,
Transport::Loopback(client) => client.read(&mut bytes)?,
Transport::WebSocket(_) => {
return Err("WebSocket entered the raw reader".into());
}
};
if length == 0 {
return Err("mount peer closed before its application frame".into());
}
self.buffer
.extend_from_slice(bytes.get(..length).ok_or("invalid peer read length")?);
}
Err(error) => return Err(error.into()),
}
}
}
fn subscribe(&mut self, channel: &str) -> TestResult {
self.send(&Frame::Subscribe {
flags: 0,
stream_id: 7,
channel: channel.to_owned(),
accepted_schemas: Vec::new(),
max_in_flight: 16,
})?;
assert!(matches!(
self.read()?,
Frame::SubscribeAck { stream_id: 7, .. }
));
Ok(())
}
fn delivery(&mut self) -> TestResult<String> {
match self.read()? {
Frame::Deliver {
stream_id: 7,
delivery_seq,
..
} => Ok(delivery_seq.to_string()),
frame => Err(format!("expected mount delivery, received {frame:?}").into()),
}
}
fn barrier(&mut self) -> TestResult {
self.send(&Frame::new_ping(0)?)?;
loop {
match self.read()? {
Frame::Pong { .. } => return Ok(()),
Frame::Deliver { stream_id: 7, .. } => {}
frame => return Err(format!("expected mount barrier, received {frame:?}").into()),
}
}
}
fn stop_listener(&mut self) -> TestResult {
if let Some(listener) = self.listener.take() {
listener.shutdown()?;
}
Ok(())
}
}
#[derive(Clone, Copy)]
enum Cell {
Open,
CleanClose,
AbruptClose,
Accept,
Defer,
Reject,
Suppressed,
}
fn metric(name: &str) -> TestResult<u64> {
use liminal::metrics::{MetricValue, global_registry};
let snapshot = global_registry()
.ok_or("server metrics registry is absent")?
.snapshot();
let counter = snapshot
.metrics()
.iter()
.find(|metric| metric.name == name)
.ok_or_else(|| format!("mount metric is absent: {name}"))?;
assert!(counter.labels.is_empty());
match counter.value {
MetricValue::Counter(value) => Ok(value),
_ => Err(format!("mount metric is not a counter: {name}").into()),
}
}
fn assert_event(rows: &[Value], pid: u64, kind: &str, response: Option<&str>, expected: usize) {
let events: Vec<_> = rows
.iter()
.filter(|row| row["details"]["kind"] == kind)
.collect();
assert_eq!(
events.len(),
expected,
"{kind} rows {}, expected {expected}",
events.len()
);
for event in events {
assert_eq!(event["details"]["connection"].as_u64(), Some(pid));
if let Some(response) = response {
assert_eq!(event["details"]["response"], response);
}
}
}
fn pressure_frame(cell: Cell, reference: MessageId) -> TestResult<(Frame, &'static str)> {
Ok(match cell {
Cell::Accept | Cell::Suppressed => (
Frame::Accept {
flags: 0,
stream_id: 7,
referenced_message_id: reference,
},
"accept",
),
Cell::Defer => (
Frame::Defer {
flags: 0,
stream_id: 7,
referenced_message_id: reference,
reason: None,
},
"defer",
),
Cell::Reject => (
Frame::Reject {
flags: 0,
stream_id: 7,
referenced_message_id: reference,
reason: None,
},
"reject",
),
_ => return Err("non-pressure cell entered pressure dispatch".into()),
})
}
fn pressure_pin(fixture: &Fixture, peer: &mut Peer, pid: u64, cell: Cell) -> TestResult {
let suppressed_cell = matches!(cell, Cell::Suppressed);
peer.subscribe(if suppressed_cell {
"liminal.events"
} else {
"pressure.source"
})?;
if suppressed_cell {
let seed_start = fixture.head()?;
fixture.sink.enqueue_event_for_test(ServerEvent::Pressure {
connection: pid,
response: "fixture_seed",
})?;
let seed = fixture.rows_since(seed_start)?;
assert_event(&seed, pid, "pressure", Some("fixture_seed"), 1);
} else {
fixture.services.publish(
pid,
"pressure.source",
&liminal::protocol::MessageEnvelope::new(
liminal::protocol::SchemaId::from_schema_bytes(b"{}"),
liminal::protocol::CausalContext::default(),
b"{}".to_vec(),
),
None,
)?;
}
let reference = MessageId::new(peer.delivery()?);
let start = fixture.head()?;
let suppressed = metric("liminal_self_streams_suppressed_pressure_total")?;
let unmatched = metric("liminal_self_streams_unmatched_pressure_total")?;
let (frame, response) = pressure_frame(cell, reference)?;
peer.send(&frame)?;
peer.barrier()?;
let rows = fixture.rows_since(start)?;
assert_event(
&rows,
pid,
"pressure",
Some(response),
usize::from(!suppressed_cell),
);
assert_eq!(
metric("liminal_self_streams_unmatched_pressure_total")? - unmatched,
0
);
assert_eq!(
metric("liminal_self_streams_suppressed_pressure_total")? - suppressed,
u64::from(suppressed_cell)
);
Ok(())
}
fn pin(mount: MountKind, cell: Cell) -> TestResult {
crate::metrics::init();
let fixture = Fixture::new()?;
let beginning = fixture.head()?;
let mut peer = Peer::connect(&fixture, mount)?;
let pids = fixture.supervisor.active_connection_pids();
let [pid] = pids.as_slice() else {
return Err("mount fixture did not admit exactly one connection".into());
};
let pid = *pid;
assert_eq!(fixture.supervisor.connection_mount(pid), Some(mount));
assert_event(
&fixture.rows_since(beginning)?,
pid,
"connection_open",
None,
1,
);
let start = fixture.head()?;
match cell {
Cell::Open => {}
Cell::CleanClose | Cell::AbruptClose => {
let ExitWatchState::Live(watch) = fixture.supervisor.scheduler().watch_exit(pid) else {
return Err("mount connection was not live before peer closure".into());
};
peer.stop_listener()?;
if matches!(cell, Cell::CleanClose) {
peer.send(&Frame::Disconnect { flags: 0 })?;
}
drop(peer);
let (exited, _) = watch.recv()?;
assert_eq!(exited, pid);
assert_eq!(fixture.supervisor.connection_mount(pid), None);
assert_event(
&fixture.rows_since(start)?,
pid,
"connection_close",
None,
1,
);
assert_event(
&fixture.rows_since(beginning)?,
pid,
"connection_open",
None,
1,
);
return Ok(());
}
Cell::Accept | Cell::Defer | Cell::Reject | Cell::Suppressed => {
pressure_pin(&fixture, &mut peer, pid, cell)?;
}
}
peer.stop_listener()?;
drop(peer);
Ok(())
}
macro_rules! mount_pin {
($name:ident, $mount:ident, $cell:ident) => {
#[test]
fn $name() -> TestResult {
pin(MountKind::$mount, Cell::$cell)
}
};
}
mount_pin!(tcp_open, Tcp, Open);
mount_pin!(tcp_clean_close, Tcp, CleanClose);
mount_pin!(tcp_abrupt_close, Tcp, AbruptClose);
mount_pin!(tcp_accept, Tcp, Accept);
mount_pin!(tcp_defer, Tcp, Defer);
mount_pin!(tcp_reject, Tcp, Reject);
mount_pin!(tcp_suppressed, Tcp, Suppressed);
mount_pin!(websocket_open, WebSocket, Open);
mount_pin!(websocket_clean_close, WebSocket, CleanClose);
mount_pin!(websocket_abrupt_close, WebSocket, AbruptClose);
mount_pin!(websocket_accept, WebSocket, Accept);
mount_pin!(websocket_defer, WebSocket, Defer);
mount_pin!(websocket_reject, WebSocket, Reject);
mount_pin!(websocket_suppressed, WebSocket, Suppressed);
mount_pin!(loopback_open, Loopback, Open);
mount_pin!(loopback_clean_close, Loopback, CleanClose);
mount_pin!(loopback_abrupt_close, Loopback, AbruptClose);
mount_pin!(loopback_accept, Loopback, Accept);
mount_pin!(loopback_defer, Loopback, Defer);
mount_pin!(loopback_reject, Loopback, Reject);
mount_pin!(loopback_suppressed, Loopback, Suppressed);