qail_qdrant/
pool.rs

1//! Connection pool for Qdrant gRPC driver.
2//!
3//! Provides efficient connection pooling with semaphore-based concurrency limiting.
4
5use crate::error::{QdrantError, QdrantResult};
6use crate::driver::QdrantDriver;
7use std::sync::Arc;
8use tokio::sync::Semaphore;
9
10/// Configuration for the connection pool.
11#[derive(Debug, Clone)]
12pub struct PoolConfig {
13    /// Maximum number of concurrent connections.
14    pub max_connections: usize,
15    /// Host to connect to.
16    pub host: String,
17    /// gRPC port (default 6334).
18    pub port: u16,
19}
20
21impl PoolConfig {
22    /// Create a new pool configuration.
23    pub fn new(host: impl Into<String>, port: u16) -> Self {
24        Self {
25            max_connections: 10,
26            host: host.into(),
27            port,
28        }
29    }
30
31    /// Set maximum connections.
32    pub fn max_connections(mut self, max: usize) -> Self {
33        self.max_connections = max;
34        self
35    }
36}
37
38impl Default for PoolConfig {
39    fn default() -> Self {
40        Self {
41            max_connections: 10,
42            host: "localhost".to_string(),
43            port: 6334,
44        }
45    }
46}
47
48/// Connection pool for Qdrant gRPC driver.
49///
50/// Uses a semaphore to limit concurrent connections. Each connection
51/// is independent and can be used concurrently.
52///
53/// # Example
54/// ```ignore
55/// use qail_qdrant::{QdrantPool, PoolConfig};
56///
57/// let pool = QdrantPool::new(
58///     PoolConfig::new("localhost", 6334).max_connections(20)
59/// ).await?;
60///
61/// // Get a connection from the pool
62/// let mut conn = pool.get().await?;
63/// let results = conn.search("products", &embedding, 10, None).await?;
64/// ```
65#[derive(Clone)]
66pub struct QdrantPool {
67    config: Arc<PoolConfig>,
68    semaphore: Arc<Semaphore>,
69}
70
71impl QdrantPool {
72    /// Create a new connection pool.
73    pub async fn new(config: PoolConfig) -> QdrantResult<Self> {
74        Ok(Self {
75            semaphore: Arc::new(Semaphore::new(config.max_connections)),
76            config: Arc::new(config),
77        })
78    }
79
80    /// Get a connection from the pool.
81    ///
82    /// This acquires a permit from the semaphore, limiting concurrency.
83    /// The connection is created lazily when acquired.
84    pub async fn get(&self) -> QdrantResult<PooledConnection> {
85        let permit = self.semaphore
86            .clone()
87            .acquire_owned()
88            .await
89            .map_err(|e| QdrantError::Connection(e.to_string()))?;
90        
91        // Create a new connection (lazy)
92        let driver = QdrantDriver::connect(&self.config.host, self.config.port).await?;
93        
94        Ok(PooledConnection {
95            driver,
96            _permit: permit,
97        })
98    }
99
100    /// Number of available permits (connections not in use).
101    pub fn available(&self) -> usize {
102        self.semaphore.available_permits()
103    }
104
105    /// Maximum number of connections.
106    pub fn max_connections(&self) -> usize {
107        self.config.max_connections
108    }
109}
110
111/// A pooled connection that releases back to the pool on drop.
112pub struct PooledConnection {
113    driver: QdrantDriver,
114    _permit: tokio::sync::OwnedSemaphorePermit,
115}
116
117impl std::ops::Deref for PooledConnection {
118    type Target = QdrantDriver;
119
120    fn deref(&self) -> &Self::Target {
121        &self.driver
122    }
123}
124
125impl std::ops::DerefMut for PooledConnection {
126    fn deref_mut(&mut self) -> &mut Self::Target {
127        &mut self.driver
128    }
129}
130
131#[cfg(test)]
132mod tests {
133    use super::*;
134
135    #[test]
136    fn test_pool_config_builder() {
137        let config = PoolConfig::new("localhost", 6334)
138            .max_connections(20);
139        
140        assert_eq!(config.max_connections, 20);
141        assert_eq!(config.host, "localhost");
142        assert_eq!(config.port, 6334);
143    }
144
145    #[test]
146    fn test_pool_config_default() {
147        let config = PoolConfig::default();
148        assert_eq!(config.max_connections, 10);
149        assert_eq!(config.host, "localhost");
150        assert_eq!(config.port, 6334);
151    }
152}
qail_qdrant/pool.rs

qail_qdrant/
pool.rs
