liminal_sdk/connection/
pool.rs

1use alloc::format;
2use alloc::vec::Vec;
3
4use crate::SdkError;
5
6use super::{
7    ConnectionEvent, ConnectionLifecycle, ResumeRequest, SubscriptionId, SubscriptionRecovery,
8};
9
10/// Caller-supplied connection pool sizing and resource configuration.
11///
12/// This type deliberately has no [`Default`] implementation: pool sizing must be
13/// supplied by the caller, builder, or runtime so the SDK never bakes in a hidden
14/// connection-count default.
15#[derive(Clone, Copy, Debug, PartialEq, Eq)]
16pub struct ConnectionPoolConfig {
17    /// Maximum number of remote connections managed by this pool.
18    pub max_connections: usize,
19    /// Per-connection operation timeout, in milliseconds.
20    pub timeout_millis: u64,
21    /// Per-connection inbound buffer size.
22    pub buffer_size: usize,
23}
24
25impl ConnectionPoolConfig {
26    /// Creates pool configuration from caller-supplied values.
27    #[must_use]
28    pub const fn new(max_connections: usize, timeout_millis: u64, buffer_size: usize) -> Self {
29        Self {
30            max_connections,
31            timeout_millis,
32            buffer_size,
33        }
34    }
35
36    /// Validates caller-supplied pool configuration.
37    ///
38    /// # Errors
39    ///
40    /// Returns [`SdkError`] when no connection can be allocated.
41    pub fn validate(self) -> Result<Self, SdkError> {
42        if self.max_connections == 0 {
43            return Err(SdkError::Connection {
44                description: "connection pool max_connections must be greater than zero".into(),
45            });
46        }
47
48        Ok(self)
49    }
50}
51
52/// Stable identifier for an internally managed pooled connection.
53#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)]
54pub struct PoolConnectionId(usize);
55
56impl PoolConnectionId {
57    /// Creates a pooled connection identifier.
58    #[must_use]
59    pub const fn new(value: usize) -> Self {
60        Self(value)
61    }
62
63    /// Returns the numeric connection slot.
64    #[must_use]
65    pub const fn get(self) -> usize {
66        self.0
67    }
68}
69
70/// Assignment returned when a subscription is placed on a pooled connection.
71#[derive(Clone, Copy, Debug, PartialEq, Eq)]
72pub struct SubscriptionAssignment {
73    /// Subscription assigned to the pool.
74    pub subscription_id: SubscriptionId,
75    /// Connection that owns the subscription.
76    pub connection_id: PoolConnectionId,
77}
78
79#[derive(Debug)]
80struct PoolConnection {
81    id: PoolConnectionId,
82    subscription_count: usize,
83    lifecycle: ConnectionLifecycle,
84    recovery: SubscriptionRecovery,
85}
86
87impl PoolConnection {
88    fn new(id: PoolConnectionId) -> Self {
89        Self {
90            id,
91            subscription_count: 0,
92            lifecycle: ConnectionLifecycle::new(),
93            recovery: SubscriptionRecovery::new(),
94        }
95    }
96}
97
98/// Configurable pool for remote SDK connections.
99#[derive(Debug)]
100pub struct ConnectionPool {
101    config: ConnectionPoolConfig,
102    connections: Vec<PoolConnection>,
103}
104
105impl ConnectionPool {
106    /// Creates a pool with exactly the caller-supplied connection count.
107    ///
108    /// # Errors
109    ///
110    /// Returns [`SdkError`] if the supplied configuration is invalid.
111    pub fn new(config: ConnectionPoolConfig) -> Result<Self, SdkError> {
112        let config = config.validate()?;
113        let mut connections = Vec::with_capacity(config.max_connections);
114
115        for slot in 0..config.max_connections {
116            connections.push(PoolConnection::new(PoolConnectionId::new(slot)));
117        }
118
119        Ok(Self {
120            config,
121            connections,
122        })
123    }
124
125    /// Returns the caller-supplied pool configuration.
126    #[must_use]
127    pub const fn config(&self) -> ConnectionPoolConfig {
128        self.config
129    }
130
131    /// Returns the caller-supplied maximum connection count.
132    #[must_use]
133    pub const fn max_connections(&self) -> usize {
134        self.config.max_connections
135    }
136
137    /// Returns the number of managed connection slots.
138    #[must_use]
139    pub fn connection_count(&self) -> usize {
140        self.connections.len()
141    }
142
143    /// Assigns a subscription to the least-loaded pooled connection.
144    ///
145    /// # Errors
146    ///
147    /// Returns [`SdkError`] if the pool has no available connection entries.
148    pub fn assign_subscription(
149        &mut self,
150        subscription_id: SubscriptionId,
151    ) -> Result<SubscriptionAssignment, SdkError> {
152        if let Some(existing) = self.connection_for_subscription(subscription_id) {
153            return Ok(SubscriptionAssignment {
154                subscription_id,
155                connection_id: existing,
156            });
157        }
158
159        let connection = self
160            .connections
161            .iter_mut()
162            .min_by_key(|connection| (connection.subscription_count, connection.id))
163            .ok_or_else(|| SdkError::Connection {
164                description: "connection pool has no connections".into(),
165            })?;
166
167        connection.subscription_count = connection.subscription_count.saturating_add(1);
168        connection.recovery.track_subscription(subscription_id);
169
170        Ok(SubscriptionAssignment {
171            subscription_id,
172            connection_id: connection.id,
173        })
174    }
175
176    /// Records an acknowledged sequence for the connection that owns a subscription.
177    ///
178    /// # Errors
179    ///
180    /// Returns [`SdkError`] when the subscription is not active in this pool.
181    pub fn acknowledge(
182        &mut self,
183        subscription_id: SubscriptionId,
184        sequence: u64,
185    ) -> Result<(), SdkError> {
186        let connection = self.connection_for_subscription_mut(subscription_id)?;
187        connection.recovery.acknowledge(subscription_id, sequence);
188        Ok(())
189    }
190
191    /// Removes a subscription assignment and its recovery state.
192    ///
193    /// # Errors
194    ///
195    /// Returns [`SdkError`] when the subscription is not active in this pool.
196    pub fn unsubscribe(&mut self, subscription_id: SubscriptionId) -> Result<(), SdkError> {
197        let connection = self.connection_for_subscription_mut(subscription_id)?;
198        connection.recovery.unsubscribe(subscription_id);
199        connection.subscription_count = connection.subscription_count.saturating_sub(1);
200        Ok(())
201    }
202
203    /// Builds subscription resume requests for every pooled connection on reconnect.
204    ///
205    /// # Errors
206    ///
207    /// Returns [`SdkError`] if any active subscription cannot compute a resume sequence.
208    pub fn resume_requests_for_transition(
209        &self,
210        event: &ConnectionEvent,
211    ) -> Result<Vec<ResumeRequest>, SdkError> {
212        let mut requests = Vec::new();
213        for connection in &self.connections {
214            requests.extend(connection.recovery.resume_requests_for_transition(event)?);
215        }
216        Ok(requests)
217    }
218
219    /// Returns the connection assigned to a subscription, if it is active.
220    #[must_use]
221    pub fn connection_for_subscription(
222        &self,
223        subscription_id: SubscriptionId,
224    ) -> Option<PoolConnectionId> {
225        self.connections
226            .iter()
227            .find(|connection| connection.recovery.is_active(subscription_id))
228            .map(|connection| connection.id)
229    }
230
231    /// Returns the number of active subscriptions on a connection.
232    ///
233    /// # Errors
234    ///
235    /// Returns [`SdkError`] if the connection identifier is not part of this pool.
236    pub fn subscription_count(&self, connection_id: PoolConnectionId) -> Result<usize, SdkError> {
237        self.connection(connection_id)
238            .map(|connection| connection.subscription_count)
239    }
240
241    /// Returns recovery state for a pooled connection.
242    ///
243    /// # Errors
244    ///
245    /// Returns [`SdkError`] if the connection identifier is not part of this pool.
246    pub fn recovery(
247        &self,
248        connection_id: PoolConnectionId,
249    ) -> Result<&SubscriptionRecovery, SdkError> {
250        self.connection(connection_id)
251            .map(|connection| &connection.recovery)
252    }
253
254    /// Returns lifecycle state for a pooled connection.
255    ///
256    /// # Errors
257    ///
258    /// Returns [`SdkError`] if the connection identifier is not part of this pool.
259    pub fn lifecycle(
260        &self,
261        connection_id: PoolConnectionId,
262    ) -> Result<&ConnectionLifecycle, SdkError> {
263        self.connection(connection_id)
264            .map(|connection| &connection.lifecycle)
265    }
266
267    fn connection(&self, connection_id: PoolConnectionId) -> Result<&PoolConnection, SdkError> {
268        self.connections
269            .iter()
270            .find(|connection| connection.id == connection_id)
271            .ok_or_else(|| SdkError::Connection {
272                description: format!("unknown pooled connection {}", connection_id.get()),
273            })
274    }
275
276    fn connection_for_subscription_mut(
277        &mut self,
278        subscription_id: SubscriptionId,
279    ) -> Result<&mut PoolConnection, SdkError> {
280        self.connections
281            .iter_mut()
282            .find(|connection| connection.recovery.is_active(subscription_id))
283            .ok_or_else(|| SdkError::Connection {
284                description: format!("unknown subscription {}", subscription_id.get()),
285            })
286    }
287}
288
289#[cfg(test)]
290mod tests {
291    use super::*;
292    use crate::connection::{ConnectionState, DisconnectReason};
293
294    #[test]
295    fn invalid_pool_size_is_rejected() {
296        let config = ConnectionPoolConfig::new(0, 10, 16);
297
298        assert!(ConnectionPool::new(config).is_err());
299    }
300
301    #[test]
302    fn subscriptions_are_distributed_across_connections() -> Result<(), SdkError> {
303        let config = ConnectionPoolConfig::new(2, 10, 16);
304        let mut pool = ConnectionPool::new(config)?;
305
306        let first = pool.assign_subscription(SubscriptionId::new(1))?;
307        let second = pool.assign_subscription(SubscriptionId::new(2))?;
308        let third = pool.assign_subscription(SubscriptionId::new(3))?;
309
310        assert_ne!(first.connection_id, second.connection_id);
311        assert_eq!(third.connection_id, first.connection_id);
312        assert_eq!(pool.subscription_count(first.connection_id)?, 2);
313        assert_eq!(pool.subscription_count(second.connection_id)?, 1);
314        Ok(())
315    }
316
317    #[test]
318    fn multiple_subscriptions_share_configured_connections() -> Result<(), SdkError> {
319        let config = ConnectionPoolConfig::new(1, 20, 32);
320        let mut pool = ConnectionPool::new(config)?;
321
322        let first = pool.assign_subscription(SubscriptionId::new(10))?;
323        let second = pool.assign_subscription(SubscriptionId::new(11))?;
324
325        assert_eq!(first.connection_id, second.connection_id);
326        assert_eq!(pool.max_connections(), 1);
327        assert_eq!(pool.subscription_count(first.connection_id)?, 2);
328        Ok(())
329    }
330
331    #[test]
332    fn pooled_recovery_builds_resume_requests_on_reconnect() -> Result<(), SdkError> {
333        let config = ConnectionPoolConfig::new(2, 10, 16);
334        let mut pool = ConnectionPool::new(config)?;
335        let first = SubscriptionId::new(21);
336        let second = SubscriptionId::new(22);
337        let event = ConnectionEvent::new(ConnectionState::Reconnecting, ConnectionState::Connected);
338
339        pool.assign_subscription(first)?;
340        pool.assign_subscription(second)?;
341        pool.acknowledge(first, 4)?;
342
343        let requests = pool.resume_requests_for_transition(&event)?;
344
345        assert_eq!(
346            requests,
347            vec![ResumeRequest::new(first, 5), ResumeRequest::new(second, 0)]
348        );
349        Ok(())
350    }
351
352    #[test]
353    fn unsubscribe_removes_assignment() -> Result<(), SdkError> {
354        let config = ConnectionPoolConfig::new(2, 10, 16);
355        let mut pool = ConnectionPool::new(config)?;
356        let subscription_id = SubscriptionId::new(31);
357        let assignment = pool.assign_subscription(subscription_id)?;
358
359        pool.unsubscribe(subscription_id)?;
360
361        assert_eq!(pool.connection_for_subscription(subscription_id), None);
362        assert_eq!(pool.subscription_count(assignment.connection_id)?, 0);
363        assert!(pool.unsubscribe(subscription_id).is_err());
364        Ok(())
365    }
366
367    #[test]
368    fn non_reconnect_transition_does_not_resume() -> Result<(), SdkError> {
369        let config = ConnectionPoolConfig::new(2, 10, 16);
370        let mut pool = ConnectionPool::new(config)?;
371        let event = ConnectionEvent::new(
372            ConnectionState::Connected,
373            ConnectionState::Disconnected {
374                reason: DisconnectReason::Normal,
375            },
376        );
377
378        pool.assign_subscription(SubscriptionId::new(41))?;
379
380        assert!(pool.resume_requests_for_transition(&event)?.is_empty());
381        Ok(())
382    }
383}
liminal_sdk/connection/pool.rs

liminal_sdk/connection/
pool.rs
