1use std::collections::BTreeSet;
2use std::sync::Arc;
3use std::time::Duration;
4
5use unb_core::{
6 CoreInput, Detail, DiscoverEvent, DiscoverPlan, DiscoverWalk, Mode, Scope, StreamKey,
7 WalkInput, WalkOutput, DEFAULT_HOPS,
8};
9use unb_runtime::{CancellationToken, ProtocolCoreHandle};
10use tokio::sync::mpsc;
11
12use crate::node::Node;
13
14const DISCOVER_TIMEOUT: Duration = Duration::from_secs(5);
15
16impl Node {
17 pub(crate) fn query_discovery_neighbor(
18 self: &Arc<Self>,
19 stream: StreamKey,
20 peer: String,
21 plan: DiscoverPlan,
22 handle: ProtocolCoreHandle,
23 ) {
24 let node = self.clone();
25 tokio::spawn(async move {
26 let (input, mut output) = mpsc::channel(64);
27 let cancel = node.cancellation.child_token();
28 let timeout = plan
29 .timeout_ms
30 .map(Duration::from_millis)
31 .unwrap_or(DISCOVER_TIMEOUT);
32 let leg = node.run_discovery_leg(&peer, plan, &input, &cancel, timeout);
33 tokio::pin!(leg);
34 loop {
35 tokio::select! {
36 biased;
37 event = output.recv() => match event {
38 Some(WalkInput::NeighborEvent { event, .. }) => {
39 if handle.submit(CoreInput::DiscoveryNeighborEvent {
40 stream: stream.clone(),
41 peer: peer.clone(),
42 event,
43 }).await.is_err() {
44 return;
45 }
46 }
47 _ => return,
48 },
49 completed = &mut leg => {
50 while let Ok(WalkInput::NeighborEvent { event, .. }) = output.try_recv() {
51 if handle.submit(CoreInput::DiscoveryNeighborEvent {
52 stream: stream.clone(),
53 peer: peer.clone(),
54 event,
55 }).await.is_err() {
56 return;
57 }
58 }
59 let input = if completed {
60 CoreInput::DiscoveryNeighborDone {
61 stream,
62 peer: peer.clone(),
63 }
64 } else {
65 CoreInput::DiscoveryNeighborTimeout {
66 stream,
67 peer: peer.clone(),
68 }
69 };
70 let _ = handle.submit(input).await;
71 return;
72 }
73 }
74 }
75 });
76 }
77
78 pub async fn discover_events(
79 self: &Arc<Self>,
80 detail: Detail,
81 scope: Scope,
82 ) -> Vec<DiscoverEvent> {
83 let plan = DiscoverPlan {
84 discover_id: String::new(),
85 detail,
86 scope,
87 hops: DEFAULT_HOPS,
88 visited: BTreeSet::new(),
89 timeout_ms: None,
90 mode: Mode::PartialOk,
91 };
92 let discover_id = plan.discover_id.clone();
93 let snapshot = self
94 .snapshot
95 .load()
96 .node_core
97 .catalog_snapshot(plan.detail.is_full());
98 let candidates = self.discovery_candidates(&plan.visited).await;
99 let mut walk = DiscoverWalk::start(snapshot, plan, candidates);
100 let (input_tx, mut input_rx) = mpsc::channel::<WalkInput>(64);
101 let cancel = self.cancellation.child_token();
102 let mut events = Vec::new();
103
104 loop {
105 while let Some(output) = walk.drain() {
106 match output {
107 WalkOutput::Emit(event) => events.push(event),
108 WalkOutput::AskNeighbor { peer, plan } => self.spawn_discovery_leg(
109 peer,
110 plan,
111 input_tx.clone(),
112 DISCOVER_TIMEOUT,
113 cancel.child_token(),
114 ),
115 WalkOutput::Finish => {
116 events.push(DiscoverEvent::Done {
117 discover_id: discover_id.clone(),
118 });
119 cancel.cancel();
120 return events;
121 }
122 }
123 }
124 tokio::select! {
125 biased;
126 () = cancel.cancelled() => return events,
127 input = input_rx.recv() => match input {
128 Some(input) => walk.handle(input),
129 None => return events,
130 }
131 }
132 }
133 }
134
135 pub(crate) async fn discovery_candidates(&self, visited: &BTreeSet<String>) -> Vec<String> {
136 self.peers
137 .read()
138 .await
139 .keys()
140 .filter(|name| !visited.contains(*name))
141 .cloned()
142 .collect()
143 }
144
145 fn spawn_discovery_leg(
146 self: &Arc<Self>,
147 peer: String,
148 plan: DiscoverPlan,
149 input: mpsc::Sender<WalkInput>,
150 timeout: Duration,
151 cancel: CancellationToken,
152 ) {
153 let node = self.clone();
154 tokio::spawn(async move {
155 let completed = node
156 .run_discovery_leg(&peer, plan, &input, &cancel, timeout)
157 .await;
158 let feedback = if completed {
159 WalkInput::NeighborDone { peer }
160 } else {
161 WalkInput::NeighborTimeout { peer }
162 };
163 let _ = input.send(feedback).await;
164 });
165 }
166
167 async fn run_discovery_leg(
168 &self,
169 peer: &str,
170 plan: DiscoverPlan,
171 input: &mpsc::Sender<WalkInput>,
172 cancel: &CancellationToken,
173 timeout: Duration,
174 ) -> bool {
175 let Ok(Ok(_permit)) =
176 tokio::time::timeout(timeout, self.dispatch_slots.clone().acquire_owned()).await
177 else {
178 return false;
179 };
180 let Some(link) = self.peer(peer).await else {
181 return false;
182 };
183 let mut stream = match link.wire.client_session().discover(plan).await {
184 Ok(stream) => stream,
185 Err(_) => return false,
186 };
187 let deadline = tokio::time::sleep(timeout);
188 tokio::pin!(deadline);
189 let completed = loop {
190 tokio::select! {
191 biased;
192 () = cancel.cancelled() => break false,
193 () = &mut deadline => break false,
194 msg = stream.next() => match msg {
195 Ok(Some(frame)) => match frame.kind {
196 unb_core::Kind::Event => {
197 if let Ok(event) =
198 serde_json::from_slice::<DiscoverEvent>(&frame.payload)
199 {
200 if input
201 .send(WalkInput::NeighborEvent {
202 peer: peer.to_string(),
203 event,
204 })
205 .await
206 .is_err()
207 {
208 break false;
209 }
210 }
211 }
212 unb_core::Kind::Response => break true,
213 _ => {}
214 },
215 Ok(None) => break false,
216 Err(_) => break false,
217 }
218 }
219 };
220 completed
221 }
222}