use std::collections::{HashMap, VecDeque};
use crate::communication::{Channel, Message, MessageProcessingError};
use crate::node_set::NodeSet;
use crate::topology::{TopologyError, TopologyTree};
use thiserror::Error;
#[derive(Error, Debug)]
pub enum RouteResolvingError {
#[error("Invalid root or gateway node: {0}")]
InvalidRoot(String),
#[error("No route available to {0}")]
NoRoute(String),
#[error("Invalid destination: {0}, host is unreachable")]
Unreachable(String),
#[error("Invalid resolution request: {0} -> {1}")]
SameSourceDest(String, String),
#[error("Topology error: {0}")]
Topology(#[from] TopologyError),
#[error("NodeSet error: {0}")]
NodeSet(#[from] crate::node_set::NodeSetError),
}
#[derive(Debug, Clone)]
pub struct RouteEntry {
network: NodeSet,
nexthops: NodeSet,
}
#[derive(Debug)]
pub struct PropagationTreeRouter {
root: String,
fanout: usize,
nodes_fanin: HashMap<String, usize>,
table: Vec<RouteEntry>,
unreachable_hosts: NodeSet,
}
impl PropagationTreeRouter {
pub fn new(
root: &str,
topology: &TopologyTree,
fanout: usize,
) -> Result<Self, RouteResolvingError> {
let mut router = Self {
root: root.to_string(),
fanout,
nodes_fanin: HashMap::new(),
table: Vec::new(),
unreachable_hosts: NodeSet::new(),
};
router.table_generate(root, topology)?;
Ok(router)
}
fn table_generate(
&mut self,
root: &str,
topology: &TopologyTree,
) -> Result<(), RouteResolvingError> {
let _root_group = topology
.find_nodegroup(root)
.map_err(|_| RouteResolvingError::InvalidRoot(root.to_string()))?;
self.table.clear();
let root_tng = topology
.root()
.ok_or_else(|| RouteResolvingError::InvalidRoot(root.to_string()))?;
let subtree_root = Self::find_group_containing(root_tng, root)
.ok_or_else(|| RouteResolvingError::InvalidRoot(root.to_string()))?;
for child in subtree_root.children() {
let mut network = NodeSet::new();
Self::collect_reachable(child, &mut network);
self.table.push(RouteEntry {
network,
nexthops: child.nodeset.clone(),
});
}
Ok(())
}
fn find_group_containing<'a>(
group: &'a crate::topology::TopologyNodeGroup,
node: &str,
) -> Option<&'a crate::topology::TopologyNodeGroup> {
if group.nodeset.contains(node) {
return Some(group);
}
for child in group.children() {
if let Some(found) = Self::find_group_containing(child, node) {
return Some(found);
}
}
None
}
fn collect_reachable(group: &crate::topology::TopologyNodeGroup, dest: &mut NodeSet) {
for child in group.children() {
dest.update(&child.nodeset);
Self::collect_reachable(child, dest);
}
}
pub fn dispatch(&self, dst: &NodeSet) -> Vec<(NodeSet, NodeSet)> {
let mut remaining = dst.clone();
let mut results = Vec::new();
for entry in &self.table {
let dst_inter = entry.network.intersection(&remaining);
if !dst_inter.is_empty() {
remaining.difference_update(&dst_inter);
for host_name in dst_inter.iter() {
if let Ok(hop) = self.next_hop_str(&host_name) {
let mut hop_ns = NodeSet::new();
let _ = hop_ns.update_str(&hop);
let mut host_ns = NodeSet::new();
let _ = host_ns.update_str(&host_name);
results.push((hop_ns, host_ns));
}
}
}
}
if !remaining.is_empty() {
results.push((remaining.clone(), remaining));
}
results
}
pub fn next_hop(&mut self, dst: &NodeSet) -> Result<NodeSet, RouteResolvingError> {
let dst_str = dst.to_string();
if !self.unreachable_hosts.intersection(dst).is_empty() {
return Err(RouteResolvingError::Unreachable(dst_str));
}
if self.root == dst_str {
return Err(RouteResolvingError::SameSourceDest(
self.root.clone(),
dst_str,
));
}
for entry in &self.table {
if !entry.network.intersection(dst).is_empty() {
let res = self.best_next_hop(&entry.nexthops)?;
let res_str = res.to_string();
*self.nodes_fanin.entry(res_str).or_insert(0) += dst.len();
return Ok(res);
}
if !entry.nexthops.intersection(dst).is_empty() {
return Ok(dst.clone());
}
}
Err(RouteResolvingError::NoRoute(dst_str))
}
fn next_hop_str(&self, dst: &str) -> Result<String, RouteResolvingError> {
if self.unreachable_hosts.contains(dst) {
return Err(RouteResolvingError::Unreachable(dst.to_string()));
}
if self.root == dst {
return Err(RouteResolvingError::SameSourceDest(
self.root.clone(),
dst.to_string(),
));
}
let dst_ns = NodeSet::parse(dst)?;
for entry in &self.table {
if !entry.network.intersection(&dst_ns).is_empty() {
let res = self.best_next_hop(&entry.nexthops)?;
return Ok(res.to_string());
}
if !entry.nexthops.intersection(&dst_ns).is_empty() {
return Ok(dst.to_string());
}
}
Err(RouteResolvingError::NoRoute(dst.to_string()))
}
pub fn mark_unreachable(&mut self, dst: &str) {
let _ = self.unreachable_hosts.update_str(dst);
}
fn best_next_hop(&self, candidates: &NodeSet) -> Result<NodeSet, RouteResolvingError> {
let available = candidates.difference(&self.unreachable_hosts);
if available.is_empty() {
return Err(RouteResolvingError::NoRoute(
"all candidates unreachable".into(),
));
}
let mut best: Option<String> = None;
let mut best_connections = usize::MAX;
for host in available.iter() {
let connections = self.nodes_fanin.get(&host).copied().unwrap_or(0);
if connections < best_connections {
best_connections = connections;
best = Some(host);
}
}
match best {
Some(host) => {
let mut ns = NodeSet::new();
let _ = ns.update_str(&host);
Ok(ns)
}
None => Err(RouteResolvingError::NoRoute(
"no candidate available".into(),
)),
}
}
pub fn root(&self) -> &str {
&self.root
}
pub fn fanout(&self) -> usize {
self.fanout
}
pub fn table(&self) -> &[RouteEntry] {
&self.table
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ChannelState {
Init,
Config,
Control,
Gather,
Closed,
}
#[derive(Debug)]
pub struct PropagationChannel {
pub channel: Channel,
pub gateway: String,
pub state: ChannelState,
send_queue: VecDeque<Message>,
pub last_rc: Option<i32>,
write_history: VecDeque<(u64, String, usize)>,
pub results: Vec<PropagationResult>,
}
#[derive(Debug, Clone)]
pub enum PropagationResult {
StdOut { node: String, data: Vec<u8> },
StdErr { node: String, data: Vec<u8> },
Retcode { node: String, rc: i32 },
Timeout { node: String },
Routing {
event: String,
gateway: String,
targets: String,
},
}
impl PropagationChannel {
pub fn new(gateway: &str) -> Self {
Self {
channel: Channel::new(true), gateway: gateway.to_string(),
state: ChannelState::Init,
send_queue: VecDeque::new(),
last_rc: None,
write_history: VecDeque::new(),
results: Vec::new(),
}
}
pub fn start(&mut self, topology_bytes: &[u8]) {
self.channel.opened = true;
let mut cfg = Message::configuration(&self.gateway);
cfg.data_encode(topology_bytes);
self.channel.send(&cfg);
self.state = ChannelState::Config;
}
pub fn send_queued(&mut self, msg: Message) {
if self.channel.setup && self.send_queue.is_empty() {
self.channel.send(&msg);
} else {
self.send_queue.push_back(msg);
}
}
fn send_dequeue(&mut self) {
if let Some(msg) = self.send_queue.pop_front() {
self.channel.send(&msg);
}
}
pub fn shell(&mut self, nodes: &str, command: &[u8], worker_id: u64) {
let mut ctl = Message::control(worker_id);
if let Message::Control {
ref mut action,
ref mut target,
..
} = ctl
{
*action = "shell".to_string();
*target = nodes.to_string();
}
ctl.data_encode(command);
self.send_queued(ctl);
}
pub fn write(&mut self, nodes: &str, buf: &[u8], worker_id: u64) {
let mut ctl = Message::control(worker_id);
if let Message::Control {
ref mut action,
ref mut target,
..
} = ctl
{
*action = "write".to_string();
*target = nodes.to_string();
}
let msgid = ctl.msgid();
ctl.data_encode(buf);
self.write_history
.push_back((msgid, nodes.to_string(), buf.len()));
self.send_queued(ctl);
}
pub fn set_write_eof(&mut self, nodes: &str, worker_id: u64) {
let mut ctl = Message::control(worker_id);
if let Message::Control {
ref mut action,
ref mut target,
..
} = ctl
{
*action = "eof".to_string();
*target = nodes.to_string();
}
self.send_queued(ctl);
}
pub fn recv(&mut self, msg: Message) -> Result<Vec<PropagationResult>, MessageProcessingError> {
let mut new_results = Vec::new();
match msg.ident() {
"END" => {
self.state = ChannelState::Closed;
return Ok(new_results);
}
"CHA" => {
self.channel.opened = true;
return Ok(new_results);
}
_ => {}
}
match self.state {
ChannelState::Config => {
if msg.ident() == "ACK" {
self.channel.setup = true;
self.state = ChannelState::Gather;
self.send_dequeue();
}
}
ChannelState::Control | ChannelState::Gather => {
match msg.ident() {
"ACK" => {
if let Message::Ack { ack, .. } = &msg {
if let Some(back) = self.write_history.back() {
if back.0 == *ack {
let _ = self.write_history.pop_back();
}
}
}
self.send_dequeue();
}
"OUT" => {
if let Some(nodes_str) = msg.nodes() {
if let Ok(decoded) = msg.data_decode() {
let ns = NodeSet::parse(nodes_str).unwrap_or_default();
for node in ns.iter() {
let result = PropagationResult::StdOut {
node,
data: decoded.clone(),
};
new_results.push(result);
}
}
}
}
"SER" => {
if let Some(nodes_str) = msg.nodes() {
if let Ok(decoded) = msg.data_decode() {
let ns = NodeSet::parse(nodes_str).unwrap_or_default();
for node in ns.iter() {
let result = PropagationResult::StdErr {
node,
data: decoded.clone(),
};
new_results.push(result);
}
}
}
}
"RET" => {
if let Message::Retcode { retcode, nodes, .. } = &msg {
let ns = NodeSet::parse(nodes).unwrap_or_default();
for node in ns.iter() {
new_results.push(PropagationResult::Retcode { node, rc: *retcode });
}
}
}
"TIM" => {
if let Some(nodes_str) = msg.nodes() {
let ns = NodeSet::parse(nodes_str).unwrap_or_default();
for node in ns.iter() {
new_results.push(PropagationResult::Timeout { node });
}
}
}
"RTR" => {
if let Message::Routing {
event,
gateway,
targets,
..
} = &msg
{
new_results.push(PropagationResult::Routing {
event: event.clone(),
gateway: gateway.clone(),
targets: targets.clone(),
});
}
}
"ERR" => {
if let Message::Error { reason, .. } = &msg {
return Err(MessageProcessingError::ParseError(format!(
"{}: {}",
self.gateway, reason
)));
}
}
_ => {
}
}
}
_ => {}
}
self.results.extend(new_results.clone());
Ok(new_results)
}
pub fn ev_hup(&mut self, rc: i32) {
self.last_rc = Some(rc);
}
pub fn ev_close(&mut self) {
self.state = ChannelState::Closed;
}
pub fn current_state(&self) -> ChannelState {
self.state
}
pub fn is_gathering(&self) -> bool {
self.state == ChannelState::Gather
}
pub fn is_closed(&self) -> bool {
self.state == ChannelState::Closed
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::topology::{TopologyGraph, TopologyTree};
fn build_test_topology() -> TopologyTree {
let mut graph = TopologyGraph::new();
let admin = NodeSet::parse("admin").unwrap();
let gw1 = NodeSet::parse("gw1").unwrap();
let gw2 = NodeSet::parse("gw2").unwrap();
let nodes1 = NodeSet::parse("node[1-4]").unwrap();
let nodes2 = NodeSet::parse("node[5-8]").unwrap();
graph.add_route(admin.clone(), gw1.clone()).unwrap();
graph.add_route(admin, gw2.clone()).unwrap();
graph.add_route(gw1, nodes1).unwrap();
graph.add_route(gw2, nodes2).unwrap();
graph.to_tree("admin").unwrap()
}
#[test]
fn test_error_display() {
let err = RouteResolvingError::InvalidRoot("badnode".into());
assert!(err.to_string().contains("badnode"));
let err = RouteResolvingError::NoRoute("node99".into());
assert!(err.to_string().contains("node99"));
let err = RouteResolvingError::Unreachable("gw1".into());
assert!(err.to_string().contains("gw1"));
let err = RouteResolvingError::SameSourceDest("a".into(), "a".into());
assert!(err.to_string().contains("a"));
}
#[test]
fn test_router_creation() {
let tree = build_test_topology();
let router = PropagationTreeRouter::new("admin", &tree, 0).unwrap();
assert_eq!(router.root(), "admin");
assert_eq!(router.fanout(), 0);
}
#[test]
fn test_router_invalid_root() {
let tree = build_test_topology();
let result = PropagationTreeRouter::new("nonexistent", &tree, 0);
assert!(result.is_err());
}
#[test]
fn test_router_next_hop() {
let tree = build_test_topology();
let mut router = PropagationTreeRouter::new("admin", &tree, 0).unwrap();
let dst = NodeSet::parse("node1").unwrap();
let hop = router.next_hop(&dst).unwrap();
assert_eq!(hop.to_string(), "gw1");
let dst = NodeSet::parse("node5").unwrap();
let hop = router.next_hop(&dst).unwrap();
assert_eq!(hop.to_string(), "gw2");
}
#[test]
fn test_router_same_source_dest() {
let tree = build_test_topology();
let mut router = PropagationTreeRouter::new("admin", &tree, 0).unwrap();
let dst = NodeSet::parse("admin").unwrap();
let result = router.next_hop(&dst);
assert!(result.is_err());
}
#[test]
fn test_router_mark_unreachable() {
let tree = build_test_topology();
let mut router = PropagationTreeRouter::new("admin", &tree, 0).unwrap();
router.mark_unreachable("gw1");
let dst = NodeSet::parse("node1").unwrap();
let result = router.next_hop(&dst);
assert!(result.is_err());
}
#[test]
fn test_router_dispatch_directly_connected() {
let tree = build_test_topology();
let router = PropagationTreeRouter::new("admin", &tree, 0).unwrap();
let dst = NodeSet::parse("localnode").unwrap();
let results = router.dispatch(&dst);
assert_eq!(results.len(), 1);
assert_eq!(results[0].0.to_string(), "localnode");
}
#[test]
fn test_router_dispatch_remote() {
let tree = build_test_topology();
let router = PropagationTreeRouter::new("admin", &tree, 0).unwrap();
let dst = NodeSet::parse("node[1-2]").unwrap();
let results = router.dispatch(&dst);
assert!(!results.is_empty());
for (hop, _target) in &results {
assert_eq!(hop.to_string(), "gw1");
}
}
#[test]
fn test_router_load_balancing() {
let mut graph = TopologyGraph::new();
let admin = NodeSet::parse("admin").unwrap();
let gws = NodeSet::parse("gw[1-2]").unwrap();
let nodes = NodeSet::parse("node[1-8]").unwrap();
graph.add_route(admin, gws).unwrap();
let gw1 = NodeSet::parse("gw1").unwrap();
let gw2 = NodeSet::parse("gw2").unwrap();
graph
.add_route(gw1, NodeSet::parse("node[1-4]").unwrap())
.unwrap();
graph
.add_route(gw2, NodeSet::parse("node[5-8]").unwrap())
.unwrap();
let tree = graph.to_tree("admin").unwrap();
let mut router = PropagationTreeRouter::new("admin", &tree, 0).unwrap();
let d1 = NodeSet::parse("node1").unwrap();
let _h1 = router.next_hop(&d1).unwrap();
let d2 = NodeSet::parse("node5").unwrap();
let _h2 = router.next_hop(&d2).unwrap();
}
#[test]
fn test_channel_creation() {
let ch = PropagationChannel::new("gw1");
assert_eq!(ch.gateway, "gw1");
assert_eq!(ch.current_state(), ChannelState::Init);
assert!(!ch.is_gathering());
assert!(!ch.is_closed());
}
#[test]
fn test_channel_start() {
let mut ch = PropagationChannel::new("gw1");
ch.start(b"topology data");
assert_eq!(ch.current_state(), ChannelState::Config);
assert!(ch.channel.opened);
let outgoing = ch.channel.take_outgoing().unwrap();
let xml = String::from_utf8(outgoing).unwrap();
assert!(xml.contains("type=\"CFG\""));
assert!(xml.contains("gateway=\"gw1\""));
}
#[test]
fn test_channel_config_ack() {
let mut ch = PropagationChannel::new("gw1");
ch.start(b"topo");
assert_eq!(ch.state, ChannelState::Config);
let ack = Message::ack(0);
let results = ch.recv(ack).unwrap();
assert!(results.is_empty());
assert_eq!(ch.state, ChannelState::Gather);
assert!(ch.channel.setup);
}
#[test]
fn test_channel_shell_command() {
let mut ch = PropagationChannel::new("gw1");
ch.start(b"topo");
ch.recv(Message::ack(0)).unwrap();
ch.shell("node[1-4]", b"uname -a", 42);
let outgoing = ch.channel.take_outgoing().unwrap();
let outgoing2 = ch.channel.take_outgoing().unwrap();
let xml = String::from_utf8(outgoing2).unwrap();
assert!(xml.contains("type=\"CTL\""));
assert!(xml.contains("action=\"shell\""));
assert!(xml.contains("target=\"node[1-4]\""));
}
#[test]
fn test_channel_recv_stdout() {
let mut ch = PropagationChannel::new("gw1");
ch.start(b"topo");
ch.recv(Message::ack(0)).unwrap();
let mut stdout = Message::stdout("node[1-2]", 1);
stdout.data_encode(b"hello world");
let results = ch.recv(stdout).unwrap();
assert_eq!(results.len(), 2);
for r in &results {
if let PropagationResult::StdOut { data, .. } = r {
assert_eq!(data, b"hello world");
} else {
panic!("expected StdOut");
}
}
}
#[test]
fn test_channel_recv_stderr() {
let mut ch = PropagationChannel::new("gw1");
ch.start(b"topo");
ch.recv(Message::ack(0)).unwrap();
let mut stderr = Message::stderr("node1", 1);
stderr.data_encode(b"error output");
let results = ch.recv(stderr).unwrap();
assert_eq!(results.len(), 1);
if let PropagationResult::StdErr { node, data } = &results[0] {
assert_eq!(node, "node1");
assert_eq!(data, b"error output");
} else {
panic!("expected StdErr");
}
}
#[test]
fn test_channel_recv_retcode() {
let mut ch = PropagationChannel::new("gw1");
ch.start(b"topo");
ch.recv(Message::ack(0)).unwrap();
let ret = Message::retcode("node[1-3]", 0, 1);
let results = ch.recv(ret).unwrap();
assert_eq!(results.len(), 3);
for r in &results {
if let PropagationResult::Retcode { rc, .. } = r {
assert_eq!(*rc, 0);
} else {
panic!("expected Retcode");
}
}
}
#[test]
fn test_channel_recv_timeout() {
let mut ch = PropagationChannel::new("gw1");
ch.start(b"topo");
ch.recv(Message::ack(0)).unwrap();
let tim = Message::timeout("node[5-6]", 1);
let results = ch.recv(tim).unwrap();
assert_eq!(results.len(), 2);
for r in &results {
assert!(matches!(r, PropagationResult::Timeout { .. }));
}
}
#[test]
fn test_channel_recv_routing() {
let mut ch = PropagationChannel::new("gw1");
ch.start(b"topo");
ch.recv(Message::ack(0)).unwrap();
let rtr = Message::routing("connect", "gw2", "node[10-20]", 1);
let results = ch.recv(rtr).unwrap();
assert_eq!(results.len(), 1);
if let PropagationResult::Routing {
event,
gateway,
targets,
} = &results[0]
{
assert_eq!(event, "connect");
assert_eq!(gateway, "gw2");
assert_eq!(targets, "node[10-20]");
} else {
panic!("expected Routing");
}
}
#[test]
fn test_channel_recv_error() {
let mut ch = PropagationChannel::new("gw1");
ch.start(b"topo");
ch.recv(Message::ack(0)).unwrap();
let err = Message::error("something went wrong");
let result = ch.recv(err);
assert!(result.is_err());
}
#[test]
fn test_channel_recv_end() {
let mut ch = PropagationChannel::new("gw1");
ch.start(b"topo");
ch.recv(Message::ack(0)).unwrap();
let end = Message::end();
ch.recv(end).unwrap();
assert!(ch.is_closed());
}
#[test]
fn test_channel_ev_hup() {
let mut ch = PropagationChannel::new("gw1");
ch.ev_hup(1);
assert_eq!(ch.last_rc, Some(1));
}
#[test]
fn test_channel_ev_close() {
let mut ch = PropagationChannel::new("gw1");
ch.ev_close();
assert!(ch.is_closed());
}
#[test]
fn test_channel_send_queue() {
let mut ch = PropagationChannel::new("gw1");
ch.start(b"topo");
ch.shell("node1", b"echo hi", 1);
let _cfg = ch.channel.take_outgoing().unwrap();
assert!(ch.channel.take_outgoing().is_none());
ch.recv(Message::ack(0)).unwrap();
let ctl = ch.channel.take_outgoing().unwrap();
let xml = String::from_utf8(ctl).unwrap();
assert!(xml.contains("type=\"CTL\""));
}
#[test]
fn test_channel_write_and_eof() {
let mut ch = PropagationChannel::new("gw1");
ch.start(b"topo");
ch.recv(Message::ack(0)).unwrap();
ch.write("node1", b"input data", 1);
ch.set_write_eof("node1", 1);
let _ = ch.channel.take_outgoing(); let write_msg = ch.channel.take_outgoing().unwrap();
let xml = String::from_utf8(write_msg).unwrap();
assert!(xml.contains("action=\"write\""));
let eof_msg = ch.channel.take_outgoing().unwrap();
let xml2 = String::from_utf8(eof_msg).unwrap();
assert!(xml2.contains("action=\"eof\""));
}
#[test]
fn test_channel_results_accumulate() {
let mut ch = PropagationChannel::new("gw1");
ch.start(b"topo");
ch.recv(Message::ack(0)).unwrap();
let mut out1 = Message::stdout("node1", 1);
out1.data_encode(b"line1");
ch.recv(out1).unwrap();
let mut out2 = Message::stdout("node2", 1);
out2.data_encode(b"line2");
ch.recv(out2).unwrap();
let ret = Message::retcode("node[1-2]", 0, 1);
ch.recv(ret).unwrap();
assert_eq!(ch.results.len(), 4);
}
}