chore(app): handle bootstrap background tasks via dedicated supervisor (#1538)

## Summary

This PR updates the bootstrap process/helpers to use supervision trees to handle all necessary background async tasks that arise out of the initialization of various subsystems.

The bootstrap phase involves initializes various fundamental subsystems -- logging, metrics, TLS, etc -- prior to the application doing useful work. In many cases, these subsystems require background tasks to run in order to perform routine upkeep, or manage requests to alter behavior, and so on. Currently, these background tasks are all spawned directly on the ambient runtime where they cannot be restarted if they failed, and generally can't participate in the advanced mechanisms available to supervision trees, such as communicating via dataspaces.

This PR updates the bootstrap process, and associated helper types, to build a single supervisor that contains all necessary workers (tasks) that must be driven in order to support the subsystems being initialized. The goal is two-fold:

- move more async tasks into supervision trees (for reasons of reliability and observability)
- gain access to advanced runtime features like dataspaces

The work here is meant as a stepping stone towards supporting dynamic log level overrides that come from changes to `log_level` in ADP's primary configuration. By switching the relevant background task that handles log level overrides currently to run in a supervision tree, we're opening up the ability to change how interested code operates with it and communicates with it, allowing us to decouple ADP-specific needs from generic Saluki code.

We've worker-ified four distinct background tasks:

- `MetricsFlusherWorker`: bridge between our `metrics` recorder and emitting internal metrics to the `internal_metrics` source
- `RuntimeMetricsWorker`: collects Tokio runtime metrics and re-emits them into our internal metrics
- `MetricsOverrideWorker`: handles API requests to dynamically override/reset the current metrics level
- `LoggingOverrideWorker`: handles API requests to dynamically override/reset the current log level

## Change Type

- [ ] Bug fix
- [ ] New feature
- [x] Non-functional (chore, refactoring, docs)
- [ ] Performance

## How did you test this PR?

Existing unit and integration tests.

## References

DADP-52


Co-authored-by: toby.lawrence <toby.lawrence@datadoghq.com>

Author: tobz

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

@@ -31,7 +31,7 @@ use saluki_components::{
 };
 use saluki_config::{ConfigurationLoader, GenericConfiguration};
 use saluki_core::health::HealthRegistry;
-use saluki_core::runtime::SupervisorError;
+use saluki_core::runtime::{Supervisor, SupervisorError};
 use saluki_core::topology::TopologyBlueprint;
 use saluki_env::EnvironmentProvider as _;
 use saluki_error::{generic_error, ErrorContext as _, GenericError};
@@ -62,6 +62,7 @@ pub struct RunCommand {
 /// Entrypoint for the `run` commands.
 pub async fn handle_run_command(
     started: Instant, bootstrap_config: GenericConfiguration, bootstrap_guard: &mut BootstrapGuard,
+    bootstrap_supervisor: Supervisor,
 ) -> Result<(), GenericError> {
     let app_details = saluki_metadata::get_app_details();
     info!(
@@ -179,6 +180,11 @@ pub async fn handle_run_command(
     .await
     .error_context("Failed to create internal supervisor.")?;
 
+    // Attach the bootstrap supervisor as a child so its workers (logging/metrics override processors,
+    // metrics flusher, runtime metrics collector) are driven and shut down alongside the rest of the
+    // internal supervision tree.
+    internal_supervisor.add_worker(bootstrap_supervisor);
+
     // Create shutdown channel for the internal supervisor - we'll drive it in the main select loop
     let (internal_shutdown_tx, internal_shutdown_rx) = tokio::sync::oneshot::channel();
     let internal_supervisor_fut = internal_supervisor.run_with_shutdown(internal_shutdown_rx).fuse();

bin/agent-data-plane/src/main.rs

@@ -7,9 +7,10 @@
 #![deny(missing_docs)]
 use std::time::Instant;
 
-use saluki_app::bootstrap::{AppBootstrapper, BootstrapGuard};
+use saluki_app::bootstrap::{AppBootstrapper, Bootstrap, BootstrapGuard};
 use saluki_components::config::{DatadogRemapper, KEY_ALIASES};
 use saluki_config::{ConfigurationLoader, GenericConfiguration};
+use saluki_core::runtime::Supervisor;
 use saluki_error::{ErrorContext as _, GenericError};
 use tracing::{error, info, warn};
 
@@ -64,13 +65,25 @@ async fn main() -> Result<(), GenericError> {
         .error_context("Failed to parse bootstrap configuration during bootstrap phase.")?
         .with_metrics_prefix("adp")
         .with_logging_configuration(bootstrap_logging_config);
-    let mut bootstrap_guard = bootstrapper
+    let Bootstrap {
+        supervisor: bootstrap_supervisor,
+        guard: mut bootstrap_guard,
+    } = bootstrapper
         .bootstrap()
         .await
         .error_context("Failed to complete bootstrap phase.")?;
 
-    // Run the given subcommand.
-    let maybe_exit_code = run_inner(cli.action, started, bootstrap_config, &mut bootstrap_guard).await?;
+    // Run the given subcommand. The bootstrap supervisor is forwarded by value; only the long-lived `run`
+    // subcommand actually drives it (it is added as a child of the internal supervisor inside
+    // `handle_run_command`). All other subcommands drop it on entry.
+    let maybe_exit_code = run_inner(
+        cli.action,
+        started,
+        bootstrap_config,
+        &mut bootstrap_guard,
+        bootstrap_supervisor,
+    )
+    .await?;
 
     // Drop the bootstrap guard to ensure logs are flushed, etc.
     drop(bootstrap_guard);
@@ -85,6 +98,7 @@ async fn main() -> Result<(), GenericError> {
 
 async fn run_inner(
     action: Action, started: Instant, bootstrap_config: GenericConfiguration, bootstrap_guard: &mut BootstrapGuard,
+    bootstrap_supervisor: Supervisor,
 ) -> Result<Option<i32>, GenericError> {
     match action {
         Action::Run(cmd) => {
@@ -97,16 +111,17 @@ async fn run_inner(
                 }
             }
 
-            let exit_code = match handle_run_command(started, bootstrap_config, bootstrap_guard).await {
-                Ok(()) => {
-                    info!("Agent Data Plane stopped.");
-                    None
-                }
-                Err(e) => {
-                    error!("{:?}", e);
-                    Some(1)
-                }
-            };
+            let exit_code =
+                match handle_run_command(started, bootstrap_config, bootstrap_guard, bootstrap_supervisor).await {
+                    Ok(()) => {
+                        info!("Agent Data Plane stopped.");
+                        None
+                    }
+                    Err(e) => {
+                        error!("{:?}", e);
+                        Some(1)
+                    }
+                };
 
             // Remove the PID file, if configured.
             if let Some(pid_file) = &cmd.pid_file {

lib/saluki-app/src/bootstrap.rs

@@ -1,6 +1,7 @@
 //! Bootstrap utilities.
 
 use saluki_config::GenericConfiguration;
+use saluki_core::runtime::Supervisor;
 use saluki_error::{ErrorContext as _, GenericError};
 
 use crate::{
@@ -9,6 +10,20 @@ use crate::{
     tls::initialize_tls,
 };
 
+/// The result of running [`AppBootstrapper::bootstrap`].
+///
+/// Bundles together the [`BootstrapGuard`] that must be held for the lifetime of the application with the
+/// [`Supervisor`] that drives the background workers spawned during bootstrap. Callers must arrange for the
+/// supervisor to be run (either directly or by adding it to a parent supervisor) for those workers to make
+/// progress.
+pub struct Bootstrap {
+    /// Supervisor populated with workers for all background async tasks created during bootstrap.
+    pub supervisor: Supervisor,
+
+    /// Drop guard for resources acquired during bootstrap.
+    pub guard: BootstrapGuard,
+}
+
 /// A drop guard for ensuring deferred cleanup of resources acquired during bootstrap.
 pub struct BootstrapGuard {
     logging_guard: LoggingGuard,
@@ -79,26 +94,40 @@ impl AppBootstrapper {
 
     /// Executes the bootstrap operation, initializing all configured subsystems.
     ///
-    /// A [`BootstrapGuard`] is returned, which must be held until the application is ready to shut down. This guard
-    /// ensures that all relevant resources created during the bootstrap phase are properly cleaned up/flushed before
-    /// the application exits.
+    /// Returns a [`Bootstrap`] containing both a [`BootstrapGuard`] (which must be held until the application is
+    /// ready to shut down) and a [`Supervisor`] populated with workers for all background async tasks created
+    /// during bootstrap. Callers must arrange for the supervisor to run (typically by adding it to a parent
+    /// supervisor or calling [`Supervisor::run_with_shutdown`]) for those workers to make progress.
     ///
     /// # Errors
     ///
     /// If any of the bootstrap steps fail, an error will be returned.
-    pub async fn bootstrap(self) -> Result<BootstrapGuard, GenericError> {
+    pub async fn bootstrap(self) -> Result<Bootstrap, GenericError> {
         // Initialize the logging subsystem first, since we want to make it possible to get any logs from the rest of
         // the bootstrap process.
-        let logging_guard = initialize_logging(self.logging_config)
+        let (logging_guard, logging_override) = initialize_logging(self.logging_config)
             .await
             .error_context("Failed to initialize logging subsystem.")?;
 
         // Initialize everything else.
         initialize_tls().error_context("Failed to initialize TLS subsystem.")?;
-        initialize_metrics(self.metrics_config)
+        let metrics_workers = initialize_metrics(self.metrics_config)
             .await
             .error_context("Failed to initialize metrics subsystem.")?;
 
-        Ok(BootstrapGuard { logging_guard })
+        // Build the supervisor for all bootstrap-spawned background workers. The default ambient runtime mode is
+        // appropriate here: these are lightweight tasks that share the parent runtime, and the runtime metrics
+        // worker has already eagerly captured the parent's `Handle` so it always describes the right runtime.
+        let mut supervisor =
+            Supervisor::new("app-bootstrap").error_context("Failed to construct app bootstrap supervisor.")?;
+        supervisor.add_worker(logging_override);
+        supervisor.add_worker(metrics_workers.flusher);
+        supervisor.add_worker(metrics_workers.runtime);
+        supervisor.add_worker(metrics_workers.override_processor);
+
+        Ok(Bootstrap {
+            supervisor,
+            guard: BootstrapGuard { logging_guard },
+        })
     }
 }

lib/saluki-app/src/logging/api.rs

@@ -3,13 +3,15 @@ use std::{
     time::Duration,
 };
 
+use async_trait::async_trait;
 use saluki_api::{
     extract::{Query, State},
     response::IntoResponse,
     routing::{post, Router},
     APIHandler, StatusCode,
 };
-use saluki_common::task::spawn_traced_named;
+use saluki_core::runtime::{InitializationError, ProcessShutdown, Supervisable, SupervisorFuture};
+use saluki_error::generic_error;
 use serde::Deserialize;
 use tokio::{select, sync::mpsc, time::sleep};
 use tracing::{error, info};
@@ -59,17 +61,28 @@ pub struct LoggingAPIHandler {
 }
 
 impl LoggingAPIHandler {
-    pub(super) fn new(base_filter: Arc<Mutex<EnvFilter>>, reload_handle: Handle<EnvFilter, Registry>) -> Self {
-        // Spawn our background task that will handle override requests.
+    /// Creates a new `LoggingAPIHandler` and a paired [`LoggingOverrideWorker`].
+    ///
+    /// The worker must be added to a [`Supervisor`][saluki_core::runtime::Supervisor] for the handler's
+    /// override/reset routes to take effect; without it, requests are accepted but never applied.
+    pub(super) fn new(
+        base_filter: Arc<Mutex<EnvFilter>>, reload_handle: Handle<EnvFilter, Registry>,
+    ) -> (Self, LoggingOverrideWorker) {
         let (override_tx, override_rx) = mpsc::channel(1);
-        spawn_traced_named(
-            "dynamic-logging-override-processor",
-            process_override_requests(base_filter, reload_handle, override_rx),
-        );
+        let worker = LoggingOverrideWorker {
+            state: Mutex::new(Some(LoggingOverrideWorkerState {
+                base_filter,
+                reload_handle,
+                override_rx,
+            })),
+        };
 
-        Self {
-            state: LoggingHandlerState { override_tx },
-        }
+        (
+            Self {
+                state: LoggingHandlerState { override_tx },
+            },
+            worker,
+        )
     }
 
     async fn override_handler(
@@ -125,17 +138,66 @@ impl APIHandler for LoggingAPIHandler {
     }
 }
 
+/// A worker that processes dynamic log filter override requests sent via [`LoggingAPIHandler`].
+///
+/// Holds the receiving half of the override channel; the corresponding sender is held by the API handler
+/// (which is itself stored in a static after the logging subsystem is initialized). The worker exits cleanly
+/// on either supervisor shutdown or channel close.
+///
+/// One-shot: a successful initialization consumes the receiver. Restart by the supervisor will fail with an
+/// initialization error, propagating up to bring the supervisor down.
+pub struct LoggingOverrideWorker {
+    state: Mutex<Option<LoggingOverrideWorkerState>>,
+}
+
+struct LoggingOverrideWorkerState {
+    base_filter: Arc<Mutex<EnvFilter>>,
+    reload_handle: Handle<EnvFilter, Registry>,
+    override_rx: mpsc::Receiver<Option<(Duration, EnvFilter)>>,
+}
+
+#[async_trait]
+impl Supervisable for LoggingOverrideWorker {
+    fn name(&self) -> &str {
+        "dynamic-logging-override-processor"
+    }
+
+    async fn initialize(&self, process_shutdown: ProcessShutdown) -> Result<SupervisorFuture, InitializationError> {
+        let LoggingOverrideWorkerState {
+            base_filter,
+            reload_handle,
+            override_rx,
+        } = self
+            .state
+            .lock()
+            .unwrap()
+            .take()
+            .ok_or_else(|| InitializationError::Failed {
+                source: generic_error!("logging override worker has already been initialized"),
+            })?;
+
+        Ok(Box::pin(async move {
+            process_override_requests(base_filter, reload_handle, override_rx, process_shutdown).await;
+            Ok(())
+        }))
+    }
+}
+
 async fn process_override_requests(
     base_filter: Arc<Mutex<EnvFilter>>, reload_handle: Handle<EnvFilter, Registry>,
-    mut rx: mpsc::Receiver<Option<(Duration, EnvFilter)>>,
+    mut rx: mpsc::Receiver<Option<(Duration, EnvFilter)>>, mut process_shutdown: ProcessShutdown,
 ) {
     let mut override_active = false;
     let override_timeout = sleep(Duration::from_secs(3600));
 
     tokio::pin!(override_timeout);
 
+    let shutdown = process_shutdown.wait_for_shutdown();
+    tokio::pin!(shutdown);
+
     loop {
         select! {
+            _ = &mut shutdown => break,
             maybe_override = rx.recv() => match maybe_override {
                 Some(Some((duration, new_filter))) => {
                     info!(directives = %new_filter, "Overriding existing log filtering directives for {} seconds...", duration.as_secs());
@@ -232,6 +294,7 @@ mod tests {
             base_filter.clone(),
             reload_handle.clone(),
             rx,
+            ProcessShutdown::noop(),
         ));
 
         tx.send(Some((Duration::from_secs(60), EnvFilter::new("hyper=warn"))))
@@ -258,6 +321,7 @@ mod tests {
             base_filter.clone(),
             reload_handle.clone(),
             rx,
+            ProcessShutdown::noop(),
         ));
 
         tx.send(Some((Duration::from_millis(100), EnvFilter::new("hyper=warn"))))

lib/saluki-app/src/logging/mod.rs

@@ -20,7 +20,7 @@ use tracing_subscriber::{layer::SubscriberExt as _, reload, util::SubscriberInit
 
 mod api;
 use self::api::set_logging_api_handler;
-pub use self::api::{acquire_logging_api_handler, LoggingAPIHandler};
+pub use self::api::{acquire_logging_api_handler, LoggingAPIHandler, LoggingOverrideWorker};
 
 mod config;
 pub use self::config::{LogLevel, LoggingConfiguration};
@@ -76,13 +76,16 @@ impl LoggingGuard {
 /// An API handler can be acquired (via [`acquire_logging_api_handler`]) to install the API routes which allow for
 /// dynamically controlling the logging level filtering. See [`LoggingAPIHandler`] for more information.
 ///
-/// Returns a [`LoggingGuard`] which must be held until the application is about to shutdown, ensuring that any
-/// configured logging backends are able to completely flush any pending logs before the application exits.
+/// Returns a [`LoggingGuard`] which must be held until the application is about to shutdown, plus a
+/// [`LoggingOverrideWorker`] that must be added to a [`Supervisor`][saluki_core::runtime::Supervisor] to drive
+/// the dynamic override processor; without the worker running, override requests are accepted but never applied.
 ///
 /// # Errors
 ///
 /// If the logging subsystem was already initialized, an error will be returned.
-pub(crate) async fn initialize_logging(config: LoggingConfiguration) -> Result<LoggingGuard, GenericError> {
+pub(crate) async fn initialize_logging(
+    config: LoggingConfiguration,
+) -> Result<(LoggingGuard, LoggingOverrideWorker), GenericError> {
     // TODO: Support for logging to syslog.
 
     // Build the initial output stack from the supplied configuration. This is later swappable via
@@ -97,18 +100,22 @@ pub(crate) async fn initialize_logging(config: LoggingConfiguration) -> Result<L
     // The base filter is the level the override-restore should land on after a `/logging/override` timeout. It starts
     // as the bootstrap level, and is updated by `LoggingGuard::reload` once the Agent's configuration is applied.
     let base_filter = Arc::new(Mutex::new(level_filter));
-    set_logging_api_handler(LoggingAPIHandler::new(base_filter.clone(), filter_handle.clone()));
+    let (api_handler, override_worker) = LoggingAPIHandler::new(base_filter.clone(), filter_handle.clone());
+    set_logging_api_handler(api_handler);
 
     tracing_subscriber::registry()
         .with(output_layer.with_filter(filter_layer))
         .try_init()?;
 
-    Ok(LoggingGuard {
-        worker_guards,
-        stack_handle,
-        filter_handle,
-        base_filter,
-    })
+    Ok((
+        LoggingGuard {
+            worker_guards,
+            stack_handle,
+            filter_handle,
+            base_filter,
+        },
+        override_worker,
+    ))
 }
 
 fn build_output_stack(config: &LoggingConfiguration) -> Result<(OutputStack, Vec<WorkerGuard>), GenericError> {