batch recv for udp proxy (#552)

easytier/src/easytier-core.rs

@@ -185,7 +185,7 @@ struct Cli {
     #[arg(
         long,
         help = t!("core_clap.multi_thread").to_string(),
-        default_value = "false"
+        default_value = "true"
     )]
     multi_thread: bool,
 

easytier/src/gateway/udp_proxy.rs

@@ -4,6 +4,7 @@ use std::{
     time::Duration,
 };
 
+use bytes::{BufMut, BytesMut};
 use cidr::Ipv4Inet;
 use crossbeam::atomic::AtomicCell;
 use dashmap::DashMap;
@@ -24,11 +25,11 @@ use tokio::{
 use tracing::Level;
 
 use crate::{
-    common::{error::Error, global_ctx::ArcGlobalCtx, PeerId},
+    common::{error::Error, global_ctx::ArcGlobalCtx, scoped_task::ScopedTask, PeerId},
     gateway::ip_reassembler::compose_ipv4_packet,
     peers::{peer_manager::PeerManager, PeerPacketFilter},
     tunnel::{
-        common::setup_sokcet2,
+        common::{reserve_buf, setup_sokcet2},
         packet_def::{PacketType, ZCPacket},
     },
 };
@@ -139,14 +140,28 @@ impl UdpNatEntry {
         mut packet_sender: Sender<ZCPacket>,
         virtual_ipv4: Ipv4Addr,
     ) {
-        let mut buf = [0u8; 65536];
+        let (s, mut r) = tachyonix::channel(128);
-        let mut udp_body: &mut [u8] = unsafe { std::mem::transmute(&mut buf[20 + 8..]) };
-        let mut ip_id = 1;
 
+        let self_clone = self.clone();
+        let recv_task = ScopedTask::from(tokio::spawn(async move {
+            let mut cur_buf = BytesMut::new();
             loop {
+                if self_clone
+                    .stopped
+                    .load(std::sync::atomic::Ordering::Relaxed)
+                {
+                    break;
+                }
+
+                reserve_buf(&mut cur_buf, 64 * 1024 + 28, 128 * 1024 + 28);
+                assert_eq!(cur_buf.len(), 0);
+                unsafe {
+                    cur_buf.advance_mut(28);
+                }
+
                 let (len, src_socket) = match timeout(
                     Duration::from_secs(120),
-                self.socket.recv_from(&mut udp_body),
+                    self_clone.socket.recv_buf_from(&mut cur_buf),
                 )
                 .await
                 {
@@ -163,27 +178,32 @@ impl UdpNatEntry {
 
                 tracing::trace!(?len, ?src_socket, "udp nat packet response received");
 
-            if self.stopped.load(std::sync::atomic::Ordering::Relaxed) {
+                let ret_buf = cur_buf.split();
-                break;
+                s.send((ret_buf, len, src_socket)).await.unwrap();
             }
+        }));
 
+        let self_clone = self.clone();
+        let send_task = ScopedTask::from(tokio::spawn(async move {
+            let mut ip_id = 1;
+            while let Ok((mut packet, len, src_socket)) = r.recv().await {
                 let SocketAddr::V4(mut src_v4) = src_socket else {
                     continue;
                 };
 
-            self.mark_active();
+                self_clone.mark_active();
 
                 if src_v4.ip().is_loopback() {
                     src_v4.set_ip(virtual_ipv4);
                 }
 
                 let Ok(_) = Self::compose_ipv4_packet(
-                &self,
+                    &self_clone,
                     &mut packet_sender,
-                &mut buf,
+                    &mut packet,
                     &src_v4,
                     len,
-                1200,
+                    1280,
                     ip_id,
                 )
                 .await
@@ -192,6 +212,9 @@ impl UdpNatEntry {
                 };
                 ip_id = ip_id.wrapping_add(1);
             }
+        }));
+
+        let _ = tokio::join!(recv_task, send_task);
 
         self.stop();
     }