1use std::collections::HashMap;
2use std::collections::VecDeque;
3use std::future::Future;
4use std::pin::Pin;
5use std::sync::atomic::{AtomicBool, Ordering};
6use std::sync::Arc;
7
8use unb_core::{
9 ClientDelivery as CoreClientDelivery, ClientOperationId, CoreEffect, CoreInput, EffectId,
10 Envelope, ProtocolCore, RetirementReason, SendResult, SessionId,
11};
12use bytes::Bytes;
13use tokio::sync::{mpsc, oneshot, OwnedSemaphorePermit, Semaphore};
14use web_time::Instant;
15
16use crate::cancellation::CancellationToken;
17use crate::client::{ClientDelivery, ClientSession};
18use crate::error::WsError;
19use crate::transport::{Pipe, PipeReader, PipeWriter};
20use crate::wire::{Directive, Wire};
21use crate::{BodyStream, WireBody};
22
23const CORE_QUEUE: usize = 256;
24const SESSION_QUEUE: usize = 64;
25const OUTBOUND_QUEUE: usize = 8_192;
26const OUTBOUND_BYTE_BUDGET: usize = 64 * 1024 * 1024;
27const BODY_WINDOW_BYTES: usize = 256 * 1024;
28const WRITE_BATCH: usize = 8;
29const RELAY_BACKPRESSURE_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(5);
30const TERMINAL_SEND_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(30);
31const CLIENT_DELIVERY_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(30);
32const ESTABLISHMENT_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(5);
33const PENDING_MODE_WRITES_MAX: usize = 128;
34const OPERATION_WRITE_QUEUE: usize = OUTBOUND_QUEUE;
35const OPERATION_ACTIVE_MAX: usize = 128;
36const OPERATION_BODY_CHUNK_BYTES: usize = 64 * 1024;
37
38#[cfg(not(all(target_arch = "wasm32", target_os = "unknown")))]
39fn clock_now() -> Instant {
40 n0_future::time::Instant::now().into_std()
41}
42
43#[cfg(not(all(target_arch = "wasm32", target_os = "unknown")))]
44async fn sleep_to(deadline: Instant) {
45 n0_future::time::sleep_until(n0_future::time::Instant::from_std(deadline)).await
46}
47
48#[cfg(all(target_arch = "wasm32", target_os = "unknown"))]
49fn clock_now() -> Instant {
50 n0_future::time::Instant::now()
51}
52
53#[cfg(all(target_arch = "wasm32", target_os = "unknown"))]
54async fn sleep_to(deadline: Instant) {
55 n0_future::time::sleep_until(deadline).await
56}
57
58#[cfg(not(all(target_arch = "wasm32", target_os = "unknown")))]
59pub use tokio::runtime::Handle as RuntimeHandle;
60
61#[cfg(all(target_arch = "wasm32", target_os = "unknown"))]
62#[derive(Clone, Debug)]
63pub struct RuntimeHandle;
64
65#[cfg(all(target_arch = "wasm32", target_os = "unknown"))]
66impl RuntimeHandle {
67 pub fn current() -> Self {
68 RuntimeHandle
69 }
70
71 pub fn spawn<F>(&self, future: F)
72 where
73 F: Future<Output = ()> + 'static,
74 {
75 n0_future::task::spawn(future);
76 }
77}
78
79pub type EffectFuture = Pin<Box<dyn Future<Output = Option<CoreInput>> + Send>>;
80
81pub trait EffectExecutor: Send + Sync + 'static {
82 fn execute(&self, effect: CoreEffect, handle: ProtocolCoreHandle) -> EffectFuture;
83}
84
85#[derive(Clone, Debug)]
86pub enum SessionOutcome {
87 Established,
88 Retired(RetirementReason),
89}
90
91pub trait SessionHandler: Send + 'static {
92 fn deliver(&mut self, envelope: Envelope) -> impl Future<Output = ()> + Send;
93
94 fn stream_closed(&mut self, _operation: ClientOperationId) -> impl Future<Output = ()> + Send {
95 async {}
96 }
97}
98
99impl SessionHandler for () {
100 async fn deliver(&mut self, _envelope: Envelope) {}
101}
102
103impl SessionHandler for mpsc::Sender<Envelope> {
104 async fn deliver(&mut self, envelope: Envelope) {
105 let _ = self.send(envelope).await;
106 }
107}
108
109impl SessionHandler for tokio::sync::broadcast::Sender<Envelope> {
110 async fn deliver(&mut self, envelope: Envelope) {
111 let _ = self.send(envelope);
112 }
113}
114
115type TokenBodies = Arc<std::sync::Mutex<HashMap<(SessionId, String), WireBody>>>;
116
117type BodyPumps = Arc<std::sync::Mutex<HashMap<(SessionId, String), (String, CancellationToken)>>>;
118
119const SESSION_BODY_TOKENS_MAX: usize = OUTBOUND_QUEUE;
120const SESSION_BODY_BYTES_MAX: usize = OUTBOUND_BYTE_BUDGET;
121
122fn mint_body_token() -> String {
123 use std::hash::{BuildHasher, Hasher};
124 type TokenKeys = (
125 std::collections::hash_map::RandomState,
126 std::collections::hash_map::RandomState,
127 );
128 static KEYS: std::sync::OnceLock<TokenKeys> = std::sync::OnceLock::new();
129 static NEXT: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
130 let (low, high) = KEYS.get_or_init(TokenKeys::default);
131 let serial = NEXT.fetch_add(1, Ordering::Relaxed);
132 let mut a = low.build_hasher();
133 a.write_u64(serial);
134 let mut b = high.build_hasher();
135 b.write_u64(serial);
136 format!("b{:016x}{:016x}", a.finish(), b.finish())
137}
138
139fn application_core_input(
140 handle: &ProtocolCoreHandle,
141 session: &SessionId,
142 mut envelope: Envelope,
143) -> Result<CoreInput, unb_core::CoreError> {
144 if envelope.kind == unb_core::Kind::Discover {
145 return Ok(CoreInput::FrameReceived {
146 session: session.clone(),
147 envelope,
148 });
149 }
150 if !(envelope.kind.is_application_request()
151 || (envelope.corr.is_some()
152 && (envelope.kind.is_application_response()
153 || envelope.kind == unb_core::Kind::Cancel)))
154 {
155 return Ok(CoreInput::FrameReceived {
156 session: session.clone(),
157 envelope,
158 });
159 }
160 if envelope.body_token.is_some() && !envelope.payload.is_empty() {
161 return Err(unb_core::CoreError::Malformed(
162 "application frame cannot carry both inline and referenced bodies".into(),
163 ));
164 }
165 if envelope.body_token.is_none()
166 && envelope.kind != unb_core::Kind::Error
167 && !envelope.payload.is_empty()
168 {
169 let payload = std::mem::take(&mut envelope.payload);
170 envelope.body_token = Some(
171 handle
172 .register_body(session, WireBody::Bytes(payload))?
173 .to_string(),
174 );
175 }
176 Ok(CoreInput::ApplicationFrameReceived {
177 session: session.clone(),
178 frame: unb_core::ApplicationFrame::from_envelope(&envelope)?,
179 })
180}
181
182fn register_outbound_body(
183 handle: &ProtocolCoreHandle,
184 session: &SessionId,
185 payload: Bytes,
186 body: Option<BodyStream>,
187) -> Result<Option<unb_core::BodyId>, unb_core::CoreError> {
188 if body.is_some() && !payload.is_empty() {
189 return Err(unb_core::CoreError::Malformed(
190 "an operation cannot provide both unary bytes and a streaming body".into(),
191 ));
192 }
193 let body = match body {
194 Some(body) => Some(WireBody::Stream(body)),
195 None if !payload.is_empty() => Some(WireBody::Bytes(payload)),
196 None => None,
197 };
198 body.map(|body| handle.register_body(session, body))
199 .transpose()
200}
201
202#[derive(Clone)]
203pub struct ProtocolCoreHandle {
204 inputs: mpsc::Sender<ActorInput>,
205 cancellation: CancellationToken,
206 token_bodies: TokenBodies,
207 body_pumps: BodyPumps,
208}
209
210impl ProtocolCoreHandle {
211 pub fn spawn(
212 core: ProtocolCore,
213 executor: Arc<dyn EffectExecutor>,
214 cancellation: CancellationToken,
215 runtime: &RuntimeHandle,
216 ) -> Self {
217 let (inputs, receiver) = mpsc::channel(CORE_QUEUE);
218 let handle = Self {
219 inputs,
220 cancellation,
221 token_bodies: Arc::new(std::sync::Mutex::new(HashMap::new())),
222 body_pumps: Arc::new(std::sync::Mutex::new(HashMap::new())),
223 };
224 runtime.spawn(CoreRuntime::new(core, receiver, handle.clone(), executor).run());
225 handle
226 }
227
228 pub fn claim_body(&self, session: &SessionId, token: &str) -> Option<WireBody> {
229 self.token_bodies
230 .lock()
231 .expect("body token registry")
232 .remove(&(session.clone(), token.to_string()))
233 }
234
235 pub fn claim_body_stream(&self, session: &SessionId, token: &str) -> Option<BodyStream> {
236 self.claim_body(session, token).map(wire_body_stream)
237 }
238
239 pub fn register_body(
240 &self,
241 session: &SessionId,
242 body: WireBody,
243 ) -> Result<unb_core::BodyId, unb_core::CoreError> {
244 let mut registry = self.token_bodies.lock().expect("body token registry");
245 let owned = registry
246 .keys()
247 .filter(|(owner, _)| owner == session)
248 .count();
249 if owned >= SESSION_BODY_TOKENS_MAX {
250 return Err(unb_core::CoreError::EffectQueueFull);
251 }
252 let incoming = match &body {
253 WireBody::Bytes(payload) => payload.len(),
254 WireBody::Stream(_) => 0,
255 };
256 let retained = registry
257 .iter()
258 .filter_map(|((owner, _), body)| {
259 (owner == session).then_some(match body {
260 WireBody::Bytes(payload) => payload.len(),
261 WireBody::Stream(_) => 0,
262 })
263 })
264 .sum::<usize>();
265 if retained.saturating_add(incoming) > SESSION_BODY_BYTES_MAX {
266 return Err(unb_core::CoreError::BodyRegistryFull(
267 SESSION_BODY_BYTES_MAX,
268 ));
269 }
270 let token = mint_body_token();
271 registry.insert((session.clone(), token.clone()), body);
272 Ok(unb_core::BodyId::from(token))
273 }
274
275 fn register_body_pump(
276 &self,
277 session: &SessionId,
278 corr: &str,
279 token: &str,
280 cancellation: CancellationToken,
281 ) {
282 self.body_pumps.lock().expect("body pump registry").insert(
283 (session.clone(), corr.to_string()),
284 (token.to_string(), cancellation),
285 );
286 }
287
288 fn cancel_body_pump(&self, session: &SessionId, corr: &str) {
289 let cancelled = self
290 .body_pumps
291 .lock()
292 .expect("body pump registry")
293 .remove(&(session.clone(), corr.to_string()));
294 if let Some((_, cancellation)) = cancelled {
295 cancellation.cancel();
296 }
297 }
298
299 fn release_body_pump(&self, session: &SessionId, corr: &str, token: &str) {
300 let mut pumps = self.body_pumps.lock().expect("body pump registry");
301 let key = (session.clone(), corr.to_string());
302 if pumps.get(&key).is_some_and(|(owner, _)| owner == token) {
303 pumps.remove(&key);
304 }
305 }
306
307 pub async fn submit(&self, input: CoreInput) -> Result<(), WsError> {
308 self.inputs
309 .send(ActorInput::Core(input))
310 .await
311 .map_err(|_| WsError::Gone)
312 }
313
314 pub async fn submit_batch(&self, inputs: Vec<CoreInput>) -> Result<(), WsError> {
315 self.inputs
316 .send(ActorInput::CoreBatch(inputs))
317 .await
318 .map_err(|_| WsError::Gone)
319 }
320
321 pub async fn install<F>(&self, input: CoreInput, publish: F) -> Result<(), WsError>
322 where
323 F: FnOnce() + Send + 'static,
324 {
325 let (reply, response) = oneshot::channel();
326 self.inputs
327 .send(ActorInput::Install {
328 input,
329 publish: Box::new(publish),
330 reply,
331 })
332 .await
333 .map_err(|_| WsError::Gone)?;
334 response.await.map_err(|_| WsError::Gone)?
335 }
336
337 pub async fn attach<H: SessionHandler>(
338 &self,
339 session: SessionId,
340 transport: Pipe,
341 expected_peer: Option<String>,
342 handler: H,
343 ) -> Result<Arc<Wire>, WsError> {
344 self.attach_with_ceiling(
345 session,
346 transport,
347 expected_peer,
348 handler,
349 unb_transport::DEFAULT_MAX_FRAME_SIZE,
350 )
351 .await
352 }
353
354 pub async fn attach_with_ceiling<H: SessionHandler>(
355 &self,
356 session: SessionId,
357 transport: Pipe,
358 expected_peer: Option<String>,
359 handler: H,
360 ws_collect_ceiling: usize,
361 ) -> Result<Arc<Wire>, WsError> {
362 let cancellation = self.cancellation.child_token();
363 let (wire, directives, outcome, routes) = Wire::channel(cancellation.clone());
364 let (reply, response) = oneshot::channel();
365 self.inputs
366 .send(ActorInput::Attach {
367 session,
368 transport,
369 expected_peer,
370 directives,
371 handler: handler_channel(handler),
372 client: wire.client_session(),
373 cancellation,
374 establish_peer: true,
375 outcome,
376 routes,
377 ws_collect_ceiling,
378 reply,
379 })
380 .await
381 .map_err(|_| WsError::Gone)?;
382 response.await.map_err(|_| WsError::Gone)??;
383 Ok(wire)
384 }
385
386 pub fn shutdown(&self) {
387 self.cancellation.cancel();
388 }
389
390 pub(crate) fn open(transport: Pipe, runtime: &RuntimeHandle) -> (Arc<Wire>, Self) {
391 struct StandaloneExecutor;
392
393 impl EffectExecutor for StandaloneExecutor {
394 fn execute(&self, _effect: CoreEffect, _handle: ProtocolCoreHandle) -> EffectFuture {
395 Box::pin(async { None })
396 }
397 }
398
399 let cancellation = CancellationToken::new();
400 let handle = Self::spawn(
401 ProtocolCore::new("client"),
402 Arc::new(StandaloneExecutor),
403 cancellation.clone(),
404 runtime,
405 );
406 let session_cancellation = cancellation.child_token();
407 let (wire, directives, outcome, routes) =
408 Wire::standalone_channel(session_cancellation.clone());
409 let (reply, _response) = oneshot::channel();
410 handle
411 .inputs
412 .try_send(ActorInput::Attach {
413 session: SessionId::from("client-session"),
414 transport,
415 expected_peer: None,
416 directives,
417 handler: handler_channel(wire.observation_sender()),
418 client: wire.client_session(),
419 cancellation: session_cancellation,
420 establish_peer: false,
421 outcome,
422 routes,
423 ws_collect_ceiling: unb_transport::DEFAULT_MAX_FRAME_SIZE,
424 reply,
425 })
426 .expect("new protocol core input queue");
427 (wire, handle)
428 }
429}
430
431fn wire_body_stream(body: WireBody) -> BodyStream {
432 match body {
433 WireBody::Bytes(bytes) => Box::pin(futures_util::stream::once(async move { Ok(bytes) })),
434 WireBody::Stream(stream) => stream,
435 }
436}
437
438enum ActorInput {
439 Core(CoreInput),
440 CoreBatch(Vec<CoreInput>),
441 Install {
442 input: CoreInput,
443 publish: Box<dyn FnOnce() + Send>,
444 reply: oneshot::Sender<Result<(), WsError>>,
445 },
446 Directive(SessionId, Directive),
447 Attach {
448 session: SessionId,
449 transport: Pipe,
450 expected_peer: Option<String>,
451 directives: mpsc::Receiver<Directive>,
452 handler: mpsc::Sender<HandlerInput>,
453 client: ClientSession,
454 cancellation: CancellationToken,
455 establish_peer: bool,
456 outcome: tokio::sync::watch::Sender<Option<SessionOutcome>>,
457 routes: tokio::sync::watch::Sender<bool>,
458 ws_collect_ceiling: usize,
459 reply: oneshot::Sender<Result<(), WsError>>,
460 },
461 Write {
462 session: SessionId,
463 envelope: Envelope,
464 completion: Option<EffectId>,
465 relay: bool,
466 },
467 TransportModeSelected {
468 session: SessionId,
469 mode: unb_transport::TransportMode,
470 },
471}
472
473struct SessionIo {
474 outbound: mpsc::Sender<WriteRequest>,
475 byte_budget: Arc<Semaphore>,
476 ws_collect_ceiling: usize,
477 bodies: crate::transport::SessionStreams,
478 handler: mpsc::Sender<HandlerInput>,
479 client: ClientSession,
480 cancellation: CancellationToken,
481 reservation_pending: Arc<AtomicBool>,
482 establish_peer: bool,
483 deadline: tokio::sync::watch::Sender<Option<Instant>>,
484 routes: tokio::sync::watch::Sender<bool>,
485 transport_mode: Option<unb_transport::TransportMode>,
486 pending_mode_writes: VecDeque<PendingModeWrite>,
487 operation_lanes: OperationLanes,
488 operation_slots: Arc<Semaphore>,
489}
490
491struct PendingModeWrite {
492 envelope: Envelope,
493 completion: Option<EffectId>,
494 relay: bool,
495}
496
497type OperationLanes = Arc<std::sync::Mutex<HashMap<String, OperationLane>>>;
498
499#[derive(Clone)]
500struct OperationLane {
501 outbound: Option<mpsc::Sender<OperationWriteRequest>>,
502 control: Option<unb_transport::operation::OperationControl>,
503}
504
505struct OperationWriteRequest {
506 envelope: Envelope,
507 completion: Option<EffectId>,
508 relay: bool,
509 body: Option<BodyStream>,
510 _budget: OwnedSemaphorePermit,
511}
512
513struct WriteRequest {
514 envelope: Envelope,
515 completion: Option<EffectId>,
516 relay: bool,
517 _budget: OwnedSemaphorePermit,
518}
519
520impl WriteRequest {
521 fn weight(envelope: &Envelope) -> u32 {
522 let head = envelope.subject.len()
523 + envelope.id.len()
524 + envelope.corr.as_ref().map_or(0, String::len)
525 + envelope.path.iter().map(String::len).sum::<usize>()
526 + envelope
527 .headers
528 .iter()
529 .map(|(name, value)| name.len() + value.as_str().map_or(16, str::len))
530 .sum::<usize>();
531 (envelope.payload.len() + head).clamp(1, OUTBOUND_BYTE_BUDGET) as u32
532 }
533}
534
535#[derive(Debug)]
536pub(crate) struct MeteredChunk {
537 pub(crate) bytes: Bytes,
538 _window: OwnedSemaphorePermit,
539 _budget: OwnedSemaphorePermit,
540}
541
542pub(crate) struct MeteredBody {
543 stream: BodyStream,
544 budget: Arc<Semaphore>,
545 window: Arc<Semaphore>,
546 pending: Option<Bytes>,
547 cancellation: CancellationToken,
548}
549
550impl MeteredBody {
551 pub(crate) fn new(
552 stream: BodyStream,
553 budget: Arc<Semaphore>,
554 cancellation: CancellationToken,
555 ) -> MeteredBody {
556 MeteredBody {
557 stream,
558 budget,
559 window: Arc::new(Semaphore::new(BODY_WINDOW_BYTES)),
560 pending: None,
561 cancellation,
562 }
563 }
564
565 pub(crate) async fn next_chunk(&mut self) -> Result<Option<MeteredChunk>, WsError> {
566 let chunk = match self.pending.take() {
567 Some(chunk) => chunk,
568 None => {
569 let next = tokio::select! {
570 biased;
571 () = self.cancellation.cancelled() => return Ok(None),
572 next = std::future::poll_fn(|cx| self.stream.as_mut().poll_next(cx)) => next,
573 };
574 match next {
575 Some(Ok(chunk)) => chunk,
576 Some(Err(error)) => return Err(error.into()),
577 None => return Ok(None),
578 }
579 }
580 };
581 let len = chunk.len();
582 if len > unb_transport::DEFAULT_MAX_FRAME_SIZE {
583 return Err(unb_core::CoreError::BodyTooLarge(
584 unb_transport::DEFAULT_MAX_FRAME_SIZE,
585 )
586 .into());
587 }
588 self.pending = Some(chunk);
589 let window_weight = len.clamp(1, BODY_WINDOW_BYTES) as u32;
590 let budget_weight = len.max(1) as u32;
591 let acquire = async {
592 let window = self
593 .window
594 .clone()
595 .acquire_many_owned(window_weight)
596 .await
597 .map_err(|_| unb_core::CoreError::BodyStalled)?;
598 let budget = self
599 .budget
600 .clone()
601 .acquire_many_owned(budget_weight)
602 .await
603 .map_err(|_| unb_core::CoreError::BodyStalled)?;
604 Ok::<_, unb_core::CoreError>((window, budget))
605 };
606 let acquired = tokio::select! {
607 biased;
608 () = self.cancellation.cancelled() => return Ok(None),
609 acquired = n0_future::time::timeout(RELAY_BACKPRESSURE_TIMEOUT, acquire) => acquired,
610 };
611 match acquired {
612 Ok(Ok((window, budget))) => Ok(Some(MeteredChunk {
613 bytes: self.pending.take().expect("pending chunk present"),
614 _window: window,
615 _budget: budget,
616 })),
617 Ok(Err(error)) => Err(error.into()),
618 Err(_) => Err(unb_core::CoreError::BodyStalled.into()),
619 }
620 }
621}
622
623struct CoreRuntime {
624 core: ProtocolCore,
625 inputs: mpsc::Receiver<ActorInput>,
626 handle: ProtocolCoreHandle,
627 executor: Arc<dyn EffectExecutor>,
628 sessions: HashMap<SessionId, SessionIo>,
629 outcomes: HashMap<SessionId, tokio::sync::watch::Sender<Option<SessionOutcome>>>,
630 client_starts: HashMap<SessionId, PendingClientStart>,
631 pending_bodies: std::sync::Mutex<HashMap<(SessionId, String), BodyStream>>,
632}
633
634struct PendingClientStart {
635 sender: mpsc::Sender<ClientDelivery>,
636 reply: oneshot::Sender<Result<ClientOperationId, unb_core::CoreError>>,
637}
638
639impl CoreRuntime {
640 fn new(
641 core: ProtocolCore,
642 inputs: mpsc::Receiver<ActorInput>,
643 handle: ProtocolCoreHandle,
644 executor: Arc<dyn EffectExecutor>,
645 ) -> Self {
646 Self {
647 core,
648 inputs,
649 handle,
650 executor,
651 sessions: HashMap::new(),
652 outcomes: HashMap::new(),
653 client_starts: HashMap::new(),
654 pending_bodies: std::sync::Mutex::new(HashMap::new()),
655 }
656 }
657
658 async fn run(mut self) {
659 loop {
660 let input = tokio::select! {
661 biased;
662 () = self.handle.cancellation.cancelled() => return,
663 input = self.inputs.recv() => input,
664 };
665 let Some(input) = input else { return };
666 match input {
667 ActorInput::Core(input) => {
668 let _ = self.core.handle(clock_now(), input);
669 }
670 ActorInput::CoreBatch(inputs) => {
671 let now = clock_now();
672 for input in inputs {
673 let _ = self.core.handle(now, input);
674 }
675 }
676 ActorInput::Install {
677 input,
678 publish,
679 reply,
680 } => {
681 let result = self.core.handle(clock_now(), input).map_err(WsError::from);
682 if result.is_ok() {
683 publish();
684 }
685 let _ = reply.send(result);
686 }
687 ActorInput::Directive(session, directive) => {
688 self.directive(&session, directive).await;
689 }
690 ActorInput::Attach {
691 session,
692 transport,
693 expected_peer,
694 directives,
695 handler,
696 client,
697 cancellation,
698 establish_peer,
699 outcome,
700 routes,
701 ws_collect_ceiling,
702 reply,
703 } => {
704 let result = self
705 .attach(
706 session,
707 transport,
708 expected_peer,
709 directives,
710 handler,
711 client,
712 cancellation,
713 establish_peer,
714 outcome,
715 routes,
716 ws_collect_ceiling,
717 )
718 .await;
719 let _ = reply.send(result);
720 }
721 ActorInput::Write {
722 session,
723 envelope,
724 completion,
725 relay,
726 } => {
727 self.write(&session, envelope, completion, relay).await;
728 }
729 ActorInput::TransportModeSelected { session, mode } => {
730 let pending = if let Some(io) = self.sessions.get_mut(&session) {
731 io.transport_mode = Some(mode);
732 std::mem::take(&mut io.pending_mode_writes)
733 } else {
734 VecDeque::new()
735 };
736 if mode == unb_transport::TransportMode::OperationStreamsV1 {
737 self.start_operation_acceptor(&session);
738 }
739 for pending in pending {
740 self.write(
741 &session,
742 pending.envelope,
743 pending.completion,
744 pending.relay,
745 )
746 .await;
747 }
748 }
749 }
750 self.drain().await;
751 }
752 }
753
754 #[allow(clippy::too_many_arguments)]
755 async fn attach(
756 &mut self,
757 session: SessionId,
758 transport: Pipe,
759 expected_peer: Option<String>,
760 directives: mpsc::Receiver<Directive>,
761 handler: mpsc::Sender<HandlerInput>,
762 client: ClientSession,
763 cancellation: CancellationToken,
764 establish_peer: bool,
765 outcome: tokio::sync::watch::Sender<Option<SessionOutcome>>,
766 routes: tokio::sync::watch::Sender<bool>,
767 ws_collect_ceiling: usize,
768 ) -> Result<(), WsError> {
769 let initiator = transport.initiator();
770 let bodies = transport.streams();
771 let (reader, writer) = transport.split();
772 let (outbound, outbound_rx) = mpsc::channel(OUTBOUND_QUEUE);
773 let (deadline, deadline_rx) = tokio::sync::watch::channel(None);
774 self.core.handle(
775 clock_now(),
776 CoreInput::SessionOpened {
777 session: session.clone(),
778 initiator,
779 establish_peer,
780 expected_peer,
781 },
782 )?;
783 self.sessions.insert(
784 session.clone(),
785 SessionIo {
786 outbound,
787 byte_budget: Arc::new(Semaphore::new(OUTBOUND_BYTE_BUDGET)),
788 ws_collect_ceiling,
789 bodies: bodies.clone(),
790 handler,
791 client,
792 cancellation: cancellation.clone(),
793 reservation_pending: Arc::new(AtomicBool::new(false)),
794 establish_peer,
795 deadline,
796 routes,
797 transport_mode: None,
798 pending_mode_writes: VecDeque::new(),
799 operation_lanes: Arc::new(std::sync::Mutex::new(HashMap::new())),
800 operation_slots: Arc::new(Semaphore::new(OPERATION_ACTIVE_MAX)),
801 },
802 );
803 self.outcomes.insert(session.clone(), outcome);
804 spawn_session_deadline(
805 session.clone(),
806 deadline_rx,
807 self.handle.clone(),
808 cancellation.clone(),
809 );
810 spawn_directives(
811 session.clone(),
812 directives,
813 self.handle.clone(),
814 cancellation.clone(),
815 );
816 spawn_reader(
817 session.clone(),
818 reader,
819 self.handle.clone(),
820 cancellation.clone(),
821 bodies,
822 );
823 spawn_writer(
824 session.clone(),
825 writer,
826 outbound_rx,
827 self.handle.clone(),
828 cancellation,
829 );
830 if establish_peer {
831 spawn_establishment_timeout(session, self.handle.clone());
832 }
833 Ok(())
834 }
835
836 async fn directive(&mut self, session: &SessionId, directive: Directive) {
837 match directive {
838 Directive::StartClientOperation {
839 subject,
840 kind,
841 payload,
842 hops,
843 headers,
844 body,
845 timeout,
846 sender,
847 reply,
848 } => {
849 let input = if kind == unb_core::Kind::Discover {
850 if body.is_some() {
851 Err(unb_core::CoreError::Malformed(
852 "discovery does not accept a streaming body".into(),
853 ))
854 } else {
855 serde_json::from_slice(&payload)
856 .map(unb_core::OperationInput::Discovery)
857 .map_err(|error| unb_core::CoreError::Malformed(error.to_string()))
858 }
859 } else {
860 register_outbound_body(&self.handle, session, payload, body)
861 .map(unb_core::OperationInput::Body)
862 };
863 let input = match input {
864 Ok(input) => input,
865 Err(error) => {
866 let _ = reply.send(Err(error));
867 return;
868 }
869 };
870 let registered_body = match &input {
871 unb_core::OperationInput::Body(body) => body.clone(),
872 unb_core::OperationInput::Discovery(_) => None,
873 };
874 let result = self.core.handle(
875 clock_now(),
876 CoreInput::StartClientOperation {
877 session: session.clone(),
878 subject,
879 kind,
880 input,
881 hops,
882 headers,
883 timeout,
884 },
885 );
886 if let Err(error) = result {
887 if let Some(body) = ®istered_body {
888 let _ = self.handle.claim_body(session, body.as_str());
889 }
890 let _ = reply.send(Err(error));
891 } else {
892 self.client_starts
893 .insert(session.clone(), PendingClientStart { sender, reply });
894 }
895 }
896 Directive::OpenStream {
897 subject,
898 kind,
899 payload,
900 hops,
901 headers,
902 body,
903 opened,
904 reply,
905 } => {
906 let result = match register_outbound_body(&self.handle, session, payload, body) {
907 Ok(body) => {
908 let result = self.core.open_stream_body(
909 session,
910 &subject,
911 kind,
912 body.clone(),
913 hops,
914 headers,
915 );
916 if result.is_err() {
917 if let Some(body) = &body {
918 let _ = self.handle.claim_body(session, body.as_str());
919 }
920 }
921 result
922 }
923 Err(error) => Err(error),
924 };
925 if let (Ok(corr), Some(opened)) = (&result, opened) {
926 opened(corr.clone()).await;
927 }
928 let _ = reply.send(result);
929 }
930 Directive::Send { corr, payload } => {
931 match register_outbound_body(&self.handle, session, payload, None) {
932 Ok(body) => {
933 if self.core.send_body(session, &corr, body.clone()).is_err() {
934 if let Some(body) = &body {
935 let _ = self.handle.claim_body(session, body.as_str());
936 }
937 }
938 }
939 Err(error) => {
940 let _ = self.core.fail(
941 session,
942 &corr,
943 unb_core::ErrorCode::Busy,
944 &error.to_string(),
945 );
946 }
947 }
948 }
949 Directive::Respond {
950 corr,
951 payload,
952 headers,
953 body,
954 } => match register_outbound_body(&self.handle, session, payload, body) {
955 Ok(body) => {
956 if self
957 .core
958 .respond_body(session, &corr, body.clone(), headers)
959 .is_err()
960 {
961 if let Some(body) = &body {
962 let _ = self.handle.claim_body(session, body.as_str());
963 }
964 }
965 }
966 Err(error) => {
967 let _ = self.core.fail(
968 session,
969 &corr,
970 unb_core::ErrorCode::Busy,
971 &error.to_string(),
972 );
973 }
974 },
975 Directive::Fail {
976 corr,
977 code,
978 message,
979 } => {
980 let _ = self.core.fail(session, &corr, code, &message);
981 }
982 Directive::Cancel { corr } => {
983 let _ = self.core.cancel(session, &corr);
984 }
985 Directive::CancelClientOperation { operation } => {
986 let _ = self.core.handle(
987 clock_now(),
988 CoreInput::CancelClientOperation {
989 session: session.clone(),
990 operation,
991 },
992 );
993 }
994 Directive::Control { kind, payload } => {
995 let _ = self.core.control(session, kind, payload);
996 }
997 }
998 }
999
1000 async fn drain(&mut self) {
1001 while let Some(effect) = self.core.poll_effect() {
1002 match effect {
1003 CoreEffect::SendFrame { session, envelope } => {
1004 self.write(&session, envelope, None, false).await
1005 }
1006 CoreEffect::HandshakeEstablished { session, .. } => {
1007 if let Some(bodies) =
1008 self.sessions.get(&session).and_then(|io| io.bodies.clone())
1009 {
1010 let inputs = self.handle.inputs.clone();
1011 let mode_session = session.clone();
1012 n0_future::task::spawn(async move {
1013 let mode = bodies.negotiate().await;
1014 let _ = inputs
1015 .send(ActorInput::TransportModeSelected {
1016 session: mode_session,
1017 mode,
1018 })
1019 .await;
1020 });
1021 }
1022 if self
1023 .sessions
1024 .get(&session)
1025 .is_some_and(|io| !io.establish_peer)
1026 {
1027 if let Some(outcome) = self.outcomes.get(&session) {
1028 let _ = outcome.send(Some(SessionOutcome::Established));
1029 }
1030 }
1031 }
1032 CoreEffect::DeliverClient {
1033 session,
1034 operation,
1035 delivery,
1036 } => {
1037 let body = match &delivery {
1038 CoreClientDelivery::Item(frame) | CoreClientDelivery::Terminal(frame) => {
1039 frame
1040 .body
1041 .as_ref()
1042 .and_then(|body| self.handle.claim_body(&session, body.as_str()))
1043 }
1044 _ => None,
1045 };
1046 let terminal = matches!(
1047 delivery,
1048 CoreClientDelivery::Terminal(_)
1049 | CoreClientDelivery::Cancelled
1050 | CoreClientDelivery::TimedOut
1051 | CoreClientDelivery::SessionClosed
1052 );
1053 let client = self.sessions.get(&session).map(|io| io.client.clone());
1054 if let Some(client) = client {
1055 let delivered = n0_future::time::timeout(
1056 CLIENT_DELIVERY_TIMEOUT,
1057 client.deliver(&operation, delivery, body),
1058 )
1059 .await
1060 .is_ok();
1061 if !delivered {
1062 client.abandon(&operation);
1063 let _ = self.core.handle(
1064 clock_now(),
1065 CoreInput::CancelClientOperation {
1066 session: session.clone(),
1067 operation: operation.clone(),
1068 },
1069 );
1070 }
1071 }
1072 if terminal {
1073 self.handler(&session, HandlerInput::StreamClosed(operation))
1074 .await;
1075 }
1076 }
1077 CoreEffect::RegisterClient { session, operation } => {
1078 let Some(start) = self.client_starts.remove(&session) else {
1079 continue;
1080 };
1081 let client = self.sessions.get(&session).map(|io| io.client.clone());
1082 let result = if let Some(client) = &client {
1083 client
1084 .register(operation.clone(), start.sender)
1085 .map(|()| operation)
1086 .map_err(|error| {
1087 unb_core::CoreError::UnknownStream(error.to_string())
1088 })
1089 } else {
1090 Err(unb_core::CoreError::UnknownSession(session.to_string()))
1091 };
1092 if let Err(Ok(operation)) = start.reply.send(result) {
1093 if let Some(client) = client {
1094 client.abandon(&operation);
1095 }
1096 let _ = self.core.handle(
1097 clock_now(),
1098 CoreInput::CancelClientOperation { session, operation },
1099 );
1100 }
1101 }
1102 CoreEffect::Deliver {
1103 session,
1104 mut envelope,
1105 } => {
1106 if let Some(token) = envelope.body_token.take() {
1107 match self.handle.claim_body(&session, &token) {
1108 Some(WireBody::Bytes(payload)) => {
1109 envelope.payload = payload;
1110 }
1111 Some(body @ WireBody::Stream(_)) => {
1112 let ceiling = self
1113 .sessions
1114 .get(&session)
1115 .map_or(unb_transport::DEFAULT_MAX_FRAME_SIZE, |io| {
1116 io.ws_collect_ceiling
1117 });
1118 match body.collect_to(ceiling).await {
1119 Ok(payload) => envelope.payload = payload,
1120 Err(error) => {
1121 envelope.kind = unb_core::Kind::Error;
1122 envelope.payload =
1123 Envelope::encode_payload(&serde_json::json!({
1124 "code": unb_core::ErrorCode::PayloadTooLarge,
1125 "message": error.to_string(),
1126 }));
1127 }
1128 }
1129 }
1130 None => {}
1131 }
1132 }
1133 self.handler(&session, HandlerInput::Deliver(envelope))
1134 .await
1135 }
1136 CoreEffect::ReleaseBody { session, body } => {
1137 let _ = self.handle.claim_body(&session, body.as_str());
1138 }
1139 CoreEffect::StreamClosed { session, operation } => {
1140 self.handler(&session, HandlerInput::StreamClosed(operation))
1141 .await
1142 }
1143 CoreEffect::CloseTransport {
1144 session,
1145 code,
1146 message,
1147 } => {
1148 let _ = (code, message);
1149 if let Some(io) = self.sessions.get(&session) {
1150 io.cancellation.cancel();
1151 }
1152 }
1153 CoreEffect::ScheduleSessionDeadline { session, deadline } => {
1154 if let Some(io) = self.sessions.get(&session) {
1155 io.deadline.send_replace(Some(deadline));
1156 }
1157 }
1158 CoreEffect::RouteExportAcked { session } => {
1159 if let Some(io) = self.sessions.get(&session) {
1160 io.routes.send_replace(true);
1161 }
1162 }
1163 CoreEffect::ScheduleClientDeadline {
1164 session,
1165 operation,
1166 deadline,
1167 } => spawn_client_timeout(
1168 session.clone(),
1169 operation,
1170 deadline,
1171 self.handle.clone(),
1172 self.sessions
1173 .get(&session)
1174 .map(|io| io.cancellation.clone()),
1175 ),
1176 CoreEffect::Send {
1177 effect,
1178 session,
1179 frame,
1180 } => {
1181 if frame.head.kind == unb_core::Kind::Error {
1182 let _ = self.core.handle(
1183 clock_now(),
1184 CoreInput::SendCompleted {
1185 effect,
1186 result: SendResult::Written,
1187 },
1188 );
1189 self.write(&session, frame.into_envelope(), None, false)
1190 .await;
1191 } else {
1192 self.write(&session, frame.into_envelope(), Some(effect), false)
1193 .await;
1194 }
1195 }
1196 CoreEffect::SendProtocol {
1197 effect,
1198 session,
1199 envelope,
1200 } => {
1201 self.write(&session, envelope, Some(effect), false).await;
1202 }
1203 CoreEffect::ContinueDiscovery { stream } => {
1204 let _ = self
1205 .core
1206 .handle(clock_now(), CoreInput::ContinueDiscovery { stream });
1207 }
1208 CoreEffect::SessionRetired { session, reason } => {
1209 #[cfg(feature = "observability")]
1210 {
1211 metrics::counter!("unb_sessions_retired", "reason" => format!("{reason:?}"))
1212 .increment(1);
1213 tracing::debug!(session = %session, reason = ?reason, "session retired");
1214 }
1215 if let Some(outcome) = self.outcomes.remove(&session) {
1216 let _ = outcome.send(Some(SessionOutcome::Retired(reason)));
1217 }
1218 if let Some(io) = self.sessions.remove(&session) {
1219 io.cancellation.cancel();
1220 }
1221 self.pending_bodies
1222 .lock()
1223 .expect("pending body registry")
1224 .retain(|(owner, _), _| owner != &session);
1225 self.handle
1226 .token_bodies
1227 .lock()
1228 .expect("body token registry")
1229 .retain(|(owner, _), _| owner != &session);
1230 self.execute(CoreEffect::SessionRetired { session, reason });
1231 }
1232 CoreEffect::SessionEstablished { session, peer } => {
1233 let effect = CoreEffect::SessionEstablished {
1234 session: session.clone(),
1235 peer,
1236 };
1237 if let Some(input) = self.executor.execute(effect, self.handle.clone()).await {
1238 let _ = self.core.handle(clock_now(), input);
1239 }
1240 if let Some(outcome) = self.outcomes.get(&session) {
1241 let _ = outcome.send(Some(SessionOutcome::Established));
1242 }
1243 }
1244 CoreEffect::ForwardRelay {
1245 effect,
1246 source,
1247 target,
1248 mut frame,
1249 ..
1250 } => {
1251 if let Some(body_id) = frame.body.take() {
1252 let Some(body) = self.handle.claim_body(&source.session, body_id.as_str())
1253 else {
1254 let _ = self.core.handle(
1255 clock_now(),
1256 CoreInput::RelayForwardCompleted {
1257 effect,
1258 result: SendResult::Refused {
1259 code: unb_core::ErrorCode::Protocol,
1260 message: "relay body capability is unavailable".into(),
1261 },
1262 },
1263 );
1264 continue;
1265 };
1266 match self.handle.register_body(&target.session, body) {
1267 Ok(target_body) => frame.body = Some(target_body),
1268 Err(error) => {
1269 let _ = self.core.handle(
1270 clock_now(),
1271 CoreInput::RelayForwardCompleted {
1272 effect,
1273 result: SendResult::Refused {
1274 code: unb_core::ErrorCode::Busy,
1275 message: error.to_string(),
1276 },
1277 },
1278 );
1279 continue;
1280 }
1281 }
1282 }
1283 self.write(&target.session, frame.into_envelope(), Some(effect), true)
1284 .await;
1285 }
1286 effect @ (CoreEffect::RouteSnapshotApplied { .. }
1287 | CoreEffect::RouteDeltaApplied { .. }
1288 | CoreEffect::RouteSessionWithdrawn { .. }) => {
1289 if let Some(input) = self.executor.execute(effect, self.handle.clone()).await {
1290 let _ = self.core.handle(clock_now(), input);
1291 }
1292 }
1293 effect => self.execute(effect),
1294 }
1295 }
1296 }
1297
1298 fn execute(&self, effect: CoreEffect) {
1299 let executor = self.executor.clone();
1300 let handle = self.handle.clone();
1301 n0_future::task::spawn(async move {
1302 if let Some(input) = executor.execute(effect, handle.clone()).await {
1303 let _ = handle.submit(input).await;
1304 }
1305 });
1306 }
1307
1308 async fn write(
1309 &mut self,
1310 session: &SessionId,
1311 mut envelope: Envelope,
1312 completion: Option<EffectId>,
1313 relay: bool,
1314 ) {
1315 if is_operation_envelope(&envelope) {
1316 let mode = self.sessions.get(session).and_then(|io| io.transport_mode);
1317 match mode {
1318 None if self
1319 .sessions
1320 .get(session)
1321 .is_some_and(|io| io.bodies.is_some()) =>
1322 {
1323 let queued = self.sessions.get_mut(session).is_some_and(|io| {
1324 if io.pending_mode_writes.len() >= PENDING_MODE_WRITES_MAX {
1325 false
1326 } else {
1327 io.pending_mode_writes.push_back(PendingModeWrite {
1328 envelope,
1329 completion,
1330 relay,
1331 });
1332 true
1333 }
1334 });
1335 if !queued {
1336 if let Some(effect) = completion {
1337 self.complete_send(effect, relay, SendResult::ReservationTimedOut);
1338 }
1339 }
1340 return;
1341 }
1342 Some(unb_transport::TransportMode::OperationStreamsV1) => {
1343 self.write_operation(session, envelope, completion, relay)
1344 .await;
1345 return;
1346 }
1347 _ => {}
1348 }
1349 }
1350 if envelope.kind == unb_core::Kind::Cancel {
1351 if let Some(corr) = &envelope.corr {
1352 self.handle.cancel_body_pump(session, corr);
1353 }
1354 }
1355 let body = match (&envelope.body_token, &envelope.corr) {
1356 (Some(token), _) => self.handle.claim_body(session, token),
1357 (None, Some(corr))
1358 if envelope.kind.is_application_request()
1359 || envelope.kind == unb_core::Kind::Response =>
1360 {
1361 let pending = self
1362 .pending_bodies
1363 .lock()
1364 .expect("pending body registry")
1365 .remove(&(session.clone(), corr.clone()));
1366 if let Some(stream) = pending {
1367 envelope.body_token = Some(mint_body_token());
1368 Some(WireBody::Stream(stream))
1369 } else {
1370 None
1371 }
1372 }
1373 _ => None,
1374 };
1375 let Some(io) = self.sessions.get(session) else {
1376 if let Some(effect) = completion {
1377 self.complete_send(effect, relay, SendResult::Closed);
1378 }
1379 return;
1380 };
1381 if let Some(WireBody::Bytes(payload)) = body {
1382 envelope.body_token = None;
1383 envelope.payload = payload;
1384 } else if let Some(WireBody::Stream(stream)) = body {
1385 if let Some(peer) = io.bodies.clone().filter(|bodies| bodies.can_open()) {
1386 let token = envelope.body_token.clone().expect("tokenized body");
1387 let pump_cancellation = io.cancellation.child_token();
1388 if let Some(corr) = &envelope.corr {
1389 self.handle.register_body_pump(
1390 session,
1391 corr,
1392 &token,
1393 pump_cancellation.clone(),
1394 );
1395 }
1396 spawn_body_pump(
1397 peer,
1398 token,
1399 stream,
1400 io.byte_budget.clone(),
1401 pump_cancellation,
1402 self.handle.clone(),
1403 envelope.corr.clone().map(|corr| (session.clone(), corr)),
1404 );
1405 } else {
1406 envelope.body_token = None;
1407 spawn_body_collect(
1408 self.handle.clone(),
1409 session.clone(),
1410 envelope,
1411 stream,
1412 io.ws_collect_ceiling,
1413 completion,
1414 relay,
1415 io.cancellation.clone(),
1416 );
1417 return;
1418 }
1419 }
1420 let envelope = match io
1421 .byte_budget
1422 .clone()
1423 .try_acquire_many_owned(WriteRequest::weight(&envelope))
1424 {
1425 Ok(budget) => {
1426 let request = WriteRequest {
1427 envelope,
1428 completion,
1429 relay,
1430 _budget: budget,
1431 };
1432 match io.outbound.try_send(request) {
1433 Ok(()) => {
1434 if let Some(effect) = completion {
1435 self.complete_send(effect, relay, SendResult::Reserved);
1436 }
1437 return;
1438 }
1439 Err(mpsc::error::TrySendError::Closed(_)) => {
1440 if let Some(effect) = completion {
1441 self.complete_send(effect, relay, SendResult::Closed);
1442 }
1443 return;
1444 }
1445 Err(mpsc::error::TrySendError::Full(request)) => {
1446 #[cfg(feature = "observability")]
1447 {
1448 metrics::counter!("unb_write_queue_full").increment(1);
1449 tracing::warn!(session = %session, "outbound write queue is full");
1450 }
1451 request.envelope
1452 }
1453 }
1454 }
1455 Err(_) => {
1456 #[cfg(feature = "observability")]
1457 {
1458 metrics::counter!("unb_write_budget_exhausted").increment(1);
1459 tracing::warn!(session = %session, "outbound byte budget is exhausted");
1460 }
1461 envelope
1462 }
1463 };
1464 if completion.is_none()
1465 || io
1466 .reservation_pending
1467 .compare_exchange(false, true, Ordering::AcqRel, Ordering::Acquire)
1468 .is_ok()
1469 {
1470 let timeout = if matches!(
1471 envelope.kind,
1472 unb_core::Kind::Response
1473 | unb_core::Kind::Error
1474 | unb_core::Kind::Cancel
1475 ) {
1476 TERMINAL_SEND_TIMEOUT
1477 } else {
1478 RELAY_BACKPRESSURE_TIMEOUT
1479 };
1480 reserve_write(
1481 io.outbound.clone(),
1482 io.byte_budget.clone(),
1483 envelope,
1484 completion,
1485 relay,
1486 self.handle.clone(),
1487 io.cancellation.clone(),
1488 completion.map(|_| io.reservation_pending.clone()),
1489 timeout,
1490 );
1491 } else if let Some(effect) = completion {
1492 self.complete_send(effect, relay, SendResult::ReservationTimedOut);
1493 }
1494 }
1495
1496 fn start_operation_acceptor(&self, session: &SessionId) {
1497 let Some(io) = self.sessions.get(session) else {
1498 return;
1499 };
1500 let Some(bodies) = io.bodies.clone() else {
1501 return;
1502 };
1503 spawn_operation_acceptor(
1504 session.clone(),
1505 bodies,
1506 io.operation_lanes.clone(),
1507 io.operation_slots.clone(),
1508 io.byte_budget.clone(),
1509 self.handle.clone(),
1510 io.cancellation.clone(),
1511 );
1512 }
1513
1514 async fn write_operation(
1515 &mut self,
1516 session: &SessionId,
1517 mut envelope: Envelope,
1518 completion: Option<EffectId>,
1519 relay: bool,
1520 ) {
1521 let Some(corr) = envelope.corr.clone() else {
1522 if let Some(effect) = completion {
1523 self.complete_send(
1524 effect,
1525 relay,
1526 SendResult::WriteFailed("missing unb-corr".into()),
1527 );
1528 }
1529 return;
1530 };
1531 let body = match &envelope.body_token {
1532 Some(token) => self.handle.claim_body(session, token),
1533 None if envelope.kind.is_application_request()
1534 || envelope.kind == unb_core::Kind::Response =>
1535 {
1536 self.pending_bodies
1537 .lock()
1538 .expect("pending body registry")
1539 .remove(&(session.clone(), corr.clone()))
1540 .map(WireBody::Stream)
1541 }
1542 None => None,
1543 };
1544 let body = match body {
1545 Some(WireBody::Bytes(payload)) => {
1546 envelope.payload = payload;
1547 None
1548 }
1549 Some(WireBody::Stream(stream)) => Some(stream),
1550 None => None,
1551 };
1552 envelope.body_token = None;
1553 let Some(io) = self.sessions.get(session) else {
1554 if let Some(effect) = completion {
1555 self.complete_send(effect, relay, SendResult::Closed);
1556 }
1557 return;
1558 };
1559 if envelope.kind == unb_core::Kind::Cancel {
1560 let lane = io
1561 .operation_lanes
1562 .lock()
1563 .expect("operation lane registry")
1564 .remove(&corr);
1565 if let Some(control) = lane.and_then(|lane| lane.control) {
1566 let _ = control.reset(unb_transport::operation::RESET_CANCELLED);
1567 let _ = control.stop(unb_transport::operation::RESET_CANCELLED);
1568 }
1569 if let Some(effect) = completion {
1570 self.complete_send(effect, relay, SendResult::Written);
1571 }
1572 return;
1573 }
1574 if envelope.kind.is_application_request() {
1575 spawn_outbound_operation(
1576 session.clone(),
1577 corr,
1578 envelope,
1579 body,
1580 completion,
1581 relay,
1582 io.bodies
1583 .clone()
1584 .expect("operation mode has stream facility"),
1585 io.operation_lanes.clone(),
1586 io.operation_slots.clone(),
1587 io.byte_budget.clone(),
1588 self.handle.clone(),
1589 io.cancellation.clone(),
1590 );
1591 return;
1592 }
1593 let lane = io
1594 .operation_lanes
1595 .lock()
1596 .expect("operation lane registry")
1597 .get(&corr)
1598 .cloned();
1599 let Some(sender) = lane.and_then(|lane| lane.outbound) else {
1600 if let Some(effect) = completion {
1601 self.complete_send(effect, relay, SendResult::Written);
1602 }
1603 return;
1604 };
1605 let budget = match io
1606 .byte_budget
1607 .clone()
1608 .try_acquire_many_owned(WriteRequest::weight(&envelope))
1609 {
1610 Ok(budget) => budget,
1611 Err(_) => {
1612 if let Some(effect) = completion {
1613 self.complete_send(effect, relay, SendResult::ReservationTimedOut);
1614 }
1615 return;
1616 }
1617 };
1618 let request = OperationWriteRequest {
1619 envelope,
1620 completion,
1621 relay,
1622 body,
1623 _budget: budget,
1624 };
1625 match sender.try_send(request) {
1626 Ok(()) => {
1627 if let Some(effect) = completion {
1628 self.complete_send(effect, relay, SendResult::Reserved);
1629 }
1630 }
1631 Err(_) => {
1632 if let Some(effect) = completion {
1633 self.complete_send(effect, relay, SendResult::ReservationTimedOut);
1634 }
1635 }
1636 }
1637 }
1638
1639 fn complete_send(&mut self, effect: EffectId, relay: bool, result: SendResult) {
1640 let _ = self
1641 .core
1642 .handle(clock_now(), send_completion(effect, relay, result));
1643 }
1644
1645 async fn handler(&self, session: &SessionId, event: HandlerInput) {
1646 if let Some(io) = self.sessions.get(session) {
1647 let _ = io.handler.send(event).await;
1648 }
1649 }
1650}
1651
1652enum HandlerInput {
1653 Deliver(Envelope),
1654 StreamClosed(ClientOperationId),
1655}
1656
1657fn handler_channel<H: SessionHandler>(mut handler: H) -> mpsc::Sender<HandlerInput> {
1658 let (events, mut receiver) = mpsc::channel(SESSION_QUEUE);
1659 n0_future::task::spawn(async move {
1660 while let Some(event) = receiver.recv().await {
1661 match event {
1662 HandlerInput::Deliver(envelope) => handler.deliver(envelope).await,
1663 HandlerInput::StreamClosed(operation) => handler.stream_closed(operation).await,
1664 }
1665 }
1666 });
1667 events
1668}
1669
1670fn spawn_directives(
1671 session: SessionId,
1672 mut directives: mpsc::Receiver<Directive>,
1673 handle: ProtocolCoreHandle,
1674 cancellation: CancellationToken,
1675) {
1676 n0_future::task::spawn(async move {
1677 loop {
1678 let directive = tokio::select! {
1679 biased;
1680 () = cancellation.cancelled() => return,
1681 directive = directives.recv() => directive,
1682 };
1683 let Some(directive) = directive else { return };
1684 if handle
1685 .inputs
1686 .send(ActorInput::Directive(session.clone(), directive))
1687 .await
1688 .is_err()
1689 {
1690 return;
1691 }
1692 }
1693 });
1694}
1695
1696fn is_operation_envelope(envelope: &Envelope) -> bool {
1697 envelope.corr.is_some()
1698 && (envelope.kind.is_application_request()
1699 || envelope.kind.is_application_response()
1700 || envelope.kind == unb_core::Kind::Cancel)
1701}
1702
1703fn spawn_operation_acceptor(
1704 session: SessionId,
1705 bodies: Arc<unb_transport::ConnectionStreams>,
1706 lanes: OperationLanes,
1707 operation_slots: Arc<Semaphore>,
1708 byte_budget: Arc<Semaphore>,
1709 handle: ProtocolCoreHandle,
1710 cancellation: CancellationToken,
1711) {
1712 n0_future::task::spawn(async move {
1713 loop {
1714 let operation = tokio::select! {
1715 biased;
1716 () = cancellation.cancelled() => return,
1717 operation = bodies.accept_operation() => operation,
1718 };
1719 let Ok(operation) = operation else {
1720 return;
1721 };
1722 let Ok(operation_slot) = operation_slots.clone().try_acquire_owned() else {
1723 if let Some(control) = operation.control {
1724 let _ = control.reset(unb_transport::operation::RESET_BUSY);
1725 let _ = control.stop(unb_transport::operation::RESET_BUSY);
1726 }
1727 continue;
1728 };
1729 let corr = operation.corr.clone();
1730 let (outbound, receiver) = mpsc::channel(OPERATION_WRITE_QUEUE);
1731 let lane = OperationLane {
1732 outbound: Some(outbound),
1733 control: operation.control.clone(),
1734 };
1735 let admitted = {
1736 let mut lanes = lanes.lock().expect("operation lane registry");
1737 if lanes.contains_key(&corr) {
1738 false
1739 } else {
1740 lanes.insert(corr.clone(), lane);
1741 true
1742 }
1743 };
1744 if !admitted {
1745 if let Some(control) = operation.control {
1746 let _ = control.reset(unb_transport::operation::RESET_PROTOCOL);
1747 }
1748 continue;
1749 }
1750 spawn_accepted_operation(
1751 session.clone(),
1752 operation,
1753 receiver,
1754 lanes.clone(),
1755 byte_budget.clone(),
1756 handle.clone(),
1757 cancellation.child_token(),
1758 operation_slot,
1759 );
1760 }
1761 });
1762}
1763
1764#[allow(clippy::too_many_arguments)]
1765fn spawn_outbound_operation(
1766 session: SessionId,
1767 corr: String,
1768 envelope: Envelope,
1769 body: Option<BodyStream>,
1770 completion: Option<EffectId>,
1771 relay: bool,
1772 bodies: Arc<unb_transport::ConnectionStreams>,
1773 lanes: OperationLanes,
1774 operation_slots: Arc<Semaphore>,
1775 byte_budget: Arc<Semaphore>,
1776 handle: ProtocolCoreHandle,
1777 cancellation: CancellationToken,
1778) {
1779 n0_future::task::spawn(async move {
1780 let open_byte_budget = byte_budget.clone();
1781 let open = async {
1782 let operation_slot = operation_slots
1783 .acquire_owned()
1784 .await
1785 .map_err(|_| unb_transport::TransportError::Closed)?;
1786 let budget = open_byte_budget
1787 .acquire_many_owned(WriteRequest::weight(&envelope))
1788 .await
1789 .map_err(|_| unb_transport::TransportError::Closed)?;
1790 let operation = bodies
1791 .open_operation(&corr)
1792 .await?
1793 .ok_or(unb_transport::TransportError::Closed)?;
1794 Ok::<_, unb_transport::TransportError>((operation, budget, operation_slot))
1795 };
1796 let opened = tokio::select! {
1797 biased;
1798 () = cancellation.cancelled() => return,
1799 opened = n0_future::time::timeout(RELAY_BACKPRESSURE_TIMEOUT, open) => opened,
1800 };
1801 let Ok(Ok((operation, _budget, _operation_slot))) = opened else {
1802 if let Some(effect) = completion {
1803 let _ = handle
1804 .submit(send_completion(
1805 effect,
1806 relay,
1807 SendResult::ReservationTimedOut,
1808 ))
1809 .await;
1810 }
1811 return;
1812 };
1813 lanes.lock().expect("operation lane registry").insert(
1814 corr.clone(),
1815 OperationLane {
1816 outbound: None,
1817 control: operation.control.clone(),
1818 },
1819 );
1820 if let Some(effect) = completion {
1821 if handle
1822 .submit(send_completion(effect, relay, SendResult::Reserved))
1823 .await
1824 .is_err()
1825 {
1826 return;
1827 }
1828 }
1829 let mut writer = unb_transport::operation::OperationRecordWriter::new(operation.write);
1830 let sent = write_operation_envelope(
1831 &mut writer,
1832 envelope,
1833 body,
1834 byte_budget.clone(),
1835 cancellation.clone(),
1836 )
1837 .await;
1838 let sent = match sent {
1839 Ok(()) => writer.finish().await,
1840 Err(error) => Err(error),
1841 };
1842 if let Some(effect) = completion {
1843 let _ = handle
1844 .submit(send_completion(effect, relay, SendResult::Written))
1845 .await;
1846 }
1847 if let Err(error) = sent {
1848 if !operation_peer_stopped_cancelled(&error) {
1849 if let Some(control) = &operation.control {
1850 let _ = control.reset(error.reset_code());
1851 }
1852 let _ = submit_operation_end(
1853 &handle,
1854 &session,
1855 &corr,
1856 OperationDirection::Opening,
1857 operation_outcome(&error),
1858 )
1859 .await;
1860 lanes.lock().expect("operation lane registry").remove(&corr);
1861 return;
1862 }
1863 }
1864 let result = read_operation_direction(
1865 &session,
1866 &corr,
1867 operation.read,
1868 OperationDirection::Return,
1869 &handle,
1870 &cancellation,
1871 )
1872 .await;
1873 if let Err(error) = result {
1874 if let Some(control) = &operation.control {
1875 let _ = control.reset(error.reset_code());
1876 let _ = control.stop(error.reset_code());
1877 }
1878 let _ = submit_operation_end(
1879 &handle,
1880 &session,
1881 &corr,
1882 OperationDirection::Return,
1883 operation_outcome(&error),
1884 )
1885 .await;
1886 }
1887 lanes.lock().expect("operation lane registry").remove(&corr);
1888 });
1889}
1890
1891#[allow(clippy::too_many_arguments)]
1892fn spawn_accepted_operation(
1893 session: SessionId,
1894 operation: unb_transport::operation::OperationIo,
1895 receiver: mpsc::Receiver<OperationWriteRequest>,
1896 lanes: OperationLanes,
1897 byte_budget: Arc<Semaphore>,
1898 handle: ProtocolCoreHandle,
1899 cancellation: CancellationToken,
1900 operation_slot: OwnedSemaphorePermit,
1901) {
1902 let corr = operation.corr.clone();
1903 let control = operation.control.clone();
1904 let read_corr = corr.clone();
1905 let read_lanes = lanes.clone();
1906 let read_handle = handle.clone();
1907 let read_cancellation = cancellation.clone();
1908 let writer_session = session.clone();
1909 n0_future::task::spawn(async move {
1910 let result = read_operation_direction(
1911 &session,
1912 &read_corr,
1913 operation.read,
1914 OperationDirection::Opening,
1915 &read_handle,
1916 &read_cancellation,
1917 )
1918 .await;
1919 if let Err(error) = result {
1920 if operation_body_consumer_closed(&error) {
1921 if let Some(control) = control {
1922 let _ = control.stop(unb_transport::operation::RESET_CANCELLED);
1923 }
1924 return;
1925 }
1926 if let Some(control) = control {
1927 let _ = control.reset(error.reset_code());
1928 let _ = control.stop(error.reset_code());
1929 }
1930 let _ = submit_operation_end(
1931 &read_handle,
1932 &session,
1933 &read_corr,
1934 OperationDirection::Opening,
1935 operation_outcome(&error),
1936 )
1937 .await;
1938 read_lanes
1939 .lock()
1940 .expect("operation lane registry")
1941 .remove(&read_corr);
1942 }
1943 });
1944 spawn_operation_writer(
1945 writer_session,
1946 corr,
1947 operation.write,
1948 operation.control,
1949 receiver,
1950 lanes,
1951 byte_budget,
1952 handle,
1953 cancellation,
1954 operation_slot,
1955 );
1956}
1957
1958#[allow(clippy::too_many_arguments)]
1959fn spawn_operation_writer(
1960 session: SessionId,
1961 corr: String,
1962 write: Box<dyn tokio::io::AsyncWrite + Send + Unpin>,
1963 control: Option<unb_transport::operation::OperationControl>,
1964 mut receiver: mpsc::Receiver<OperationWriteRequest>,
1965 lanes: OperationLanes,
1966 byte_budget: Arc<Semaphore>,
1967 handle: ProtocolCoreHandle,
1968 cancellation: CancellationToken,
1969 _operation_slot: OwnedSemaphorePermit,
1970) {
1971 n0_future::task::spawn(async move {
1972 let mut writer = unb_transport::operation::OperationRecordWriter::new(write);
1973 loop {
1974 let request = tokio::select! {
1975 biased;
1976 () = cancellation.cancelled() => break,
1977 request = receiver.recv() => request,
1978 };
1979 let Some(request) = request else {
1980 break;
1981 };
1982 let terminal = matches!(
1983 request.envelope.kind,
1984 unb_core::Kind::Response | unb_core::Kind::Error
1985 );
1986 let result = write_operation_envelope(
1987 &mut writer,
1988 request.envelope,
1989 request.body,
1990 byte_budget.clone(),
1991 cancellation.clone(),
1992 )
1993 .await;
1994 let result = if result.is_ok() && terminal {
1995 writer.finish().await
1996 } else {
1997 result
1998 };
1999 if let Some(effect) = request.completion {
2000 let _ = handle
2001 .submit(send_completion(effect, request.relay, SendResult::Written))
2002 .await;
2003 }
2004 if result.is_err() || terminal {
2005 if let Err(error) = result {
2006 if let Some(control) = &control {
2007 let _ = control.reset(error.reset_code());
2008 }
2009 let _ = submit_operation_end(
2010 &handle,
2011 &session,
2012 &corr,
2013 OperationDirection::Return,
2014 operation_outcome(&error),
2015 )
2016 .await;
2017 }
2018 break;
2019 }
2020 }
2021 lanes.lock().expect("operation lane registry").remove(&corr);
2022 });
2023}
2024
2025async fn write_operation_envelope<W>(
2026 writer: &mut unb_transport::operation::OperationRecordWriter<W>,
2027 mut envelope: Envelope,
2028 body: Option<BodyStream>,
2029 byte_budget: Arc<Semaphore>,
2030 cancellation: CancellationToken,
2031) -> Result<(), unb_transport::operation::OperationCodecError>
2032where
2033 W: tokio::io::AsyncWrite + Unpin,
2034{
2035 let Some(body) = body else {
2036 return writer
2037 .write_record(
2038 unb_transport::operation::OperationRecordKind::CompleteMessage,
2039 &envelope.encode(),
2040 )
2041 .await;
2042 };
2043 envelope.body_token = None;
2044 envelope.payload = Bytes::new();
2045 writer
2046 .write_record(
2047 unb_transport::operation::OperationRecordKind::Head,
2048 &envelope.encode(),
2049 )
2050 .await?;
2051 let mut body = MeteredBody::new(body, byte_budget, cancellation.clone());
2052 loop {
2053 let chunk = body.next_chunk().await.map_err(|error| {
2054 unb_transport::operation::OperationCodecError::Io(error.to_string())
2055 })?;
2056 let Some(chunk) = chunk else {
2057 if cancellation.is_cancelled() {
2058 return Err(unb_transport::operation::OperationCodecError::Io(
2059 "operation body cancelled".into(),
2060 ));
2061 }
2062 break;
2063 };
2064 for piece in chunk.bytes.chunks(OPERATION_BODY_CHUNK_BYTES) {
2065 writer
2066 .write_record(
2067 unb_transport::operation::OperationRecordKind::BodyChunk,
2068 piece,
2069 )
2070 .await?;
2071 }
2072 }
2073 writer
2074 .write_record(
2075 unb_transport::operation::OperationRecordKind::BodyEnd,
2076 &[],
2077 )
2078 .await
2079}
2080
2081#[derive(Clone, Copy)]
2082enum OperationDirection {
2083 Opening,
2084 Return,
2085}
2086
2087async fn read_operation_direction(
2088 session: &SessionId,
2089 corr: &str,
2090 read: Box<dyn tokio::io::AsyncRead + Send + Unpin>,
2091 direction: OperationDirection,
2092 handle: &ProtocolCoreHandle,
2093 cancellation: &CancellationToken,
2094) -> Result<(), unb_transport::operation::OperationCodecError> {
2095 let mut reader = unb_transport::operation::OperationRecordReader::new(read);
2096 let mut records = 0usize;
2097 let mut terminal = false;
2098 let mut body_sender: Option<mpsc::Sender<Result<Bytes, unb_core::CoreError>>> = None;
2099 loop {
2100 let kind = tokio::select! {
2101 biased;
2102 () = cancellation.cancelled() => {
2103 if let Some(sender) = body_sender.take() {
2104 let _ = sender.send(Err(unb_core::CoreError::Malformed(
2105 "operation body cancelled".into(),
2106 ))).await;
2107 }
2108 return Ok(())
2109 },
2110 record = reader.read_record() => record,
2111 };
2112 let (kind, payload) = match kind {
2113 Ok(record) => record,
2114 Err(error) => {
2115 if let Some(sender) = body_sender.take() {
2116 let _ = sender
2117 .send(Err(unb_core::CoreError::Malformed(error.to_string())))
2118 .await;
2119 }
2120 return Err(error);
2121 }
2122 };
2123 match kind {
2124 unb_transport::operation::OperationRecordKind::CompleteMessage => {
2125 if body_sender.is_some() {
2126 return Err(unb_transport::operation::OperationCodecError::Protocol(
2127 "complete message arrived before body end".into(),
2128 ));
2129 }
2130 let envelope = decode_operation_message(payload, corr, direction, records)?;
2131 terminal |= matches!(
2132 envelope.kind,
2133 unb_core::Kind::Response | unb_core::Kind::Error
2134 );
2135 records += 1;
2136 let input = application_core_input(handle, session, envelope).map_err(|error| {
2137 unb_transport::operation::OperationCodecError::Protocol(error.to_string())
2138 })?;
2139 handle.submit(input).await.map_err(|error| {
2140 unb_transport::operation::OperationCodecError::Io(error.to_string())
2141 })?;
2142 }
2143 unb_transport::operation::OperationRecordKind::Head => {
2144 if body_sender.is_some() {
2145 return Err(unb_transport::operation::OperationCodecError::Protocol(
2146 "operation head arrived before body end".into(),
2147 ));
2148 }
2149 let mut envelope = decode_operation_message(payload, corr, direction, records)?;
2150 let (token, sender) = register_operation_body(handle, session)?;
2151 envelope.body_token = Some(token);
2152 terminal |= matches!(
2153 envelope.kind,
2154 unb_core::Kind::Response | unb_core::Kind::Error
2155 );
2156 records += 1;
2157 body_sender = Some(sender);
2158 let input = application_core_input(handle, session, envelope).map_err(|error| {
2159 unb_transport::operation::OperationCodecError::Protocol(error.to_string())
2160 })?;
2161 handle.submit(input).await.map_err(|error| {
2162 unb_transport::operation::OperationCodecError::Io(error.to_string())
2163 })?;
2164 }
2165 unb_transport::operation::OperationRecordKind::BodyChunk => {
2166 let Some(sender) = &body_sender else {
2167 return Err(unb_transport::operation::OperationCodecError::Protocol(
2168 "body chunk arrived without a head".into(),
2169 ));
2170 };
2171 sender.send(Ok(payload)).await.map_err(|_| {
2172 unb_transport::operation::OperationCodecError::Io(
2173 "operation body consumer closed".into(),
2174 )
2175 })?;
2176 }
2177 unb_transport::operation::OperationRecordKind::BodyEnd => {
2178 let Some(sender) = body_sender.take() else {
2179 return Err(unb_transport::operation::OperationCodecError::Protocol(
2180 "body end arrived without a head".into(),
2181 ));
2182 };
2183 drop(sender);
2184 }
2185 unb_transport::operation::OperationRecordKind::DirectionEnd => {
2186 if body_sender.is_some() {
2187 return Err(
2188 unb_transport::operation::OperationCodecError::Truncated(
2189 "direction ended before body end".into(),
2190 ),
2191 );
2192 }
2193 if records == 0 || matches!(direction, OperationDirection::Return) && !terminal {
2194 return Err(unb_transport::operation::OperationCodecError::Protocol(
2195 "operation direction ended before its required message".into(),
2196 ));
2197 }
2198 reader.expect_fin().await?;
2199 submit_operation_end(
2200 handle,
2201 session,
2202 corr,
2203 direction,
2204 unb_core::OperationStreamOutcome::Clean,
2205 )
2206 .await
2207 .map_err(|error| {
2208 unb_transport::operation::OperationCodecError::Io(error.to_string())
2209 })?;
2210 return Ok(());
2211 }
2212 }
2213 }
2214}
2215
2216fn register_operation_body(
2217 handle: &ProtocolCoreHandle,
2218 session: &SessionId,
2219) -> Result<
2220 (String, mpsc::Sender<Result<Bytes, unb_core::CoreError>>),
2221 unb_transport::operation::OperationCodecError,
2222> {
2223 let token = mint_body_token();
2224 let (sender, receiver) = mpsc::channel(8);
2225 let stream: BodyStream = Box::pin(futures_util::stream::unfold(
2226 receiver,
2227 |mut receiver| async move { receiver.recv().await.map(|chunk| (chunk, receiver)) },
2228 ));
2229 let mut registry = handle.token_bodies.lock().expect("body token registry");
2230 let owned = registry
2231 .keys()
2232 .filter(|(owner, _)| owner == session)
2233 .count();
2234 if owned >= SESSION_BODY_TOKENS_MAX {
2235 return Err(unb_transport::operation::OperationCodecError::Protocol(
2236 "session operation body limit reached".into(),
2237 ));
2238 }
2239 registry.insert((session.clone(), token.clone()), WireBody::Stream(stream));
2240 Ok((token, sender))
2241}
2242
2243fn operation_outcome(
2244 error: &unb_transport::operation::OperationCodecError,
2245) -> unb_core::OperationStreamOutcome {
2246 match error {
2247 unb_transport::operation::OperationCodecError::Protocol(message) => {
2248 unb_core::OperationStreamOutcome::Protocol(message.clone())
2249 }
2250 unb_transport::operation::OperationCodecError::Truncated(_) => {
2251 unb_core::OperationStreamOutcome::Truncated
2252 }
2253 unb_transport::operation::OperationCodecError::Oversized { size, max } => {
2254 unb_core::OperationStreamOutcome::Protocol(format!(
2255 "operation record of {size} bytes exceeds max {max}"
2256 ))
2257 }
2258 unb_transport::operation::OperationCodecError::Io(message) => {
2259 unb_core::OperationStreamOutcome::Transport(message.clone())
2260 }
2261 }
2262}
2263
2264fn operation_body_consumer_closed(
2265 error: &unb_transport::operation::OperationCodecError,
2266) -> bool {
2267 matches!(
2268 error,
2269 unb_transport::operation::OperationCodecError::Io(message)
2270 if message == "operation body consumer closed"
2271 )
2272}
2273
2274fn operation_peer_stopped_cancelled(
2275 error: &unb_transport::operation::OperationCodecError,
2276) -> bool {
2277 matches!(
2278 error,
2279 unb_transport::operation::OperationCodecError::Io(message)
2280 if message.contains(&unb_transport::operation::RESET_CANCELLED.to_string())
2281 )
2282}
2283
2284async fn submit_operation_end(
2285 handle: &ProtocolCoreHandle,
2286 session: &SessionId,
2287 corr: &str,
2288 direction: OperationDirection,
2289 outcome: unb_core::OperationStreamOutcome,
2290) -> Result<(), WsError> {
2291 handle
2292 .submit(CoreInput::OperationStreamEnded {
2293 session: session.clone(),
2294 corr: unb_core::CorrelationId::from(corr.to_owned()),
2295 direction: match direction {
2296 OperationDirection::Opening => unb_core::OperationStreamDirection::Opening,
2297 OperationDirection::Return => unb_core::OperationStreamDirection::Return,
2298 },
2299 outcome,
2300 })
2301 .await
2302}
2303
2304fn decode_operation_message(
2305 payload: Bytes,
2306 corr: &str,
2307 direction: OperationDirection,
2308 records: usize,
2309) -> Result<Envelope, unb_transport::operation::OperationCodecError> {
2310 let envelope = Envelope::decode(payload).map_err(|error| {
2311 unb_transport::operation::OperationCodecError::Protocol(error.to_string())
2312 })?;
2313 if envelope.corr.as_deref() != Some(corr) {
2314 return Err(unb_transport::operation::OperationCodecError::Protocol(
2315 "operation record changed unb-corr".into(),
2316 ));
2317 }
2318 match direction {
2319 OperationDirection::Opening if records != 0 || !envelope.kind.is_application_request() => {
2320 Err(unb_transport::operation::OperationCodecError::Protocol(
2321 "opening direction requires exactly one application opener".into(),
2322 ))
2323 }
2324 OperationDirection::Return if !envelope.kind.is_application_response() => {
2325 Err(unb_transport::operation::OperationCodecError::Protocol(
2326 "return direction requires application responses".into(),
2327 ))
2328 }
2329 _ => Ok(envelope),
2330 }
2331}
2332
2333fn spawn_body_pump(
2334 peer: std::sync::Arc<unb_transport::ConnectionStreams>,
2335 token: String,
2336 stream: BodyStream,
2337 budget: Arc<Semaphore>,
2338 cancellation: CancellationToken,
2339 handle: ProtocolCoreHandle,
2340 registration: Option<(SessionId, String)>,
2341) {
2342 n0_future::task::spawn(async move {
2343 let pump = async {
2344 let Ok(mut writer) = peer.open_body(&token).await else {
2345 return;
2346 };
2347 let mut metered = MeteredBody::new(stream, budget, cancellation.clone());
2348 loop {
2349 match metered.next_chunk().await {
2350 Ok(Some(chunk)) => {
2351 if writer.write_chunk(&chunk.bytes).await.is_err() {
2352 return;
2353 }
2354 }
2355 Ok(None) => {
2356 if cancellation.is_cancelled() {
2357 writer.abort().await;
2358 } else {
2359 let _ = writer.finish().await;
2360 }
2361 return;
2362 }
2363 Err(_) => {
2364 writer.abort().await;
2365 return;
2366 }
2367 }
2368 }
2369 };
2370 pump.await;
2371 if let Some((session, corr)) = registration {
2372 handle.release_body_pump(&session, &corr, &token);
2373 }
2374 });
2375}
2376
2377#[allow(clippy::too_many_arguments)]
2378fn spawn_body_collect(
2379 handle: ProtocolCoreHandle,
2380 session: SessionId,
2381 mut envelope: Envelope,
2382 stream: BodyStream,
2383 ceiling: usize,
2384 completion: Option<EffectId>,
2385 relay: bool,
2386 cancellation: CancellationToken,
2387) {
2388 n0_future::task::spawn(async move {
2389 let collected = tokio::select! {
2390 biased;
2391 () = cancellation.cancelled() => {
2392 if let Some(effect) = completion {
2393 let _ = handle
2394 .submit(send_completion(effect, relay, SendResult::Cancelled))
2395 .await;
2396 }
2397 return;
2398 }
2399 collected = WireBody::Stream(stream).collect_to(ceiling) => collected,
2400 };
2401 match collected {
2402 Ok(payload) => {
2403 envelope.payload = payload;
2404 let _ = handle
2405 .inputs
2406 .send(ActorInput::Write {
2407 session,
2408 envelope,
2409 completion,
2410 relay,
2411 })
2412 .await;
2413 }
2414 Err(error) => {
2415 let (code, message) = match error {
2416 unb_core::CoreError::BodyTooLarge(_) => (
2417 unb_core::ErrorCode::PayloadTooLarge,
2418 "streaming body exceeds this route's WS collect ceiling; raise the \
2419 ceiling or reach this node over WebTransport"
2420 .to_string(),
2421 ),
2422 other => (unb_core::ErrorCode::Protocol, other.to_string()),
2423 };
2424 if let Some(effect) = completion {
2425 let refused = SendResult::Refused { code, message };
2426 let _ = handle.submit(send_completion(effect, relay, refused)).await;
2427 } else if envelope.kind.is_application_request() {
2428 if let Some(corr) = envelope.corr.clone() {
2429 let outcome = if code == unb_core::ErrorCode::PayloadTooLarge {
2430 unb_core::OperationStreamOutcome::PayloadTooLarge(message)
2431 } else {
2432 unb_core::OperationStreamOutcome::Protocol(message)
2433 };
2434 let _ = handle
2435 .submit(CoreInput::OperationStreamEnded {
2436 session,
2437 corr: unb_core::CorrelationId::from(corr),
2438 direction: unb_core::OperationStreamDirection::Opening,
2439 outcome,
2440 })
2441 .await;
2442 }
2443 } else if envelope.kind == unb_core::Kind::Response {
2444 envelope.kind = unb_core::Kind::Error;
2445 envelope.payload = Envelope::encode_payload(&serde_json::json!({
2446 "code": code,
2447 "message": message,
2448 }));
2449 let _ = handle
2450 .inputs
2451 .send(ActorInput::Write {
2452 session,
2453 envelope,
2454 completion: None,
2455 relay: false,
2456 })
2457 .await;
2458 }
2459 }
2460 }
2461 });
2462}
2463
2464fn spawn_reader(
2465 session: SessionId,
2466 mut reader: PipeReader,
2467 handle: ProtocolCoreHandle,
2468 cancellation: CancellationToken,
2469 bodies: crate::transport::SessionStreams,
2470) {
2471 n0_future::task::spawn(async move {
2472 loop {
2473 let received = tokio::select! {
2474 biased;
2475 () = cancellation.cancelled() => {
2476 let _ = handle.submit(CoreInput::SessionClosed { session }).await;
2477 return;
2478 },
2479 received = reader.recv() => received,
2480 };
2481 let envelope = match received {
2482 Ok(Some((envelope, _body))) => {
2483 if let Some(token) = &envelope.body_token {
2484 if !(envelope.kind.is_application_request()
2485 || envelope.kind == unb_core::Kind::Response)
2486 {
2487 let _ = handle.submit(CoreInput::SessionFailed { session }).await;
2488 return;
2489 }
2490 {
2491 let Some(peer) = bodies.clone() else {
2492 let _ = handle.submit(CoreInput::SessionFailed { session }).await;
2493 return;
2494 };
2495 let lazy_token = token.clone();
2496 let accepted = futures_util::stream::once(async move {
2497 peer.accept_body(&lazy_token, RELAY_BACKPRESSURE_TIMEOUT)
2498 .await
2499 });
2500 let stream: BodyStream =
2501 Box::pin(futures_util::StreamExt::flat_map(accepted, |outcome| {
2502 match outcome {
2503 Ok(reader) => futures_util::StreamExt::boxed(
2504 futures_util::StreamExt::map(reader, |chunk| {
2505 chunk.map_err(|error| {
2506 unb_core::CoreError::Malformed(
2507 error.to_string(),
2508 )
2509 })
2510 }),
2511 ),
2512 Err(error) => futures_util::StreamExt::boxed(
2513 futures_util::stream::once(async move {
2514 Err(unb_core::CoreError::Malformed(
2515 error.to_string(),
2516 ))
2517 }),
2518 ),
2519 }
2520 }));
2521 let rejected = {
2522 let mut registry =
2523 handle.token_bodies.lock().expect("body token registry");
2524 let key = (session.clone(), token.clone());
2525 let owned = registry
2526 .keys()
2527 .filter(|(owner, _)| owner == &session)
2528 .count();
2529 if registry.contains_key(&key) || owned >= SESSION_BODY_TOKENS_MAX {
2530 true
2531 } else {
2532 registry.insert(key, WireBody::Stream(stream));
2533 false
2534 }
2535 };
2536 if rejected {
2537 let _ = handle.submit(CoreInput::SessionFailed { session }).await;
2538 return;
2539 }
2540 }
2541 }
2542 envelope
2543 }
2544 Ok(None) => {
2545 let _ = handle.submit(CoreInput::SessionClosed { session }).await;
2546 return;
2547 }
2548 Err(_) => {
2549 #[cfg(feature = "observability")]
2550 {
2551 metrics::counter!("unb_transport_failures").increment(1);
2552 tracing::warn!(session = %session, "transport read failed");
2553 }
2554 let _ = handle.submit(CoreInput::SessionFailed { session }).await;
2555 return;
2556 }
2557 };
2558 if envelope.kind == unb_core::Kind::Cancel {
2559 if let Some(corr) = &envelope.corr {
2560 handle.cancel_body_pump(&session, corr);
2561 }
2562 }
2563 let input = match application_core_input(&handle, &session, envelope) {
2564 Ok(input) => input,
2565 Err(_) => {
2566 let _ = handle.submit(CoreInput::SessionFailed { session }).await;
2567 return;
2568 }
2569 };
2570 if handle.submit(input).await.is_err() {
2571 return;
2572 }
2573 }
2574 });
2575}
2576
2577fn spawn_writer(
2578 session: SessionId,
2579 mut writer: PipeWriter,
2580 mut outbound: mpsc::Receiver<WriteRequest>,
2581 handle: ProtocolCoreHandle,
2582 cancellation: CancellationToken,
2583) {
2584 n0_future::task::spawn(async move {
2585 loop {
2586 let request = tokio::select! {
2587 biased;
2588 () = cancellation.cancelled() => {
2589 let farewell = async {
2590 while let Ok(request) = outbound.try_recv() {
2591 if writer.feed(request.envelope).await.is_err() {
2592 break;
2593 }
2594 }
2595 let _ = writer.flush().await;
2596 let _ = writer.close().await;
2597 };
2598 let _ = n0_future::time::timeout(TERMINAL_SEND_TIMEOUT, farewell).await;
2599 return;
2600 },
2601 request = outbound.recv() => request,
2602 };
2603 let Some(request) = request else {
2604 let _ = handle
2605 .submit(CoreInput::SessionClosed {
2606 session: session.clone(),
2607 })
2608 .await;
2609 return;
2610 };
2611 let mut batch = Vec::with_capacity(WRITE_BATCH);
2612 batch.push(request);
2613 while batch.len() < WRITE_BATCH {
2614 match outbound.try_recv() {
2615 Ok(request) => batch.push(request),
2616 Err(_) => break,
2617 }
2618 }
2619 let completions = batch
2620 .iter()
2621 .filter_map(|request| request.completion.map(|effect| (effect, request.relay)))
2622 .collect::<Vec<_>>();
2623 let result = async {
2624 for request in batch {
2625 writer.feed(request.envelope).await?;
2626 }
2627 writer.flush().await
2628 }
2629 .await;
2630 let send_result = result.as_ref().map_or_else(
2631 |error| SendResult::WriteFailed(error.to_string()),
2632 |_| SendResult::Written,
2633 );
2634 let completions = completions
2635 .into_iter()
2636 .map(|(effect, relay)| send_completion(effect, relay, send_result.clone()))
2637 .collect::<Vec<_>>();
2638 if !completions.is_empty()
2639 && handle
2640 .inputs
2641 .send(ActorInput::CoreBatch(completions))
2642 .await
2643 .is_err()
2644 {
2645 return;
2646 }
2647 if result.is_err() {
2648 let _ = handle
2649 .submit(CoreInput::SessionClosed {
2650 session: session.clone(),
2651 })
2652 .await;
2653 cancellation.cancel();
2654 let _ = writer.close().await;
2655 return;
2656 }
2657 }
2658 });
2659}
2660
2661#[allow(clippy::too_many_arguments)]
2662fn reserve_write(
2663 sender: mpsc::Sender<WriteRequest>,
2664 byte_budget: Arc<Semaphore>,
2665 envelope: Envelope,
2666 completion: Option<EffectId>,
2667 relay: bool,
2668 handle: ProtocolCoreHandle,
2669 cancellation: CancellationToken,
2670 reservation_pending: Option<Arc<AtomicBool>>,
2671 timeout: std::time::Duration,
2672) {
2673 n0_future::task::spawn(async move {
2674 let effect = completion;
2675 let acquire = async {
2676 let budget = byte_budget
2677 .acquire_many_owned(WriteRequest::weight(&envelope))
2678 .await
2679 .map_err(|_| ())?;
2680 let slot = sender.reserve_owned().await.map_err(|_| ())?;
2681 Ok::<_, ()>((budget, slot))
2682 };
2683 let reserved = tokio::select! {
2684 biased;
2685 () = cancellation.cancelled() => {
2686 if let Some(pending) = &reservation_pending {
2687 pending.store(false, Ordering::Release);
2688 }
2689 if let Some(effect) = effect {
2690 let _ = handle.submit(send_completion(effect, relay, SendResult::Cancelled)).await;
2691 }
2692 return;
2693 }
2694 reserved = n0_future::time::timeout(timeout, acquire) => reserved,
2695 };
2696 if let Some(pending) = &reservation_pending {
2697 pending.store(false, Ordering::Release);
2698 }
2699 match reserved {
2700 Ok(Ok((budget, slot))) => {
2701 if let Some(effect) = effect {
2702 if handle
2703 .submit(send_completion(effect, relay, SendResult::Reserved))
2704 .await
2705 .is_err()
2706 {
2707 return;
2708 }
2709 }
2710 slot.send(WriteRequest {
2711 envelope,
2712 completion,
2713 relay,
2714 _budget: budget,
2715 });
2716 }
2717 Ok(Err(())) => {
2718 if let Some(effect) = effect {
2719 let _ = handle
2720 .submit(send_completion(effect, relay, SendResult::Closed))
2721 .await;
2722 }
2723 }
2724 Err(_) => {
2725 if let Some(effect) = effect {
2726 let _ = handle
2727 .submit(send_completion(
2728 effect,
2729 relay,
2730 SendResult::ReservationTimedOut,
2731 ))
2732 .await;
2733 }
2734 }
2735 }
2736 });
2737}
2738
2739fn send_completion(effect: EffectId, relay: bool, result: SendResult) -> CoreInput {
2740 if relay {
2741 CoreInput::RelayForwardCompleted { effect, result }
2742 } else {
2743 CoreInput::SendCompleted { effect, result }
2744 }
2745}
2746
2747fn spawn_session_deadline(
2748 session: SessionId,
2749 mut deadlines: tokio::sync::watch::Receiver<Option<Instant>>,
2750 handle: ProtocolCoreHandle,
2751 cancellation: CancellationToken,
2752) {
2753 n0_future::task::spawn(async move {
2754 loop {
2755 let deadline = *deadlines.borrow_and_update();
2756 match deadline {
2757 Some(deadline) if clock_now() < deadline => {
2758 tokio::select! {
2759 biased;
2760 () = cancellation.cancelled() => return,
2761 () = sleep_to(deadline) => {}
2762 }
2763 }
2764 Some(_) => {
2765 if handle
2766 .submit(CoreInput::SessionTimeout {
2767 session: session.clone(),
2768 })
2769 .await
2770 .is_err()
2771 {
2772 return;
2773 }
2774 tokio::select! {
2775 biased;
2776 () = cancellation.cancelled() => return,
2777 changed = deadlines.changed() => {
2778 if changed.is_err() {
2779 return;
2780 }
2781 }
2782 }
2783 }
2784 None => {
2785 tokio::select! {
2786 biased;
2787 () = cancellation.cancelled() => return,
2788 changed = deadlines.changed() => {
2789 if changed.is_err() {
2790 return;
2791 }
2792 }
2793 }
2794 }
2795 }
2796 }
2797 });
2798}
2799
2800fn spawn_client_timeout(
2801 session: SessionId,
2802 operation: ClientOperationId,
2803 deadline: Instant,
2804 handle: ProtocolCoreHandle,
2805 cancellation: Option<CancellationToken>,
2806) {
2807 let Some(cancellation) = cancellation else {
2808 return;
2809 };
2810 n0_future::task::spawn(async move {
2811 tokio::select! {
2812 biased;
2813 () = cancellation.cancelled() => {}
2814 () = sleep_to(deadline) => {
2815 let _ = handle.submit(CoreInput::ClientOperationTimeout { session, operation }).await;
2816 }
2817 }
2818 });
2819}
2820
2821fn spawn_establishment_timeout(session: SessionId, handle: ProtocolCoreHandle) {
2822 n0_future::task::spawn(async move {
2823 n0_future::time::sleep(ESTABLISHMENT_TIMEOUT).await;
2824 let _ = handle
2825 .submit(CoreInput::EstablishmentTimeout { session })
2826 .await;
2827 });
2828}
2829
2830#[cfg(test)]
2831mod tests {
2832 use super::*;
2833 use unb_core::{DiscoverPlan, Kind, RetirementReason, PROTOCOL_VERSION};
2834 use bytes::Bytes;
2835 use std::sync::atomic::{AtomicBool, Ordering};
2836
2837 struct RecordingExecutor {
2838 effects: mpsc::Sender<CoreEffect>,
2839 }
2840
2841 impl EffectExecutor for RecordingExecutor {
2842 fn execute(&self, effect: CoreEffect, _handle: ProtocolCoreHandle) -> EffectFuture {
2843 let effects = self.effects.clone();
2844 Box::pin(async move {
2845 let _ = effects.send(effect).await;
2846 None
2847 })
2848 }
2849 }
2850
2851 struct PublicationExecutor {
2852 published: Arc<AtomicBool>,
2853 effects: mpsc::Sender<bool>,
2854 }
2855
2856 struct ClosingHandler {
2857 delivered: mpsc::Sender<Envelope>,
2858 closed: mpsc::Sender<ClientOperationId>,
2859 }
2860
2861 impl SessionHandler for ClosingHandler {
2862 async fn deliver(&mut self, envelope: Envelope) {
2863 let _ = self.delivered.send(envelope).await;
2864 }
2865
2866 async fn stream_closed(&mut self, operation: ClientOperationId) {
2867 let _ = self.closed.send(operation).await;
2868 }
2869 }
2870
2871 impl EffectExecutor for PublicationExecutor {
2872 fn execute(&self, effect: CoreEffect, _handle: ProtocolCoreHandle) -> EffectFuture {
2873 let published = self.published.clone();
2874 let effects = self.effects.clone();
2875 Box::pin(async move {
2876 if matches!(effect, CoreEffect::SessionRetired { .. }) {
2877 let _ = effects.send(published.load(Ordering::SeqCst)).await;
2878 }
2879 None
2880 })
2881 }
2882 }
2883
2884 fn local_pipe() -> (Pipe, mpsc::Sender<Envelope>, mpsc::Receiver<Envelope>) {
2885 let (incoming, rx) = mpsc::channel(16);
2886 let (tx, outgoing) = mpsc::channel(16);
2887 (
2888 Pipe::Local {
2889 rx,
2890 tx,
2891 initiator: true,
2892 },
2893 incoming,
2894 outgoing,
2895 )
2896 }
2897
2898 fn frame(kind: Kind, corr: Option<String>, payload: serde_json::Value) -> Envelope {
2899 Envelope {
2900 v: PROTOCOL_VERSION,
2901 id: format!("{kind:?}"),
2902 subject: String::new(),
2903 kind,
2904 corr,
2905 seq: None,
2906 hops: None,
2907 body_token: None,
2908 payload: Envelope::encode_payload(&payload),
2909 path: Vec::new(),
2910 headers: Default::default(),
2911 }
2912 }
2913
2914 async fn establish(incoming: &mpsc::Sender<Envelope>, outgoing: &mut mpsc::Receiver<Envelope>) {
2915 assert_eq!(outgoing.recv().await.unwrap().kind, Kind::Hello);
2916 incoming
2917 .send(frame(
2918 Kind::Welcome,
2919 None,
2920 serde_json::json!({ "version": 1 }),
2921 ))
2922 .await
2923 .unwrap();
2924 }
2925
2926 fn request(subject: &str, body: &'static [u8]) -> http::Request<Bytes> {
2927 http::Request::post(format!("/{subject}"))
2928 .body(Bytes::from_static(body))
2929 .unwrap()
2930 }
2931
2932 async fn attach_client(handle: &ProtocolCoreHandle, name: &str, pipe: Pipe) -> Arc<Wire> {
2933 attach_client_with_ceiling(
2934 handle,
2935 name,
2936 pipe,
2937 unb_transport::DEFAULT_MAX_FRAME_SIZE,
2938 )
2939 .await
2940 }
2941
2942 async fn attach_client_with_ceiling(
2943 handle: &ProtocolCoreHandle,
2944 name: &str,
2945 pipe: Pipe,
2946 ws_collect_ceiling: usize,
2947 ) -> Arc<Wire> {
2948 let session_cancellation = handle.cancellation.child_token();
2949 let (wire, directives, outcome, routes) =
2950 Wire::standalone_channel(session_cancellation.clone());
2951 let (reply, response) = oneshot::channel();
2952 handle
2953 .inputs
2954 .send(ActorInput::Attach {
2955 session: SessionId::from(name),
2956 transport: pipe,
2957 expected_peer: None,
2958 directives,
2959 handler: handler_channel(wire.observation_sender()),
2960 client: wire.client_session(),
2961 cancellation: session_cancellation,
2962 establish_peer: false,
2963 outcome,
2964 routes,
2965 ws_collect_ceiling,
2966 reply,
2967 })
2968 .await
2969 .unwrap();
2970 response.await.unwrap().unwrap();
2971 wire
2972 }
2973
2974 fn spawn_core() -> ProtocolCoreHandle {
2975 ProtocolCoreHandle::spawn(
2976 ProtocolCore::new("node"),
2977 Arc::new(RecordingExecutor {
2978 effects: mpsc::channel(1).0,
2979 }),
2980 CancellationToken::new(),
2981 &tokio::runtime::Handle::current(),
2982 )
2983 }
2984
2985 fn stalled_pipe() -> (Pipe, mpsc::Sender<Envelope>, mpsc::Receiver<Envelope>) {
2986 let (incoming, rx) = mpsc::channel(16);
2987 let (tx, outgoing) = mpsc::channel(1);
2988 (
2989 Pipe::Local {
2990 rx,
2991 tx,
2992 initiator: true,
2993 },
2994 incoming,
2995 outgoing,
2996 )
2997 }
2998
2999 fn chunk_stream(chunks: Vec<Result<Bytes, unb_core::CoreError>>) -> BodyStream {
3000 Box::pin(futures_util::stream::iter(chunks))
3001 }
3002
3003 #[tokio::test(start_paused = true)]
3004 async fn a_metered_body_stall_past_the_deadline_is_busy_shaped() {
3005 let budget = Arc::new(Semaphore::new(4));
3006 let hold = budget.clone().try_acquire_many_owned(4).unwrap();
3007 let mut body = MeteredBody::new(
3008 chunk_stream(vec![Ok(Bytes::from_static(b"x"))]),
3009 budget,
3010 CancellationToken::new(),
3011 );
3012 let outcome = body.next_chunk().await;
3013 assert!(
3014 matches!(
3015 outcome,
3016 Err(WsError::Core(unb_core::CoreError::BodyStalled))
3017 ),
3018 "budget starvation past the deadline must be BUSY-shaped: {outcome:?}"
3019 );
3020 drop(hold);
3021 }
3022
3023 #[tokio::test]
3024 async fn a_window_saturated_body_leaves_a_sibling_pump_undisturbed() {
3025 let budget = Arc::new(Semaphore::new(OUTBOUND_BYTE_BUDGET));
3026 let mut saturated = MeteredBody::new(
3027 chunk_stream(vec![
3028 Ok(Bytes::from(vec![0u8; BODY_WINDOW_BYTES])),
3029 Ok(Bytes::from_static(b"tail")),
3030 ]),
3031 budget.clone(),
3032 CancellationToken::new(),
3033 );
3034 let held = saturated.next_chunk().await.unwrap().unwrap();
3035 assert!(
3036 tokio::time::timeout(std::time::Duration::from_millis(50), saturated.next_chunk())
3037 .await
3038 .is_err(),
3039 "a full per-body window must block the next chunk"
3040 );
3041 let mut sibling = MeteredBody::new(
3042 chunk_stream(vec![Ok(Bytes::from_static(b"y"))]),
3043 budget,
3044 CancellationToken::new(),
3045 );
3046 let chunk = tokio::time::timeout(std::time::Duration::from_secs(1), sibling.next_chunk())
3047 .await
3048 .expect("a sibling body must progress while another window is saturated")
3049 .unwrap()
3050 .unwrap();
3051 assert_eq!(chunk.bytes, Bytes::from_static(b"y"));
3052 drop(held);
3053 let tail = tokio::time::timeout(std::time::Duration::from_secs(1), saturated.next_chunk())
3054 .await
3055 .expect("draining the window must unblock the body")
3056 .unwrap()
3057 .unwrap();
3058 assert_eq!(tail.bytes, Bytes::from_static(b"tail"));
3059 }
3060
3061 #[test]
3062 fn body_tokens_are_unique_and_carry_entropy() {
3063 let first = mint_body_token();
3064 let second = mint_body_token();
3065 assert_ne!(first, second);
3066 assert!(
3067 first.len() >= 33,
3068 "a guessable serial token invites body substitution: {first}"
3069 );
3070 }
3071
3072 #[tokio::test]
3073 async fn an_oversized_source_chunk_is_rejected_before_retention() {
3074 let budget = Arc::new(Semaphore::new(OUTBOUND_BYTE_BUDGET));
3075 let mut body = MeteredBody::new(
3076 chunk_stream(vec![Ok(Bytes::from(vec![
3077 0u8;
3078 unb_transport::DEFAULT_MAX_FRAME_SIZE
3079 + 1
3080 ]))]),
3081 budget.clone(),
3082 CancellationToken::new(),
3083 );
3084 let outcome = body.next_chunk().await;
3085 assert!(
3086 matches!(
3087 outcome,
3088 Err(WsError::Core(unb_core::CoreError::BodyTooLarge(_)))
3089 ),
3090 "a chunk beyond the frame cap must be refused before retention: {outcome:?}"
3091 );
3092 assert_eq!(
3093 budget.available_permits(),
3094 OUTBOUND_BYTE_BUDGET,
3095 "a refused chunk must not consume budget"
3096 );
3097 }
3098
3099 #[tokio::test]
3100 async fn concurrent_bodies_stay_within_the_session_budget() {
3101 let budget = Arc::new(Semaphore::new(100));
3102 let mut first = MeteredBody::new(
3103 chunk_stream(vec![Ok(Bytes::from(vec![0u8; 60]))]),
3104 budget.clone(),
3105 CancellationToken::new(),
3106 );
3107 let mut second = MeteredBody::new(
3108 chunk_stream(vec![Ok(Bytes::from(vec![0u8; 60]))]),
3109 budget.clone(),
3110 CancellationToken::new(),
3111 );
3112 let held = first.next_chunk().await.unwrap().unwrap();
3113 assert_eq!(budget.available_permits(), 40);
3114 assert!(
3115 tokio::time::timeout(std::time::Duration::from_millis(50), second.next_chunk())
3116 .await
3117 .is_err(),
3118 "combined queued bytes must not exceed the session budget"
3119 );
3120 drop(held);
3121 let chunk = tokio::time::timeout(std::time::Duration::from_secs(1), second.next_chunk())
3122 .await
3123 .expect("released budget must admit the waiting body")
3124 .unwrap()
3125 .unwrap();
3126 assert_eq!(chunk.bytes.len(), 60);
3127 }
3128
3129 #[tokio::test(start_paused = true)]
3130 async fn a_steady_body_streams_past_the_terminal_deadline_without_timing_out() {
3131 let budget = Arc::new(Semaphore::new(OUTBOUND_BYTE_BUDGET));
3132 let chunks: Vec<Result<Bytes, unb_core::CoreError>> =
3133 (0..8).map(|_| Ok(Bytes::from_static(b"chunk"))).collect();
3134 let mut body = MeteredBody::new(chunk_stream(chunks), budget, CancellationToken::new());
3135 for _ in 0..8 {
3136 let chunk = body
3137 .next_chunk()
3138 .await
3139 .expect("a steadily progressing body never times out")
3140 .expect("chunks remain");
3141 drop(chunk);
3142 tokio::time::advance(std::time::Duration::from_secs(6)).await;
3143 }
3144 assert!(body.next_chunk().await.unwrap().is_none());
3145 }
3146
3147 #[tokio::test]
3148 async fn pipe_recv_yields_whole_message_envelopes_with_no_body_stream() {
3149 let mut envelope = frame(
3150 Kind::Request,
3151 Some("c1".into()),
3152 serde_json::json!({"n": 1}),
3153 );
3154 envelope.subject = "echo".into();
3155 let encoded = envelope.encode();
3156 let pipe = Pipe::Piped {
3157 pipe: Box::pin(ScriptedPipe {
3158 frames: std::collections::VecDeque::from([Ok(encoded.clone())]),
3159 }),
3160 initiator: true,
3161 };
3162 let (mut reader, _writer) = pipe.split();
3163 let (received, body) = reader.recv().await.unwrap().unwrap();
3164 assert!(body.is_none(), "whole-message carriers collect the body");
3165 assert_eq!(received, Envelope::decode(encoded).unwrap());
3166
3167 let (incoming, rx) = mpsc::channel(4);
3168 let (tx, _outgoing) = mpsc::channel(4);
3169 let pipe = Pipe::Local {
3170 rx,
3171 tx,
3172 initiator: true,
3173 };
3174 let (mut reader, _writer) = pipe.split();
3175 incoming.send(envelope.clone()).await.unwrap();
3176 let (received, body) = reader.recv().await.unwrap().unwrap();
3177 assert!(body.is_none());
3178 assert_eq!(received, envelope);
3179 }
3180
3181 #[test]
3182 fn write_weight_counts_head_and_payload_and_clamps_to_the_whole_budget() {
3183 let mut envelope = frame(Kind::Request, None, serde_json::Value::Null);
3184 envelope.id = String::new();
3185 envelope.payload = Bytes::new();
3186 assert_eq!(WriteRequest::weight(&envelope), 1);
3187 envelope
3188 .headers
3189 .insert("origin".into(), serde_json::Value::String("x".repeat(64)));
3190 assert_eq!(WriteRequest::weight(&envelope), 70);
3191 envelope.payload = Bytes::from(vec![0u8; OUTBOUND_BYTE_BUDGET + 1]);
3192 assert_eq!(WriteRequest::weight(&envelope), OUTBOUND_BYTE_BUDGET as u32);
3193 }
3194
3195 #[tokio::test]
3196 async fn a_byte_saturated_outbound_ends_with_busy_and_recovers_on_drain() {
3197 let handle = spawn_core();
3198 let (pipe, incoming, mut outgoing) = stalled_pipe();
3199 let wire = attach_client(&handle, "stalled", pipe).await;
3200 establish(&incoming, &mut outgoing).await;
3201 let client = wire.client_session();
3202 let payload = Bytes::from(vec![0u8; 16 * 1024 * 1024]);
3203
3204 let mut flood = Vec::new();
3205 for _ in 0..6 {
3206 flood.push(
3207 client
3208 .start(
3209 "flood",
3210 Kind::Request,
3211 payload.clone(),
3212 None,
3213 Default::default(),
3214 )
3215 .await
3216 .unwrap(),
3217 );
3218 }
3219 let mut rejected = client
3220 .start(
3221 "flood",
3222 Kind::Request,
3223 payload.clone(),
3224 None,
3225 Default::default(),
3226 )
3227 .await
3228 .unwrap();
3229
3230 let refused = rejected.operation().as_str().to_owned();
3231 let refusal = tokio::time::timeout(std::time::Duration::from_secs(2), rejected.next())
3232 .await
3233 .expect("an over-budget write must be refused promptly, not silently queued");
3234 assert!(
3235 refusal.is_err(),
3236 "queued outbound bytes past the budget must refuse the write: {refusal:?}"
3237 );
3238
3239 let mut saw_busy = false;
3240 loop {
3241 match tokio::time::timeout(std::time::Duration::from_millis(300), outgoing.recv()).await
3242 {
3243 Ok(Some(envelope)) => {
3244 if envelope.kind == Kind::Error
3245 && envelope.corr.as_deref() == Some(refused.as_str())
3246 {
3247 assert_eq!(envelope.payload_json()["code"], "BUSY");
3248 saw_busy = true;
3249 }
3250 }
3251 Ok(None) => panic!("the stalled pipe closed while draining"),
3252 Err(_) => break,
3253 }
3254 }
3255 assert!(saw_busy, "the refused write must surface as a BUSY error");
3256
3257 let probe = client
3258 .start(
3259 "probe",
3260 Kind::Request,
3261 payload.clone(),
3262 None,
3263 Default::default(),
3264 )
3265 .await
3266 .unwrap();
3267 let opening = tokio::time::timeout(std::time::Duration::from_secs(5), async {
3268 loop {
3269 let envelope = outgoing.recv().await.unwrap();
3270 if envelope.subject == "probe" {
3271 break envelope;
3272 }
3273 }
3274 })
3275 .await
3276 .expect("draining the writer must release budget for later writes");
3277 incoming
3278 .send(frame(
3279 Kind::Response,
3280 opening.corr,
3281 serde_json::json!({ "ok": true }),
3282 ))
3283 .await
3284 .unwrap();
3285 let mut probe = probe;
3286 let response = tokio::time::timeout(std::time::Duration::from_secs(5), probe.next())
3287 .await
3288 .expect("a post-drain write must complete end to end")
3289 .unwrap()
3290 .unwrap();
3291 assert_eq!(response.payload_json(), serde_json::json!({ "ok": true }));
3292 handle.shutdown();
3293 }
3294
3295 #[tokio::test]
3296 async fn a_byte_saturated_session_leaves_a_sibling_session_undisturbed() {
3297 let handle = spawn_core();
3298 let (saturated_pipe, stalled_incoming, mut stalled_outgoing) = stalled_pipe();
3299 let stalled = attach_client(&handle, "stalled", saturated_pipe).await;
3300 establish(&stalled_incoming, &mut stalled_outgoing).await;
3301 let (live_pipe, live_incoming, mut live_outgoing) = local_pipe();
3302 let live = attach_client(&handle, "live", live_pipe).await;
3303 establish(&live_incoming, &mut live_outgoing).await;
3304
3305 let payload = Bytes::from(vec![0u8; 16 * 1024 * 1024]);
3306 let mut flood = Vec::new();
3307 for _ in 0..6 {
3308 flood.push(
3309 stalled
3310 .client_session()
3311 .start(
3312 "flood",
3313 Kind::Request,
3314 payload.clone(),
3315 None,
3316 Default::default(),
3317 )
3318 .await
3319 .unwrap(),
3320 );
3321 }
3322
3323 let mut sibling = live
3324 .client_session()
3325 .start(
3326 "sibling",
3327 Kind::Request,
3328 Bytes::new(),
3329 None,
3330 Default::default(),
3331 )
3332 .await
3333 .unwrap();
3334 let opening = tokio::time::timeout(std::time::Duration::from_secs(1), live_outgoing.recv())
3335 .await
3336 .expect("a sibling session must keep writing while another is byte-saturated")
3337 .unwrap();
3338 assert_eq!(opening.subject, "sibling");
3339 live_incoming
3340 .send(frame(
3341 Kind::Response,
3342 opening.corr,
3343 serde_json::json!({ "ok": true }),
3344 ))
3345 .await
3346 .unwrap();
3347 let response = tokio::time::timeout(std::time::Duration::from_secs(2), sibling.next())
3348 .await
3349 .expect("a sibling response must be delivered while another session is saturated")
3350 .unwrap()
3351 .unwrap();
3352 assert_eq!(response.payload_json(), serde_json::json!({ "ok": true }));
3353 handle.shutdown();
3354 }
3355
3356 #[tokio::test]
3357 async fn terminal_client_delivery_closes_handler_stream_once() {
3358 let cancellation = CancellationToken::new();
3359 let handle = ProtocolCoreHandle::spawn(
3360 ProtocolCore::new("node"),
3361 Arc::new(RecordingExecutor {
3362 effects: mpsc::channel(1).0,
3363 }),
3364 cancellation,
3365 &tokio::runtime::Handle::current(),
3366 );
3367 let (pipe, incoming, mut outgoing) = local_pipe();
3368 let (delivered, _deliveries) = mpsc::channel(4);
3369 let (closed, mut closures) = mpsc::channel(4);
3370 let session_cancellation = handle.cancellation.child_token();
3371 let (wire, directives, outcome, routes) =
3372 Wire::standalone_channel(session_cancellation.clone());
3373 let (reply, response) = oneshot::channel();
3374 handle
3375 .inputs
3376 .send(ActorInput::Attach {
3377 session: SessionId::from("session"),
3378 transport: pipe,
3379 expected_peer: None,
3380 directives,
3381 handler: handler_channel(ClosingHandler { delivered, closed }),
3382 client: wire.client_session(),
3383 cancellation: session_cancellation,
3384 establish_peer: false,
3385 outcome,
3386 routes,
3387 ws_collect_ceiling: unb_transport::DEFAULT_MAX_FRAME_SIZE,
3388 reply,
3389 })
3390 .await
3391 .unwrap();
3392 response.await.unwrap().unwrap();
3393 establish(&incoming, &mut outgoing).await;
3394 let mut stream = wire
3395 .client_session()
3396 .start(
3397 "echo",
3398 Kind::Request,
3399 Bytes::new(),
3400 None,
3401 Default::default(),
3402 )
3403 .await
3404 .unwrap();
3405 let opening = outgoing.recv().await.unwrap();
3406
3407 incoming
3408 .send(frame(
3409 Kind::Response,
3410 opening.corr,
3411 serde_json::json!({ "ok": true }),
3412 ))
3413 .await
3414 .unwrap();
3415
3416 assert!(stream.next().await.unwrap().is_some());
3417 assert_eq!(closures.recv().await, Some(stream.operation().clone()));
3418 assert!(
3419 tokio::time::timeout(std::time::Duration::from_millis(25), closures.recv())
3420 .await
3421 .is_err()
3422 );
3423 handle.shutdown();
3424 }
3425
3426 #[tokio::test]
3427 async fn generic_terminal_delivery_closes_handler_stream_once() {
3428 let cancellation = CancellationToken::new();
3429 let handle = ProtocolCoreHandle::spawn(
3430 ProtocolCore::new("node"),
3431 Arc::new(RecordingExecutor {
3432 effects: mpsc::channel(1).0,
3433 }),
3434 cancellation,
3435 &tokio::runtime::Handle::current(),
3436 );
3437 let (pipe, incoming, mut outgoing) = local_pipe();
3438 let (delivered, mut deliveries) = mpsc::channel(4);
3439 let (closed, mut closures) = mpsc::channel(4);
3440 let session_cancellation = handle.cancellation.child_token();
3441 let (wire, directives, outcome, routes) =
3442 Wire::standalone_channel(session_cancellation.clone());
3443 let (reply, response) = oneshot::channel();
3444 handle
3445 .inputs
3446 .send(ActorInput::Attach {
3447 session: SessionId::from("session"),
3448 transport: pipe,
3449 expected_peer: None,
3450 directives,
3451 handler: handler_channel(ClosingHandler { delivered, closed }),
3452 client: wire.client_session(),
3453 cancellation: session_cancellation,
3454 establish_peer: false,
3455 outcome,
3456 routes,
3457 ws_collect_ceiling: unb_transport::DEFAULT_MAX_FRAME_SIZE,
3458 reply,
3459 })
3460 .await
3461 .unwrap();
3462 response.await.unwrap().unwrap();
3463 establish(&incoming, &mut outgoing).await;
3464 let corr = wire
3465 .open_stream("echo", Kind::Request, serde_json::Value::Null)
3466 .await
3467 .unwrap();
3468 let opening = outgoing.recv().await.unwrap();
3469 assert_eq!(opening.corr.as_deref(), Some(corr.as_str()));
3470
3471 incoming
3472 .send(frame(
3473 Kind::Response,
3474 Some(corr.clone()),
3475 serde_json::json!({ "ok": true }),
3476 ))
3477 .await
3478 .unwrap();
3479
3480 assert_eq!(
3481 deliveries.recv().await.unwrap().corr.as_deref(),
3482 Some(corr.as_str())
3483 );
3484 assert_eq!(closures.recv().await, Some(ClientOperationId::from(corr)));
3485 assert!(
3486 tokio::time::timeout(std::time::Duration::from_millis(25), closures.recv())
3487 .await
3488 .is_err()
3489 );
3490 handle.shutdown();
3491 }
3492
3493 #[tokio::test]
3494 async fn immediate_loopback_response_finds_registered_client_operation() {
3495 let (pipe, incoming, mut outgoing) = local_pipe();
3496 let wire = Wire::open(pipe);
3497 establish(&incoming, &mut outgoing).await;
3498 let loopback = tokio::spawn(async move {
3499 let opening = outgoing.recv().await.unwrap();
3500 let corr = opening.corr.unwrap();
3501 incoming
3502 .send(frame(
3503 Kind::Response,
3504 Some(corr),
3505 serde_json::json!({ "ok": true }),
3506 ))
3507 .await
3508 .unwrap();
3509 });
3510
3511 let mut operation = wire
3512 .client_session()
3513 .start(
3514 "echo",
3515 Kind::Request,
3516 Bytes::new(),
3517 None,
3518 Default::default(),
3519 )
3520 .await
3521 .unwrap();
3522
3523 loopback.await.unwrap();
3524 assert_eq!(
3525 operation
3526 .try_next()
3527 .unwrap()
3528 .unwrap()
3529 .unwrap()
3530 .payload_json(),
3531 serde_json::json!({ "ok": true })
3532 );
3533 assert_eq!(operation.try_next().unwrap(), Some(None));
3534 }
3535
3536 #[tokio::test]
3537 async fn opening_send_failure_cleans_up_client_operation() {
3538 let (pipe, incoming, mut outgoing) = local_pipe();
3539 let wire = Wire::open(pipe);
3540 establish(&incoming, &mut outgoing).await;
3541 drop(outgoing);
3542
3543 let mut operation = wire
3544 .client_session()
3545 .start(
3546 "echo",
3547 Kind::Request,
3548 Bytes::new(),
3549 None,
3550 Default::default(),
3551 )
3552 .await
3553 .unwrap();
3554
3555 assert!(matches!(
3556 operation.next().await,
3557 Err(crate::ClientError::SessionClosed(_))
3558 ));
3559 }
3560
3561 #[tokio::test]
3562 async fn fetch_projects_response_and_ignores_late_terminal_frame() {
3563 let (pipe, incoming, mut outgoing) = local_pipe();
3564 let wire = Wire::open(pipe);
3565 establish(&incoming, &mut outgoing).await;
3566 let responder = tokio::spawn(async move {
3567 let opening = outgoing.recv().await.unwrap();
3568 let corr = opening.corr.unwrap();
3569 let mut response = frame(
3570 Kind::Response,
3571 Some(corr.clone()),
3572 serde_json::json!({ "ok": true }),
3573 );
3574 response.headers.insert("x-source".into(), "owner".into());
3575 incoming.send(response.clone()).await.unwrap();
3576 incoming.send(response).await.unwrap();
3577 });
3578
3579 let response = wire
3580 .client_session()
3581 .fetch(
3582 request("echo", b"request"),
3583 std::time::Duration::from_secs(1),
3584 )
3585 .await
3586 .unwrap();
3587
3588 assert_eq!(response.status(), http::StatusCode::OK);
3589 assert_eq!(response.headers()["x-source"], "owner");
3590 assert_eq!(response.body(), &Bytes::from_static(br#"{"ok":true}"#));
3591 responder.await.unwrap();
3592 }
3593
3594 #[tokio::test]
3595 async fn subscribe_and_discovery_are_isolated() {
3596 let (pipe, incoming, mut outgoing) = local_pipe();
3597 let wire = Wire::open(pipe);
3598 establish(&incoming, &mut outgoing).await;
3599 let client = wire.client_session();
3600 let mut subscription = client
3601 .subscribe(request("events", b""), None)
3602 .await
3603 .unwrap();
3604 let mut discovery = client
3605 .discover(DiscoverPlan::decode(&Bytes::new()).unwrap())
3606 .await
3607 .unwrap();
3608 let first = outgoing.recv().await.unwrap();
3609 let second = outgoing.recv().await.unwrap();
3610 let (subscribe_corr, discover_corr) = if first.kind == Kind::Subscribe {
3611 (first.corr.unwrap(), second.corr.unwrap())
3612 } else {
3613 (second.corr.unwrap(), first.corr.unwrap())
3614 };
3615
3616 incoming
3617 .send(frame(
3618 Kind::Event,
3619 Some(discover_corr.clone()),
3620 serde_json::json!({ "discover": 1 }),
3621 ))
3622 .await
3623 .unwrap();
3624 incoming
3625 .send(frame(
3626 Kind::Event,
3627 Some(subscribe_corr.clone()),
3628 serde_json::json!({ "event": 1 }),
3629 ))
3630 .await
3631 .unwrap();
3632 incoming
3633 .send(frame(
3634 Kind::Response,
3635 Some(discover_corr),
3636 serde_json::json!({ "done": true }),
3637 ))
3638 .await
3639 .unwrap();
3640
3641 assert_eq!(
3642 subscription.next().await.unwrap().unwrap().payload_json()["event"],
3643 1
3644 );
3645 assert_eq!(
3646 discovery.next().await.unwrap().unwrap().payload_json()["discover"],
3647 1
3648 );
3649 assert_eq!(
3650 discovery.next().await.unwrap().unwrap().payload_json()["done"],
3651 true
3652 );
3653 }
3654
3655 #[tokio::test]
3656 async fn dropping_a_saturated_operation_frees_its_session_for_idle_delivery() {
3657 let (pipe, incoming, mut outgoing) = local_pipe();
3658 let wire = Wire::open(pipe);
3659 establish(&incoming, &mut outgoing).await;
3660 let client = wire.client_session();
3661 let busy = client.subscribe(request("busy", b""), None).await.unwrap();
3662 let mut idle = client.subscribe(request("idle", b""), None).await.unwrap();
3663 let busy_open = outgoing.recv().await.unwrap();
3664 let idle_open = outgoing.recv().await.unwrap();
3665
3666 let burst = tokio::spawn(async move {
3667 for sequence in 0..257 {
3668 incoming
3669 .send(frame(
3670 Kind::Event,
3671 busy_open.corr.clone(),
3672 serde_json::json!({ "sequence": sequence }),
3673 ))
3674 .await
3675 .unwrap();
3676 }
3677 incoming
3678 .send(frame(
3679 Kind::Response,
3680 idle_open.corr,
3681 serde_json::json!({ "done": true }),
3682 ))
3683 .await
3684 .unwrap();
3685 incoming
3686 });
3687 tokio::time::sleep(std::time::Duration::from_millis(25)).await;
3688
3689 drop(busy);
3690
3691 let _incoming = tokio::time::timeout(std::time::Duration::from_secs(5), burst)
3692 .await
3693 .expect("dropping the saturated stream must unblock its session")
3694 .unwrap();
3695 assert_eq!(
3696 idle.next().await.unwrap().unwrap().payload_json()["done"],
3697 true
3698 );
3699 }
3700
3701 #[tokio::test]
3702 async fn out_of_order_unary_responses_reach_their_operations() {
3703 let (pipe, incoming, mut outgoing) = local_pipe();
3704 let wire = Wire::open(pipe);
3705 establish(&incoming, &mut outgoing).await;
3706 let client = wire.client_session();
3707 let mut first = client
3708 .start(
3709 "first",
3710 Kind::Request,
3711 Bytes::new(),
3712 None,
3713 Default::default(),
3714 )
3715 .await
3716 .unwrap();
3717 let mut second = client
3718 .start(
3719 "second",
3720 Kind::Request,
3721 Bytes::new(),
3722 None,
3723 Default::default(),
3724 )
3725 .await
3726 .unwrap();
3727 let first_open = outgoing.recv().await.unwrap();
3728 let second_open = outgoing.recv().await.unwrap();
3729
3730 incoming
3731 .send(frame(
3732 Kind::Response,
3733 second_open.corr,
3734 serde_json::json!({ "response": "second" }),
3735 ))
3736 .await
3737 .unwrap();
3738 incoming
3739 .send(frame(
3740 Kind::Response,
3741 first_open.corr,
3742 serde_json::json!({ "response": "first" }),
3743 ))
3744 .await
3745 .unwrap();
3746
3747 assert_eq!(
3748 first.next().await.unwrap().unwrap().payload_json()["response"],
3749 "first"
3750 );
3751 assert_eq!(
3752 second.next().await.unwrap().unwrap().payload_json()["response"],
3753 "second"
3754 );
3755 }
3756
3757 #[tokio::test]
3758 async fn mixed_unary_subscription_and_discovery_are_isolated() {
3759 let (pipe, incoming, mut outgoing) = local_pipe();
3760 let wire = Wire::open(pipe);
3761 establish(&incoming, &mut outgoing).await;
3762 let client = wire.client_session();
3763 let mut unary = client
3764 .start(
3765 "unary",
3766 Kind::Request,
3767 Bytes::new(),
3768 None,
3769 Default::default(),
3770 )
3771 .await
3772 .unwrap();
3773 let mut subscription = client
3774 .subscribe(request("events", b""), None)
3775 .await
3776 .unwrap();
3777 let mut discovery = client
3778 .discover(DiscoverPlan::decode(&Bytes::new()).unwrap())
3779 .await
3780 .unwrap();
3781 let mut unary_corr = None;
3782 let mut subscription_corr = None;
3783 let mut discovery_corr = None;
3784 for _ in 0..3 {
3785 let opening = outgoing.recv().await.unwrap();
3786 match opening.kind {
3787 Kind::Request => unary_corr = opening.corr,
3788 Kind::Subscribe => subscription_corr = opening.corr,
3789 Kind::Discover => discovery_corr = opening.corr,
3790 kind => panic!("unexpected opening kind {kind:?}"),
3791 }
3792 }
3793
3794 for (kind, corr, value) in [
3795 (Kind::Event, discovery_corr, "discover"),
3796 (Kind::Response, unary_corr, "unary"),
3797 (Kind::Event, subscription_corr, "subscription"),
3798 ] {
3799 incoming
3800 .send(frame(kind, corr, serde_json::json!({ "owner": value })))
3801 .await
3802 .unwrap();
3803 }
3804
3805 assert_eq!(
3806 unary.next().await.unwrap().unwrap().payload_json()["owner"],
3807 "unary"
3808 );
3809 assert_eq!(
3810 subscription.next().await.unwrap().unwrap().payload_json()["owner"],
3811 "subscription"
3812 );
3813 assert_eq!(
3814 discovery.next().await.unwrap().unwrap().payload_json()["owner"],
3815 "discover"
3816 );
3817 }
3818
3819 #[tokio::test]
3820 async fn cancelling_one_operation_does_not_complete_another() {
3821 let (pipe, incoming, mut outgoing) = local_pipe();
3822 let wire = Wire::open(pipe);
3823 establish(&incoming, &mut outgoing).await;
3824 let client = wire.client_session();
3825 let cancelled = client
3826 .subscribe(request("cancelled", b""), None)
3827 .await
3828 .unwrap();
3829 let mut live = client.subscribe(request("live", b""), None).await.unwrap();
3830 let cancelled_open = outgoing.recv().await.unwrap();
3831 let live_open = outgoing.recv().await.unwrap();
3832
3833 drop(cancelled);
3834 let cancel = outgoing.recv().await.unwrap();
3835 assert_eq!(cancel.kind, Kind::Cancel);
3836 assert_eq!(cancel.corr, cancelled_open.corr);
3837 incoming
3838 .send(frame(
3839 Kind::Event,
3840 live_open.corr,
3841 serde_json::json!({ "live": true }),
3842 ))
3843 .await
3844 .unwrap();
3845 assert_eq!(
3846 live.next().await.unwrap().unwrap().payload_json()["live"],
3847 true
3848 );
3849 }
3850
3851 #[tokio::test]
3852 async fn stream_drop_sends_exactly_one_cancel() {
3853 let (pipe, incoming, mut outgoing) = local_pipe();
3854 let wire = Wire::open(pipe);
3855 establish(&incoming, &mut outgoing).await;
3856 let stream = wire
3857 .client_session()
3858 .subscribe(request("events", b""), None)
3859 .await
3860 .unwrap();
3861 let opening = outgoing.recv().await.unwrap();
3862 let corr = opening.corr.unwrap();
3863
3864 drop(stream);
3865
3866 let cancel = tokio::time::timeout(std::time::Duration::from_secs(1), outgoing.recv())
3867 .await
3868 .unwrap()
3869 .unwrap();
3870 assert_eq!(cancel.kind, Kind::Cancel);
3871 assert_eq!(cancel.corr.as_deref(), Some(corr.as_str()));
3872 assert!(
3873 tokio::time::timeout(std::time::Duration::from_millis(25), outgoing.recv())
3874 .await
3875 .is_err()
3876 );
3877 }
3878
3879 #[tokio::test]
3880 async fn operation_timeout_is_typed_and_cancels_once() {
3881 let (pipe, incoming, mut outgoing) = local_pipe();
3882 let wire = Wire::open(pipe);
3883 establish(&incoming, &mut outgoing).await;
3884 let mut stream = wire
3885 .client_session()
3886 .subscribe(
3887 request("events", b""),
3888 Some(std::time::Duration::from_millis(10)),
3889 )
3890 .await
3891 .unwrap();
3892 let opening = outgoing.recv().await.unwrap();
3893
3894 assert!(matches!(
3895 stream.next().await,
3896 Err(crate::ClientError::Timeout(_))
3897 ));
3898 let cancel = outgoing.recv().await.unwrap();
3899 assert_eq!(cancel.kind, Kind::Cancel);
3900 assert_eq!(cancel.corr, opening.corr);
3901 drop(stream);
3902 assert!(
3903 tokio::time::timeout(std::time::Duration::from_millis(25), outgoing.recv())
3904 .await
3905 .is_err()
3906 );
3907 }
3908
3909 #[tokio::test]
3910 async fn session_close_completes_every_operation_once() {
3911 let (pipe, incoming, mut outgoing) = local_pipe();
3912 let wire = Wire::open(pipe);
3913 establish(&incoming, &mut outgoing).await;
3914 let client = wire.client_session();
3915 let mut first = client.subscribe(request("first", b""), None).await.unwrap();
3916 let mut second = client
3917 .subscribe(request("second", b""), None)
3918 .await
3919 .unwrap();
3920 outgoing.recv().await.unwrap();
3921 outgoing.recv().await.unwrap();
3922
3923 drop(incoming);
3924
3925 assert!(matches!(
3926 first.next().await,
3927 Err(crate::ClientError::SessionClosed(_))
3928 ));
3929 assert!(matches!(
3930 second.next().await,
3931 Err(crate::ClientError::SessionClosed(_))
3932 ));
3933 assert_eq!(first.next().await.unwrap(), None);
3934 assert_eq!(second.next().await.unwrap(), None);
3935 }
3936
3937 #[tokio::test]
3938 async fn a_session_the_core_rejects_closes_its_transport() {
3939 let (pipe, incoming, mut outgoing) = local_pipe();
3940 let _wire = Wire::open(pipe);
3941 assert_eq!(outgoing.recv().await.unwrap().kind, Kind::Hello);
3942 incoming
3943 .send(frame(
3944 Kind::Welcome,
3945 None,
3946 serde_json::json!({ "version": 999 }),
3947 ))
3948 .await
3949 .unwrap();
3950 let mut saw_error = false;
3951 let closed = loop {
3952 match tokio::time::timeout(std::time::Duration::from_millis(500), outgoing.recv()).await
3953 {
3954 Ok(Some(frame)) => {
3955 if frame.kind == Kind::Error {
3956 saw_error = true;
3957 }
3958 }
3959 Ok(None) => break true,
3960 Err(_) => break false,
3961 }
3962 };
3963 assert!(
3964 saw_error,
3965 "the peer must receive the close-reason error frame before the transport closes"
3966 );
3967 assert!(
3968 closed,
3969 "core deciding a session invalid must close the transport, not leave it open"
3970 );
3971 }
3972
3973 struct ScriptedPipe {
3974 frames: std::collections::VecDeque<Result<Bytes, unb_transport::TransportError>>,
3975 }
3976
3977 impl futures_util::Stream for ScriptedPipe {
3978 type Item = Result<Bytes, unb_transport::TransportError>;
3979
3980 fn poll_next(
3981 mut self: Pin<&mut Self>,
3982 _cx: &mut std::task::Context<'_>,
3983 ) -> std::task::Poll<Option<Self::Item>> {
3984 std::task::Poll::Ready(self.frames.pop_front())
3985 }
3986 }
3987
3988 impl futures_util::Sink<Bytes> for ScriptedPipe {
3989 type Error = unb_transport::TransportError;
3990
3991 fn poll_ready(
3992 self: Pin<&mut Self>,
3993 _cx: &mut std::task::Context<'_>,
3994 ) -> std::task::Poll<Result<(), Self::Error>> {
3995 std::task::Poll::Ready(Ok(()))
3996 }
3997
3998 fn start_send(self: Pin<&mut Self>, _item: Bytes) -> Result<(), Self::Error> {
3999 Ok(())
4000 }
4001
4002 fn poll_flush(
4003 self: Pin<&mut Self>,
4004 _cx: &mut std::task::Context<'_>,
4005 ) -> std::task::Poll<Result<(), Self::Error>> {
4006 std::task::Poll::Ready(Ok(()))
4007 }
4008
4009 fn poll_close(
4010 self: Pin<&mut Self>,
4011 _cx: &mut std::task::Context<'_>,
4012 ) -> std::task::Poll<Result<(), Self::Error>> {
4013 std::task::Poll::Ready(Ok(()))
4014 }
4015 }
4016
4017 #[tokio::test]
4018 async fn a_malformed_frame_retires_the_session_with_a_typed_reason() {
4019 let cancellation = CancellationToken::new();
4020 let (effect_tx, mut effects) = mpsc::channel(64);
4021 let handle = ProtocolCoreHandle::spawn(
4022 ProtocolCore::new("node"),
4023 Arc::new(RecordingExecutor { effects: effect_tx }),
4024 cancellation,
4025 &tokio::runtime::Handle::current(),
4026 );
4027 let pipe = Pipe::Piped {
4028 pipe: Box::pin(ScriptedPipe {
4029 frames: std::collections::VecDeque::from([Ok(Bytes::from_static(
4030 b"\xff\xff\xff\xffgarbage",
4031 ))]),
4032 }),
4033 initiator: true,
4034 };
4035 let (events, _events_rx) = mpsc::channel(16);
4036 handle
4037 .attach(SessionId::from("session"), pipe, None, events)
4038 .await
4039 .unwrap();
4040 let reason = tokio::time::timeout(std::time::Duration::from_secs(2), async {
4041 loop {
4042 if let CoreEffect::SessionRetired { reason, .. } = effects.recv().await.unwrap() {
4043 break reason;
4044 }
4045 }
4046 })
4047 .await
4048 .expect("a malformed frame must retire the session");
4049 assert_eq!(reason, RetirementReason::TransportFailed);
4050 handle.shutdown();
4051 }
4052
4053 #[tokio::test]
4054 async fn a_transport_read_error_retires_the_session_with_a_typed_reason() {
4055 let cancellation = CancellationToken::new();
4056 let (effect_tx, mut effects) = mpsc::channel(64);
4057 let handle = ProtocolCoreHandle::spawn(
4058 ProtocolCore::new("node"),
4059 Arc::new(RecordingExecutor { effects: effect_tx }),
4060 cancellation,
4061 &tokio::runtime::Handle::current(),
4062 );
4063 let pipe = Pipe::Piped {
4064 pipe: Box::pin(ScriptedPipe {
4065 frames: std::collections::VecDeque::from([Err(
4066 unb_transport::TransportError::Io("connection reset".into()),
4067 )]),
4068 }),
4069 initiator: true,
4070 };
4071 let (events, _events_rx) = mpsc::channel(16);
4072 handle
4073 .attach(SessionId::from("session"), pipe, None, events)
4074 .await
4075 .unwrap();
4076 let reason = tokio::time::timeout(std::time::Duration::from_secs(2), async {
4077 loop {
4078 if let CoreEffect::SessionRetired { reason, .. } = effects.recv().await.unwrap() {
4079 break reason;
4080 }
4081 }
4082 })
4083 .await
4084 .expect("a transport read error must retire the session");
4085 assert_eq!(reason, RetirementReason::TransportFailed);
4086 handle.shutdown();
4087 }
4088
4089 #[tokio::test]
4090 async fn routes_two_sessions_and_submits_transport_closure() {
4091 let cancellation = CancellationToken::new();
4092 let (effect_tx, mut effects) = mpsc::channel(16);
4093 let handle = ProtocolCoreHandle::spawn(
4094 ProtocolCore::new("node"),
4095 Arc::new(RecordingExecutor { effects: effect_tx }),
4096 cancellation,
4097 &tokio::runtime::Handle::current(),
4098 );
4099 let (first_pipe, first_incoming, mut first_outgoing) = local_pipe();
4100 let (second_pipe, _second_incoming, mut second_outgoing) = local_pipe();
4101 let (first_events, _first_events_rx) = mpsc::channel(16);
4102 let (second_events, _second_events_rx) = mpsc::channel(16);
4103
4104 let first = handle
4105 .attach(SessionId::from("first"), first_pipe, None, first_events)
4106 .await
4107 .unwrap();
4108 let second = handle
4109 .attach(SessionId::from("second"), second_pipe, None, second_events)
4110 .await
4111 .unwrap();
4112
4113 assert!(first_outgoing.recv().await.is_some());
4114 assert!(second_outgoing.recv().await.is_some());
4115 assert!(Arc::ptr_eq(&first, &first));
4116 assert!(!Arc::ptr_eq(&first, &second));
4117
4118 drop(first_incoming);
4119 let retired = loop {
4120 let effect = effects.recv().await.unwrap();
4121 if matches!(
4122 effect,
4123 CoreEffect::SessionRetired {
4124 ref session,
4125 reason: RetirementReason::SessionClosed,
4126 } if session == &SessionId::from("first")
4127 ) {
4128 break effect;
4129 }
4130 };
4131 assert!(matches!(
4132 retired,
4133 CoreEffect::SessionRetired {
4134 reason: RetirementReason::SessionClosed,
4135 ..
4136 }
4137 ));
4138
4139 second.shutdown();
4140 handle.shutdown();
4141 }
4142
4143 #[tokio::test]
4144 async fn install_publishes_after_acceptance_and_before_effect_execution() {
4145 let cancellation = CancellationToken::new();
4146 let published = Arc::new(AtomicBool::new(false));
4147 let (effect_tx, mut effects) = mpsc::channel(1);
4148 let handle = ProtocolCoreHandle::spawn(
4149 ProtocolCore::new("node"),
4150 Arc::new(PublicationExecutor {
4151 published: published.clone(),
4152 effects: effect_tx,
4153 }),
4154 cancellation,
4155 &tokio::runtime::Handle::current(),
4156 );
4157 let (pipe, _incoming, mut outgoing) = local_pipe();
4158 let (events, _events_rx) = mpsc::channel(16);
4159 handle
4160 .attach(SessionId::from("session"), pipe, None, events)
4161 .await
4162 .unwrap();
4163 assert!(outgoing.recv().await.is_some());
4164
4165 handle
4166 .install(
4167 CoreInput::SessionClosed {
4168 session: SessionId::from("session"),
4169 },
4170 {
4171 let published = published.clone();
4172 move || published.store(true, Ordering::SeqCst)
4173 },
4174 )
4175 .await
4176 .unwrap();
4177
4178 assert!(effects.recv().await.unwrap());
4179 handle.shutdown();
4180 }
4181
4182 #[tokio::test]
4183 async fn install_failure_does_not_publish() {
4184 let cancellation = CancellationToken::new();
4185 let (effect_tx, _effects) = mpsc::channel(1);
4186 let handle = ProtocolCoreHandle::spawn(
4187 ProtocolCore::new("node"),
4188 Arc::new(RecordingExecutor { effects: effect_tx }),
4189 cancellation,
4190 &tokio::runtime::Handle::current(),
4191 );
4192 let input = CoreInput::SessionOpened {
4193 session: SessionId::from("duplicate"),
4194 initiator: true,
4195 establish_peer: false,
4196 expected_peer: None,
4197 };
4198 handle.submit(input.clone()).await.unwrap();
4199 let published = Arc::new(AtomicBool::new(false));
4200
4201 let result = handle
4202 .install(input, {
4203 let published = published.clone();
4204 move || published.store(true, Ordering::SeqCst)
4205 })
4206 .await;
4207
4208 assert!(matches!(result, Err(WsError::Core(_))));
4209 assert!(!published.load(Ordering::SeqCst));
4210 handle.shutdown();
4211 }
4212
4213 #[test]
4214 fn operation_records_cannot_change_their_bound_correlation() {
4215 let response = frame(
4216 Kind::Response,
4217 Some("s2".into()),
4218 serde_json::json!({"ok": true}),
4219 );
4220 let error =
4221 decode_operation_message(response.encode(), "s1", OperationDirection::Return, 0)
4222 .unwrap_err();
4223 assert!(matches!(
4224 error,
4225 unb_transport::operation::OperationCodecError::Protocol(_)
4226 ));
4227 assert_eq!(
4228 error.reset_code(),
4229 unb_transport::operation::RESET_PROTOCOL
4230 );
4231 }
4232
4233 #[tokio::test]
4234 async fn streaming_operation_body_uses_head_chunks_and_no_wire_body_token() {
4235 let mut envelope = frame(Kind::Response, Some("s1".into()), serde_json::json!(null));
4236 envelope.body_token = Some("legacy-token-must-not-leak".into());
4237 let body: BodyStream = Box::pin(futures_util::stream::iter([
4238 Ok(Bytes::from_static(b"first")),
4239 Ok(Bytes::from_static(b"second")),
4240 ]));
4241 let (write, read) = tokio::io::duplex(4096);
4242 let sender = async move {
4243 let mut writer = unb_transport::operation::OperationRecordWriter::new(write);
4244 write_operation_envelope(
4245 &mut writer,
4246 envelope,
4247 Some(body),
4248 Arc::new(Semaphore::new(OUTBOUND_BYTE_BUDGET)),
4249 CancellationToken::new(),
4250 )
4251 .await
4252 .unwrap();
4253 writer.finish().await.unwrap();
4254 };
4255 let receiver = async move {
4256 let mut reader = unb_transport::operation::OperationRecordReader::new(read);
4257 let (kind, payload) = reader.read_record().await.unwrap();
4258 assert_eq!(
4259 kind,
4260 unb_transport::operation::OperationRecordKind::Head
4261 );
4262 let head = Envelope::decode(payload).unwrap();
4263 assert_eq!(head.corr.as_deref(), Some("s1"));
4264 assert!(head.body_token.is_none());
4265 assert!(head.payload.is_empty());
4266 let mut collected = Vec::new();
4267 loop {
4268 let (kind, payload) = reader.read_record().await.unwrap();
4269 match kind {
4270 unb_transport::operation::OperationRecordKind::BodyChunk => {
4271 collected.extend_from_slice(&payload)
4272 }
4273 unb_transport::operation::OperationRecordKind::BodyEnd => break,
4274 other => panic!("unexpected body record {other:?}"),
4275 }
4276 }
4277 assert_eq!(collected, b"firstsecond");
4278 let (kind, payload) = reader.read_record().await.unwrap();
4279 assert_eq!(
4280 kind,
4281 unb_transport::operation::OperationRecordKind::DirectionEnd
4282 );
4283 assert!(payload.is_empty());
4284 reader.expect_fin().await.unwrap();
4285 };
4286 tokio::join!(sender, receiver);
4287 }
4288
4289 #[tokio::test]
4290 async fn cancelling_a_client_operation_releases_a_late_response_body() {
4291 let handle = spawn_core();
4292 let (pipe, incoming, mut outgoing) = local_pipe();
4293 let wire = attach_client(&handle, "cancel-body", pipe).await;
4294 establish(&incoming, &mut outgoing).await;
4295
4296 let call = wire
4297 .client_session()
4298 .start(
4299 "echo",
4300 Kind::Request,
4301 Bytes::from_static(b"request"),
4302 None,
4303 Default::default(),
4304 )
4305 .await
4306 .unwrap();
4307 let corr = call.operation().as_str().to_owned();
4308 assert_eq!(
4309 outgoing.recv().await.unwrap().corr.as_deref(),
4310 Some(corr.as_str())
4311 );
4312 drop(call);
4313 let cancel = outgoing.recv().await.unwrap();
4314 assert_eq!(cancel.kind, Kind::Cancel);
4315 assert_eq!(cancel.corr.as_deref(), Some(corr.as_str()));
4316
4317 incoming
4318 .send(frame(
4319 Kind::Response,
4320 Some(corr),
4321 serde_json::json!({ "late": true }),
4322 ))
4323 .await
4324 .unwrap();
4325 tokio::time::timeout(std::time::Duration::from_secs(1), async {
4326 loop {
4327 let retained = handle
4328 .token_bodies
4329 .lock()
4330 .expect("body token registry")
4331 .keys()
4332 .any(|(owner, _)| owner.as_str() == "cancel-body");
4333 if !retained {
4334 break;
4335 }
4336 tokio::task::yield_now().await;
4337 }
4338 })
4339 .await
4340 .expect("late response body must be released");
4341 handle.shutdown();
4342 }
4343
4344 #[tokio::test]
4345 async fn session_retirement_releases_unclaimed_bodies() {
4346 let handle = spawn_core();
4347 let (pipe, incoming, mut outgoing) = local_pipe();
4348 let _wire = attach_client(&handle, "retired-body", pipe).await;
4349 establish(&incoming, &mut outgoing).await;
4350 handle
4351 .register_body(
4352 &SessionId::from("retired-body"),
4353 WireBody::Bytes(Bytes::from_static(b"retained")),
4354 )
4355 .unwrap();
4356 assert!(handle
4357 .token_bodies
4358 .lock()
4359 .expect("body token registry")
4360 .keys()
4361 .any(|(owner, _)| owner.as_str() == "retired-body"));
4362
4363 drop(incoming);
4364 tokio::time::timeout(std::time::Duration::from_secs(1), async {
4365 loop {
4366 let retained = handle
4367 .token_bodies
4368 .lock()
4369 .expect("body token registry")
4370 .keys()
4371 .any(|(owner, _)| owner.as_str() == "retired-body");
4372 if !retained {
4373 break;
4374 }
4375 tokio::task::yield_now().await;
4376 }
4377 })
4378 .await
4379 .expect("session retirement must release body capabilities");
4380 handle.shutdown();
4381 }
4382}