use crate::grpc_network::GrpcError;
use dashmap::DashMap;
use std::time::Duration;
use tonic::transport::{Channel, Endpoint};
#[derive(Clone, Default)]
pub struct ConnectionPool {
connect_timeout: Option<Duration>,
timeout: Option<Duration>,
channels: DashMap<String, Channel>,
}
impl ConnectionPool {
pub fn with_connect_timeout(mut self, connect_timeout: impl Into<Option<Duration>>) -> Self {
self.connect_timeout = connect_timeout.into();
self
}
pub fn with_timeout(mut self, timeout: impl Into<Option<Duration>>) -> Self {
self.timeout = timeout.into();
self
}
pub fn channel(&self, address: String) -> Result<Channel, GrpcError> {
let channel = self
.channels
.entry(address.clone())
.or_try_insert_with(|| {
let mut endpoint = Endpoint::from_shared(address)?;
if let Some(timeout) = self.connect_timeout {
endpoint = endpoint.connect_timeout(timeout);
}
if let Some(timeout) = self.timeout {
endpoint = endpoint.timeout(timeout);
}
Ok::<_, GrpcError>(endpoint.connect_lazy())
})?;
Ok(channel.clone())
}
}