use std::collections::{BTreeSet, VecDeque};
use bytes::Bytes;
use serde::{Deserialize, Serialize};
use serde_json::Value;
use crate::{CoreError, DEFAULT_HOPS};
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum Detail {
#[default]
Index,
Full,
}
impl Detail {
pub fn is_full(self) -> bool {
matches!(self, Detail::Full)
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum Scope {
Local,
Neighbors,
#[default]
Reachable,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum Mode {
#[default]
PartialOk,
Strict,
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(tag = "type", rename_all = "snake_case")]
pub enum DiscoverEvent {
NodeCatalog {
node: String,
instance_id: String,
revision: u64,
fingerprint: String,
subjects: Vec<Value>,
},
Edge {
from: String,
to: String,
},
Warning {
node: String,
message: String,
},
Done {
discover_id: String,
},
}
#[derive(Debug, Clone)]
pub struct NodeCatalogSnapshot {
pub node: String,
pub instance_id: String,
pub revision: u64,
pub fingerprint: String,
pub subjects: Vec<Value>,
}
impl NodeCatalogSnapshot {
fn into_event(self) -> DiscoverEvent {
DiscoverEvent::NodeCatalog {
node: self.node,
instance_id: self.instance_id,
revision: self.revision,
fingerprint: self.fingerprint,
subjects: self.subjects,
}
}
}
fn default_hops() -> u8 {
DEFAULT_HOPS
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct DiscoverPlan {
#[serde(default)]
pub discover_id: String,
#[serde(default)]
pub detail: Detail,
#[serde(default)]
pub scope: Scope,
#[serde(default = "default_hops")]
pub hops: u8,
#[serde(default)]
pub visited: BTreeSet<String>,
#[serde(default)]
pub timeout_ms: Option<u64>,
#[serde(default)]
pub mode: Mode,
}
impl DiscoverPlan {
pub fn decode(payload: &Bytes) -> Result<DiscoverPlan, CoreError> {
let value: Value = if payload.is_empty() {
Value::Object(serde_json::Map::new())
} else {
serde_json::from_slice(payload)
.map_err(|error| CoreError::Malformed(error.to_string()))?
};
serde_json::from_value(value).map_err(|error| CoreError::Malformed(error.to_string()))
}
}
#[derive(Debug, Clone, PartialEq)]
pub enum WalkInput {
NeighborEvent { peer: String, event: DiscoverEvent },
NeighborDone { peer: String },
NeighborTimeout { peer: String },
}
#[derive(Debug, Clone, PartialEq)]
pub enum WalkOutput {
Emit(DiscoverEvent),
AskNeighbor { peer: String, plan: DiscoverPlan },
Finish,
}
pub struct DiscoverWalk {
discover_id: String,
mode: Mode,
visited: BTreeSet<String>,
outstanding: BTreeSet<String>,
outputs: VecDeque<WalkOutput>,
strict_failure: Option<String>,
finished: bool,
}
impl DiscoverWalk {
pub fn start(
node: NodeCatalogSnapshot,
plan: DiscoverPlan,
candidates: Vec<String>,
) -> DiscoverWalk {
let self_node = node.node.clone();
let mut visited = plan.visited.clone();
visited.insert(self_node.clone());
let effective_hops = match plan.scope {
Scope::Local => 0,
Scope::Neighbors => plan.hops.min(1),
Scope::Reachable => plan.hops,
};
let mut outputs = VecDeque::new();
outputs.push_back(WalkOutput::Emit(node.into_event()));
let mut outstanding = BTreeSet::new();
if effective_hops > 0 {
let to_ask: Vec<String> = candidates
.into_iter()
.filter(|peer| !visited.contains(peer))
.collect();
let child_visited: BTreeSet<String> =
visited.iter().chain(to_ask.iter()).cloned().collect();
for peer in to_ask {
outputs.push_back(WalkOutput::Emit(DiscoverEvent::Edge {
from: self_node.clone(),
to: peer.clone(),
}));
let child_plan = DiscoverPlan {
discover_id: plan.discover_id.clone(),
detail: plan.detail,
scope: Scope::Reachable,
hops: effective_hops - 1,
visited: child_visited.clone(),
timeout_ms: plan.timeout_ms,
mode: plan.mode,
};
outstanding.insert(peer.clone());
visited.insert(peer.clone());
outputs.push_back(WalkOutput::AskNeighbor {
peer,
plan: child_plan,
});
}
}
let finished = outstanding.is_empty();
if finished {
outputs.push_back(WalkOutput::Finish);
}
DiscoverWalk {
discover_id: plan.discover_id,
mode: plan.mode,
visited,
outstanding,
outputs,
strict_failure: None,
finished,
}
}
pub fn discover_id(&self) -> &str {
&self.discover_id
}
pub fn handle(&mut self, input: WalkInput) {
if self.finished {
return;
}
match input {
WalkInput::NeighborEvent { event, .. } => {
self.outputs.push_back(WalkOutput::Emit(event));
}
WalkInput::NeighborDone { peer } => {
self.outstanding.remove(&peer);
self.finish_if_drained();
}
WalkInput::NeighborTimeout { peer } => {
if !self.outstanding.remove(&peer) {
return;
}
match self.mode {
Mode::PartialOk => {
self.outputs
.push_back(WalkOutput::Emit(DiscoverEvent::Warning {
node: peer,
message: "discovery timed out".to_string(),
}));
self.finish_if_drained();
}
Mode::Strict => {
if self.strict_failure.is_none() {
self.strict_failure = Some(peer);
}
self.finish_now();
}
}
}
}
}
pub fn drain(&mut self) -> Option<WalkOutput> {
self.outputs.pop_front()
}
pub fn has_output(&self) -> bool {
!self.outputs.is_empty()
}
pub fn strict_failure(&self) -> Option<&str> {
self.strict_failure.as_deref()
}
pub fn visited(&self) -> &BTreeSet<String> {
&self.visited
}
fn finish_if_drained(&mut self) {
if !self.finished && self.outstanding.is_empty() {
self.finish_now();
}
}
fn finish_now(&mut self) {
if !self.finished {
self.finished = true;
self.outputs.push_back(WalkOutput::Finish);
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use serde_json::json;
fn snapshot(node: &str, subjects: &[&str]) -> NodeCatalogSnapshot {
NodeCatalogSnapshot {
node: node.to_string(),
instance_id: node.to_string(),
revision: 1,
fingerprint: "fp".to_string(),
subjects: subjects.iter().map(|s| json!({ "subject": s })).collect(),
}
}
fn plan(scope: Scope, hops: u8) -> DiscoverPlan {
DiscoverPlan {
discover_id: "d1".to_string(),
detail: Detail::Index,
scope,
hops,
visited: BTreeSet::new(),
timeout_ms: None,
mode: Mode::PartialOk,
}
}
fn drain_all(walk: &mut DiscoverWalk) -> Vec<WalkOutput> {
let mut out = Vec::new();
while let Some(output) = walk.drain() {
out.push(output);
}
out
}
#[test]
fn a_local_scope_walk_emits_one_node_catalog_then_finishes() {
let mut walk = DiscoverWalk::start(
snapshot("hub", &["chess"]),
plan(Scope::Local, 8),
vec!["a".into(), "b".into()],
);
let outputs = drain_all(&mut walk);
assert_eq!(outputs.len(), 2);
assert!(matches!(
outputs[0],
WalkOutput::Emit(DiscoverEvent::NodeCatalog { .. })
));
assert_eq!(outputs[1], WalkOutput::Finish);
}
#[test]
fn zero_hops_stops_descent_even_with_candidates() {
let mut walk = DiscoverWalk::start(
snapshot("hub", &[]),
plan(Scope::Reachable, 0),
vec!["a".into()],
);
let outputs = drain_all(&mut walk);
assert_eq!(outputs.len(), 2);
assert!(matches!(
outputs[0],
WalkOutput::Emit(DiscoverEvent::NodeCatalog { .. })
));
assert_eq!(outputs[1], WalkOutput::Finish);
}
#[test]
fn a_fan_out_walk_asks_each_candidate_and_finishes_after_all_done() {
let mut walk = DiscoverWalk::start(
snapshot("hub", &[]),
plan(Scope::Reachable, 8),
vec!["a".into(), "b".into()],
);
let opened = drain_all(&mut walk);
assert!(matches!(
opened[0],
WalkOutput::Emit(DiscoverEvent::NodeCatalog { .. })
));
let asks: Vec<&WalkOutput> = opened
.iter()
.filter(|o| matches!(o, WalkOutput::AskNeighbor { .. }))
.collect();
assert_eq!(asks.len(), 2, "both candidates asked");
let edges = opened
.iter()
.filter(|o| matches!(o, WalkOutput::Emit(DiscoverEvent::Edge { .. })))
.count();
assert_eq!(edges, 2, "one edge per asked candidate");
assert!(!opened.iter().any(|o| matches!(o, WalkOutput::Finish)));
if let WalkOutput::AskNeighbor { plan, .. } = asks[0] {
assert!(plan.visited.contains("hub"));
assert!(plan.visited.contains("a"));
assert!(plan.visited.contains("b"));
assert_eq!(plan.hops, 7, "hops decrements per edge crossed");
}
walk.handle(WalkInput::NeighborEvent {
peer: "a".into(),
event: DiscoverEvent::NodeCatalog {
node: "a".into(),
instance_id: "a".into(),
revision: 1,
fingerprint: "fp".into(),
subjects: vec![],
},
});
walk.handle(WalkInput::NeighborDone { peer: "a".into() });
let mid = drain_all(&mut walk);
assert!(matches!(
mid[0],
WalkOutput::Emit(DiscoverEvent::NodeCatalog { .. })
));
assert!(
!mid.iter().any(|o| matches!(o, WalkOutput::Finish)),
"b still outstanding"
);
walk.handle(WalkInput::NeighborDone { peer: "b".into() });
assert_eq!(walk.drain(), Some(WalkOutput::Finish));
}
#[test]
fn a_candidate_already_in_visited_is_not_asked() {
let mut plan = plan(Scope::Reachable, 8);
plan.visited.insert("a".into());
let mut walk =
DiscoverWalk::start(snapshot("hub", &[]), plan, vec!["a".into(), "b".into()]);
let outputs = drain_all(&mut walk);
let asked: Vec<String> = outputs
.iter()
.filter_map(|o| match o {
WalkOutput::AskNeighbor { peer, .. } => Some(peer.clone()),
_ => None,
})
.collect();
assert_eq!(
asked,
vec!["b".to_string()],
"the visited candidate is skipped"
);
}
#[test]
fn duplicate_subjects_from_two_nodes_both_emit_no_dedup() {
let mut walk = DiscoverWalk::start(
snapshot("hub", &[]),
plan(Scope::Reachable, 8),
vec!["a".into()],
);
let _ = drain_all(&mut walk);
for node in ["a", "phantom"] {
walk.handle(WalkInput::NeighborEvent {
peer: "a".into(),
event: DiscoverEvent::NodeCatalog {
node: node.into(),
instance_id: node.into(),
revision: 1,
fingerprint: "fp".into(),
subjects: vec![json!({ "subject": "todo" })],
},
});
}
let emitted = drain_all(&mut walk);
let todo_events = emitted
.iter()
.filter(|o| matches!(o, WalkOutput::Emit(DiscoverEvent::NodeCatalog { .. })))
.count();
assert_eq!(
todo_events, 2,
"node-scoped: the same subject under two nodes emits twice"
);
}
#[test]
fn partial_ok_timeout_emits_a_warning_and_still_finishes() {
let mut walk = DiscoverWalk::start(
snapshot("hub", &[]),
plan(Scope::Reachable, 8),
vec!["slow".into()],
);
let _ = drain_all(&mut walk);
walk.handle(WalkInput::NeighborTimeout {
peer: "slow".into(),
});
let outputs = drain_all(&mut walk);
assert!(matches!(
outputs[0],
WalkOutput::Emit(DiscoverEvent::Warning { .. })
));
assert_eq!(outputs[1], WalkOutput::Finish);
assert!(walk.strict_failure().is_none());
}
#[test]
fn strict_timeout_records_a_failure_and_finishes_immediately() {
let mut strict = plan(Scope::Reachable, 8);
strict.mode = Mode::Strict;
let mut walk =
DiscoverWalk::start(snapshot("hub", &[]), strict, vec!["a".into(), "b".into()]);
let _ = drain_all(&mut walk);
walk.handle(WalkInput::NeighborTimeout { peer: "a".into() });
assert_eq!(
walk.drain(),
Some(WalkOutput::Finish),
"strict finishes without waiting for b"
);
assert_eq!(walk.strict_failure(), Some("a"));
walk.handle(WalkInput::NeighborDone { peer: "b".into() });
assert_eq!(walk.drain(), None);
}
#[test]
fn decode_defaults_an_empty_payload_and_parses_fields() {
let empty = DiscoverPlan::decode(&Bytes::new()).unwrap();
assert_eq!(empty.detail, Detail::Index);
assert_eq!(empty.scope, Scope::Reachable);
assert_eq!(empty.mode, Mode::PartialOk);
assert_eq!(empty.hops, DEFAULT_HOPS);
let full = DiscoverPlan::decode(&Bytes::from_static(
br#"{"discover_id":"d9","detail":"full","scope":"local","mode":"strict","hops":3}"#,
))
.unwrap();
assert_eq!(full.discover_id, "d9");
assert!(full.detail.is_full());
assert_eq!(full.scope, Scope::Local);
assert_eq!(full.mode, Mode::Strict);
assert_eq!(full.hops, 3);
}
}