wrap token in a handle and expand service error handling

Author: redsh4de

binaries/cuprated/src/blockchain/manager.rs

@@ -72,7 +72,7 @@ pub async fn init_blockchain_manager(
         Arc::clone(&stop_current_block_downloader),
         block_downloader_config,
         Arc::clone(&synced_notify),
-        task.cancellation_token.clone(),
+        task.shutdown_handle.token(),
     );
     task.spawn_critical(
         async move {
@@ -95,7 +95,7 @@ pub async fn init_blockchain_manager(
         broadcast_svc: clearnet_interface.broadcast_svc(),
     };
 
-    let shutdown_token = task.cancellation_token.clone();
+    let shutdown_token = task.shutdown_handle.token();
     task.spawn_critical(manager.run(batch_rx, command_rx, shutdown_token), || {
         tracing::info!("Blockchain manager shut down.");
     });

binaries/cuprated/src/commands.rs

@@ -5,7 +5,6 @@ use std::{io, thread::sleep, time::Duration};
 
 use clap::{builder::TypedValueParser, Parser, ValueEnum};
 use tokio::sync::mpsc;
-use tokio_util::sync::CancellationToken;
 use tower::{Service, ServiceExt};
 use tracing::level_filters::LevelFilter;
 
@@ -17,6 +16,7 @@ use cuprate_helper::time::secs_to_hms;
 use crate::{
     logging::{self, CupratedTracingFilter},
     statics,
+    supervisor::ShutdownHandle,
 };
 
 /// A command received from [`io::stdin`].
@@ -102,12 +102,12 @@ pub fn command_listener(incoming_commands: mpsc::Sender<Command>) {
 pub async fn io_loop(
     mut incoming_commands: mpsc::Receiver<Command>,
     mut context_service: BlockchainContextService,
-    shutdown_token: CancellationToken,
+    shutdown_handle: ShutdownHandle,
 ) {
     loop {
         let command = tokio::select! {
             biased;
-            () = shutdown_token.cancelled() => break,
+            () = shutdown_handle.cancelled() => break,
             cmd = incoming_commands.recv() => {
                 let Some(cmd) = cmd else { break };
                 cmd
@@ -157,7 +157,7 @@ pub async fn io_loop(
                 }
             }
             Command::Exit => {
-                crate::supervisor::trigger_shutdown(&shutdown_token);
+                shutdown_handle.trigger_shutdown();
                 break;
             }
         }

binaries/cuprated/src/main.rs

@@ -98,8 +98,8 @@ fn main() {
 
     rt.block_on(async move {
         // Create the supervisor and task handle, spawn signal handlers.
-        let (mut supervisor, task) = supervisor::new();
-        supervisor::spawn_signal_handler(supervisor.cancellation_token.clone());
+        let (supervisor, task) = supervisor::new();
+        supervisor::spawn_signal_handler(supervisor.shutdown_handle.clone());
 
         // TODO: Add an argument/option for keeping alt blocks between restart.
         blockchain_write_handle
@@ -135,7 +135,7 @@ fn main() {
             blockchain_read_handle.clone(),
             txpool_read_handle.clone(),
             &tor_context,
-            task.cancellation_token.clone(),
+            task.shutdown_handle.clone(),
         )
         .await;
 
@@ -190,13 +190,13 @@ fn main() {
         if tor_enabled {
             info!("Tor P2P zone will start after sync.");
             let context_svc = context_svc.clone();
-            let cancellation_token = task.cancellation_token.clone();
+            let shutdown_handle = task.shutdown_handle.clone();
 
             task.task_tracker.spawn(async move {
                 // Wait for the node to synchronize with the network, or shutdown.
                 tokio::select! {
                     () = synced_notify.notified() => {}
-                    () = cancellation_token.cancelled() => {
+                    () = shutdown_handle.cancelled() => {
                         return;
                     }
                 }
@@ -208,7 +208,7 @@ fn main() {
                     blockchain_read_handle,
                     txpool_read_handle,
                     tor_context,
-                    cancellation_token.clone(),
+                    shutdown_handle.clone(),
                 )
                 .await;
 
@@ -236,13 +236,13 @@ fn main() {
             task.task_tracker.spawn(commands::io_loop(
                 command_rx,
                 context_svc,
-                task.cancellation_token.clone(),
+                task.shutdown_handle.clone(),
             ));
         } else {
             info!("Terminal/TTY not detected, disabling STDIN commands");
         }
         // Wait for shutdown (signal, command, or critical task failure).
-        supervisor.cancellation_token.cancelled().await;
+        supervisor.shutdown_handle.cancelled().await;
         supervisor.task_tracker.close();
         supervisor.task_tracker.wait().await;
     });

binaries/cuprated/src/p2p.rs

@@ -8,7 +8,6 @@ use anyhow::anyhow;
 use futures::{FutureExt, TryFutureExt};
 use serde::{Deserialize, Serialize};
 use tokio::sync::oneshot::{self, Sender};
-use tokio_util::sync::CancellationToken;
 use tower::{Service, ServiceExt};
 
 use cuprate_blockchain::service::{BlockchainReadHandle, BlockchainWriteHandle};
@@ -24,7 +23,7 @@ use cuprate_types::blockchain::BlockchainWriteRequest;
 use crate::{
     blockchain,
     config::Config,
-    constants::PANIC_CRITICAL_SERVICE_ERROR,
+    supervisor::ShutdownHandle,
     tor::{transport_clearnet_daemon_config, transport_daemon_config, TorContext, TorMode},
     txpool::{self, IncomingTxHandler},
 };
@@ -100,7 +99,7 @@ pub async fn initialize_clearnet_p2p(
     blockchain_read_handle: BlockchainReadHandle,
     txpool_read_handle: TxpoolReadHandle,
     tor_ctx: &TorContext,
-    shutdown_token: CancellationToken,
+    shutdown_handle: ShutdownHandle,
 ) -> (NetworkInterface<ClearNet>, Sender<IncomingTxHandler>) {
     match config.p2p.clear_net.proxy {
         ProxySettings::Tor => match tor_ctx.mode {
@@ -113,7 +112,7 @@ pub async fn initialize_clearnet_p2p(
                     txpool_read_handle,
                     config.clearnet_p2p_config(),
                     transport_clearnet_arti_config(tor_ctx),
-                    shutdown_token,
+                    shutdown_handle,
                 )
                 .await
                 .unwrap()
@@ -124,7 +123,7 @@ pub async fn initialize_clearnet_p2p(
                 txpool_read_handle,
                 config.clearnet_p2p_config(),
                 transport_clearnet_daemon_config(config),
-                shutdown_token,
+                shutdown_handle,
             )
             .await
             .unwrap(),
@@ -138,7 +137,7 @@ pub async fn initialize_clearnet_p2p(
                     txpool_read_handle,
                     config.clearnet_p2p_config(),
                     config.p2p.clear_net.tcp_transport_config(config.network),
-                    shutdown_token.clone(),
+                    shutdown_handle.clone(),
                 )
                 .await
                 .unwrap()
@@ -152,7 +151,7 @@ pub async fn initialize_clearnet_p2p(
                         client_config: socks_proxy_str_to_config(s).unwrap(),
                         server_config: None,
                     },
-                    shutdown_token.clone(),
+                    shutdown_handle.clone(),
                 )
                 .await
                 .unwrap()
@@ -169,7 +168,7 @@ pub async fn start_tor_p2p(
     blockchain_read_handle: BlockchainReadHandle,
     txpool_read_handle: TxpoolReadHandle,
     tor_ctx: TorContext,
-    shutdown_token: CancellationToken,
+    shutdown_handle: ShutdownHandle,
 ) -> (NetworkInterface<Tor>, Sender<IncomingTxHandler>) {
     match tor_ctx.mode {
         TorMode::Daemon => start_zone_p2p::<Tor, Daemon>(
@@ -178,7 +177,7 @@ pub async fn start_tor_p2p(
             txpool_read_handle,
             config.tor_p2p_config(&tor_ctx),
             transport_daemon_config(config),
-            shutdown_token,
+            shutdown_handle,
         )
         .await
         .unwrap(),
@@ -189,7 +188,7 @@ pub async fn start_tor_p2p(
             txpool_read_handle,
             config.tor_p2p_config(&tor_ctx),
             transport_arti_config(config, tor_ctx),
-            shutdown_token,
+            shutdown_handle,
         )
         .await
         .unwrap(),
@@ -207,7 +206,7 @@ pub async fn start_zone_p2p<N, T>(
     txpool_read_handle: TxpoolReadHandle,
     config: P2PConfig<N>,
     transport_config: TransportConfig<N, T>,
-    shutdown_token: CancellationToken,
+    shutdown_handle: ShutdownHandle,
 ) -> Result<(NetworkInterface<N>, Sender<IncomingTxHandler>), tower::BoxError>
 where
     N: NetworkZone,
@@ -221,7 +220,7 @@ where
         blockchain_read_handle,
         blockchain_context_service: blockchain_context_service.clone(),
         txpool_read_handle,
-        shutdown_token,
+        shutdown_handle,
         incoming_tx_handler: None,
         incoming_tx_handler_fut: incoming_tx_handler_rx.shared(),
     };