fix(pool): preserve capacity across replenishment races (#181)

Author: tisonkun

asyncband/src/pool/bounded.rs

@@ -190,20 +190,39 @@ impl<M: ManageObject> Pool<M> {
     ///
     /// The pool reserves only capacity that is immediately available and never waits for
     /// checked-out objects. Existing idle objects count toward the target, and the pool's
-    /// maximum size is never exceeded. Concurrent calls and checkouts can change the observed
-    /// idle count while this method is running, so the target is best effort rather than a
-    /// postcondition.
+    /// maximum size is never exceeded. Targets above the maximum size are treated as the maximum.
+    /// Concurrent calls and checkouts can change the observed idle count while this method is
+    /// running, so the target is best effort rather than a postcondition.
     ///
     /// Returns the number of objects created. If [`ManageObject::create`] fails, objects created by
     /// this call before the failure remain in the pool and the error is returned.
     pub async fn replenish_to(&self, target_idle: usize) -> Result<usize, M::Error> {
-        let Some(mut permit) = self.permits.clone().try_acquire_up_to_owned(target_idle) else {
+        let target_idle = target_idle.min(self.config.max_size);
+        let Some(mut reservation) = ReplenishReservation::reserve_up_to(&self.permits, target_idle)
+        else {
             return Ok(0);
         };
 
-        let idle_count = self.slots.lock().idle_count();
-        let to_create = target_idle.saturating_sub(idle_count).min(permit.permits());
-        permit.release(permit.permits() - to_create);
+        let (idle_count, available_slots) = {
+            let slots = self.slots.lock();
+            let idle_count = slots.idle_count();
+
+            // Idle objects occupy pool slots without holding permits. Available permits plus this
+            // reservation represent capacity not committed to other checkouts, creations, or
+            // replenishments; subtracting idle objects leaves the slots this call may create.
+            let uncommitted_capacity = self
+                .permits
+                .available_permits()
+                .checked_add(reservation.permits())
+                .expect("invariant broken: semaphore capacity must not overflow");
+            let available_slots = uncommitted_capacity.saturating_sub(idle_count);
+            (idle_count, available_slots)
+        };
+        let to_create = target_idle
+            .saturating_sub(idle_count)
+            .min(reservation.permits())
+            .min(available_slots);
+        reservation.release(reservation.permits() - to_create);
 
         let mut replenished = 0;
         for _ in 0..to_create {
@@ -213,7 +232,7 @@ impl<M: ManageObject> Pool<M> {
                 slots.add_idle(ObjectState::new(object));
             }
             replenished += 1;
-            permit.release(1);
+            reservation.release(1);
         }
 
         Ok(replenished)
@@ -354,6 +373,40 @@ impl<M: ManageObject> Pool<M> {
     }
 }
 
+// Temporarily removes capacity while `replenish_to` creates objects. Idle objects do not consume
+// semaphore permits, so successful insertions release their reservation. Dropping the guard
+// restores any unfinished capacity after an error or cancellation.
+struct ReplenishReservation<'a> {
+    semaphore: &'a Semaphore,
+    permits: usize,
+}
+
+impl<'a> ReplenishReservation<'a> {
+    fn reserve_up_to(semaphore: &'a Semaphore, up_to: usize) -> Option<Self> {
+        let permits = semaphore.drain_permits(up_to);
+        (permits != 0).then_some(Self { semaphore, permits })
+    }
+
+    fn permits(&self) -> usize {
+        self.permits
+    }
+
+    fn release(&mut self, permits: usize) {
+        assert!(
+            permits <= self.permits,
+            "cannot release more permits than this reservation holds"
+        );
+        self.permits -= permits;
+        self.semaphore.release(permits);
+    }
+}
+
+impl Drop for ReplenishReservation<'_> {
+    fn drop(&mut self) {
+        self.semaphore.release(self.permits);
+    }
+}
+
 /// A wrapper of the actual pooled object.
 ///
 /// This object implements [`Deref`] and [`DerefMut`]. You can use it as if it was of type

asyncband/src/pool/common.rs

@@ -105,7 +105,7 @@ pub trait ManageObject: Send + Sync {
     fn on_detached(&self, _o: &mut Self::Object) {}
 }
 
-/// Queue strategy when dequeuing objects from the object pool.
+/// Strategy for dequeuing objects from the object pool.
 #[derive(Debug, Default, Clone, Copy)]
 pub enum QueueStrategy {
     /// First in first out.

asyncband/src/pool/mod.rs

@@ -171,14 +171,14 @@
 //! }
 //! ```
 
-pub use common::ManageObject;
-pub use common::ObjectStatus;
-pub use common::QueueStrategy;
-pub use common::RecycleCancelledStrategy;
-pub use common::RetainResult;
-
 mod common;
 mod state;
 
 pub mod bounded;
 pub mod unbounded;
+
+pub use self::common::ManageObject;
+pub use self::common::ObjectStatus;
+pub use self::common::QueueStrategy;
+pub use self::common::RecycleCancelledStrategy;
+pub use self::common::RetainResult;

asyncband/src/pool/state.rs

@@ -22,13 +22,13 @@ use crate::pool::QueueStrategy;
 use crate::pool::RetainResult;
 
 #[derive(Debug)]
-pub(crate) struct ObjectState<T> {
-    pub(crate) o: T,
-    pub(crate) status: ObjectStatus,
+pub struct ObjectState<T> {
+    pub o: T,
+    pub status: ObjectStatus,
 }
 
 impl<T> ObjectState<T> {
-    pub(crate) fn new(o: T) -> Self {
+    pub fn new(o: T) -> Self {
         Self {
             o,
             status: ObjectStatus::default(),
@@ -37,59 +37,56 @@ impl<T> ObjectState<T> {
 }
 
 #[derive(Debug)]
-pub(crate) struct PoolState<T> {
+pub struct PoolState<T> {
     idle: VecDeque<ObjectState<T>>,
     current_size: usize,
 }
 
 impl<T> PoolState<T> {
-    pub(crate) const fn new() -> Self {
+    pub const fn new() -> Self {
         Self {
             idle: VecDeque::new(),
             current_size: 0,
         }
     }
 
-    pub(crate) fn current_size(&self) -> usize {
+    pub fn current_size(&self) -> usize {
         self.current_size
     }
 
-    pub(crate) fn idle_count(&self) -> usize {
+    pub fn idle_count(&self) -> usize {
         self.idle.len()
     }
 
-    pub(crate) fn pop(&mut self, strategy: QueueStrategy) -> Option<ObjectState<T>> {
+    pub fn pop(&mut self, strategy: QueueStrategy) -> Option<ObjectState<T>> {
         match strategy {
             QueueStrategy::Fifo => self.idle.pop_front(),
             QueueStrategy::Lifo => self.idle.pop_back(),
         }
     }
 
-    pub(crate) fn add_idle(&mut self, state: ObjectState<T>) {
+    pub fn add_idle(&mut self, state: ObjectState<T>) {
         self.current_size += 1;
         self.idle.push_back(state);
     }
 
-    pub(crate) fn add_active(&mut self) {
+    pub fn add_active(&mut self) {
         self.current_size += 1;
     }
 
-    pub(crate) fn return_idle(&mut self, state: ObjectState<T>) {
+    pub fn return_idle(&mut self, state: ObjectState<T>) {
         self.idle.push_back(state);
     }
 
-    pub(crate) fn detach(&mut self) {
+    pub fn detach(&mut self) {
         self.current_size = self
             .current_size
             .checked_sub(1)
             .expect("detached object must belong to the pool");
     }
 
     /// Retains matching idle objects without losing any object if the predicate panics.
-    pub(crate) fn retain(
-        &mut self,
-        mut f: impl FnMut(&mut T, ObjectStatus) -> bool,
-    ) -> RetainResult<T> {
+    pub fn retain(&mut self, mut f: impl FnMut(&mut T, ObjectStatus) -> bool) -> RetainResult<T> {
         let len = self.idle.len();
         let mut retained = 0;
         let mut current = 0;

asyncband/src/semaphore/mod.rs

@@ -310,15 +310,6 @@ impl Semaphore {
         }
     }
 
-    #[cfg(feature = "pool")]
-    pub(crate) fn try_acquire_up_to_owned(
-        self: Arc<Self>,
-        up_to: usize,
-    ) -> Option<OwnedSemaphorePermit> {
-        let permits = self.s.drain_permits(up_to);
-        (permits != 0).then_some(OwnedSemaphorePermit { sem: self, permits })
-    }
-
     /// Acquires `n` permits from the semaphore.
     ///
     /// The semaphore must be wrapped in an [`Arc`] to call this method.
@@ -518,16 +509,6 @@ pub struct OwnedSemaphorePermit {
 }
 
 impl OwnedSemaphorePermit {
-    #[cfg(feature = "pool")]
-    pub(crate) fn release(&mut self, permits: usize) {
-        assert!(
-            permits <= self.permits,
-            "cannot release more permits than this permit holds"
-        );
-        self.permits -= permits;
-        self.sem.release(permits);
-    }
-
     /// Forgets the permit **without** releasing it back to the semaphore.
     ///
     /// This can be used to permanently reduce the number of permits available

tests-integration/tests/pool_replenish_test.rs

@@ -162,6 +162,39 @@ impl ManageObject for ControlledManager {
     }
 }
 
+#[tokio::test]
+async fn concurrent_replenish_to_calls_respect_capacity() {
+    let calls = Arc::new(AtomicUsize::new(0));
+    let allow_create = Arc::new(AtomicBool::new(false));
+    let pool = Pool::new(
+        PoolConfig::new(2),
+        ControlledManager {
+            calls: calls.clone(),
+            allow_create: allow_create.clone(),
+        },
+    );
+
+    assert_eq!(pool.replenish_to(1).await, Ok(1));
+
+    let mut first = Box::pin(pool.replenish_to(2));
+    assert!(tests_integration::poll_once(first.as_mut()).is_pending());
+
+    let mut second = Box::pin(pool.replenish_to(2));
+    assert_eq!(
+        tests_integration::poll_once(second.as_mut()),
+        Poll::Ready(Ok(0))
+    );
+
+    allow_create.store(true, Ordering::Release);
+    assert_eq!(
+        tests_integration::poll_once(first.as_mut()),
+        Poll::Ready(Ok(1))
+    );
+    assert_eq!(calls.load(Ordering::Relaxed), 2);
+    assert_eq!(pool.status().current_size, 2);
+    assert_eq!(pool.status().idle_count, 2);
+}
+
 #[tokio::test]
 async fn concurrent_get_and_replenish_to_respect_capacity() {
     let calls = Arc::new(AtomicUsize::new(0));
@@ -200,3 +233,74 @@ async fn concurrent_get_and_replenish_to_respect_capacity() {
     drop((first, second));
     assert_eq!(pool.status().idle_count, 2);
 }
+
+struct BlockingManager {
+    allow_create: Arc<AtomicBool>,
+}
+
+impl ManageObject for BlockingManager {
+    type Object = ();
+    type Error = Infallible;
+
+    async fn create(&self) -> Result<Self::Object, Self::Error> {
+        poll_fn(|_| {
+            if self.allow_create.load(Ordering::Acquire) {
+                Poll::Ready(())
+            } else {
+                Poll::Pending
+            }
+        })
+        .await;
+        Ok(())
+    }
+
+    async fn is_recyclable(
+        &self,
+        _object: &mut Self::Object,
+        _status: &ObjectStatus,
+    ) -> Result<(), Self::Error> {
+        Ok(())
+    }
+}
+
+#[tokio::test]
+async fn replenish_to_respects_max_size_with_active_and_idle_objects() {
+    let pool = Pool::new(
+        PoolConfig::new(2),
+        BlockingManager {
+            allow_create: Arc::new(AtomicBool::new(true)),
+        },
+    );
+
+    assert_eq!(pool.replenish_to(2).await, Ok(2));
+    let active = pool.get().await.unwrap();
+    assert_eq!(pool.status().current_size, 2);
+    assert_eq!(pool.status().idle_count, 1);
+
+    assert_eq!(pool.replenish_to(usize::MAX).await, Ok(0));
+    assert_eq!(pool.status().current_size, 2);
+    assert_eq!(pool.status().idle_count, 1);
+
+    drop(active);
+    assert_eq!(pool.status().idle_count, 2);
+}
+
+#[tokio::test]
+async fn cancelling_replenish_to_releases_reserved_capacity() {
+    let allow_create = Arc::new(AtomicBool::new(false));
+    let pool = Pool::new(
+        PoolConfig::new(1),
+        BlockingManager {
+            allow_create: allow_create.clone(),
+        },
+    );
+
+    let mut replenish = Box::pin(pool.replenish_to(1));
+    assert!(tests_integration::poll_once(replenish.as_mut()).is_pending());
+    drop(replenish);
+
+    allow_create.store(true, Ordering::Release);
+    let mut get = Box::pin(pool.get());
+    assert!(tests_integration::poll_once(get.as_mut()).is_ready());
+    assert_eq!(pool.status().idle_count, 1);
+}