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) verified_peers: Mutex<HashMap<SessionId, VerifiedPeer>>,
202 pub(crate) candidate_identities:
203 Mutex<HashMap<SessionId, tokio::sync::watch::Sender<Option<NodeIdentity>>>>,
204 pub(crate) active: Arc<Mutex<HashMap<(SessionId, String), CancellationToken>>>,
205 pub(crate) protocol: ProtocolCoreHandle,
206 pub(crate) identity: NodeIdentity,
207 pub(crate) peer_layers: Arc<[Arc<dyn PeerLayer>]>,
208 pub(crate) dial_policy: unb_client::Peers,
209 pub(crate) next_session: AtomicU64,
210 pub(crate) ws_collect_ceiling: usize,
211 pub(crate) cancellation: CancellationToken,
212 pub(crate) _shutdown: DropGuard,
213}
214
215impl Node {
216 pub fn cancellation(&self) -> &CancellationToken {
217 &self.cancellation
218 }
219
220 pub fn identity(&self) -> &NodeIdentity {
221 &self.identity
222 }
223
224 pub fn shutdown(&self) {
225 for connection in self
226 .connections
227 .read()
228 .unwrap_or_else(|poisoned| poisoned.into_inner())
229 .values()
230 {
231 connection.node_shutdown();
232 }
233 self.cancellation.cancel();
234 }
235
236 pub fn reachable_names(&self) -> Vec<String> {
237 self.snapshot.load().node_core.reachable_names()
238 }
239
240 pub fn catalog_revision(&self) -> u64 {
241 self.snapshot.load().node_core.catalog_revision()
242 }
243
244 pub fn local_catalog(&self, detail_full: bool) -> Value {
245 self.snapshot.load().node_core.catalog(detail_full)
246 }
247
248 pub async fn remove_subject(&self, subject: &str) -> Result<(), WsError> {
249 let _gate = self.mutation_gate.lock().await;
250 let mut services = self.snapshot.load().services.clone();
251 services.remove(subject);
252 self.install_services(services).await
253 }
254
255 pub async fn add_service(&self, handler: impl Handler) -> Result<(), WsError> {
256 let _gate = self.mutation_gate.lock().await;
257 let mut services = self.snapshot.load().services.clone();
258 SubjectServices::register(
259 &mut services,
260 handler.into_service(),
261 &[],
262 self.global_layers.iter().cloned().collect(),
263 &self.states,
264 )
265 .map_err(WsError::Connect)?;
266 self.install_services(services).await
267 }
268
269 pub async fn remove_operation(
270 &self,
271 subject: &str,
272 operation: Operation,
273 ) -> Result<(), WsError> {
274 let _gate = self.mutation_gate.lock().await;
275 let mut services = self.snapshot.load().services.clone();
276 let Some(existing) = services.get(subject).cloned() else {
277 return Ok(());
278 };
279 let mut entry = existing.as_ref().clone();
280 let slot = match operation {
281 Operation::Unary => &mut entry.unary,
282 Operation::Streaming => &mut entry.streaming,
283 };
284 if slot.take().is_none() {
285 return Ok(());
286 }
287 let empty = entry.unary.is_none() && entry.streaming.is_none();
288 if empty {
289 services.remove(subject);
290 } else {
291 services.insert(subject.to_string(), Arc::new(entry));
292 }
293 self.install_services(services).await
294 }
295
296 async fn install_services(
297 &self,
298 services: BTreeMap<String, Arc<SubjectServices>>,
299 ) -> Result<(), WsError> {
300 let current = self.snapshot.load();
301 if current.same_service_contract(&services) {
302 return Ok(());
303 }
304 let snapshot = NodeSnapshot::new(services, current.node_core.clone());
305 let capabilities = snapshot.capabilities.clone();
306 let publication = self.snapshot.clone();
307 self.protocol
308 .install(
309 CoreInput::LocalCapabilitiesInstalled {
310 capabilities: capabilities.entries().clone(),
311 },
312 move || publication.store(Arc::new(snapshot)),
313 )
314 .await
315 }
316
317 pub(crate) async fn peer(&self, name: &str) -> Option<PeerLink> {
318 self.peers.read().await.get(name).cloned()
319 }
320
321 pub(crate) async fn session(&self, id: &str) -> Option<Arc<Wire>> {
322 self.sessions.read().await.get(id).cloned()
323 }
324
325 pub(crate) fn connection(&self, peer: &str) -> Option<PeerConnection> {
326 self.connections
327 .read()
328 .unwrap_or_else(|poisoned| poisoned.into_inner())
329 .get(peer)
330 .cloned()
331 }
332
333 pub(crate) fn route_changes(&self) -> tokio::sync::watch::Receiver<u64> {
334 self.routes_changed.subscribe()
335 }
336
337 pub(crate) fn publish_route_change(&self) {
338 self.routes_changed.send_modify(|revision| {
339 *revision = revision
340 .checked_add(1)
341 .expect("route change revision overflow");
342 });
343 }
344
345 pub(crate) fn readiness_waits_for_destination(
346 &self,
347 destination: &str,
348 ) -> Vec<crate::connection::ReadinessWait> {
349 self.connections
350 .read()
351 .unwrap_or_else(|poisoned| poisoned.into_inner())
352 .values()
353 .filter_map(|connection| {
354 if connection.is_terminal() || !connection.carried_destination(destination) {
355 return None;
356 }
357 connection.readiness_wait()
358 })
359 .collect()
360 }
361}
362
363impl Drop for Node {
364 fn drop(&mut self) {
365 for connection in self
366 .connections
367 .read()
368 .unwrap_or_else(|poisoned| poisoned.into_inner())
369 .values()
370 {
371 connection.node_shutdown();
372 }
373 }
374}
375
376#[cfg(test)]
377mod tests {
378 use super::*;
379 use crate::service::OperationContract;
380
381 fn operation(contract: Value) -> CompiledOperation {
382 CompiledOperation {
383 call: Arc::new(|_| Box::pin(async { unreachable!() })),
384 layers: Arc::from([]),
385 contract,
386 }
387 }
388
389 fn service(subject: &str, operation: Operation, result: &'static str) -> HandlerService {
390 HandlerService::declare(
391 subject,
392 None,
393 Some(result),
394 operation,
395 OperationContract::unknown(),
396 move |_| {
397 Ok(Arc::new(move |_| {
398 Box::pin(async move {
399 Ok(http::Response::new(crate::layer::ServiceBody::Unary(
400 unb_core::Envelope::encode_payload(&Value::String(result.into())),
401 )))
402 })
403 }))
404 },
405 )
406 }
407
408 #[test]
409 fn node_snapshot_capabilities_match_service_subjects_and_operations() {
410 let mut services = BTreeMap::new();
411 services.insert(
412 "chess.move".into(),
413 Arc::new(SubjectServices {
414 unary: Some(operation(serde_json::json!({ "input": "Move" }))),
415 streaming: Some(operation(serde_json::json!({ "event": "Position" }))),
416 metadata: Some(serde_json::json!({ "one_line": "Play a move", "tier": 1 })),
417 one_line: None,
418 }),
419 );
420 services.insert(
421 "chess.state".into(),
422 Arc::new(SubjectServices {
423 streaming: Some(operation(serde_json::json!({ "event": "Position" }))),
424 one_line: Some("Watch the board".into()),
425 ..SubjectServices::default()
426 }),
427 );
428
429 let snapshot = NodeSnapshot::new(services, NodeCore::new("snapshot-test"));
430
431 assert_eq!(
432 snapshot.capabilities.entries().keys().collect::<Vec<_>>(),
433 snapshot.services.keys().collect::<Vec<_>>()
434 );
435 assert_eq!(
436 snapshot.capabilities.entries()["chess.move"],
437 serde_json::json!({
438 "one_line": "Play a move",
439 "metadata": { "one_line": "Play a move", "tier": 1 },
440 "operations": {
441 "unary": { "input": "Move" },
442 "streaming": { "event": "Position" }
443 }
444 })
445 );
446 assert_eq!(
447 snapshot.capabilities.entries()["chess.state"],
448 serde_json::json!({
449 "one_line": "Watch the board",
450 "operations": { "streaming": { "event": "Position" } }
451 })
452 );
453 }
454
455 #[tokio::test]
456 async fn mutation_updates_catalog_without_churning_node_routes() {
457 let node = Node::builder("snapshot-test")
458 .insecure_accept_declared_peer_identities()
459 .build()
460 .unwrap();
461 let empty_fingerprint = node.snapshot.load().node_core.fingerprint();
462
463 node.add_service(
464 service("chess", Operation::Unary, "move")
465 .describe(serde_json::json!({ "one_line": "Play chess", "tier": 1 })),
466 )
467 .await
468 .unwrap();
469 assert_eq!(node.catalog_revision(), 1);
470 let export = node.snapshot.load().node_core.export_for("peer");
471 assert_eq!(export.len(), 1);
472 assert_eq!(export[0].destination, "snapshot-test");
473 assert_eq!(export[0].owner_revision, 0);
474 assert_ne!(
475 node.snapshot.load().node_core.fingerprint(),
476 empty_fingerprint
477 );
478 assert_eq!(
479 node.local_catalog(true)["subjects"][0],
480 serde_json::json!({
481 "subject": "chess",
482 "target_path": "/snapshot-test/chess",
483 "one_line": "Play chess",
484 "metadata": { "one_line": "Play chess", "tier": 1 },
485 "operations": {
486 "unary": {
487 "input_schema": { "unknown": true },
488 "output_schema": { "unknown": true }
489 }
490 }
491 })
492 );
493
494 node.add_service(service("chess", Operation::Streaming, "watch"))
495 .await
496 .unwrap();
497 assert_eq!(node.catalog_revision(), 2);
498 assert_eq!(
499 node.snapshot.load().node_core.export_for("peer")[0].owner_revision,
500 0
501 );
502 assert_eq!(
503 node.snapshot.load().node_core.resolve("chess"),
504 unb_core::Resolution::Unknown
505 );
506 assert_eq!(
507 node.snapshot.load().node_core.resolve("snapshot-test"),
508 unb_core::Resolution::Local
509 );
510 assert!(node.local_catalog(true)["subjects"][0]["operations"]["streaming"].is_object());
511
512 node.remove_operation("chess", Operation::Unary)
513 .await
514 .unwrap();
515 assert_eq!(node.catalog_revision(), 3);
516 assert_eq!(
517 node.snapshot.load().node_core.export_for("peer")[0].owner_revision,
518 0
519 );
520 assert_eq!(
521 node.snapshot.load().node_core.resolve("chess"),
522 unb_core::Resolution::Unknown
523 );
524 assert!(node.local_catalog(true)["subjects"][0]["operations"]["unary"].is_null());
525
526 node.remove_operation("chess", Operation::Streaming)
527 .await
528 .unwrap();
529 assert_eq!(node.catalog_revision(), 4);
530 assert_eq!(node.snapshot.load().node_core.export_for("peer").len(), 1);
531 assert_eq!(
532 node.snapshot.load().node_core.fingerprint(),
533 empty_fingerprint
534 );
535 assert_eq!(node.reachable_names(), ["snapshot-test"]);
536 assert!(node.local_catalog(true)["subjects"]
537 .as_array()
538 .unwrap()
539 .is_empty());
540 }
541
542 #[tokio::test]
543 async fn missing_removals_are_snapshot_no_ops_and_replacement_is_effective() {
544 let node = Node::builder("snapshot-no-op")
545 .insecure_accept_declared_peer_identities()
546 .build()
547 .unwrap();
548 node.remove_subject("missing").await.unwrap();
549 node.remove_operation("missing", Operation::Unary)
550 .await
551 .unwrap();
552 let initial = node.snapshot.load_full();
553 assert_eq!(node.catalog_revision(), 0);
554 node.remove_subject("missing").await.unwrap();
555 assert!(Arc::ptr_eq(&initial, &node.snapshot.load_full()));
556
557 node.add_service(service("replace", Operation::Unary, "first"))
558 .await
559 .unwrap();
560 node.remove_operation("replace", Operation::Unary)
561 .await
562 .unwrap();
563 node.add_service(service("replace", Operation::Unary, "second"))
564 .await
565 .unwrap();
566 assert_eq!(node.catalog_revision(), 3);
567 let call = node.snapshot.load().services["replace"]
568 .unary
569 .as_ref()
570 .unwrap()
571 .call
572 .clone();
573 let response = call(http::Request::new(bytes::Bytes::new())).await.unwrap();
574 let crate::layer::ServiceBody::Unary(payload) = response.into_body() else {
575 panic!("expected unary response");
576 };
577 let value: Value = serde_json::from_slice(&payload).unwrap();
578 assert_eq!(value, Value::String("second".into()));
579
580 let revision = node.catalog_revision();
581 let current = node.snapshot.load_full();
582 node.remove_operation("replace", Operation::Streaming)
583 .await
584 .unwrap();
585 assert_eq!(node.catalog_revision(), revision);
586 assert!(Arc::ptr_eq(¤t, &node.snapshot.load_full()));
587 }
588
589 #[tokio::test]
590 async fn published_snapshot_keeps_services_and_capabilities_consistent() {
591 let node = Node::builder("snapshot-consistency")
592 .insecure_accept_declared_peer_identities()
593 .build()
594 .unwrap();
595
596 node.add_service(service("chess", Operation::Unary, "move"))
597 .await
598 .unwrap();
599 let snapshot = node.snapshot.load_full();
600
601 assert_eq!(
602 snapshot.services.keys().collect::<Vec<_>>(),
603 snapshot.capabilities.entries().keys().collect::<Vec<_>>()
604 );
605 assert_eq!(
606 snapshot.node_core.resolve("chess"),
607 unb_core::Resolution::Unknown
608 );
609 assert_eq!(
610 snapshot.node_core.resolve("snapshot-consistency"),
611 unb_core::Resolution::Local
612 );
613 assert_eq!(snapshot.node_core.catalog_revision(), 1);
614 }
615}