1use std::sync::{Arc, Weak};
2
3use bytes::Bytes;
4use futures_util::{FutureExt, StreamExt};
5use serde_json::Value;
6use unb_core::{
7 ApplicationFailure, ApplicationInvocation, ApplicationOrigin, ApplicationResponse,
8 ApplicationResult, CapacityResult, CoreEffect, CoreInput, EffectId, Envelope, ErrorCode,
9 PeerAdmission, RelayOpenResult, RetirementReason, SessionId,
10};
11use unb_runtime::{
12 EffectExecutor, EffectFuture, Pipe, ProtocolCoreHandle, SessionHandler, SessionOutcome, Wire,
13 WsError,
14};
15
16use crate::layer::{Origin, ServiceBody};
17use crate::node::{Node, PeerLink};
18use crate::peer::{PeerNext, PeerRequest, VerifiedPeer};
19
20const STREAM_BATCH: usize = 8;
21pub(crate) const ROUTE_SYNC_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(5);
22
23pub(crate) struct CandidateSession {
24 pub wire: Arc<Wire>,
25 pub cleaned: tokio::sync::oneshot::Receiver<()>,
26 identity: tokio::sync::watch::Receiver<Option<unb_core::NodeIdentity>>,
27}
28
29impl CandidateSession {
30 pub async fn outcome(&self, peer: &str) -> Result<CandidateOutcome, WsError> {
31 self.observed_outcome()
32 .await
33 .map_err(|failure| match failure {
34 CandidateFailure::Session(error) => error,
35 CandidateFailure::Retired { reason, .. } => retirement_error(peer, reason),
36 CandidateFailure::MissingIdentity => WsError::Connect(format!(
37 "connection to {peer:?} completed without an admitted identity"
38 )),
39 })
40 }
41
42 pub(crate) async fn observed_outcome(&self) -> Result<CandidateOutcome, CandidateFailure> {
43 match self.wire.session_outcome().await? {
44 SessionOutcome::Established => {
45 Ok(CandidateOutcome::Promoted(self.observed_identity().await?))
46 }
47 SessionOutcome::Retired(
48 RetirementReason::DuplicateSession | RetirementReason::DuplicateSessionReplaced,
49 ) => Ok(CandidateOutcome::Duplicate(self.observed_identity().await?)),
50 SessionOutcome::Retired(reason) => Err(CandidateFailure::Retired {
51 reason,
52 identity: self.identity.borrow().clone(),
53 }),
54 }
55 }
56
57 async fn observed_identity(&self) -> Result<unb_core::NodeIdentity, CandidateFailure> {
58 let mut identity = self.identity.clone();
59 loop {
60 if let Some(identity) = identity.borrow().clone() {
61 return Ok(identity);
62 }
63 if identity.changed().await.is_err() {
64 return Err(CandidateFailure::MissingIdentity);
65 }
66 }
67 }
68}
69
70pub(crate) enum CandidateFailure {
71 Session(WsError),
72 Retired {
73 reason: RetirementReason,
74 identity: Option<unb_core::NodeIdentity>,
75 },
76 MissingIdentity,
77}
78
79impl From<WsError> for CandidateFailure {
80 fn from(error: WsError) -> Self {
81 CandidateFailure::Session(error)
82 }
83}
84
85#[derive(Debug)]
86pub(crate) enum CandidateOutcome {
87 Promoted(unb_core::NodeIdentity),
88 Duplicate(unb_core::NodeIdentity),
89}
90
91pub(crate) struct ServerEffectExecutor {
92 pub(crate) node: Weak<Node>,
93}
94
95impl EffectExecutor for ServerEffectExecutor {
96 fn execute(&self, effect: CoreEffect, handle: ProtocolCoreHandle) -> EffectFuture {
97 let node = self.node.clone();
98 Box::pin(async move {
99 let node = node.upgrade()?;
100 match effect {
101 CoreEffect::RequestPeerAdmission {
102 effect,
103 session,
104 remote,
105 } => Some(CoreInput::PeerAdmissionCompleted {
106 effect,
107 result: node.admit_peer(session, remote).await,
108 }),
109 CoreEffect::CheckDispatchCapacity { effect, .. } => {
110 let result = match node.dispatch_slots.clone().try_acquire_owned() {
111 Ok(permit) => {
112 node.dispatch_permits.lock().await.insert(effect, permit);
113 CapacityResult::Available
114 }
115 Err(_) => {
116 #[cfg(feature = "observability")]
117 metrics::counter!("unb_dispatch_busy").increment(1);
118 CapacityResult::Busy
119 }
120 };
121 Some(CoreInput::CapacityChecked { effect, result })
122 }
123 CoreEffect::InvokeApplication { effect, invocation } => {
124 node.invoke_application(effect, invocation, handle).await
125 }
126 CoreEffect::OpenRelay {
127 effect,
128 source,
129 peer: peer_name,
130 frame,
131 } => {
132 let body = frame
133 .body
134 .as_ref()
135 .and_then(|body| handle.claim_body(&source.session, body.as_str()));
136 let (payload, body) = match body {
137 Some(unb_runtime::WireBody::Bytes(payload)) => (payload, None),
138 Some(unb_runtime::WireBody::Stream(body)) => (Bytes::new(), Some(body)),
139 None => (Bytes::new(), None),
140 };
141 let forwarded = frame.into_envelope();
142 let result = if let Some(link) = node.peer(&peer_name).await {
143 match link
144 .wire
145 .open_forward_with(
146 &forwarded.subject,
147 forwarded.kind,
148 payload,
149 forwarded.hops,
150 forwarded.headers,
151 body,
152 |_| async {},
153 )
154 .await
155 {
156 Ok(corr) => RelayOpenResult::Opened(unb_core::StreamKey {
157 session: link.session_id.into(),
158 corr: corr.into(),
159 }),
160 Err(error) => RelayOpenResult::Failed(ApplicationFailure {
161 code: ErrorCode::PeerUnreachable,
162 message: error.to_string(),
163 }),
164 }
165 } else {
166 RelayOpenResult::Failed(ApplicationFailure {
167 code: ErrorCode::PeerUnreachable,
168 message: format!("no live connection to peer {peer_name:?}"),
169 })
170 };
171 Some(CoreInput::RelayOpenCompleted { effect, result })
172 }
173 CoreEffect::QueryDiscoveryNeighbor { stream, peer, plan } => {
174 node.query_discovery_neighbor(stream, peer, plan, handle);
175 None
176 }
177 CoreEffect::SessionEstablished { session, peer } => {
178 node.session_established(session.clone(), peer).await;
179 None
180 }
181 CoreEffect::RouteSnapshotApplied {
182 session,
183 peer,
184 snapshot,
185 ..
186 } => {
187 node.replace_route_hints(&peer, &snapshot);
188 node.snapshot.rcu(|snapshot_state| {
189 let mut next = (**snapshot_state).clone();
190 let _ = next
191 .node_core
192 .apply_snapshot(session.as_str(), &peer, &snapshot);
193 next
194 });
195 None
196 }
197 CoreEffect::RouteDeltaApplied {
198 session,
199 peer,
200 delta,
201 ..
202 } => {
203 node.apply_route_hint_delta(&peer, &delta);
204 node.snapshot.rcu(|snapshot_state| {
205 let mut next = (**snapshot_state).clone();
206 let _ = next.node_core.apply_delta(session.as_str(), &peer, &delta);
207 next
208 });
209 None
210 }
211 CoreEffect::RouteSessionWithdrawn { session, .. } => {
212 node.snapshot.rcu(|snapshot_state| {
213 let mut next = (**snapshot_state).clone();
214 next.node_core.leave(session.as_str());
215 next
216 });
217 None
218 }
219 CoreEffect::AbortDispatch { effect, .. } => {
220 if let Some(abort) = node.dispatching.lock().await.remove(&effect) {
221 abort.cancel();
222 }
223 None
224 }
225 CoreEffect::SessionRetired { session, reason } => {
226 node.session_retired(session, reason).await;
227 None
228 }
229 _ => None,
230 }
231 })
232 }
233}
234
235struct SessionBridge {
236 node: Weak<Node>,
237 session: SessionId,
238}
239
240impl SessionHandler for SessionBridge {
241 async fn deliver(&mut self, _envelope: Envelope) {}
242
243 async fn stream_closed(&mut self, operation: unb_core::ClientOperationId) {
244 let Some(node) = self.node.upgrade() else {
245 return;
246 };
247 if let Some(cancel) = node
248 .active
249 .lock()
250 .await
251 .remove(&(self.session.clone(), operation.as_str().to_owned()))
252 {
253 cancel.cancel();
254 };
255 }
256}
257
258impl Node {
259 async fn admit_peer(
260 &self,
261 session: SessionId,
262 remote: unb_core::NodeIdentity,
263 ) -> PeerAdmission {
264 let request = PeerRequest::new(self.identity.clone(), remote.clone());
265 match PeerNext::root(self.peer_layers.clone())
266 .admit(request)
267 .await
268 {
269 Ok(admitted) => match admitted.verified() {
270 Some(verified) => {
271 self.verified_peers
272 .lock()
273 .await
274 .insert(session.clone(), verified);
275 if let Some(observation) =
276 self.candidate_identities.lock().await.remove(&session)
277 {
278 observation.send_replace(Some(remote.clone()));
279 }
280 PeerAdmission::Admitted(remote)
281 }
282 None => PeerAdmission::Rejected("peer admission produced no VerifiedPeer".into()),
283 },
284 Err(error) => PeerAdmission::Rejected(error.message),
285 }
286 }
287
288 async fn invoke_application(
289 self: &Arc<Self>,
290 effect: EffectId,
291 invocation: ApplicationInvocation,
292 handle: ProtocolCoreHandle,
293 ) -> Option<CoreInput> {
294 let abort = self.cancellation.child_token();
295 self.dispatching.lock().await.insert(effect, abort.clone());
296 let permit = self
297 .dispatch_permits
298 .lock()
299 .await
300 .remove(&invocation.reservation);
301 let Some(_permit) = permit else {
302 self.dispatching.lock().await.remove(&effect);
303 return Some(dispatch_failure(
304 effect,
305 ErrorCode::Busy,
306 "capacity reservation expired",
307 ));
308 };
309 if abort.is_cancelled() {
310 self.dispatching.lock().await.remove(&effect);
311 return Some(dispatch_failure(
312 effect,
313 ErrorCode::Cancelled,
314 "request cancelled",
315 ));
316 }
317 let origin = match invocation.origin {
318 ApplicationOrigin::Client { session } => Origin::Client {
319 session: session.to_string(),
320 },
321 ApplicationOrigin::Peer { session, peer } => Origin::Peer {
322 peer: self
323 .verified_peers
324 .lock()
325 .await
326 .get(&session)
327 .cloned()
328 .unwrap_or_else(|| VerifiedPeer::from_identity(&peer)),
329 session: session.to_string(),
330 },
331 };
332 let mut envelope = invocation.frame.clone().into_envelope();
333 let streaming_body = if let Some(body) = &invocation.frame.body {
334 let Some(body) = handle.claim_body(&invocation.stream.session, body.as_str()) else {
335 self.dispatching.lock().await.remove(&effect);
336 return Some(dispatch_failure(
337 effect,
338 ErrorCode::Protocol,
339 "application body unavailable",
340 ));
341 };
342 match body {
343 unb_runtime::WireBody::Bytes(payload) => {
344 envelope.payload = payload;
345 None
346 }
347 unb_runtime::WireBody::Stream(stream) => Some(stream),
348 }
349 } else {
350 None
351 };
352 let snapshot = self.snapshot.load_full();
353 let mut request = match Self::inbound_request(&envelope) {
354 Ok(request) => request,
355 Err(error) => {
356 self.dispatching.lock().await.remove(&effect);
357 return Some(dispatch_failure(effect, error.code, error.message));
358 }
359 };
360 if let Some(stream) = streaming_body {
361 request
362 .extensions_mut()
363 .insert(crate::service::StreamingBody(std::sync::Arc::new(
364 std::sync::Mutex::new(Some(stream)),
365 )));
366 }
367 let outcome = tokio::select! {
368 biased;
369 () = abort.cancelled() => {
370 self.dispatching.lock().await.remove(&effect);
371 return Some(dispatch_failure(effect, ErrorCode::Cancelled, "request cancelled"));
372 }
373 outcome = self.run_service(snapshot.clone(), request, origin) => outcome,
374 };
375 self.dispatching.lock().await.remove(&effect);
376 let outcome = match outcome {
377 Some(Ok(outcome)) => outcome,
378 Some(Err(error)) => return Some(dispatch_failure(effect, error.code, error.message)),
379 None => {
380 let error = Self::teach_unknown_subject(&snapshot, &invocation.frame.head.subject);
381 return Some(dispatch_failure(effect, error.code, error.message));
382 }
383 };
384 let (parts, body) = outcome.into_parts();
385 match body {
386 ServiceBody::Unary(payload) => Some(CoreInput::DispatchCompleted {
387 effect,
388 result: unary_result(parts, payload, &handle, &invocation.stream.session),
389 }),
390 ServiceBody::Stream(mut stream) => {
391 let key = (
392 invocation.stream.session.clone(),
393 invocation.stream.corr.as_str().to_string(),
394 );
395 let cancel = self.cancellation.child_token();
396 self.active.lock().await.insert(key.clone(), cancel.clone());
397 let active = self.active.clone();
398 let response_handle = handle.clone();
399 let response_session = invocation.stream.session.clone();
400 unb_runtime::RuntimeHandle::current().spawn(async move {
401 'pump: loop {
402 let item = tokio::select! {
403 biased;
404 () = cancel.cancelled() => break,
405 item = stream.next() => item,
406 };
407 let mut result =
408 stream_result(item, &parts, &response_handle, &response_session);
409 let mut batch = Vec::new();
410 let terminal = loop {
411 let terminal = !matches!(result, Ok(ApplicationResult::Event(_)));
412 batch.push(CoreInput::DispatchCompleted { effect, result });
413 if terminal || batch.len() >= STREAM_BATCH {
414 break terminal;
415 }
416 match stream.next().now_or_never() {
417 Some(item) => {
418 result = stream_result(
419 item,
420 &parts,
421 &response_handle,
422 &response_session,
423 )
424 }
425 None => break false,
426 }
427 };
428 if handle.submit_batch(batch).await.is_err() || terminal {
429 break 'pump;
430 }
431 }
432 active.lock().await.remove(&key);
433 });
434 None
435 }
436 }
437 }
438
439 async fn session_established(
440 self: &Arc<Self>,
441 session: SessionId,
442 peer: unb_core::NodeIdentity,
443 ) {
444 let wire = loop {
445 if let Some(wire) = self.session(session.as_str()).await {
446 break wire;
447 }
448 if self.cancellation.is_cancelled() {
449 return;
450 }
451 tokio::task::yield_now().await;
452 };
453 let outbound = self.outbound_sessions.lock().await.contains(&session);
454 self.session_peers
455 .write()
456 .await
457 .insert(session.to_string(), peer.node_id.clone());
458 let replaced = self.peers.write().await.insert(
459 peer.node_id.clone(),
460 PeerLink {
461 session_id: session.to_string(),
462 wire: wire.clone(),
463 instance_id: peer.instance_id.clone(),
464 outbound,
465 },
466 );
467 if let Some(old) = replaced {
468 if old.session_id != session.as_str() {
469 old.wire.shutdown();
470 }
471 }
472 if !outbound {
473 if let Some(connection) = self.connection(&peer.node_id) {
474 connection.bind(peer, session.to_string(), wire);
475 }
476 }
477 }
478
479 async fn session_retired(&self, session: SessionId, reason: RetirementReason) {
480 self.verified_peers.lock().await.remove(&session);
481 self.candidate_identities.lock().await.remove(&session);
482 self.outbound_sessions.lock().await.remove(&session);
483 let peer = self.session_peers.write().await.remove(session.as_str());
484 if let Some(connection) = peer.and_then(|peer| self.connection(&peer)) {
485 connection.retire(session.as_str(), reason);
486 }
487 if let Some(wire) = self.sessions.write().await.remove(session.as_str()) {
488 wire.shutdown();
489 }
490 self.cleanup_session(session.as_str()).await;
491 }
492
493 #[cfg(feature = "hosting")]
494 pub fn serve_ws_upgrade(
495 self: &Arc<Self>,
496 upgrade: axum::extract::ws::WebSocketUpgrade,
497 ) -> axum::response::Response {
498 let node = self.clone();
499 upgrade
500 .max_message_size(unb_transport::DEFAULT_MAX_FRAME_SIZE)
501 .max_frame_size(unb_transport::DEFAULT_MAX_FRAME_SIZE)
502 .on_upgrade(move |socket| async move {
503 let (pipe, initiator) = unb_transport::ws::accept(socket);
504 let _ = node.attach(Pipe::Piped { pipe, initiator }, None).await;
505 })
506 }
507
508 #[cfg(feature = "hosting")]
509 pub async fn serve_webtransport(
510 self: &Arc<Self>,
511 connection: unb_transport::webtransport::wtransport::Connection,
512 ) -> Result<Arc<Wire>, WsError> {
513 let (pipe, initiator, bodies) = unb_transport::webtransport::accept(connection).await?;
514 Ok(self
515 .attach(Pipe::piped_with_streams(pipe, initiator, bodies), None)
516 .await
517 .0)
518 }
519
520 pub async fn serve_transport(self: &Arc<Self>, transport: Pipe) -> Arc<Wire> {
521 self.attach(transport, None).await.0
522 }
523
524 pub async fn connect_transport(
525 self: &Arc<Self>,
526 peer: &str,
527 transport: Pipe,
528 ) -> Result<Arc<Wire>, WsError> {
529 let candidate = self.establish(transport, Some(peer.to_string())).await;
530 match candidate.outcome(peer).await {
531 Ok(CandidateOutcome::Promoted(_)) => {
532 let _ = n0_future::time::timeout(ROUTE_SYNC_TIMEOUT, candidate.wire.routes_acked())
533 .await;
534 Ok(candidate.wire)
535 }
536 Ok(CandidateOutcome::Duplicate(_)) => Ok(candidate.wire),
537 Err(error) => {
538 candidate.wire.shutdown();
539 Err(error)
540 }
541 }
542 }
543
544 pub(crate) async fn establish(
545 self: &Arc<Self>,
546 transport: Pipe,
547 expected_peer: Option<String>,
548 ) -> CandidateSession {
549 let (wire, _session, cleaned, identity) = self.attach_inner(transport, expected_peer).await;
550 CandidateSession {
551 wire,
552 cleaned,
553 identity,
554 }
555 }
556
557 pub(crate) async fn attach(
558 self: &Arc<Self>,
559 transport: Pipe,
560 expected_peer: Option<String>,
561 ) -> (Arc<Wire>, SessionId, tokio::sync::oneshot::Receiver<()>) {
562 let session = self.next_session_id();
563 let (wire, cleaned) = self
564 .attach_session(session.clone(), transport, expected_peer)
565 .await;
566 (wire, session, cleaned)
567 }
568
569 async fn attach_inner(
570 self: &Arc<Self>,
571 transport: Pipe,
572 expected_peer: Option<String>,
573 ) -> (
574 Arc<Wire>,
575 SessionId,
576 tokio::sync::oneshot::Receiver<()>,
577 tokio::sync::watch::Receiver<Option<unb_core::NodeIdentity>>,
578 ) {
579 let session = self.next_session_id();
580 self.outbound_sessions.lock().await.insert(session.clone());
581 let (identity_tx, identity) = tokio::sync::watch::channel(None);
582 self.candidate_identities
583 .lock()
584 .await
585 .insert(session.clone(), identity_tx);
586 let (wire, cleaned) = self
587 .attach_session(session.clone(), transport, expected_peer)
588 .await;
589 (wire, session, cleaned, identity)
590 }
591
592 fn next_session_id(&self) -> SessionId {
593 SessionId::from(format!(
594 "sess-{}",
595 self.next_session
596 .fetch_add(1, std::sync::atomic::Ordering::Relaxed)
597 + 1
598 ))
599 }
600
601 async fn attach_session(
602 self: &Arc<Self>,
603 session: SessionId,
604 transport: Pipe,
605 expected_peer: Option<String>,
606 ) -> (Arc<Wire>, tokio::sync::oneshot::Receiver<()>) {
607 let bridge = SessionBridge {
608 node: Arc::downgrade(self),
609 session: session.clone(),
610 };
611 let wire = self
612 .protocol
613 .attach_with_ceiling(
614 session.clone(),
615 transport,
616 expected_peer,
617 bridge,
618 self.ws_collect_ceiling,
619 )
620 .await
621 .expect("protocol core actor unavailable");
622 self.sessions
623 .write()
624 .await
625 .insert(session.to_string(), wire.clone());
626 let (cleaned_tx, cleaned_rx) = tokio::sync::oneshot::channel();
627 let cancellation = self.cancellation.child_token();
628 let closed = wire.clone();
629 unb_runtime::RuntimeHandle::current().spawn(async move {
630 tokio::select! {
631 biased;
632 () = cancellation.cancelled() => {}
633 () = closed.closed() => {}
634 }
635 let _ = cleaned_tx.send(());
636 });
637 (wire, cleaned_rx)
638 }
639
640 async fn cleanup_session(&self, session: &str) {
641 self.peers
642 .write()
643 .await
644 .retain(|_, link| link.session_id != session);
645 self.active.lock().await.retain(|(owner, _), cancel| {
646 let keep = owner.as_str() != session;
647 if !keep {
648 cancel.cancel();
649 }
650 keep
651 });
652 }
653}
654
655fn unary_result(
656 parts: http::response::Parts,
657 payload: Bytes,
658 handle: &ProtocolCoreHandle,
659 session: &SessionId,
660) -> Result<ApplicationResult, ApplicationFailure> {
661 if parts.status.is_client_error() || parts.status.is_server_error() {
662 let value: Value = serde_json::from_slice(&payload).unwrap_or(Value::Null);
663 let code = value
664 .get("code")
665 .and_then(|code| serde_json::from_value(code.clone()).ok())
666 .unwrap_or_else(|| ErrorCode::from_status(parts.status));
667 let message = value
668 .get("message")
669 .and_then(Value::as_str)
670 .map(str::to_string)
671 .unwrap_or_else(|| String::from_utf8_lossy(&payload).into_owned());
672 return Err(ApplicationFailure { code, message });
673 }
674 Ok(ApplicationResult::Response(application_response(
675 parts.status,
676 &parts.headers,
677 register_response_body(handle, session, payload)?,
678 )))
679}
680
681fn stream_result(
682 item: Option<Result<Bytes, crate::handler::HandlerError>>,
683 parts: &http::response::Parts,
684 handle: &ProtocolCoreHandle,
685 session: &SessionId,
686) -> Result<ApplicationResult, ApplicationFailure> {
687 match item {
688 Some(Ok(payload)) => Ok(ApplicationResult::Event(application_response(
689 parts.status,
690 &parts.headers,
691 register_response_body(handle, session, payload)?,
692 ))),
693 Some(Err(error)) => Err(ApplicationFailure {
694 code: error.code,
695 message: error.message,
696 }),
697 None => Ok(ApplicationResult::Finished(application_response(
698 parts.status,
699 &parts.headers,
700 None,
701 ))),
702 }
703}
704
705fn register_response_body(
706 handle: &ProtocolCoreHandle,
707 session: &SessionId,
708 payload: Bytes,
709) -> Result<Option<unb_core::BodyId>, ApplicationFailure> {
710 if payload.is_empty() {
711 return Ok(None);
712 }
713 handle
714 .register_body(session, unb_runtime::WireBody::Bytes(payload))
715 .map(Some)
716 .map_err(|error| ApplicationFailure {
717 code: ErrorCode::Busy,
718 message: error.to_string(),
719 })
720}
721
722fn application_response(
723 status: http::StatusCode,
724 headers: &http::HeaderMap,
725 body: Option<unb_core::BodyId>,
726) -> ApplicationResponse {
727 let mut head = http::Response::new(());
728 *head.status_mut() = status;
729 *head.headers_mut() = headers.clone();
730 ApplicationResponse { head, body }
731}
732
733fn dispatch_failure(effect: EffectId, code: ErrorCode, message: impl Into<String>) -> CoreInput {
734 CoreInput::DispatchCompleted {
735 effect,
736 result: Err(ApplicationFailure {
737 code,
738 message: message.into(),
739 }),
740 }
741}
742
743pub(crate) fn retirement_error(peer: &str, reason: RetirementReason) -> WsError {
744 WsError::Connect(format!(
745 "connection to {peer:?} retired during establishment: {reason:?}"
746 ))
747}