Trying to improve responsiveness by batch sending data

Power2All · Power2All · commit b14cf3a1be91 · 2025-08-31T10:56:51.000+02:00
diff --git a/src/udp/impls.rs b/src/udp/impls.rs
@@ -3,4 +3,5 @@ pub mod request;
 pub mod response;
 pub mod ipv4_addr;
 pub mod ipv6_addr;
-pub mod udp_server;
+pub mod udp_server;
+pub mod response_batch_manager;
diff --git a/src/udp/impls/response_batch_manager.rs b/src/udp/impls/response_batch_manager.rs
@@ -0,0 +1,67 @@
+use std::net::SocketAddr;
+use std::sync::{Arc, OnceLock};
+use std::time::Duration;
+use std::collections::{HashMap, VecDeque};
+use tokio::net::UdpSocket;
+use tokio::sync::{mpsc, RwLock};
+use tokio::time::interval;
+use crate::udp::structs::queued_response::QueuedResponse;
+use crate::udp::structs::response_batch_manager::ResponseBatchManager;
+
+static BATCH_MANAGERS: OnceLock<RwLock<HashMap<String, Arc<ResponseBatchManager>>>> = OnceLock::new();
+
+impl ResponseBatchManager {
+    fn new(socket: Arc<UdpSocket>) -> Self {
+        let (sender, mut receiver) = mpsc::unbounded_channel::<QueuedResponse>();
+
+        tokio::spawn(async move {
+            let mut buffer = VecDeque::with_capacity(1000);
+            let mut timer = interval(Duration::from_millis(5));
+
+            loop {
+                tokio::select! {
+                    Some(response) = receiver.recv() => {
+                        buffer.push_back(response);
+
+                        if buffer.len() >= 500 {
+                            Self::flush_buffer(&socket, &mut buffer).await;
+                        }
+                    }
+
+                    _ = timer.tick() => {
+                        if !buffer.is_empty() {
+                            Self::flush_buffer(&socket, &mut buffer).await;
+                        }
+                    }
+                }
+            }
+        });
+
+        Self { sender }
+    }
+
+    pub(crate) fn queue_response(&self, remote_addr: SocketAddr, payload: Vec<u8>) {
+        let _ = self.sender.send(QueuedResponse { remote_addr, payload });
+    }
+
+    async fn flush_buffer(socket: &UdpSocket, buffer: &mut VecDeque<QueuedResponse>) {
+        while let Some(response) = buffer.pop_front() {
+            let _ = socket.send_to(&response.payload, &response.remote_addr).await;
+        }
+    }
+
+    pub(crate) async fn get_for_socket(socket: Arc<UdpSocket>) -> Arc<ResponseBatchManager> {
+        let socket_key = format!("{:?}", socket.local_addr().unwrap_or_else(|_| "0.0.0.0:0".parse().unwrap()));
+
+        let registry = BATCH_MANAGERS.get_or_init(|| RwLock::new(HashMap::new()));
+
+        if let Some(manager) = registry.read().await.get(&socket_key) {
+            return manager.clone();
+        }
+
+        let manager = Arc::new(ResponseBatchManager::new(socket));
+        registry.write().await.insert(socket_key, manager.clone());
+
+        manager
+    }
+}
diff --git a/src/udp/impls/udp_server.rs b/src/udp/impls/udp_server.rs
@@ -25,6 +25,7 @@ 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::port::Port;
+use crate::udp::structs::response_batch_manager::ResponseBatchManager;
 use crate::udp::structs::response_peer::ResponsePeer;
 use crate::udp::structs::scrape_request::ScrapeRequest;
 use crate::udp::structs::scrape_response::ScrapeResponse;
@@ -46,7 +47,6 @@ impl UdpServer {
         socket.bind(&bind_address.into()).map_err(tokio::io::Error::other)?;
         socket.set_nonblocking(true).map_err(tokio::io::Error::other)?;
 
-        // Convert to std::net::UdpSocket, then to tokio::net::UdpSocket
         let std_socket: std::net::UdpSocket = socket.into();
         let tokio_socket = UdpSocket::from_std(std_socket)?;
 
@@ -82,7 +82,6 @@ impl UdpServer {
 
                             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;
@@ -102,15 +101,16 @@ impl UdpServer {
         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
+        let mut buffer = Vec::with_capacity(512);
         let mut cursor = Cursor::new(&mut buffer);
 
         match response.write(&mut cursor) {
             Ok(_) => {
                 let position = cursor.position() as usize;
                 debug!("Response bytes: {:?}", &buffer[..position]);
-                UdpServer::send_packet(socket, &remote_addr, &buffer[..position]).await;
+
+                let batch_manager = ResponseBatchManager::get_for_socket(socket).await;
+                batch_manager.queue_response(remote_addr, buffer[..position].to_vec());
             }
             Err(error) => {
                 sentry::capture_error(&error);
@@ -214,7 +214,6 @@ impl UdpServer {
     pub async fn handle_udp_announce(remote_addr: SocketAddr, request: &AnnounceRequest, tracker: Arc<TorrentTracker>) -> Result<Response, ServerError> {
         let config = tracker.config.tracker_config.clone();
 
-        // Whitelist/Blacklist checks
         if config.whitelist_enabled && !tracker.check_whitelist(InfoHash(request.info_hash.0)) {
             debug!("[UDP ERROR] Torrent Not Whitelisted");
             return Err(ServerError::TorrentNotWhitelisted);
@@ -224,7 +223,6 @@ impl UdpServer {
             return Err(ServerError::TorrentBlacklisted);
         }
 
-        // Key validation
         if config.keys_enabled {
             if request.path.len() < 50 {
                 debug!("[UDP ERROR] Unknown Key");
@@ -248,7 +246,6 @@ impl UdpServer {
             }
         }
 
-        // User key validation
         let user_key = if config.users_enabled {
             let user_key_path_extract = if request.path.len() >= 91 {
                 Some(&request.path[51..=91])
@@ -281,7 +278,6 @@ impl UdpServer {
             return Err(ServerError::PeerKeyNotValid);
         }
 
-        // Handle announce
         let torrent = match tracker.handle_announce(tracker.clone(), AnnounceQueryRequest {
             info_hash: InfoHash(request.info_hash.0),
             peer_id: PeerId(request.peer_id.0),
@@ -302,15 +298,13 @@ impl UdpServer {
             }
         };
 
-        // Get peers efficiently
         let torrent_peers = tracker.get_torrent_peers(request.info_hash, 72, TorrentPeersType::All, Some(remote_addr.ip()));
 
         let (peers, peers6) = if let Some(torrent_peers_unwrapped) = torrent_peers {
             let mut peers = Vec::with_capacity(72);
             let mut peers6 = Vec::with_capacity(72);
             let mut count = 0;
 
-            // Only collect peers if not completed download
             if request.bytes_left.0 != 0 {
                 if remote_addr.is_ipv4() {
                     for torrent_peer in torrent_peers_unwrapped.seeds_ipv4.values().take(72) {
@@ -331,7 +325,6 @@ impl UdpServer {
                 }
             }
 
-            // Collect regular peers
             if remote_addr.is_ipv4() {
                 for torrent_peer in torrent_peers_unwrapped.peers_ipv4.values().take(72 - count) {
                     if let Ok(ip) = torrent_peer.peer_addr.ip().to_string().parse::<Ipv4Addr>() {
@@ -351,7 +344,6 @@ impl UdpServer {
             (Vec::new(), Vec::new())
         };
 
-        // Create response
         let response = if remote_addr.is_ipv6() {
             Response::from(AnnounceResponse {
                 transaction_id: request.transaction_id,
@@ -370,7 +362,6 @@ impl UdpServer {
             })
         };
 
-        // Update stats
         let stats_event = if remote_addr.is_ipv4() {
             StatsEvent::Udp4AnnouncesHandled
         } else {
diff --git a/src/udp/structs.rs b/src/udp/structs.rs
@@ -14,4 +14,6 @@ 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 queued_response;
+pub mod response_batch_manager;
diff --git a/src/udp/structs/queued_response.rs b/src/udp/structs/queued_response.rs
@@ -0,0 +1,7 @@
+use std::net::SocketAddr;
+
+#[derive(Debug, Clone)]
+pub struct QueuedResponse {
+    pub(crate) remote_addr: SocketAddr,
+    pub(crate) payload: Vec<u8>,
+}
diff --git a/src/udp/structs/response_batch_manager.rs b/src/udp/structs/response_batch_manager.rs
@@ -0,0 +1,7 @@
+use crate::udp::structs::queued_response::QueuedResponse;
+use mpsc::UnboundedSender;
+use tokio::sync::mpsc;
+
+pub struct ResponseBatchManager {
+    pub(crate) sender: UnboundedSender<QueuedResponse>,
+}
