mockforge_grpc/reflection/
connection_pool.rs

1//! Connection pool for gRPC clients
2
3use std::collections::HashMap;
4use std::sync::Arc;
5use tokio::sync::Mutex;
6use tonic::transport::Channel;
7use tracing::{debug, trace};
8
9/// A simple connection pool for gRPC channels
10pub struct ConnectionPool {
11    /// Map of endpoint URIs to channels
12    channels: Arc<Mutex<HashMap<String, Channel>>>,
13}
14
15impl Default for ConnectionPool {
16    fn default() -> Self {
17        Self::new()
18    }
19}
20
21impl ConnectionPool {
22    /// Create a new connection pool
23    pub fn new() -> Self {
24        Self {
25            channels: Arc::new(Mutex::new(HashMap::new())),
26        }
27    }
28
29    /// Get a channel for the given endpoint URI
30    /// If a channel already exists for this endpoint, it will be reused
31    /// Otherwise, a new channel will be created and added to the pool
32    pub async fn get_channel(
33        &self,
34        endpoint_uri: &str,
35    ) -> Result<Channel, Box<dyn std::error::Error + Send + Sync>> {
36        trace!("Getting channel for endpoint: {}", endpoint_uri);
37
38        // Lock the channels map
39        let mut channels = self.channels.lock().await;
40
41        // Check if we already have a channel for this endpoint
42        if let Some(channel) = channels.get(endpoint_uri) {
43            debug!("Reusing existing channel for endpoint: {}", endpoint_uri);
44            return Ok(channel.clone());
45        }
46
47        // Create a new channel for this endpoint
48        debug!("Creating new channel for endpoint: {}", endpoint_uri);
49        let channel = Channel::from_shared(endpoint_uri.to_string())?.connect().await?;
50
51        // Add the new channel to the pool
52        channels.insert(endpoint_uri.to_string(), channel.clone());
53
54        Ok(channel)
55    }
56}
57
58impl Clone for ConnectionPool {
59    fn clone(&self) -> Self {
60        Self {
61            channels: self.channels.clone(),
62        }
63    }
64}
65
66#[cfg(test)]
67mod tests {
68
69    #[test]
70    fn test_module_compiles() {}
71}
mockforge_grpc/reflection/connection_pool.rs

mockforge_grpc/reflection/
connection_pool.rs
