Skip to main content

linera_rpc/grpc/
pool.rs

1// Copyright (c) Zefchain Labs, Inc.
2// SPDX-License-Identifier: Apache-2.0
3
4use std::sync::Arc;
5
6use linera_base::time::Duration;
7
8use super::{transport, GrpcError};
9
10/// A pool of transport channels to be used by gRPC.
11#[derive(Clone, Default)]
12pub struct GrpcConnectionPool {
13    options: transport::Options,
14    channels: Arc<papaya::HashMap<String, transport::Channel>>,
15}
16
17impl GrpcConnectionPool {
18    /// Creates a new connection pool with the given transport options.
19    pub fn new(options: transport::Options) -> Self {
20        Self {
21            options,
22            channels: Arc::new(papaya::HashMap::default()),
23        }
24    }
25
26    /// Sets the connection timeout for channels created by this pool.
27    pub fn with_connect_timeout(mut self, connect_timeout: impl Into<Option<Duration>>) -> Self {
28        self.options.connect_timeout = connect_timeout.into();
29        self
30    }
31
32    /// Sets the request timeout for channels created by this pool.
33    pub fn with_timeout(mut self, timeout: impl Into<Option<Duration>>) -> Self {
34        self.options.timeout = timeout.into();
35        self
36    }
37
38    /// Obtains a channel for the current address. Either clones an existing one (thereby
39    /// reusing the connection), or creates one if needed. New channels do not create a
40    /// connection immediately.
41    pub fn channel(&self, address: String) -> Result<transport::Channel, GrpcError> {
42        let pinned = self.channels.pin();
43        if let Some(channel) = pinned.get(&address) {
44            return Ok(channel.clone());
45        }
46        let channel = transport::create_channel(address.clone(), &self.options)?;
47        Ok(pinned.get_or_insert(address, channel).clone())
48    }
49}