use std::collections::HashMap;
use std::net::SocketAddr;
use std::str::FromStr;
use std::sync::Arc;
use std::time::Duration;
use parking_lot::RwLock;
use tonic::transport::{Channel, Endpoint, Error};
pub const TIMEOUT_LIMIT: Duration = Duration::from_secs(2);
pub const CONNECT_TIMEOUT_LIMIT: Duration = Duration::from_secs(2);
#[derive(Clone, Default)]
pub struct RpcNetwork {
clients: Arc<RwLock<HashMap<SocketAddr, Channel>>>,
}
impl RpcNetwork {
pub async fn get_or_connect(&self, addr: SocketAddr) -> Result<Channel, Error> {
{
let guard = self.clients.read();
if let Some(channel) = guard.get(&addr) {
return Ok(channel.clone());
}
}
trace!(addr = %addr, "Connect client to network.");
self.connect(addr).await
}
pub async fn connect(&self, addr: SocketAddr) -> Result<Channel, Error> {
let uri = format!("http://{}", addr);
let channel = Endpoint::from_str(&uri)
.unwrap()
.timeout(TIMEOUT_LIMIT)
.connect_timeout(CONNECT_TIMEOUT_LIMIT)
.connect()
.await?;
{
let mut guard = self.clients.write();
guard.insert(addr, channel.clone());
}
Ok(channel)
}
pub fn get_or_connect_lazy(&self, addr: SocketAddr) -> Channel {
{
let guard = self.clients.read();
if let Some(channel) = guard.get(&addr) {
return channel.clone();
}
}
self.connect_lazy(addr)
}
pub fn connect_lazy(&self, addr: SocketAddr) -> Channel {
let uri = format!("http://{}", addr);
let channel = Endpoint::from_str(&uri)
.unwrap()
.timeout(TIMEOUT_LIMIT)
.connect_timeout(CONNECT_TIMEOUT_LIMIT)
.connect_lazy();
{
let mut guard = self.clients.write();
guard.insert(addr, channel.clone());
}
channel
}
pub fn disconnect(&self, addr: SocketAddr) {
let mut guard = self.clients.write();
guard.remove(&addr);
}
}