Add IoRef::with_rw_buf() helper

ntex-io/CHANGES.md

@@ -1,5 +1,9 @@
 # Changes
 
+## [0.2.6] - 2023-01-27
+
+* Add IoRef::with_rw_buf() helper
+
 ## [0.2.5] - 2023-01-27
 
 * Custom panic message for nested buffer borrow

ntex-io/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "ntex-io"
-version = "0.2.5"
+version = "0.2.6"
 authors = ["ntex contributors <team@ntex.rs>"]
 description = "Utilities for encoding and decoding frames"
 keywords = ["network", "framework", "async", "futures"]

ntex-io/src/buf.rs

@@ -131,6 +131,20 @@ impl Stack {
             .unwrap_or(0)
     }
 
+    pub(crate) fn with_rw_buf<F, R>(&mut self, io: &IoRef, f: F) -> R
+    where
+        F: FnOnce(&mut BytesVec, &mut BytesVec) -> R,
+    {
+        let lvl = self.get_first_level();
+        if lvl.0.is_none() {
+            lvl.0 = Some(io.memory_pool().get_read_buf());
+        }
+        if lvl.1.is_none() {
+            lvl.1 = Some(io.memory_pool().get_write_buf());
+        }
+        f(lvl.0.as_mut().unwrap(), lvl.1.as_mut().unwrap())
+    }
+
     pub(crate) fn first_read_buf(&mut self) -> &mut Option<BytesVec> {
         &mut self.get_first_level().0
     }

ntex-io/src/ioref.rs

@@ -1,9 +1,9 @@
-use std::{any, cell, fmt, hash, io, time};
+use std::{any, fmt, hash, io, time, cell};
 
 use ntex_bytes::{BytesVec, PoolRef};
 use ntex_codec::{Decoder, Encoder};
 
-use super::{buf::Stack, io::Flags, timer, types, Filter, IoRef, OnDisconnect, WriteBuf};
+use super::{io::Flags, buf::Stack, timer, types, Filter, IoRef, OnDisconnect, WriteBuf};
 
 impl IoRef {
     #[inline]
@@ -198,6 +198,15 @@ impl IoRef {
         f(buf.as_mut().unwrap())
     }
 
+    #[inline]
+    /// Get mut access to read and write buffer
+    pub fn with_rw_buf<F, R>(&self, f: F) -> R
+    where
+        F: FnOnce(&mut BytesVec, &mut BytesVec) -> R,
+    {
+        borrow_buffer(&self.0.buffer).with_rw_buf(self, f)
+    }
+
     #[inline]
     /// Start keep-alive timer
     pub fn start_keepalive_timer(&self, timeout: time::Duration) {