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
This commit is contained in:
Jon Häggblad
2022-06-10 11:36:48 +02:00
committed by GitHub
parent e1b5407613
commit 8df24b8ce2
7 changed files with 221 additions and 106 deletions
Generated
+1
View File
@@ -2852,6 +2852,7 @@ dependencies = [
"nymsphinx-types",
"rand 0.8.5",
"serde",
"task",
"tokio",
"tokio-util 0.7.3",
"url",
+1
View File
@@ -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" }
@@ -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")
}
}
}
}
+64 -32
View File
@@ -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<ConnectionHandler>,
shutdown: ShutdownListener,
}
impl PacketListener {
pub(crate) fn new(address: SocketAddr, identity: Arc<identity::KeyPair>) -> Self {
pub(crate) fn new(
address: SocketAddr,
identity: Arc<identity::KeyPair>,
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<Self>, conn: TcpStream, remote: SocketAddr) {
pub(crate) async fn handle_connection(
self: Arc<Self>,
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");
}
}
}
}
+71 -26
View File
@@ -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<PacketSender>,
packet_listener: Arc<PacketListener>,
shutdown_listener: ShutdownListener,
currently_used_api: usize,
@@ -177,7 +181,11 @@ pub struct VerlocMeasurer {
}
impl VerlocMeasurer {
pub fn new(mut config: Config, identity: Arc<identity::KeyPair>) -> Self {
pub fn new(
mut config: Config,
identity: Arc<identity::KeyPair>,
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<TestedNode>) {
async fn perform_measurement(&self, nodes_to_test: Vec<TestedNode>) -> 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::<FuturesUnordered<_>>();
// 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,
}
+72 -44
View File
@@ -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<identity::KeyPair>,
// 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<Self>,
tested_node: TestedNode,
) -> Result<Measurement, RttError> {
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...
+7 -4
View File
@@ -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!");