ntex: optimize allocations

Author: fafhrd91

frameworks/Rust/ntex/Cargo.toml

@@ -1,5 +1,5 @@
 [package]
-name = "ntex-bench"
+name = "ntex-framework"
 version = "3.0.0"
 edition = "2018"
 
@@ -60,17 +60,15 @@ tokio = ["ntex/tokio"]
 # compio runtime
 compio = ["ntex/compio"]
 
-# neon runtime
-neon = ["ntex/neon"]
+# neon polling runtime
+neon = ["ntex/neon-polling"]
 
-# neon runtime
+# neon io-uring runtime
 neon-uring = ["ntex/neon-uring"]
 
 [dependencies]
-ntex = "3.0.0-pre.5"
-ntex-neon = "0.1.35"
-ntex-net = "3.0.0"
-ntex-bytes = { version = "1", features=["simd"] }
+ntex = "3.0.0-pre.10"
+ntex-bytes = { version = "1.2", features=["simd"] }
 mimalloc = { version = "0.1.25", default-features = false }
 snmalloc-rs = { version = "0.3.3", features = ["native-cpu"] }
 yarte = { version = "0.15", features = ["bytes-buf", "json"] }

frameworks/Rust/ntex/src/db.rs

@@ -1,7 +1,7 @@
-use std::cell::RefCell;
+use std::cell::Cell;
 
 use nanorand::{Rng, WyRand};
-use ntex::util::{Bytes, BytesMut};
+use ntex::util::{Bytes, BytesVec, BufMut};
 use smallvec::SmallVec;
 use tokio_postgres::{connect, Client, Statement};
 use yarte::TemplateBytesTrait;
@@ -33,7 +33,7 @@ pub struct PgConnection {
     world: Statement,
     rng: WyRand,
     updates: Statement,
-    buf: RefCell<BytesMut>,
+    buf: Cell<Option<BytesVec>>,
 }
 
 impl PgConnection {
@@ -55,7 +55,7 @@ impl PgConnection {
             world,
             updates,
             rng: WyRand::new(),
-            buf: RefCell::new(BytesMut::with_capacity(10 * 1024 * 1024)),
+            buf: Cell::new(Some(BytesVec::with_capacity(10 * 1024 * 1024))),
         }
     }
 }
@@ -66,17 +66,18 @@ impl PgConnection {
 
         let row = self.cl.query_one(&self.world, &[&random_id]).await.unwrap();
 
-        let mut body = self.buf.borrow_mut();
-        utils::reserve(&mut body, 1024);
+        let mut body = self.buf.take().unwrap();
         sonic_rs::to_writer(
-            utils::BytesWriter(&mut body),
+            utils::BVecWriter::new(&mut body),
             &World {
                 id: row.get(0),
                 randomnumber: row.get(1),
             },
         )
         .unwrap();
-        body.split().freeze()
+        let result = body.take_bytes();
+        self.buf.set(Some(body));
+        result
     }
 
     pub async fn get_worlds(&self, num: usize) -> Bytes {
@@ -96,10 +97,11 @@ impl PgConnection {
             })
         }
 
-        let mut body = self.buf.borrow_mut();
-        utils::reserve(&mut body, 2 * 1024);
-        sonic_rs::to_writer(utils::BytesWriter(&mut body), &worlds[..]).unwrap();
-        body.split().freeze()
+        let mut body = self.buf.take().unwrap();
+        sonic_rs::to_writer(utils::BVecWriter::new(&mut body), &worlds[..]).unwrap();
+        let result = body.take_bytes();
+        self.buf.set(Some(body));
+        result
     }
 
     pub async fn update(&self, num: usize) -> Bytes {
@@ -131,10 +133,11 @@ impl PgConnection {
 
         update.await.unwrap();
 
-        let mut body = self.buf.borrow_mut();
-        utils::reserve(&mut body, 2 * 1024);
-        sonic_rs::to_writer(utils::BytesWriter(&mut body), &worlds[..]).unwrap();
-        body.split().freeze()
+        let mut body = self.buf.take().unwrap();
+        sonic_rs::to_writer(utils::BVecWriter::new(&mut body), &worlds[..]).unwrap();
+        let result = body.take_bytes();
+        self.buf.set(Some(body));
+        result
     }
 
     pub async fn tell_fortune(&self) -> Bytes {
@@ -151,17 +154,20 @@ impl PgConnection {
         }));
         fortunes.sort_by(|it, next| it.message.cmp(&next.message));
 
-        let mut body = std::mem::replace(&mut *self.buf.borrow_mut(), BytesMut::new());
-        utils::reserve(&mut body, 4 * 1024);
+        let mut body = self.buf.take().unwrap();
+        let remaining = body.remaining_mut();
+        if remaining < 4 * 1024 {
+            body.reserve(128 * 1024);
+        }
 
         FortunesTemplate {
             fortunes: &fortunes,
         }
         .write_call(&mut body);
         fortunes.clear();
 
-        let result = body.split().freeze();
-        let _ = std::mem::replace(&mut *self.buf.borrow_mut(), body);
+        let result = body.take_bytes();
+        self.buf.set(Some(body));
         result
     }
 }

frameworks/Rust/ntex/src/utils.rs

@@ -3,7 +3,8 @@ use std::{cmp, io, io::Write, mem::MaybeUninit, slice::from_raw_parts_mut};
 
 use atoi::FromRadix10;
 use ntex::http::{header::HeaderValue, HttpServiceConfig, KeepAlive};
-use ntex::{io::IoConfig, time::Seconds, util::BufMut, util::Bytes, util::BytesMut, SharedCfg};
+use ntex::util::{BufMut, Bytes, BytesMut, BytesVec};
+use ntex::{io::IoConfig, time::Seconds, SharedCfg};
 use sonic_rs::writer::WriteExt;
 
 pub const HDR_SERVER: HeaderValue = HeaderValue::from_static("N");
@@ -97,3 +98,47 @@ impl WriteExt for BytesWriter<'_> {
         Ok(())
     }
 }
+
+pub struct BVecWriter<'a>(&'a mut BytesVec);
+
+impl<'a> BVecWriter<'a> {
+    pub fn new(buf: &'a mut BytesVec) -> BVecWriter<'a> {
+        let remaining = buf.remaining_mut();
+        if remaining < 2048 {
+            buf.reserve(HW);
+        }
+        Self(buf)
+    }
+}
+
+impl Write for BVecWriter<'_> {
+    fn write(&mut self, src: &[u8]) -> Result<usize, io::Error> {
+        self.0.extend_from_slice(src);
+        Ok(src.len())
+    }
+
+    fn flush(&mut self) -> Result<(), io::Error> {
+        Ok(())
+    }
+}
+
+impl WriteExt for BVecWriter<'_> {
+    #[inline(always)]
+    fn reserve_with(&mut self, additional: usize) -> Result<&mut [MaybeUninit<u8>], io::Error> {
+        self.0.reserve(additional);
+
+        unsafe {
+            let ptr = self.0.as_mut_ptr().add(self.0.len()) as *mut MaybeUninit<u8>;
+            Ok(from_raw_parts_mut(ptr, additional))
+        }
+    }
+
+    #[inline(always)]
+    unsafe fn flush_len(&mut self, additional: usize) -> io::Result<()> {
+        unsafe {
+            let new_len = self.0.len() + additional;
+            self.0.set_len(new_len);
+        }
+        Ok(())
+    }
+}