From 3f50734ab759426d87ccc81d2a2e2f4ee63d3e64 Mon Sep 17 00:00:00 2001 From: tison Date: Mon, 24 Aug 2026 03:23:19 +0800 Subject: [PATCH 1/3] refactor(pool): align internal organization --- asyncband/src/pool/bounded.rs | 2 +- asyncband/src/pool/common.rs | 2 +- asyncband/src/pool/mod.rs | 12 ++++++------ asyncband/src/pool/state.rs | 31 ++++++++++++++----------------- 4 files changed, 22 insertions(+), 25 deletions(-) diff --git a/asyncband/src/pool/bounded.rs b/asyncband/src/pool/bounded.rs index 25ea11f..b6f96ad 100644 --- a/asyncband/src/pool/bounded.rs +++ b/asyncband/src/pool/bounded.rs @@ -191,7 +191,7 @@ impl Pool { /// 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 + /// idle count while this method is running, so the target is the best effort rather than a /// postcondition. /// /// Returns the number of objects created. If [`ManageObject::create`] fails, objects created by diff --git a/asyncband/src/pool/common.rs b/asyncband/src/pool/common.rs index d6f54f2..c12ef2e 100644 --- a/asyncband/src/pool/common.rs +++ b/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. +/// Queue strategy when dequeue objects from the object pool. #[derive(Debug, Default, Clone, Copy)] pub enum QueueStrategy { /// First in first out. diff --git a/asyncband/src/pool/mod.rs b/asyncband/src/pool/mod.rs index 46e5567..40a1deb 100644 --- a/asyncband/src/pool/mod.rs +++ b/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; diff --git a/asyncband/src/pool/state.rs b/asyncband/src/pool/state.rs index 1a57427..fc53614 100644 --- a/asyncband/src/pool/state.rs +++ b/asyncband/src/pool/state.rs @@ -22,13 +22,13 @@ use crate::pool::QueueStrategy; use crate::pool::RetainResult; #[derive(Debug)] -pub(crate) struct ObjectState { - pub(crate) o: T, - pub(crate) status: ObjectStatus, +pub struct ObjectState { + pub o: T, + pub status: ObjectStatus, } impl ObjectState { - pub(crate) fn new(o: T) -> Self { + pub fn new(o: T) -> Self { Self { o, status: ObjectStatus::default(), @@ -37,48 +37,48 @@ impl ObjectState { } #[derive(Debug)] -pub(crate) struct PoolState { +pub struct PoolState { idle: VecDeque>, current_size: usize, } impl PoolState { - 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> { + pub fn pop(&mut self, strategy: QueueStrategy) -> Option> { match strategy { QueueStrategy::Fifo => self.idle.pop_front(), QueueStrategy::Lifo => self.idle.pop_back(), } } - pub(crate) fn add_idle(&mut self, state: ObjectState) { + pub fn add_idle(&mut self, state: ObjectState) { 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) { + pub fn return_idle(&mut self, state: ObjectState) { 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) @@ -86,10 +86,7 @@ impl PoolState { } /// 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 { + pub fn retain(&mut self, mut f: impl FnMut(&mut T, ObjectStatus) -> bool) -> RetainResult { let len = self.idle.len(); let mut retained = 0; let mut current = 0; From 01727d8bd7b16fc58eee0c91db3db698981b054b Mon Sep 17 00:00:00 2001 From: tison Date: Mon, 24 Aug 2026 03:30:54 +0800 Subject: [PATCH 2/3] refactor(pool): keep replenishment reservations local --- asyncband/src/pool/bounded.rs | 47 ++++++++++++++++-- asyncband/src/pool/common.rs | 2 +- asyncband/src/semaphore/mod.rs | 19 ------- .../tests/pool_replenish_test.rs | 49 +++++++++++++++++++ 4 files changed, 92 insertions(+), 25 deletions(-) diff --git a/asyncband/src/pool/bounded.rs b/asyncband/src/pool/bounded.rs index b6f96ad..1f61b66 100644 --- a/asyncband/src/pool/bounded.rs +++ b/asyncband/src/pool/bounded.rs @@ -191,19 +191,22 @@ impl Pool { /// 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 the best effort rather than a + /// 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 { - let Some(mut permit) = self.permits.clone().try_acquire_up_to_owned(target_idle) else { + 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 to_create = target_idle + .saturating_sub(idle_count) + .min(reservation.permits()); + reservation.release(reservation.permits() - to_create); let mut replenished = 0; for _ in 0..to_create { @@ -213,7 +216,7 @@ impl Pool { slots.add_idle(ObjectState::new(object)); } replenished += 1; - permit.release(1); + reservation.release(1); } Ok(replenished) @@ -354,6 +357,40 @@ impl Pool { } } +// 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 { + 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 diff --git a/asyncband/src/pool/common.rs b/asyncband/src/pool/common.rs index c12ef2e..90dffef 100644 --- a/asyncband/src/pool/common.rs +++ b/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 dequeue 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. diff --git a/asyncband/src/semaphore/mod.rs b/asyncband/src/semaphore/mod.rs index b6fe842..c812f51 100644 --- a/asyncband/src/semaphore/mod.rs +++ b/asyncband/src/semaphore/mod.rs @@ -310,15 +310,6 @@ impl Semaphore { } } - #[cfg(feature = "pool")] - pub(crate) fn try_acquire_up_to_owned( - self: Arc, - up_to: usize, - ) -> Option { - 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 diff --git a/tests-integration/tests/pool_replenish_test.rs b/tests-integration/tests/pool_replenish_test.rs index 08c8e22..0952108 100644 --- a/tests-integration/tests/pool_replenish_test.rs +++ b/tests-integration/tests/pool_replenish_test.rs @@ -200,3 +200,52 @@ async fn concurrent_get_and_replenish_to_respect_capacity() { drop((first, second)); assert_eq!(pool.status().idle_count, 2); } + +struct BlockingManager { + allow_create: Arc, +} + +impl ManageObject for BlockingManager { + type Object = (); + type Error = Infallible; + + async fn create(&self) -> Result { + 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 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); +} From 9568e5c0d7687f9b3de504643ae468b4b406c351 Mon Sep 17 00:00:00 2001 From: tison Date: Mon, 24 Aug 2026 04:16:34 +0800 Subject: [PATCH 3/3] fix(pool): preserve capacity across replenishment races --- asyncband/src/pool/bounded.rs | 26 +++++++-- .../tests/pool_replenish_test.rs | 55 +++++++++++++++++++ 2 files changed, 76 insertions(+), 5 deletions(-) diff --git a/asyncband/src/pool/bounded.rs b/asyncband/src/pool/bounded.rs index 1f61b66..a8385f7 100644 --- a/asyncband/src/pool/bounded.rs +++ b/asyncband/src/pool/bounded.rs @@ -190,22 +190,38 @@ impl Pool { /// /// 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 { + 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 (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(reservation.permits()) + .min(available_slots); reservation.release(reservation.permits() - to_create); let mut replenished = 0; diff --git a/tests-integration/tests/pool_replenish_test.rs b/tests-integration/tests/pool_replenish_test.rs index 0952108..280e7c9 100644 --- a/tests-integration/tests/pool_replenish_test.rs +++ b/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)); @@ -230,6 +263,28 @@ impl ManageObject for BlockingManager { } } +#[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));