mixnode: graceful shutdown on ctrl-c (#1304)

* mixnode: add graceful notification to most tasks

* Add note about remaining work

* task/shutdown: add shutdown timer

* changelog: add entry for shutdown

* mixnode: revert some temp changes

* common/task: make sure to use latest tokio

mixnode/src/node/listener/connection_handler/mod.rs

@@ -5,6 +5,7 @@ use crate::node::listener::connection_handler::packet_processing::{
     MixProcessingResult, PacketProcessor,
 };
 use crate::node::packet_delayforwarder::PacketDelayForwardSender;
+use crate::node::ShutdownListener;
 use futures::StreamExt;
 use log::{error, info};
 use nymsphinx::forwarding::packet::MixPacket;
@@ -69,28 +70,40 @@ impl ConnectionHandler {
         }
     }
 
-    pub(crate) async fn handle_connection(self, conn: TcpStream, remote: SocketAddr) {
+    pub(crate) async fn handle_connection(
+        self,
+        conn: TcpStream,
+        remote: SocketAddr,
+        mut shutdown: ShutdownListener,
+    ) {
         debug!("Starting connection handler for {:?}", remote);
         let mut framed_conn = Framed::new(conn, SphinxCodec);
-        while let Some(framed_sphinx_packet) = framed_conn.next().await {
-            match framed_sphinx_packet {
-                Ok(framed_sphinx_packet) => {
-                    // TODO: benchmark spawning tokio task with full processing vs just processing it
-                    // synchronously (without delaying inside of course,
-                    // delay is moved to a global DelayQueue)
-                    // under higher load in single and multi-threaded situation.
+        while !shutdown.is_shutdown() {
+            tokio::select! {
+                Some(framed_sphinx_packet) = framed_conn.next() => {
+                    match framed_sphinx_packet {
+                        Ok(framed_sphinx_packet) => {
+                            // TODO: benchmark spawning tokio task with full processing vs just processing it
+                            // synchronously (without delaying inside of course,
+                            // delay is moved to a global DelayQueue)
+                            // under higher load in single and multi-threaded situation.
 
-                    // in theory we could process multiple sphinx packet from the same connection in parallel,
-                    // but we already handle multiple concurrent connections so if anything, making
-                    // that change would only slow things down
-                    self.handle_received_packet(framed_sphinx_packet);
-                }
-                Err(err) => {
-                    error!(
-                        "The socket connection got corrupted with error: {:?}. Closing the socket",
-                        err
-                    );
-                    return;
+                            // in theory we could process multiple sphinx packet from the same connection in parallel,
+                            // but we already handle multiple concurrent connections so if anything, making
+                            // that change would only slow things down
+                            self.handle_received_packet(framed_sphinx_packet);
+                        }
+                        Err(err) => {
+                            error!(
+                                "The socket connection got corrupted with error: {:?}. Closing the socket",
+                                err
+                            );
+                            return;
+                        }
+                    }
+                },
+                _ = shutdown.recv() => {
+                    log::trace!("ConnectionHandler: received shutdown");
                 }
             }
         }
@@ -99,5 +112,6 @@ impl ConnectionHandler {
             "Closing connection from {:?}",
             framed_conn.into_inner().peer_addr()
         );
+        log::trace!("ConnectionHandler: Exiting");
     }
 }

mixnode/src/node/listener/mod.rs

@@ -8,18 +8,22 @@ use std::process;
 use tokio::net::TcpListener;
 use tokio::task::JoinHandle;
 
+use super::ShutdownListener;
+
 pub(crate) mod connection_handler;
 
 pub(crate) struct Listener {
     address: SocketAddr,
+    shutdown: ShutdownListener,
 }
 
 impl Listener {
-    pub(crate) fn new(address: SocketAddr) -> Self {
-        Listener { address }
+    pub(crate) fn new(address: SocketAddr, shutdown: ShutdownListener) -> Self {
+        Listener { address, shutdown }
     }
 
     async fn run(&mut self, connection_handler: ConnectionHandler) {
+        log::trace!("Starting Listener");
         let listener = match TcpListener::bind(self.address).await {
             Ok(listener) => listener,
             Err(err) => {
@@ -28,15 +32,23 @@ impl Listener {
             }
         };
 
-        loop {
-            match listener.accept().await {
-                Ok((socket, remote_addr)) => {
-                    let handler = connection_handler.clone();
-                    tokio::spawn(handler.handle_connection(socket, remote_addr));
+        while !self.shutdown.is_shutdown() {
+            tokio::select! {
+                connection = listener.accept() => {
+                    match connection {
+                        Ok((socket, remote_addr)) => {
+                            let handler = connection_handler.clone();
+                            tokio::spawn(handler.handle_connection(socket, remote_addr, self.shutdown.clone()));
+                        }
+                        Err(err) => warn!("Failed to accept incoming connection - {:?}", err),
+                    }
+                },
+                _ = self.shutdown.recv() => {
+                    log::trace!("Listener: Received shutdown");
                 }
-                Err(err) => warn!("Failed to accept incoming connection - {:?}", err),
-            }
+            };
         }
+        log::trace!("Listener: Exiting");
     }
 
     pub(crate) fn start(mut self, connection_handler: ConnectionHandler) -> JoinHandle<()> {

mixnode/src/node/mod.rs

@@ -25,6 +25,7 @@ use rand::thread_rng;
 use std::net::SocketAddr;
 use std::process;
 use std::sync::Arc;
+use task::{ShutdownListener, ShutdownNotifier};
 use version_checker::parse_version;
 
 mod http;
@@ -149,11 +150,15 @@ impl MixNode {
         });
     }
 
-    fn start_node_stats_controller(&self) -> (SharedNodeStats, node_statistics::UpdateSender) {
+    fn start_node_stats_controller(
+        &self,
+        shutdown: ShutdownListener,
+    ) -> (SharedNodeStats, node_statistics::UpdateSender) {
         info!("Starting node stats controller...");
         let controller = node_statistics::Controller::new(
             self.config.get_node_stats_logging_delay(),
             self.config.get_node_stats_updating_delay(),
+            shutdown,
         );
         let node_stats_pointer = controller.get_node_stats_data_pointer();
         let update_sender = controller.start();
@@ -165,6 +170,7 @@ impl MixNode {
         &self,
         node_stats_update_sender: node_statistics::UpdateSender,
         delay_forwarding_channel: PacketDelayForwardSender,
+        shutdown: ShutdownListener,
     ) {
         info!("Starting socket listener...");
 
@@ -178,12 +184,13 @@ impl MixNode {
             self.config.get_mix_port(),
         );
 
-        Listener::new(listening_address).start(connection_handler);
+        Listener::new(listening_address, shutdown).start(connection_handler);
     }
 
     fn start_packet_delay_forwarder(
         &mut self,
         node_stats_update_sender: node_statistics::UpdateSender,
+        shutdown: ShutdownListener,
     ) -> PacketDelayForwardSender {
         info!("Starting packet delay-forwarder...");
 
@@ -197,6 +204,7 @@ impl MixNode {
         let mut packet_forwarder = DelayForwarder::new(
             mixnet_client::Client::new(client_config),
             node_stats_update_sender,
+            shutdown,
         );
 
         let packet_sender = packet_forwarder.sender();
@@ -259,7 +267,10 @@ impl MixNode {
         let existing_nodes = match validator_client.get_cached_mixnodes().await {
             Ok(nodes) => nodes,
             Err(err) => {
-                error!("failed to grab initial network mixnodes - {}\n Please try to startup again in few minutes", err);
+                error!(
+                    "failed to grab initial network mixnodes - {err}\n \
+                    Please try to startup again in few minutes",
+                );
                 process::exit(1);
             }
         };
@@ -272,16 +283,26 @@ impl MixNode {
             .map(|node| node.mix_node.identity_key.clone())
     }
 
-    async fn wait_for_interrupt(&self) {
+    async fn wait_for_interrupt(&self, mut shutdown: ShutdownNotifier) {
         if let Err(e) = tokio::signal::ctrl_c().await {
             error!(
-                "There was an error while capturing SIGINT - {:?}. We will terminate regardless",
+                "There was an error while capturing SIGINT - {:?}. \
+                We will terminate regardless",
                 e
             );
         }
         println!(
-            "Received SIGINT - the mixnode will terminate now (threads are not yet nicely stopped, if you see stack traces that's alright)."
+            "Received SIGINT - the mixnode will terminate now \
+            (threads are not yet nicely stopped, if you see stack traces that's alright)."
         );
+
+        log::info!("Sending shutdown");
+        shutdown.signal_shutdown().ok();
+
+        log::info!("Waiting for tasks to finish... (Press ctrl-c to force)");
+        shutdown.wait_for_shutdown().await;
+
+        log::info!("Stopping nym mixnode");
     }
 
     pub async fn run(&mut self) {
@@ -299,15 +320,23 @@ impl MixNode {
             }
         }
 
-        let (node_stats_pointer, node_stats_update_sender) = self.start_node_stats_controller();
-        let delay_forwarding_channel =
-            self.start_packet_delay_forwarder(node_stats_update_sender.clone());
-        self.start_socket_listener(node_stats_update_sender, delay_forwarding_channel);
+        let shutdown = ShutdownNotifier::default();
 
+        let (node_stats_pointer, node_stats_update_sender) =
+            self.start_node_stats_controller(shutdown.subscribe());
+        let delay_forwarding_channel = self
+            .start_packet_delay_forwarder(node_stats_update_sender.clone(), shutdown.subscribe());
+        self.start_socket_listener(
+            node_stats_update_sender,
+            delay_forwarding_channel,
+            shutdown.subscribe(),
+        );
+
+        // TODO: these two also needs to be shutdown
         let atomic_verloc_results = self.start_verloc_measurements();
         self.start_http_api(atomic_verloc_results, node_stats_pointer);
 
         info!("Finished nym mixnode startup procedure - it should now be able to receive mix traffic!");
-        self.wait_for_interrupt().await
+        self.wait_for_interrupt(shutdown).await
     }
 }

mixnode/src/node/node_statistics.rs

@@ -9,6 +9,8 @@ use std::sync::Arc;
 use std::time::{Duration, SystemTime};
 use tokio::sync::{RwLock, RwLockReadGuard};
 
+use super::ShutdownListener;
+
 // convenience aliases
 type PacketsMap = HashMap<String, u64>;
 type PacketDataReceiver = mpsc::UnboundedReceiver<PacketEvent>;
@@ -209,28 +211,45 @@ impl CurrentPacketData {
 struct UpdateHandler {
     current_data: CurrentPacketData,
     update_receiver: PacketDataReceiver,
+    shutdown: ShutdownListener,
 }
 
 impl UpdateHandler {
-    fn new(current_data: CurrentPacketData, update_receiver: PacketDataReceiver) -> Self {
+    fn new(
+        current_data: CurrentPacketData,
+        update_receiver: PacketDataReceiver,
+        shutdown: ShutdownListener,
+    ) -> Self {
         UpdateHandler {
             current_data,
             update_receiver,
+            shutdown,
         }
     }
 
     async fn run(&mut self) {
-        while let Some(packet_data) = self.update_receiver.next().await {
-            match packet_data {
-                PacketEvent::Received => self.current_data.increment_received(),
-                PacketEvent::Sent(destination) => {
-                    self.current_data.increment_sent(destination).await
+        log::trace!("Starting UpdateHandler");
+        while !self.shutdown.is_shutdown() {
+            tokio::select! {
+                Some(packet_data) = self.update_receiver.next() => {
+                    match packet_data {
+                        PacketEvent::Received => self.current_data.increment_received(),
+                        PacketEvent::Sent(destination) => {
+                            self.current_data.increment_sent(destination).await
+                        }
+                        PacketEvent::Dropped(destination) => {
+                            self.current_data.increment_dropped(destination).await
+                        }
+                    }
                 }
-                PacketEvent::Dropped(destination) => {
-                    self.current_data.increment_dropped(destination).await
+                _ = self.shutdown.recv() => {
+                    log::trace!("UpdateHandler: Received shutdown");
+                    break;
                 }
             }
         }
+
+        log::trace!("UpdateHandler: Exiting");
     }
 }
 
@@ -274,6 +293,7 @@ struct StatsUpdater {
     updating_delay: Duration,
     current_packet_data: CurrentPacketData,
     current_stats: SharedNodeStats,
+    shutdown: ShutdownListener,
 }
 
 impl StatsUpdater {
@@ -281,11 +301,13 @@ impl StatsUpdater {
         updating_delay: Duration,
         current_packet_data: CurrentPacketData,
         current_stats: SharedNodeStats,
+        shutdown: ShutdownListener,
     ) -> Self {
         StatsUpdater {
             updating_delay,
             current_packet_data,
             current_stats,
+            shutdown,
         }
     }
 
@@ -295,11 +317,16 @@ impl StatsUpdater {
         self.current_stats.update(received, sent, dropped).await;
     }
 
-    async fn run(&self) {
-        loop {
-            tokio::time::sleep(self.updating_delay).await;
-            self.update_stats().await
+    async fn run(&mut self) {
+        while !self.shutdown.is_shutdown() {
+            tokio::select! {
+                _ = tokio::time::sleep(self.updating_delay) => self.update_stats().await,
+                _ = self.shutdown.recv() => {
+                    log::trace!("StatsUpdater: Received shutdown");
+                }
+            }
         }
+        log::trace!("StatsUpdater: Exiting");
     }
 }
 
@@ -308,13 +335,15 @@ impl StatsUpdater {
 struct PacketStatsConsoleLogger {
     logging_delay: Duration,
     stats: SharedNodeStats,
+    shutdown: ShutdownListener,
 }
 
 impl PacketStatsConsoleLogger {
-    fn new(logging_delay: Duration, stats: SharedNodeStats) -> Self {
+    fn new(logging_delay: Duration, stats: SharedNodeStats, shutdown: ShutdownListener) -> Self {
         PacketStatsConsoleLogger {
             logging_delay,
             stats,
+            shutdown,
         }
     }
 
@@ -387,10 +416,16 @@ impl PacketStatsConsoleLogger {
     }
 
     async fn run(&mut self) {
-        loop {
-            tokio::time::sleep(self.logging_delay).await;
-            self.log_running_stats().await;
+        log::trace!("Starting PacketStatsConsoleLogger");
+        while !self.shutdown.is_shutdown() {
+            tokio::select! {
+                _ = tokio::time::sleep(self.logging_delay) => self.log_running_stats().await,
+                _ = self.shutdown.recv() => {
+                    log::trace!("PacketStatsConsoleLogger: Received shutdown");
+                }
+            };
         }
+        log::trace!("PacketStatsConsoleLogger: Exiting");
     }
 }
 
@@ -413,19 +448,32 @@ pub struct Controller {
 }
 
 impl Controller {
-    pub(crate) fn new(logging_delay: Duration, stats_updating_delay: Duration) -> Self {
+    pub(crate) fn new(
+        logging_delay: Duration,
+        stats_updating_delay: Duration,
+        shutdown: ShutdownListener,
+    ) -> Self {
         let (sender, receiver) = mpsc::unbounded();
         let shared_packet_data = CurrentPacketData::new();
         let shared_node_stats = SharedNodeStats::new();
 
         Controller {
-            update_handler: UpdateHandler::new(shared_packet_data.clone(), receiver),
+            update_handler: UpdateHandler::new(
+                shared_packet_data.clone(),
+                receiver,
+                shutdown.clone(),
+            ),
             update_sender: UpdateSender::new(sender),
-            console_logger: PacketStatsConsoleLogger::new(logging_delay, shared_node_stats.clone()),
+            console_logger: PacketStatsConsoleLogger::new(
+                logging_delay,
+                shared_node_stats.clone(),
+                shutdown.clone(),
+            ),
             stats_updater: StatsUpdater::new(
                 stats_updating_delay,
                 shared_packet_data,
                 shared_node_stats.clone(),
+                shutdown,
             ),
             node_stats: shared_node_stats,
         }
@@ -441,7 +489,7 @@ impl Controller {
     pub(crate) fn start(self) -> UpdateSender {
         // move out of self
         let mut update_handler = self.update_handler;
-        let stats_updater = self.stats_updater;
+        let mut stats_updater = self.stats_updater;
         let mut console_logger = self.console_logger;
 
         tokio::spawn(async move { update_handler.run().await });
@@ -455,12 +503,15 @@ impl Controller {
 #[cfg(test)]
 mod tests {
     use super::*;
+    use task::ShutdownNotifier;
 
     #[tokio::test]
     async fn node_stats_reported_are_received() {
         let logging_delay = Duration::from_millis(20);
         let stats_updating_delay = Duration::from_millis(10);
-        let node_stats_controller = Controller::new(logging_delay, stats_updating_delay);
+        let shutdown = ShutdownNotifier::default();
+        let node_stats_controller =
+            Controller::new(logging_delay, stats_updating_delay, shutdown.subscribe());
 
         let node_stats_pointer = node_stats_controller.get_node_stats_data_pointer();
         let update_sender = node_stats_controller.start();

mixnode/src/node/packet_delayforwarder.rs

@@ -9,6 +9,8 @@ use nymsphinx::forwarding::packet::MixPacket;
 use std::io;
 use tokio::time::Instant;
 
+use super::ShutdownListener;
+
 // Delay + MixPacket vs Instant + MixPacket
 
 // rather than using Duration directly, we use an Instant, this way we minimise skew due to
@@ -26,13 +28,18 @@ where
     packet_sender: PacketDelayForwardSender,
     packet_receiver: PacketDelayForwardReceiver,
     node_stats_update_sender: UpdateSender,
+    shutdown: ShutdownListener,
 }
 
 impl<C> DelayForwarder<C>
 where
     C: mixnet_client::SendWithoutResponse,
 {
-    pub(crate) fn new(client: C, node_stats_update_sender: UpdateSender) -> DelayForwarder<C> {
+    pub(crate) fn new(
+        client: C,
+        node_stats_update_sender: UpdateSender,
+        shutdown: ShutdownListener,
+    ) -> DelayForwarder<C> {
         let (packet_sender, packet_receiver) = mpsc::unbounded();
 
         DelayForwarder::<C> {
@@ -41,6 +48,7 @@ where
             packet_sender,
             packet_receiver,
             node_stats_update_sender,
+            shutdown,
         }
     }
 
@@ -97,6 +105,7 @@ where
     }
 
     pub(crate) async fn run(&mut self) {
+        log::trace!("Starting DelayForwarder");
         loop {
             tokio::select! {
                 delayed = self.delay_queue.next() => {
@@ -107,8 +116,13 @@ where
                     // and hence it can't happen that ALL senders are dropped
                     self.handle_new_packet(new_packet.unwrap())
                 }
+                _ = self.shutdown.recv() => {
+                    log::trace!("DelayForwarder: Received shutdown");
+                    break;
+                }
             }
         }
+        log::trace!("DelayForwarder: Exiting");
     }
 }
 
@@ -120,6 +134,8 @@ mod tests {
     use std::sync::{Arc, Mutex};
     use std::time::Duration;
 
+    use task::ShutdownNotifier;
+
     use nymsphinx::addressing::nodes::NymNodeRoutingAddress;
     use nymsphinx_params::packet_sizes::PacketSize;
     use nymsphinx_params::PacketMode;
@@ -189,7 +205,9 @@ mod tests {
         let node_stats_update_sender = UpdateSender::new(stats_sender);
         let client = TestClient::default();
         let client_packets_sent = client.packets_sent.clone();
-        let mut delay_forwarder = DelayForwarder::new(client, node_stats_update_sender);
+        let shutdown = ShutdownNotifier::default();
+        let mut delay_forwarder =
+            DelayForwarder::new(client, node_stats_update_sender, shutdown.subscribe());
         let packet_sender = delay_forwarder.sender();
 
         // Spawn the worker, listening on packet_sender channel