use std::collections::HashMap;
#[derive(Debug, Clone)]
#[non_exhaustive]
pub struct DiscoveredNode {
pub address: String,
pub port: u16,
pub metadata: HashMap<String, String>,
}
impl DiscoveredNode {
pub fn new(address: impl Into<String>, port: u16) -> Self {
Self {
address: address.into(),
port,
metadata: HashMap::new(),
}
}
pub fn with_metadata(mut self, key: impl Into<String>, value: impl Into<String>) -> Self {
self.metadata.insert(key.into(), value.into());
self
}
#[must_use]
#[inline]
pub fn socket_addr(&self) -> String {
format!("{}:{}", self.address, self.port)
}
}
pub trait DiscoveryBackend: Send + Sync {
fn discover(&self) -> impl std::future::Future<Output = Vec<DiscoveredNode>> + Send;
}
#[derive(Debug, Clone)]
pub struct StaticDiscovery {
nodes: Vec<DiscoveredNode>,
}
impl StaticDiscovery {
pub fn new(nodes: Vec<DiscoveredNode>) -> Self {
Self { nodes }
}
}
impl DiscoveryBackend for StaticDiscovery {
async fn discover(&self) -> Vec<DiscoveredNode> {
tracing::debug!(
count = self.nodes.len(),
"static discovery returning fixed node list"
);
self.nodes.clone()
}
}
#[derive(Debug, Clone)]
pub struct DnsDiscovery {
pub srv_name: String,
}
impl DnsDiscovery {
pub fn new(srv_name: impl Into<String>) -> Self {
Self {
srv_name: srv_name.into(),
}
}
}
impl DiscoveryBackend for DnsDiscovery {
async fn discover(&self) -> Vec<DiscoveredNode> {
tracing::warn!(
srv_name = %self.srv_name,
"DNS SRV discovery is a stub — no DNS client dependency available"
);
Vec::new()
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn discovered_node_new() {
let node = DiscoveredNode::new("10.0.0.1", 8080);
assert_eq!(node.address, "10.0.0.1");
assert_eq!(node.port, 8080);
assert!(node.metadata.is_empty());
}
#[test]
fn discovered_node_with_metadata() {
let node = DiscoveredNode::new("10.0.0.1", 8080)
.with_metadata("region", "us-east-1")
.with_metadata("gpu", "A100");
assert_eq!(node.metadata.len(), 2);
assert_eq!(node.metadata.get("region").unwrap(), "us-east-1");
assert_eq!(node.metadata.get("gpu").unwrap(), "A100");
}
#[test]
fn discovered_node_socket_addr() {
let node = DiscoveredNode::new("10.0.0.1", 8080);
assert_eq!(node.socket_addr(), "10.0.0.1:8080");
}
#[tokio::test]
async fn static_discovery_returns_all_nodes() {
let nodes = vec![
DiscoveredNode::new("10.0.0.1", 8080),
DiscoveredNode::new("10.0.0.2", 8080),
DiscoveredNode::new("10.0.0.3", 9090),
];
let backend = StaticDiscovery::new(nodes);
let discovered = backend.discover().await;
assert_eq!(discovered.len(), 3);
assert_eq!(discovered[0].address, "10.0.0.1");
assert_eq!(discovered[1].address, "10.0.0.2");
assert_eq!(discovered[2].port, 9090);
}
#[tokio::test]
async fn static_discovery_empty() {
let backend = StaticDiscovery::new(vec![]);
let discovered = backend.discover().await;
assert!(discovered.is_empty());
}
#[tokio::test]
async fn dns_discovery_stub_returns_empty() {
let backend = DnsDiscovery::new("_agnosai._tcp.fleet.example.com");
let discovered = backend.discover().await;
assert!(discovered.is_empty());
}
#[test]
fn dns_discovery_stores_srv_name() {
let backend = DnsDiscovery::new("_agnosai._tcp.fleet.example.com");
assert_eq!(backend.srv_name, "_agnosai._tcp.fleet.example.com");
}
}