1use std::collections::{BTreeSet, VecDeque};
2
3use bytes::Bytes;
4use serde::{Deserialize, Serialize};
5use serde_json::Value;
6
7use crate::{CoreError, DEFAULT_HOPS};
8
9#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Serialize, Deserialize)]
10#[serde(rename_all = "snake_case")]
11pub enum Detail {
12 #[default]
13 Index,
14 Full,
15}
16
17impl Detail {
18 pub fn is_full(self) -> bool {
19 matches!(self, Detail::Full)
20 }
21}
22
23#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Serialize, Deserialize)]
24#[serde(rename_all = "snake_case")]
25pub enum Scope {
26 Local,
27 Neighbors,
28 #[default]
29 Reachable,
30}
31
32#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Serialize, Deserialize)]
33#[serde(rename_all = "snake_case")]
34pub enum Mode {
35 #[default]
36 PartialOk,
37 Strict,
38}
39
40#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
41#[serde(tag = "type", rename_all = "snake_case")]
42pub enum DiscoverEvent {
43 NodeCatalog {
44 node: String,
45 instance_id: String,
46 revision: u64,
47 fingerprint: String,
48 subjects: Vec<Value>,
49 },
50 Edge {
51 from: String,
52 to: String,
53 },
54 Warning {
55 node: String,
56 message: String,
57 },
58 Done {
59 discover_id: String,
60 },
61}
62
63#[derive(Debug, Clone)]
64pub struct NodeCatalogSnapshot {
65 pub node: String,
66 pub instance_id: String,
67 pub revision: u64,
68 pub fingerprint: String,
69 pub subjects: Vec<Value>,
70}
71
72impl NodeCatalogSnapshot {
73 fn into_event(self) -> DiscoverEvent {
74 DiscoverEvent::NodeCatalog {
75 node: self.node,
76 instance_id: self.instance_id,
77 revision: self.revision,
78 fingerprint: self.fingerprint,
79 subjects: self.subjects,
80 }
81 }
82}
83
84fn default_hops() -> u8 {
85 DEFAULT_HOPS
86}
87
88#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
89pub struct DiscoverPlan {
90 #[serde(default)]
91 pub discover_id: String,
92 #[serde(default)]
93 pub detail: Detail,
94 #[serde(default)]
95 pub scope: Scope,
96 #[serde(default = "default_hops")]
97 pub hops: u8,
98 #[serde(default)]
99 pub visited: BTreeSet<String>,
100 #[serde(default)]
101 pub timeout_ms: Option<u64>,
102 #[serde(default)]
103 pub mode: Mode,
104}
105
106impl DiscoverPlan {
107 pub fn decode(payload: &Bytes) -> Result<DiscoverPlan, CoreError> {
108 let value: Value = if payload.is_empty() {
109 Value::Object(serde_json::Map::new())
110 } else {
111 serde_json::from_slice(payload)
112 .map_err(|error| CoreError::Malformed(error.to_string()))?
113 };
114 serde_json::from_value(value).map_err(|error| CoreError::Malformed(error.to_string()))
115 }
116}
117
118#[derive(Debug, Clone, PartialEq)]
119pub enum WalkInput {
120 NeighborEvent { peer: String, event: DiscoverEvent },
121 NeighborDone { peer: String },
122 NeighborTimeout { peer: String },
123}
124
125#[derive(Debug, Clone, PartialEq)]
126pub enum WalkOutput {
127 Emit(DiscoverEvent),
128 AskNeighbor { peer: String, plan: DiscoverPlan },
129 Finish,
130}
131
132pub struct DiscoverWalk {
133 discover_id: String,
134 mode: Mode,
135 visited: BTreeSet<String>,
136 outstanding: BTreeSet<String>,
137 outputs: VecDeque<WalkOutput>,
138 strict_failure: Option<String>,
139 finished: bool,
140}
141
142impl DiscoverWalk {
143 pub fn start(
144 node: NodeCatalogSnapshot,
145 plan: DiscoverPlan,
146 candidates: Vec<String>,
147 ) -> DiscoverWalk {
148 let self_node = node.node.clone();
149 let mut visited = plan.visited.clone();
150 visited.insert(self_node.clone());
151
152 let effective_hops = match plan.scope {
153 Scope::Local => 0,
154 Scope::Neighbors => plan.hops.min(1),
155 Scope::Reachable => plan.hops,
156 };
157
158 let mut outputs = VecDeque::new();
159 outputs.push_back(WalkOutput::Emit(node.into_event()));
160
161 let mut outstanding = BTreeSet::new();
162 if effective_hops > 0 {
163 let to_ask: Vec<String> = candidates
164 .into_iter()
165 .filter(|peer| !visited.contains(peer))
166 .collect();
167 let child_visited: BTreeSet<String> =
168 visited.iter().chain(to_ask.iter()).cloned().collect();
169 for peer in to_ask {
170 outputs.push_back(WalkOutput::Emit(DiscoverEvent::Edge {
171 from: self_node.clone(),
172 to: peer.clone(),
173 }));
174 let child_plan = DiscoverPlan {
175 discover_id: plan.discover_id.clone(),
176 detail: plan.detail,
177 scope: Scope::Reachable,
178 hops: effective_hops - 1,
179 visited: child_visited.clone(),
180 timeout_ms: plan.timeout_ms,
181 mode: plan.mode,
182 };
183 outstanding.insert(peer.clone());
184 visited.insert(peer.clone());
185 outputs.push_back(WalkOutput::AskNeighbor {
186 peer,
187 plan: child_plan,
188 });
189 }
190 }
191
192 let finished = outstanding.is_empty();
193 if finished {
194 outputs.push_back(WalkOutput::Finish);
195 }
196
197 DiscoverWalk {
198 discover_id: plan.discover_id,
199 mode: plan.mode,
200 visited,
201 outstanding,
202 outputs,
203 strict_failure: None,
204 finished,
205 }
206 }
207
208 pub fn discover_id(&self) -> &str {
209 &self.discover_id
210 }
211
212 pub fn handle(&mut self, input: WalkInput) {
213 if self.finished {
214 return;
215 }
216 match input {
217 WalkInput::NeighborEvent { event, .. } => {
218 self.outputs.push_back(WalkOutput::Emit(event));
219 }
220 WalkInput::NeighborDone { peer } => {
221 self.outstanding.remove(&peer);
222 self.finish_if_drained();
223 }
224 WalkInput::NeighborTimeout { peer } => {
225 if !self.outstanding.remove(&peer) {
226 return;
227 }
228 match self.mode {
229 Mode::PartialOk => {
230 self.outputs
231 .push_back(WalkOutput::Emit(DiscoverEvent::Warning {
232 node: peer,
233 message: "discovery timed out".to_string(),
234 }));
235 self.finish_if_drained();
236 }
237 Mode::Strict => {
238 if self.strict_failure.is_none() {
239 self.strict_failure = Some(peer);
240 }
241 self.finish_now();
242 }
243 }
244 }
245 }
246 }
247
248 pub fn drain(&mut self) -> Option<WalkOutput> {
249 self.outputs.pop_front()
250 }
251
252 pub fn has_output(&self) -> bool {
253 !self.outputs.is_empty()
254 }
255
256 pub fn strict_failure(&self) -> Option<&str> {
257 self.strict_failure.as_deref()
258 }
259
260 pub fn visited(&self) -> &BTreeSet<String> {
261 &self.visited
262 }
263
264 fn finish_if_drained(&mut self) {
265 if !self.finished && self.outstanding.is_empty() {
266 self.finish_now();
267 }
268 }
269
270 fn finish_now(&mut self) {
271 if !self.finished {
272 self.finished = true;
273 self.outputs.push_back(WalkOutput::Finish);
274 }
275 }
276}
277
278#[cfg(test)]
279mod tests {
280 use super::*;
281 use serde_json::json;
282
283 fn snapshot(node: &str, subjects: &[&str]) -> NodeCatalogSnapshot {
284 NodeCatalogSnapshot {
285 node: node.to_string(),
286 instance_id: node.to_string(),
287 revision: 1,
288 fingerprint: "fp".to_string(),
289 subjects: subjects.iter().map(|s| json!({ "subject": s })).collect(),
290 }
291 }
292
293 fn plan(scope: Scope, hops: u8) -> DiscoverPlan {
294 DiscoverPlan {
295 discover_id: "d1".to_string(),
296 detail: Detail::Index,
297 scope,
298 hops,
299 visited: BTreeSet::new(),
300 timeout_ms: None,
301 mode: Mode::PartialOk,
302 }
303 }
304
305 fn drain_all(walk: &mut DiscoverWalk) -> Vec<WalkOutput> {
306 let mut out = Vec::new();
307 while let Some(output) = walk.drain() {
308 out.push(output);
309 }
310 out
311 }
312
313 #[test]
314 fn a_local_scope_walk_emits_one_node_catalog_then_finishes() {
315 let mut walk = DiscoverWalk::start(
316 snapshot("hub", &["chess"]),
317 plan(Scope::Local, 8),
318 vec!["a".into(), "b".into()],
319 );
320 let outputs = drain_all(&mut walk);
321 assert_eq!(outputs.len(), 2);
322 assert!(matches!(
323 outputs[0],
324 WalkOutput::Emit(DiscoverEvent::NodeCatalog { .. })
325 ));
326 assert_eq!(outputs[1], WalkOutput::Finish);
327 }
328
329 #[test]
330 fn zero_hops_stops_descent_even_with_candidates() {
331 let mut walk = DiscoverWalk::start(
332 snapshot("hub", &[]),
333 plan(Scope::Reachable, 0),
334 vec!["a".into()],
335 );
336 let outputs = drain_all(&mut walk);
337 assert_eq!(outputs.len(), 2);
338 assert!(matches!(
339 outputs[0],
340 WalkOutput::Emit(DiscoverEvent::NodeCatalog { .. })
341 ));
342 assert_eq!(outputs[1], WalkOutput::Finish);
343 }
344
345 #[test]
346 fn a_fan_out_walk_asks_each_candidate_and_finishes_after_all_done() {
347 let mut walk = DiscoverWalk::start(
348 snapshot("hub", &[]),
349 plan(Scope::Reachable, 8),
350 vec!["a".into(), "b".into()],
351 );
352 let opened = drain_all(&mut walk);
353 assert!(matches!(
354 opened[0],
355 WalkOutput::Emit(DiscoverEvent::NodeCatalog { .. })
356 ));
357 let asks: Vec<&WalkOutput> = opened
358 .iter()
359 .filter(|o| matches!(o, WalkOutput::AskNeighbor { .. }))
360 .collect();
361 assert_eq!(asks.len(), 2, "both candidates asked");
362 let edges = opened
363 .iter()
364 .filter(|o| matches!(o, WalkOutput::Emit(DiscoverEvent::Edge { .. })))
365 .count();
366 assert_eq!(edges, 2, "one edge per asked candidate");
367 assert!(!opened.iter().any(|o| matches!(o, WalkOutput::Finish)));
368
369 if let WalkOutput::AskNeighbor { plan, .. } = asks[0] {
370 assert!(plan.visited.contains("hub"));
371 assert!(plan.visited.contains("a"));
372 assert!(plan.visited.contains("b"));
373 assert_eq!(plan.hops, 7, "hops decrements per edge crossed");
374 }
375
376 walk.handle(WalkInput::NeighborEvent {
377 peer: "a".into(),
378 event: DiscoverEvent::NodeCatalog {
379 node: "a".into(),
380 instance_id: "a".into(),
381 revision: 1,
382 fingerprint: "fp".into(),
383 subjects: vec![],
384 },
385 });
386 walk.handle(WalkInput::NeighborDone { peer: "a".into() });
387 let mid = drain_all(&mut walk);
388 assert!(matches!(
389 mid[0],
390 WalkOutput::Emit(DiscoverEvent::NodeCatalog { .. })
391 ));
392 assert!(
393 !mid.iter().any(|o| matches!(o, WalkOutput::Finish)),
394 "b still outstanding"
395 );
396
397 walk.handle(WalkInput::NeighborDone { peer: "b".into() });
398 assert_eq!(walk.drain(), Some(WalkOutput::Finish));
399 }
400
401 #[test]
402 fn a_candidate_already_in_visited_is_not_asked() {
403 let mut plan = plan(Scope::Reachable, 8);
404 plan.visited.insert("a".into());
405 let mut walk =
406 DiscoverWalk::start(snapshot("hub", &[]), plan, vec!["a".into(), "b".into()]);
407 let outputs = drain_all(&mut walk);
408 let asked: Vec<String> = outputs
409 .iter()
410 .filter_map(|o| match o {
411 WalkOutput::AskNeighbor { peer, .. } => Some(peer.clone()),
412 _ => None,
413 })
414 .collect();
415 assert_eq!(
416 asked,
417 vec!["b".to_string()],
418 "the visited candidate is skipped"
419 );
420 }
421
422 #[test]
423 fn duplicate_subjects_from_two_nodes_both_emit_no_dedup() {
424 let mut walk = DiscoverWalk::start(
425 snapshot("hub", &[]),
426 plan(Scope::Reachable, 8),
427 vec!["a".into()],
428 );
429 let _ = drain_all(&mut walk);
430 for node in ["a", "phantom"] {
431 walk.handle(WalkInput::NeighborEvent {
432 peer: "a".into(),
433 event: DiscoverEvent::NodeCatalog {
434 node: node.into(),
435 instance_id: node.into(),
436 revision: 1,
437 fingerprint: "fp".into(),
438 subjects: vec![json!({ "subject": "todo" })],
439 },
440 });
441 }
442 let emitted = drain_all(&mut walk);
443 let todo_events = emitted
444 .iter()
445 .filter(|o| matches!(o, WalkOutput::Emit(DiscoverEvent::NodeCatalog { .. })))
446 .count();
447 assert_eq!(
448 todo_events, 2,
449 "node-scoped: the same subject under two nodes emits twice"
450 );
451 }
452
453 #[test]
454 fn partial_ok_timeout_emits_a_warning_and_still_finishes() {
455 let mut walk = DiscoverWalk::start(
456 snapshot("hub", &[]),
457 plan(Scope::Reachable, 8),
458 vec!["slow".into()],
459 );
460 let _ = drain_all(&mut walk);
461 walk.handle(WalkInput::NeighborTimeout {
462 peer: "slow".into(),
463 });
464 let outputs = drain_all(&mut walk);
465 assert!(matches!(
466 outputs[0],
467 WalkOutput::Emit(DiscoverEvent::Warning { .. })
468 ));
469 assert_eq!(outputs[1], WalkOutput::Finish);
470 assert!(walk.strict_failure().is_none());
471 }
472
473 #[test]
474 fn strict_timeout_records_a_failure_and_finishes_immediately() {
475 let mut strict = plan(Scope::Reachable, 8);
476 strict.mode = Mode::Strict;
477 let mut walk =
478 DiscoverWalk::start(snapshot("hub", &[]), strict, vec!["a".into(), "b".into()]);
479 let _ = drain_all(&mut walk);
480 walk.handle(WalkInput::NeighborTimeout { peer: "a".into() });
481 assert_eq!(
482 walk.drain(),
483 Some(WalkOutput::Finish),
484 "strict finishes without waiting for b"
485 );
486 assert_eq!(walk.strict_failure(), Some("a"));
487 walk.handle(WalkInput::NeighborDone { peer: "b".into() });
488 assert_eq!(walk.drain(), None);
489 }
490
491 #[test]
492 fn decode_defaults_an_empty_payload_and_parses_fields() {
493 let empty = DiscoverPlan::decode(&Bytes::new()).unwrap();
494 assert_eq!(empty.detail, Detail::Index);
495 assert_eq!(empty.scope, Scope::Reachable);
496 assert_eq!(empty.mode, Mode::PartialOk);
497 assert_eq!(empty.hops, DEFAULT_HOPS);
498
499 let full = DiscoverPlan::decode(&Bytes::from_static(
500 br#"{"discover_id":"d9","detail":"full","scope":"local","mode":"strict","hops":3}"#,
501 ))
502 .unwrap();
503 assert_eq!(full.discover_id, "d9");
504 assert!(full.detail.is_full());
505 assert_eq!(full.scope, Scope::Local);
506 assert_eq!(full.mode, Mode::Strict);
507 assert_eq!(full.hops, 3);
508 }
509}