use std::collections::BTreeSet;
use std::sync::Arc;
use std::time::Duration;
use tokio::sync::mpsc;
use unb_core::{
CoreInput, Detail, DiscoverEvent, DiscoverPlan, DiscoverWalk, Mode, Scope, StreamKey,
TargetPath, WalkInput, WalkOutput, DEFAULT_HOPS,
};
use unb_runtime::{CancellationToken, ProtocolCoreHandle};
use crate::node::Node;
use crate::HandlerError;
const DISCOVER_TIMEOUT: Duration = Duration::from_secs(5);
impl Node {
pub async fn discover_target_events(
self: &Arc<Self>,
target: &str,
detail: Detail,
scope: Scope,
) -> Result<Vec<DiscoverEvent>, HandlerError> {
let target_path = TargetPath::discovery(target)
.map_err(|error| {
HandlerError::new(unb_core::ErrorCode::InvalidInput, error.to_string())
})?
.to_string();
if target == self.identity().node_id {
return Ok(self.discover_events(detail, scope).await);
}
let deadline = n0_future::time::Instant::now() + DISCOVER_TIMEOUT;
let (snapshot, resolution) = self.resolve_unary_until(target, deadline).await?;
let peer = match resolution {
unb_core::Resolution::Route(peer) => peer,
unb_core::Resolution::Unknown => {
return Err(Self::teach_unknown_target(&snapshot, target))
}
unb_core::Resolution::Conflicted { owners } => {
return Err(HandlerError::new(
unb_core::ErrorCode::PeerUnreachable,
format!(
"target node {target:?} has multiple live incarnations: {}",
owners.join(", ")
),
))
}
unb_core::Resolution::Local => return Ok(self.discover_events(detail, scope).await),
};
let plan = DiscoverPlan {
discover_id: String::new(),
detail,
scope,
hops: DEFAULT_HOPS,
visited: BTreeSet::new(),
timeout_ms: Some(DISCOVER_TIMEOUT.as_millis() as u64),
mode: Mode::PartialOk,
};
let link = self.route_link(&peer).await?;
let mut stream = link
.wire
.client_session()
.discover(&target_path, plan)
.await
.map_err(|error| {
HandlerError::new(unb_core::ErrorCode::PeerUnreachable, error.to_string())
})?;
let mut events = Vec::new();
while let Some(envelope) = stream.next().await.map_err(|error| {
HandlerError::new(unb_core::ErrorCode::PeerUnreachable, error.to_string())
})? {
let event =
serde_json::from_slice::<DiscoverEvent>(&envelope.payload).map_err(|error| {
HandlerError::new(unb_core::ErrorCode::Protocol, error.to_string())
})?;
let done = matches!(event, DiscoverEvent::Done { .. });
events.push(event);
if done {
break;
}
}
Ok(events)
}
pub(crate) fn query_discovery_target(
self: &Arc<Self>,
stream: StreamKey,
peer: String,
target_path: String,
plan: DiscoverPlan,
handle: ProtocolCoreHandle,
) {
let node = self.clone();
unb_runtime::RuntimeHandle::current().spawn(async move {
let timeout = plan
.timeout_ms
.map(Duration::from_millis)
.unwrap_or(DISCOVER_TIMEOUT);
let outcome = async {
let _permit = node
.dispatch_slots
.clone()
.acquire_owned()
.await
.map_err(|_| "discovery dispatch capacity closed".to_string())?;
let link = node
.peer(&peer)
.await
.ok_or_else(|| format!("no live connection to peer {peer:?}"))?;
let mut downstream = link
.wire
.client_session()
.discover(&target_path, plan)
.await
.map_err(|error| error.to_string())?;
loop {
let envelope = downstream
.next()
.await
.map_err(|error| error.to_string())?
.ok_or_else(|| "target discovery ended without completion".to_string())?;
let event = serde_json::from_slice::<DiscoverEvent>(&envelope.payload)
.map_err(|error| error.to_string())?;
let terminal = matches!(event, DiscoverEvent::Done { .. });
handle
.submit(CoreInput::DiscoveryTargetEvent {
stream: stream.clone(),
event,
})
.await
.map_err(|error| error.to_string())?;
if terminal {
return Ok::<(), String>(());
}
}
};
let outcome = tokio::select! {
biased;
() = node.cancellation.cancelled() => Err("node shut down during target discovery".to_string()),
result = n0_future::time::timeout(timeout, outcome) => result
.map_err(|_| format!("discovery target {target_path:?} timed out"))
.and_then(|result| result),
};
if let Err(message) = outcome {
let _ = handle
.submit(CoreInput::DiscoveryTargetFailed { stream, message })
.await;
}
});
}
pub(crate) fn query_discovery_neighbor(
self: &Arc<Self>,
stream: StreamKey,
peer: String,
plan: DiscoverPlan,
handle: ProtocolCoreHandle,
) {
let node = self.clone();
unb_runtime::RuntimeHandle::current().spawn(async move {
let (input, mut output) = mpsc::channel(64);
let cancel = node.cancellation.child_token();
let timeout = plan
.timeout_ms
.map(Duration::from_millis)
.unwrap_or(DISCOVER_TIMEOUT);
let leg = node.run_discovery_leg(&peer, plan, &input, &cancel, timeout);
tokio::pin!(leg);
loop {
tokio::select! {
biased;
event = output.recv() => match event {
Some(WalkInput::NeighborEvent { event, .. }) => {
if handle.submit(CoreInput::DiscoveryNeighborEvent {
stream: stream.clone(),
peer: peer.clone(),
event,
}).await.is_err() {
return;
}
}
_ => return,
},
completed = &mut leg => {
while let Ok(WalkInput::NeighborEvent { event, .. }) = output.try_recv() {
if handle.submit(CoreInput::DiscoveryNeighborEvent {
stream: stream.clone(),
peer: peer.clone(),
event,
}).await.is_err() {
return;
}
}
let input = if completed {
CoreInput::DiscoveryNeighborDone {
stream,
peer: peer.clone(),
}
} else {
CoreInput::DiscoveryNeighborTimeout {
stream,
peer: peer.clone(),
}
};
let _ = handle.submit(input).await;
return;
}
}
}
});
}
pub async fn discover_events(
self: &Arc<Self>,
detail: Detail,
scope: Scope,
) -> Vec<DiscoverEvent> {
let plan = DiscoverPlan {
discover_id: String::new(),
detail,
scope,
hops: DEFAULT_HOPS,
visited: BTreeSet::new(),
timeout_ms: None,
mode: Mode::PartialOk,
};
let discover_id = plan.discover_id.clone();
let snapshot = self
.snapshot
.load()
.node_core
.catalog_snapshot(plan.detail.is_full());
let candidates = self.discovery_candidates(&plan.visited).await;
let mut walk = DiscoverWalk::start(snapshot, plan, candidates);
let (input_tx, mut input_rx) = mpsc::channel::<WalkInput>(64);
let cancel = self.cancellation.child_token();
let mut events = Vec::new();
loop {
while let Some(output) = walk.drain() {
match output {
WalkOutput::Emit(event) => events.push(event),
WalkOutput::AskNeighbor { peer, plan } => self.spawn_discovery_leg(
peer,
plan,
input_tx.clone(),
DISCOVER_TIMEOUT,
cancel.child_token(),
),
WalkOutput::Finish => {
events.push(DiscoverEvent::Done {
discover_id: discover_id.clone(),
});
cancel.cancel();
return events;
}
}
}
tokio::select! {
biased;
() = cancel.cancelled() => return events,
input = input_rx.recv() => match input {
Some(input) => walk.handle(input),
None => return events,
}
}
}
}
pub(crate) async fn discovery_candidates(&self, visited: &BTreeSet<String>) -> Vec<String> {
self.peers
.read()
.await
.keys()
.filter(|name| !visited.contains(*name))
.cloned()
.collect()
}
fn spawn_discovery_leg(
self: &Arc<Self>,
peer: String,
plan: DiscoverPlan,
input: mpsc::Sender<WalkInput>,
timeout: Duration,
cancel: CancellationToken,
) {
let node = self.clone();
unb_runtime::RuntimeHandle::current().spawn(async move {
let completed = node
.run_discovery_leg(&peer, plan, &input, &cancel, timeout)
.await;
let feedback = if completed {
WalkInput::NeighborDone { peer }
} else {
WalkInput::NeighborTimeout { peer }
};
let _ = input.send(feedback).await;
});
}
async fn run_discovery_leg(
&self,
peer: &str,
plan: DiscoverPlan,
input: &mpsc::Sender<WalkInput>,
cancel: &CancellationToken,
timeout: Duration,
) -> bool {
let Ok(Ok(_permit)) =
n0_future::time::timeout(timeout, self.dispatch_slots.clone().acquire_owned()).await
else {
return false;
};
let Some(link) = self.peer(peer).await else {
return false;
};
let target_path = match TargetPath::discovery(peer) {
Ok(path) => path.to_string(),
Err(_) => return false,
};
let mut stream = match link
.wire
.client_session()
.discover(&target_path, plan)
.await
{
Ok(stream) => stream,
Err(_) => return false,
};
let deadline = n0_future::time::sleep(timeout);
tokio::pin!(deadline);
let completed = loop {
tokio::select! {
biased;
() = cancel.cancelled() => break false,
() = &mut deadline => break false,
msg = stream.next() => match msg {
Ok(Some(frame)) => match frame.kind {
unb_core::Kind::Event => {
if let Ok(event) =
serde_json::from_slice::<DiscoverEvent>(&frame.payload)
{
if input
.send(WalkInput::NeighborEvent {
peer: peer.to_string(),
event,
})
.await
.is_err()
{
break false;
}
}
}
unb_core::Kind::Response => break true,
_ => {}
},
Ok(None) => break false,
Err(_) => break false,
}
}
};
completed
}
}