orbit_client_pool/

pool.rs

use std::collections::{HashMap, VecDeque};
use std::future::Future;
use std::ops::{Deref, DerefMut};
use std::sync::{Arc, Mutex, MutexGuard};
use std::time::Instant;

use async_trait::async_trait;
use tokio::sync::Notify;
use tokio::time::Instant as TokioInstant;
use tokio_util::sync::CancellationToken;

use crate::{
    AcquireError, ClientPolicy, OptionsError, PoolDirective, PoolOptions,
};

#[path = "reservation.rs"]
mod reservation;
use reservation::Pending;

/// Creates and validates one protocol-specific outbound client.
///
/// These futures may be dropped before completion. The manager owns cleanup
/// of partially created protocol resources; the pool returns its reservation.
#[async_trait]
pub trait ClientManager: Send + Sync + 'static {
    type Client: Send + 'static;
    type Error: std::error::Error + Send + Sync + 'static;

    /// Establish a fresh client. Called only after capacity is reserved.
    async fn create(&self) -> Result<Self::Client, Self::Error>;

    /// Prove that an idle client may be handed out again.
    async fn recycle(
        &self,
        client: &mut Self::Client,
    ) -> Result<(), Self::Error>;
}

/// Current local ownership and queue counts.
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct PoolStats {
    pub live: usize,
    pub idle: usize,
    pub leased: usize,
    pub creating: usize,
    pub waiting: usize,
    /// Requests that currently hold a client, including a client being created.
    pub in_flight: usize,
    pub state: PoolState,
}

/// Admission and lifecycle state of one pool.
#[derive(Clone, Debug, Eq, PartialEq)]
pub enum PoolState {
    Open,
    Paused { reason: Arc<str> },
    Draining,
    Closed,
}

/// Shutdown behavior declared by the owner of one active request.
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum ActiveRequestPolicy {
    /// Let the request finish naturally while pool drain waits for its lease.
    Preserve,
    /// Signal the request on drain, immediately or after the declared grace.
    CancelOnDrain(Option<std::time::Duration>),
}

struct ActiveCancellation {
    policy: ActiveRequestPolicy,
    token: CancellationToken,
}

struct Entry<C> {
    client: C,
    created_at: Instant,
    idle_since: Instant,
}

struct State<C> {
    idle: VecDeque<Entry<C>>,
    live: usize,
    creating: usize,
    waiting: usize,
    in_flight: usize,
    status: PoolState,
    next_active_id: u64,
    drain_started: Option<Instant>,
    drain_cancellations: HashMap<u64, ActiveCancellation>,
}

struct Inner<M: ClientManager> {
    manager: M,
    options: PoolOptions,
    state: Mutex<State<M::Client>>,
    changed: Notify,
    admission_closed: CancellationToken,
    active_cancellation: CancellationToken,
}

/// A bounded process-local pool of protocol-specific clients.
pub struct ClientPool<M: ClientManager> {
    inner: Arc<Inner<M>>,
}

impl<M: ClientManager> Clone for ClientPool<M> {
    fn clone(&self) -> Self {
        Self { inner: Arc::clone(&self.inner) }
    }
}

impl<M: ClientManager> ClientPool<M> {
    /// Construct an empty pool. This performs no network IO.
    pub fn new(
        manager: M,
        options: PoolOptions,
    ) -> Result<Self, OptionsError> {
        options.validate()?;
        Ok(Self {
            inner: Arc::new(Inner {
                manager,
                options,
                state: Mutex::new(State {
                    idle: VecDeque::new(),
                    live: 0,
                    creating: 0,
                    waiting: 0,
                    in_flight: 0,
                    status: PoolState::Open,
                    next_active_id: 0,
                    drain_started: None,
                    drain_cancellations: HashMap::new(),
                }),
                changed: Notify::new(),
                admission_closed: CancellationToken::new(),
                active_cancellation: CancellationToken::new(),
            }),
        })
    }

    /// Acquire one client, creating it lazily or waiting for returned capacity.
    pub async fn acquire(
        &self,
        active_policy: ActiveRequestPolicy,
    ) -> Result<ClientLease<M>, AcquireError<M::Error>> {
        self.acquire_inner(active_policy, None).await
    }

    /// Acquire one client while also observing caller-owned cancellation.
    pub async fn acquire_cancellable(
        &self,
        active_policy: ActiveRequestPolicy,
        cancellation: &CancellationToken,
    ) -> Result<ClientLease<M>, AcquireError<M::Error>> {
        self.acquire_inner(active_policy, Some(cancellation)).await
    }

    async fn acquire_inner(
        &self,
        active_policy: ActiveRequestPolicy,
        cancellation: Option<&CancellationToken>,
    ) -> Result<ClientLease<M>, AcquireError<M::Error>> {
        let deadline = self
            .inner
            .options
            .acquire_timeout
            .map(|duration| TokioInstant::now() + duration);
        let mut waiter = Waiter::new(self);

        loop {
            if cancellation.is_some_and(CancellationToken::is_cancelled) {
                return Err(AcquireError::Cancelled);
            }
            let notified = self.inner.changed.notified();
            let action = {
                let mut state = self.state();
                self.prune_expired(&mut state);
                match &state.status {
                    PoolState::Draining | PoolState::Closed => {
                        return Err(AcquireError::Closed);
                    }
                    PoolState::Paused { .. } => {
                        waiter.enter(&mut state)?;
                        Action::Wait
                    }
                    PoolState::Open => {
                        if state.in_flight >= self.inner.options.max_concurrency
                        {
                            waiter.enter(&mut state)?;
                            Action::Wait
                        } else if let Some(entry) = state.idle.pop_front() {
                            state.in_flight += 1;
                            waiter.leave(&mut state);
                            Action::Reuse(entry)
                        } else if state.live < self.inner.options.max_live {
                            state.live += 1;
                            state.creating += 1;
                            state.in_flight += 1;
                            waiter.leave(&mut state);
                            Action::Create
                        } else {
                            waiter.enter(&mut state)?;
                            Action::Wait
                        }
                    }
                }
            };

            match action {
                Action::Reuse(entry) => {
                    let mut pending = Pending::recycle(self, entry);
                    match self
                        .operation(
                            self.inner.manager.recycle(pending.client()),
                            cancellation,
                        )
                        .await
                    {
                        Ok(Ok(())) => {}
                        Ok(Err(error)) => {
                            return Err(AcquireError::Manager(error));
                        }
                        Err(interrupted) => {
                            return Err(interrupted.acquire_error());
                        }
                    }
                    return Ok(pending.lease(active_policy));
                }
                Action::Create => {
                    let pending = Pending::create(self, true);
                    match self
                        .operation(self.inner.manager.create(), cancellation)
                        .await
                    {
                        Ok(Ok(client)) => {
                            return Ok(pending.created(client, active_policy));
                        }
                        Ok(Err(error)) => {
                            return Err(AcquireError::Manager(error));
                        }
                        Err(interrupted) => {
                            return Err(interrupted.acquire_error());
                        }
                    }
                }
                Action::Wait => {
                    self.wait_for_change(notified, deadline, cancellation)
                        .await?
                }
            }
        }
    }

    #[cfg(feature = "fleet")]
    pub(crate) async fn create(
        &self,
        active_policy: ActiveRequestPolicy,
    ) -> Result<ClientLease<M>, AcquireError<M::Error>> {
        self.create_inner(active_policy, None).await
    }

    #[cfg(feature = "fleet")]
    pub(crate) async fn create_cancellable(
        &self,
        active_policy: ActiveRequestPolicy,
        cancellation: &CancellationToken,
    ) -> Result<ClientLease<M>, AcquireError<M::Error>> {
        self.create_inner(active_policy, Some(cancellation)).await
    }

    #[cfg(feature = "fleet")]
    async fn create_inner(
        &self,
        active_policy: ActiveRequestPolicy,
        cancellation: Option<&CancellationToken>,
    ) -> Result<ClientLease<M>, AcquireError<M::Error>> {
        {
            let mut state = self.state();
            self.prune_expired(&mut state);
            if !matches!(&state.status, PoolState::Open) {
                return Err(AcquireError::Closed);
            }
            if state.live >= self.inner.options.max_live {
                return Err(AcquireError::AtCapacity);
            }
            if state.in_flight >= self.inner.options.max_concurrency {
                return Err(AcquireError::AtCapacity);
            }
            state.live += 1;
            state.creating += 1;
            state.in_flight += 1;
        }
        let pending = Pending::create(self, true);
        match self.operation(self.inner.manager.create(), cancellation).await {
            Ok(Ok(client)) => Ok(pending.created(client, active_policy)),
            Ok(Err(error)) => Err(AcquireError::Manager(error)),
            Err(interrupted) => Err(interrupted.acquire_error()),
        }
    }

    #[cfg(feature = "fleet")]
    pub(crate) async fn acquire_idle_where(
        &self,
        predicate: impl Fn(&M::Client) -> bool,
        active_policy: ActiveRequestPolicy,
    ) -> Result<Option<ClientLease<M>>, AcquireError<M::Error>> {
        self.acquire_idle_where_inner(predicate, active_policy, None).await
    }

    #[cfg(feature = "fleet")]
    pub(crate) async fn acquire_idle_where_cancellable(
        &self,
        predicate: impl Fn(&M::Client) -> bool,
        active_policy: ActiveRequestPolicy,
        cancellation: &CancellationToken,
    ) -> Result<Option<ClientLease<M>>, AcquireError<M::Error>> {
        self.acquire_idle_where_inner(
            predicate,
            active_policy,
            Some(cancellation),
        )
        .await
    }

    #[cfg(feature = "fleet")]
    async fn acquire_idle_where_inner(
        &self,
        predicate: impl Fn(&M::Client) -> bool,
        active_policy: ActiveRequestPolicy,
        cancellation: Option<&CancellationToken>,
    ) -> Result<Option<ClientLease<M>>, AcquireError<M::Error>> {
        let entry = {
            let mut state = self.state();
            self.prune_expired(&mut state);
            if !matches!(&state.status, PoolState::Open) {
                return Err(AcquireError::Closed);
            }
            if state.in_flight >= self.inner.options.max_concurrency {
                return Err(AcquireError::AtCapacity);
            }
            let Some(index) =
                state.idle.iter().position(|entry| predicate(&entry.client))
            else {
                return Ok(None);
            };
            state.in_flight += 1;
            state.idle.remove(index).expect("idle index came from this queue")
        };
        let mut pending = Pending::recycle(self, entry);
        match self
            .operation(
                self.inner.manager.recycle(pending.client()),
                cancellation,
            )
            .await
        {
            Ok(Ok(())) => {}
            Ok(Err(error)) => {
                return Err(AcquireError::Manager(error));
            }
            Err(interrupted) => {
                return Err(interrupted.acquire_error());
            }
        }
        Ok(Some(pending.lease(active_policy)))
    }

    #[cfg(feature = "fleet")]
    pub(crate) fn discard_idle_where(
        &self,
        predicate: impl Fn(&M::Client) -> bool,
    ) -> bool {
        let mut state = self.state();
        let Some(index) =
            state.idle.iter().position(|entry| predicate(&entry.client))
        else {
            return false;
        };
        let retired =
            state.idle.remove(index).expect("idle index came from this queue");
        state.live = state.live.saturating_sub(1);
        drop(state);
        drop(retired);
        self.inner.changed.notify_waiters();
        true
    }

    /// Maintain the declared warm floor until cancellation, then close and drain.
    pub async fn run(
        &self,
        cancellation: CancellationToken,
    ) -> Result<(), M::Error> {
        loop {
            if cancellation.is_cancelled() {
                self.drain().await;
                return Ok(());
            }

            self.prune_now();
            if matches!(
                self.stats().state,
                PoolState::Draining | PoolState::Closed
            ) {
                self.drain().await;
                return Ok(());
            }
            if self.reserve_warm_create() {
                let pending = Pending::create(self, false);
                match self
                    .operation(self.inner.manager.create(), Some(&cancellation))
                    .await
                {
                    Ok(Ok(client)) => pending.warm(client),
                    Ok(Err(error)) => {
                        return Err(error);
                    }
                    Err(_) => {
                        // Release the unfinished creation before drain waits
                        // for the pool's remaining live clients.
                        drop(pending);
                        self.drain().await;
                        return Ok(());
                    }
                }
                continue;
            }

            let notified = self.inner.changed.notified();
            let deadline = self.next_retirement_deadline();
            match deadline {
                Some(deadline) => {
                    tokio::select! {
                        _ = cancellation.cancelled() => {},
                        _ = notified => {},
                        _ = tokio::time::sleep_until(TokioInstant::from_std(deadline)) => {}
                    }
                }
                None => {
                    tokio::select! {
                        _ = cancellation.cancelled() => {},
                        _ = notified => {}
                    }
                }
            }
        }
    }

    /// Reject new acquisitions and drop retained idle clients.
    pub fn close(&self) {
        let mut state = self.state();
        state.status = PoolState::Closed;
        let idle = state.idle.len();
        let retired: Vec<_> = state.idle.drain(..).collect();
        state.live = state.live.saturating_sub(idle);
        drop(state);
        drop(retired);
        self.inner.admission_closed.cancel();
        self.inner.changed.notify_waiters();
    }

    /// Stop admission and wait indefinitely for every leased or creating client to end.
    pub async fn drain(&self) {
        let immediate = {
            let mut state = self.state();
            if !matches!(&state.status, PoolState::Closed) {
                state.status = PoolState::Draining;
            }
            let drain_started =
                *state.drain_started.get_or_insert_with(Instant::now);
            let idle = state.idle.len();
            let retired: Vec<_> = state.idle.drain(..).collect();
            state.live = state.live.saturating_sub(idle);
            let immediate = take_due_cancellations(&mut state, drain_started);
            drop(state);
            drop(retired);
            immediate
        };
        for cancellation in immediate {
            cancellation.cancel();
        }
        self.inner.admission_closed.cancel();
        self.inner.changed.notify_waiters();
        loop {
            let notified = self.inner.changed.notified();
            let (due, deadline) = {
                let mut state = self.state();
                if state.live == 0 {
                    state.status = PoolState::Closed;
                    drop(state);
                    self.inner.changed.notify_waiters();
                    return;
                }
                let now = Instant::now();
                let due = take_due_cancellations(&mut state, now);
                let deadline = next_active_cancellation_deadline(&state);
                (due, deadline)
            };
            for cancellation in due {
                cancellation.cancel();
            }
            match deadline {
                Some(deadline) => {
                    tokio::select! {
                        _ = notified => {},
                        _ = tokio::time::sleep_until(TokioInstant::from_std(deadline)) => {}
                    }
                }
                None => notified.await,
            }
        }
    }

    /// Pause new acquisitions without retiring retained or leased clients.
    pub fn pause(
        &self,
        reason: impl Into<Arc<str>>,
    ) {
        let mut state = self.state();
        if matches!(&state.status, PoolState::Open | PoolState::Paused { .. }) {
            state.status = PoolState::Paused { reason: reason.into() };
        }
        drop(state);
        self.inner.changed.notify_waiters();
    }

    /// Resume a policy-paused pool. Closing and draining are irreversible.
    pub fn resume(&self) -> bool {
        let mut state = self.state();
        if !matches!(&state.status, PoolState::Paused { .. }) {
            return false;
        }
        state.status = PoolState::Open;
        drop(state);
        self.inner.changed.notify_waiters();
        true
    }

    /// Apply a protocol-neutral admission consequence.
    pub fn apply(
        &self,
        directive: PoolDirective,
    ) {
        match directive {
            PoolDirective::Continue => {}
            PoolDirective::Pause { reason } => self.pause(reason),
        }
    }

    /// Explicitly cancel work holding a lease after admission has stopped.
    ///
    /// Graceful [`Self::drain`] never calls this. The runtime may call it only
    /// when its own policy has decided that active requests should no longer
    /// be preserved. Calling it while admission is open has no effect.
    pub fn cancel_active(&self) -> bool {
        if !matches!(
            &self.state().status,
            PoolState::Draining | PoolState::Closed
        ) {
            return false;
        }
        self.inner.active_cancellation.cancel();
        true
    }

    fn register_active(
        &self,
        policy: ActiveRequestPolicy,
    ) -> (CancellationToken, Option<u64>) {
        let token = self.inner.active_cancellation.child_token();
        if policy == ActiveRequestPolicy::Preserve {
            return (token, None);
        }

        let mut state = self.state();
        let id = state.next_active_id;
        state.next_active_id = state.next_active_id.wrapping_add(1);
        let cancel_now = match (policy, state.drain_started) {
            (ActiveRequestPolicy::CancelOnDrain(None), Some(_)) => true,
            (
                ActiveRequestPolicy::CancelOnDrain(Some(duration)),
                Some(started),
            ) => started
                .checked_add(duration)
                .is_some_and(|deadline| deadline <= Instant::now()),
            _ => false,
        };
        if cancel_now {
            token.cancel();
            (token, None)
        } else {
            state.drain_cancellations.insert(
                id,
                ActiveCancellation { policy, token: token.clone() },
            );
            drop(state);
            self.inner.changed.notify_waiters();
            (token, Some(id))
        }
    }

    fn unregister_active(
        &self,
        registration: Option<u64>,
    ) {
        let Some(id) = registration else {
            return;
        };
        self.state().drain_cancellations.remove(&id);
        self.inner.changed.notify_waiters();
    }

    #[cfg(feature = "fleet")]
    pub(crate) async fn wait_until_open(
        &self,
        deadline: Option<TokioInstant>,
        cancellation: Option<&CancellationToken>,
    ) -> Result<(), AcquireError<M::Error>> {
        let mut waiter = Waiter::new(self);
        loop {
            if cancellation.is_some_and(CancellationToken::is_cancelled) {
                return Err(AcquireError::Cancelled);
            }
            let notified = self.inner.changed.notified();
            {
                let mut state = self.state();
                match &state.status {
                    PoolState::Open => {
                        waiter.leave(&mut state);
                        return Ok(());
                    }
                    PoolState::Paused { .. } => waiter.enter(&mut state)?,
                    PoolState::Draining | PoolState::Closed => {
                        return Err(AcquireError::Closed);
                    }
                }
            }
            self.wait_for_change(notified, deadline, cancellation).await?;
        }
    }

    #[cfg(feature = "fleet")]
    pub(crate) async fn wait_until_open_for_maintenance(
        &self,
        cancellation: &CancellationToken,
    ) -> bool {
        loop {
            let notified = self.inner.changed.notified();
            match &self.state().status {
                PoolState::Open => return true,
                PoolState::Draining | PoolState::Closed => return false,
                PoolState::Paused { .. } => {}
            }
            tokio::select! {
                _ = notified => {},
                _ = cancellation.cancelled() => return false
            }
        }
    }

    pub fn stats(&self) -> PoolStats {
        let state = self.state();
        PoolStats {
            live: state.live,
            idle: state.idle.len(),
            leased: state
                .live
                .saturating_sub(state.idle.len() + state.creating),
            creating: state.creating,
            waiting: state.waiting,
            in_flight: state.in_flight,
            state: state.status.clone(),
        }
    }

    fn state(&self) -> MutexGuard<'_, State<M::Client>> {
        self.inner.state.lock().unwrap_or_else(|poison| poison.into_inner())
    }

    fn finish_create(
        &self,
        succeeded: bool,
        admitted: bool,
    ) {
        let mut state = self.state();
        state.creating = state.creating.saturating_sub(1);
        if !succeeded {
            state.live = state.live.saturating_sub(1);
            if admitted {
                state.in_flight = state.in_flight.saturating_sub(1);
            }
        }
        drop(state);
        self.inner.changed.notify_waiters();
    }

    fn finish_warm_create(
        &self,
        client: M::Client,
    ) {
        let now = Instant::now();
        let mut state = self.state();
        state.creating = state.creating.saturating_sub(1);
        if matches!(&state.status, PoolState::Draining | PoolState::Closed)
            || state.idle.len() >= self.inner.options.max_idle
        {
            state.live = state.live.saturating_sub(1);
        } else {
            state.idle.push_back(Entry {
                client,
                created_at: now,
                idle_since: now,
            });
        }
        drop(state);
        self.inner.changed.notify_waiters();
    }

    fn reserve_warm_create(&self) -> bool {
        let mut state = self.state();
        if !matches!(&state.status, PoolState::Open)
            || state.idle.len() + state.creating >= self.inner.options.min_idle
            || state.live >= self.inner.options.max_live
        {
            return false;
        }
        state.live += 1;
        state.creating += 1;
        true
    }

    fn retire_one(&self) {
        let mut state = self.state();
        state.live = state.live.saturating_sub(1);
        state.in_flight = state.in_flight.saturating_sub(1);
        drop(state);
        self.inner.changed.notify_waiters();
    }

    fn return_entry(
        &self,
        mut entry: Entry<M::Client>,
    ) {
        let now = Instant::now();
        let mut state = self.state();
        state.in_flight = state.in_flight.saturating_sub(1);
        let reusable =
            matches!(&state.status, PoolState::Open | PoolState::Paused { .. })
                && state.idle.len() < self.inner.options.max_idle
                && !expired(&entry, now, &self.inner.options);
        if reusable {
            entry.idle_since = now;
            state.idle.push_back(entry);
        } else {
            state.live = state.live.saturating_sub(1);
        }
        drop(state);
        self.inner.changed.notify_waiters();
    }

    fn prune_now(&self) {
        let mut state = self.state();
        self.prune_expired(&mut state);
    }

    fn prune_expired(
        &self,
        state: &mut State<M::Client>,
    ) {
        let now = Instant::now();
        let before = state.idle.len();
        state.idle.retain(|entry| !expired(entry, now, &self.inner.options));
        state.live = state.live.saturating_sub(before - state.idle.len());
    }

    fn next_retirement_deadline(&self) -> Option<Instant> {
        let state = self.state();
        state
            .idle
            .iter()
            .filter_map(|entry| retirement_deadline(entry, &self.inner.options))
            .min()
    }

    async fn wait_for_change<F>(
        &self,
        notified: F,
        deadline: Option<TokioInstant>,
        cancellation: Option<&CancellationToken>,
    ) -> Result<(), AcquireError<M::Error>>
    where
        F: Future<Output = ()>,
    {
        match (deadline, cancellation) {
            (Some(deadline), Some(cancellation)) => {
                tokio::select! {
                    _ = notified => Ok(()),
                    _ = cancellation.cancelled() => Err(AcquireError::Cancelled),
                    _ = tokio::time::sleep_until(deadline) => Err(AcquireError::Timeout(
                        self.inner.options.acquire_timeout.expect("deadline has duration")
                    ))
                }
            }
            (Some(deadline), None) => {
                tokio::select! {
                    _ = notified => Ok(()),
                    _ = tokio::time::sleep_until(deadline) => Err(AcquireError::Timeout(
                        self.inner.options.acquire_timeout.expect("deadline has duration")
                    ))
                }
            }
            (None, Some(cancellation)) => {
                tokio::select! {
                    _ = notified => Ok(()),
                    _ = cancellation.cancelled() => Err(AcquireError::Cancelled)
                }
            }
            (None, None) => {
                notified.await;
                Ok(())
            }
        }
    }

    async fn operation<F, T>(
        &self,
        operation: F,
        cancellation: Option<&CancellationToken>,
    ) -> Result<T, OperationInterrupted>
    where
        F: Future<Output = T>,
    {
        match cancellation {
            Some(cancellation) => tokio::select! {
                result = operation => Ok(result),
                _ = cancellation.cancelled() => Err(OperationInterrupted::Caller),
                _ = self.inner.admission_closed.cancelled() => {
                    Err(OperationInterrupted::AdmissionClosed)
                }
            },
            None => tokio::select! {
                result = operation => Ok(result),
                _ = self.inner.admission_closed.cancelled() => {
                    Err(OperationInterrupted::AdmissionClosed)
                }
            },
        }
    }
}

enum OperationInterrupted {
    Caller,
    AdmissionClosed,
}

impl OperationInterrupted {
    fn acquire_error<E>(self) -> AcquireError<E> {
        match self {
            Self::Caller => AcquireError::Cancelled,
            Self::AdmissionClosed => AcquireError::Closed,
        }
    }
}

enum Action<C> {
    Reuse(Entry<C>),
    Create,
    Wait,
}

struct Waiter<'a, M: ClientManager> {
    pool: &'a ClientPool<M>,
    entered: bool,
}

impl<'a, M: ClientManager> Waiter<'a, M> {
    fn new(pool: &'a ClientPool<M>) -> Self {
        Self { pool, entered: false }
    }

    fn enter(
        &mut self,
        state: &mut State<M::Client>,
    ) -> Result<(), AcquireError<M::Error>> {
        if self.entered {
            return Ok(());
        }
        if state.waiting >= self.pool.inner.options.max_waiting {
            return Err(AcquireError::QueueFull {
                max_waiting: self.pool.inner.options.max_waiting,
            });
        }
        state.waiting += 1;
        self.entered = true;
        Ok(())
    }

    fn leave(
        &mut self,
        state: &mut State<M::Client>,
    ) {
        if self.entered {
            state.waiting = state.waiting.saturating_sub(1);
            self.entered = false;
        }
    }
}

impl<M: ClientManager> Drop for Waiter<'_, M> {
    fn drop(&mut self) {
        if !self.entered {
            return;
        }
        let mut state = self.pool.state();
        state.waiting = state.waiting.saturating_sub(1);
        drop(state);
        self.pool.inner.changed.notify_waiters();
    }
}

/// Exclusive checkout of one local client.
pub struct ClientLease<M: ClientManager> {
    pool: ClientPool<M>,
    entry: Option<Entry<M::Client>>,
    active_cancellation: CancellationToken,
    active_registration: Option<u64>,
}

impl<M: ClientManager> ClientLease<M> {
    fn new(
        pool: ClientPool<M>,
        entry: Entry<M::Client>,
        active_policy: ActiveRequestPolicy,
    ) -> Self {
        let (active_cancellation, active_registration) =
            pool.register_active(active_policy);
        Self {
            pool,
            entry: Some(entry),
            active_cancellation,
            active_registration,
        }
    }

    /// Retire this client instead of returning it to the idle pool.
    pub fn discard(mut self) {
        self.entry.take();
        self.pool.unregister_active(self.active_registration.take());
        self.pool.retire_one();
    }

    /// Apply the application policy for one protocol-specific outcome.
    pub fn observe<O, P: ClientPolicy<O>>(
        &self,
        policy: &P,
        outcome: &O,
    ) {
        self.pool.apply(policy.classify(outcome));
    }

    /// Optional cancellation for active work holding this lease.
    ///
    /// Its request policy controls whether drain signals it immediately, after
    /// an explicit grace period, or only through [`ClientPool::cancel_active`].
    pub fn cancellation_token(&self) -> CancellationToken {
        self.active_cancellation.clone()
    }
}

impl<M: ClientManager> Deref for ClientLease<M> {
    type Target = M::Client;

    fn deref(&self) -> &Self::Target {
        &self.entry.as_ref().expect("client lease is live").client
    }
}

impl<M: ClientManager> DerefMut for ClientLease<M> {
    fn deref_mut(&mut self) -> &mut Self::Target {
        &mut self.entry.as_mut().expect("client lease is live").client
    }
}

impl<M: ClientManager> Drop for ClientLease<M> {
    fn drop(&mut self) {
        self.pool.unregister_active(self.active_registration.take());
        if let Some(entry) = self.entry.take() {
            self.pool.return_entry(entry);
        }
    }
}

fn take_due_cancellations<C>(
    state: &mut State<C>,
    now: Instant,
) -> Vec<CancellationToken> {
    let Some(started) = state.drain_started else {
        return Vec::new();
    };
    let due: Vec<_> = state
        .drain_cancellations
        .iter()
        .filter_map(|(id, active)| {
            let due = match active.policy {
                ActiveRequestPolicy::Preserve => false,
                ActiveRequestPolicy::CancelOnDrain(None) => true,
                ActiveRequestPolicy::CancelOnDrain(Some(grace)) => started
                    .checked_add(grace)
                    .is_some_and(|deadline| deadline <= now),
            };
            due.then_some(*id)
        })
        .collect();
    due.into_iter()
        .filter_map(|id| {
            state.drain_cancellations.remove(&id).map(|active| active.token)
        })
        .collect()
}

fn next_active_cancellation_deadline<C>(state: &State<C>) -> Option<Instant> {
    let started = state.drain_started?;
    state
        .drain_cancellations
        .values()
        .filter_map(|active| match active.policy {
            ActiveRequestPolicy::CancelOnDrain(Some(grace)) => {
                started.checked_add(grace)
            }
            ActiveRequestPolicy::Preserve
            | ActiveRequestPolicy::CancelOnDrain(None) => None,
        })
        .min()
}

fn expired<C>(
    entry: &Entry<C>,
    now: Instant,
    options: &PoolOptions,
) -> bool {
    retirement_deadline(entry, options).is_some_and(|deadline| deadline <= now)
}

fn retirement_deadline<C>(
    entry: &Entry<C>,
    options: &PoolOptions,
) -> Option<Instant> {
    let idle = options
        .idle_timeout
        .and_then(|duration| entry.idle_since.checked_add(duration));
    let lifetime = options
        .max_lifetime
        .and_then(|duration| entry.created_at.checked_add(duration));
    match (idle, lifetime) {
        (Some(idle), Some(lifetime)) => Some(idle.min(lifetime)),
        (Some(deadline), None) | (None, Some(deadline)) => Some(deadline),
        (None, None) => None,
    }
}

#[cfg(test)]
#[path = "../tests/pool.rs"]
mod tests;

#[cfg(test)]
#[path = "../tests/cancellation.rs"]
mod cancellation_tests;