saluki_app/logging/

api.rs

use std::{pin::Pin, sync::Arc, time::Duration};

use async_trait::async_trait;
use futures::{Stream, StreamExt as _};
use saluki_api::{
    extract::{Query, State},
    response::IntoResponse,
    routing::{post, Router},
    APIHandler, DynamicRoute, EndpointType, StatusCode,
};
use saluki_common::{logging::parse_filter_directives, sync::shutdown::ShutdownHandle};
use saluki_core::runtime::{state::DataspaceRegistry, InitializationError, Supervisable, SupervisorFuture};
use saluki_error::{generic_error, GenericError};
use serde::Deserialize;
use tokio::{
    pin, select,
    sync::{mpsc, Mutex},
    time::sleep,
};
use tracing::{debug, error, info, warn};
use tracing_subscriber::{filter::Targets, reload::Handle, Registry};

use super::LogLevel;

#[derive(Deserialize)]
struct OverrideQueryParams {
    time_secs: u64,
}

/// An action driven through a [`LoggingOverrideController`] and processed by [`LoggingOverrideWorker`].
enum LoggingOverrideAction {
    /// Apply `filter` as a temporary override on top of the current base, lasting `duration`.
    ///
    /// When the duration elapses, or a [`Reset`][LoggingOverrideAction::Reset] arrives, the worker restores the
    /// current base filter.
    Override { duration: Duration, filter: Targets },

    /// Clear any active override, immediately restoring the current base filter.
    Reset,

    /// Replace the current base filter.
    ///
    /// Applied immediately if no override is active. If an override is active, the new base is stored and applied
    /// once the override expires (or is reset), so we don't clobber an in-flight override.
    UpdateBase(Targets),
}

/// Controls the dynamic logging filter at runtime.
///
/// All filter mutations—temporary overrides, resets, and base updates—flow through here to the
/// [`LoggingOverrideWorker`] that owns the underlying `tracing` reload handle. Cloning is cheap; hand clones to
/// any caller that needs to drive filter changes (HTTP handlers, configuration watchers, etc.).
#[derive(Clone)]
pub struct LoggingOverrideController {
    tx: mpsc::Sender<LoggingOverrideAction>,
}

impl LoggingOverrideController {
    /// Applies `filter` as a temporary override for `duration`, after which the base filter is restored.
    pub(crate) async fn override_for(&self, duration: Duration, filter: Targets) -> Result<(), GenericError> {
        self.send(LoggingOverrideAction::Override { duration, filter }).await
    }

    /// Clears any active override, immediately restoring the current base filter.
    pub(crate) async fn reset(&self) -> Result<(), GenericError> {
        self.send(LoggingOverrideAction::Reset).await
    }

    /// Replaces the current base filter.
    ///
    /// Applied immediately if no override is active; otherwise stored and applied when the override expires.
    pub async fn update_base(&self, filter: Targets) -> Result<(), GenericError> {
        self.send(LoggingOverrideAction::UpdateBase(filter)).await
    }

    async fn send(&self, action: LoggingOverrideAction) -> Result<(), GenericError> {
        self.tx
            .send(action)
            .await
            .map_err(|_| generic_error!("logging override worker is no longer running"))
    }
}

/// State used for the logging API handler.
#[derive(Clone)]
pub struct LoggingHandlerState {
    controller: LoggingOverrideController,
}

/// An API handler for updating log filtering directives at runtime.
///
/// This handler exposes two main routes -- `/logging/override` and `/logging/reset` -- which allow for overriding the
/// default log filtering directives (configured at startup) at runtime, and then resetting them once the override is no
/// longer needed.
///
/// As this has the potential for incredibly verbose logging at runtime, the override is set with a specific duration in
/// which it will apply. Once an override has been active for the configured duration, it will automatically be reset
/// unless the override is refreshed before the duration elapses.
///
/// The maximum duration for an override is 10 minutes.
pub struct LoggingAPIHandler {
    state: LoggingHandlerState,
}

impl LoggingAPIHandler {
    /// Creates a new `LoggingAPIHandler` driven by the given controller.
    fn new(controller: LoggingOverrideController) -> Self {
        Self {
            state: LoggingHandlerState { controller },
        }
    }

    async fn override_handler(
        State(state): State<LoggingHandlerState>, params: Query<OverrideQueryParams>, body: String,
    ) -> impl IntoResponse {
        // Make sure the override length is within the acceptable range.
        const MAXIMUM_OVERRIDE_LENGTH_SECS: u64 = 600;
        if params.time_secs > MAXIMUM_OVERRIDE_LENGTH_SECS {
            return (
                StatusCode::BAD_REQUEST,
                format!(
                    "override time cannot be greater than {} seconds",
                    MAXIMUM_OVERRIDE_LENGTH_SECS
                ),
            );
        }

        // Parse the override duration and create a new filter from the body.
        let duration = Duration::from_secs(params.time_secs);
        let new_filter = match parse_filter_directives(&body) {
            Ok(filter) => filter,
            Err(e) => {
                return (
                    StatusCode::BAD_REQUEST,
                    format!("failed to parse override filter: {}", e),
                )
            }
        };

        // Instruct the override worker to apply the new log filtering directives for the given duration.
        let _ = state.controller.override_for(duration, new_filter).await;

        (StatusCode::OK, "acknowledged".to_string())
    }

    async fn reset_handler(State(state): State<LoggingHandlerState>) {
        // Instruct the override worker to immediately reset back to the current base log filtering directives.
        let _ = state.controller.reset().await;
    }
}

impl APIHandler for LoggingAPIHandler {
    type State = LoggingHandlerState;

    fn generate_initial_state(&self) -> Self::State {
        self.state.clone()
    }

    fn generate_routes(&self) -> Router<Self::State> {
        Router::new()
            .route("/logging/override", post(Self::override_handler))
            .route("/logging/reset", post(Self::reset_handler))
    }
}

/// A worker that processes log filter directive overrides.
///
/// When running, the worker asserts a set of routes (based on [`LoggingAPIHandler`]) that allow triggering
/// an override of the current logging filter directives as well as clearing (resetting) the active override.
pub struct LoggingOverrideWorker {
    handler: LoggingAPIHandler,
    state: Arc<Mutex<LoggingOverrideWorkerState>>,
}

struct LoggingOverrideWorkerState {
    reload_handle: Handle<Targets, Registry>,
    rx: mpsc::Receiver<LoggingOverrideAction>,
}

impl LoggingOverrideWorker {
    /// Creates a new `LoggingOverrideWorker` driving the given reload handle.
    ///
    /// When the worker starts, the filter directives present in the reload handle will be used as the "base" filter:
    /// the filter that's reapplied after an override expires or is reset. This base filter can then be subsequently
    /// updated through the [`LoggingOverrideController`] handle that's returned.
    pub(super) fn new(reload_handle: Handle<Targets, Registry>) -> (Self, LoggingOverrideController) {
        let (tx, rx) = mpsc::channel(1);
        let controller = LoggingOverrideController { tx };
        let handler = LoggingAPIHandler::new(controller.clone());
        let worker = Self {
            handler,
            state: Arc::new(Mutex::new(LoggingOverrideWorkerState { reload_handle, rx })),
        };

        (worker, controller)
    }
}

#[async_trait]
impl Supervisable for LoggingOverrideWorker {
    fn name(&self) -> &str {
        "level-override-processor"
    }

    async fn initialize(&self, process_shutdown: ShutdownHandle) -> Result<SupervisorFuture, InitializationError> {
        let mut state = self.state.clone().lock_owned().await;
        let logging_route = DynamicRoute::http(EndpointType::Privileged, &self.handler);

        Ok(Box::pin(async move {
            DataspaceRegistry::try_current()
                .ok_or_else(|| generic_error!("Dataspace not available."))?
                .assert(logging_route, "logging-api");

            process_override_actions(&mut state, process_shutdown).await;
            Ok(())
        }))
    }
}

async fn process_override_actions(state: &mut LoggingOverrideWorkerState, process_shutdown: ShutdownHandle) {
    // Seed the canonical base filter from the reload handle. If the underlying reload layer has been dropped, we
    // cannot perform overrides or resets meaningfully: just emit an error and wait for shutdown so that we don't
    // exit prematurely and drive the supervisor into an infinite restart loop.
    let mut base_filter = match state.reload_handle.clone_current() {
        Some(filter) => filter,
        None => {
            error!("Logging subsystem is in an indeterminate state; dynamic log filtering will not be available.");

            process_shutdown.await;
            return;
        }
    };

    let mut override_active = false;
    let override_timeout = sleep(Duration::from_secs(3600));

    pin!(override_timeout, process_shutdown);

    loop {
        select! {
            _ = &mut process_shutdown => break,
            maybe_action = state.rx.recv() => match maybe_action {
                Some(LoggingOverrideAction::Override { duration, filter }) => {
                    info!(directives = %filter, "Overriding existing log filtering directives for {} seconds...", duration.as_secs());

                    match state.reload_handle.reload(filter) {
                        Ok(()) => {
                            // We were able to successfully reload the filter, so mark ourselves as having an active
                            // override and update the override timeout.
                            override_active = true;
                            override_timeout.as_mut().reset(tokio::time::Instant::now() + duration);
                        },
                        Err(e) => error!(error = %e, "Failed to override log filtering directives."),
                    }
                },

                Some(LoggingOverrideAction::Reset) => {
                    // We've been instructed to immediately reset the filter back to the base one, so simply update
                    // the override timeout to fire as soon as possible.
                    override_timeout.as_mut().reset(tokio::time::Instant::now());
                },

                Some(LoggingOverrideAction::UpdateBase(new_base)) => {
                    // Before replacing the base filter, check if the new base is different from the current one before we trigger a reload
                    // and log a big, noisy message.
                    if new_base == base_filter {
                        continue;
                    }

                    base_filter = new_base;

                    if override_active {
                        // An override is currently active. Don't clobber it -- the new base will take effect when
                        // the override expires.
                        info!(
                            directives = %base_filter,
                            "Updated base log filtering directives; application deferred until active override expires."
                        );
                    } else {
                        let new_directives = base_filter.to_string();
                        match state.reload_handle.reload(base_filter.clone()) {
                            Ok(()) => info!(directives = %new_directives, "Updated base log filtering directives."),
                            Err(e) => error!(error = %e, "Failed to update base log filtering directives."),
                        }
                    }
                },

                // Our sender has dropped, so there's no more actions for us to handle.
                None => break,
            },
            _ = &mut override_timeout => {
                // Our override timeout has fired. If we have an active override, reset it back to the base filter.
                //
                // Otherwise, this just means that we've been running for a while without any overrides, so we can just
                // reset the timeout with a long duration.
                if override_active {
                    override_active = false;

                    let restore_directives = base_filter.to_string();
                    if let Err(e) = state.reload_handle.reload(base_filter.clone()) {
                        error!(error = %e, "Failed to reset log filtering directives.");
                    }

                    info!(directives = %restore_directives, "Restored base log filtering directives.");
                }

                override_timeout.as_mut().reset(tokio::time::Instant::now() + Duration::from_secs(3600));
            }
        }
    }

    // Reset our filter to the base one before we exit.
    if let Err(e) = state.reload_handle.reload(base_filter) {
        error!(error = %e, "Failed to reset log filtering directives before override handler shutdown.");
    }
}

type LogLevelStream = Pin<Box<dyn Stream<Item = LogLevel> + Send>>;

/// A worker that applies log level changes as they arrive.
///
/// Each [`LogLevel`] from the stream replaces the base filter through a [`LoggingOverrideController`], so an active
/// override stays in place until it expires or is reset. The stream carries levels that are already translated: the
/// application turns its own configuration changes into `LogLevel`s, and decides what to do with values that it can't
/// translate.
///
/// The worker keeps the stream across restarts, so a restarted worker picks up where the previous run left off. If the
/// stream ends, the worker waits for shutdown rather than returning, so that it doesn't drive the supervisor into an
/// infinite restart loop.
pub struct DynamicLogLevelWorker {
    levels: Arc<Mutex<LogLevelStream>>,
    controller: LoggingOverrideController,
}

impl DynamicLogLevelWorker {
    /// Creates a new `DynamicLogLevelWorker` that applies each log level from `levels` through `controller`.
    pub fn new<S>(levels: S, controller: LoggingOverrideController) -> Self
    where
        S: Stream<Item = LogLevel> + Send + 'static,
    {
        Self {
            levels: Arc::new(Mutex::new(Box::pin(levels))),
            controller,
        }
    }
}

#[async_trait]
impl Supervisable for DynamicLogLevelWorker {
    fn name(&self) -> &str {
        "dynamic-log-level"
    }

    async fn initialize(&self, process_shutdown: ShutdownHandle) -> Result<SupervisorFuture, InitializationError> {
        let mut levels = self.levels.clone().lock_owned().await;
        let controller = self.controller.clone();

        Ok(Box::pin(async move {
            pin!(process_shutdown);

            debug!("Dynamic log level worker started.");

            loop {
                select! {
                    _ = &mut process_shutdown => break,
                    maybe_level = levels.next() => match maybe_level {
                        Some(log_level) => {
                            if let Err(e) = controller.update_base(log_level.as_targets()).await {
                                warn!(error = %e, %log_level, "Failed to apply updated log level.");
                            }
                        }
                        None => {
                            debug!("Log level stream ended. No further log level changes will be applied.");
                            (&mut process_shutdown).await;
                            break;
                        }
                    },
                }
            }

            debug!("Dynamic log level worker stopped.");

            Ok(())
        }))
    }
}

#[cfg(test)]
mod tests {
    use std::time::Duration;

    use tokio::{sync::mpsc, time::sleep};
    use tracing_subscriber::{reload, Registry};

    use super::*;

    fn filter(directives: &str) -> Targets {
        parse_filter_directives(directives).expect("valid filter directives")
    }

    fn current_filter(handle: &reload::Handle<Targets, Registry>) -> String {
        handle
            .clone_current()
            .expect("reload layer should be alive")
            .to_string()
    }

    async fn wait_for_filter(handle: &reload::Handle<Targets, Registry>, expected: &str) {
        let deadline = tokio::time::Instant::now() + Duration::from_secs(1);

        loop {
            let actual = current_filter(handle);
            if actual == expected {
                return;
            }

            assert!(
                tokio::time::Instant::now() < deadline,
                "filter did not become `{expected}`; last value was `{actual}`"
            );

            sleep(Duration::from_millis(10)).await;
        }
    }

    fn spawn_processor(
        base_filter: Targets,
    ) -> (
        LoggingOverrideController,
        reload::Handle<Targets, Registry>,
        tokio::task::JoinHandle<()>,
        reload::Layer<Targets, Registry>,
    ) {
        let (filter_layer, reload_handle) = reload::Layer::new(base_filter);
        let (tx, rx) = mpsc::channel(1);
        let controller = LoggingOverrideController { tx };
        let mut state = LoggingOverrideWorkerState {
            reload_handle: reload_handle.clone(),
            rx,
        };
        let processor = tokio::spawn(async move {
            process_override_actions(&mut state, ShutdownHandle::noop()).await;
        });

        // The reload layer must outlive the handles -- callers should bind it (e.g., `let _layer = ...`) to keep
        // the handle alive for the duration of the test.
        (controller, reload_handle, processor, filter_layer)
    }

    fn log_level(directives: &str) -> LogLevel {
        LogLevel::try_from(directives.to_string()).expect("valid log level")
    }

    /// Returns a stream that yields each log level sent on the returned sender.
    fn log_level_stream() -> (mpsc::Sender<LogLevel>, impl Stream<Item = LogLevel> + Send + 'static) {
        let (tx, rx) = mpsc::channel(4);
        let levels = futures::stream::unfold(rx, |mut rx| async move { rx.recv().await.map(|level| (level, rx)) });
        (tx, levels)
    }

    async fn next_base_update(actions: &mut mpsc::Receiver<LoggingOverrideAction>) -> String {
        let action = tokio::time::timeout(Duration::from_secs(2), actions.recv())
            .await
            .expect("worker should send an action")
            .expect("worker should still be running");
        match action {
            LoggingOverrideAction::UpdateBase(filter) => filter.to_string(),
            _ => panic!("worker should only update the base filter"),
        }
    }

    #[tokio::test]
    async fn dynamic_log_level_worker_applies_each_level_as_the_base_filter() {
        let (tx, mut actions) = mpsc::channel(1);
        let (levels_tx, levels) = log_level_stream();
        let worker = DynamicLogLevelWorker::new(levels, LoggingOverrideController { tx });

        let (coordinator, shutdown) = ShutdownHandle::paired();
        let run = tokio::spawn(worker.initialize(shutdown).await.expect("worker initializes"));

        for directives in ["saluki=debug", "saluki=warn"] {
            levels_tx.send(log_level(directives)).await.unwrap();
            assert_eq!(next_base_update(&mut actions).await, directives);
        }

        coordinator.shutdown();
        run.await.expect("worker task completes").expect("worker exits cleanly");
    }

    #[tokio::test]
    async fn dynamic_log_level_worker_waits_for_shutdown_after_the_stream_ends() {
        let (tx, _actions) = mpsc::channel(1);
        let worker = DynamicLogLevelWorker::new(futures::stream::empty(), LoggingOverrideController { tx });

        let (coordinator, shutdown) = ShutdownHandle::paired();
        let mut run = tokio::spawn(worker.initialize(shutdown).await.expect("worker initializes"));

        // Returning early would get the worker restarted, so it has to keep running until shutdown. Nothing else can
        // wake it, so this bound is deterministic rather than timing-dependent.
        assert!(tokio::time::timeout(Duration::from_millis(100), &mut run)
            .await
            .is_err());

        coordinator.shutdown();
        run.await.expect("worker task completes").expect("worker exits cleanly");
    }

    #[tokio::test]
    async fn dynamic_log_level_worker_keeps_the_stream_across_restarts() {
        let (tx, mut actions) = mpsc::channel(1);
        let (levels_tx, levels) = log_level_stream();
        let worker = DynamicLogLevelWorker::new(levels, LoggingOverrideController { tx });

        // Each pass is a separate run of the same worker, as the supervisor would start after a restart.
        for directives in ["saluki=debug", "saluki=warn"] {
            let (coordinator, shutdown) = ShutdownHandle::paired();
            let run = tokio::spawn(worker.initialize(shutdown).await.expect("worker initializes"));

            levels_tx.send(log_level(directives)).await.unwrap();
            assert_eq!(next_base_update(&mut actions).await, directives);

            coordinator.shutdown();
            run.await.expect("worker task completes").expect("worker exits cleanly");
        }
    }

    #[tokio::test]
    async fn reset_restores_current_base_filter() {
        let (controller, reload_handle, processor, _filter_layer) = spawn_processor(filter("agent_data_plane=info"));

        controller
            .override_for(Duration::from_secs(60), filter("hyper=warn"))
            .await
            .expect("send override");
        wait_for_filter(&reload_handle, "hyper=warn").await;

        controller
            .update_base(filter("agent_data_plane=debug"))
            .await
            .expect("send update_base");
        controller.reset().await.expect("send reset");
        wait_for_filter(&reload_handle, "agent_data_plane=debug").await;

        drop(controller);
        processor.await.expect("override processor should exit cleanly");
    }

    #[tokio::test]
    async fn override_expiration_restores_current_base_filter() {
        let (controller, reload_handle, processor, _filter_layer) = spawn_processor(filter("agent_data_plane=info"));

        controller
            .override_for(Duration::from_millis(100), filter("hyper=warn"))
            .await
            .expect("send override");
        wait_for_filter(&reload_handle, "hyper=warn").await;

        controller
            .update_base(filter("agent_data_plane=warn"))
            .await
            .expect("send update_base");
        wait_for_filter(&reload_handle, "agent_data_plane=warn").await;

        drop(controller);
        processor.await.expect("override processor should exit cleanly");
    }

    #[tokio::test]
    async fn update_base_applies_immediately_when_no_override_active() {
        let (controller, reload_handle, processor, _filter_layer) = spawn_processor(filter("agent_data_plane=info"));

        controller
            .update_base(filter("agent_data_plane=debug"))
            .await
            .expect("send update_base");
        wait_for_filter(&reload_handle, "agent_data_plane=debug").await;

        drop(controller);
        processor.await.expect("override processor should exit cleanly");
    }

    #[tokio::test]
    async fn override_handler_rejects_span_and_field_filters() {
        let (tx, mut rx) = mpsc::channel(1);
        let state = LoggingHandlerState {
            controller: LoggingOverrideController { tx },
        };

        let response = LoggingAPIHandler::override_handler(
            State(state),
            Query(OverrideQueryParams { time_secs: 60 }),
            "saluki[span]=debug".to_string(),
        )
        .await
        .into_response();

        assert_eq!(response.status(), StatusCode::BAD_REQUEST);
        assert!(rx.try_recv().is_err(), "a rejected override must not reach the worker");
    }
}