Update Service trait and use unified Counter (#455)

ntex-server/Cargo.toml

@@ -20,7 +20,7 @@ ntex-bytes = "0.1"
 ntex-net = "2"
 ntex-service = "3.3"
 ntex-rt = "0.4"
-ntex-util = "2"
+ntex-util = "2.5"
 
 async-channel = "2"
 async-broadcast = "0.7"

ntex-server/src/net/counter.rs

@@ -1,109 +0,0 @@
-use std::{cell::Cell, future::poll_fn, rc::Rc, task};
-
-use ntex_util::task::LocalWaker;
-
-#[derive(Debug)]
-/// Simple counter with ability to notify task on reaching specific number
-///
-/// Counter could be cloned, total count is shared across all clones.
-pub(super) struct Counter(Rc<CounterInner>);
-
-#[derive(Debug)]
-struct CounterInner {
-    count: Cell<usize>,
-    capacity: usize,
-    task: LocalWaker,
-}
-
-impl Counter {
-    /// Create `Counter` instance and set max value.
-    pub(super) fn new(capacity: usize) -> Self {
-        Counter(Rc::new(CounterInner {
-            capacity,
-            count: Cell::new(0),
-            task: LocalWaker::new(),
-        }))
-    }
-
-    /// Get counter guard.
-    pub(super) fn get(&self) -> CounterGuard {
-        CounterGuard::new(self.0.clone())
-    }
-
-    pub(crate) fn is_available(&self) -> bool {
-        self.0.count.get() < self.0.capacity
-    }
-
-    /// Check if counter is not at capacity. If counter at capacity
-    /// it registers notification for current task.
-    pub(crate) async fn available(&self) {
-        poll_fn(|cx| {
-            if self.0.available(cx) {
-                task::Poll::Ready(())
-            } else {
-                task::Poll::Pending
-            }
-        })
-        .await
-    }
-
-    pub(crate) async fn unavailable(&self) {
-        poll_fn(|cx| {
-            if self.0.available(cx) {
-                task::Poll::Pending
-            } else {
-                task::Poll::Ready(())
-            }
-        })
-        .await
-    }
-
-    /// Get total number of acquired counts
-    pub(super) fn total(&self) -> usize {
-        self.0.count.get()
-    }
-
-    pub(super) fn priv_clone(&self) -> Self {
-        Counter(self.0.clone())
-    }
-}
-
-pub(super) struct CounterGuard(Rc<CounterInner>);
-
-impl CounterGuard {
-    fn new(inner: Rc<CounterInner>) -> Self {
-        inner.inc();
-        CounterGuard(inner)
-    }
-}
-
-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 {
-            self.task.wake();
-        }
-    }
-
-    fn dec(&self) {
-        let num = self.count.get();
-        self.count.set(num - 1);
-        if num == self.capacity {
-            self.task.wake();
-        }
-    }
-
-    fn available(&self, cx: &mut task::Context<'_>) -> bool {
-        self.task.register(cx.waker());
-        self.count.get() < self.capacity
-    }
-}

ntex-server/src/net/mod.rs

@@ -1,10 +1,10 @@
 //! General purpose tcp server
+use ntex_util::services::Counter;
 use std::sync::atomic::{AtomicUsize, Ordering};
 
 mod accept;
 mod builder;
 mod config;
-mod counter;
 mod factory;
 mod service;
 mod socket;
@@ -56,8 +56,7 @@ pub enum SslError<E> {
 static MAX_CONNS: AtomicUsize = AtomicUsize::new(25600);
 
 thread_local! {
-    static MAX_CONNS_COUNTER: self::counter::Counter =
-        self::counter::Counter::new(MAX_CONNS.load(Ordering::Relaxed));
+    static MAX_CONNS_COUNTER: Counter = Counter::new(MAX_CONNS.load(Ordering::Relaxed));
 }
 
 /// Sets the maximum per-worker number of concurrent connections.
@@ -68,6 +67,7 @@ thread_local! {
 /// By default max connections is set to a 25k per worker.
 pub(super) fn max_concurrent_connections(num: usize) {
     MAX_CONNS.store(num, Ordering::Relaxed);
+    MAX_CONNS_COUNTER.with(|conns| conns.set_capacity(num));
 }
 
 pub(super) fn num_connections() -> usize {

ntex-server/src/net/service.rs

@@ -3,12 +3,11 @@ use std::{fmt, future::poll_fn, future::Future, pin::Pin, task::Poll};
 use ntex_bytes::{Pool, PoolRef};
 use ntex_net::Io;
 use ntex_service::{boxed, Service, ServiceCtx, ServiceFactory};
-use ntex_util::{future::join_all, HashMap};
+use ntex_util::{future::join_all, services::Counter, HashMap};
 
 use crate::ServerConfiguration;
 
 use super::accept::{AcceptNotify, AcceptorCommand};
-use super::counter::Counter;
 use super::factory::{FactoryServiceType, NetService, OnWorkerStart};
 use super::{socket::Connection, Token, MAX_CONNS_COUNTER};
 
@@ -135,7 +134,7 @@ impl ServiceFactory<Connection> for StreamService {
         Ok(StreamServiceImpl {
             tokens,
             services,
-            conns: MAX_CONNS_COUNTER.with(|conns| conns.priv_clone()),
+            conns: MAX_CONNS_COUNTER.with(|conns| conns.clone()),
         })
     }
 }