Migrate from Tokio Graceful Shutdown to Kameo actor framework

Cargo.toml

@@ -58,7 +58,7 @@ serde_json5 = "0.2.1"
 serde_with = { version = "3.12", features = ["chrono"] }
 thiserror.workspace = true
 tokio = { workspace = true, features = ["full"] }
-tokio-graceful-shutdown = "0.16.0"
+kameo = "0.19"
 tokio-stream = { workspace = true, features = ["time"] }
 toml = { version = "0.8.8", features = ["preserve_order"] }
 multi-core = { path = "crates/multi-core" }

src/cmd/run/canary_mode.rs

@@ -1,14 +1,12 @@
 use async_trait::async_trait;
 use miette::Result;
-use tokio::time::Duration;
-use tokio_graceful_shutdown::{IntoSubsystem as _, SubsystemBuilder, Toplevel};
+use tokio::signal;
 use tracing::{debug, info};
 
-use crate::ControllerSubsystem;
 use crate::adapters::{BackendClient, BoxedIngress, BoxedMonitor, BoxedPlatform, RolloutMetadata};
-use crate::subsystems::CONTROLLER_SUBSYSTEM_NAME;
+use crate::ControllerSubsystem;
 
-use super::{DEFAULT_SHUTDOWN_TIMEOUT, DeploymentMode};
+use super::DeploymentMode;
 
 /// Canary deployment mode - runs the full canary analysis subsystems
 pub struct CanaryMode {
@@ -52,17 +50,27 @@ impl DeploymentMode for CanaryMode {
 
         info!("Starting the rollout...");
 
-        // Let's capture the shutdown signal from the OS.
-        Toplevel::new(|s| async move {
-            // • Start the action listener subsystem.
-            s.start(SubsystemBuilder::new(
-                CONTROLLER_SUBSYSTEM_NAME,
-                controller.into_subsystem(),
-            ));
-        })
-        .catch_signals()
-        .handle_shutdown_requests(Duration::from_millis(DEFAULT_SHUTDOWN_TIMEOUT))
-        .await
-        .map_err(Into::into)
+        // Spawn the controller actor
+        let controller_ref = controller.spawn();
+
+        // Wait for either the controller to finish or a shutdown signal
+        tokio::select! {
+            // Wait for the controller to stop (either normally or due to error)
+            _ = controller_ref.wait_for_shutdown() => {
+                debug!("Controller stopped");
+            }
+            // Handle Ctrl+C signal
+            _ = signal::ctrl_c() => {
+                info!("Received shutdown signal, stopping...");
+                if let Err(e) = controller_ref.stop_gracefully().await {
+                    debug!("Error during graceful shutdown: {:?}", e);
+                }
+                // Wait for the actor to finish shutting down
+                controller_ref.wait_for_shutdown().await;
+                debug!("Controller stopped after shutdown signal");
+            }
+        }
+
+        Ok(())
     }
 }

src/cmd/run/mod.rs

@@ -22,6 +22,7 @@ pub use force_mode::ForceMode;
 
 /// The amount of time, in miliseconds, each subsystem has
 /// to gracefully shutdown before being forcably shutdown.
+#[allow(dead_code)]
 pub(super) const DEFAULT_SHUTDOWN_TIMEOUT: u64 = 5000;
 
 /// Trait defining different deployment modes for the MultiTool CLI

src/stats/categorical.rs

@@ -9,9 +9,9 @@
 /// You can think of a [Categorical] as a hashmap with fixed integer keys. When the map is
 /// created, its keys must already be known and completely cover the range `[0, N)`.
 ///
-/// ```rust
+/// ```rust,ignore
 /// use std::collections::HashSet;
-/// use canary::stats::Categorical;
+/// use multitool::stats::Categorical;
 ///
 /// #[derive(PartialEq, Eq, Debug, Hash)]
 /// enum Coin {

src/stats/contingency.rs

@@ -97,6 +97,7 @@ impl<const N: usize, C: Categorical<N>> Default for ContingencyTable<N, C> {
 }
 
 #[cfg(test)]
+#[allow(unused_imports)]
 pub(crate) use tests::Coin;
 
 #[cfg(test)]