From 6abc3731ae0b6acda25caa212155d83a6a29768c Mon Sep 17 00:00:00 2001 From: Luna Yao <40349250+ZnqbuZ@users.noreply.github.com> Date: Mon, 5 Oct 2026 05:57:40 +0200 Subject: [PATCH] perf(wireguard): rename WG_MAX_PACKET_SIZE and optimize scratch buffers with FixedBufPool (#2627) --- easytier/src/tunnel/wireguard.rs | 24 +++++++---- easytier/src/utils/buf.rs | 70 +++++++++++++++++++++++++++++++- 2 files changed, 84 insertions(+), 10 deletions(-) diff --git a/easytier/src/tunnel/wireguard.rs b/easytier/src/tunnel/wireguard.rs index 8268da09..8336a529 100644 --- a/easytier/src/tunnel/wireguard.rs +++ b/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> = + 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) => { diff --git a/easytier/src/utils/buf.rs b/easytier/src/utils/buf.rs index fe4ed932..31173141 100644 --- a/easytier/src/utils/buf.rs +++ b/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 { + queue: crossbeam::queue::ArrayQueue>, +} + +impl FixedBufPool { + 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, + #[deref] + #[deref_mut] + #[as_ref([u8])] + #[as_mut([u8])] + buf: Vec, +} + +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); + } }