perf(wireguard): rename WG_MAX_PACKET_SIZE and optimize scratch buffers with FixedBufPool (#2627)

easytier/src/tunnel/wireguard.rs

@@ -43,7 +43,15 @@ use tokio::{
     task::JoinSet,
 };
 
-const MAX_PACKET: usize = 2048;
+use crate::utils::buf::FixedBufPool;
+use std::sync::LazyLock;
+
+/// Maximum size of a single WireGuard UDP packet (payload + headers + auth tag).
+/// Standard Ethernet MTU is 1500; 2048 safely accommodates any WireGuard datagram.
+pub const WG_MAX_PACKET_SIZE: usize = 2048;
+
+static WG_BUF_POOL: LazyLock<FixedBufPool<WG_MAX_PACKET_SIZE>> =
+    LazyLock::new(|| FixedBufPool::new(64));
 
 #[derive(Debug, Clone, Copy, PartialEq, Eq)]
 enum WgType {
@@ -166,7 +174,7 @@ impl Debug for WgPeerData {
 impl WgPeerData {
     #[tracing::instrument]
     async fn handle_one_packet_from_me(&self, zc_packet: ZCPacket) -> Result<(), anyhow::Error> {
-        let mut send_buf = vec![0u8; MAX_PACKET];
+        let mut send_buf = WG_BUF_POOL.acquire();
 
         let packet = if self.internal_use {
             let mut zc_packet = zc_packet.convert_type(ZCPacketType::WG);
@@ -224,7 +232,7 @@ impl WgPeerData {
         recv_buf: &[u8],
     ) {
         self.access_time.store(Instant::now());
-        let mut send_buf = vec![0u8; MAX_PACKET];
+        let mut send_buf = WG_BUF_POOL.acquire();
         let data = recv_buf;
         let decapsulate_result = {
             let mut peer = self.tunn.lock().await;
@@ -247,7 +255,7 @@ impl WgPeerData {
                 };
                 let mut peer = self.tunn.lock().await;
                 loop {
-                    let mut send_buf = vec![0u8; MAX_PACKET];
+                    let mut send_buf = WG_BUF_POOL.acquire();
                     match peer.decapsulate(None, &[], &mut send_buf) {
                         TunnResult::WriteToNetwork(packet) => {
                             match self.session.send(packet).await {
@@ -318,7 +326,7 @@ impl WgPeerData {
             TunnResult::Err(WireGuardError::ConnectionExpired) => {
                 tracing::warn!("Wireguard handshake has expired!");
 
-                let mut buf = vec![0u8; MAX_PACKET];
+                let mut buf = WG_BUF_POOL.acquire();
                 let result = self
                     .tunn
                     .lock()
@@ -347,7 +355,7 @@ impl WgPeerData {
     /// WireGuard Routine task. Handles Handshake, keep-alive, etc.
     pub async fn routine_task(self) {
         loop {
-            let mut send_buf = vec![0u8; MAX_PACKET];
+            let mut send_buf = WG_BUF_POOL.acquire();
             let tun_result = { self.tunn.lock().await.update_timers(&mut send_buf) };
             self.handle_routine_tun_result(tun_result).await;
         }
@@ -489,7 +497,7 @@ impl WgPeer {
                 data.handle_one_packet_from_peer(&mut sink, &packet).await;
             }
 
-            let mut buf = vec![0u8; MAX_PACKET];
+            let mut buf = vec![0u8; WG_MAX_PACKET_SIZE];
             loop {
                 let n = match session.recv(&mut buf).await {
                     Ok(n) => n,
@@ -708,7 +716,7 @@ pub(crate) async fn upgrade_connected(
     // do handshake here so we will return after receive first packet
     let handshake = wg_peer.create_handshake_init().await;
     session.send(&handshake).await?;
-    let mut buf = [0u8; MAX_PACKET];
+    let mut buf = [0u8; WG_MAX_PACKET_SIZE];
     let n = match session.recv(&mut buf).await {
         Ok(ret) => ret,
         Err(e) => {

easytier/src/utils/buf.rs

@@ -1,6 +1,6 @@
 use bytes::{BufMut, BytesMut};
-use derive_more::{From, Into};
-use std::mem::MaybeUninit;
+use derive_more::{AsMut, AsRef, Deref, DerefMut, From, Into};
+use std::mem::{MaybeUninit, take};
 use std::ptr::copy_nonoverlapping;
 
 pub use easytier_core::tunnel::buf::BufList;
@@ -119,6 +119,42 @@ impl<'t> BufPoolWriter<'t> {
     }
 }
 
+/// A lock-free object pool for fixed-capacity reusable scratch buffers.
+#[derive(Debug)]
+pub struct FixedBufPool<const SIZE: usize> {
+    queue: crossbeam::queue::ArrayQueue<Vec<u8>>,
+}
+
+impl<const SIZE: usize> FixedBufPool<SIZE> {
+    pub fn new(capacity: usize) -> Self {
+        Self {
+            queue: crossbeam::queue::ArrayQueue::new(capacity),
+        }
+    }
+
+    pub fn acquire(&self) -> FixedBufGuard<'_, SIZE> {
+        let buf = self.queue.pop().unwrap_or_else(|| vec![0u8; SIZE]);
+        FixedBufGuard { pool: self, buf }
+    }
+}
+
+#[derive(Debug, Deref, DerefMut, AsRef, AsMut)]
+pub struct FixedBufGuard<'p, const SIZE: usize> {
+    pool: &'p FixedBufPool<SIZE>,
+    #[deref]
+    #[deref_mut]
+    #[as_ref([u8])]
+    #[as_mut([u8])]
+    buf: Vec<u8>,
+}
+
+impl<'p, const SIZE: usize> Drop for FixedBufGuard<'p, SIZE> {
+    #[inline(always)]
+    fn drop(&mut self) {
+        let _ = self.pool.queue.push(take(&mut self.buf));
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -156,4 +192,34 @@ mod tests {
         assert_eq!(buf.len(), 10 + data.len() + 6);
         assert_eq!(&buf[10..10 + data.len()], data);
     }
+
+    #[test]
+    fn test_fixed_buf_pool() {
+        let pool = FixedBufPool::<2048>::new(2);
+
+        // 1. Acquire and verify size
+        let ptr1;
+        {
+            let mut buf1 = pool.acquire();
+            assert_eq!(buf1.len(), 2048);
+            buf1[0] = 42;
+            ptr1 = buf1.as_ptr();
+        }
+
+        // 2. Re-acquire: should reuse the recycled buffer from the pool
+        {
+            let buf2 = pool.acquire();
+            assert_eq!(buf2.len(), 2048);
+            assert_eq!(buf2[0], 42); // Same underlying memory was recycled
+            assert_eq!(buf2.as_ptr(), ptr1);
+        }
+
+        // 3. Exceed pool capacity
+        let b1 = pool.acquire();
+        let b2 = pool.acquire();
+        let b3 = pool.acquire(); // exceeds capacity=2, allocates on demand
+        assert_eq!(b1.len(), 2048);
+        assert_eq!(b2.len(), 2048);
+        assert_eq!(b3.len(), 2048);
+    }
 }