feat: gracefully shutdown udp servers

Author: mickvandijke

Cargo.lock

@@ -31,3 +31,4 @@ config = "0.11"
 derive_more = "0.99"
 thiserror = "1.0"
 aquatic_udp_protocol = { git = "https://github.com/greatest-ape/aquatic" }
+futures = "0.3.21"

Cargo.toml

@@ -7,20 +7,21 @@ use torrust_tracker::torrust_http_tracker::server::HttpServer;
 
 #[tokio::main]
 async fn main() {
+    // torrust config
     let config = match Configuration::load_from_file() {
         Ok(config) => Arc::new(config),
         Err(error) => {
             panic!("{}", error)
         }
     };
 
-    logging::setup_logging(&config);
-
     // the singleton torrent tracker that gets passed to the HTTP and UDP server
     let tracker = Arc::new(TorrentTracker::new(config.clone()).unwrap_or_else(|e| {
         panic!("{}", e)
     }));
 
+    logging::setup_logging(&config);
+
     // load persistent torrents if enabled
     if config.persistence {
         info!("Loading persistent torrents into memory...");
@@ -38,10 +39,17 @@ async fn main() {
         let _api_server = start_api_server(&config.http_api, tracker.clone());
     }
 
+    let (tx, rx) = tokio::sync::watch::channel(false);
+    let mut udp_server_handles = Vec::new();
+
     // start the udp blocks
     for udp_tracker in &config.udp_trackers {
+        // used to send kill signal to thread
+
         if udp_tracker.enabled {
-            let _ = start_udp_tracker_server(&udp_tracker, tracker.clone()).await;
+            udp_server_handles.push(
+                start_udp_tracker_server(&udp_tracker, tracker.clone(), rx.clone()).await
+            )
         }
     }
 
@@ -52,10 +60,16 @@ async fn main() {
     }
 
     // handle the signals here
-    let ctrl_c = tokio::signal::ctrl_c();
     tokio::select! {
-        _ = ctrl_c => {
+        _ = tokio::signal::ctrl_c() => {
             info!("Torrust shutting down..");
+
+            // send kill signal
+            let _ = tx.send(true);
+
+            // await for all udp servers to shutdown
+            futures::future::join_all(udp_server_handles).await;
+
             // Save torrents if enabled
             if config.persistence {
                 info!("Saving torrents into SQL from memory...");
@@ -118,13 +132,13 @@ fn start_http_tracker_server(config: &HttpTrackerConfig, tracker: Arc<TorrentTra
     })
 }
 
-async fn start_udp_tracker_server(config: &UdpTrackerConfig, tracker: Arc<TorrentTracker>) -> JoinHandle<()> {
-    let udp_server = UdpServer::new(tracker, config).await.unwrap_or_else(|e| {
+async fn start_udp_tracker_server(config: &UdpTrackerConfig, tracker: Arc<TorrentTracker>, rx: tokio::sync::watch::Receiver<bool>) -> JoinHandle<()> {
+    let udp_server = UdpServer::new(tracker, &config.bind_address).await.unwrap_or_else(|e| {
         panic!("Could not start UDP server: {}", e);
     });
 
     info!("Starting UDP server on: {}", config.bind_address);
     tokio::spawn(async move {
-        udp_server.start().await;
+        udp_server.start(rx).await;
     })
 }

src/main.rs

@@ -24,7 +24,7 @@ pub async fn authenticate(info_hash: &InfoHash, tracker: Arc<TorrentTracker>) ->
     }
 }
 
-pub async fn handle_packet(remote_addr: SocketAddr, payload: &[u8], tracker: Arc<TorrentTracker>) -> Response {
+pub async fn handle_packet(remote_addr: SocketAddr, payload: Vec<u8>, tracker: Arc<TorrentTracker>) -> Response {
     match Request::from_bytes(&payload[..payload.len()], MAX_SCRAPE_TORRENTS).map_err(|_| ServerError::InternalServerError) {
         Ok(request) => {
             let transaction_id = match &request {

src/torrust_udp_tracker/handlers.rs

@@ -1,41 +1,54 @@
+use std::future::Future;
 use std::io::Cursor;
 use std::net::{SocketAddr};
 use std::sync::Arc;
 use aquatic_udp_protocol::{Response};
-use log::debug;
+use log::{debug, info};
 use tokio::net::UdpSocket;
-use crate::{TorrentTracker, UdpTrackerConfig};
+use crate::{TorrentTracker};
 use crate::torrust_udp_tracker::{handle_packet, MAX_PACKET_SIZE};
 
 pub struct UdpServer {
-    socket: UdpSocket,
+    socket: Arc<UdpSocket>,
     tracker: Arc<TorrentTracker>,
 }
 
 impl UdpServer {
-    pub async fn new(tracker: Arc<TorrentTracker>, config: &UdpTrackerConfig) -> Result<UdpServer, std::io::Error> {
-        let srv = UdpSocket::bind(&config.bind_address).await?;
+    pub async fn new(tracker: Arc<TorrentTracker>, bind_address: &str) -> tokio::io::Result<UdpServer> {
+        let socket = UdpSocket::bind(bind_address).await?;
 
         Ok(UdpServer {
-            socket: srv,
+            socket: Arc::new(socket),
             tracker,
         })
     }
 
-    pub async fn start(&self) {
+    pub async fn start(&self, rx: tokio::sync::watch::Receiver<bool>) {
         loop {
+            let mut rx = rx.clone();
             let mut data = [0; MAX_PACKET_SIZE];
-            if let Ok((valid_bytes, remote_addr)) = self.socket.recv_from(&mut data).await {
-                let data = &data[..valid_bytes];
-                debug!("Received {} bytes from {}", data.len(), remote_addr);
-                debug!("{:?}", data);
-                let response = handle_packet(remote_addr, data, self.tracker.clone()).await;
-                self.send_response(remote_addr, response).await;
+            let socket = self.socket.clone();
+            let tracker = self.tracker.clone();
+
+            tokio::select! {
+                _ = rx.changed() => {
+                    info!("Stopping UDP server: {}...", socket.local_addr().unwrap());
+                    break;
+                }
+                Ok((valid_bytes, remote_addr)) = socket.recv_from(&mut data) => {
+                    let payload = data[..valid_bytes].to_vec();
+
+                    debug!("Received {} bytes from {}", payload.len(), remote_addr);
+                    debug!("{:?}", payload);
+
+                    let response = handle_packet(remote_addr, payload, tracker).await;
+                    UdpServer::send_response(socket, remote_addr, response).await;
+                }
             }
         }
     }
 
-    async fn send_response(&self, remote_addr: SocketAddr, response: Response) {
+    async fn send_response(socket: Arc<UdpSocket>, remote_addr: SocketAddr, response: Response) {
         debug!("sending response to: {:?}", &remote_addr);
 
         let buffer = vec![0u8; MAX_PACKET_SIZE];
@@ -47,14 +60,14 @@ impl UdpServer {
                 let inner = cursor.get_ref();
 
                 debug!("{:?}", &inner[..position]);
-                self.send_packet(&remote_addr, &inner[..position]).await;
+                UdpServer::send_packet(socket, &remote_addr, &inner[..position]).await;
             }
             Err(_) => { debug!("could not write response to bytes."); }
         }
     }
 
-    async fn send_packet(&self, remote_addr: &SocketAddr, payload: &[u8]) {
+    async fn send_packet(socket: Arc<UdpSocket>, remote_addr: &SocketAddr, payload: &[u8]) {
         // doesn't matter if it reaches or not
-        let _ = self.socket.send_to(payload, remote_addr).await;
+        let _ = socket.send_to(payload, remote_addr).await;
     }
 }