diff --git a/clients/native/src/websocket/handler.rs b/clients/native/src/websocket/handler.rs index 3185d71900..3d199cc152 100644 --- a/clients/native/src/websocket/handler.rs +++ b/clients/native/src/websocket/handler.rs @@ -422,7 +422,7 @@ impl Handler { ) { // We don't want a crash in the connection handler to trigger a shutdown of the whole // process. - task_client.mark_as_success(); + task_client.disarm(); let ws_stream = match accept_async(socket).await { Ok(ws_stream) => ws_stream, diff --git a/common/client-core/src/client/base_client/mod.rs b/common/client-core/src/client/base_client/mod.rs index 7d494bce35..c91a7569fc 100644 --- a/common/client-core/src/client/base_client/mod.rs +++ b/common/client-core/src/client/base_client/mod.rs @@ -455,7 +455,7 @@ where Err(ClientCoreError::CustomGatewaySelectionExpected) } else { // and make sure to invalidate the task client so we wouldn't cause premature shutdown - shutdown.mark_as_success(); + shutdown.disarm(); custom_gateway_transceiver.set_packet_router(packet_router)?; Ok(custom_gateway_transceiver) }; @@ -562,7 +562,7 @@ where if topology_config.disable_refreshing { // if we're not spawning the refresher, don't cause shutdown immediately info!("The topology refesher is not going to be started"); - shutdown.mark_as_success(); + shutdown.disarm(); } else { // don't spawn the refresher if we don't want to be refreshing the topology. // only use the initial values obtained diff --git a/common/client-libs/gateway-client/src/packet_router.rs b/common/client-libs/gateway-client/src/packet_router.rs index cd3cc23d51..ab94ecbac3 100644 --- a/common/client-libs/gateway-client/src/packet_router.rs +++ b/common/client-libs/gateway-client/src/packet_router.rs @@ -70,8 +70,8 @@ impl PacketRouter { Ok(()) } - pub fn mark_as_success(&mut self) { - self.shutdown.mark_as_success(); + pub fn disarm(&mut self) { + self.shutdown.disarm(); } } diff --git a/common/client-libs/gateway-client/src/socket_state.rs b/common/client-libs/gateway-client/src/socket_state.rs index 9b9982f852..d7754c5341 100644 --- a/common/client-libs/gateway-client/src/socket_state.rs +++ b/common/client-libs/gateway-client/src/socket_state.rs @@ -113,8 +113,8 @@ impl PartiallyDelegatedRouter { let return_res = match ret { Err(err) => self.stream_return.send(Err(err)), Ok(_) => { - self.packet_router.mark_as_success(); - task_client.mark_as_success(); + self.packet_router.disarm(); + task_client.disarm(); self.stream_return.send(Ok(split_stream)) } }; diff --git a/common/mixnode-common/src/verloc/listener.rs b/common/mixnode-common/src/verloc/listener.rs index fada769a76..aa58b56b8e 100644 --- a/common/mixnode-common/src/verloc/listener.rs +++ b/common/mixnode-common/src/verloc/listener.rs @@ -57,7 +57,7 @@ impl PacketListener { // cloning the arc as each accepted socket is handled in separate task let connection_handler = Arc::clone(&self.connection_handler); let mut handler_shutdown_listener = self.shutdown.clone(); - handler_shutdown_listener.mark_as_success(); + handler_shutdown_listener.disarm(); tokio::select! { socket = listener.accept() => { diff --git a/common/mixnode-common/src/verloc/mod.rs b/common/mixnode-common/src/verloc/mod.rs index 25c7bd8881..68004105d2 100644 --- a/common/mixnode-common/src/verloc/mod.rs +++ b/common/mixnode-common/src/verloc/mod.rs @@ -245,7 +245,7 @@ impl VerlocMeasurer { } let mut shutdown_listener = self.shutdown_listener.clone().named("VerlocMeasurement"); - shutdown_listener.mark_as_success(); + shutdown_listener.disarm(); for chunk in nodes_to_test.chunks(self.config.tested_nodes_batch_size) { let mut chunk_results = Vec::with_capacity(chunk.len()); diff --git a/common/mixnode-common/src/verloc/sender.rs b/common/mixnode-common/src/verloc/sender.rs index 1549af972f..5887a22a6e 100644 --- a/common/mixnode-common/src/verloc/sender.rs +++ b/common/mixnode-common/src/verloc/sender.rs @@ -84,7 +84,7 @@ impl PacketSender { tested_node: TestedNode, ) -> Result { let mut shutdown_listener = self.shutdown_listener.fork(tested_node.address.to_string()); - shutdown_listener.mark_as_success(); + shutdown_listener.disarm(); let mut conn = match tokio::time::timeout( self.connection_timeout, diff --git a/common/socks5-client-core/src/socks/client.rs b/common/socks5-client-core/src/socks/client.rs index cca9cd6b53..572849d296 100644 --- a/common/socks5-client-core/src/socks/client.rs +++ b/common/socks5-client-core/src/socks/client.rs @@ -218,7 +218,7 @@ impl SocksClient { packet_type: Option, ) -> Self { // If this task fails and exits, we don't want to send shutdown signal - shutdown_listener.mark_as_success(); + shutdown_listener.disarm(); let connection_id = Self::generate_random(); @@ -294,7 +294,7 @@ impl SocksClient { .shutdown() .await .map_err(|source| SocksProxyError::SocketShutdownFailure { source })?; - self.shutdown_listener.mark_as_success(); + self.shutdown_listener.disarm(); Ok(()) } diff --git a/common/socks5/proxy-helpers/src/proxy_runner/inbound.rs b/common/socks5/proxy-helpers/src/proxy_runner/inbound.rs index 82e6396a5f..e63df37851 100644 --- a/common/socks5/proxy-helpers/src/proxy_runner/inbound.rs +++ b/common/socks5/proxy-helpers/src/proxy_runner/inbound.rs @@ -172,6 +172,6 @@ where trace!("{} - inbound closed", connection_id); shutdown_notify.notify_one(); - shutdown_listener.mark_as_success(); + shutdown_listener.disarm(); reader } diff --git a/common/socks5/proxy-helpers/src/proxy_runner/mod.rs b/common/socks5/proxy-helpers/src/proxy_runner/mod.rs index 38499ea59c..970378d811 100644 --- a/common/socks5/proxy-helpers/src/proxy_runner/mod.rs +++ b/common/socks5/proxy-helpers/src/proxy_runner/mod.rs @@ -148,7 +148,7 @@ where } pub fn into_inner(mut self) -> (TcpStream, ConnectionReceiver) { - self.shutdown_listener.mark_as_success(); + self.shutdown_listener.disarm(); ( self.socket.take().unwrap(), self.mix_receiver.take().unwrap(), diff --git a/common/socks5/proxy-helpers/src/proxy_runner/outbound.rs b/common/socks5/proxy-helpers/src/proxy_runner/outbound.rs index 72d8824d2b..77f5ae2b59 100644 --- a/common/socks5/proxy-helpers/src/proxy_runner/outbound.rs +++ b/common/socks5/proxy-helpers/src/proxy_runner/outbound.rs @@ -90,6 +90,6 @@ pub(super) async fn run_outbound( trace!("{} - outbound closed", connection_id); shutdown_notify.notify_one(); - shutdown_listener.mark_as_success(); + shutdown_listener.disarm(); (writer, mix_receiver) } diff --git a/common/task/src/manager.rs b/common/task/src/manager.rs index ab03870dc0..f880629bac 100644 --- a/common/task/src/manager.rs +++ b/common/task/src/manager.rs @@ -470,12 +470,8 @@ impl TaskClient { // This listener should to *not* notify the ShutdownNotifier to shutdown when dropped. For // example when we clone the listener for a task handling connections, we often want to drop // without signal failure. - pub fn mark_as_success(&mut self) { - self.mode.set_should_not_signal_on_drop(); - } - pub fn disarm(&mut self) { - self.mark_as_success(); + self.mode.set_should_not_signal_on_drop(); } pub fn send_we_stopped(&mut self, err: SentError) { diff --git a/gateway/src/node/client_handling/websocket/connection_handler/mod.rs b/gateway/src/node/client_handling/websocket/connection_handler/mod.rs index ecb0aa48d7..aeebdb02fa 100644 --- a/gateway/src/node/client_handling/websocket/connection_handler/mod.rs +++ b/gateway/src/node/client_handling/websocket/connection_handler/mod.rs @@ -95,7 +95,7 @@ pub(crate) async fn handle_connection( St: Storage + Clone + 'static, { // If the connection handler abruptly stops, we shouldn't signal global shutdown - shutdown.mark_as_success(); + shutdown.disarm(); match shutdown .run_future(handle.perform_websocket_handshake()) diff --git a/gateway/src/node/mixnet_handling/receiver/connection_handler.rs b/gateway/src/node/mixnet_handling/receiver/connection_handler.rs index 7d6649530f..4ee18c77d5 100644 --- a/gateway/src/node/mixnet_handling/receiver/connection_handler.rs +++ b/gateway/src/node/mixnet_handling/receiver/connection_handler.rs @@ -203,7 +203,7 @@ impl ConnectionHandler { mut shutdown: TaskClient, ) { debug!("Starting connection handler for {:?}", remote); - shutdown.mark_as_success(); + shutdown.disarm(); let mut framed_conn = Framed::new(conn, NymCodec); while !shutdown.is_shutdown() { tokio::select! { diff --git a/mixnode/src/node/listener/connection_handler/mod.rs b/mixnode/src/node/listener/connection_handler/mod.rs index 08edd7d298..f4144c4f67 100644 --- a/mixnode/src/node/listener/connection_handler/mod.rs +++ b/mixnode/src/node/listener/connection_handler/mod.rs @@ -81,7 +81,7 @@ impl ConnectionHandler { mut shutdown: TaskClient, ) { debug!("Starting connection handler for {:?}", remote); - shutdown.mark_as_success(); + shutdown.disarm(); let mut framed_conn = Framed::new(conn, NymCodec); while !shutdown.is_shutdown() { tokio::select! { diff --git a/service-providers/network-requester/src/core.rs b/service-providers/network-requester/src/core.rs index 788064fcd5..f9a6b3e5aa 100644 --- a/service-providers/network-requester/src/core.rs +++ b/service-providers/network-requester/src/core.rs @@ -495,7 +495,7 @@ impl NRServiceProvider { .send(error_msg) .await .expect("InputMessageReceiver has stopped receiving!"); - shutdown.mark_as_success(); + shutdown.disarm(); return; }