60 lines
1.9 KiB
Rust
60 lines
1.9 KiB
Rust
use std::collections::HashSet;
|
|
use crate::node_service::node_service_client::NodeServiceClient;
|
|
use crate::node_service::{CollectTopologyRequest, HealthCheckRequest, Topology as TopologyProto};
|
|
use crate::topology::{DeviceCapabilities, Topology};
|
|
use std::net::SocketAddr;
|
|
use tonic::codec::CompressionEncoding;
|
|
|
|
pub struct PeerHandle {
|
|
pub node_id: String,
|
|
pub address: SocketAddr,
|
|
pub address_priority: u8,
|
|
pub description: Option<String>,
|
|
pub client: tokio::sync::Mutex<NodeServiceClient<tonic::transport::Channel>>,
|
|
pub device_capabilities: DeviceCapabilities,
|
|
}
|
|
|
|
impl PeerHandle {
|
|
pub async fn new(
|
|
node_id: String,
|
|
address: SocketAddr,
|
|
address_priority: u8,
|
|
description: Option<String>,
|
|
device_capabilities: DeviceCapabilities,
|
|
) -> Result<Self, tonic::transport::Error> {
|
|
let endpoint = format!("http://{}", address);
|
|
let client = NodeServiceClient::connect(endpoint)
|
|
.await?
|
|
.accept_compressed(CompressionEncoding::Gzip);
|
|
|
|
Ok(Self {
|
|
node_id,
|
|
description,
|
|
address_priority,
|
|
address,
|
|
client: tokio::sync::Mutex::new(client),
|
|
device_capabilities,
|
|
})
|
|
}
|
|
|
|
pub async fn is_healthy(&self) -> bool {
|
|
self.client
|
|
.lock()
|
|
.await
|
|
.health_check(HealthCheckRequest::default())
|
|
.await
|
|
.ok()
|
|
.map(|x| x.into_inner().is_healthy)
|
|
.unwrap_or(false)
|
|
}
|
|
|
|
pub async fn collect_topology(&self, visited: HashSet<String>, max_depth: u8) -> Topology {
|
|
let response = self.client.lock().await.collect_topology(CollectTopologyRequest {
|
|
visited: visited.clone().into_iter().collect(),
|
|
max_depth: max_depth as i32
|
|
}).await.unwrap().into_inner();
|
|
|
|
response.into()
|
|
}
|
|
}
|