Spaces:
Build error
Build error
File size: 2,402 Bytes
84d2a97 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 |
use std::time::Duration;
use parking_lot::Mutex;
use tonic::transport::{Channel, ClientTlsConfig, Error as TonicError, Uri};
use crate::grpc::dynamic_pool::{CountedItem, DynamicPool};
pub async fn make_grpc_channel(
timeout: Duration,
connection_timeout: Duration,
uri: Uri,
tls_config: Option<ClientTlsConfig>,
) -> Result<Channel, TonicError> {
let mut endpoint = Channel::builder(uri)
.timeout(timeout)
.connect_timeout(connection_timeout);
if let Some(config) = tls_config {
endpoint = endpoint.tls_config(config)?;
}
// `connect` is using the `Reconnect` network service internally to handle dropped connections
endpoint.connect().await
}
pub struct DynamicChannelPool {
pool: Mutex<DynamicPool<Channel>>,
uri: Uri,
timeout: Duration,
connection_timeout: Duration,
tls_config: Option<ClientTlsConfig>,
}
impl DynamicChannelPool {
pub async fn new(
uri: Uri,
timeout: Duration,
connection_timeout: Duration,
tls_config: Option<ClientTlsConfig>,
usage_per_channel: usize,
min_channels: usize,
) -> Result<Self, TonicError> {
let mut channels = Vec::with_capacity(min_channels);
for _ in 0..min_channels {
let channel =
make_grpc_channel(timeout, connection_timeout, uri.clone(), tls_config.clone())
.await?;
channels.push(channel);
}
let pool = DynamicPool::new(channels, usage_per_channel, min_channels);
Ok(Self {
pool: Mutex::new(pool),
uri,
timeout,
connection_timeout,
tls_config,
})
}
pub async fn choose(&self) -> Result<CountedItem<Channel>, TonicError> {
let channel = self.pool.lock().choose();
let channel = match channel {
None => {
let channel = make_grpc_channel(
self.timeout,
self.connection_timeout,
self.uri.clone(),
self.tls_config.clone(),
)
.await?;
self.pool.lock().add(channel)
}
Some(channel) => channel,
};
Ok(channel)
}
pub fn drop_channel(&self, channel: CountedItem<Channel>) {
self.pool.lock().drop_item(channel);
}
}
|