Merge pull request #177 from QueryaHub/perf-buffer-pool-160

perf(memory): implement thread-local buffer pool for request bodies (#160)

Author: ZhuchkaTriplesix

src/buffer_pool.rs

@@ -0,0 +1,84 @@
+//! Thread-local buffer pool for request bodies to prevent heap fragmentation (issue #160).
+
+use std::cell::RefCell;
+use std::ops::{Deref, DerefMut};
+
+const POOL_CAPACITY_LIMIT: usize = 128 * 1024;
+const INITIAL_BUFFER_CAPACITY: usize = 16 * 1024;
+const MAX_POOL_SIZE: usize = 64;
+
+thread_local! {
+    static BODY_POOL: RefCell<Vec<Vec<u8>>> = const { RefCell::new(Vec::new()) };
+}
+
+/// RAII wrapper around a pooled `Vec<u8>` that automatically returns itself to the
+/// thread-local buffer pool on `Drop` if capacity <= 128 KB.
+pub struct PooledBuffer(Vec<u8>);
+
+impl PooledBuffer {
+    pub fn new() -> Self {
+        let buf = BODY_POOL
+            .with(|pool| pool.borrow_mut().pop())
+            .unwrap_or_else(|| Vec::with_capacity(INITIAL_BUFFER_CAPACITY));
+        Self(buf)
+    }
+
+    pub fn take(&mut self) -> Vec<u8> {
+        std::mem::take(&mut self.0)
+    }
+}
+
+impl Default for PooledBuffer {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+impl Deref for PooledBuffer {
+    type Target = Vec<u8>;
+
+    #[inline]
+    fn deref(&self) -> &Self::Target {
+        &self.0
+    }
+}
+
+impl DerefMut for PooledBuffer {
+    #[inline]
+    fn deref_mut(&mut self) -> &mut Self::Target {
+        &mut self.0
+    }
+}
+
+impl Drop for PooledBuffer {
+    fn drop(&mut self) {
+        if self.0.capacity() >= INITIAL_BUFFER_CAPACITY && self.0.capacity() <= POOL_CAPACITY_LIMIT {
+            self.0.clear();
+            let buf = std::mem::take(&mut self.0);
+            BODY_POOL.with(|pool| {
+                let mut p = pool.borrow_mut();
+                if p.len() < MAX_POOL_SIZE {
+                    p.push(buf);
+                }
+            });
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_pooled_buffer_recycle() {
+        {
+            let mut buf = PooledBuffer::new();
+            buf.extend_from_slice(b"hello world");
+            assert_eq!(&*buf, b"hello world");
+        }
+        // Dropped -> returned to pool.
+        let buf2 = PooledBuffer::new();
+        assert!(buf2.is_empty());
+        assert!(buf2.capacity() >= INITIAL_BUFFER_CAPACITY);
+    }
+}

src/dispatch.rs

@@ -10,6 +10,7 @@ use pyo3::types::{PyBytes, PyDict, PyList, PyString, PyTuple};
 use pyo3::IntoPyObjectExt;
 use serde_json::Value as JsonValue;
 
+use crate::buffer_pool::PooledBuffer;
 use crate::config;
 use crate::form::{self, ParsedFile};
 use crate::params::{build_request_context, header_get_lax, parse_query, value_for_path_param};
@@ -699,22 +700,26 @@ pub async fn run_rsgi(
             routes_arc[route_idx].extra.path_template.clone(),
         )
     });
-    let mut body_bytes: Vec<u8> = if should_read_body {
+    let mut body_bytes: PooledBuffer = if should_read_body {
         let read_fut = Python::with_gil(|py| {
             let p = protocol.bind(py);
             let aw: Bound<PyAny> = p.call0()?;
             pyo3_async_runtimes::tokio::into_future(aw)
         })?;
         let body_obj: PyObject = read_fut.await?;
-        let body = Python::with_gil(|py| -> PyResult<Vec<u8>> {
+        let mut body = PooledBuffer::new();
+        Python::with_gil(|py| -> PyResult<()> {
             let b = body_obj.bind(py);
-            if let Ok(x) = b.extract::<Vec<u8>>() {
-                return Ok(x);
-            }
-            if let Ok(s) = b.str() {
-                return Ok(s.to_string().into_bytes());
+            if let Ok(py_bytes) = b.downcast::<pyo3::types::PyBytes>() {
+                body.extend_from_slice(py_bytes.as_bytes());
+            } else if let Ok(py_str) = b.downcast::<pyo3::types::PyString>() {
+                if let Ok(s) = py_str.to_str() {
+                    body.extend_from_slice(s.as_bytes());
+                }
+            } else if let Ok(bytes_vec) = b.extract::<Vec<u8>>() {
+                body.extend_from_slice(&bytes_vec);
             }
-            Ok(Vec::new())
+            Ok(())
         })?;
         let max = form::max_body_bytes();
         if (body.len() as u64) > max {
@@ -728,7 +733,7 @@ pub async fn run_rsgi(
         }
         body
     } else {
-        Vec::new()
+        PooledBuffer::new()
     };
     let (auth, cookie_raw): (Option<String>, Option<String>) = if require_jwt {
         Python::with_gil(|py| -> PyResult<(Option<String>, Option<String>)> {
@@ -886,7 +891,7 @@ pub async fn run_rsgi(
                         .await
                     }
                 };
-                let multipart_body = std::mem::take(&mut body_bytes);
+                let multipart_body = body_bytes.take();
                 let parsed = match form::parse_multipart(multipart_body, &boundary).await {
                     Ok(p) => p,
                     Err(e) => {

src/lib.rs

@@ -8,6 +8,7 @@ use pyo3::prelude::*;
 use pyo3::types::{PyDict, PyList, PyTuple};
 use serde_json::json;
 
+mod buffer_pool;
 mod config;
 mod db;
 mod dispatch;