rs_dapi_client/
connection_pool.rs

1use std::{
2    fmt::Display,
3    sync::{Arc, Mutex},
4};
5
6use lru::LruCache;
7
8use crate::{
9    request_settings::AppliedRequestSettings,
10    transport::{CoreGrpcClient, PlatformGrpcClient},
11    Uri,
12};
13
14/// ConnectionPool represents pool of connections to DAPI nodes.
15///
16/// It can be cloned and shared between threads.
17/// Cloning the pool will create a new reference to the same pool.
18#[derive(Debug, Clone)]
19pub struct ConnectionPool {
20    inner: Arc<Mutex<LruCache<String, PoolItem>>>,
21}
22
23impl ConnectionPool {
24    /// Create a new pool with a given capacity.
25    /// The pool will evict the least recently used item when the capacity is reached.
26    ///
27    /// # Panics
28    ///
29    /// Panics if the capacity is zero.
30    pub fn new(capacity: usize) -> Self {
31        Self {
32            inner: Arc::new(Mutex::new(LruCache::new(
33                capacity.try_into().expect("must be non-zero"),
34            ))),
35        }
36    }
37}
38
39impl Default for ConnectionPool {
40    fn default() -> Self {
41        Self::new(50)
42    }
43}
44
45impl ConnectionPool {
46    /// Get item from the pool for the given uri and settings.
47    ///
48    /// # Arguments
49    /// * `prefix` -  Prefix for the item in the pool. Used to distinguish between Core and Platform clients.
50    /// * `uri` - URI of the node.
51    /// * `settings` - Applied request settings.
52    pub fn get(
53        &self,
54        prefix: PoolPrefix,
55        uri: &Uri,
56        settings: Option<&AppliedRequestSettings>,
57    ) -> Option<PoolItem> {
58        let key = Self::key(prefix, uri, settings);
59        self.inner.lock().expect("must lock").get(&key).cloned()
60    }
61
62    /// Get value from cache or create it using provided closure.
63    /// If value is already in the cache, it will be returned.
64    /// If value is not in the cache, it will be created by calling `create()` and stored in the cache.
65    ///
66    /// # Arguments
67    /// * `prefix` -  Prefix for the item in the pool. Used to distinguish between Core and Platform clients.
68    /// * `uri` - URI of the node.
69    /// * `settings` - Applied request settings.
70    pub fn get_or_create<E>(
71        &self,
72        prefix: PoolPrefix,
73        uri: &Uri,
74        settings: Option<&AppliedRequestSettings>,
75        create: impl FnOnce() -> Result<PoolItem, E>,
76    ) -> Result<PoolItem, E> {
77        if let Some(cli) = self.get(prefix, uri, settings) {
78            return Ok(cli);
79        }
80
81        let cli = create();
82        if let Ok(cli) = &cli {
83            self.put(uri, settings, cli.clone());
84        }
85        cli
86    }
87
88    /// Put item into the pool for the given uri and settings.
89    pub fn put(&self, uri: &Uri, settings: Option<&AppliedRequestSettings>, value: PoolItem) {
90        let key = Self::key(&value, uri, settings);
91        self.inner.lock().expect("must lock").put(key, value);
92    }
93
94    fn key<C: Into<PoolPrefix>>(
95        class: C,
96        uri: &Uri,
97        settings: Option<&AppliedRequestSettings>,
98    ) -> String {
99        let prefix: PoolPrefix = class.into();
100        // Only connection-affecting settings participate in the key (see
101        // `AppliedRequestSettings::connection_key`), so requests differing only
102        // in per-request knobs (timeout, retries, banning) share a connection.
103        // The settings segment is always present (and contains no `:`), so the
104        // two branches cannot produce colliding shapes even for a URI whose
105        // path mimics a key fragment.
106        match settings {
107            Some(settings) => format!("{}:{}:{}", prefix, uri, settings.connection_key()),
108            None => format!("{}:{}:none", prefix, uri),
109        }
110    }
111}
112
113/// Item stored in the pool.
114///
115/// We use an enum as we need to represent two different types of clients.
116#[derive(Clone, Debug)]
117pub enum PoolItem {
118    Core(CoreGrpcClient),
119    Platform(PlatformGrpcClient),
120}
121
122impl From<PlatformGrpcClient> for PoolItem {
123    fn from(client: PlatformGrpcClient) -> Self {
124        Self::Platform(client)
125    }
126}
127impl From<CoreGrpcClient> for PoolItem {
128    fn from(client: CoreGrpcClient) -> Self {
129        Self::Core(client)
130    }
131}
132
133impl From<PoolItem> for PlatformGrpcClient {
134    fn from(client: PoolItem) -> Self {
135        match client {
136            PoolItem::Platform(client) => client,
137            _ => {
138                tracing::error!(
139                    ?client,
140                    "invalid connection fetched from pool: expected platform client"
141                );
142                panic!("ClientType is not Platform: {:?}", client)
143            }
144        }
145    }
146}
147
148impl From<PoolItem> for CoreGrpcClient {
149    fn from(client: PoolItem) -> Self {
150        match client {
151            PoolItem::Core(client) => client,
152            _ => {
153                tracing::error!(
154                    ?client,
155                    "invalid connection fetched from pool: expected core client"
156                );
157                panic!("ClientType is not Core: {:?}", client)
158            }
159        }
160    }
161}
162
163/// Prefix for the item in the pool. Used to distinguish between Core and Platform clients.
164pub enum PoolPrefix {
165    Core,
166    Platform,
167}
168impl Display for PoolPrefix {
169    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
170        match self {
171            PoolPrefix::Core => write!(f, "Core"),
172            PoolPrefix::Platform => write!(f, "Platform"),
173        }
174    }
175}
176impl From<&PoolItem> for PoolPrefix {
177    fn from(item: &PoolItem) -> Self {
178        match item {
179            PoolItem::Core(_) => PoolPrefix::Core,
180            PoolItem::Platform(_) => PoolPrefix::Platform,
181        }
182    }
183}
184
185#[cfg(test)]
186mod tests {
187    use super::*;
188    use crate::RequestSettings;
189    use dapi_grpc::tonic::transport::Channel;
190    use std::str::FromStr;
191    use std::time::Duration;
192
193    fn test_uri() -> Uri {
194        Uri::from_str("http://127.0.0.1:3000").unwrap()
195    }
196
197    fn make_platform_pool_item() -> PoolItem {
198        let channel = Channel::builder(test_uri()).connect_lazy();
199        PoolItem::Platform(PlatformGrpcClient::new(channel))
200    }
201
202    fn make_core_pool_item() -> PoolItem {
203        let channel = Channel::builder(test_uri()).connect_lazy();
204        PoolItem::Core(CoreGrpcClient::new(channel))
205    }
206
207    #[test]
208    fn test_connection_pool_new() {
209        let pool = ConnectionPool::new(10);
210        let result = pool.get(PoolPrefix::Platform, &test_uri(), None);
211        assert!(result.is_none());
212    }
213
214    #[test]
215    fn test_connection_pool_default() {
216        let pool = ConnectionPool::default();
217        let result = pool.get(PoolPrefix::Core, &test_uri(), None);
218        assert!(result.is_none());
219    }
220
221    #[tokio::test]
222    async fn test_connection_pool_put_and_get_platform() {
223        let pool = ConnectionPool::new(10);
224        let uri = test_uri();
225        let item = make_platform_pool_item();
226
227        pool.put(&uri, None, item);
228
229        let result = pool.get(PoolPrefix::Platform, &uri, None);
230        assert!(result.is_some());
231        assert!(matches!(result.unwrap(), PoolItem::Platform(_)));
232    }
233
234    #[tokio::test]
235    async fn test_connection_pool_put_and_get_core() {
236        let pool = ConnectionPool::new(10);
237        let uri = test_uri();
238        let item = make_core_pool_item();
239
240        pool.put(&uri, None, item);
241
242        let result = pool.get(PoolPrefix::Core, &uri, None);
243        assert!(result.is_some());
244        assert!(matches!(result.unwrap(), PoolItem::Core(_)));
245    }
246
247    #[tokio::test]
248    async fn test_connection_pool_get_or_create_creates_new() {
249        let pool = ConnectionPool::new(10);
250        let uri = test_uri();
251
252        let result: Result<PoolItem, String> =
253            pool.get_or_create(PoolPrefix::Platform, &uri, None, || {
254                Ok(make_platform_pool_item())
255            });
256
257        assert!(result.is_ok());
258
259        // Second call should return cached version
260        let mut create_called = false;
261        let result2: Result<PoolItem, String> =
262            pool.get_or_create(PoolPrefix::Platform, &uri, None, || {
263                create_called = true;
264                Ok(make_platform_pool_item())
265            });
266
267        assert!(result2.is_ok());
268        assert!(
269            !create_called,
270            "create should not be called for cached item"
271        );
272    }
273
274    #[test]
275    fn test_connection_pool_get_or_create_error_not_cached() {
276        let pool = ConnectionPool::new(10);
277        let uri = test_uri();
278
279        let result: Result<PoolItem, String> =
280            pool.get_or_create(PoolPrefix::Platform, &uri, None, || {
281                Err("creation failed".to_string())
282            });
283
284        assert!(result.is_err());
285
286        // Pool should still be empty after failed creation
287        let cached = pool.get(PoolPrefix::Platform, &uri, None);
288        assert!(cached.is_none());
289    }
290
291    #[test]
292    fn test_pool_prefix_display() {
293        assert_eq!(format!("{}", PoolPrefix::Core), "Core");
294        assert_eq!(format!("{}", PoolPrefix::Platform), "Platform");
295    }
296
297    #[tokio::test]
298    async fn test_pool_prefix_from_pool_item() {
299        let platform_item = make_platform_pool_item();
300        let prefix: PoolPrefix = (&platform_item).into();
301        assert!(matches!(prefix, PoolPrefix::Platform));
302
303        let core_item = make_core_pool_item();
304        let prefix: PoolPrefix = (&core_item).into();
305        assert!(matches!(prefix, PoolPrefix::Core));
306    }
307
308    #[tokio::test]
309    async fn test_pool_item_from_platform_client() {
310        let channel = Channel::builder(test_uri()).connect_lazy();
311        let client = PlatformGrpcClient::new(channel);
312        let item: PoolItem = client.into();
313        assert!(matches!(item, PoolItem::Platform(_)));
314    }
315
316    #[tokio::test]
317    async fn test_pool_item_from_core_client() {
318        let channel = Channel::builder(test_uri()).connect_lazy();
319        let client = CoreGrpcClient::new(channel);
320        let item: PoolItem = client.into();
321        assert!(matches!(item, PoolItem::Core(_)));
322    }
323
324    #[tokio::test]
325    async fn test_pool_item_into_platform_client() {
326        let item = make_platform_pool_item();
327        let _client: PlatformGrpcClient = item.into();
328    }
329
330    #[tokio::test]
331    async fn test_pool_item_into_core_client() {
332        let item = make_core_pool_item();
333        let _client: CoreGrpcClient = item.into();
334    }
335
336    #[tokio::test]
337    #[should_panic(expected = "ClientType is not Platform")]
338    async fn test_pool_item_core_into_platform_panics() {
339        let item = make_core_pool_item();
340        let _client: PlatformGrpcClient = item.into();
341    }
342
343    #[tokio::test]
344    #[should_panic(expected = "ClientType is not Core")]
345    async fn test_pool_item_platform_into_core_panics() {
346        let item = make_platform_pool_item();
347        let _client: CoreGrpcClient = item.into();
348    }
349
350    #[tokio::test]
351    async fn test_connection_pool_shares_client_across_per_request_settings() {
352        let pool = ConnectionPool::new(10);
353        let uri = test_uri();
354
355        // Settings differing only in per-request knobs (timeout, retries,
356        // banning) must map to the same pooled connection...
357        let stored = RequestSettings {
358            timeout: Some(Duration::from_secs(30)),
359            retries: Some(3),
360            ban_failed_address: Some(false),
361            ..RequestSettings::default()
362        }
363        .finalize();
364        pool.put(&uri, Some(&stored), make_platform_pool_item());
365
366        let default = RequestSettings::default().finalize();
367        assert!(
368            pool.get(PoolPrefix::Platform, &uri, Some(&default))
369                .is_some(),
370            "per-request settings must not split pooled connections"
371        );
372
373        // ...while connection-affecting settings still get their own entry.
374        let connect = RequestSettings {
375            connect_timeout: Some(Duration::from_secs(3)),
376            ..RequestSettings::default()
377        }
378        .finalize();
379        assert!(
380            pool.get(PoolPrefix::Platform, &uri, Some(&connect))
381                .is_none(),
382            "connection-affecting settings must key separate connections"
383        );
384    }
385
386    #[tokio::test]
387    async fn test_connection_pool_different_prefixes_different_keys() {
388        let pool = ConnectionPool::new(10);
389        let uri = test_uri();
390
391        pool.put(&uri, None, make_platform_pool_item());
392
393        // Core prefix should not find a Platform item
394        let result = pool.get(PoolPrefix::Core, &uri, None);
395        assert!(result.is_none());
396
397        // Platform prefix should find it
398        let result = pool.get(PoolPrefix::Platform, &uri, None);
399        assert!(result.is_some());
400    }
401
402    #[tokio::test]
403    async fn test_connection_pool_clone_shares_data() {
404        let pool = ConnectionPool::new(10);
405        let pool_clone = pool.clone();
406        let uri = test_uri();
407
408        pool.put(&uri, None, make_platform_pool_item());
409
410        // Clone should see the same data
411        let result = pool_clone.get(PoolPrefix::Platform, &uri, None);
412        assert!(result.is_some());
413    }
414}
rs_dapi_client/connection_pool.rs

rs_dapi_client/
connection_pool.rs
