Trying to make UDP more efficient

Power2All · Power2All · commit ff9601355d12 · 2025-08-31T14:56:09.000+02:00
diff --git a/src/config/impls/configuration.rs b/src/config/impls/configuration.rs
@@ -125,7 +125,10 @@ impl Configuration {
                     threads: available_parallelism().unwrap().get() as u64,
                     receive_buffer_size: 134217728,
                     send_buffer_size: 67108864,
-                    reuse_address: true
+                    reuse_address: true,
+                    receiver_threads: 2,
+                    worker_threads: 16,
+                    queue_size: 10000
                 }
             ),
             api_server: vec!(
diff --git a/src/config/structs/udp_trackers_config.rs b/src/config/structs/udp_trackers_config.rs
@@ -7,5 +7,8 @@ pub struct UdpTrackersConfig {
     pub threads: u64,
     pub receive_buffer_size: usize,
     pub send_buffer_size: usize,
-    pub reuse_address: bool
+    pub reuse_address: bool,
+    pub receiver_threads: usize,
+    pub worker_threads: usize,
+    pub queue_size: usize
 }
diff --git a/src/udp/impls/udp_server.rs b/src/udp/impls/udp_server.rs
@@ -2,9 +2,10 @@ use std::io::Cursor;
 use std::net::{Ipv4Addr, Ipv6Addr, SocketAddr};
 use std::sync::Arc;
 use std::time::SystemTime;
-use log::{debug, info};
+use log::{debug, info, warn};
 use socket2::{Socket, Domain, Type, Protocol};
 use tokio::net::UdpSocket;
+use tokio::sync::mpsc;
 use crate::stats::enums::stats_event::StatsEvent;
 use crate::tracker::enums::torrent_peers_type::TorrentPeersType;
 use crate::tracker::structs::announce_query_request::AnnounceQueryRequest;
@@ -24,6 +25,7 @@ use crate::udp::structs::connection_id::ConnectionId;
 use crate::udp::structs::error_response::ErrorResponse;
 use crate::udp::structs::number_of_downloads::NumberOfDownloads;
 use crate::udp::structs::number_of_peers::NumberOfPeers;
+use crate::udp::structs::packet_job::PacketJob;
 use crate::udp::structs::port::Port;
 use crate::udp::structs::response_peer::ResponsePeer;
 use crate::udp::structs::scrape_request::ScrapeRequest;
@@ -35,7 +37,16 @@ use crate::udp::udp::MAX_SCRAPE_TORRENTS;
 
 impl UdpServer {
     #[tracing::instrument(level = "debug")]
-    pub async fn new(tracker: Arc<TorrentTracker>, bind_address: SocketAddr, threads: u64, recv_buffer_size: usize, send_buffer_size: usize, reuse_address: bool) -> tokio::io::Result<UdpServer>
+    pub async fn new(
+        tracker: Arc<TorrentTracker>,
+        bind_address: SocketAddr,
+        recv_buffer_size: usize,
+        send_buffer_size: usize,
+        reuse_address: bool,
+        receiver_threads: usize,
+        worker_threads: usize,
+        queue_size: usize
+    ) -> tokio::io::Result<UdpServer>
     {
         let domain = if bind_address.is_ipv4() { Domain::IPV4 } else { Domain::IPV6 };
         let socket = Socket::new(domain, Type::DGRAM, Some(Protocol::UDP))?;
@@ -52,58 +63,126 @@ impl UdpServer {
 
         Ok(UdpServer {
             socket: Arc::new(tokio_socket),
-            threads,
             tracker,
+            receiver_threads: receiver_threads as u64,
+            worker_threads: worker_threads as u64,
+            queue_size: queue_size as u64
         })
     }
 
     #[tracing::instrument(level = "debug")]
-    pub async fn start(&self, rx: tokio::sync::watch::Receiver<bool>)
-    {
-        let threads = self.threads;
-        for _index in 0..=threads {
+    pub async fn start(&self, rx: tokio::sync::watch::Receiver<bool>) {
+        let (packet_tx, packet_rx) = mpsc::channel::<PacketJob>(self.queue_size as usize);
+        let packet_rx = Arc::new(tokio::sync::Mutex::new(packet_rx));
+
+        let receiver_threads = self.receiver_threads as usize;
+        let worker_threads = self.worker_threads as usize;
+
+        for thread_id in 0..receiver_threads {
             let socket_clone = self.socket.clone();
             let tracker = self.tracker.clone();
-            let mut rx = rx.clone();
-            let mut data = [0; 1496];
+            let mut shutdown_rx = rx.clone();
+            let packet_tx = packet_tx.clone();
+
             tokio::spawn(async move {
+                info!("Starting UDP receiver thread {}", thread_id);
+                let mut data = [0; 1496];
+
                 loop {
-                    let udp_sock = socket_clone.local_addr().unwrap();
                     tokio::select! {
-                        _ = rx.changed() => {
-                            info!("Stopping UDP server: {udp_sock}...");
+                        _ = shutdown_rx.changed() => {
+                            info!("Stopping UDP receiver thread {}...", thread_id);
                             break;
                         }
                         Ok((valid_bytes, remote_addr)) = socket_clone.recv_from(&mut data) => {
-                            let payload = &data[..valid_bytes];
-
-                            debug!("Received {} bytes from {}", payload.len(), remote_addr);
-                            debug!("{payload:?}");
-
-                            let tracker_cloned = tracker.clone();
-                            let socket_cloned = socket_clone.clone();
-                            // Use payload slice instead of cloning the entire Vec
-                            let payload_vec = payload.to_vec();
-                            tokio::spawn(async move {
-                                let response = UdpServer::handle_packet(remote_addr, payload_vec, tracker_cloned.clone()).await;
-                                UdpServer::send_response(tracker_cloned.clone(), socket_cloned.clone(), remote_addr, response).await;
-                            });
+                            let payload = data[..valid_bytes].to_vec();
+
+                            debug!("Receiver {} got {} bytes from {}", thread_id, payload.len(), remote_addr);
+
+                            let job = PacketJob {
+                                data: payload,
+                                remote_addr,
+                            };
+
+                            if let Err(e) = packet_tx.try_send(job) {
+                                warn!("Packet queue full, dropping packet: {}", e);
+                                match remote_addr {
+                                    SocketAddr::V4(_) => tracker.update_stats(StatsEvent::Udp4BadRequest, 1),
+                                    SocketAddr::V6(_) => tracker.update_stats(StatsEvent::Udp6BadRequest, 1),
+                                };
+                            }
+                        }
+                    }
+                }
+            });
+        }
+
+        for thread_id in 0..worker_threads {
+            let socket_clone = self.socket.clone();
+            let tracker = self.tracker.clone();
+            let mut shutdown_rx = rx.clone();
+            let packet_rx = packet_rx.clone();
+
+            tokio::spawn(async move {
+                info!("Starting UDP worker thread {}", thread_id);
+
+                loop {
+                    tokio::select! {
+                        _ = shutdown_rx.changed() => {
+                            info!("Stopping UDP worker thread {}...", thread_id);
+                            break;
+                        }
+                        job = async {
+                            let mut rx = packet_rx.lock().await;
+                            rx.recv().await
+                        } => {
+                            if let Some(PacketJob { data, remote_addr }) = job {
+                                debug!("Worker {} processing packet from {}", thread_id, remote_addr);
+
+                                let response = UdpServer::handle_packet(
+                                    remote_addr,
+                                    data,
+                                    tracker.clone()
+                                ).await;
+
+                                UdpServer::send_response(
+                                    tracker.clone(),
+                                    socket_clone.clone(),
+                                    remote_addr,
+                                    response
+                                ).await;
+                            }
                         }
                     }
                 }
             });
         }
+
+        drop(packet_tx);
     }
 
+    // Optimized send_response with pre-sized buffer
     #[tracing::instrument(level = "debug")]
-    pub async fn send_response(tracker: Arc<TorrentTracker>, socket: Arc<UdpSocket>, remote_addr: SocketAddr, response: Response)
-    {
+    pub async fn send_response(
+        tracker: Arc<TorrentTracker>,
+        socket: Arc<UdpSocket>,
+        remote_addr: SocketAddr,
+        response: Response
+    ) {
         debug!("sending response to: {:?}", &remote_addr);
         let sentry = sentry::TransactionContext::new("udp server", "send response");
         let transaction = sentry::start_transaction(sentry);
 
-        // Pre-allocate buffer with exact capacity instead of MAX_PACKET_SIZE
-        let mut buffer = Vec::with_capacity(512); // Most responses are much smaller than MAX_PACKET_SIZE
+        // Optimize buffer allocation based on response type
+        let estimated_size = match &response {
+            Response::Connect(_) => 16,
+            Response::AnnounceIpv4(_) => 20 + 6 * 72,  // header + max IPv4 peers (6 bytes each)
+            Response::AnnounceIpv6(_) => 20 + 18 * 72, // header + max IPv6 peers (18 bytes each)
+            Response::Scrape(_) => 8 + 12 * 74,        // header + max torrents
+            Response::Error(_) => 128,                 // reasonable max for error message
+        };
+
+        let mut buffer = Vec::with_capacity(estimated_size);
         let mut cursor = Cursor::new(&mut buffer);
 
         match response.write(&mut cursor) {
@@ -125,6 +204,7 @@ impl UdpServer {
         transaction.finish();
     }
 
+    // Rest of the methods remain the same...
     #[tracing::instrument(level = "debug")]
     pub async fn send_packet(socket: Arc<UdpSocket>, remote_addr: &SocketAddr, payload: &[u8]) {
         let _ = socket.send_to(payload, remote_addr).await;
@@ -195,7 +275,7 @@ impl UdpServer {
     #[tracing::instrument(level = "debug")]
     pub async fn handle_udp_connect(remote_addr: SocketAddr, request: &ConnectRequest, tracker: Arc<TorrentTracker>) -> Result<Response, ServerError> {
         let connection_id = UdpServer::get_connection_id(&remote_addr).await;
-        let response = Response::from(ConnectResponse {
+        let response = Response::Connect(ConnectResponse {
             transaction_id: request.transaction_id,
             connection_id
         });
@@ -353,15 +433,15 @@ impl UdpServer {
 
         // Create response
         let response = if remote_addr.is_ipv6() {
-            Response::from(AnnounceResponse {
+            Response::AnnounceIpv6(AnnounceResponse {
                 transaction_id: request.transaction_id,
                 announce_interval: AnnounceInterval(config.request_interval as i32),
                 leechers: NumberOfPeers(torrent.peers.len() as i32),
                 seeders: NumberOfPeers(torrent.seeds.len() as i32),
                 peers: peers6,
             })
         } else {
-            Response::from(AnnounceResponse {
+            Response::AnnounceIpv4(AnnounceResponse {
                 transaction_id: request.transaction_id,
                 announce_interval: AnnounceInterval(config.request_interval as i32),
                 leechers: NumberOfPeers(torrent.peers.len() as i32),
@@ -408,15 +488,15 @@ impl UdpServer {
         };
         tracker.update_stats(stats_event, 1);
 
-        Ok(Response::from(ScrapeResponse {
+        Ok(Response::Scrape(ScrapeResponse {
             transaction_id: request.transaction_id,
             torrent_stats,
         }))
     }
 
     #[tracing::instrument(level = "debug")]
     pub async fn handle_udp_error(e: ServerError, transaction_id: TransactionId) -> Response {
-        Response::from(ErrorResponse {
+        Response::Error(ErrorResponse {
             transaction_id,
             message: e.to_string().into()
         })
diff --git a/src/udp/structs.rs b/src/udp/structs.rs
@@ -14,4 +14,5 @@ pub mod number_of_downloads;
 pub mod port;
 pub mod peer_key;
 pub mod response_peer;
-pub mod udp_server;
+pub mod udp_server;
+pub mod packet_job;
diff --git a/src/udp/structs/packet_job.rs b/src/udp/structs/packet_job.rs
@@ -0,0 +1,6 @@
+use std::net::SocketAddr;
+
+pub struct PacketJob {
+    pub(crate) data: Vec<u8>,
+    pub(crate) remote_addr: SocketAddr,
+}
diff --git a/src/udp/structs/udp_server.rs b/src/udp/structs/udp_server.rs
@@ -5,6 +5,8 @@ use crate::tracker::structs::torrent_tracker::TorrentTracker;
 #[derive(Debug)]
 pub struct UdpServer {
     pub(crate) socket: Arc<UdpSocket>,
-    pub(crate) threads: u64,
     pub(crate) tracker: Arc<TorrentTracker>,
+    pub(crate) receiver_threads: u64,
+    pub(crate) worker_threads: u64,
+    pub(crate) queue_size: u64
 }
diff --git a/src/udp/udp.rs b/src/udp/udp.rs
@@ -12,14 +12,34 @@ pub const MAX_SCRAPE_TORRENTS: u8 = 74;
 pub const MAX_PACKET_SIZE: usize = 1496;
 
 #[allow(clippy::too_many_arguments)]
-pub async fn udp_service(addr: SocketAddr, threads: u64, recv_buffer_size: usize, send_buffer_size: usize, reuse_address: bool, data: Arc<TorrentTracker>, rx: tokio::sync::watch::Receiver<bool>, tokio_udp: Arc<Runtime>) -> JoinHandle<()>
+pub async fn udp_service(
+    addr: SocketAddr,
+    recv_buffer_size: usize,
+    send_buffer_size: usize,
+    reuse_address: bool,
+    receiver_threads: usize,
+    worker_threads: usize,
+    queue_size: usize,
+    data: Arc<TorrentTracker>,
+    rx: tokio::sync::watch::Receiver<bool>,
+    tokio_udp: Arc<Runtime>
+) -> JoinHandle<()>
 {
-    let udp_server = UdpServer::new(data, addr, threads, recv_buffer_size, send_buffer_size, reuse_address).await.unwrap_or_else(|e| {
+    let udp_server = UdpServer::new(
+        data,
+        addr,
+        recv_buffer_size,
+        send_buffer_size,
+        reuse_address,
+        receiver_threads,
+        worker_threads,
+        queue_size
+    ).await.unwrap_or_else(|e| {
         error!("Could not listen to the UDP port: {e}");
         exit(1);
     });
 
-    info!("[UDP] Starting server listener on {addr} with {threads} threads");
+    info!("[UDP] Starting server listener on {addr}");
     tokio_udp.spawn(async move {
         udp_server.start(rx).await;
     })
