use std::collections::{BTreeMap, HashMap, HashSet};
use std::sync::atomic::AtomicU64;
use std::sync::{Arc, RwLock as StdRwLock};
#[cfg(feature = "hosting")]
use std::time::Duration;
use arc_swap::ArcSwap;
use serde_json::Value;
use tokio::sync::{Mutex, RwLock, Semaphore};
use unb_core::{CoreCapabilitySnapshot, CoreInput, EffectId, NodeCore, NodeIdentity, SessionId};
use unb_runtime::{CancellationToken, DropGuard, ProtocolCoreHandle, Wire, WsError};
use crate::layer::{ErasedCall, Layer};
use crate::peer::{PeerLayer, VerifiedPeer};
use crate::service::{Handler, HandlerService, Operation, StateMap, States};
use crate::PeerConnection;
#[cfg(feature = "hosting")]
pub(crate) const WEBTRANSPORT_ACCEPT_TIMEOUT: Duration = Duration::from_secs(5);
pub(crate) struct PeerLink {
pub session_id: String,
pub wire: Arc<Wire>,
pub instance_id: String,
pub outbound: bool,
}
impl Clone for PeerLink {
fn clone(&self) -> PeerLink {
PeerLink {
session_id: self.session_id.clone(),
wire: self.wire.clone(),
instance_id: self.instance_id.clone(),
outbound: self.outbound,
}
}
}
#[derive(Clone)]
pub(crate) struct CompiledOperation {
pub(crate) call: ErasedCall,
pub(crate) layers: Arc<[Arc<dyn Layer>]>,
pub(crate) contract: Value,
}
#[derive(Clone, Default)]
pub(crate) struct SubjectServices {
pub(crate) unary: Option<CompiledOperation>,
pub(crate) streaming: Option<CompiledOperation>,
pub(crate) metadata: Option<Value>,
pub(crate) one_line: Option<String>,
}
impl SubjectServices {
pub(crate) fn register(
services: &mut BTreeMap<String, Arc<SubjectServices>>,
service: HandlerService,
scopes: &[String],
layers: Vec<Arc<dyn Layer>>,
states: &StateMap,
) -> Result<(String, Value), String> {
let subject = service.effective_subject(scopes)?;
let mut entry = services
.get(&subject)
.map(|existing| existing.as_ref().clone())
.unwrap_or_default();
let slot = match service.operation {
Operation::Unary => &mut entry.unary,
Operation::Streaming => &mut entry.streaming,
};
if slot.is_some() {
return Err(format!(
"subject {subject:?} already serves a {:?} operation",
service.operation
));
}
if service.metadata.is_some() && entry.metadata.is_some() {
return Err(format!(
"subject {subject:?} already carries metadata; attach it to one registration"
));
}
let call = (service.build)(&States(states))?;
*slot = Some(CompiledOperation {
call,
layers: Arc::from(layers),
contract: service.contract.to_json(service.operation),
});
if service.metadata.is_some() {
entry.metadata = service.metadata;
}
if entry.one_line.is_none() {
entry.one_line = service.one_line;
}
let catalog_entry = entry.catalog_entry();
services.insert(subject.clone(), Arc::new(entry));
Ok((subject, catalog_entry))
}
pub(crate) fn catalog_entry(&self) -> Value {
let mut operations = serde_json::Map::new();
if let Some(unary) = &self.unary {
operations.insert("unary".into(), unary.contract.clone());
}
if let Some(streaming) = &self.streaming {
operations.insert("streaming".into(), streaming.contract.clone());
}
let mut entry = serde_json::Map::new();
if let Some(metadata) = &self.metadata {
if let Some(one_line) = metadata.get("one_line") {
entry.insert("one_line".into(), one_line.clone());
}
entry.insert("metadata".into(), metadata.clone());
}
if !entry.contains_key("one_line") {
if let Some(one_line) = &self.one_line {
entry.insert("one_line".into(), Value::String(one_line.clone()));
}
}
entry.insert("operations".into(), Value::Object(operations));
Value::Object(entry)
}
fn same_contract(&self, other: &Self) -> bool {
fn same_operation(
left: &Option<CompiledOperation>,
right: &Option<CompiledOperation>,
) -> bool {
match (left, right) {
(Some(left), Some(right)) => {
Arc::ptr_eq(&left.call, &right.call)
&& left.contract == right.contract
&& left.layers.len() == right.layers.len()
&& left
.layers
.iter()
.zip(right.layers.iter())
.all(|(left, right)| Arc::ptr_eq(left, right))
}
(None, None) => true,
_ => false,
}
}
same_operation(&self.unary, &other.unary)
&& same_operation(&self.streaming, &other.streaming)
&& self.metadata == other.metadata
&& self.one_line == other.one_line
}
}
#[derive(Clone)]
pub(crate) struct NodeSnapshot {
pub(crate) services: BTreeMap<String, Arc<SubjectServices>>,
pub(crate) capabilities: CoreCapabilitySnapshot,
pub(crate) node_core: NodeCore,
}
impl NodeSnapshot {
pub(crate) fn new(
services: BTreeMap<String, Arc<SubjectServices>>,
mut node_core: NodeCore,
) -> Self {
let capabilities = CoreCapabilitySnapshot::new(
services
.iter()
.map(|(subject, services)| (subject.clone(), services.catalog_entry()))
.collect(),
);
node_core.install_local_capabilities(capabilities.entries().clone());
Self {
services,
capabilities,
node_core,
}
}
fn same_service_contract(&self, services: &BTreeMap<String, Arc<SubjectServices>>) -> bool {
self.services.len() == services.len()
&& self.services.iter().all(|(subject, current)| {
services
.get(subject)
.is_some_and(|result| current.same_contract(result))
})
}
}
pub struct Node {
pub(crate) snapshot: Arc<ArcSwap<NodeSnapshot>>,
pub(crate) states: StateMap,
pub(crate) global_layers: Arc<[Arc<dyn Layer>]>,
pub(crate) mutation_gate: Mutex<()>,
pub(crate) peers: RwLock<HashMap<String, PeerLink>>,
pub(crate) sessions: RwLock<HashMap<String, Arc<Wire>>>,
pub(crate) connections: StdRwLock<HashMap<String, PeerConnection>>,
pub(crate) routes_changed: tokio::sync::watch::Sender<u64>,
pub(crate) session_peers: RwLock<HashMap<String, String>>,
pub(crate) outbound_sessions: Mutex<HashSet<SessionId>>,
pub(crate) dispatch_slots: Arc<Semaphore>,
pub(crate) dispatch_permits: Mutex<HashMap<EffectId, tokio::sync::OwnedSemaphorePermit>>,
pub(crate) dispatching: Mutex<HashMap<EffectId, CancellationToken>>,
pub(crate) peer_admissions: Mutex<HashMap<SessionId, CancellationToken>>,
pub(crate) verified_peers: Mutex<HashMap<SessionId, VerifiedPeer>>,
pub(crate) candidate_identities:
Mutex<HashMap<SessionId, tokio::sync::watch::Sender<Option<NodeIdentity>>>>,
pub(crate) active: Arc<Mutex<HashMap<(SessionId, String), CancellationToken>>>,
pub(crate) protocol: ProtocolCoreHandle,
pub(crate) identity: NodeIdentity,
pub(crate) peer_layers: Arc<[Arc<dyn PeerLayer>]>,
pub(crate) dial_policy: unb_client::Peers,
pub(crate) next_session: AtomicU64,
pub(crate) ws_collect_ceiling: usize,
pub(crate) cancellation: CancellationToken,
pub(crate) _shutdown: DropGuard,
}
impl Node {
pub fn cancellation(&self) -> &CancellationToken {
&self.cancellation
}
pub fn identity(&self) -> &NodeIdentity {
&self.identity
}
pub fn shutdown(&self) {
for connection in self
.connections
.read()
.unwrap_or_else(|poisoned| poisoned.into_inner())
.values()
{
connection.node_shutdown();
}
self.cancellation.cancel();
}
pub fn reachable_names(&self) -> Vec<String> {
self.snapshot.load().node_core.reachable_names()
}
pub fn catalog_revision(&self) -> u64 {
self.snapshot.load().node_core.catalog_revision()
}
pub fn local_catalog(&self, detail_full: bool) -> Value {
self.snapshot.load().node_core.catalog(detail_full)
}
pub async fn remove_subject(&self, subject: &str) -> Result<(), WsError> {
let _gate = self.mutation_gate.lock().await;
let mut services = self.snapshot.load().services.clone();
services.remove(subject);
self.install_services(services).await
}
pub async fn add_service(&self, handler: impl Handler) -> Result<(), WsError> {
let _gate = self.mutation_gate.lock().await;
let mut services = self.snapshot.load().services.clone();
SubjectServices::register(
&mut services,
handler.into_service(),
&[],
self.global_layers.iter().cloned().collect(),
&self.states,
)
.map_err(WsError::Connect)?;
self.install_services(services).await
}
pub async fn remove_operation(
&self,
subject: &str,
operation: Operation,
) -> Result<(), WsError> {
let _gate = self.mutation_gate.lock().await;
let mut services = self.snapshot.load().services.clone();
let Some(existing) = services.get(subject).cloned() else {
return Ok(());
};
let mut entry = existing.as_ref().clone();
let slot = match operation {
Operation::Unary => &mut entry.unary,
Operation::Streaming => &mut entry.streaming,
};
if slot.take().is_none() {
return Ok(());
}
let empty = entry.unary.is_none() && entry.streaming.is_none();
if empty {
services.remove(subject);
} else {
services.insert(subject.to_string(), Arc::new(entry));
}
self.install_services(services).await
}
async fn install_services(
&self,
services: BTreeMap<String, Arc<SubjectServices>>,
) -> Result<(), WsError> {
let current = self.snapshot.load();
if current.same_service_contract(&services) {
return Ok(());
}
let snapshot = NodeSnapshot::new(services, current.node_core.clone());
let capabilities = snapshot.capabilities.clone();
let publication = self.snapshot.clone();
self.protocol
.install(
CoreInput::LocalCapabilitiesInstalled {
capabilities: capabilities.entries().clone(),
},
move || publication.store(Arc::new(snapshot)),
)
.await
}
pub(crate) async fn peer(&self, name: &str) -> Option<PeerLink> {
self.peers.read().await.get(name).cloned()
}
pub(crate) async fn session(&self, id: &str) -> Option<Arc<Wire>> {
self.sessions.read().await.get(id).cloned()
}
pub(crate) fn connection(&self, peer: &str) -> Option<PeerConnection> {
self.connections
.read()
.unwrap_or_else(|poisoned| poisoned.into_inner())
.get(peer)
.cloned()
}
pub(crate) fn route_changes(&self) -> tokio::sync::watch::Receiver<u64> {
self.routes_changed.subscribe()
}
pub(crate) fn publish_route_change(&self) {
self.routes_changed.send_modify(|revision| {
*revision = revision
.checked_add(1)
.expect("route change revision overflow");
});
}
pub(crate) fn readiness_waits_for_destination(
&self,
destination: &str,
) -> Vec<crate::connection::ReadinessWait> {
self.connections
.read()
.unwrap_or_else(|poisoned| poisoned.into_inner())
.values()
.filter_map(|connection| {
if connection.is_terminal() || !connection.carried_destination(destination) {
return None;
}
connection.readiness_wait()
})
.collect()
}
}
impl Drop for Node {
fn drop(&mut self) {
for connection in self
.connections
.read()
.unwrap_or_else(|poisoned| poisoned.into_inner())
.values()
{
connection.node_shutdown();
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::service::OperationContract;
fn operation(contract: Value) -> CompiledOperation {
CompiledOperation {
call: Arc::new(|_| Box::pin(async { unreachable!() })),
layers: Arc::from([]),
contract,
}
}
fn service(subject: &str, operation: Operation, result: &'static str) -> HandlerService {
HandlerService::declare(
subject,
None,
Some(result),
operation,
OperationContract::unknown(),
move |_| {
Ok(Arc::new(move |_| {
Box::pin(async move {
Ok(http::Response::new(crate::layer::ServiceBody::Unary(
unb_core::Envelope::encode_payload(&Value::String(result.into())),
)))
})
}))
},
)
}
#[test]
fn node_snapshot_capabilities_match_service_subjects_and_operations() {
let mut services = BTreeMap::new();
services.insert(
"chess.move".into(),
Arc::new(SubjectServices {
unary: Some(operation(serde_json::json!({ "input": "Move" }))),
streaming: Some(operation(serde_json::json!({ "event": "Position" }))),
metadata: Some(serde_json::json!({ "one_line": "Play a move", "tier": 1 })),
one_line: None,
}),
);
services.insert(
"chess.state".into(),
Arc::new(SubjectServices {
streaming: Some(operation(serde_json::json!({ "event": "Position" }))),
one_line: Some("Watch the board".into()),
..SubjectServices::default()
}),
);
let snapshot = NodeSnapshot::new(services, NodeCore::new("snapshot-test"));
assert_eq!(
snapshot.capabilities.entries().keys().collect::<Vec<_>>(),
snapshot.services.keys().collect::<Vec<_>>()
);
assert_eq!(
snapshot.capabilities.entries()["chess.move"],
serde_json::json!({
"one_line": "Play a move",
"metadata": { "one_line": "Play a move", "tier": 1 },
"operations": {
"unary": { "input": "Move" },
"streaming": { "event": "Position" }
}
})
);
assert_eq!(
snapshot.capabilities.entries()["chess.state"],
serde_json::json!({
"one_line": "Watch the board",
"operations": { "streaming": { "event": "Position" } }
})
);
}
#[tokio::test]
async fn mutation_updates_catalog_without_churning_node_routes() {
let node = Node::builder("snapshot-test")
.insecure_accept_declared_peer_identities()
.build()
.unwrap();
let empty_fingerprint = node.snapshot.load().node_core.fingerprint();
node.add_service(
service("chess", Operation::Unary, "move")
.describe(serde_json::json!({ "one_line": "Play chess", "tier": 1 })),
)
.await
.unwrap();
assert_eq!(node.catalog_revision(), 1);
let export = node.snapshot.load().node_core.export_for("peer");
assert_eq!(export.len(), 1);
assert_eq!(export[0].destination, "snapshot-test");
assert_eq!(export[0].owner_revision, 0);
assert_ne!(
node.snapshot.load().node_core.fingerprint(),
empty_fingerprint
);
assert_eq!(
node.local_catalog(true)["subjects"][0],
serde_json::json!({
"subject": "chess",
"target_path": "/snapshot-test/chess",
"one_line": "Play chess",
"metadata": { "one_line": "Play chess", "tier": 1 },
"operations": {
"unary": {
"input_schema": { "unknown": true },
"output_schema": { "unknown": true }
}
}
})
);
node.add_service(service("chess", Operation::Streaming, "watch"))
.await
.unwrap();
assert_eq!(node.catalog_revision(), 2);
assert_eq!(
node.snapshot.load().node_core.export_for("peer")[0].owner_revision,
0
);
assert_eq!(
node.snapshot.load().node_core.resolve("chess"),
unb_core::Resolution::Unknown
);
assert_eq!(
node.snapshot.load().node_core.resolve("snapshot-test"),
unb_core::Resolution::Local
);
assert!(node.local_catalog(true)["subjects"][0]["operations"]["streaming"].is_object());
node.remove_operation("chess", Operation::Unary)
.await
.unwrap();
assert_eq!(node.catalog_revision(), 3);
assert_eq!(
node.snapshot.load().node_core.export_for("peer")[0].owner_revision,
0
);
assert_eq!(
node.snapshot.load().node_core.resolve("chess"),
unb_core::Resolution::Unknown
);
assert!(node.local_catalog(true)["subjects"][0]["operations"]["unary"].is_null());
node.remove_operation("chess", Operation::Streaming)
.await
.unwrap();
assert_eq!(node.catalog_revision(), 4);
assert_eq!(node.snapshot.load().node_core.export_for("peer").len(), 1);
assert_eq!(
node.snapshot.load().node_core.fingerprint(),
empty_fingerprint
);
assert_eq!(node.reachable_names(), ["snapshot-test"]);
assert!(node.local_catalog(true)["subjects"]
.as_array()
.unwrap()
.is_empty());
}
#[tokio::test]
async fn missing_removals_are_snapshot_no_ops_and_replacement_is_effective() {
let node = Node::builder("snapshot-no-op")
.insecure_accept_declared_peer_identities()
.build()
.unwrap();
node.remove_subject("missing").await.unwrap();
node.remove_operation("missing", Operation::Unary)
.await
.unwrap();
let initial = node.snapshot.load_full();
assert_eq!(node.catalog_revision(), 0);
node.remove_subject("missing").await.unwrap();
assert!(Arc::ptr_eq(&initial, &node.snapshot.load_full()));
node.add_service(service("replace", Operation::Unary, "first"))
.await
.unwrap();
node.remove_operation("replace", Operation::Unary)
.await
.unwrap();
node.add_service(service("replace", Operation::Unary, "second"))
.await
.unwrap();
assert_eq!(node.catalog_revision(), 3);
let call = node.snapshot.load().services["replace"]
.unary
.as_ref()
.unwrap()
.call
.clone();
let response = call(http::Request::new(bytes::Bytes::new())).await.unwrap();
let crate::layer::ServiceBody::Unary(payload) = response.into_body() else {
panic!("expected unary response");
};
let value: Value = serde_json::from_slice(&payload).unwrap();
assert_eq!(value, Value::String("second".into()));
let revision = node.catalog_revision();
let current = node.snapshot.load_full();
node.remove_operation("replace", Operation::Streaming)
.await
.unwrap();
assert_eq!(node.catalog_revision(), revision);
assert!(Arc::ptr_eq(¤t, &node.snapshot.load_full()));
}
#[tokio::test]
async fn published_snapshot_keeps_services_and_capabilities_consistent() {
let node = Node::builder("snapshot-consistency")
.insecure_accept_declared_peer_identities()
.build()
.unwrap();
node.add_service(service("chess", Operation::Unary, "move"))
.await
.unwrap();
let snapshot = node.snapshot.load_full();
assert_eq!(
snapshot.services.keys().collect::<Vec<_>>(),
snapshot.capabilities.entries().keys().collect::<Vec<_>>()
);
assert_eq!(
snapshot.node_core.resolve("chess"),
unb_core::Resolution::Unknown
);
assert_eq!(
snapshot.node_core.resolve("snapshot-consistency"),
unb_core::Resolution::Local
);
assert_eq!(snapshot.node_core.catalog_revision(), 1);
}
}