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