ntex-util 4.9.0

use std::{
    cell::Cell, cell::OnceCell, cell::RefCell, future::poll_fn, rc::Rc, task::Context, task::Poll,
};

use ntex_service::{BindCtx, ReadyNotifier};

use crate::{future::OnDropFn, task::LocalWaker};

/// A shared count with an asynchronous capacity notification.
///
/// Clones use the same count and capacity, but each clone registers its own
/// waiting task.
#[derive(Debug)]
pub struct Counter(usize, Rc<CounterInner>);

#[derive(Debug)]
struct CounterInner {
    count: Cell<usize>,
    capacity: Cell<usize>,
    tasks: RefCell<slab::Slab<LocalWaker>>,
    changed: OnceCell<ReadyNotifier>,
}

impl Counter {
    /// Creates a counter with the specified capacity.
    pub fn new(capacity: usize) -> Self {
        let mut tasks = slab::Slab::new();
        let idx = tasks.insert(LocalWaker::new());

        Counter(
            idx,
            Rc::new(CounterInner {
                count: Cell::new(0),
                capacity: Cell::new(capacity),
                tasks: RefCell::new(tasks),
                changed: OnceCell::new(),
            }),
        )
    }

    /// Acquires one count and returns a guard that releases it on drop.
    ///
    /// This does not check the capacity, so the count may exceed it. Use
    /// [`acquire`](Self::acquire) or [`try_get`](Self::try_get) when the
    /// capacity must not be exceeded.
    pub fn get(&self) -> CounterGuard {
        CounterGuard::new(self.1.clone())
    }

    /// Acquires one count if the counter has free capacity.
    pub fn try_get(&self) -> Option<CounterGuard> {
        if self.is_available() {
            Some(self.get())
        } else {
            None
        }
    }

    /// Waits until the counter has free capacity and acquires one count.
    ///
    /// The count is taken in the same poll that observes free capacity, so
    /// tasks acquiring concurrently never exceed the capacity. Concurrent
    /// waiters must use separate clones, each clone wakes only its most recent
    /// waiting task. Cancelling the future acquires nothing.
    pub async fn acquire(&self) -> CounterGuard {
        poll_fn(|cx| {
            if self.poll_available(cx) {
                Poll::Ready(self.get())
            } else {
                Poll::Pending
            }
        })
        .await
    }

    /// Changes the capacity and wakes waiting tasks.
    pub fn set_capacity(&self, cap: usize) {
        self.1.capacity.set(cap);
        self.1.notify();
    }

    /// Returns `true` if another count can be acquired without exceeding the capacity.
    pub fn is_available(&self) -> bool {
        self.1.count.get() < self.1.capacity.get()
    }

    /// Waits until the counter has free capacity.
    ///
    /// Returns immediately if there is capacity available. Otherwise,
    /// registers the current task for wakeup and waits until a slot is freed.
    ///
    /// This does not reserve the slot: another task may take it before a
    /// later [`get`](Self::get), unless `get` is called without awaiting in
    /// between. Use [`acquire`](Self::acquire) to wait and take a slot at once.
    pub async fn available(&self) {
        poll_fn(|cx| {
            if self.poll_available(cx) {
                Poll::Ready(())
            } else {
                Poll::Pending
            }
        })
        .await;
    }

    /// Waits until the counter reaches its capacity (i.e., becomes unavailable).
    pub async fn unavailable(&self) {
        poll_fn(|cx| {
            if self.is_available() {
                self.1.tasks.borrow()[self.0].register(cx.waker());
                Poll::Pending
            } else {
                Poll::Ready(())
            }
        })
        .await;
    }

    /// Check if counter is not at capacity. If counter at capacity
    /// it registers notification for current task.
    fn poll_available(&self, cx: &mut Context<'_>) -> bool {
        if self.1.count.get() < self.1.capacity.get() {
            true
        } else {
            let tasks = self.1.tasks.borrow();
            tasks[self.0].register(cx.waker());
            false
        }
    }

    /// Returns the number of currently held guards.
    pub fn total(&self) -> usize {
        self.1.count.get()
    }

    /// Subscribes `ctx` to capacity changes.
    ///
    /// Services whose readiness depends on this counter must call this in
    /// [`Service::bind`](ntex_service::Service::bind), otherwise their cached
    /// readiness is not invalidated when the counter fills or frees a slot.
    pub fn bind(&self, ctx: &BindCtx<'_>) {
        ctx.subscribe(self.1.changed.get_or_init(ReadyNotifier::new));
    }
}

impl Clone for Counter {
    fn clone(&self) -> Self {
        let idx = self.1.tasks.borrow_mut().insert(LocalWaker::new());
        Self(idx, self.1.clone())
    }
}

impl Drop for Counter {
    fn drop(&mut self) {
        self.1.tasks.borrow_mut().remove(self.0);
    }
}

#[derive(Debug)]
/// An acquired counter slot.
///
/// Dropping the guard releases the slot and wakes availability waiters.
pub struct CounterGuard(Rc<CounterInner>);

impl CounterGuard {
    fn new(inner: Rc<CounterInner>) -> Self {
        let guard = CounterGuard(inner);
        guard.0.inc();
        guard
    }
}

impl Unpin for CounterGuard {}

impl Drop for CounterGuard {
    fn drop(&mut self) {
        self.0.dec();
    }
}

impl CounterInner {
    fn inc(&self) {
        let num = self.count.get() + 1;
        self.count.set(num);
        if num == self.capacity.get() {
            self.notify();
        }
    }

    fn dec(&self) {
        let num = self.count.get();
        self.count.set(num - 1);
        if num == self.capacity.get() {
            self.notify();
        }
    }

    fn notify(&self) {
        let _wake = OnDropFn::new(|| {
            let tasks = self.tasks.borrow();
            for (_, task) in &*tasks {
                task.wake();
            }
        });
        if let Some(changed) = self.changed.get() {
            changed.notify();
        }
    }
}

#[cfg(test)]
mod tests {
    use std::time::Duration;

    use super::*;
    use crate::time::sleep;

    /// STATEMENT: Unavailability waiters wake only when the counter reaches capacity.
    ///
    /// HOW:
    /// - Spawn an unavailability waiter on a capacity-two counter.
    /// - Acquire guards one at a time and verify completion only after the second.
    #[ntex::test]
    async fn test_unavailable_is_woken() {
        let counter = Counter::new(2);
        let done = Rc::new(Cell::new(false));

        let (c, d) = (counter.clone(), done.clone());
        crate::spawn(async move {
            c.unavailable().await;
            d.set(true);
        });
        sleep(Duration::from_millis(10)).await;
        assert!(!done.get());

        let _g1 = counter.get();
        sleep(Duration::from_millis(10)).await;
        assert!(!done.get());

        let _g2 = counter.get();
        sleep(Duration::from_millis(10)).await;
        assert!(done.get());
    }

    /// STATEMENT: Releasing a guard wakes an availability waiter.
    ///
    /// HOW:
    /// - Fill a capacity-one counter and verify an availability waiter stays pending.
    /// - Drop the guard and observe the waiter complete.
    #[ntex::test]
    async fn test_available_is_woken() {
        let counter = Counter::new(1);
        let guard = counter.get();
        let done = Rc::new(Cell::new(false));

        let (c, d) = (counter.clone(), done.clone());
        crate::spawn(async move {
            c.available().await;
            d.set(true);
        });
        sleep(Duration::from_millis(10)).await;
        assert!(!done.get());

        drop(guard);
        sleep(Duration::from_millis(10)).await;
        assert!(done.get());
    }

    /// STATEMENT: Increasing capacity wakes waiters without changing the active count.
    ///
    /// HOW:
    /// - Hold the only slot and verify an availability task is pending.
    /// - Raise capacity, await the task, and drop the guard to check the count returns to zero.
    #[ntex::test]
    async fn test_set_capacity() {
        let counter = Counter::new(1);
        let guard = counter.get();
        assert_eq!(counter.total(), 1);
        assert!(!counter.is_available());

        let counter2 = counter.clone();
        let hnd = crate::spawn(async move { counter2.available().await });
        sleep(Duration::from_millis(10)).await;
        assert!(!hnd.is_finished());

        // raising the capacity wakes waiters
        counter.set_capacity(2);
        assert!(counter.is_available());
        crate::time::timeout(Duration::from_secs(1), hnd)
            .await
            .unwrap()
            .unwrap();

        drop(guard);
        assert_eq!(counter.total(), 0);
    }

    /// STATEMENT: Concurrent `acquire` calls never exceed the capacity, even when
    /// the acquiring tasks yield while holding their guards.
    ///
    /// HOW:
    /// - Hold the only slot of a capacity-one counter and spawn two tasks that `acquire`
    ///   through their own clones, record `total()`, sleep, then release.
    /// - Check `try_get` returns `None` while full, release the slot, and check both tasks
    ///   acquired with a recorded maximum `total()` of one, and `try_get` succeeds afterwards.
    #[ntex::test]
    async fn acquire_reserves_slot() {
        let counter = Counter::new(1);
        let guard = counter.get();
        assert!(counter.try_get().is_none());
        assert_eq!(counter.total(), 1);

        let max = Rc::new(Cell::new(0));
        let mut handles = Vec::new();
        for _ in 0..2 {
            let (c, max) = (counter.clone(), max.clone());
            handles.push(crate::spawn(async move {
                let _guard = c.acquire().await;
                max.set(max.get().max(c.total()));
                sleep(Duration::from_millis(10)).await;
            }));
        }
        sleep(Duration::from_millis(10)).await;
        assert_eq!(max.get(), 0);

        drop(guard);
        for hnd in handles {
            crate::time::timeout(Duration::from_secs(1), hnd)
                .await
                .unwrap()
                .unwrap();
        }
        assert_eq!(max.get(), 1);
        assert_eq!(counter.total(), 0);

        let guard = counter.try_get().unwrap();
        assert_eq!(counter.total(), 1);
        drop(guard);
        assert_eq!(counter.total(), 0);
    }
}