use async_trait::async_trait;
use serde::{Deserialize, Serialize};
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct DiscoveredPeer {
pub node_id: String,
pub address: String,
pub platform: String,
pub source: DiscoverySource,
}
#[non_exhaustive]
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum DiscoverySource {
SeedPeer,
Mdns,
Kademlia,
PeerExchange,
Manual,
}
#[async_trait]
pub trait DiscoveryBackend: Send + Sync + 'static {
fn name(&self) -> &str;
async fn start(&mut self) -> Result<(), DiscoveryError>;
async fn poll(&mut self) -> Vec<DiscoveredPeer>;
async fn stop(&mut self) -> Result<(), DiscoveryError>;
}
#[non_exhaustive]
#[derive(Debug, thiserror::Error)]
pub enum DiscoveryError {
#[error("discovery backend error: {0}")]
Backend(String),
#[error("invalid seed address: {0}")]
InvalidAddress(String),
#[error("discovery not started")]
NotStarted,
}
pub struct DiscoveryCoordinator {
backends: Vec<Box<dyn DiscoveryBackend>>,
known_peers: std::collections::HashSet<String>,
}
impl DiscoveryCoordinator {
pub fn new() -> Self {
Self {
backends: Vec::new(),
known_peers: std::collections::HashSet::new(),
}
}
pub fn add_backend(&mut self, backend: Box<dyn DiscoveryBackend>) {
self.backends.push(backend);
}
pub async fn start_all(&mut self) -> Result<(), DiscoveryError> {
for backend in &mut self.backends {
backend.start().await?;
}
Ok(())
}
pub async fn poll_all(&mut self) -> Vec<DiscoveredPeer> {
let mut new_peers = Vec::new();
for backend in &mut self.backends {
let peers = backend.poll().await;
for peer in peers {
if self.known_peers.insert(peer.node_id.clone()) {
new_peers.push(peer);
}
}
}
new_peers
}
pub async fn stop_all(&mut self) -> Result<(), DiscoveryError> {
for backend in &mut self.backends {
backend.stop().await?;
}
Ok(())
}
pub fn known_peer_count(&self) -> usize {
self.known_peers.len()
}
pub fn is_known(&self, node_id: &str) -> bool {
self.known_peers.contains(node_id)
}
}
impl Default for DiscoveryCoordinator {
fn default() -> Self {
Self::new()
}
}
#[cfg(test)]
mod tests {
use super::*;
struct MockBackend {
name: &'static str,
peers: Vec<DiscoveredPeer>,
started: bool,
}
impl MockBackend {
fn new(name: &'static str, peers: Vec<DiscoveredPeer>) -> Self {
Self {
name,
peers,
started: false,
}
}
}
#[async_trait]
impl DiscoveryBackend for MockBackend {
fn name(&self) -> &str {
self.name
}
async fn start(&mut self) -> Result<(), DiscoveryError> {
self.started = true;
Ok(())
}
async fn poll(&mut self) -> Vec<DiscoveredPeer> {
std::mem::take(&mut self.peers)
}
async fn stop(&mut self) -> Result<(), DiscoveryError> {
self.started = false;
Ok(())
}
}
fn make_peer(id: &str, source: DiscoverySource) -> DiscoveredPeer {
DiscoveredPeer {
node_id: id.to_string(),
address: format!("127.0.0.1:900{}", id.len()),
platform: "linux".into(),
source,
}
}
#[tokio::test]
async fn coordinator_add_and_start() {
let mut coord = DiscoveryCoordinator::new();
assert_eq!(coord.known_peer_count(), 0);
let backend = MockBackend::new("mock", vec![make_peer("a", DiscoverySource::Manual)]);
coord.add_backend(Box::new(backend));
coord.start_all().await.unwrap();
}
#[tokio::test]
async fn coordinator_poll_deduplicates() {
let mut coord = DiscoveryCoordinator::new();
let b1 = MockBackend::new(
"b1",
vec![
make_peer("node-1", DiscoverySource::SeedPeer),
make_peer("node-2", DiscoverySource::SeedPeer),
],
);
let b2 = MockBackend::new(
"b2",
vec![
make_peer("node-2", DiscoverySource::Mdns), make_peer("node-3", DiscoverySource::Mdns),
],
);
coord.add_backend(Box::new(b1));
coord.add_backend(Box::new(b2));
coord.start_all().await.unwrap();
let peers = coord.poll_all().await;
assert_eq!(peers.len(), 3);
assert_eq!(coord.known_peer_count(), 3);
assert!(coord.is_known("node-1"));
assert!(coord.is_known("node-2"));
assert!(coord.is_known("node-3"));
}
#[tokio::test]
async fn coordinator_second_poll_empty() {
let mut coord = DiscoveryCoordinator::new();
let b = MockBackend::new("b", vec![make_peer("x", DiscoverySource::Manual)]);
coord.add_backend(Box::new(b));
coord.start_all().await.unwrap();
let first = coord.poll_all().await;
assert_eq!(first.len(), 1);
let second = coord.poll_all().await;
assert!(second.is_empty());
}
#[tokio::test]
async fn coordinator_stop_all() {
let mut coord = DiscoveryCoordinator::new();
let b = MockBackend::new("b", vec![]);
coord.add_backend(Box::new(b));
coord.start_all().await.unwrap();
coord.stop_all().await.unwrap();
}
#[test]
fn discovered_peer_serde_roundtrip() {
let peer = make_peer("serde-test", DiscoverySource::Kademlia);
let json = serde_json::to_string(&peer).unwrap();
let back: DiscoveredPeer = serde_json::from_str(&json).unwrap();
assert_eq!(back.node_id, "serde-test");
assert_eq!(back.source, DiscoverySource::Kademlia);
}
#[test]
fn discovery_source_serde_roundtrip() {
let sources = [
DiscoverySource::SeedPeer,
DiscoverySource::Mdns,
DiscoverySource::Kademlia,
DiscoverySource::PeerExchange,
DiscoverySource::Manual,
];
for src in &sources {
let json = serde_json::to_string(src).unwrap();
let back: DiscoverySource = serde_json::from_str(&json).unwrap();
assert_eq!(&back, src);
}
}
#[test]
fn discovery_error_display() {
let e = DiscoveryError::Backend("timeout".into());
assert!(e.to_string().contains("timeout"));
let e = DiscoveryError::InvalidAddress("bad:addr".into());
assert!(e.to_string().contains("bad:addr"));
let e = DiscoveryError::NotStarted;
assert!(e.to_string().contains("not started"));
}
}