1use std::collections::{BTreeMap, HashMap, HashSet};
2use std::sync::atomic::AtomicU64;
3use std::sync::{Arc, RwLock as StdRwLock};
4#[cfg(feature = "hosting")]
5use std::time::Duration;
6
7use arc_swap::ArcSwap;
8use serde_json::Value;
9use tokio::sync::{Mutex, RwLock, Semaphore};
10use unb_core::{
11 CoreCapabilitySnapshot, CoreInput, EffectId, NodeCore, NodeIdentity, RouteAdvertisement,
12 RouteDelta, RouteSnapshot, SessionId,
13};
14use unb_runtime::{CancellationToken, DropGuard, ProtocolCoreHandle, Wire, WsError};
15
16use crate::layer::{ErasedCall, Layer};
17use crate::peer::{PeerLayer, VerifiedPeer};
18use crate::service::{Handler, HandlerService, Operation, StateMap, States};
19use crate::PeerConnection;
20
21#[cfg(feature = "hosting")]
22pub(crate) const WEBTRANSPORT_ACCEPT_TIMEOUT: Duration = Duration::from_secs(5);
23
24pub(crate) struct PeerLink {
25 pub session_id: String,
26 pub wire: Arc<Wire>,
27 pub instance_id: String,
28 pub outbound: bool,
29}
30
31#[derive(Clone, Eq, Ord, PartialEq, PartialOrd)]
32pub(crate) struct RouteHint {
33 subject: String,
34 owner: String,
35 owner_instance: String,
36 owner_epoch: u64,
37}
38
39impl From<&RouteAdvertisement> for RouteHint {
40 fn from(route: &RouteAdvertisement) -> Self {
41 RouteHint {
42 subject: route.subject.clone(),
43 owner: route.owner.clone(),
44 owner_instance: route.owner_instance.clone(),
45 owner_epoch: route.owner_epoch,
46 }
47 }
48}
49
50impl Clone for PeerLink {
51 fn clone(&self) -> PeerLink {
52 PeerLink {
53 session_id: self.session_id.clone(),
54 wire: self.wire.clone(),
55 instance_id: self.instance_id.clone(),
56 outbound: self.outbound,
57 }
58 }
59}
60
61#[derive(Clone)]
62pub(crate) struct CompiledOperation {
63 pub(crate) call: ErasedCall,
64 pub(crate) layers: Arc<[Arc<dyn Layer>]>,
65 pub(crate) contract: Value,
66}
67
68#[derive(Clone, Default)]
69pub(crate) struct SubjectServices {
70 pub(crate) unary: Option<CompiledOperation>,
71 pub(crate) streaming: Option<CompiledOperation>,
72 pub(crate) metadata: Option<Value>,
73 pub(crate) one_line: Option<String>,
74}
75
76impl SubjectServices {
77 pub(crate) fn register(
78 services: &mut BTreeMap<String, Arc<SubjectServices>>,
79 service: HandlerService,
80 scopes: &[String],
81 layers: Vec<Arc<dyn Layer>>,
82 states: &StateMap,
83 ) -> Result<(String, Value), String> {
84 let subject = service.effective_subject(scopes)?;
85 let mut entry = services
86 .get(&subject)
87 .map(|existing| existing.as_ref().clone())
88 .unwrap_or_default();
89 let slot = match service.operation {
90 Operation::Unary => &mut entry.unary,
91 Operation::Streaming => &mut entry.streaming,
92 };
93 if slot.is_some() {
94 return Err(format!(
95 "subject {subject:?} already serves a {:?} operation",
96 service.operation
97 ));
98 }
99 if service.metadata.is_some() && entry.metadata.is_some() {
100 return Err(format!(
101 "subject {subject:?} already carries metadata; attach it to one registration"
102 ));
103 }
104 let call = (service.build)(&States(states))?;
105 *slot = Some(CompiledOperation {
106 call,
107 layers: Arc::from(layers),
108 contract: service.contract.to_json(service.operation),
109 });
110 if service.metadata.is_some() {
111 entry.metadata = service.metadata;
112 }
113 if entry.one_line.is_none() {
114 entry.one_line = service.one_line;
115 }
116 let catalog_entry = entry.catalog_entry();
117 services.insert(subject.clone(), Arc::new(entry));
118 Ok((subject, catalog_entry))
119 }
120
121 pub(crate) fn catalog_entry(&self) -> Value {
122 let mut operations = serde_json::Map::new();
123 if let Some(unary) = &self.unary {
124 operations.insert("unary".into(), unary.contract.clone());
125 }
126 if let Some(streaming) = &self.streaming {
127 operations.insert("streaming".into(), streaming.contract.clone());
128 }
129 let mut entry = serde_json::Map::new();
130 if let Some(metadata) = &self.metadata {
131 if let Some(one_line) = metadata.get("one_line") {
132 entry.insert("one_line".into(), one_line.clone());
133 }
134 entry.insert("metadata".into(), metadata.clone());
135 }
136 if !entry.contains_key("one_line") {
137 if let Some(one_line) = &self.one_line {
138 entry.insert("one_line".into(), Value::String(one_line.clone()));
139 }
140 }
141 entry.insert("operations".into(), Value::Object(operations));
142 Value::Object(entry)
143 }
144
145 fn same_contract(&self, other: &Self) -> bool {
146 fn same_operation(
147 left: &Option<CompiledOperation>,
148 right: &Option<CompiledOperation>,
149 ) -> bool {
150 match (left, right) {
151 (Some(left), Some(right)) => {
152 Arc::ptr_eq(&left.call, &right.call)
153 && left.contract == right.contract
154 && left.layers.len() == right.layers.len()
155 && left
156 .layers
157 .iter()
158 .zip(right.layers.iter())
159 .all(|(left, right)| Arc::ptr_eq(left, right))
160 }
161 (None, None) => true,
162 _ => false,
163 }
164 }
165
166 same_operation(&self.unary, &other.unary)
167 && same_operation(&self.streaming, &other.streaming)
168 && self.metadata == other.metadata
169 && self.one_line == other.one_line
170 }
171}
172
173#[derive(Clone)]
174pub(crate) struct NodeSnapshot {
175 pub(crate) services: BTreeMap<String, Arc<SubjectServices>>,
176 pub(crate) capabilities: CoreCapabilitySnapshot,
177 pub(crate) node_core: NodeCore,
178}
179
180impl NodeSnapshot {
181 pub(crate) fn new(
182 services: BTreeMap<String, Arc<SubjectServices>>,
183 mut node_core: NodeCore,
184 ) -> Self {
185 let capabilities = CoreCapabilitySnapshot::new(
186 services
187 .iter()
188 .map(|(subject, services)| (subject.clone(), services.catalog_entry()))
189 .collect(),
190 );
191 node_core.install_local_capabilities(capabilities.entries().clone());
192 Self {
193 services,
194 capabilities,
195 node_core,
196 }
197 }
198
199 fn same_service_contract(&self, services: &BTreeMap<String, Arc<SubjectServices>>) -> bool {
200 self.services.len() == services.len()
201 && self.services.iter().all(|(subject, current)| {
202 services
203 .get(subject)
204 .is_some_and(|result| current.same_contract(result))
205 })
206 }
207}
208
209pub struct Node {
210 pub(crate) snapshot: Arc<ArcSwap<NodeSnapshot>>,
211 pub(crate) states: StateMap,
212 pub(crate) global_layers: Arc<[Arc<dyn Layer>]>,
213 pub(crate) mutation_gate: Mutex<()>,
214 pub(crate) peers: RwLock<HashMap<String, PeerLink>>,
215 pub(crate) sessions: RwLock<HashMap<String, Arc<Wire>>>,
216 pub(crate) connections: StdRwLock<HashMap<String, PeerConnection>>,
217 pub(crate) route_hints: StdRwLock<HashMap<String, std::collections::BTreeSet<RouteHint>>>,
218 pub(crate) session_peers: RwLock<HashMap<String, String>>,
219 pub(crate) outbound_sessions: Mutex<HashSet<SessionId>>,
220 pub(crate) dispatch_slots: Arc<Semaphore>,
221 pub(crate) dispatch_permits: Mutex<HashMap<EffectId, tokio::sync::OwnedSemaphorePermit>>,
222 pub(crate) dispatching: Mutex<HashMap<EffectId, CancellationToken>>,
223 pub(crate) verified_peers: Mutex<HashMap<SessionId, VerifiedPeer>>,
224 pub(crate) candidate_identities:
225 Mutex<HashMap<SessionId, tokio::sync::watch::Sender<Option<NodeIdentity>>>>,
226 pub(crate) active: Arc<Mutex<HashMap<(SessionId, String), CancellationToken>>>,
227 pub(crate) protocol: ProtocolCoreHandle,
228 pub(crate) identity: NodeIdentity,
229 pub(crate) peer_layers: Arc<[Arc<dyn PeerLayer>]>,
230 pub(crate) dial_policy: unb_client::Peers,
231 pub(crate) next_session: AtomicU64,
232 pub(crate) ws_collect_ceiling: usize,
233 pub(crate) cancellation: CancellationToken,
234 pub(crate) _shutdown: DropGuard,
235}
236
237impl Node {
238 pub fn cancellation(&self) -> &CancellationToken {
239 &self.cancellation
240 }
241
242 pub fn identity(&self) -> &NodeIdentity {
243 &self.identity
244 }
245
246 pub fn shutdown(&self) {
247 for connection in self
248 .connections
249 .read()
250 .unwrap_or_else(|poisoned| poisoned.into_inner())
251 .values()
252 {
253 connection.node_shutdown();
254 }
255 self.cancellation.cancel();
256 }
257
258 pub fn reachable_names(&self) -> Vec<String> {
259 self.snapshot.load().node_core.reachable_names()
260 }
261
262 pub fn catalog_revision(&self) -> u64 {
263 self.snapshot.load().node_core.catalog_revision()
264 }
265
266 pub fn local_catalog(&self, detail_full: bool) -> Value {
267 self.snapshot.load().node_core.catalog(detail_full)
268 }
269
270 pub async fn remove_subject(&self, subject: &str) -> Result<(), WsError> {
271 let _gate = self.mutation_gate.lock().await;
272 let mut services = self.snapshot.load().services.clone();
273 services.remove(subject);
274 self.install_services(services).await
275 }
276
277 pub async fn add_service(&self, handler: impl Handler) -> Result<(), WsError> {
278 let _gate = self.mutation_gate.lock().await;
279 let mut services = self.snapshot.load().services.clone();
280 SubjectServices::register(
281 &mut services,
282 handler.into_service(),
283 &[],
284 self.global_layers.iter().cloned().collect(),
285 &self.states,
286 )
287 .map_err(WsError::Connect)?;
288 self.install_services(services).await
289 }
290
291 pub async fn remove_operation(
292 &self,
293 subject: &str,
294 operation: Operation,
295 ) -> Result<(), WsError> {
296 let _gate = self.mutation_gate.lock().await;
297 let mut services = self.snapshot.load().services.clone();
298 let Some(existing) = services.get(subject).cloned() else {
299 return Ok(());
300 };
301 let mut entry = existing.as_ref().clone();
302 let slot = match operation {
303 Operation::Unary => &mut entry.unary,
304 Operation::Streaming => &mut entry.streaming,
305 };
306 if slot.take().is_none() {
307 return Ok(());
308 }
309 let empty = entry.unary.is_none() && entry.streaming.is_none();
310 if empty {
311 services.remove(subject);
312 } else {
313 services.insert(subject.to_string(), Arc::new(entry));
314 }
315 self.install_services(services).await
316 }
317
318 async fn install_services(
319 &self,
320 services: BTreeMap<String, Arc<SubjectServices>>,
321 ) -> Result<(), WsError> {
322 let current = self.snapshot.load();
323 if current.same_service_contract(&services) {
324 return Ok(());
325 }
326 let snapshot = NodeSnapshot::new(services, current.node_core.clone());
327 let capabilities = snapshot.capabilities.clone();
328 let publication = self.snapshot.clone();
329 self.protocol
330 .install(
331 CoreInput::LocalCapabilitiesInstalled {
332 capabilities: capabilities.entries().clone(),
333 },
334 move || publication.store(Arc::new(snapshot)),
335 )
336 .await
337 }
338
339 pub(crate) async fn peer(&self, name: &str) -> Option<PeerLink> {
340 self.peers.read().await.get(name).cloned()
341 }
342
343 pub(crate) async fn session(&self, id: &str) -> Option<Arc<Wire>> {
344 self.sessions.read().await.get(id).cloned()
345 }
346
347 pub(crate) fn connection(&self, peer: &str) -> Option<PeerConnection> {
348 self.connections
349 .read()
350 .unwrap_or_else(|poisoned| poisoned.into_inner())
351 .get(peer)
352 .cloned()
353 }
354
355 pub(crate) fn replace_route_hints(&self, peer: &str, snapshot: &RouteSnapshot) {
356 self.route_hints
357 .write()
358 .unwrap_or_else(|poisoned| poisoned.into_inner())
359 .insert(
360 peer.to_string(),
361 snapshot.routes.iter().map(RouteHint::from).collect(),
362 );
363 }
364
365 pub(crate) fn apply_route_hint_delta(&self, peer: &str, delta: &RouteDelta) {
366 let mut hints = self
367 .route_hints
368 .write()
369 .unwrap_or_else(|poisoned| poisoned.into_inner());
370 let peer_hints = hints.entry(peer.to_string()).or_default();
371 for withdrawal in &delta.withdraw {
372 peer_hints.retain(|hint| {
373 hint.subject != withdrawal.subject
374 || hint.owner != withdrawal.owner
375 || hint.owner_instance != withdrawal.owner_instance
376 || hint.owner_epoch != withdrawal.owner_epoch
377 });
378 }
379 peer_hints.extend(delta.upsert.iter().map(RouteHint::from));
380 }
381
382 pub(crate) fn reconnects_for_subject(
383 &self,
384 subject: &str,
385 ) -> Vec<crate::connection::ReconnectWait> {
386 let hints = self
387 .route_hints
388 .read()
389 .unwrap_or_else(|poisoned| poisoned.into_inner());
390 self.connections
391 .read()
392 .unwrap_or_else(|poisoned| poisoned.into_inner())
393 .iter()
394 .filter_map(|(peer, connection)| {
395 if connection.status() == crate::ConnectionStatus::Connecting
396 && hints
397 .get(peer)
398 .is_some_and(|routes| routes.iter().any(|route| route.subject == subject))
399 {
400 connection.active_reconnect()
401 } else {
402 None
403 }
404 })
405 .collect()
406 }
407}
408
409impl Drop for Node {
410 fn drop(&mut self) {
411 for connection in self
412 .connections
413 .read()
414 .unwrap_or_else(|poisoned| poisoned.into_inner())
415 .values()
416 {
417 connection.node_shutdown();
418 }
419 }
420}
421
422#[cfg(test)]
423mod tests {
424 use super::*;
425 use crate::service::OperationContract;
426
427 fn operation(contract: Value) -> CompiledOperation {
428 CompiledOperation {
429 call: Arc::new(|_| Box::pin(async { unreachable!() })),
430 layers: Arc::from([]),
431 contract,
432 }
433 }
434
435 fn service(subject: &str, operation: Operation, result: &'static str) -> HandlerService {
436 HandlerService::declare(
437 subject,
438 None,
439 Some(result),
440 operation,
441 OperationContract::unknown(),
442 move |_| {
443 Ok(Arc::new(move |_| {
444 Box::pin(async move {
445 Ok(http::Response::new(crate::layer::ServiceBody::Unary(
446 unb_core::Envelope::encode_payload(&Value::String(result.into())),
447 )))
448 })
449 }))
450 },
451 )
452 }
453
454 #[test]
455 fn node_snapshot_capabilities_match_service_subjects_and_operations() {
456 let mut services = BTreeMap::new();
457 services.insert(
458 "chess.move".into(),
459 Arc::new(SubjectServices {
460 unary: Some(operation(serde_json::json!({ "input": "Move" }))),
461 streaming: Some(operation(serde_json::json!({ "event": "Position" }))),
462 metadata: Some(serde_json::json!({ "one_line": "Play a move", "tier": 1 })),
463 one_line: None,
464 }),
465 );
466 services.insert(
467 "chess.state".into(),
468 Arc::new(SubjectServices {
469 streaming: Some(operation(serde_json::json!({ "event": "Position" }))),
470 one_line: Some("Watch the board".into()),
471 ..SubjectServices::default()
472 }),
473 );
474
475 let snapshot = NodeSnapshot::new(services, NodeCore::new("snapshot-test"));
476
477 assert_eq!(
478 snapshot.capabilities.entries().keys().collect::<Vec<_>>(),
479 snapshot.services.keys().collect::<Vec<_>>()
480 );
481 assert_eq!(
482 snapshot.capabilities.entries()["chess.move"],
483 serde_json::json!({
484 "one_line": "Play a move",
485 "metadata": { "one_line": "Play a move", "tier": 1 },
486 "operations": {
487 "unary": { "input": "Move" },
488 "streaming": { "event": "Position" }
489 }
490 })
491 );
492 assert_eq!(
493 snapshot.capabilities.entries()["chess.state"],
494 serde_json::json!({
495 "one_line": "Watch the board",
496 "operations": { "streaming": { "event": "Position" } }
497 })
498 );
499 }
500
501 #[tokio::test]
502 async fn mutation_derives_operations_metadata_fingerprint_and_revision_once() {
503 let node = Node::builder("snapshot-test")
504 .insecure_accept_declared_peer_identities()
505 .build()
506 .unwrap();
507 let empty_fingerprint = node.snapshot.load().node_core.fingerprint();
508
509 node.add_service(
510 service("chess", Operation::Unary, "move")
511 .describe(serde_json::json!({ "one_line": "Play chess", "tier": 1 })),
512 )
513 .await
514 .unwrap();
515 assert_eq!(node.catalog_revision(), 1);
516 let export = node.snapshot.load().node_core.export_for("peer");
517 assert_eq!(export.len(), 1);
518 assert_eq!(export[0].subject, "chess");
519 assert_eq!(export[0].owner_revision, 1);
520 assert_ne!(
521 node.snapshot.load().node_core.fingerprint(),
522 empty_fingerprint
523 );
524 assert_eq!(
525 node.local_catalog(true)["subjects"][0],
526 serde_json::json!({
527 "subject": "chess",
528 "one_line": "Play chess",
529 "metadata": { "one_line": "Play chess", "tier": 1 },
530 "operations": {
531 "unary": {
532 "input_schema": { "unknown": true },
533 "output_schema": { "unknown": true }
534 }
535 }
536 })
537 );
538
539 node.add_service(service("chess", Operation::Streaming, "watch"))
540 .await
541 .unwrap();
542 assert_eq!(node.catalog_revision(), 2);
543 assert_eq!(
544 node.snapshot.load().node_core.export_for("peer")[0].owner_revision,
545 2
546 );
547 assert_eq!(
548 node.snapshot.load().node_core.resolve("chess"),
549 unb_core::Resolution::Local
550 );
551 assert!(node.local_catalog(true)["subjects"][0]["operations"]["streaming"].is_object());
552
553 node.remove_operation("chess", Operation::Unary)
554 .await
555 .unwrap();
556 assert_eq!(node.catalog_revision(), 3);
557 assert_eq!(
558 node.snapshot.load().node_core.export_for("peer")[0].owner_revision,
559 3
560 );
561 assert_eq!(
562 node.snapshot.load().node_core.resolve("chess"),
563 unb_core::Resolution::Local
564 );
565 assert!(node.local_catalog(true)["subjects"][0]["operations"]["unary"].is_null());
566
567 node.remove_operation("chess", Operation::Streaming)
568 .await
569 .unwrap();
570 assert_eq!(node.catalog_revision(), 4);
571 assert!(node.snapshot.load().node_core.export_for("peer").is_empty());
572 assert_eq!(
573 node.snapshot.load().node_core.fingerprint(),
574 empty_fingerprint
575 );
576 assert!(node.reachable_names().is_empty());
577 assert!(node.local_catalog(true)["subjects"]
578 .as_array()
579 .unwrap()
580 .is_empty());
581 }
582
583 #[tokio::test]
584 async fn missing_removals_are_snapshot_no_ops_and_replacement_is_effective() {
585 let node = Node::builder("snapshot-no-op")
586 .insecure_accept_declared_peer_identities()
587 .build()
588 .unwrap();
589 node.remove_subject("missing").await.unwrap();
590 node.remove_operation("missing", Operation::Unary)
591 .await
592 .unwrap();
593 let initial = node.snapshot.load_full();
594 assert_eq!(node.catalog_revision(), 0);
595 node.remove_subject("missing").await.unwrap();
596 assert!(Arc::ptr_eq(&initial, &node.snapshot.load_full()));
597
598 node.add_service(service("replace", Operation::Unary, "first"))
599 .await
600 .unwrap();
601 node.remove_operation("replace", Operation::Unary)
602 .await
603 .unwrap();
604 node.add_service(service("replace", Operation::Unary, "second"))
605 .await
606 .unwrap();
607 assert_eq!(node.catalog_revision(), 3);
608 let call = node.snapshot.load().services["replace"]
609 .unary
610 .as_ref()
611 .unwrap()
612 .call
613 .clone();
614 let response = call(http::Request::new(bytes::Bytes::new())).await.unwrap();
615 let crate::layer::ServiceBody::Unary(payload) = response.into_body() else {
616 panic!("expected unary response");
617 };
618 let value: Value = serde_json::from_slice(&payload).unwrap();
619 assert_eq!(value, Value::String("second".into()));
620
621 let revision = node.catalog_revision();
622 let current = node.snapshot.load_full();
623 node.remove_operation("replace", Operation::Streaming)
624 .await
625 .unwrap();
626 assert_eq!(node.catalog_revision(), revision);
627 assert!(Arc::ptr_eq(¤t, &node.snapshot.load_full()));
628 }
629
630 #[tokio::test]
631 async fn published_snapshot_keeps_services_and_capabilities_consistent() {
632 let node = Node::builder("snapshot-consistency")
633 .insecure_accept_declared_peer_identities()
634 .build()
635 .unwrap();
636
637 node.add_service(service("chess", Operation::Unary, "move"))
638 .await
639 .unwrap();
640 let snapshot = node.snapshot.load_full();
641
642 assert_eq!(
643 snapshot.services.keys().collect::<Vec<_>>(),
644 snapshot.capabilities.entries().keys().collect::<Vec<_>>()
645 );
646 assert_eq!(
647 snapshot.node_core.resolve("chess"),
648 unb_core::Resolution::Local
649 );
650 assert_eq!(snapshot.node_core.catalog_revision(), 1);
651 }
652}