use std::io::{self, Write};
use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
use std::sync::{Arc, Mutex};
use beamr::native::native_process::{NativeContext, NativeHandler, NativeOutcome};
use beamr::process::ExitReason;
use beamr::scheduler::Interest;
use super::{ConnectionRuntime, ConnectionSupervisor, recover_lock};
use crate::ServerError;
use crate::server::connection::process::{
ConnectionTransport, ReadStatus, TransportConnectionProcess,
};
use crate::server::mount::MountKind;
use crate::server::participant::ParticipantSemanticHandler;
type TestResult = Result<(), Box<dyn std::error::Error>>;
fn supervisor() -> Result<ConnectionSupervisor, Box<dyn std::error::Error>> {
use crate::config::{LimitsConfig, ServerConfig, ServicesConfig};
use crate::server::connection::services::LiminalConnectionServices;
use crate::server::participant::ParticipantMemoryStore;
let config = ServerConfig {
listen_address: "127.0.0.1:0".parse()?,
extra_listen_addresses: Vec::new(),
health_listen_address: "127.0.0.1:0".parse()?,
drain_timeout_ms: 0,
channels: Vec::new(),
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::new(ParticipantMemoryStore::default()),
)?);
Ok(ConnectionSupervisor::with_services(services)?)
}
fn assert_shutdown_fates(supervisor: &ConnectionSupervisor, count: usize) -> TestResult {
use crate::server::participant::incarnation_stream::IncarnationStream;
use liminal::durability::bridge::block_on;
let store = supervisor
.inner
.runtime
.participant_service()
.ok_or("missing participant service")?
.durable_store();
let current = supervisor
.current_server_incarnation()?
.ok_or("missing incarnation")?;
let rows = block_on(store.read_from(&IncarnationStream::epoch_stream_key(current), 0, 1024))??;
let opens = rows
.iter()
.filter(|row| row.payload.get(5) == Some(&3))
.collect::<Vec<_>>();
assert_eq!(opens.len(), count, "one durable fate per connection");
assert!(
opens.iter().all(|row| row.payload.get(26) == Some(&2)),
"every fate is ServerShutdown"
);
assert_eq!(
rows.iter()
.filter(|row| row.payload.get(5) == Some(&4))
.count(),
count,
"every shutdown fate completed before the wait returned"
);
Ok(())
}
#[derive(Debug, Default)]
struct Counts {
armed: AtomicUsize,
reads: AtomicUsize,
exits: AtomicUsize,
crashes: AtomicUsize,
close_on_turn: AtomicBool,
frame_waiting: AtomicBool,
close_during_read: AtomicBool,
written: Mutex<Vec<u8>>,
log: crate::test_log::CapturedLog,
}
#[derive(Debug)]
struct QuietTransport {
counts: Arc<Counts>,
connected: bool,
supervisor: ConnectionSupervisor,
}
impl Write for QuietTransport {
fn write(&mut self, bytes: &[u8]) -> io::Result<usize> {
recover_lock(&self.counts.written).extend_from_slice(bytes);
Ok(bytes.len())
}
fn flush(&mut self) -> io::Result<()> {
Ok(())
}
}
impl ConnectionTransport for QuietTransport {
const MOUNT: MountKind = MountKind::Tcp;
fn is_connected(&self) -> bool {
self.connected
}
fn read_available(&mut self, buffer: &mut Vec<u8>) -> Result<ReadStatus, ServerError> {
self.counts.reads.fetch_add(1, Ordering::SeqCst);
if self.counts.close_during_read.swap(false, Ordering::SeqCst) {
self.supervisor.force_close_active_connections();
return Ok(ReadStatus::Closed);
}
if self.counts.frame_waiting.swap(false, Ordering::SeqCst) {
let frame = liminal::protocol::Frame::LinkOpen {
flags: 0,
stream_id: 7,
offer: liminal::protocol::LinkOffer {
link_id: "shutdown-link".to_owned(),
end: liminal::protocol::LinkEnd::Export,
machine: [7; 32],
handshake: vec![1, 2, 3],
},
};
let mut bytes = vec![
0;
liminal::protocol::encoded_len(&frame).map_err(|error| {
ServerError::ListenerAccept {
message: error.to_string(),
}
})?
];
liminal::protocol::encode(&frame, &mut bytes).map_err(|error| {
ServerError::ListenerAccept {
message: error.to_string(),
}
})?;
buffer.extend_from_slice(&bytes);
Ok(ReadStatus::Read)
} else {
Ok(ReadStatus::WouldBlock)
}
}
fn sink(&mut self) -> Option<&mut dyn Write> {
Some(self)
}
fn probe_inbound(&self) -> Result<bool, ServerError> {
Ok(false)
}
fn arm_readiness(
&mut self,
_: u64,
_: &NativeContext<'_>,
interest: Interest,
_: &ConnectionRuntime,
) -> Result<(), ServerError> {
assert_eq!(interest, Interest::READABLE);
self.counts.armed.fetch_add(1, Ordering::SeqCst);
Ok(())
}
fn install_wake(&mut self, _: u64, _: &ConnectionRuntime) -> Result<(), ServerError> {
Ok(())
}
fn release(&mut self) {
self.connected = false;
}
fn note_process_drop(&mut self, _: &ConnectionRuntime) {}
}
struct InFlightControl {
process: TransportConnectionProcess<QuietTransport>,
supervisor: ConnectionSupervisor,
counts: Arc<Counts>,
}
impl NativeHandler for InFlightControl {
fn handle(&mut self, ctx: &mut NativeContext<'_>) -> NativeOutcome {
if self.counts.close_on_turn.swap(false, Ordering::SeqCst) {
self.supervisor.force_close_active_connections();
}
let outcome = self.counts.log.capture(|| self.process.handle(ctx));
match &outcome {
NativeOutcome::Stop(ExitReason::Normal) => {
self.counts.exits.fetch_add(1, Ordering::SeqCst);
}
NativeOutcome::Stop(_) => {
self.counts.crashes.fetch_add(1, Ordering::SeqCst);
}
_ => {}
}
outcome
}
}
fn parked(
supervisor: &ConnectionSupervisor,
) -> Result<(u64, Arc<Counts>), Box<dyn std::error::Error>> {
let counts = Arc::new(Counts::default());
let observed = Arc::clone(&counts);
let owner = supervisor.clone();
let runtime = Arc::clone(&supervisor.inner.runtime);
supervisor.inner.runtime.try_reserve_admission()?;
let incarnation = supervisor
.inner
.allocate_connection_incarnation()?
.ok_or("missing connection incarnation")?;
let conversation_id = incarnation.connection_ordinal;
let mut conversations =
crate::server::participant::ParticipantConnectionConversations::default();
let service = supervisor
.inner
.runtime
.participant_service()
.ok_or("missing participant service")?;
service.handle(
crate::server::participant::ParticipantConnectionContext::new(incarnation, MountKind::Tcp),
&mut conversations,
liminal_protocol::wire::ClientRequest::Enrollment(
liminal_protocol::wire::EnrollmentRequest {
conversation_id,
enrollment_token: liminal_protocol::wire::EnrollmentToken::new([7; 16]),
},
),
)?;
assert_eq!(
conversations.occupied(),
1,
"the connection tracks its real enrolled conversation"
);
let pid = supervisor.inner.spawn_native_connection(
Box::new(move || {
let mut process = TransportConnectionProcess::over_transport(
Arc::clone(&runtime),
None,
QuietTransport {
counts: Arc::clone(&observed),
connected: true,
supervisor: owner.clone(),
},
Some(incarnation),
);
process.track_conversation_for_tests(conversation_id);
Box::new(InFlightControl {
process,
supervisor: owner.clone(),
counts: Arc::clone(&observed),
})
}),
"failed to spawn the native connection fixture",
)?;
supervisor.inner.runtime.register_connection(
pid,
None,
Some(incarnation),
MountKind::Tcp,
None,
)?;
supervisor.observe_settled_park(pid).recv()?;
assert_eq!(counts.armed.load(Ordering::SeqCst), 1);
counts.reads.store(0, Ordering::SeqCst);
counts.armed.store(0, Ordering::SeqCst);
Ok((pid, counts))
}
fn close_in_flight(frame_waiting: bool) -> TestResult {
let supervisor = supervisor()?;
let (pid, counts) = parked(&supervisor)?;
let beamr::scheduler::ExitWatchState::Live(exit) = supervisor.inner.scheduler.watch_exit(pid)
else {
return Err("connection not live".into());
};
counts.frame_waiting.store(frame_waiting, Ordering::SeqCst);
counts.close_on_turn.store(true, Ordering::SeqCst);
assert!(supervisor.resume_test_process(pid));
let log = crate::test_log::CapturedLog::default();
log.capture(|| crate::server::shutdown::wait_after_force_close(&supervisor));
assert!(!log.contents().contains("connections remained active"));
let active = supervisor.active_connection_count();
assert!(
supervisor
.inner
.scheduler
.process_table()
.get(pid)
.is_none(),
"zero native processes before teardown"
);
let (exited_pid, _) = exit.recv()?;
assert_shutdown_fates(&supervisor, 1)?;
supervisor.shutdown();
assert_eq!(exited_pid, pid);
assert_eq!(counts.exits.load(Ordering::SeqCst), 1, "one shutdown exit");
assert_eq!(
counts.crashes.load(Ordering::SeqCst),
0,
"zero crash records"
);
assert_eq!(
counts.reads.load(Ordering::SeqCst),
0,
"ForceClose must run before inbound service"
);
assert_eq!(
counts.armed.load(Ordering::SeqCst),
0,
"zero readiness rearms"
);
assert_eq!(active, 0, "zero registered processes at teardown");
assert!(
!counts.log.contents().contains("inbound link open refused"),
"zero link refusal events"
);
let written = recover_lock(&counts.written);
let (frame, consumed) = liminal::protocol::decode(&written)?;
assert!(matches!(frame, liminal::protocol::Frame::Disconnect { .. }));
assert_eq!(
consumed,
written.len(),
"only the shutdown frame was served"
);
Ok(())
}
#[test]
fn a_connection_parked_on_readiness_ends_on_force_close_without_a_socket_event() -> TestResult {
close_in_flight(false)?;
let supervisor = supervisor()?;
let (pid, counts) = parked(&supervisor)?;
supervisor.force_close_active_connections();
crate::server::shutdown::wait_after_force_close(&supervisor);
assert!(
supervisor
.inner
.scheduler
.process_table()
.get(pid)
.is_none()
);
assert_eq!(counts.exits.load(Ordering::SeqCst), 1);
assert_eq!(counts.crashes.load(Ordering::SeqCst), 0);
assert_eq!(counts.reads.load(Ordering::SeqCst), 0);
assert_eq!(counts.armed.load(Ordering::SeqCst), 0);
assert_shutdown_fates(&supervisor, 1)?;
supervisor.shutdown();
Ok(())
}
#[test]
fn a_force_closed_connection_never_rearms_readiness() -> TestResult {
close_in_flight(false)
}
#[test]
fn a_frame_arriving_after_force_close_is_not_served() -> TestResult {
close_in_flight(true)
}
#[test]
fn the_wedged_stop_sees_no_active_connection_after_closing_n_live_links() -> TestResult {
let supervisor = supervisor()?;
let mut connections = Vec::new();
for _ in 0..3 {
let (pid, counts) = parked(&supervisor)?;
let beamr::scheduler::ExitWatchState::Live(exit) =
supervisor.inner.scheduler.watch_exit(pid)
else {
return Err("connection not live".into());
};
connections.push((pid, counts, exit));
}
let (pid, counts, _) = connections.first().ok_or("missing native connection")?;
counts.close_on_turn.store(true, Ordering::SeqCst);
assert!(supervisor.resume_test_process(*pid));
let log = crate::test_log::CapturedLog::default();
log.capture(|| crate::server::shutdown::wait_after_force_close(&supervisor));
assert!(!log.contents().contains("connections remained active"));
let active = supervisor.active_connection_count();
assert_shutdown_fates(&supervisor, 3)?;
for (pid, _, _) in &connections {
assert!(
supervisor
.inner
.scheduler
.process_table()
.get(*pid)
.is_none(),
"zero native processes before teardown"
);
}
supervisor.shutdown();
assert_eq!(active, 0);
for (pid, counts, exit) in connections {
assert!(
supervisor
.inner
.scheduler
.process_table()
.get(pid)
.is_none()
);
exit.recv()?;
assert_eq!(counts.exits.load(Ordering::SeqCst), 1);
assert_eq!(counts.crashes.load(Ordering::SeqCst), 0);
assert_eq!(
counts.armed.load(Ordering::SeqCst),
0,
"each close ends on its next slice"
);
}
Ok(())
}
#[test]
fn a_force_close_published_during_a_read_keeps_the_shutdown_fate() -> TestResult {
let supervisor = supervisor()?;
let (pid, counts) = parked(&supervisor)?;
counts.close_during_read.store(true, Ordering::SeqCst);
assert!(supervisor.resume_test_process(pid));
crate::server::shutdown::wait_after_force_close(&supervisor);
assert!(
supervisor
.inner
.scheduler
.process_table()
.get(pid)
.is_none()
);
assert_eq!(
counts.reads.load(Ordering::SeqCst),
1,
"one read already started before publication"
);
assert_eq!(counts.armed.load(Ordering::SeqCst), 0);
assert_eq!(counts.exits.load(Ordering::SeqCst), 1);
assert_eq!(counts.crashes.load(Ordering::SeqCst), 0);
let fate = assert_shutdown_fates(&supervisor, 1);
supervisor.shutdown();
fate
}