76 lines
2.5 KiB
Rust
76 lines
2.5 KiB
Rust
use super::challenge::NodeUpdate;
|
|
use crate::datastore::Store;
|
|
use crate::grpc::challenge::update_client::UpdateClient;
|
|
use std::sync::Arc;
|
|
use tokio::sync::broadcast::Sender;
|
|
use tokio::time::{sleep, Duration};
|
|
use tokio_stream::wrappers::BroadcastStream;
|
|
use tokio_stream::StreamExt;
|
|
|
|
#[derive(Clone)]
|
|
pub struct ConnManager {
|
|
ds: Arc<Store>,
|
|
tx: Sender<NodeUpdate>,
|
|
}
|
|
|
|
impl ConnManager {
|
|
pub fn init(ds: Arc<Store>, tx: Sender<NodeUpdate>) -> Self {
|
|
Self { ds, tx }
|
|
}
|
|
|
|
pub async fn start_with_node(self, node_ip: String) {
|
|
self.connect_wrapper(node_ip).await;
|
|
}
|
|
|
|
pub async fn start(self) {
|
|
loop {
|
|
if let Some(node) = self.ds.get_random_node().await {
|
|
if node != "localhost" {
|
|
self.connect_wrapper(node.clone()).await;
|
|
}
|
|
}
|
|
sleep(Duration::from_secs(3)).await;
|
|
}
|
|
}
|
|
|
|
async fn connect_wrapper(&self, node_ip: String) {
|
|
let ds = self.ds.clone();
|
|
ds.add_conn(&node_ip).await;
|
|
if let Err(e) = self.connect(node_ip.clone()).await {
|
|
println!("Client connection for {node_ip} failed: {e:?}");
|
|
}
|
|
ds.delete_conn(&node_ip).await;
|
|
}
|
|
|
|
async fn connect(&self, node_ip: String) -> Result<(), Box<dyn std::error::Error>> {
|
|
println!("Connecting to {node_ip}...");
|
|
let mut client = UpdateClient::connect(format!("http://{node_ip}:31373")).await?;
|
|
|
|
let rx = self.tx.subscribe();
|
|
let rx_stream = BroadcastStream::new(rx).filter_map(|n| n.ok());
|
|
let response = client.get_updates(rx_stream).await?;
|
|
let mut resp_stream = response.into_inner();
|
|
|
|
let _ = self.tx.send(self.ds.get_localhost().await);
|
|
|
|
while let Some(mut update) = resp_stream.message().await? {
|
|
// "localhost" IPs need to be changed to the real IP of the counterpart
|
|
if update.ip == "localhost" {
|
|
update.ip = node_ip.clone();
|
|
// since we are connecting TO this server, we have a guarantee that this
|
|
// server is not behind NAT, so we can set it public
|
|
update.public = true;
|
|
}
|
|
|
|
// update the entire network in case the information is new
|
|
if self.ds.process_node_update(update.clone()).await {
|
|
if let Err(_) = self.tx.send(update.clone()) {
|
|
println!("tokio broadcast receivers had an issue consuming the channel");
|
|
}
|
|
};
|
|
}
|
|
|
|
Ok(())
|
|
}
|
|
}
|