1use std::collections::BTreeSet;
2use std::sync::Arc;
3use std::time::Duration;
4
5use tokio::sync::mpsc;
6use unb_core::{
7 CoreInput, Detail, DiscoverEvent, DiscoverPlan, DiscoverWalk, Mode, Scope, StreamKey,
8 TargetPath, WalkInput, WalkOutput, DEFAULT_HOPS,
9};
10use unb_runtime::{CancellationToken, ProtocolCoreHandle};
11
12use crate::node::Node;
13use crate::HandlerError;
14
15const DISCOVER_TIMEOUT: Duration = Duration::from_secs(5);
16
17impl Node {
18 pub async fn discover_target_events(
20 self: &Arc<Self>,
21 target: &str,
22 detail: Detail,
23 scope: Scope,
24 ) -> Result<Vec<DiscoverEvent>, HandlerError> {
25 let target_path = TargetPath::discovery(target)
26 .map_err(|error| {
27 HandlerError::new(unb_core::ErrorCode::InvalidInput, error.to_string())
28 })?
29 .to_string();
30 if target == self.identity().node_id {
31 return Ok(self.discover_events(detail, scope).await);
32 }
33 let deadline = n0_future::time::Instant::now() + DISCOVER_TIMEOUT;
34 let (snapshot, resolution) = self.resolve_unary_until(target, deadline).await?;
35 let peer = match resolution {
36 unb_core::Resolution::Route(peer) => peer,
37 unb_core::Resolution::Unknown => {
38 return Err(Self::teach_unknown_target(&snapshot, target))
39 }
40 unb_core::Resolution::Conflicted { owners } => {
41 return Err(HandlerError::new(
42 unb_core::ErrorCode::PeerUnreachable,
43 format!(
44 "target node {target:?} has multiple live incarnations: {}",
45 owners.join(", ")
46 ),
47 ))
48 }
49 unb_core::Resolution::Local => return Ok(self.discover_events(detail, scope).await),
50 };
51 let plan = DiscoverPlan {
52 discover_id: String::new(),
53 detail,
54 scope,
55 hops: DEFAULT_HOPS,
56 visited: BTreeSet::new(),
57 timeout_ms: Some(DISCOVER_TIMEOUT.as_millis() as u64),
58 mode: Mode::PartialOk,
59 };
60 let link = self.route_link(&peer).await?;
61 let mut stream = link
62 .wire
63 .client_session()
64 .discover(&target_path, plan)
65 .await
66 .map_err(|error| {
67 HandlerError::new(unb_core::ErrorCode::PeerUnreachable, error.to_string())
68 })?;
69 let mut events = Vec::new();
70 while let Some(envelope) = stream.next().await.map_err(|error| {
71 HandlerError::new(unb_core::ErrorCode::PeerUnreachable, error.to_string())
72 })? {
73 let event =
74 serde_json::from_slice::<DiscoverEvent>(&envelope.payload).map_err(|error| {
75 HandlerError::new(unb_core::ErrorCode::Protocol, error.to_string())
76 })?;
77 let done = matches!(event, DiscoverEvent::Done { .. });
78 events.push(event);
79 if done {
80 break;
81 }
82 }
83 Ok(events)
84 }
85
86 pub(crate) fn query_discovery_target(
87 self: &Arc<Self>,
88 stream: StreamKey,
89 peer: String,
90 target_path: String,
91 plan: DiscoverPlan,
92 handle: ProtocolCoreHandle,
93 ) {
94 let node = self.clone();
95 unb_runtime::RuntimeHandle::current().spawn(async move {
96 let timeout = plan
97 .timeout_ms
98 .map(Duration::from_millis)
99 .unwrap_or(DISCOVER_TIMEOUT);
100 let outcome = async {
101 let _permit = node
102 .dispatch_slots
103 .clone()
104 .acquire_owned()
105 .await
106 .map_err(|_| "discovery dispatch capacity closed".to_string())?;
107 let link = node
108 .peer(&peer)
109 .await
110 .ok_or_else(|| format!("no live connection to peer {peer:?}"))?;
111 let mut downstream = link
112 .wire
113 .client_session()
114 .discover(&target_path, plan)
115 .await
116 .map_err(|error| error.to_string())?;
117 loop {
118 let envelope = downstream
119 .next()
120 .await
121 .map_err(|error| error.to_string())?
122 .ok_or_else(|| "target discovery ended without completion".to_string())?;
123 let event = serde_json::from_slice::<DiscoverEvent>(&envelope.payload)
124 .map_err(|error| error.to_string())?;
125 let terminal = matches!(event, DiscoverEvent::Done { .. });
126 handle
127 .submit(CoreInput::DiscoveryTargetEvent {
128 stream: stream.clone(),
129 event,
130 })
131 .await
132 .map_err(|error| error.to_string())?;
133 if terminal {
134 return Ok::<(), String>(());
135 }
136 }
137 };
138 let outcome = tokio::select! {
139 biased;
140 () = node.cancellation.cancelled() => Err("node shut down during target discovery".to_string()),
141 result = n0_future::time::timeout(timeout, outcome) => result
142 .map_err(|_| format!("discovery target {target_path:?} timed out"))
143 .and_then(|result| result),
144 };
145 if let Err(message) = outcome {
146 let _ = handle
147 .submit(CoreInput::DiscoveryTargetFailed { stream, message })
148 .await;
149 }
150 });
151 }
152
153 pub(crate) fn query_discovery_neighbor(
154 self: &Arc<Self>,
155 stream: StreamKey,
156 peer: String,
157 plan: DiscoverPlan,
158 handle: ProtocolCoreHandle,
159 ) {
160 let node = self.clone();
161 unb_runtime::RuntimeHandle::current().spawn(async move {
162 let (input, mut output) = mpsc::channel(64);
163 let cancel = node.cancellation.child_token();
164 let timeout = plan
165 .timeout_ms
166 .map(Duration::from_millis)
167 .unwrap_or(DISCOVER_TIMEOUT);
168 let leg = node.run_discovery_leg(&peer, plan, &input, &cancel, timeout);
169 tokio::pin!(leg);
170 loop {
171 tokio::select! {
172 biased;
173 event = output.recv() => match event {
174 Some(WalkInput::NeighborEvent { event, .. }) => {
175 if handle.submit(CoreInput::DiscoveryNeighborEvent {
176 stream: stream.clone(),
177 peer: peer.clone(),
178 event,
179 }).await.is_err() {
180 return;
181 }
182 }
183 _ => return,
184 },
185 completed = &mut leg => {
186 while let Ok(WalkInput::NeighborEvent { event, .. }) = output.try_recv() {
187 if handle.submit(CoreInput::DiscoveryNeighborEvent {
188 stream: stream.clone(),
189 peer: peer.clone(),
190 event,
191 }).await.is_err() {
192 return;
193 }
194 }
195 let input = if completed {
196 CoreInput::DiscoveryNeighborDone {
197 stream,
198 peer: peer.clone(),
199 }
200 } else {
201 CoreInput::DiscoveryNeighborTimeout {
202 stream,
203 peer: peer.clone(),
204 }
205 };
206 let _ = handle.submit(input).await;
207 return;
208 }
209 }
210 }
211 });
212 }
213
214 pub async fn discover_events(
215 self: &Arc<Self>,
216 detail: Detail,
217 scope: Scope,
218 ) -> Vec<DiscoverEvent> {
219 let plan = DiscoverPlan {
220 discover_id: String::new(),
221 detail,
222 scope,
223 hops: DEFAULT_HOPS,
224 visited: BTreeSet::new(),
225 timeout_ms: None,
226 mode: Mode::PartialOk,
227 };
228 let discover_id = plan.discover_id.clone();
229 let snapshot = self
230 .snapshot
231 .load()
232 .node_core
233 .catalog_snapshot(plan.detail.is_full());
234 let candidates = self.discovery_candidates(&plan.visited).await;
235 let mut walk = DiscoverWalk::start(snapshot, plan, candidates);
236 let (input_tx, mut input_rx) = mpsc::channel::<WalkInput>(64);
237 let cancel = self.cancellation.child_token();
238 let mut events = Vec::new();
239
240 loop {
241 while let Some(output) = walk.drain() {
242 match output {
243 WalkOutput::Emit(event) => events.push(event),
244 WalkOutput::AskNeighbor { peer, plan } => self.spawn_discovery_leg(
245 peer,
246 plan,
247 input_tx.clone(),
248 DISCOVER_TIMEOUT,
249 cancel.child_token(),
250 ),
251 WalkOutput::Finish => {
252 events.push(DiscoverEvent::Done {
253 discover_id: discover_id.clone(),
254 });
255 cancel.cancel();
256 return events;
257 }
258 }
259 }
260 tokio::select! {
261 biased;
262 () = cancel.cancelled() => return events,
263 input = input_rx.recv() => match input {
264 Some(input) => walk.handle(input),
265 None => return events,
266 }
267 }
268 }
269 }
270
271 pub(crate) async fn discovery_candidates(&self, visited: &BTreeSet<String>) -> Vec<String> {
272 self.peers
273 .read()
274 .await
275 .keys()
276 .filter(|name| !visited.contains(*name))
277 .cloned()
278 .collect()
279 }
280
281 fn spawn_discovery_leg(
282 self: &Arc<Self>,
283 peer: String,
284 plan: DiscoverPlan,
285 input: mpsc::Sender<WalkInput>,
286 timeout: Duration,
287 cancel: CancellationToken,
288 ) {
289 let node = self.clone();
290 unb_runtime::RuntimeHandle::current().spawn(async move {
291 let completed = node
292 .run_discovery_leg(&peer, plan, &input, &cancel, timeout)
293 .await;
294 let feedback = if completed {
295 WalkInput::NeighborDone { peer }
296 } else {
297 WalkInput::NeighborTimeout { peer }
298 };
299 let _ = input.send(feedback).await;
300 });
301 }
302
303 async fn run_discovery_leg(
304 &self,
305 peer: &str,
306 plan: DiscoverPlan,
307 input: &mpsc::Sender<WalkInput>,
308 cancel: &CancellationToken,
309 timeout: Duration,
310 ) -> bool {
311 let Ok(Ok(_permit)) =
312 n0_future::time::timeout(timeout, self.dispatch_slots.clone().acquire_owned()).await
313 else {
314 return false;
315 };
316 let Some(link) = self.peer(peer).await else {
317 return false;
318 };
319 let target_path = match TargetPath::discovery(peer) {
320 Ok(path) => path.to_string(),
321 Err(_) => return false,
322 };
323 let mut stream = match link
324 .wire
325 .client_session()
326 .discover(&target_path, plan)
327 .await
328 {
329 Ok(stream) => stream,
330 Err(_) => return false,
331 };
332 let deadline = n0_future::time::sleep(timeout);
333 tokio::pin!(deadline);
334 let completed = loop {
335 tokio::select! {
336 biased;
337 () = cancel.cancelled() => break false,
338 () = &mut deadline => break false,
339 msg = stream.next() => match msg {
340 Ok(Some(frame)) => match frame.kind {
341 unb_core::Kind::Event => {
342 if let Ok(event) =
343 serde_json::from_slice::<DiscoverEvent>(&frame.payload)
344 {
345 if input
346 .send(WalkInput::NeighborEvent {
347 peer: peer.to_string(),
348 event,
349 })
350 .await
351 .is_err()
352 {
353 break false;
354 }
355 }
356 }
357 unb_core::Kind::Response => break true,
358 _ => {}
359 },
360 Ok(None) => break false,
361 Err(_) => break false,
362 }
363 }
364 };
365 completed
366 }
367}