better api usability

ntex-bytes/CHANGELOG.md

@@ -1,5 +1,9 @@
 # Changes
 
+## 0.1.6 (2021-12-03)
+
+* Better api usability
+
 ## 0.1.5 (2021-12-02)
 
 * Split,freeze,truncate operations produce inline Bytes object if possible

ntex-bytes/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "ntex-bytes"
-version = "0.1.5"
+version = "0.1.6"
 license = "MIT"
 authors = ["Nikolay Kim <fafhrd91@gmail.com>", "Carl Lerche <me@carllerche.com>"]
 description = "Types and traits for working with bytes (bytes crate fork)"

ntex-bytes/src/bytes.rs

@@ -5,7 +5,7 @@ use std::sync::atomic::Ordering::{Acquire, Relaxed, Release};
 use std::sync::atomic::{self, AtomicUsize};
 use std::{cmp, fmt, hash, mem, ptr, ptr::NonNull, slice, usize};
 
-use crate::pool::{AsPoolRef, PoolId, PoolRef};
+use crate::pool::{PoolId, PoolRef};
 use crate::{buf::IntoIter, buf::UninitSlice, debug, Buf, BufMut};
 
 /// A reference counted contiguous slice of memory.
@@ -473,14 +473,17 @@ impl Bytes {
     }
 
     /// Creates `Bytes` instance from slice, by copying it.
-    pub fn copy_from_slice_in<T: AsPoolRef>(data: &[u8], pool: T) -> Self {
+    pub fn copy_from_slice_in<T>(data: &[u8], pool: T) -> Self
+    where
+        PoolRef: From<T>,
+    {
         if data.len() <= INLINE_CAP {
             Bytes {
                 inner: Inner::from_slice_inline(data),
             }
         } else {
             Bytes {
-                inner: BytesMut::copy_from_slice_in(data, pool.pool_ref()).inner,
+                inner: BytesMut::copy_from_slice_in(data, pool.into()).inner,
             }
         }
     }
@@ -1072,9 +1075,12 @@ impl BytesMut {
     /// assert!(PoolId::P1.pool_ref().allocated() > 0);
     /// ```
     #[inline]
-    pub fn with_capacity_in<T: AsPoolRef>(capacity: usize, pool: T) -> BytesMut {
+    pub fn with_capacity_in<T>(capacity: usize, pool: T) -> BytesMut
+    where
+        PoolRef: From<T>,
+    {
         BytesMut {
-            inner: Inner::with_capacity(capacity, pool.pool_ref()),
+            inner: Inner::with_capacity(capacity, pool.into()),
         }
     }
 
@@ -1086,8 +1092,11 @@ impl BytesMut {
     }
 
     /// Creates a new `BytesMut` from slice, by copying it.
-    pub fn copy_from_slice_in<T: AsPoolRef>(src: &[u8], pool: T) -> Self {
-        let mut bytes = BytesMut::with_capacity_in(src.len(), pool.pool_ref());
+    pub fn copy_from_slice_in<T>(src: &[u8], pool: T) -> Self
+    where
+        PoolRef: From<T>,
+    {
+        let mut bytes = BytesMut::with_capacity_in(src.len(), pool.into());
         bytes.extend_from_slice(src);
         bytes
     }
@@ -1100,9 +1109,12 @@ impl BytesMut {
 
     #[inline]
     /// Convert a `Vec` into a `BytesMut`
-    pub fn from_vec<T: AsPoolRef>(src: Vec<u8>, pool: T) -> BytesMut {
+    pub fn from_vec<T>(src: Vec<u8>, pool: T) -> BytesMut
+    where
+        PoolRef: From<T>,
+    {
         BytesMut {
-            inner: Inner::from_vec(src, pool.pool_ref()),
+            inner: Inner::from_vec(src, pool.into()),
         }
     }
 

ntex-bytes/src/pool.rs

@@ -19,10 +19,6 @@ pub struct PoolRef(&'static MemoryPool);
 #[derive(Copy, Clone, Debug)]
 pub struct PoolId(u8);
 
-pub trait AsPoolRef {
-    fn pool_ref(&self) -> PoolRef;
-}
-
 #[derive(Copy, Clone)]
 pub struct BufParams {
     pub high: u16,
@@ -94,8 +90,44 @@ impl PoolId {
         POOLS.with(|pools| PoolRef(pools[self.0 as usize]))
     }
 
+    #[inline]
+    /// Set max pool size
+    pub fn set_pool_size(self, size: usize) -> Self {
+        self.pool_ref().set_pool_size(size);
+        self
+    }
+
+    #[doc(hidden)]
+    #[inline]
+    pub fn set_read_params(self, h: u16, l: u16) -> Self {
+        self.pool_ref().set_read_params(h, l);
+        self
+    }
+
+    #[doc(hidden)]
+    #[inline]
+    pub fn set_write_params(self, h: u16, l: u16) -> Self {
+        self.pool_ref().set_write_params(h, l);
+        self
+    }
+
     /// Set future spawn fn
-    pub fn set_spawn_fn<T>(f: T)
+    pub fn set_spawn_fn<T>(self, f: T) -> Self
+    where
+        T: Fn(Pin<Box<dyn Future<Output = ()>>>) + 'static,
+    {
+        let spawn: Rc<dyn Fn(Pin<Box<dyn Future<Output = ()>>>)> =
+            Rc::new(move |fut| f(fut));
+
+        POOLS.with(move |pools| {
+            *pools[self.0 as usize].spawn.borrow_mut() = Some(spawn.clone());
+        });
+
+        self
+    }
+
+    /// Set future spawn fn to all pools
+    pub fn set_spawn_fn_all<T>(f: T)
     where
         T: Fn(Pin<Box<dyn Future<Output = ()>>>) + 'static,
     {
@@ -110,13 +142,6 @@ impl PoolId {
     }
 }
 
-impl AsPoolRef for PoolId {
-    #[inline]
-    fn pool_ref(&self) -> PoolRef {
-        POOLS.with(|pools| PoolRef(pools[self.0 as usize]))
-    }
-}
-
 thread_local! {
     static POOLS: [&'static MemoryPool; 16] = [
         MemoryPool::create(PoolId::P0),
@@ -171,6 +196,7 @@ impl PoolRef {
         BytesMut::with_capacity_in_priv(cap, self)
     }
 
+    #[doc(hidden)]
     #[inline]
     /// Set max pool size
     pub fn set_pool_size(self, size: usize) -> Self {
@@ -329,10 +355,17 @@ impl Default for PoolRef {
     }
 }
 
-impl AsPoolRef for PoolRef {
+impl From<PoolId> for PoolRef {
     #[inline]
-    fn pool_ref(&self) -> PoolRef {
-        *self
+    fn from(pid: PoolId) -> Self {
+        pid.pool_ref()
+    }
+}
+
+impl<'a> From<&'a Pool> for PoolRef {
+    #[inline]
+    fn from(pool: &'a Pool) -> Self {
+        PoolRef(pool.inner)
     }
 }
 
@@ -386,6 +419,7 @@ impl BufParams {
 }
 
 impl Clone for Pool {
+    #[inline]
     fn clone(&self) -> Pool {
         Pool {
             idx: Cell::new(0),
@@ -394,6 +428,20 @@ impl Clone for Pool {
     }
 }
 
+impl From<PoolId> for Pool {
+    #[inline]
+    fn from(pid: PoolId) -> Self {
+        pid.pool()
+    }
+}
+
+impl From<PoolRef> for Pool {
+    #[inline]
+    fn from(pref: PoolRef) -> Self {
+        pref.pool()
+    }
+}
+
 impl Drop for Pool {
     fn drop(&mut self) {
         let idx = self.idx.get();

ntex-bytes/tests/test_bytes.rs

@@ -655,11 +655,11 @@ fn pool() {
 async fn pool_usage() {
     use ntex::{time, util};
 
-    PoolId::set_spawn_fn(|f| {
+    PoolId::set_spawn_fn_all(|f| {
         let _ = ntex::rt::spawn(f);
     });
 
-    let p_ref = PoolId::P1.pool_ref().set_pool_size(10 * 1024);
+    let p_ref = PoolId::P1.set_pool_size(10 * 1024);
     let p1 = p_ref.pool();
     let p2 = p_ref.pool();