use std::collections::HashMap;
use std::collections::VecDeque;
use std::future::Future;
use std::pin::Pin;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::Arc;
use bytes::Bytes;
use tokio::sync::{mpsc, oneshot, OwnedSemaphorePermit, Semaphore};
use unb_core::{
ClientDelivery as CoreClientDelivery, ClientOperationId, CoreEffect, CoreInput, EffectId,
Envelope, ProtocolCore, RetirementReason, SendResult, SessionId,
};
use web_time::Instant;
use crate::cancellation::CancellationToken;
use crate::client::{ClientDelivery, ClientSession};
use crate::error::WsError;
use crate::transport::{Pipe, PipeReader, PipeWriter};
use crate::wire::{Directive, Wire};
use crate::{BodyStream, WireBody};
const CORE_QUEUE: usize = 256;
const SESSION_QUEUE: usize = 64;
const OUTBOUND_QUEUE: usize = 8_192;
const OUTBOUND_BYTE_BUDGET: usize = 64 * 1024 * 1024;
const BODY_WINDOW_BYTES: usize = 256 * 1024;
const WRITE_BATCH: usize = 8;
const RELAY_BACKPRESSURE_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(5);
const TERMINAL_SEND_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(30);
const CLIENT_DELIVERY_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(30);
const ESTABLISHMENT_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(5);
const PENDING_MODE_WRITES_MAX: usize = 128;
const OPERATION_WRITE_QUEUE: usize = OUTBOUND_QUEUE;
const OPERATION_ACTIVE_MAX: usize = 128;
const OPERATION_BODY_CHUNK_BYTES: usize = 64 * 1024;
#[cfg(not(all(target_arch = "wasm32", target_os = "unknown")))]
fn clock_now() -> Instant {
n0_future::time::Instant::now().into_std()
}
#[cfg(not(all(target_arch = "wasm32", target_os = "unknown")))]
async fn sleep_to(deadline: Instant) {
n0_future::time::sleep_until(n0_future::time::Instant::from_std(deadline)).await
}
#[cfg(all(target_arch = "wasm32", target_os = "unknown"))]
fn clock_now() -> Instant {
n0_future::time::Instant::now()
}
#[cfg(all(target_arch = "wasm32", target_os = "unknown"))]
async fn sleep_to(deadline: Instant) {
n0_future::time::sleep_until(deadline).await
}
#[cfg(not(all(target_arch = "wasm32", target_os = "unknown")))]
pub use tokio::runtime::Handle as RuntimeHandle;
#[cfg(all(target_arch = "wasm32", target_os = "unknown"))]
#[derive(Clone, Debug)]
pub struct RuntimeHandle;
#[cfg(all(target_arch = "wasm32", target_os = "unknown"))]
impl RuntimeHandle {
pub fn current() -> Self {
RuntimeHandle
}
pub fn spawn<F>(&self, future: F)
where
F: Future<Output = ()> + 'static,
{
n0_future::task::spawn(future);
}
}
pub type EffectFuture = Pin<Box<dyn Future<Output = Option<CoreInput>> + Send>>;
pub trait EffectExecutor: Send + Sync + 'static {
fn execute(&self, effect: CoreEffect, handle: ProtocolCoreHandle) -> EffectFuture;
}
#[derive(Clone, Debug)]
pub enum SessionOutcome {
Established,
Retired(RetirementReason),
}
pub trait SessionHandler: Send + 'static {
fn deliver(&mut self, envelope: Envelope) -> impl Future<Output = ()> + Send;
fn stream_closed(&mut self, _operation: ClientOperationId) -> impl Future<Output = ()> + Send {
async {}
}
}
impl SessionHandler for () {
async fn deliver(&mut self, _envelope: Envelope) {}
}
impl SessionHandler for mpsc::Sender<Envelope> {
async fn deliver(&mut self, envelope: Envelope) {
let _ = self.send(envelope).await;
}
}
impl SessionHandler for tokio::sync::broadcast::Sender<Envelope> {
async fn deliver(&mut self, envelope: Envelope) {
let _ = self.send(envelope);
}
}
type TokenBodies = Arc<std::sync::Mutex<HashMap<(SessionId, String), WireBody>>>;
type BodyPumps = Arc<std::sync::Mutex<HashMap<(SessionId, String), (String, CancellationToken)>>>;
const SESSION_BODY_TOKENS_MAX: usize = OUTBOUND_QUEUE;
const SESSION_BODY_BYTES_MAX: usize = OUTBOUND_BYTE_BUDGET;
fn mint_body_token() -> String {
use std::hash::{BuildHasher, Hasher};
type TokenKeys = (
std::collections::hash_map::RandomState,
std::collections::hash_map::RandomState,
);
static KEYS: std::sync::OnceLock<TokenKeys> = std::sync::OnceLock::new();
static NEXT: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
let (low, high) = KEYS.get_or_init(TokenKeys::default);
let serial = NEXT.fetch_add(1, Ordering::Relaxed);
let mut a = low.build_hasher();
a.write_u64(serial);
let mut b = high.build_hasher();
b.write_u64(serial);
format!("b{:016x}{:016x}", a.finish(), b.finish())
}
fn application_core_input(
handle: &ProtocolCoreHandle,
session: &SessionId,
mut envelope: Envelope,
) -> Result<CoreInput, unb_core::CoreError> {
if envelope.kind == unb_core::Kind::Discover {
return Ok(CoreInput::FrameReceived {
session: session.clone(),
envelope,
});
}
if !(envelope.kind.is_application_request()
|| (envelope.corr.is_some()
&& (envelope.kind.is_application_response()
|| envelope.kind == unb_core::Kind::Cancel)))
{
return Ok(CoreInput::FrameReceived {
session: session.clone(),
envelope,
});
}
if envelope.body_token.is_some() && !envelope.payload.is_empty() {
return Err(unb_core::CoreError::Malformed(
"application frame cannot carry both inline and referenced bodies".into(),
));
}
if envelope.body_token.is_none()
&& envelope.kind != unb_core::Kind::Error
&& !envelope.payload.is_empty()
{
let payload = std::mem::take(&mut envelope.payload);
envelope.body_token = Some(
handle
.register_body(session, WireBody::Bytes(payload))?
.to_string(),
);
}
Ok(CoreInput::ApplicationFrameReceived {
session: session.clone(),
frame: unb_core::ApplicationFrame::from_envelope(&envelope)?,
})
}
fn register_outbound_body(
handle: &ProtocolCoreHandle,
session: &SessionId,
payload: Bytes,
body: Option<BodyStream>,
) -> Result<Option<unb_core::BodyId>, unb_core::CoreError> {
if body.is_some() && !payload.is_empty() {
return Err(unb_core::CoreError::Malformed(
"an operation cannot provide both unary bytes and a streaming body".into(),
));
}
let body = match body {
Some(body) => Some(WireBody::Stream(body)),
None if !payload.is_empty() => Some(WireBody::Bytes(payload)),
None => None,
};
body.map(|body| handle.register_body(session, body))
.transpose()
}
#[derive(Clone)]
pub struct ProtocolCoreHandle {
inputs: mpsc::Sender<ActorInput>,
cancellation: CancellationToken,
token_bodies: TokenBodies,
body_pumps: BodyPumps,
}
impl ProtocolCoreHandle {
pub fn spawn(
core: ProtocolCore,
executor: Arc<dyn EffectExecutor>,
cancellation: CancellationToken,
runtime: &RuntimeHandle,
) -> Self {
let (inputs, receiver) = mpsc::channel(CORE_QUEUE);
let handle = Self {
inputs,
cancellation,
token_bodies: Arc::new(std::sync::Mutex::new(HashMap::new())),
body_pumps: Arc::new(std::sync::Mutex::new(HashMap::new())),
};
runtime.spawn(CoreRuntime::new(core, receiver, handle.clone(), executor).run());
handle
}
pub fn claim_body(&self, session: &SessionId, token: &str) -> Option<WireBody> {
self.token_bodies
.lock()
.expect("body token registry")
.remove(&(session.clone(), token.to_string()))
}
pub fn claim_body_stream(&self, session: &SessionId, token: &str) -> Option<BodyStream> {
self.claim_body(session, token).map(wire_body_stream)
}
pub fn register_body(
&self,
session: &SessionId,
body: WireBody,
) -> Result<unb_core::BodyId, unb_core::CoreError> {
let mut registry = self.token_bodies.lock().expect("body token registry");
let owned = registry
.keys()
.filter(|(owner, _)| owner == session)
.count();
if owned >= SESSION_BODY_TOKENS_MAX {
return Err(unb_core::CoreError::EffectQueueFull);
}
let incoming = match &body {
WireBody::Bytes(payload) => payload.len(),
WireBody::Stream(_) => 0,
};
let retained = registry
.iter()
.filter_map(|((owner, _), body)| {
(owner == session).then_some(match body {
WireBody::Bytes(payload) => payload.len(),
WireBody::Stream(_) => 0,
})
})
.sum::<usize>();
if retained.saturating_add(incoming) > SESSION_BODY_BYTES_MAX {
return Err(unb_core::CoreError::BodyRegistryFull(
SESSION_BODY_BYTES_MAX,
));
}
let token = mint_body_token();
registry.insert((session.clone(), token.clone()), body);
Ok(unb_core::BodyId::from(token))
}
fn register_body_pump(
&self,
session: &SessionId,
corr: &str,
token: &str,
cancellation: CancellationToken,
) {
self.body_pumps.lock().expect("body pump registry").insert(
(session.clone(), corr.to_string()),
(token.to_string(), cancellation),
);
}
fn cancel_body_pump(&self, session: &SessionId, corr: &str) {
let cancelled = self
.body_pumps
.lock()
.expect("body pump registry")
.remove(&(session.clone(), corr.to_string()));
if let Some((_, cancellation)) = cancelled {
cancellation.cancel();
}
}
fn release_body_pump(&self, session: &SessionId, corr: &str, token: &str) {
let mut pumps = self.body_pumps.lock().expect("body pump registry");
let key = (session.clone(), corr.to_string());
if pumps.get(&key).is_some_and(|(owner, _)| owner == token) {
pumps.remove(&key);
}
}
pub async fn submit(&self, input: CoreInput) -> Result<(), WsError> {
self.inputs
.send(ActorInput::Core(input))
.await
.map_err(|_| WsError::Gone)
}
pub async fn submit_batch(&self, inputs: Vec<CoreInput>) -> Result<(), WsError> {
self.inputs
.send(ActorInput::CoreBatch(inputs))
.await
.map_err(|_| WsError::Gone)
}
pub async fn install<F>(&self, input: CoreInput, publish: F) -> Result<(), WsError>
where
F: FnOnce() + Send + 'static,
{
let (reply, response) = oneshot::channel();
self.inputs
.send(ActorInput::Install {
input,
publish: Box::new(publish),
reply,
})
.await
.map_err(|_| WsError::Gone)?;
response.await.map_err(|_| WsError::Gone)?
}
pub async fn attach<H: SessionHandler>(
&self,
session: SessionId,
transport: Pipe,
expected_peer: Option<String>,
handler: H,
) -> Result<Arc<Wire>, WsError> {
self.attach_with_ceiling(
session,
transport,
expected_peer,
handler,
unb_transport::DEFAULT_MAX_FRAME_SIZE,
)
.await
}
pub async fn attach_with_ceiling<H: SessionHandler>(
&self,
session: SessionId,
transport: Pipe,
expected_peer: Option<String>,
handler: H,
ws_collect_ceiling: usize,
) -> Result<Arc<Wire>, WsError> {
let cancellation = self.cancellation.child_token();
let (wire, directives, outcome, routes) = Wire::channel(cancellation.clone());
let (reply, response) = oneshot::channel();
self.inputs
.send(ActorInput::Attach {
session,
transport,
expected_peer,
directives,
handler: handler_channel(handler),
client: wire.client_session(),
cancellation,
establish_peer: true,
outcome,
routes,
ws_collect_ceiling,
reply,
})
.await
.map_err(|_| WsError::Gone)?;
response.await.map_err(|_| WsError::Gone)??;
Ok(wire)
}
pub fn shutdown(&self) {
self.cancellation.cancel();
}
pub(crate) fn open(transport: Pipe, runtime: &RuntimeHandle) -> (Arc<Wire>, Self) {
struct StandaloneExecutor;
impl EffectExecutor for StandaloneExecutor {
fn execute(&self, _effect: CoreEffect, _handle: ProtocolCoreHandle) -> EffectFuture {
Box::pin(async { None })
}
}
let cancellation = CancellationToken::new();
let handle = Self::spawn(
ProtocolCore::new("client"),
Arc::new(StandaloneExecutor),
cancellation.clone(),
runtime,
);
let session_cancellation = cancellation.child_token();
let (wire, directives, outcome, routes) =
Wire::standalone_channel(session_cancellation.clone());
let (reply, _response) = oneshot::channel();
handle
.inputs
.try_send(ActorInput::Attach {
session: SessionId::from("client-session"),
transport,
expected_peer: None,
directives,
handler: handler_channel(wire.observation_sender()),
client: wire.client_session(),
cancellation: session_cancellation,
establish_peer: false,
outcome,
routes,
ws_collect_ceiling: unb_transport::DEFAULT_MAX_FRAME_SIZE,
reply,
})
.expect("new protocol core input queue");
(wire, handle)
}
}
fn wire_body_stream(body: WireBody) -> BodyStream {
match body {
WireBody::Bytes(bytes) => Box::pin(futures_util::stream::once(async move { Ok(bytes) })),
WireBody::Stream(stream) => stream,
}
}
enum ActorInput {
Core(CoreInput),
CoreBatch(Vec<CoreInput>),
Install {
input: CoreInput,
publish: Box<dyn FnOnce() + Send>,
reply: oneshot::Sender<Result<(), WsError>>,
},
Directive(SessionId, Directive),
Attach {
session: SessionId,
transport: Pipe,
expected_peer: Option<String>,
directives: mpsc::Receiver<Directive>,
handler: mpsc::Sender<HandlerInput>,
client: ClientSession,
cancellation: CancellationToken,
establish_peer: bool,
outcome: tokio::sync::watch::Sender<Option<SessionOutcome>>,
routes: tokio::sync::watch::Sender<bool>,
ws_collect_ceiling: usize,
reply: oneshot::Sender<Result<(), WsError>>,
},
Write {
session: SessionId,
envelope: Envelope,
completion: Option<EffectId>,
relay: bool,
},
TransportModeSelected {
session: SessionId,
mode: unb_transport::TransportMode,
},
}
struct SessionIo {
outbound: mpsc::Sender<WriteRequest>,
byte_budget: Arc<Semaphore>,
ws_collect_ceiling: usize,
bodies: crate::transport::SessionStreams,
handler: mpsc::Sender<HandlerInput>,
client: ClientSession,
cancellation: CancellationToken,
reservation_pending: Arc<AtomicBool>,
establish_peer: bool,
deadline: tokio::sync::watch::Sender<Option<Instant>>,
routes: tokio::sync::watch::Sender<bool>,
transport_mode: Option<unb_transport::TransportMode>,
pending_mode_writes: VecDeque<PendingModeWrite>,
operation_lanes: OperationLanes,
operation_slots: Arc<Semaphore>,
}
struct PendingModeWrite {
envelope: Envelope,
completion: Option<EffectId>,
relay: bool,
}
type OperationLanes = Arc<std::sync::Mutex<HashMap<String, OperationLane>>>;
#[derive(Clone)]
struct OperationLane {
outbound: Option<mpsc::Sender<OperationWriteRequest>>,
control: Option<unb_transport::operation::OperationControl>,
}
struct OperationWriteRequest {
envelope: Envelope,
completion: Option<EffectId>,
relay: bool,
body: Option<BodyStream>,
_budget: OwnedSemaphorePermit,
}
struct WriteRequest {
envelope: Envelope,
completion: Option<EffectId>,
relay: bool,
_budget: OwnedSemaphorePermit,
}
impl WriteRequest {
fn weight(envelope: &Envelope) -> u32 {
let head = envelope.target.len()
+ envelope.subject.len()
+ envelope.id.len()
+ envelope.corr.as_ref().map_or(0, String::len)
+ envelope.path.iter().map(String::len).sum::<usize>()
+ envelope
.headers
.iter()
.map(|(name, value)| name.len() + value.as_str().map_or(16, str::len))
.sum::<usize>();
(envelope.payload.len() + head).clamp(1, OUTBOUND_BYTE_BUDGET) as u32
}
}
#[derive(Debug)]
pub(crate) struct MeteredChunk {
pub(crate) bytes: Bytes,
_window: OwnedSemaphorePermit,
_budget: OwnedSemaphorePermit,
}
pub(crate) struct MeteredBody {
stream: BodyStream,
budget: Arc<Semaphore>,
window: Arc<Semaphore>,
pending: Option<Bytes>,
cancellation: CancellationToken,
}
impl MeteredBody {
pub(crate) fn new(
stream: BodyStream,
budget: Arc<Semaphore>,
cancellation: CancellationToken,
) -> MeteredBody {
MeteredBody {
stream,
budget,
window: Arc::new(Semaphore::new(BODY_WINDOW_BYTES)),
pending: None,
cancellation,
}
}
pub(crate) async fn next_chunk(&mut self) -> Result<Option<MeteredChunk>, WsError> {
let chunk = match self.pending.take() {
Some(chunk) => chunk,
None => {
let next = tokio::select! {
biased;
() = self.cancellation.cancelled() => return Ok(None),
next = std::future::poll_fn(|cx| self.stream.as_mut().poll_next(cx)) => next,
};
match next {
Some(Ok(chunk)) => chunk,
Some(Err(error)) => return Err(error.into()),
None => return Ok(None),
}
}
};
let len = chunk.len();
if len > unb_transport::DEFAULT_MAX_FRAME_SIZE {
return Err(
unb_core::CoreError::BodyTooLarge(unb_transport::DEFAULT_MAX_FRAME_SIZE).into(),
);
}
self.pending = Some(chunk);
let window_weight = len.clamp(1, BODY_WINDOW_BYTES) as u32;
let budget_weight = len.max(1) as u32;
let acquire = async {
let window = self
.window
.clone()
.acquire_many_owned(window_weight)
.await
.map_err(|_| unb_core::CoreError::BodyStalled)?;
let budget = self
.budget
.clone()
.acquire_many_owned(budget_weight)
.await
.map_err(|_| unb_core::CoreError::BodyStalled)?;
Ok::<_, unb_core::CoreError>((window, budget))
};
let acquired = tokio::select! {
biased;
() = self.cancellation.cancelled() => return Ok(None),
acquired = n0_future::time::timeout(RELAY_BACKPRESSURE_TIMEOUT, acquire) => acquired,
};
match acquired {
Ok(Ok((window, budget))) => Ok(Some(MeteredChunk {
bytes: self.pending.take().expect("pending chunk present"),
_window: window,
_budget: budget,
})),
Ok(Err(error)) => Err(error.into()),
Err(_) => Err(unb_core::CoreError::BodyStalled.into()),
}
}
}
struct CoreRuntime {
core: ProtocolCore,
inputs: mpsc::Receiver<ActorInput>,
handle: ProtocolCoreHandle,
executor: Arc<dyn EffectExecutor>,
sessions: HashMap<SessionId, SessionIo>,
outcomes: HashMap<SessionId, tokio::sync::watch::Sender<Option<SessionOutcome>>>,
client_starts: HashMap<SessionId, PendingClientStart>,
pending_bodies: std::sync::Mutex<HashMap<(SessionId, String), BodyStream>>,
}
struct PendingClientStart {
sender: mpsc::Sender<ClientDelivery>,
reply: oneshot::Sender<Result<ClientOperationId, unb_core::CoreError>>,
}
impl CoreRuntime {
fn new(
core: ProtocolCore,
inputs: mpsc::Receiver<ActorInput>,
handle: ProtocolCoreHandle,
executor: Arc<dyn EffectExecutor>,
) -> Self {
Self {
core,
inputs,
handle,
executor,
sessions: HashMap::new(),
outcomes: HashMap::new(),
client_starts: HashMap::new(),
pending_bodies: std::sync::Mutex::new(HashMap::new()),
}
}
async fn run(mut self) {
loop {
let input = tokio::select! {
biased;
() = self.handle.cancellation.cancelled() => return,
input = self.inputs.recv() => input,
};
let Some(input) = input else { return };
match input {
ActorInput::Core(input) => {
let _ = self.core.handle(clock_now(), input);
}
ActorInput::CoreBatch(inputs) => {
let now = clock_now();
for input in inputs {
let _ = self.core.handle(now, input);
}
}
ActorInput::Install {
input,
publish,
reply,
} => {
let result = self.core.handle(clock_now(), input).map_err(WsError::from);
if result.is_ok() {
publish();
}
let _ = reply.send(result);
}
ActorInput::Directive(session, directive) => {
self.directive(&session, directive).await;
}
ActorInput::Attach {
session,
transport,
expected_peer,
directives,
handler,
client,
cancellation,
establish_peer,
outcome,
routes,
ws_collect_ceiling,
reply,
} => {
let result = self
.attach(
session,
transport,
expected_peer,
directives,
handler,
client,
cancellation,
establish_peer,
outcome,
routes,
ws_collect_ceiling,
)
.await;
let _ = reply.send(result);
}
ActorInput::Write {
session,
envelope,
completion,
relay,
} => {
self.write(&session, envelope, completion, relay).await;
}
ActorInput::TransportModeSelected { session, mode } => {
let pending = if let Some(io) = self.sessions.get_mut(&session) {
io.transport_mode = Some(mode);
std::mem::take(&mut io.pending_mode_writes)
} else {
VecDeque::new()
};
if mode == unb_transport::TransportMode::OperationStreamsV1 {
self.start_operation_acceptor(&session);
}
for pending in pending {
self.write(
&session,
pending.envelope,
pending.completion,
pending.relay,
)
.await;
}
}
}
self.drain().await;
}
}
#[allow(clippy::too_many_arguments)]
async fn attach(
&mut self,
session: SessionId,
transport: Pipe,
expected_peer: Option<String>,
directives: mpsc::Receiver<Directive>,
handler: mpsc::Sender<HandlerInput>,
client: ClientSession,
cancellation: CancellationToken,
establish_peer: bool,
outcome: tokio::sync::watch::Sender<Option<SessionOutcome>>,
routes: tokio::sync::watch::Sender<bool>,
ws_collect_ceiling: usize,
) -> Result<(), WsError> {
let initiator = transport.initiator();
let bodies = transport.streams();
let (reader, writer) = transport.split();
let (outbound, outbound_rx) = mpsc::channel(OUTBOUND_QUEUE);
let (deadline, deadline_rx) = tokio::sync::watch::channel(None);
self.core.handle(
clock_now(),
CoreInput::SessionOpened {
session: session.clone(),
initiator,
establish_peer,
expected_peer,
},
)?;
self.sessions.insert(
session.clone(),
SessionIo {
outbound,
byte_budget: Arc::new(Semaphore::new(OUTBOUND_BYTE_BUDGET)),
ws_collect_ceiling,
bodies: bodies.clone(),
handler,
client,
cancellation: cancellation.clone(),
reservation_pending: Arc::new(AtomicBool::new(false)),
establish_peer,
deadline,
routes,
transport_mode: None,
pending_mode_writes: VecDeque::new(),
operation_lanes: Arc::new(std::sync::Mutex::new(HashMap::new())),
operation_slots: Arc::new(Semaphore::new(OPERATION_ACTIVE_MAX)),
},
);
self.outcomes.insert(session.clone(), outcome);
spawn_session_deadline(
session.clone(),
deadline_rx,
self.handle.clone(),
cancellation.clone(),
);
spawn_directives(
session.clone(),
directives,
self.handle.clone(),
cancellation.clone(),
);
spawn_reader(
session.clone(),
reader,
self.handle.clone(),
cancellation.clone(),
bodies,
);
spawn_writer(
session.clone(),
writer,
outbound_rx,
self.handle.clone(),
cancellation,
);
if establish_peer {
spawn_establishment_timeout(session, self.handle.clone());
}
Ok(())
}
async fn directive(&mut self, session: &SessionId, directive: Directive) {
match directive {
Directive::StartClientOperation {
target_path,
kind,
payload,
hops,
headers,
body,
timeout,
sender,
reply,
} => {
let input = if kind == unb_core::Kind::Discover {
if body.is_some() {
Err(unb_core::CoreError::Malformed(
"discovery does not accept a streaming body".into(),
))
} else {
serde_json::from_slice(&payload)
.map(unb_core::OperationInput::Discovery)
.map_err(|error| unb_core::CoreError::Malformed(error.to_string()))
}
} else {
register_outbound_body(&self.handle, session, payload, body)
.map(unb_core::OperationInput::Body)
};
let input = match input {
Ok(input) => input,
Err(error) => {
let _ = reply.send(Err(error));
return;
}
};
let registered_body = match &input {
unb_core::OperationInput::Body(body) => body.clone(),
unb_core::OperationInput::Discovery(_) => None,
};
let result = self.core.handle(
clock_now(),
CoreInput::StartClientOperation {
session: session.clone(),
target_path,
kind,
input,
hops,
headers,
timeout,
},
);
if let Err(error) = result {
if let Some(body) = ®istered_body {
let _ = self.handle.claim_body(session, body.as_str());
}
let _ = reply.send(Err(error));
} else {
self.client_starts
.insert(session.clone(), PendingClientStart { sender, reply });
}
}
Directive::OpenStream {
target_path,
kind,
payload,
hops,
headers,
body,
opened,
reply,
} => {
let result = match register_outbound_body(&self.handle, session, payload, body) {
Ok(body) => {
let result = self.core.open_stream_body(
session,
&target_path,
kind,
body.clone(),
hops,
headers,
);
if result.is_err() {
if let Some(body) = &body {
let _ = self.handle.claim_body(session, body.as_str());
}
}
result
}
Err(error) => Err(error),
};
if let (Ok(corr), Some(opened)) = (&result, opened) {
opened(corr.clone()).await;
}
let _ = reply.send(result);
}
Directive::Send { corr, payload } => {
match register_outbound_body(&self.handle, session, payload, None) {
Ok(body) => {
if self.core.send_body(session, &corr, body.clone()).is_err() {
if let Some(body) = &body {
let _ = self.handle.claim_body(session, body.as_str());
}
}
}
Err(error) => {
let _ = self.core.fail(
session,
&corr,
unb_core::ErrorCode::Busy,
&error.to_string(),
);
}
}
}
Directive::Respond {
corr,
payload,
headers,
body,
} => match register_outbound_body(&self.handle, session, payload, body) {
Ok(body) => {
if self
.core
.respond_body(session, &corr, body.clone(), headers)
.is_err()
{
if let Some(body) = &body {
let _ = self.handle.claim_body(session, body.as_str());
}
}
}
Err(error) => {
let _ = self.core.fail(
session,
&corr,
unb_core::ErrorCode::Busy,
&error.to_string(),
);
}
},
Directive::Fail {
corr,
code,
message,
} => {
let _ = self.core.fail(session, &corr, code, &message);
}
Directive::Cancel { corr } => {
let _ = self.core.cancel(session, &corr);
}
Directive::CancelClientOperation { operation } => {
let _ = self.core.handle(
clock_now(),
CoreInput::CancelClientOperation {
session: session.clone(),
operation,
},
);
}
Directive::Control { kind, payload } => {
let _ = self.core.control(session, kind, payload);
}
}
}
async fn drain(&mut self) {
while let Some(effect) = self.core.poll_effect() {
match effect {
CoreEffect::SendFrame { session, envelope } => {
self.write(&session, envelope, None, false).await
}
CoreEffect::HandshakeEstablished { session, .. } => {
if let Some(bodies) =
self.sessions.get(&session).and_then(|io| io.bodies.clone())
{
let inputs = self.handle.inputs.clone();
let mode_session = session.clone();
n0_future::task::spawn(async move {
let mode = bodies.negotiate().await;
let _ = inputs
.send(ActorInput::TransportModeSelected {
session: mode_session,
mode,
})
.await;
});
}
if self
.sessions
.get(&session)
.is_some_and(|io| !io.establish_peer)
{
if let Some(outcome) = self.outcomes.get(&session) {
let _ = outcome.send(Some(SessionOutcome::Established));
}
}
}
CoreEffect::DeliverClient {
session,
operation,
delivery,
} => {
let body = match &delivery {
CoreClientDelivery::Item(frame) | CoreClientDelivery::Terminal(frame) => {
frame
.body
.as_ref()
.and_then(|body| self.handle.claim_body(&session, body.as_str()))
}
_ => None,
};
let terminal = matches!(
delivery,
CoreClientDelivery::Terminal(_)
| CoreClientDelivery::Cancelled
| CoreClientDelivery::TimedOut
| CoreClientDelivery::SessionClosed
);
let client = self.sessions.get(&session).map(|io| io.client.clone());
if let Some(client) = client {
let delivered = n0_future::time::timeout(
CLIENT_DELIVERY_TIMEOUT,
client.deliver(&operation, delivery, body),
)
.await
.is_ok();
if !delivered {
client.abandon(&operation);
let _ = self.core.handle(
clock_now(),
CoreInput::CancelClientOperation {
session: session.clone(),
operation: operation.clone(),
},
);
}
}
if terminal {
self.handler(&session, HandlerInput::StreamClosed(operation))
.await;
}
}
CoreEffect::RegisterClient { session, operation } => {
let Some(start) = self.client_starts.remove(&session) else {
continue;
};
let client = self.sessions.get(&session).map(|io| io.client.clone());
let result = if let Some(client) = &client {
client
.register(operation.clone(), start.sender)
.map(|()| operation)
.map_err(|error| unb_core::CoreError::UnknownStream(error.to_string()))
} else {
Err(unb_core::CoreError::UnknownSession(session.to_string()))
};
if let Err(Ok(operation)) = start.reply.send(result) {
if let Some(client) = client {
client.abandon(&operation);
}
let _ = self.core.handle(
clock_now(),
CoreInput::CancelClientOperation { session, operation },
);
}
}
CoreEffect::Deliver {
session,
mut envelope,
} => {
if let Some(token) = envelope.body_token.take() {
match self.handle.claim_body(&session, &token) {
Some(WireBody::Bytes(payload)) => {
envelope.payload = payload;
}
Some(body @ WireBody::Stream(_)) => {
let ceiling = self
.sessions
.get(&session)
.map_or(unb_transport::DEFAULT_MAX_FRAME_SIZE, |io| {
io.ws_collect_ceiling
});
match body.collect_to(ceiling).await {
Ok(payload) => envelope.payload = payload,
Err(error) => {
envelope.kind = unb_core::Kind::Error;
envelope.payload =
Envelope::encode_payload(&serde_json::json!({
"code": unb_core::ErrorCode::PayloadTooLarge,
"message": error.to_string(),
}));
}
}
}
None => {}
}
}
self.handler(&session, HandlerInput::Deliver(envelope))
.await
}
CoreEffect::ReleaseBody { session, body } => {
let _ = self.handle.claim_body(&session, body.as_str());
}
CoreEffect::StreamClosed { session, operation } => {
self.handler(&session, HandlerInput::StreamClosed(operation))
.await
}
CoreEffect::CloseTransport {
session,
code,
message,
} => {
let _ = (code, message);
if let Some(io) = self.sessions.get(&session) {
io.cancellation.cancel();
}
}
CoreEffect::ScheduleSessionDeadline { session, deadline } => {
if let Some(io) = self.sessions.get(&session) {
io.deadline.send_replace(Some(deadline));
}
}
CoreEffect::RouteExportAcked { session } => {
if let Some(io) = self.sessions.get(&session) {
io.routes.send_replace(true);
}
}
CoreEffect::ScheduleClientDeadline {
session,
operation,
deadline,
} => spawn_client_timeout(
session.clone(),
operation,
deadline,
self.handle.clone(),
self.sessions
.get(&session)
.map(|io| io.cancellation.clone()),
),
CoreEffect::Send {
effect,
session,
frame,
} => {
if frame.head.kind == unb_core::Kind::Error {
let _ = self.core.handle(
clock_now(),
CoreInput::SendCompleted {
effect,
result: SendResult::Written,
},
);
self.write(&session, frame.into_envelope(), None, false)
.await;
} else {
self.write(&session, frame.into_envelope(), Some(effect), false)
.await;
}
}
CoreEffect::SendProtocol {
effect,
session,
envelope,
} => {
self.write(&session, envelope, Some(effect), false).await;
}
CoreEffect::ContinueDiscovery { stream } => {
let _ = self
.core
.handle(clock_now(), CoreInput::ContinueDiscovery { stream });
}
CoreEffect::SessionRetired { session, reason } => {
#[cfg(feature = "observability")]
{
metrics::counter!("unb_sessions_retired", "reason" => format!("{reason:?}"))
.increment(1);
tracing::debug!(session = %session, reason = ?reason, "session retired");
}
if let Some(outcome) = self.outcomes.remove(&session) {
let _ = outcome.send(Some(SessionOutcome::Retired(reason)));
}
if let Some(io) = self.sessions.remove(&session) {
io.cancellation.cancel();
}
self.pending_bodies
.lock()
.expect("pending body registry")
.retain(|(owner, _), _| owner != &session);
self.handle
.token_bodies
.lock()
.expect("body token registry")
.retain(|(owner, _), _| owner != &session);
self.execute(CoreEffect::SessionRetired { session, reason });
}
CoreEffect::SessionEstablished { session, peer } => {
let effect = CoreEffect::SessionEstablished {
session: session.clone(),
peer,
};
if let Some(input) = self.executor.execute(effect, self.handle.clone()).await {
let _ = self.core.handle(clock_now(), input);
}
if let Some(outcome) = self.outcomes.get(&session) {
let _ = outcome.send(Some(SessionOutcome::Established));
}
}
CoreEffect::ForwardRelay {
effect,
source,
target,
mut frame,
..
} => {
if let Some(body_id) = frame.body.take() {
let Some(body) = self.handle.claim_body(&source.session, body_id.as_str())
else {
let _ = self.core.handle(
clock_now(),
CoreInput::RelayForwardCompleted {
effect,
result: SendResult::Refused {
code: unb_core::ErrorCode::Protocol,
message: "relay body capability is unavailable".into(),
},
},
);
continue;
};
match self.handle.register_body(&target.session, body) {
Ok(target_body) => frame.body = Some(target_body),
Err(error) => {
let _ = self.core.handle(
clock_now(),
CoreInput::RelayForwardCompleted {
effect,
result: SendResult::Refused {
code: unb_core::ErrorCode::Busy,
message: error.to_string(),
},
},
);
continue;
}
}
}
self.write(&target.session, frame.into_envelope(), Some(effect), true)
.await;
}
effect @ (CoreEffect::RouteSnapshotApplied { .. }
| CoreEffect::RouteDeltaApplied { .. }
| CoreEffect::RouteSessionWithdrawn { .. }) => {
if let Some(input) = self.executor.execute(effect, self.handle.clone()).await {
let _ = self.core.handle(clock_now(), input);
}
}
effect => self.execute(effect),
}
}
}
fn execute(&self, effect: CoreEffect) {
let executor = self.executor.clone();
let handle = self.handle.clone();
n0_future::task::spawn(async move {
if let Some(input) = executor.execute(effect, handle.clone()).await {
let _ = handle.submit(input).await;
}
});
}
async fn write(
&mut self,
session: &SessionId,
mut envelope: Envelope,
completion: Option<EffectId>,
relay: bool,
) {
if is_operation_envelope(&envelope) {
let mode = self.sessions.get(session).and_then(|io| io.transport_mode);
match mode {
None if self
.sessions
.get(session)
.is_some_and(|io| io.bodies.is_some()) =>
{
let queued = self.sessions.get_mut(session).is_some_and(|io| {
if io.pending_mode_writes.len() >= PENDING_MODE_WRITES_MAX {
false
} else {
io.pending_mode_writes.push_back(PendingModeWrite {
envelope,
completion,
relay,
});
true
}
});
if !queued {
if let Some(effect) = completion {
self.complete_send(effect, relay, SendResult::ReservationTimedOut);
}
}
return;
}
Some(unb_transport::TransportMode::OperationStreamsV1) => {
self.write_operation(session, envelope, completion, relay)
.await;
return;
}
_ => {}
}
}
if envelope.kind == unb_core::Kind::Cancel {
if let Some(corr) = &envelope.corr {
self.handle.cancel_body_pump(session, corr);
}
}
let body = match (&envelope.body_token, &envelope.corr) {
(Some(token), _) => self.handle.claim_body(session, token),
(None, Some(corr))
if envelope.kind.is_application_request()
|| envelope.kind == unb_core::Kind::Response =>
{
let pending = self
.pending_bodies
.lock()
.expect("pending body registry")
.remove(&(session.clone(), corr.clone()));
if let Some(stream) = pending {
envelope.body_token = Some(mint_body_token());
Some(WireBody::Stream(stream))
} else {
None
}
}
_ => None,
};
let Some(io) = self.sessions.get(session) else {
if let Some(effect) = completion {
self.complete_send(effect, relay, SendResult::Closed);
}
return;
};
if let Some(WireBody::Bytes(payload)) = body {
envelope.body_token = None;
envelope.payload = payload;
} else if let Some(WireBody::Stream(stream)) = body {
if let Some(peer) = io.bodies.clone().filter(|bodies| bodies.can_open()) {
let token = envelope.body_token.clone().expect("tokenized body");
let pump_cancellation = io.cancellation.child_token();
if let Some(corr) = &envelope.corr {
self.handle.register_body_pump(
session,
corr,
&token,
pump_cancellation.clone(),
);
}
spawn_body_pump(
peer,
token,
stream,
io.byte_budget.clone(),
pump_cancellation,
self.handle.clone(),
envelope.corr.clone().map(|corr| (session.clone(), corr)),
);
} else {
envelope.body_token = None;
spawn_body_collect(
self.handle.clone(),
session.clone(),
envelope,
stream,
io.ws_collect_ceiling,
completion,
relay,
io.cancellation.clone(),
);
return;
}
}
let envelope = match io
.byte_budget
.clone()
.try_acquire_many_owned(WriteRequest::weight(&envelope))
{
Ok(budget) => {
let request = WriteRequest {
envelope,
completion,
relay,
_budget: budget,
};
match io.outbound.try_send(request) {
Ok(()) => {
if let Some(effect) = completion {
self.complete_send(effect, relay, SendResult::Reserved);
}
return;
}
Err(mpsc::error::TrySendError::Closed(_)) => {
if let Some(effect) = completion {
self.complete_send(effect, relay, SendResult::Closed);
}
return;
}
Err(mpsc::error::TrySendError::Full(request)) => {
#[cfg(feature = "observability")]
{
metrics::counter!("unb_write_queue_full").increment(1);
tracing::warn!(session = %session, "outbound write queue is full");
}
request.envelope
}
}
}
Err(_) => {
#[cfg(feature = "observability")]
{
metrics::counter!("unb_write_budget_exhausted").increment(1);
tracing::warn!(session = %session, "outbound byte budget is exhausted");
}
envelope
}
};
if completion.is_none()
|| io
.reservation_pending
.compare_exchange(false, true, Ordering::AcqRel, Ordering::Acquire)
.is_ok()
{
let timeout = if matches!(
envelope.kind,
unb_core::Kind::Response | unb_core::Kind::Error | unb_core::Kind::Cancel
) {
TERMINAL_SEND_TIMEOUT
} else {
RELAY_BACKPRESSURE_TIMEOUT
};
reserve_write(
io.outbound.clone(),
io.byte_budget.clone(),
envelope,
completion,
relay,
self.handle.clone(),
io.cancellation.clone(),
completion.map(|_| io.reservation_pending.clone()),
timeout,
);
} else if let Some(effect) = completion {
self.complete_send(effect, relay, SendResult::ReservationTimedOut);
}
}
fn start_operation_acceptor(&self, session: &SessionId) {
let Some(io) = self.sessions.get(session) else {
return;
};
let Some(bodies) = io.bodies.clone() else {
return;
};
spawn_operation_acceptor(
session.clone(),
bodies,
io.operation_lanes.clone(),
io.operation_slots.clone(),
io.byte_budget.clone(),
self.handle.clone(),
io.cancellation.clone(),
);
}
async fn write_operation(
&mut self,
session: &SessionId,
mut envelope: Envelope,
completion: Option<EffectId>,
relay: bool,
) {
let Some(corr) = envelope.corr.clone() else {
if let Some(effect) = completion {
self.complete_send(
effect,
relay,
SendResult::WriteFailed("missing unb-corr".into()),
);
}
return;
};
let body = match &envelope.body_token {
Some(token) => self.handle.claim_body(session, token),
None if envelope.kind.is_application_request()
|| envelope.kind == unb_core::Kind::Response =>
{
self.pending_bodies
.lock()
.expect("pending body registry")
.remove(&(session.clone(), corr.clone()))
.map(WireBody::Stream)
}
None => None,
};
let body = match body {
Some(WireBody::Bytes(payload)) => {
envelope.payload = payload;
None
}
Some(WireBody::Stream(stream)) => Some(stream),
None => None,
};
envelope.body_token = None;
let Some(io) = self.sessions.get(session) else {
if let Some(effect) = completion {
self.complete_send(effect, relay, SendResult::Closed);
}
return;
};
if envelope.kind == unb_core::Kind::Cancel {
let lane = io
.operation_lanes
.lock()
.expect("operation lane registry")
.remove(&corr);
if let Some(control) = lane.and_then(|lane| lane.control) {
let _ = control.reset(unb_transport::operation::RESET_CANCELLED);
let _ = control.stop(unb_transport::operation::RESET_CANCELLED);
}
if let Some(effect) = completion {
self.complete_send(effect, relay, SendResult::Written);
}
return;
}
if envelope.kind.is_application_request() {
spawn_outbound_operation(
session.clone(),
corr,
envelope,
body,
completion,
relay,
io.bodies
.clone()
.expect("operation mode has stream facility"),
io.operation_lanes.clone(),
io.operation_slots.clone(),
io.byte_budget.clone(),
self.handle.clone(),
io.cancellation.clone(),
);
return;
}
let lane = io
.operation_lanes
.lock()
.expect("operation lane registry")
.get(&corr)
.cloned();
let Some(sender) = lane.and_then(|lane| lane.outbound) else {
if let Some(effect) = completion {
self.complete_send(effect, relay, SendResult::Written);
}
return;
};
let budget = match io
.byte_budget
.clone()
.try_acquire_many_owned(WriteRequest::weight(&envelope))
{
Ok(budget) => budget,
Err(_) => {
if let Some(effect) = completion {
self.complete_send(effect, relay, SendResult::ReservationTimedOut);
}
return;
}
};
let request = OperationWriteRequest {
envelope,
completion,
relay,
body,
_budget: budget,
};
match sender.try_send(request) {
Ok(()) => {
if let Some(effect) = completion {
self.complete_send(effect, relay, SendResult::Reserved);
}
}
Err(_) => {
if let Some(effect) = completion {
self.complete_send(effect, relay, SendResult::ReservationTimedOut);
}
}
}
}
fn complete_send(&mut self, effect: EffectId, relay: bool, result: SendResult) {
let _ = self
.core
.handle(clock_now(), send_completion(effect, relay, result));
}
async fn handler(&self, session: &SessionId, event: HandlerInput) {
if let Some(io) = self.sessions.get(session) {
let _ = io.handler.send(event).await;
}
}
}
enum HandlerInput {
Deliver(Envelope),
StreamClosed(ClientOperationId),
}
fn handler_channel<H: SessionHandler>(mut handler: H) -> mpsc::Sender<HandlerInput> {
let (events, mut receiver) = mpsc::channel(SESSION_QUEUE);
n0_future::task::spawn(async move {
while let Some(event) = receiver.recv().await {
match event {
HandlerInput::Deliver(envelope) => handler.deliver(envelope).await,
HandlerInput::StreamClosed(operation) => handler.stream_closed(operation).await,
}
}
});
events
}
fn spawn_directives(
session: SessionId,
mut directives: mpsc::Receiver<Directive>,
handle: ProtocolCoreHandle,
cancellation: CancellationToken,
) {
n0_future::task::spawn(async move {
loop {
let directive = tokio::select! {
biased;
() = cancellation.cancelled() => return,
directive = directives.recv() => directive,
};
let Some(directive) = directive else { return };
if handle
.inputs
.send(ActorInput::Directive(session.clone(), directive))
.await
.is_err()
{
return;
}
}
});
}
fn is_operation_envelope(envelope: &Envelope) -> bool {
envelope.corr.is_some()
&& (envelope.kind.is_application_request()
|| envelope.kind.is_application_response()
|| envelope.kind == unb_core::Kind::Cancel)
}
fn spawn_operation_acceptor(
session: SessionId,
bodies: Arc<unb_transport::ConnectionStreams>,
lanes: OperationLanes,
operation_slots: Arc<Semaphore>,
byte_budget: Arc<Semaphore>,
handle: ProtocolCoreHandle,
cancellation: CancellationToken,
) {
n0_future::task::spawn(async move {
loop {
let operation = tokio::select! {
biased;
() = cancellation.cancelled() => return,
operation = bodies.accept_operation() => operation,
};
let Ok(operation) = operation else {
return;
};
let Ok(operation_slot) = operation_slots.clone().try_acquire_owned() else {
if let Some(control) = operation.control {
let _ = control.reset(unb_transport::operation::RESET_BUSY);
let _ = control.stop(unb_transport::operation::RESET_BUSY);
}
continue;
};
let corr = operation.corr.clone();
let (outbound, receiver) = mpsc::channel(OPERATION_WRITE_QUEUE);
let lane = OperationLane {
outbound: Some(outbound),
control: operation.control.clone(),
};
let admitted = {
let mut lanes = lanes.lock().expect("operation lane registry");
if lanes.contains_key(&corr) {
false
} else {
lanes.insert(corr.clone(), lane);
true
}
};
if !admitted {
if let Some(control) = operation.control {
let _ = control.reset(unb_transport::operation::RESET_PROTOCOL);
}
continue;
}
spawn_accepted_operation(
session.clone(),
operation,
receiver,
lanes.clone(),
byte_budget.clone(),
handle.clone(),
cancellation.child_token(),
operation_slot,
);
}
});
}
#[allow(clippy::too_many_arguments)]
fn spawn_outbound_operation(
session: SessionId,
corr: String,
envelope: Envelope,
body: Option<BodyStream>,
completion: Option<EffectId>,
relay: bool,
bodies: Arc<unb_transport::ConnectionStreams>,
lanes: OperationLanes,
operation_slots: Arc<Semaphore>,
byte_budget: Arc<Semaphore>,
handle: ProtocolCoreHandle,
cancellation: CancellationToken,
) {
n0_future::task::spawn(async move {
let open_byte_budget = byte_budget.clone();
let open = async {
let operation_slot = operation_slots
.acquire_owned()
.await
.map_err(|_| unb_transport::TransportError::Closed)?;
let budget = open_byte_budget
.acquire_many_owned(WriteRequest::weight(&envelope))
.await
.map_err(|_| unb_transport::TransportError::Closed)?;
let operation = bodies
.open_operation(&corr)
.await?
.ok_or(unb_transport::TransportError::Closed)?;
Ok::<_, unb_transport::TransportError>((operation, budget, operation_slot))
};
let opened = tokio::select! {
biased;
() = cancellation.cancelled() => return,
opened = n0_future::time::timeout(RELAY_BACKPRESSURE_TIMEOUT, open) => opened,
};
let Ok(Ok((operation, _budget, _operation_slot))) = opened else {
if let Some(effect) = completion {
let _ = handle
.submit(send_completion(
effect,
relay,
SendResult::ReservationTimedOut,
))
.await;
}
return;
};
lanes.lock().expect("operation lane registry").insert(
corr.clone(),
OperationLane {
outbound: None,
control: operation.control.clone(),
},
);
if let Some(effect) = completion {
if handle
.submit(send_completion(effect, relay, SendResult::Reserved))
.await
.is_err()
{
return;
}
}
let mut writer = unb_transport::operation::OperationRecordWriter::new(operation.write);
let sent = write_operation_envelope(
&mut writer,
envelope,
body,
byte_budget.clone(),
cancellation.clone(),
)
.await;
let sent = match sent {
Ok(()) => writer.finish().await,
Err(error) => Err(error),
};
if let Some(effect) = completion {
let _ = handle
.submit(send_completion(effect, relay, SendResult::Written))
.await;
}
if let Err(error) = sent {
if !operation_peer_stopped_cancelled(&error) {
if let Some(control) = &operation.control {
let _ = control.reset(error.reset_code());
}
let _ = submit_operation_end(
&handle,
&session,
&corr,
OperationDirection::Opening,
operation_outcome(&error),
)
.await;
lanes.lock().expect("operation lane registry").remove(&corr);
return;
}
}
let result = read_operation_direction(
&session,
&corr,
operation.read,
OperationDirection::Return,
&handle,
&cancellation,
)
.await;
if let Err(error) = result {
if let Some(control) = &operation.control {
let _ = control.reset(error.reset_code());
let _ = control.stop(error.reset_code());
}
let _ = submit_operation_end(
&handle,
&session,
&corr,
OperationDirection::Return,
operation_outcome(&error),
)
.await;
}
lanes.lock().expect("operation lane registry").remove(&corr);
});
}
#[allow(clippy::too_many_arguments)]
fn spawn_accepted_operation(
session: SessionId,
operation: unb_transport::operation::OperationIo,
receiver: mpsc::Receiver<OperationWriteRequest>,
lanes: OperationLanes,
byte_budget: Arc<Semaphore>,
handle: ProtocolCoreHandle,
cancellation: CancellationToken,
operation_slot: OwnedSemaphorePermit,
) {
let corr = operation.corr.clone();
let control = operation.control.clone();
let read_corr = corr.clone();
let read_lanes = lanes.clone();
let read_handle = handle.clone();
let read_cancellation = cancellation.clone();
let writer_session = session.clone();
n0_future::task::spawn(async move {
let result = read_operation_direction(
&session,
&read_corr,
operation.read,
OperationDirection::Opening,
&read_handle,
&read_cancellation,
)
.await;
if let Err(error) = result {
if operation_body_consumer_closed(&error) {
if let Some(control) = control {
let _ = control.stop(unb_transport::operation::RESET_CANCELLED);
}
return;
}
if let Some(control) = control {
let _ = control.reset(error.reset_code());
let _ = control.stop(error.reset_code());
}
let _ = submit_operation_end(
&read_handle,
&session,
&read_corr,
OperationDirection::Opening,
operation_outcome(&error),
)
.await;
read_lanes
.lock()
.expect("operation lane registry")
.remove(&read_corr);
}
});
spawn_operation_writer(
writer_session,
corr,
operation.write,
operation.control,
receiver,
lanes,
byte_budget,
handle,
cancellation,
operation_slot,
);
}
#[allow(clippy::too_many_arguments)]
fn spawn_operation_writer(
session: SessionId,
corr: String,
write: Box<dyn tokio::io::AsyncWrite + Send + Unpin>,
control: Option<unb_transport::operation::OperationControl>,
mut receiver: mpsc::Receiver<OperationWriteRequest>,
lanes: OperationLanes,
byte_budget: Arc<Semaphore>,
handle: ProtocolCoreHandle,
cancellation: CancellationToken,
_operation_slot: OwnedSemaphorePermit,
) {
n0_future::task::spawn(async move {
let mut writer = unb_transport::operation::OperationRecordWriter::new(write);
loop {
let request = tokio::select! {
biased;
() = cancellation.cancelled() => break,
request = receiver.recv() => request,
};
let Some(request) = request else {
break;
};
let terminal = matches!(
request.envelope.kind,
unb_core::Kind::Response | unb_core::Kind::Error
);
let result = write_operation_envelope(
&mut writer,
request.envelope,
request.body,
byte_budget.clone(),
cancellation.clone(),
)
.await;
let result = if result.is_ok() && terminal {
writer.finish().await
} else {
result
};
if let Some(effect) = request.completion {
let _ = handle
.submit(send_completion(effect, request.relay, SendResult::Written))
.await;
}
if result.is_err() || terminal {
if let Err(error) = result {
if let Some(control) = &control {
let _ = control.reset(error.reset_code());
}
let _ = submit_operation_end(
&handle,
&session,
&corr,
OperationDirection::Return,
operation_outcome(&error),
)
.await;
}
break;
}
}
lanes.lock().expect("operation lane registry").remove(&corr);
});
}
async fn write_operation_envelope<W>(
writer: &mut unb_transport::operation::OperationRecordWriter<W>,
mut envelope: Envelope,
body: Option<BodyStream>,
byte_budget: Arc<Semaphore>,
cancellation: CancellationToken,
) -> Result<(), unb_transport::operation::OperationCodecError>
where
W: tokio::io::AsyncWrite + Unpin,
{
let Some(body) = body else {
return writer
.write_record(
unb_transport::operation::OperationRecordKind::CompleteMessage,
&envelope.encode(),
)
.await;
};
envelope.body_token = None;
envelope.payload = Bytes::new();
writer
.write_record(
unb_transport::operation::OperationRecordKind::Head,
&envelope.encode(),
)
.await?;
let mut body = MeteredBody::new(body, byte_budget, cancellation.clone());
loop {
let chunk = body.next_chunk().await.map_err(|error| {
unb_transport::operation::OperationCodecError::Io(error.to_string())
})?;
let Some(chunk) = chunk else {
if cancellation.is_cancelled() {
return Err(unb_transport::operation::OperationCodecError::Io(
"operation body cancelled".into(),
));
}
break;
};
for piece in chunk.bytes.chunks(OPERATION_BODY_CHUNK_BYTES) {
writer
.write_record(
unb_transport::operation::OperationRecordKind::BodyChunk,
piece,
)
.await?;
}
}
writer
.write_record(unb_transport::operation::OperationRecordKind::BodyEnd, &[])
.await
}
#[derive(Clone, Copy)]
enum OperationDirection {
Opening,
Return,
}
async fn read_operation_direction(
session: &SessionId,
corr: &str,
read: Box<dyn tokio::io::AsyncRead + Send + Unpin>,
direction: OperationDirection,
handle: &ProtocolCoreHandle,
cancellation: &CancellationToken,
) -> Result<(), unb_transport::operation::OperationCodecError> {
let mut reader = unb_transport::operation::OperationRecordReader::new(read);
let mut records = 0usize;
let mut terminal = false;
let mut body_sender: Option<mpsc::Sender<Result<Bytes, unb_core::CoreError>>> = None;
loop {
let kind = tokio::select! {
biased;
() = cancellation.cancelled() => {
if let Some(sender) = body_sender.take() {
let _ = sender.send(Err(unb_core::CoreError::Malformed(
"operation body cancelled".into(),
))).await;
}
return Ok(())
},
record = reader.read_record() => record,
};
let (kind, payload) = match kind {
Ok(record) => record,
Err(error) => {
if let Some(sender) = body_sender.take() {
let _ = sender
.send(Err(unb_core::CoreError::Malformed(error.to_string())))
.await;
}
return Err(error);
}
};
match kind {
unb_transport::operation::OperationRecordKind::CompleteMessage => {
if body_sender.is_some() {
return Err(unb_transport::operation::OperationCodecError::Protocol(
"complete message arrived before body end".into(),
));
}
let envelope = decode_operation_message(payload, corr, direction, records)?;
terminal |= matches!(
envelope.kind,
unb_core::Kind::Response | unb_core::Kind::Error
);
records += 1;
let input = application_core_input(handle, session, envelope).map_err(|error| {
unb_transport::operation::OperationCodecError::Protocol(error.to_string())
})?;
handle.submit(input).await.map_err(|error| {
unb_transport::operation::OperationCodecError::Io(error.to_string())
})?;
}
unb_transport::operation::OperationRecordKind::Head => {
if body_sender.is_some() {
return Err(unb_transport::operation::OperationCodecError::Protocol(
"operation head arrived before body end".into(),
));
}
let mut envelope = decode_operation_message(payload, corr, direction, records)?;
let (token, sender) = register_operation_body(handle, session)?;
envelope.body_token = Some(token);
terminal |= matches!(
envelope.kind,
unb_core::Kind::Response | unb_core::Kind::Error
);
records += 1;
body_sender = Some(sender);
let input = application_core_input(handle, session, envelope).map_err(|error| {
unb_transport::operation::OperationCodecError::Protocol(error.to_string())
})?;
handle.submit(input).await.map_err(|error| {
unb_transport::operation::OperationCodecError::Io(error.to_string())
})?;
}
unb_transport::operation::OperationRecordKind::BodyChunk => {
let Some(sender) = &body_sender else {
return Err(unb_transport::operation::OperationCodecError::Protocol(
"body chunk arrived without a head".into(),
));
};
sender.send(Ok(payload)).await.map_err(|_| {
unb_transport::operation::OperationCodecError::Io(
"operation body consumer closed".into(),
)
})?;
}
unb_transport::operation::OperationRecordKind::BodyEnd => {
let Some(sender) = body_sender.take() else {
return Err(unb_transport::operation::OperationCodecError::Protocol(
"body end arrived without a head".into(),
));
};
drop(sender);
}
unb_transport::operation::OperationRecordKind::DirectionEnd => {
if body_sender.is_some() {
return Err(unb_transport::operation::OperationCodecError::Truncated(
"direction ended before body end".into(),
));
}
if records == 0 || matches!(direction, OperationDirection::Return) && !terminal {
return Err(unb_transport::operation::OperationCodecError::Protocol(
"operation direction ended before its required message".into(),
));
}
reader.expect_fin().await?;
submit_operation_end(
handle,
session,
corr,
direction,
unb_core::OperationStreamOutcome::Clean,
)
.await
.map_err(|error| {
unb_transport::operation::OperationCodecError::Io(error.to_string())
})?;
return Ok(());
}
}
}
}
fn register_operation_body(
handle: &ProtocolCoreHandle,
session: &SessionId,
) -> Result<
(String, mpsc::Sender<Result<Bytes, unb_core::CoreError>>),
unb_transport::operation::OperationCodecError,
> {
let token = mint_body_token();
let (sender, receiver) = mpsc::channel(8);
let stream: BodyStream = Box::pin(futures_util::stream::unfold(
receiver,
|mut receiver| async move { receiver.recv().await.map(|chunk| (chunk, receiver)) },
));
let mut registry = handle.token_bodies.lock().expect("body token registry");
let owned = registry
.keys()
.filter(|(owner, _)| owner == session)
.count();
if owned >= SESSION_BODY_TOKENS_MAX {
return Err(unb_transport::operation::OperationCodecError::Protocol(
"session operation body limit reached".into(),
));
}
registry.insert((session.clone(), token.clone()), WireBody::Stream(stream));
Ok((token, sender))
}
fn operation_outcome(
error: &unb_transport::operation::OperationCodecError,
) -> unb_core::OperationStreamOutcome {
match error {
unb_transport::operation::OperationCodecError::Protocol(message) => {
unb_core::OperationStreamOutcome::Protocol(message.clone())
}
unb_transport::operation::OperationCodecError::Truncated(_) => {
unb_core::OperationStreamOutcome::Truncated
}
unb_transport::operation::OperationCodecError::Oversized { size, max } => {
unb_core::OperationStreamOutcome::Protocol(format!(
"operation record of {size} bytes exceeds max {max}"
))
}
unb_transport::operation::OperationCodecError::Io(message) => {
unb_core::OperationStreamOutcome::Transport(message.clone())
}
}
}
fn operation_body_consumer_closed(error: &unb_transport::operation::OperationCodecError) -> bool {
matches!(
error,
unb_transport::operation::OperationCodecError::Io(message)
if message == "operation body consumer closed"
)
}
fn operation_peer_stopped_cancelled(error: &unb_transport::operation::OperationCodecError) -> bool {
matches!(
error,
unb_transport::operation::OperationCodecError::Io(message)
if message.contains(&unb_transport::operation::RESET_CANCELLED.to_string())
)
}
async fn submit_operation_end(
handle: &ProtocolCoreHandle,
session: &SessionId,
corr: &str,
direction: OperationDirection,
outcome: unb_core::OperationStreamOutcome,
) -> Result<(), WsError> {
handle
.submit(CoreInput::OperationStreamEnded {
session: session.clone(),
corr: unb_core::CorrelationId::from(corr.to_owned()),
direction: match direction {
OperationDirection::Opening => unb_core::OperationStreamDirection::Opening,
OperationDirection::Return => unb_core::OperationStreamDirection::Return,
},
outcome,
})
.await
}
fn decode_operation_message(
payload: Bytes,
corr: &str,
direction: OperationDirection,
records: usize,
) -> Result<Envelope, unb_transport::operation::OperationCodecError> {
let envelope = Envelope::decode(payload).map_err(|error| {
unb_transport::operation::OperationCodecError::Protocol(error.to_string())
})?;
if envelope.corr.as_deref() != Some(corr) {
return Err(unb_transport::operation::OperationCodecError::Protocol(
"operation record changed unb-corr".into(),
));
}
match direction {
OperationDirection::Opening if records != 0 || !envelope.kind.is_application_request() => {
Err(unb_transport::operation::OperationCodecError::Protocol(
"opening direction requires exactly one application opener".into(),
))
}
OperationDirection::Return if !envelope.kind.is_application_response() => {
Err(unb_transport::operation::OperationCodecError::Protocol(
"return direction requires application responses".into(),
))
}
_ => Ok(envelope),
}
}
fn spawn_body_pump(
peer: std::sync::Arc<unb_transport::ConnectionStreams>,
token: String,
stream: BodyStream,
budget: Arc<Semaphore>,
cancellation: CancellationToken,
handle: ProtocolCoreHandle,
registration: Option<(SessionId, String)>,
) {
n0_future::task::spawn(async move {
let pump = async {
let Ok(mut writer) = peer.open_body(&token).await else {
return;
};
let mut metered = MeteredBody::new(stream, budget, cancellation.clone());
loop {
match metered.next_chunk().await {
Ok(Some(chunk)) => {
if writer.write_chunk(&chunk.bytes).await.is_err() {
return;
}
}
Ok(None) => {
if cancellation.is_cancelled() {
writer.abort().await;
} else {
let _ = writer.finish().await;
}
return;
}
Err(_) => {
writer.abort().await;
return;
}
}
}
};
pump.await;
if let Some((session, corr)) = registration {
handle.release_body_pump(&session, &corr, &token);
}
});
}
#[allow(clippy::too_many_arguments)]
fn spawn_body_collect(
handle: ProtocolCoreHandle,
session: SessionId,
mut envelope: Envelope,
stream: BodyStream,
ceiling: usize,
completion: Option<EffectId>,
relay: bool,
cancellation: CancellationToken,
) {
n0_future::task::spawn(async move {
let collected = tokio::select! {
biased;
() = cancellation.cancelled() => {
if let Some(effect) = completion {
let _ = handle
.submit(send_completion(effect, relay, SendResult::Cancelled))
.await;
}
return;
}
collected = WireBody::Stream(stream).collect_to(ceiling) => collected,
};
match collected {
Ok(payload) => {
envelope.payload = payload;
let _ = handle
.inputs
.send(ActorInput::Write {
session,
envelope,
completion,
relay,
})
.await;
}
Err(error) => {
let (code, message) = match error {
unb_core::CoreError::BodyTooLarge(_) => (
unb_core::ErrorCode::PayloadTooLarge,
"streaming body exceeds this route's WS collect ceiling; raise the \
ceiling or reach this node over WebTransport"
.to_string(),
),
other => (unb_core::ErrorCode::Protocol, other.to_string()),
};
if let Some(effect) = completion {
let refused = SendResult::Refused { code, message };
let _ = handle.submit(send_completion(effect, relay, refused)).await;
} else if envelope.kind.is_application_request() {
if let Some(corr) = envelope.corr.clone() {
let outcome = if code == unb_core::ErrorCode::PayloadTooLarge {
unb_core::OperationStreamOutcome::PayloadTooLarge(message)
} else {
unb_core::OperationStreamOutcome::Protocol(message)
};
let _ = handle
.submit(CoreInput::OperationStreamEnded {
session,
corr: unb_core::CorrelationId::from(corr),
direction: unb_core::OperationStreamDirection::Opening,
outcome,
})
.await;
}
} else if envelope.kind == unb_core::Kind::Response {
envelope.kind = unb_core::Kind::Error;
envelope.payload = Envelope::encode_payload(&serde_json::json!({
"code": code,
"message": message,
}));
let _ = handle
.inputs
.send(ActorInput::Write {
session,
envelope,
completion: None,
relay: false,
})
.await;
}
}
}
});
}
fn spawn_reader(
session: SessionId,
mut reader: PipeReader,
handle: ProtocolCoreHandle,
cancellation: CancellationToken,
bodies: crate::transport::SessionStreams,
) {
n0_future::task::spawn(async move {
loop {
let received = tokio::select! {
biased;
() = cancellation.cancelled() => {
let _ = handle.submit(CoreInput::SessionClosed { session }).await;
return;
},
received = reader.recv() => received,
};
let envelope = match received {
Ok(Some((envelope, _body))) => {
if let Some(token) = &envelope.body_token {
if !(envelope.kind.is_application_request()
|| envelope.kind == unb_core::Kind::Response)
{
let _ = handle.submit(CoreInput::SessionFailed { session }).await;
return;
}
{
let Some(peer) = bodies.clone() else {
let _ = handle.submit(CoreInput::SessionFailed { session }).await;
return;
};
let lazy_token = token.clone();
let accepted = futures_util::stream::once(async move {
peer.accept_body(&lazy_token, RELAY_BACKPRESSURE_TIMEOUT)
.await
});
let stream: BodyStream =
Box::pin(futures_util::StreamExt::flat_map(accepted, |outcome| {
match outcome {
Ok(reader) => futures_util::StreamExt::boxed(
futures_util::StreamExt::map(reader, |chunk| {
chunk.map_err(|error| {
unb_core::CoreError::Malformed(
error.to_string(),
)
})
}),
),
Err(error) => futures_util::StreamExt::boxed(
futures_util::stream::once(async move {
Err(unb_core::CoreError::Malformed(
error.to_string(),
))
}),
),
}
}));
let rejected = {
let mut registry =
handle.token_bodies.lock().expect("body token registry");
let key = (session.clone(), token.clone());
let owned = registry
.keys()
.filter(|(owner, _)| owner == &session)
.count();
if registry.contains_key(&key) || owned >= SESSION_BODY_TOKENS_MAX {
true
} else {
registry.insert(key, WireBody::Stream(stream));
false
}
};
if rejected {
let _ = handle.submit(CoreInput::SessionFailed { session }).await;
return;
}
}
}
envelope
}
Ok(None) => {
let _ = handle.submit(CoreInput::SessionClosed { session }).await;
return;
}
Err(_) => {
#[cfg(feature = "observability")]
{
metrics::counter!("unb_transport_failures").increment(1);
tracing::warn!(session = %session, "transport read failed");
}
let _ = handle.submit(CoreInput::SessionFailed { session }).await;
return;
}
};
if envelope.kind == unb_core::Kind::Cancel {
if let Some(corr) = &envelope.corr {
handle.cancel_body_pump(&session, corr);
}
}
let input = match application_core_input(&handle, &session, envelope) {
Ok(input) => input,
Err(_) => {
let _ = handle.submit(CoreInput::SessionFailed { session }).await;
return;
}
};
if handle.submit(input).await.is_err() {
return;
}
}
});
}
fn spawn_writer(
session: SessionId,
mut writer: PipeWriter,
mut outbound: mpsc::Receiver<WriteRequest>,
handle: ProtocolCoreHandle,
cancellation: CancellationToken,
) {
n0_future::task::spawn(async move {
loop {
let request = tokio::select! {
biased;
() = cancellation.cancelled() => {
let farewell = async {
while let Ok(request) = outbound.try_recv() {
if writer.feed(request.envelope).await.is_err() {
break;
}
}
let _ = writer.flush().await;
let _ = writer.close().await;
};
let _ = n0_future::time::timeout(TERMINAL_SEND_TIMEOUT, farewell).await;
return;
},
request = outbound.recv() => request,
};
let Some(request) = request else {
let _ = handle
.submit(CoreInput::SessionClosed {
session: session.clone(),
})
.await;
return;
};
let mut batch = Vec::with_capacity(WRITE_BATCH);
batch.push(request);
while batch.len() < WRITE_BATCH {
match outbound.try_recv() {
Ok(request) => batch.push(request),
Err(_) => break,
}
}
let completions = batch
.iter()
.filter_map(|request| request.completion.map(|effect| (effect, request.relay)))
.collect::<Vec<_>>();
let result = async {
for request in batch {
writer.feed(request.envelope).await?;
}
writer.flush().await
}
.await;
let send_result = result.as_ref().map_or_else(
|error| SendResult::WriteFailed(error.to_string()),
|_| SendResult::Written,
);
let completions = completions
.into_iter()
.map(|(effect, relay)| send_completion(effect, relay, send_result.clone()))
.collect::<Vec<_>>();
if !completions.is_empty()
&& handle
.inputs
.send(ActorInput::CoreBatch(completions))
.await
.is_err()
{
return;
}
if result.is_err() {
let _ = handle
.submit(CoreInput::SessionClosed {
session: session.clone(),
})
.await;
cancellation.cancel();
let _ = writer.close().await;
return;
}
}
});
}
#[allow(clippy::too_many_arguments)]
fn reserve_write(
sender: mpsc::Sender<WriteRequest>,
byte_budget: Arc<Semaphore>,
envelope: Envelope,
completion: Option<EffectId>,
relay: bool,
handle: ProtocolCoreHandle,
cancellation: CancellationToken,
reservation_pending: Option<Arc<AtomicBool>>,
timeout: std::time::Duration,
) {
n0_future::task::spawn(async move {
let effect = completion;
let acquire = async {
let budget = byte_budget
.acquire_many_owned(WriteRequest::weight(&envelope))
.await
.map_err(|_| ())?;
let slot = sender.reserve_owned().await.map_err(|_| ())?;
Ok::<_, ()>((budget, slot))
};
let reserved = tokio::select! {
biased;
() = cancellation.cancelled() => {
if let Some(pending) = &reservation_pending {
pending.store(false, Ordering::Release);
}
if let Some(effect) = effect {
let _ = handle.submit(send_completion(effect, relay, SendResult::Cancelled)).await;
}
return;
}
reserved = n0_future::time::timeout(timeout, acquire) => reserved,
};
if let Some(pending) = &reservation_pending {
pending.store(false, Ordering::Release);
}
match reserved {
Ok(Ok((budget, slot))) => {
if let Some(effect) = effect {
if handle
.submit(send_completion(effect, relay, SendResult::Reserved))
.await
.is_err()
{
return;
}
}
slot.send(WriteRequest {
envelope,
completion,
relay,
_budget: budget,
});
}
Ok(Err(())) => {
if let Some(effect) = effect {
let _ = handle
.submit(send_completion(effect, relay, SendResult::Closed))
.await;
}
}
Err(_) => {
if let Some(effect) = effect {
let _ = handle
.submit(send_completion(
effect,
relay,
SendResult::ReservationTimedOut,
))
.await;
}
}
}
});
}
fn send_completion(effect: EffectId, relay: bool, result: SendResult) -> CoreInput {
if relay {
CoreInput::RelayForwardCompleted { effect, result }
} else {
CoreInput::SendCompleted { effect, result }
}
}
fn spawn_session_deadline(
session: SessionId,
mut deadlines: tokio::sync::watch::Receiver<Option<Instant>>,
handle: ProtocolCoreHandle,
cancellation: CancellationToken,
) {
n0_future::task::spawn(async move {
loop {
let deadline = *deadlines.borrow_and_update();
match deadline {
Some(deadline) if clock_now() < deadline => {
tokio::select! {
biased;
() = cancellation.cancelled() => return,
() = sleep_to(deadline) => {}
}
}
Some(_) => {
if handle
.submit(CoreInput::SessionTimeout {
session: session.clone(),
})
.await
.is_err()
{
return;
}
tokio::select! {
biased;
() = cancellation.cancelled() => return,
changed = deadlines.changed() => {
if changed.is_err() {
return;
}
}
}
}
None => {
tokio::select! {
biased;
() = cancellation.cancelled() => return,
changed = deadlines.changed() => {
if changed.is_err() {
return;
}
}
}
}
}
}
});
}
fn spawn_client_timeout(
session: SessionId,
operation: ClientOperationId,
deadline: Instant,
handle: ProtocolCoreHandle,
cancellation: Option<CancellationToken>,
) {
let Some(cancellation) = cancellation else {
return;
};
n0_future::task::spawn(async move {
tokio::select! {
biased;
() = cancellation.cancelled() => {}
() = sleep_to(deadline) => {
let _ = handle.submit(CoreInput::ClientOperationTimeout { session, operation }).await;
}
}
});
}
fn spawn_establishment_timeout(session: SessionId, handle: ProtocolCoreHandle) {
n0_future::task::spawn(async move {
n0_future::time::sleep(ESTABLISHMENT_TIMEOUT).await;
let _ = handle
.submit(CoreInput::EstablishmentTimeout { session })
.await;
});
}
#[cfg(test)]
mod tests {
use super::*;
use bytes::Bytes;
use std::sync::atomic::{AtomicBool, Ordering};
use unb_core::{DiscoverPlan, Kind, RetirementReason, PROTOCOL_VERSION};
struct RecordingExecutor {
effects: mpsc::Sender<CoreEffect>,
}
impl EffectExecutor for RecordingExecutor {
fn execute(&self, effect: CoreEffect, _handle: ProtocolCoreHandle) -> EffectFuture {
let effects = self.effects.clone();
Box::pin(async move {
let _ = effects.send(effect).await;
None
})
}
}
struct PublicationExecutor {
published: Arc<AtomicBool>,
effects: mpsc::Sender<bool>,
}
struct ClosingHandler {
delivered: mpsc::Sender<Envelope>,
closed: mpsc::Sender<ClientOperationId>,
}
impl SessionHandler for ClosingHandler {
async fn deliver(&mut self, envelope: Envelope) {
let _ = self.delivered.send(envelope).await;
}
async fn stream_closed(&mut self, operation: ClientOperationId) {
let _ = self.closed.send(operation).await;
}
}
impl EffectExecutor for PublicationExecutor {
fn execute(&self, effect: CoreEffect, _handle: ProtocolCoreHandle) -> EffectFuture {
let published = self.published.clone();
let effects = self.effects.clone();
Box::pin(async move {
if matches!(effect, CoreEffect::SessionRetired { .. }) {
let _ = effects.send(published.load(Ordering::SeqCst)).await;
}
None
})
}
}
fn local_pipe() -> (Pipe, mpsc::Sender<Envelope>, mpsc::Receiver<Envelope>) {
let (incoming, rx) = mpsc::channel(16);
let (tx, outgoing) = mpsc::channel(16);
(
Pipe::Local {
rx,
tx,
initiator: true,
},
incoming,
outgoing,
)
}
fn frame(kind: Kind, corr: Option<String>, payload: serde_json::Value) -> Envelope {
Envelope {
v: PROTOCOL_VERSION,
id: format!("{kind:?}"),
target: String::new(),
subject: String::new(),
kind,
corr,
seq: None,
hops: None,
body_token: None,
payload: Envelope::encode_payload(&payload),
path: Vec::new(),
headers: Default::default(),
}
}
async fn establish(incoming: &mpsc::Sender<Envelope>, outgoing: &mut mpsc::Receiver<Envelope>) {
assert_eq!(outgoing.recv().await.unwrap().kind, Kind::Hello);
incoming
.send(frame(
Kind::Welcome,
None,
serde_json::json!({ "version": 1 }),
))
.await
.unwrap();
}
fn request(target_path: &str, body: &'static [u8]) -> http::Request<Bytes> {
http::Request::post(target_path)
.body(Bytes::from_static(body))
.unwrap()
}
async fn attach_client(handle: &ProtocolCoreHandle, name: &str, pipe: Pipe) -> Arc<Wire> {
attach_client_with_ceiling(handle, name, pipe, unb_transport::DEFAULT_MAX_FRAME_SIZE).await
}
async fn attach_client_with_ceiling(
handle: &ProtocolCoreHandle,
name: &str,
pipe: Pipe,
ws_collect_ceiling: usize,
) -> Arc<Wire> {
let session_cancellation = handle.cancellation.child_token();
let (wire, directives, outcome, routes) =
Wire::standalone_channel(session_cancellation.clone());
let (reply, response) = oneshot::channel();
handle
.inputs
.send(ActorInput::Attach {
session: SessionId::from(name),
transport: pipe,
expected_peer: None,
directives,
handler: handler_channel(wire.observation_sender()),
client: wire.client_session(),
cancellation: session_cancellation,
establish_peer: false,
outcome,
routes,
ws_collect_ceiling,
reply,
})
.await
.unwrap();
response.await.unwrap().unwrap();
wire
}
fn spawn_core() -> ProtocolCoreHandle {
ProtocolCoreHandle::spawn(
ProtocolCore::new("node"),
Arc::new(RecordingExecutor {
effects: mpsc::channel(1).0,
}),
CancellationToken::new(),
&tokio::runtime::Handle::current(),
)
}
fn stalled_pipe() -> (Pipe, mpsc::Sender<Envelope>, mpsc::Receiver<Envelope>) {
let (incoming, rx) = mpsc::channel(16);
let (tx, outgoing) = mpsc::channel(1);
(
Pipe::Local {
rx,
tx,
initiator: true,
},
incoming,
outgoing,
)
}
fn chunk_stream(chunks: Vec<Result<Bytes, unb_core::CoreError>>) -> BodyStream {
Box::pin(futures_util::stream::iter(chunks))
}
#[tokio::test(start_paused = true)]
async fn a_metered_body_stall_past_the_deadline_is_busy_shaped() {
let budget = Arc::new(Semaphore::new(4));
let hold = budget.clone().try_acquire_many_owned(4).unwrap();
let mut body = MeteredBody::new(
chunk_stream(vec![Ok(Bytes::from_static(b"x"))]),
budget,
CancellationToken::new(),
);
let outcome = body.next_chunk().await;
assert!(
matches!(
outcome,
Err(WsError::Core(unb_core::CoreError::BodyStalled))
),
"budget starvation past the deadline must be BUSY-shaped: {outcome:?}"
);
drop(hold);
}
#[tokio::test]
async fn a_window_saturated_body_leaves_a_sibling_pump_undisturbed() {
let budget = Arc::new(Semaphore::new(OUTBOUND_BYTE_BUDGET));
let mut saturated = MeteredBody::new(
chunk_stream(vec![
Ok(Bytes::from(vec![0u8; BODY_WINDOW_BYTES])),
Ok(Bytes::from_static(b"tail")),
]),
budget.clone(),
CancellationToken::new(),
);
let held = saturated.next_chunk().await.unwrap().unwrap();
assert!(
tokio::time::timeout(std::time::Duration::from_millis(50), saturated.next_chunk())
.await
.is_err(),
"a full per-body window must block the next chunk"
);
let mut sibling = MeteredBody::new(
chunk_stream(vec![Ok(Bytes::from_static(b"y"))]),
budget,
CancellationToken::new(),
);
let chunk = tokio::time::timeout(std::time::Duration::from_secs(1), sibling.next_chunk())
.await
.expect("a sibling body must progress while another window is saturated")
.unwrap()
.unwrap();
assert_eq!(chunk.bytes, Bytes::from_static(b"y"));
drop(held);
let tail = tokio::time::timeout(std::time::Duration::from_secs(1), saturated.next_chunk())
.await
.expect("draining the window must unblock the body")
.unwrap()
.unwrap();
assert_eq!(tail.bytes, Bytes::from_static(b"tail"));
}
#[test]
fn body_tokens_are_unique_and_carry_entropy() {
let first = mint_body_token();
let second = mint_body_token();
assert_ne!(first, second);
assert!(
first.len() >= 33,
"a guessable serial token invites body substitution: {first}"
);
}
#[tokio::test]
async fn an_oversized_source_chunk_is_rejected_before_retention() {
let budget = Arc::new(Semaphore::new(OUTBOUND_BYTE_BUDGET));
let mut body = MeteredBody::new(
chunk_stream(vec![Ok(Bytes::from(vec![
0u8;
unb_transport::DEFAULT_MAX_FRAME_SIZE
+ 1
]))]),
budget.clone(),
CancellationToken::new(),
);
let outcome = body.next_chunk().await;
assert!(
matches!(
outcome,
Err(WsError::Core(unb_core::CoreError::BodyTooLarge(_)))
),
"a chunk beyond the frame cap must be refused before retention: {outcome:?}"
);
assert_eq!(
budget.available_permits(),
OUTBOUND_BYTE_BUDGET,
"a refused chunk must not consume budget"
);
}
#[tokio::test]
async fn concurrent_bodies_stay_within_the_session_budget() {
let budget = Arc::new(Semaphore::new(100));
let mut first = MeteredBody::new(
chunk_stream(vec![Ok(Bytes::from(vec![0u8; 60]))]),
budget.clone(),
CancellationToken::new(),
);
let mut second = MeteredBody::new(
chunk_stream(vec![Ok(Bytes::from(vec![0u8; 60]))]),
budget.clone(),
CancellationToken::new(),
);
let held = first.next_chunk().await.unwrap().unwrap();
assert_eq!(budget.available_permits(), 40);
assert!(
tokio::time::timeout(std::time::Duration::from_millis(50), second.next_chunk())
.await
.is_err(),
"combined queued bytes must not exceed the session budget"
);
drop(held);
let chunk = tokio::time::timeout(std::time::Duration::from_secs(1), second.next_chunk())
.await
.expect("released budget must admit the waiting body")
.unwrap()
.unwrap();
assert_eq!(chunk.bytes.len(), 60);
}
#[tokio::test(start_paused = true)]
async fn a_steady_body_streams_past_the_terminal_deadline_without_timing_out() {
let budget = Arc::new(Semaphore::new(OUTBOUND_BYTE_BUDGET));
let chunks: Vec<Result<Bytes, unb_core::CoreError>> =
(0..8).map(|_| Ok(Bytes::from_static(b"chunk"))).collect();
let mut body = MeteredBody::new(chunk_stream(chunks), budget, CancellationToken::new());
for _ in 0..8 {
let chunk = body
.next_chunk()
.await
.expect("a steadily progressing body never times out")
.expect("chunks remain");
drop(chunk);
tokio::time::advance(std::time::Duration::from_secs(6)).await;
}
assert!(body.next_chunk().await.unwrap().is_none());
}
#[tokio::test]
async fn pipe_recv_yields_whole_message_envelopes_with_no_body_stream() {
let mut envelope = frame(
Kind::Request,
Some("c1".into()),
serde_json::json!({"n": 1}),
);
envelope.target = "server".into();
envelope.subject = "echo".into();
let encoded = envelope.encode();
let pipe = Pipe::Piped {
pipe: Box::pin(ScriptedPipe {
frames: std::collections::VecDeque::from([Ok(encoded.clone())]),
}),
initiator: true,
};
let (mut reader, _writer) = pipe.split();
let (received, body) = reader.recv().await.unwrap().unwrap();
assert!(body.is_none(), "whole-message carriers collect the body");
assert_eq!(received, Envelope::decode(encoded).unwrap());
let (incoming, rx) = mpsc::channel(4);
let (tx, _outgoing) = mpsc::channel(4);
let pipe = Pipe::Local {
rx,
tx,
initiator: true,
};
let (mut reader, _writer) = pipe.split();
incoming.send(envelope.clone()).await.unwrap();
let (received, body) = reader.recv().await.unwrap().unwrap();
assert!(body.is_none());
assert_eq!(received, envelope);
}
#[test]
fn write_weight_counts_head_and_payload_and_clamps_to_the_whole_budget() {
let mut envelope = frame(Kind::Request, None, serde_json::Value::Null);
envelope.id = String::new();
envelope.target = "server".into();
envelope.subject = "echo".into();
envelope.payload = Bytes::new();
assert_eq!(WriteRequest::weight(&envelope), 10);
envelope
.headers
.insert("origin".into(), serde_json::Value::String("x".repeat(64)));
assert_eq!(WriteRequest::weight(&envelope), 80);
envelope.payload = Bytes::from(vec![0u8; OUTBOUND_BYTE_BUDGET + 1]);
assert_eq!(WriteRequest::weight(&envelope), OUTBOUND_BYTE_BUDGET as u32);
}
#[tokio::test]
async fn a_byte_saturated_outbound_ends_with_busy_and_recovers_on_drain() {
let handle = spawn_core();
let (pipe, incoming, mut outgoing) = stalled_pipe();
let wire = attach_client(&handle, "stalled", pipe).await;
establish(&incoming, &mut outgoing).await;
let client = wire.client_session();
let payload = Bytes::from(vec![0u8; 16 * 1024 * 1024]);
let mut flood = Vec::new();
for _ in 0..6 {
flood.push(
client
.start(
"/server/flood",
Kind::Request,
payload.clone(),
None,
Default::default(),
)
.await
.unwrap(),
);
}
let mut rejected = client
.start(
"/server/flood",
Kind::Request,
payload.clone(),
None,
Default::default(),
)
.await
.unwrap();
let refused = rejected.operation().as_str().to_owned();
let refusal = tokio::time::timeout(std::time::Duration::from_secs(2), rejected.next())
.await
.expect("an over-budget write must be refused promptly, not silently queued");
assert!(
refusal.is_err(),
"queued outbound bytes past the budget must refuse the write: {refusal:?}"
);
let mut saw_busy = false;
loop {
match tokio::time::timeout(std::time::Duration::from_millis(300), outgoing.recv()).await
{
Ok(Some(envelope)) => {
if envelope.kind == Kind::Error
&& envelope.corr.as_deref() == Some(refused.as_str())
{
assert_eq!(envelope.payload_json()["code"], "BUSY");
saw_busy = true;
}
}
Ok(None) => panic!("the stalled pipe closed while draining"),
Err(_) => break,
}
}
assert!(saw_busy, "the refused write must surface as a BUSY error");
let probe = client
.start(
"/server/probe",
Kind::Request,
payload.clone(),
None,
Default::default(),
)
.await
.unwrap();
let opening = tokio::time::timeout(std::time::Duration::from_secs(5), async {
loop {
let envelope = outgoing.recv().await.unwrap();
if envelope.subject == "probe" {
break envelope;
}
}
})
.await
.expect("draining the writer must release budget for later writes");
incoming
.send(frame(
Kind::Response,
opening.corr,
serde_json::json!({ "ok": true }),
))
.await
.unwrap();
let mut probe = probe;
let response = tokio::time::timeout(std::time::Duration::from_secs(5), probe.next())
.await
.expect("a post-drain write must complete end to end")
.unwrap()
.unwrap();
assert_eq!(response.payload_json(), serde_json::json!({ "ok": true }));
handle.shutdown();
}
#[tokio::test]
async fn a_byte_saturated_session_leaves_a_sibling_session_undisturbed() {
let handle = spawn_core();
let (saturated_pipe, stalled_incoming, mut stalled_outgoing) = stalled_pipe();
let stalled = attach_client(&handle, "stalled", saturated_pipe).await;
establish(&stalled_incoming, &mut stalled_outgoing).await;
let (live_pipe, live_incoming, mut live_outgoing) = local_pipe();
let live = attach_client(&handle, "live", live_pipe).await;
establish(&live_incoming, &mut live_outgoing).await;
let payload = Bytes::from(vec![0u8; 16 * 1024 * 1024]);
let mut flood = Vec::new();
for _ in 0..6 {
flood.push(
stalled
.client_session()
.start(
"/server/flood",
Kind::Request,
payload.clone(),
None,
Default::default(),
)
.await
.unwrap(),
);
}
let mut sibling = live
.client_session()
.start(
"/server/sibling",
Kind::Request,
Bytes::new(),
None,
Default::default(),
)
.await
.unwrap();
let opening = tokio::time::timeout(std::time::Duration::from_secs(1), live_outgoing.recv())
.await
.expect("a sibling session must keep writing while another is byte-saturated")
.unwrap();
assert_eq!(opening.subject, "sibling");
live_incoming
.send(frame(
Kind::Response,
opening.corr,
serde_json::json!({ "ok": true }),
))
.await
.unwrap();
let response = tokio::time::timeout(std::time::Duration::from_secs(2), sibling.next())
.await
.expect("a sibling response must be delivered while another session is saturated")
.unwrap()
.unwrap();
assert_eq!(response.payload_json(), serde_json::json!({ "ok": true }));
handle.shutdown();
}
#[tokio::test]
async fn terminal_client_delivery_closes_handler_stream_once() {
let cancellation = CancellationToken::new();
let handle = ProtocolCoreHandle::spawn(
ProtocolCore::new("node"),
Arc::new(RecordingExecutor {
effects: mpsc::channel(1).0,
}),
cancellation,
&tokio::runtime::Handle::current(),
);
let (pipe, incoming, mut outgoing) = local_pipe();
let (delivered, _deliveries) = mpsc::channel(4);
let (closed, mut closures) = mpsc::channel(4);
let session_cancellation = handle.cancellation.child_token();
let (wire, directives, outcome, routes) =
Wire::standalone_channel(session_cancellation.clone());
let (reply, response) = oneshot::channel();
handle
.inputs
.send(ActorInput::Attach {
session: SessionId::from("session"),
transport: pipe,
expected_peer: None,
directives,
handler: handler_channel(ClosingHandler { delivered, closed }),
client: wire.client_session(),
cancellation: session_cancellation,
establish_peer: false,
outcome,
routes,
ws_collect_ceiling: unb_transport::DEFAULT_MAX_FRAME_SIZE,
reply,
})
.await
.unwrap();
response.await.unwrap().unwrap();
establish(&incoming, &mut outgoing).await;
let mut stream = wire
.client_session()
.start(
"/server/echo",
Kind::Request,
Bytes::new(),
None,
Default::default(),
)
.await
.unwrap();
let opening = outgoing.recv().await.unwrap();
incoming
.send(frame(
Kind::Response,
opening.corr,
serde_json::json!({ "ok": true }),
))
.await
.unwrap();
assert!(stream.next().await.unwrap().is_some());
assert_eq!(closures.recv().await, Some(stream.operation().clone()));
assert!(
tokio::time::timeout(std::time::Duration::from_millis(25), closures.recv())
.await
.is_err()
);
handle.shutdown();
}
#[tokio::test]
async fn generic_terminal_delivery_closes_handler_stream_once() {
let cancellation = CancellationToken::new();
let handle = ProtocolCoreHandle::spawn(
ProtocolCore::new("node"),
Arc::new(RecordingExecutor {
effects: mpsc::channel(1).0,
}),
cancellation,
&tokio::runtime::Handle::current(),
);
let (pipe, incoming, mut outgoing) = local_pipe();
let (delivered, mut deliveries) = mpsc::channel(4);
let (closed, mut closures) = mpsc::channel(4);
let session_cancellation = handle.cancellation.child_token();
let (wire, directives, outcome, routes) =
Wire::standalone_channel(session_cancellation.clone());
let (reply, response) = oneshot::channel();
handle
.inputs
.send(ActorInput::Attach {
session: SessionId::from("session"),
transport: pipe,
expected_peer: None,
directives,
handler: handler_channel(ClosingHandler { delivered, closed }),
client: wire.client_session(),
cancellation: session_cancellation,
establish_peer: false,
outcome,
routes,
ws_collect_ceiling: unb_transport::DEFAULT_MAX_FRAME_SIZE,
reply,
})
.await
.unwrap();
response.await.unwrap().unwrap();
establish(&incoming, &mut outgoing).await;
let corr = wire
.open_stream("/server/echo", Kind::Request, serde_json::Value::Null)
.await
.unwrap();
let opening = outgoing.recv().await.unwrap();
assert_eq!(opening.corr.as_deref(), Some(corr.as_str()));
incoming
.send(frame(
Kind::Response,
Some(corr.clone()),
serde_json::json!({ "ok": true }),
))
.await
.unwrap();
assert_eq!(
deliveries.recv().await.unwrap().corr.as_deref(),
Some(corr.as_str())
);
assert_eq!(closures.recv().await, Some(ClientOperationId::from(corr)));
assert!(
tokio::time::timeout(std::time::Duration::from_millis(25), closures.recv())
.await
.is_err()
);
handle.shutdown();
}
#[tokio::test]
async fn immediate_loopback_response_finds_registered_client_operation() {
let (pipe, incoming, mut outgoing) = local_pipe();
let wire = Wire::open(pipe);
establish(&incoming, &mut outgoing).await;
let loopback = tokio::spawn(async move {
let opening = outgoing.recv().await.unwrap();
let corr = opening.corr.unwrap();
incoming
.send(frame(
Kind::Response,
Some(corr),
serde_json::json!({ "ok": true }),
))
.await
.unwrap();
});
let mut operation = wire
.client_session()
.start(
"/server/echo",
Kind::Request,
Bytes::new(),
None,
Default::default(),
)
.await
.unwrap();
loopback.await.unwrap();
assert_eq!(
operation
.try_next()
.unwrap()
.unwrap()
.unwrap()
.payload_json(),
serde_json::json!({ "ok": true })
);
assert_eq!(operation.try_next().unwrap(), Some(None));
}
#[tokio::test]
async fn opening_send_failure_cleans_up_client_operation() {
let (pipe, incoming, mut outgoing) = local_pipe();
let wire = Wire::open(pipe);
establish(&incoming, &mut outgoing).await;
drop(outgoing);
let mut operation = wire
.client_session()
.start(
"/server/echo",
Kind::Request,
Bytes::new(),
None,
Default::default(),
)
.await
.unwrap();
assert!(matches!(
operation.next().await,
Err(crate::ClientError::SessionClosed(_))
));
}
#[tokio::test]
async fn fetch_projects_response_and_ignores_late_terminal_frame() {
let (pipe, incoming, mut outgoing) = local_pipe();
let wire = Wire::open(pipe);
establish(&incoming, &mut outgoing).await;
let responder = tokio::spawn(async move {
let opening = outgoing.recv().await.unwrap();
let corr = opening.corr.unwrap();
let mut response = frame(
Kind::Response,
Some(corr.clone()),
serde_json::json!({ "ok": true }),
);
response.headers.insert("x-source".into(), "owner".into());
incoming.send(response.clone()).await.unwrap();
incoming.send(response).await.unwrap();
});
let response = wire
.client_session()
.fetch(
request("/server/echo", b"request"),
std::time::Duration::from_secs(1),
)
.await
.unwrap();
assert_eq!(response.status(), http::StatusCode::OK);
assert_eq!(response.headers()["x-source"], "owner");
assert_eq!(response.body(), &Bytes::from_static(br#"{"ok":true}"#));
responder.await.unwrap();
}
#[tokio::test]
async fn subscribe_and_discovery_are_isolated() {
let (pipe, incoming, mut outgoing) = local_pipe();
let wire = Wire::open(pipe);
establish(&incoming, &mut outgoing).await;
let client = wire.client_session();
let mut subscription = client
.subscribe(request("/server/events", b""), None)
.await
.unwrap();
let mut discovery = client
.discover("/server", DiscoverPlan::decode(&Bytes::new()).unwrap())
.await
.unwrap();
let first = outgoing.recv().await.unwrap();
let second = outgoing.recv().await.unwrap();
let (subscribe_corr, discover_corr) = if first.kind == Kind::Subscribe {
(first.corr.unwrap(), second.corr.unwrap())
} else {
(second.corr.unwrap(), first.corr.unwrap())
};
incoming
.send(frame(
Kind::Event,
Some(discover_corr.clone()),
serde_json::json!({ "discover": 1 }),
))
.await
.unwrap();
incoming
.send(frame(
Kind::Event,
Some(subscribe_corr.clone()),
serde_json::json!({ "event": 1 }),
))
.await
.unwrap();
incoming
.send(frame(
Kind::Response,
Some(discover_corr),
serde_json::json!({ "done": true }),
))
.await
.unwrap();
assert_eq!(
subscription.next().await.unwrap().unwrap().payload_json()["event"],
1
);
assert_eq!(
discovery.next().await.unwrap().unwrap().payload_json()["discover"],
1
);
assert_eq!(
discovery.next().await.unwrap().unwrap().payload_json()["done"],
true
);
}
#[tokio::test]
async fn dropping_a_saturated_operation_frees_its_session_for_idle_delivery() {
let (pipe, incoming, mut outgoing) = local_pipe();
let wire = Wire::open(pipe);
establish(&incoming, &mut outgoing).await;
let client = wire.client_session();
let busy = client
.subscribe(request("/server/busy", b""), None)
.await
.unwrap();
let mut idle = client
.subscribe(request("/server/idle", b""), None)
.await
.unwrap();
let busy_open = outgoing.recv().await.unwrap();
let idle_open = outgoing.recv().await.unwrap();
let burst = tokio::spawn(async move {
for sequence in 0..257 {
incoming
.send(frame(
Kind::Event,
busy_open.corr.clone(),
serde_json::json!({ "sequence": sequence }),
))
.await
.unwrap();
}
incoming
.send(frame(
Kind::Response,
idle_open.corr,
serde_json::json!({ "done": true }),
))
.await
.unwrap();
incoming
});
tokio::time::sleep(std::time::Duration::from_millis(25)).await;
drop(busy);
let _incoming = tokio::time::timeout(std::time::Duration::from_secs(5), burst)
.await
.expect("dropping the saturated stream must unblock its session")
.unwrap();
assert_eq!(
idle.next().await.unwrap().unwrap().payload_json()["done"],
true
);
}
#[tokio::test]
async fn out_of_order_unary_responses_reach_their_operations() {
let (pipe, incoming, mut outgoing) = local_pipe();
let wire = Wire::open(pipe);
establish(&incoming, &mut outgoing).await;
let client = wire.client_session();
let mut first = client
.start(
"/server/first",
Kind::Request,
Bytes::new(),
None,
Default::default(),
)
.await
.unwrap();
let mut second = client
.start(
"/server/second",
Kind::Request,
Bytes::new(),
None,
Default::default(),
)
.await
.unwrap();
let first_open = outgoing.recv().await.unwrap();
let second_open = outgoing.recv().await.unwrap();
incoming
.send(frame(
Kind::Response,
second_open.corr,
serde_json::json!({ "response": "second" }),
))
.await
.unwrap();
incoming
.send(frame(
Kind::Response,
first_open.corr,
serde_json::json!({ "response": "first" }),
))
.await
.unwrap();
assert_eq!(
first.next().await.unwrap().unwrap().payload_json()["response"],
"first"
);
assert_eq!(
second.next().await.unwrap().unwrap().payload_json()["response"],
"second"
);
}
#[tokio::test]
async fn mixed_unary_subscription_and_discovery_are_isolated() {
let (pipe, incoming, mut outgoing) = local_pipe();
let wire = Wire::open(pipe);
establish(&incoming, &mut outgoing).await;
let client = wire.client_session();
let mut unary = client
.start(
"/server/unary",
Kind::Request,
Bytes::new(),
None,
Default::default(),
)
.await
.unwrap();
let mut subscription = client
.subscribe(request("/server/events", b""), None)
.await
.unwrap();
let mut discovery = client
.discover("/server", DiscoverPlan::decode(&Bytes::new()).unwrap())
.await
.unwrap();
let mut unary_corr = None;
let mut subscription_corr = None;
let mut discovery_corr = None;
for _ in 0..3 {
let opening = outgoing.recv().await.unwrap();
match opening.kind {
Kind::Request => unary_corr = opening.corr,
Kind::Subscribe => subscription_corr = opening.corr,
Kind::Discover => discovery_corr = opening.corr,
kind => panic!("unexpected opening kind {kind:?}"),
}
}
for (kind, corr, value) in [
(Kind::Event, discovery_corr, "discover"),
(Kind::Response, unary_corr, "unary"),
(Kind::Event, subscription_corr, "subscription"),
] {
incoming
.send(frame(kind, corr, serde_json::json!({ "owner": value })))
.await
.unwrap();
}
assert_eq!(
unary.next().await.unwrap().unwrap().payload_json()["owner"],
"unary"
);
assert_eq!(
subscription.next().await.unwrap().unwrap().payload_json()["owner"],
"subscription"
);
assert_eq!(
discovery.next().await.unwrap().unwrap().payload_json()["owner"],
"discover"
);
}
#[tokio::test]
async fn cancelling_one_operation_does_not_complete_another() {
let (pipe, incoming, mut outgoing) = local_pipe();
let wire = Wire::open(pipe);
establish(&incoming, &mut outgoing).await;
let client = wire.client_session();
let cancelled = client
.subscribe(request("/server/cancelled", b""), None)
.await
.unwrap();
let mut live = client
.subscribe(request("/server/live", b""), None)
.await
.unwrap();
let cancelled_open = outgoing.recv().await.unwrap();
let live_open = outgoing.recv().await.unwrap();
drop(cancelled);
let cancel = outgoing.recv().await.unwrap();
assert_eq!(cancel.kind, Kind::Cancel);
assert_eq!(cancel.corr, cancelled_open.corr);
incoming
.send(frame(
Kind::Event,
live_open.corr,
serde_json::json!({ "live": true }),
))
.await
.unwrap();
assert_eq!(
live.next().await.unwrap().unwrap().payload_json()["live"],
true
);
}
#[tokio::test]
async fn stream_drop_sends_exactly_one_cancel() {
let (pipe, incoming, mut outgoing) = local_pipe();
let wire = Wire::open(pipe);
establish(&incoming, &mut outgoing).await;
let stream = wire
.client_session()
.subscribe(request("/server/events", b""), None)
.await
.unwrap();
let opening = outgoing.recv().await.unwrap();
let corr = opening.corr.unwrap();
drop(stream);
let cancel = tokio::time::timeout(std::time::Duration::from_secs(1), outgoing.recv())
.await
.unwrap()
.unwrap();
assert_eq!(cancel.kind, Kind::Cancel);
assert_eq!(cancel.corr.as_deref(), Some(corr.as_str()));
assert!(
tokio::time::timeout(std::time::Duration::from_millis(25), outgoing.recv())
.await
.is_err()
);
}
#[tokio::test]
async fn operation_timeout_is_typed_and_cancels_once() {
let (pipe, incoming, mut outgoing) = local_pipe();
let wire = Wire::open(pipe);
establish(&incoming, &mut outgoing).await;
let mut stream = wire
.client_session()
.subscribe(
request("/server/events", b""),
Some(std::time::Duration::from_millis(10)),
)
.await
.unwrap();
let opening = outgoing.recv().await.unwrap();
assert!(matches!(
stream.next().await,
Err(crate::ClientError::Timeout(_))
));
let cancel = outgoing.recv().await.unwrap();
assert_eq!(cancel.kind, Kind::Cancel);
assert_eq!(cancel.corr, opening.corr);
drop(stream);
assert!(
tokio::time::timeout(std::time::Duration::from_millis(25), outgoing.recv())
.await
.is_err()
);
}
#[tokio::test]
async fn session_close_completes_every_operation_once() {
let (pipe, incoming, mut outgoing) = local_pipe();
let wire = Wire::open(pipe);
establish(&incoming, &mut outgoing).await;
let client = wire.client_session();
let mut first = client
.subscribe(request("/server/first", b""), None)
.await
.unwrap();
let mut second = client
.subscribe(request("/server/second", b""), None)
.await
.unwrap();
outgoing.recv().await.unwrap();
outgoing.recv().await.unwrap();
drop(incoming);
assert!(matches!(
first.next().await,
Err(crate::ClientError::SessionClosed(_))
));
assert!(matches!(
second.next().await,
Err(crate::ClientError::SessionClosed(_))
));
assert_eq!(first.next().await.unwrap(), None);
assert_eq!(second.next().await.unwrap(), None);
}
#[tokio::test]
async fn a_session_the_core_rejects_closes_its_transport() {
let (pipe, incoming, mut outgoing) = local_pipe();
let _wire = Wire::open(pipe);
assert_eq!(outgoing.recv().await.unwrap().kind, Kind::Hello);
incoming
.send(frame(
Kind::Welcome,
None,
serde_json::json!({ "version": 999 }),
))
.await
.unwrap();
let mut saw_error = false;
let closed = loop {
match tokio::time::timeout(std::time::Duration::from_millis(500), outgoing.recv()).await
{
Ok(Some(frame)) => {
if frame.kind == Kind::Error {
saw_error = true;
}
}
Ok(None) => break true,
Err(_) => break false,
}
};
assert!(
saw_error,
"the peer must receive the close-reason error frame before the transport closes"
);
assert!(
closed,
"core deciding a session invalid must close the transport, not leave it open"
);
}
struct ScriptedPipe {
frames: std::collections::VecDeque<Result<Bytes, unb_transport::TransportError>>,
}
impl futures_util::Stream for ScriptedPipe {
type Item = Result<Bytes, unb_transport::TransportError>;
fn poll_next(
mut self: Pin<&mut Self>,
_cx: &mut std::task::Context<'_>,
) -> std::task::Poll<Option<Self::Item>> {
std::task::Poll::Ready(self.frames.pop_front())
}
}
impl futures_util::Sink<Bytes> for ScriptedPipe {
type Error = unb_transport::TransportError;
fn poll_ready(
self: Pin<&mut Self>,
_cx: &mut std::task::Context<'_>,
) -> std::task::Poll<Result<(), Self::Error>> {
std::task::Poll::Ready(Ok(()))
}
fn start_send(self: Pin<&mut Self>, _item: Bytes) -> Result<(), Self::Error> {
Ok(())
}
fn poll_flush(
self: Pin<&mut Self>,
_cx: &mut std::task::Context<'_>,
) -> std::task::Poll<Result<(), Self::Error>> {
std::task::Poll::Ready(Ok(()))
}
fn poll_close(
self: Pin<&mut Self>,
_cx: &mut std::task::Context<'_>,
) -> std::task::Poll<Result<(), Self::Error>> {
std::task::Poll::Ready(Ok(()))
}
}
#[tokio::test]
async fn a_malformed_frame_retires_the_session_with_a_typed_reason() {
let cancellation = CancellationToken::new();
let (effect_tx, mut effects) = mpsc::channel(64);
let handle = ProtocolCoreHandle::spawn(
ProtocolCore::new("node"),
Arc::new(RecordingExecutor { effects: effect_tx }),
cancellation,
&tokio::runtime::Handle::current(),
);
let pipe = Pipe::Piped {
pipe: Box::pin(ScriptedPipe {
frames: std::collections::VecDeque::from([Ok(Bytes::from_static(
b"\xff\xff\xff\xffgarbage",
))]),
}),
initiator: true,
};
let (events, _events_rx) = mpsc::channel(16);
handle
.attach(SessionId::from("session"), pipe, None, events)
.await
.unwrap();
let reason = tokio::time::timeout(std::time::Duration::from_secs(2), async {
loop {
if let CoreEffect::SessionRetired { reason, .. } = effects.recv().await.unwrap() {
break reason;
}
}
})
.await
.expect("a malformed frame must retire the session");
assert_eq!(reason, RetirementReason::TransportFailed);
handle.shutdown();
}
#[tokio::test]
async fn a_transport_read_error_retires_the_session_with_a_typed_reason() {
let cancellation = CancellationToken::new();
let (effect_tx, mut effects) = mpsc::channel(64);
let handle = ProtocolCoreHandle::spawn(
ProtocolCore::new("node"),
Arc::new(RecordingExecutor { effects: effect_tx }),
cancellation,
&tokio::runtime::Handle::current(),
);
let pipe = Pipe::Piped {
pipe: Box::pin(ScriptedPipe {
frames: std::collections::VecDeque::from([Err(unb_transport::TransportError::Io(
"connection reset".into(),
))]),
}),
initiator: true,
};
let (events, _events_rx) = mpsc::channel(16);
handle
.attach(SessionId::from("session"), pipe, None, events)
.await
.unwrap();
let reason = tokio::time::timeout(std::time::Duration::from_secs(2), async {
loop {
if let CoreEffect::SessionRetired { reason, .. } = effects.recv().await.unwrap() {
break reason;
}
}
})
.await
.expect("a transport read error must retire the session");
assert_eq!(reason, RetirementReason::TransportFailed);
handle.shutdown();
}
#[tokio::test]
async fn routes_two_sessions_and_submits_transport_closure() {
let cancellation = CancellationToken::new();
let (effect_tx, mut effects) = mpsc::channel(16);
let handle = ProtocolCoreHandle::spawn(
ProtocolCore::new("node"),
Arc::new(RecordingExecutor { effects: effect_tx }),
cancellation,
&tokio::runtime::Handle::current(),
);
let (first_pipe, first_incoming, mut first_outgoing) = local_pipe();
let (second_pipe, _second_incoming, mut second_outgoing) = local_pipe();
let (first_events, _first_events_rx) = mpsc::channel(16);
let (second_events, _second_events_rx) = mpsc::channel(16);
let first = handle
.attach(SessionId::from("first"), first_pipe, None, first_events)
.await
.unwrap();
let second = handle
.attach(SessionId::from("second"), second_pipe, None, second_events)
.await
.unwrap();
assert!(first_outgoing.recv().await.is_some());
assert!(second_outgoing.recv().await.is_some());
assert!(Arc::ptr_eq(&first, &first));
assert!(!Arc::ptr_eq(&first, &second));
drop(first_incoming);
let retired = loop {
let effect = effects.recv().await.unwrap();
if matches!(
effect,
CoreEffect::SessionRetired {
ref session,
reason: RetirementReason::SessionClosed,
} if session == &SessionId::from("first")
) {
break effect;
}
};
assert!(matches!(
retired,
CoreEffect::SessionRetired {
reason: RetirementReason::SessionClosed,
..
}
));
second.shutdown();
handle.shutdown();
}
#[tokio::test]
async fn install_publishes_after_acceptance_and_before_effect_execution() {
let cancellation = CancellationToken::new();
let published = Arc::new(AtomicBool::new(false));
let (effect_tx, mut effects) = mpsc::channel(1);
let handle = ProtocolCoreHandle::spawn(
ProtocolCore::new("node"),
Arc::new(PublicationExecutor {
published: published.clone(),
effects: effect_tx,
}),
cancellation,
&tokio::runtime::Handle::current(),
);
let (pipe, _incoming, mut outgoing) = local_pipe();
let (events, _events_rx) = mpsc::channel(16);
handle
.attach(SessionId::from("session"), pipe, None, events)
.await
.unwrap();
assert!(outgoing.recv().await.is_some());
handle
.install(
CoreInput::SessionClosed {
session: SessionId::from("session"),
},
{
let published = published.clone();
move || published.store(true, Ordering::SeqCst)
},
)
.await
.unwrap();
assert!(effects.recv().await.unwrap());
handle.shutdown();
}
#[tokio::test]
async fn install_failure_does_not_publish() {
let cancellation = CancellationToken::new();
let (effect_tx, _effects) = mpsc::channel(1);
let handle = ProtocolCoreHandle::spawn(
ProtocolCore::new("node"),
Arc::new(RecordingExecutor { effects: effect_tx }),
cancellation,
&tokio::runtime::Handle::current(),
);
let input = CoreInput::SessionOpened {
session: SessionId::from("duplicate"),
initiator: true,
establish_peer: false,
expected_peer: None,
};
handle.submit(input.clone()).await.unwrap();
let published = Arc::new(AtomicBool::new(false));
let result = handle
.install(input, {
let published = published.clone();
move || published.store(true, Ordering::SeqCst)
})
.await;
assert!(matches!(result, Err(WsError::Core(_))));
assert!(!published.load(Ordering::SeqCst));
handle.shutdown();
}
#[test]
fn operation_records_cannot_change_their_bound_correlation() {
let response = frame(
Kind::Response,
Some("s2".into()),
serde_json::json!({"ok": true}),
);
let error =
decode_operation_message(response.encode(), "s1", OperationDirection::Return, 0)
.unwrap_err();
assert!(matches!(
error,
unb_transport::operation::OperationCodecError::Protocol(_)
));
assert_eq!(error.reset_code(), unb_transport::operation::RESET_PROTOCOL);
}
#[tokio::test]
async fn streaming_operation_body_uses_head_chunks_and_no_wire_body_token() {
let mut envelope = frame(Kind::Response, Some("s1".into()), serde_json::json!(null));
envelope.body_token = Some("legacy-token-must-not-leak".into());
let body: BodyStream = Box::pin(futures_util::stream::iter([
Ok(Bytes::from_static(b"first")),
Ok(Bytes::from_static(b"second")),
]));
let (write, read) = tokio::io::duplex(4096);
let sender = async move {
let mut writer = unb_transport::operation::OperationRecordWriter::new(write);
write_operation_envelope(
&mut writer,
envelope,
Some(body),
Arc::new(Semaphore::new(OUTBOUND_BYTE_BUDGET)),
CancellationToken::new(),
)
.await
.unwrap();
writer.finish().await.unwrap();
};
let receiver = async move {
let mut reader = unb_transport::operation::OperationRecordReader::new(read);
let (kind, payload) = reader.read_record().await.unwrap();
assert_eq!(kind, unb_transport::operation::OperationRecordKind::Head);
let head = Envelope::decode(payload).unwrap();
assert_eq!(head.corr.as_deref(), Some("s1"));
assert!(head.body_token.is_none());
assert!(head.payload.is_empty());
let mut collected = Vec::new();
loop {
let (kind, payload) = reader.read_record().await.unwrap();
match kind {
unb_transport::operation::OperationRecordKind::BodyChunk => {
collected.extend_from_slice(&payload)
}
unb_transport::operation::OperationRecordKind::BodyEnd => break,
other => panic!("unexpected body record {other:?}"),
}
}
assert_eq!(collected, b"firstsecond");
let (kind, payload) = reader.read_record().await.unwrap();
assert_eq!(
kind,
unb_transport::operation::OperationRecordKind::DirectionEnd
);
assert!(payload.is_empty());
reader.expect_fin().await.unwrap();
};
tokio::join!(sender, receiver);
}
#[tokio::test]
async fn cancelling_a_client_operation_releases_a_late_response_body() {
let handle = spawn_core();
let (pipe, incoming, mut outgoing) = local_pipe();
let wire = attach_client(&handle, "cancel-body", pipe).await;
establish(&incoming, &mut outgoing).await;
let call = wire
.client_session()
.start(
"/server/echo",
Kind::Request,
Bytes::from_static(b"request"),
None,
Default::default(),
)
.await
.unwrap();
let corr = call.operation().as_str().to_owned();
assert_eq!(
outgoing.recv().await.unwrap().corr.as_deref(),
Some(corr.as_str())
);
drop(call);
let cancel = outgoing.recv().await.unwrap();
assert_eq!(cancel.kind, Kind::Cancel);
assert_eq!(cancel.corr.as_deref(), Some(corr.as_str()));
incoming
.send(frame(
Kind::Response,
Some(corr),
serde_json::json!({ "late": true }),
))
.await
.unwrap();
tokio::time::timeout(std::time::Duration::from_secs(1), async {
loop {
let retained = handle
.token_bodies
.lock()
.expect("body token registry")
.keys()
.any(|(owner, _)| owner.as_str() == "cancel-body");
if !retained {
break;
}
tokio::task::yield_now().await;
}
})
.await
.expect("late response body must be released");
handle.shutdown();
}
#[tokio::test]
async fn session_retirement_releases_unclaimed_bodies() {
let handle = spawn_core();
let (pipe, incoming, mut outgoing) = local_pipe();
let _wire = attach_client(&handle, "retired-body", pipe).await;
establish(&incoming, &mut outgoing).await;
handle
.register_body(
&SessionId::from("retired-body"),
WireBody::Bytes(Bytes::from_static(b"retained")),
)
.unwrap();
assert!(handle
.token_bodies
.lock()
.expect("body token registry")
.keys()
.any(|(owner, _)| owner.as_str() == "retired-body"));
drop(incoming);
tokio::time::timeout(std::time::Duration::from_secs(1), async {
loop {
let retained = handle
.token_bodies
.lock()
.expect("body token registry")
.keys()
.any(|(owner, _)| owner.as_str() == "retired-body");
if !retained {
break;
}
tokio::task::yield_now().await;
}
})
.await
.expect("session retirement must release body capabilities");
handle.shutdown();
}
}