chore(agent-data-plane): share shutdown signal wait logic

bin/agent-data-plane/src/cli/dogstatsd.rs

@@ -35,6 +35,8 @@ use tokio_util::sync::CancellationToken;
 use tracing::debug;
 use tracing::{error, info};
 
+#[cfg(target_os = "linux")]
+use crate::cli::shutdown_signal::wait_for_shutdown_signal;
 use crate::cli::utils::{get_api_client_or_exit, DataPlaneAPIClient};
 
 mod top;
@@ -302,13 +304,8 @@ async fn handle_dogstatsd_replay(
     let cancel_on_signal = tokio::spawn({
         let cancel = cancel.clone();
         async move {
-            let mut sigterm = tokio::signal::unix::signal(tokio::signal::unix::SignalKind::terminate())
-                .expect("failed to install SIGTERM handler");
-
-            tokio::select! {
-                _ = tokio::signal::ctrl_c() => cancel.cancel(),
-                _ = sigterm.recv() => cancel.cancel(),
-            }
+            wait_for_shutdown_signal().await;
+            cancel.cancel();
         }
     });
 

bin/agent-data-plane/src/cli/mod.rs

@@ -18,6 +18,8 @@ mod run;
 pub use self::run::handle_run_command;
 use self::run::RunCommand;
 
+mod shutdown_signal;
+
 pub(super) mod utils;
 
 mod version;

bin/agent-data-plane/src/cli/run.rs

@@ -45,6 +45,7 @@ use saluki_error::{generic_error, ErrorContext as _, GenericError};
 use tracing::{debug, error, info, trace, warn};
 
 use crate::{
+    cli::shutdown_signal::wait_for_shutdown_signal,
     components::{
         apm_onboarding::ApmOnboardingConfiguration,
         dogstatsd_post_aggregate_filter::DogStatsDPostAggregateFilterConfiguration,
@@ -255,46 +256,6 @@ pub async fn handle_run_command(
     }
 }
 
-/// Waits for a shutdown signal.
-///
-/// On Unix, this waits for either `SIGINT` or `SIGTERM`, either of which are used to request a graceful shutdown:
-/// `SIGINT` interactively (`Ctrl+C`), and `SIGTERM` by process supervisors (systemd, container runtimes,
-/// Kubernetes) during rollouts, evictions, node drains, and container shutdown.
-///
-/// On Windows, this waits for either `CTRL_C_EVENT` (interactively) or `CTRL_BREAK_EVENT`, the latter being what
-/// `dd-procmgr` (which manages ADP as a subprocess on Windows) sends via `GenerateConsoleCtrlEvent` to request a
-/// graceful stop.
-async fn wait_for_shutdown_signal() {
-    #[cfg(unix)]
-    {
-        use tokio::signal::unix::{signal, SignalKind};
-
-        let mut sigterm = signal(SignalKind::terminate()).expect("failed to install SIGTERM handler");
-
-        tokio::select! {
-            _ = tokio::signal::ctrl_c() => info!("Received SIGINT, shutting down..."),
-            _ = sigterm.recv() => info!("Received SIGTERM, shutting down..."),
-        }
-    }
-
-    #[cfg(windows)]
-    {
-        let mut ctrl_break = tokio::signal::windows::ctrl_break().expect("failed to install CTRL_BREAK handler");
-
-        tokio::select! {
-            _ = tokio::signal::ctrl_c() => info!("Received CTRL_C, shutting down..."),
-            _ = ctrl_break.recv() => info!("Received CTRL_BREAK, shutting down..."),
-        }
-    }
-
-    #[cfg(not(any(unix, windows)))]
-    {
-        let _ = tokio::signal::ctrl_c().await;
-
-        info!("Received SIGINT, shutting down...");
-    }
-}
-
 /// Check the resolved configuration against the config registry for incompatibilities.
 ///
 /// Classifies each flattened key in `config` with the config registry `Classifier`. Returns an

bin/agent-data-plane/src/cli/shutdown_signal.rs

@@ -0,0 +1,41 @@
+use tracing::info;
+
+/// Waits for a shutdown signal.
+///
+/// On Unix, this waits for either `SIGINT` or `SIGTERM`, either of which are used to request a graceful shutdown:
+/// `SIGINT` interactively (`Ctrl+C`), and `SIGTERM` by process supervisors (systemd, container runtimes,
+/// Kubernetes) during rollouts, evictions, node drains, and container shutdown.
+///
+/// On Windows, this waits for either `CTRL_C_EVENT` (interactively) or `CTRL_BREAK_EVENT`, the latter being what
+/// `dd-procmgr` (which manages ADP as a subprocess on Windows) sends via `GenerateConsoleCtrlEvent` to request a
+/// graceful stop.
+pub async fn wait_for_shutdown_signal() {
+    #[cfg(unix)]
+    {
+        use tokio::signal::unix::{signal, SignalKind};
+
+        let mut sigterm = signal(SignalKind::terminate()).expect("failed to install SIGTERM handler");
+
+        tokio::select! {
+            _ = tokio::signal::ctrl_c() => info!("Received SIGINT, shutting down..."),
+            _ = sigterm.recv() => info!("Received SIGTERM, shutting down..."),
+        }
+    }
+
+    #[cfg(windows)]
+    {
+        let mut ctrl_break = tokio::signal::windows::ctrl_break().expect("failed to install CTRL_BREAK handler");
+
+        tokio::select! {
+            _ = tokio::signal::ctrl_c() => info!("Received CTRL_C, shutting down..."),
+            _ = ctrl_break.recv() => info!("Received CTRL_BREAK, shutting down..."),
+        }
+    }
+
+    #[cfg(not(any(unix, windows)))]
+    {
+        let _ = tokio::signal::ctrl_c().await;
+
+        info!("Received SIGINT, shutting down...");
+    }
+}