From 8df24b8ce2239dee16ca7ba4e125f670d9e96739 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jon=20H=C3=A4ggblad?= Date: Fri, 10 Jun 2022 11:36:48 +0200 Subject: [PATCH] verloc: signalling for graceful shutdown (#1323) * common/verloc: signalling for graceful shutdown * verloc: make shutdown handler not optional * verloc: logging without explicit target * rustmt * mixnode: note about rocket * verloc: remove accidental duplicate block * verloc: pass shutdown handler as argument to connection handler --- Cargo.lock | 1 + common/mixnode-common/Cargo.toml | 1 + common/mixnode-common/src/verloc/error.rs | 5 + common/mixnode-common/src/verloc/listener.rs | 96 ++++++++++----- common/mixnode-common/src/verloc/mod.rs | 97 +++++++++++----- common/mixnode-common/src/verloc/sender.rs | 116 ++++++++++++------- mixnode/src/node/mod.rs | 11 +- 7 files changed, 221 insertions(+), 106 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 48c921e8d3..37bb1e7bfd 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2852,6 +2852,7 @@ dependencies = [ "nymsphinx-types", "rand 0.8.5", "serde", + "task", "tokio", "tokio-util 0.7.3", "url", diff --git a/common/mixnode-common/Cargo.toml b/common/mixnode-common/Cargo.toml index 77c25d0455..212a3b30d5 100644 --- a/common/mixnode-common/Cargo.toml +++ b/common/mixnode-common/Cargo.toml @@ -26,5 +26,6 @@ nymsphinx-forwarding = { path = "../nymsphinx/forwarding" } nymsphinx-framing = { path = "../nymsphinx/framing" } nymsphinx-params = { path = "../nymsphinx/params" } nymsphinx-types = { path = "../nymsphinx/types" } +task = { path = "../task" } validator-client = { path = "../client-libs/validator-client" } version-checker = { path = "../version-checker" } diff --git a/common/mixnode-common/src/verloc/error.rs b/common/mixnode-common/src/verloc/error.rs index 5e1cb460da..44a408c77a 100644 --- a/common/mixnode-common/src/verloc/error.rs +++ b/common/mixnode-common/src/verloc/error.rs @@ -24,6 +24,8 @@ pub enum RttError { ConnectionWriteTimeout(String), UnexpectedReplySequence, + + ShutdownReceived, } impl Display for RttError { @@ -69,6 +71,9 @@ impl Display for RttError { f, "The received reply packet had an unexpected sequence number" ), + RttError::ShutdownReceived => { + write!(f, "Shutdown signal received") + } } } } diff --git a/common/mixnode-common/src/verloc/listener.rs b/common/mixnode-common/src/verloc/listener.rs index c2522ac223..8f49779130 100644 --- a/common/mixnode-common/src/verloc/listener.rs +++ b/common/mixnode-common/src/verloc/listener.rs @@ -11,6 +11,7 @@ use std::fmt::{Display, Formatter}; use std::net::SocketAddr; use std::sync::Arc; use std::{fmt, io, process}; +use task::ShutdownListener; use tokio::io::AsyncWriteExt; use tokio::net::{TcpListener, TcpStream}; use tokio_util::codec::{Decoder, Encoder, Framed}; @@ -18,13 +19,19 @@ use tokio_util::codec::{Decoder, Encoder, Framed}; pub(crate) struct PacketListener { address: SocketAddr, connection_handler: Arc, + shutdown: ShutdownListener, } impl PacketListener { - pub(crate) fn new(address: SocketAddr, identity: Arc) -> Self { + pub(crate) fn new( + address: SocketAddr, + identity: Arc, + shutdown: ShutdownListener, + ) -> Self { PacketListener { address, connection_handler: Arc::new(ConnectionHandler { identity }), + shutdown, } } } @@ -34,24 +41,37 @@ impl PacketListener { let listener = match TcpListener::bind(self.address).await { Ok(listener) => listener, Err(err) => { - error!("Failed to bind to {} - {}. Are you sure nothing else is running on the specified port and your user has sufficient permission to bind to the requested address?", self.address, err); + error!( + "Failed to bind to {} - {}. Are you sure nothing else is running on the specified port and your user has sufficient permission to bind to the requested address?", + self.address, err + ); process::exit(1); } }; info!("Started listening for echo packets on {}", self.address); - loop { + let mut shutdown_listener = self.shutdown.clone(); + + while !shutdown_listener.is_shutdown() { // cloning the arc as each accepted socket is handled in separate task let connection_handler = Arc::clone(&self.connection_handler); + let handler_shutdown_listener = self.shutdown.clone(); - match listener.accept().await { - Ok((socket, remote_addr)) => { - debug!("New verloc connection from {}", remote_addr); + tokio::select! { + socket = listener.accept() => { + match socket { + Ok((socket, remote_addr)) => { + debug!("New verloc connection from {}", remote_addr); - tokio::spawn(connection_handler.handle_connection(socket, remote_addr)); + tokio::spawn(connection_handler.handle_connection(socket, remote_addr, handler_shutdown_listener)); + } + Err(err) => warn!("Failed to accept incoming connection - {:?}", err), + } + }, + _ = shutdown_listener.recv() => { + log::trace!("PacketListener: Received shutdown"); } - Err(err) => warn!("Failed to accept incoming connection - {:?}", err), } } } @@ -67,34 +87,46 @@ impl ConnectionHandler { packet.construct_reply(self.identity.private_key()) } - pub(crate) async fn handle_connection(self: Arc, conn: TcpStream, remote: SocketAddr) { + pub(crate) async fn handle_connection( + self: Arc, + conn: TcpStream, + remote: SocketAddr, + mut shutdown_listener: ShutdownListener, + ) { debug!("Starting connection handler for {:?}", remote); let mut framed_conn = Framed::new(conn, EchoPacketCodec); - while let Some(echo_packet) = framed_conn.next().await { - // handle echo packet - let reply_packet = match echo_packet { - Ok(echo_packet) => self.handle_echo_packet(echo_packet), - Err(err) => { - error!( - "The socket connection got corrupted with error: {}. Closing the socket", - err - ); - return; - } - }; + while !shutdown_listener.is_shutdown() { + tokio::select! { + Some(echo_packet) = framed_conn.next() => { + // handle echo packet + let reply_packet = match echo_packet { + Ok(echo_packet) => self.handle_echo_packet(echo_packet), + Err(err) => { + error!( + "The socket connection got corrupted with error: {}. Closing the socket", + err + ); + return; + } + }; - // write back the reply (note the lack of framing) - if let Err(err) = framed_conn - .get_mut() - .write_all(reply_packet.to_bytes().as_ref()) - .await - { - error!( - "Failed to write reply packet back to the sender - {}. Closing the socket on our end", - err - ); - return; + // write back the reply (note the lack of framing) + if let Err(err) = framed_conn + .get_mut() + .write_all(reply_packet.to_bytes().as_ref()) + .await + { + error!( + "Failed to write reply packet back to the sender - {}. Closing the socket on our end", + err + ); + return; + } + }, + _ = shutdown_listener.recv() => { + trace!("ConnectionHandler: Shutdown received"); + } } } } diff --git a/common/mixnode-common/src/verloc/mod.rs b/common/mixnode-common/src/verloc/mod.rs index 623bdc9e18..610a7a1a66 100644 --- a/common/mixnode-common/src/verloc/mod.rs +++ b/common/mixnode-common/src/verloc/mod.rs @@ -2,7 +2,6 @@ // SPDX-License-Identifier: Apache-2.0 use crate::verloc::listener::PacketListener; -pub use crate::verloc::measurement::{AtomicVerlocResult, Verloc, VerlocResult}; use crate::verloc::sender::{PacketSender, TestedNode}; use crypto::asymmetric::identity; use futures::stream::FuturesUnordered; @@ -13,11 +12,14 @@ use rand::thread_rng; use std::net::{SocketAddr, ToSocketAddrs}; use std::sync::Arc; use std::time::Duration; +use task::ShutdownListener; use tokio::task::JoinHandle; use tokio::time::sleep; use url::Url; use version_checker::parse_version; +pub use crate::verloc::measurement::{AtomicVerlocResult, Verloc, VerlocResult}; + pub mod error; pub(crate) mod listener; pub(crate) mod measurement; @@ -137,9 +139,10 @@ impl ConfigBuilder { pub fn build(self) -> Config { // panics here are fine as those are only ever constructed at the initial setup - if self.0.validator_api_urls.is_empty() { - panic!("at least one validator endpoint must be provided") - } + assert!( + !self.0.validator_api_urls.is_empty(), + "at least one validator endpoint must be provided", + ); self.0 } } @@ -165,6 +168,7 @@ pub struct VerlocMeasurer { config: Config, packet_sender: Arc, packet_listener: Arc, + shutdown_listener: ShutdownListener, currently_used_api: usize, @@ -177,7 +181,11 @@ pub struct VerlocMeasurer { } impl VerlocMeasurer { - pub fn new(mut config: Config, identity: Arc) -> Self { + pub fn new( + mut config: Config, + identity: Arc, + shutdown_listener: ShutdownListener, + ) -> Self { config.validator_api_urls.shuffle(&mut thread_rng()); VerlocMeasurer { @@ -187,11 +195,14 @@ impl VerlocMeasurer { config.packet_timeout, config.connection_timeout, config.delay_between_packets, + shutdown_listener.clone(), )), packet_listener: Arc::new(PacketListener::new( config.listening_address, Arc::clone(&identity), + shutdown_listener.clone(), )), + shutdown_listener, currently_used_api: 0, validator_client: validator_client::ApiClient::new( config.validator_api_urls[0].clone(), @@ -222,7 +233,11 @@ impl VerlocMeasurer { tokio::spawn(packet_listener.run()) } - async fn perform_measurement(&self, nodes_to_test: Vec) { + async fn perform_measurement(&self, nodes_to_test: Vec) -> MeasurementOutcome { + log::trace!("Performing measurements"); + + let mut shutdown_listener = self.shutdown_listener.clone(); + for chunk in nodes_to_test.chunks(self.config.tested_nodes_batch_size) { let mut chunk_results = Vec::with_capacity(chunk.len()); @@ -246,33 +261,44 @@ impl VerlocMeasurer { .collect::>(); // exhaust the results - while let Some(result) = measurement_chunk.next().await { - // if we receive JoinError it means the task failed to get executed, so either there's a bigger issue with tokio - // or there was a panic inside the task itself. In either case, we should just terminate ourselves. - let execution_result = result.expect("the measurement task panicked!"); - let measurement_result = match execution_result.0 { - Err(err) => { - debug!( - "Failed to perform measurement for {} - {}", - execution_result.1.to_base58_string(), - err - ); - None + while !shutdown_listener.is_shutdown() { + tokio::select! { + Some(result) = measurement_chunk.next() => { + // if we receive JoinError it means the task failed to get executed, so either there's a bigger issue with tokio + // or there was a panic inside the task itself. In either case, we should just terminate ourselves. + let execution_result = result.expect("the measurement task panicked!"); + let measurement_result = match execution_result.0 { + Err(err) => { + debug!( + "Failed to perform measurement for {} - {}", + execution_result.1.to_base58_string(), + err + ); + None + } + Ok(result) => Some(result), + }; + chunk_results.push(Verloc::new(execution_result.1, measurement_result)); + }, + _ = shutdown_listener.recv() => { + trace!("Shutdown received while measuring"); + return MeasurementOutcome::Shutdown; } - Ok(result) => Some(result), - }; - chunk_results.push(Verloc::new(execution_result.1, measurement_result)); + } } // update the results vector with chunks as they become available (by default every 50 nodes) self.results.append_results(chunk_results).await; } + + MeasurementOutcome::Done } pub async fn run(&mut self) { self.start_listening(); - loop { - info!(target: "verloc", "Starting verloc measurements"); + + while !self.shutdown_listener.is_shutdown() { + info!("Starting verloc measurements"); // TODO: should we also measure gateways? let all_mixes = match self.validator_client.get_cached_mixnodes().await { @@ -322,13 +348,32 @@ impl VerlocMeasurer { // on start of each run remove old results self.results.reset_results(tested_nodes.len()).await; - self.perform_measurement(tested_nodes).await; + if let MeasurementOutcome::Shutdown = self.perform_measurement(tested_nodes).await { + log::trace!("Shutting down after aborting measurements"); + break; + } // write current time to "run finished" field self.results.finish_measurements().await; - info!(target: "verloc", "Finished performing verloc measurements. The next one will happen in {:?}", self.config.testing_interval); - sleep(self.config.testing_interval).await + info!( + "Finished performing verloc measurements. The next one will happen in {:?}", + self.config.testing_interval + ); + + tokio::select! { + _ = sleep(self.config.testing_interval) => {}, + _ = self.shutdown_listener.recv() => { + log::trace!("Shutdown received while sleeping"); + } + } } + + log::trace!("Verloc: Exiting"); } } + +enum MeasurementOutcome { + Done, + Shutdown, +} diff --git a/common/mixnode-common/src/verloc/sender.rs b/common/mixnode-common/src/verloc/sender.rs index 81fd5667d2..bed34f4474 100644 --- a/common/mixnode-common/src/verloc/sender.rs +++ b/common/mixnode-common/src/verloc/sender.rs @@ -7,10 +7,11 @@ use crate::verloc::packet::{EchoPacket, ReplyPacket}; use crypto::asymmetric::identity; use log::*; use rand::{thread_rng, Rng}; -use std::io; use std::net::SocketAddr; use std::sync::Arc; use std::time::Duration; +use std::{fmt, io}; +use task::ShutdownListener; use tokio::io::{AsyncReadExt, AsyncWriteExt}; use tokio::net::TcpStream; use tokio::time::sleep; @@ -27,6 +28,16 @@ impl TestedNode { } } +impl fmt::Display for TestedNode { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!( + f, + "TestedNode(id: {}, address: {})", + self.identity, self.address + ) + } +} + pub(crate) struct PacketSender { identity: Arc, // timeout for receiving before sending new one @@ -34,6 +45,7 @@ pub(crate) struct PacketSender { packet_timeout: Duration, connection_timeout: Duration, delay_between_packets: Duration, + shutdown_listener: ShutdownListener, } impl PacketSender { @@ -43,6 +55,7 @@ impl PacketSender { packet_timeout: Duration, connection_timeout: Duration, delay_between_packets: Duration, + shutdown_listener: ShutdownListener, ) -> Self { PacketSender { identity, @@ -50,6 +63,7 @@ impl PacketSender { packet_timeout, connection_timeout, delay_between_packets, + shutdown_listener, } } @@ -69,6 +83,8 @@ impl PacketSender { self: Arc, tested_node: TestedNode, ) -> Result { + let mut shutdown_listener = self.shutdown_listener.clone(); + let mut conn = match tokio::time::timeout( self.connection_timeout, TcpStream::connect(tested_node.address), @@ -98,35 +114,40 @@ impl PacketSender { let start = tokio::time::Instant::now(); // TODO: should we get the start time after or before actually sending the data? // there's going to definitely some scheduler and network stack bias here - match tokio::time::timeout( - self.packet_timeout, - conn.write_all(packet.to_bytes().as_ref()), - ) - .await - { - Err(_timeout) => { - let identity_string = tested_node.identity.to_base58_string(); - debug!( - "failed to write echo packet to {} within {:?}. Stopping the test.", - identity_string, self.packet_timeout - ); - return Err(RttError::UnexpectedConnectionFailureWrite( - identity_string, - io::ErrorKind::TimedOut.into(), - )); - } - Ok(Err(err)) => { - let identity_string = tested_node.identity.to_base58_string(); - debug!( - "failed to write echo packet to {} - {}. Stopping the test.", - identity_string, err - ); - return Err(RttError::UnexpectedConnectionFailureWrite( - identity_string, - err, - )); - } - Ok(Ok(_)) => {} + let packet_bytes = packet.to_bytes(); + + tokio::select! { + write = tokio::time::timeout(self.packet_timeout, conn.write_all(packet_bytes.as_ref())) => { + match write { + Err(_timeout) => { + let identity_string = tested_node.identity.to_base58_string(); + debug!( + "failed to write echo packet to {} within {:?}. Stopping the test.", + identity_string, self.packet_timeout + ); + return Err(RttError::UnexpectedConnectionFailureWrite( + identity_string, + io::ErrorKind::TimedOut.into(), + )); + } + Ok(Err(err)) => { + let identity_string = tested_node.identity.to_base58_string(); + debug!( + "failed to write echo packet to {} - {}. Stopping the test.", + identity_string, err + ); + return Err(RttError::UnexpectedConnectionFailureWrite( + identity_string, + err, + )); + } + Ok(Ok(_)) => {} + } + }, + _ = shutdown_listener.recv() => { + log::trace!("PacketSender: Received shutdown while sending"); + return Err(RttError::ShutdownReceived); + }, } // there's absolutely no need to put a codec on ReplyPackets as we know exactly @@ -147,21 +168,28 @@ impl PacketSender { ReplyPacket::try_from_bytes(&buf, &tested_node.identity) }; - let reply_packet = - match tokio::time::timeout(self.packet_timeout, reply_packet_future).await { - Ok(reply_packet) => reply_packet, - Err(_timeout) => { - // TODO: should we continue regardless (with the rest of the packets, or abandon the whole thing?) - // Note: if we decide to continue, it would increase the complexity of the whole thing - debug!( - "failed to receive reply to our echo packet within {:?}. Stopping the test", - self.packet_timeout - ); - return Err(RttError::ConnectionReadTimeout( - tested_node.identity.to_base58_string(), - )); + let reply_packet = tokio::select! { + reply = tokio::time::timeout(self.packet_timeout, reply_packet_future) => { + match reply { + Ok(reply_packet) => reply_packet, + Err(_timeout) => { + // TODO: should we continue regardless (with the rest of the packets, or abandon the whole thing?) + // Note: if we decide to continue, it would increase the complexity of the whole thing + debug!( + "failed to receive reply to our echo packet within {:?}. Stopping the test", + self.packet_timeout + ); + return Err(RttError::ConnectionReadTimeout( + tested_node.identity.to_base58_string(), + )); + } } - }; + }, + _ = shutdown_listener.recv() => { + log::trace!("PacketSender: Received shutdown while waiting for reply"); + return Err(RttError::ShutdownReceived); + } + }; let reply_packet = reply_packet?; // make sure it's actually the expected packet... diff --git a/mixnode/src/node/mod.rs b/mixnode/src/node/mod.rs index b107e698a3..75d4b47ec0 100644 --- a/mixnode/src/node/mod.rs +++ b/mixnode/src/node/mod.rs @@ -213,7 +213,7 @@ impl MixNode { packet_sender } - fn start_verloc_measurements(&self) -> AtomicVerlocResult { + fn start_verloc_measurements(&self, shutdown: ShutdownListener) -> AtomicVerlocResult { info!("Starting the round-trip-time measurer..."); // this is a sanity check to make sure we didn't mess up with the minimum version at some point @@ -246,7 +246,8 @@ impl MixNode { .validator_api_urls(self.config.get_validator_api_endpoints()) .build(); - let mut verloc_measurer = VerlocMeasurer::new(config, Arc::clone(&self.identity_keypair)); + let mut verloc_measurer = + VerlocMeasurer::new(config, Arc::clone(&self.identity_keypair), shutdown); let atomic_verloc_results = verloc_measurer.get_verloc_results_pointer(); tokio::spawn(async move { verloc_measurer.run().await }); atomic_verloc_results @@ -331,9 +332,11 @@ impl MixNode { delay_forwarding_channel, shutdown.subscribe(), ); + let atomic_verloc_results = self.start_verloc_measurements(shutdown.subscribe()); - // TODO: these two also needs to be shutdown - let atomic_verloc_results = self.start_verloc_measurements(); + // Rocket handles shutdown on it's own, but its shutdown handling should be incorporated + // with that of the rest of the tasks. + // Currently it's runtime is forcefully terminated once the mixnode exits. 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!");