aster_forge_config/notification/

supervisor.rs

use async_trait::async_trait;
use std::future::Future;
use std::sync::Arc;
use std::time::{Duration, Instant};
use tokio::sync::mpsc;
use tokio::time::Instant as RetryInstant;
use tokio_util::sync::CancellationToken;

use crate::{ConfigCoreError, Result};

use super::message::{
    ConfigChangeEvent, ConfigReloadDecision, ConfigReloadMessage, ConfigReloadWorkerConfig,
    handle_config_reload_notification,
};
use super::notifier::{ConfigChangeNotifier, ConfigNotification};

/// Observability event emitted after handling a config reload notification.
#[derive(Debug, Clone, PartialEq)]
pub struct ConfigReloadObservation {
    /// Source label suitable for low-cardinality metrics.
    pub source: &'static str,
    /// Handling decision.
    pub decision: ConfigReloadDecision,
    /// Whether the handling path succeeded.
    pub status: &'static str,
    /// Number of changed keys advertised by the reload hint.
    pub changed_keys: u64,
    /// Time spent handling this notification.
    pub duration_seconds: f64,
}

impl ConfigReloadObservation {
    fn new(
        source: &'static str,
        decision: ConfigReloadDecision,
        status: &'static str,
        changed_keys: usize,
        duration_seconds: f64,
    ) -> Self {
        Self {
            source,
            decision,
            status,
            changed_keys: u64::try_from(changed_keys).unwrap_or(u64::MAX),
            duration_seconds,
        }
    }
}

/// Connection lifecycle state for a config-sync subscription.
pub type ConfigSyncConnectionState = aster_forge_events::EventConnectionState;

/// Low-cardinality observation emitted for config-sync connection transitions.
#[derive(Debug, Clone, PartialEq)]
pub struct ConfigSyncConnectionObservation {
    /// Connection lifecycle state.
    pub state: ConfigSyncConnectionState,
    /// One-based reconnect attempt number, or zero outside reconnect attempts.
    pub reconnect_attempt: u32,
    /// Planned backoff for a reconnect attempt, or zero for other states.
    pub backoff_seconds: f64,
}

impl ConfigSyncConnectionObservation {
    fn new(state: ConfigSyncConnectionState, reconnect_attempt: u32, backoff: Duration) -> Self {
        Self {
            state,
            reconnect_attempt,
            backoff_seconds: backoff.as_secs_f64(),
        }
    }
}

/// Receives config reload observability events.
pub trait ConfigReloadObserver: Send + Sync {
    /// Records one reload observation.
    fn observe_config_reload(&self, observation: ConfigReloadObservation);
}

impl<F> ConfigReloadObserver for F
where
    F: Fn(ConfigReloadObservation) + Send + Sync,
{
    fn observe_config_reload(&self, observation: ConfigReloadObservation) {
        self(observation);
    }
}

/// Receives config-sync connection lifecycle observations.
pub trait ConfigSyncConnectionObserver: Send + Sync {
    /// Records one connection transition.
    fn observe_config_sync_connection(&self, observation: ConfigSyncConnectionObservation);
}

impl<F> ConfigSyncConnectionObserver for F
where
    F: Fn(ConfigSyncConnectionObservation) + Send + Sync,
{
    fn observe_config_sync_connection(&self, observation: ConfigSyncConnectionObservation) {
        self(observation);
    }
}

pub(super) type ConfigReloadReconnectPolicy = aster_forge_events::EventReconnectPolicy;

fn default_config_reload_reconnect_policy() -> ConfigReloadReconnectPolicy {
    ConfigReloadReconnectPolicy {
        initial_delay: Duration::from_millis(250),
        max_delay: Duration::from_secs(30),
        stable_reset_after: Duration::from_secs(30),
        jitter_min_percent: 50,
        jitter_max_percent: 100,
    }
}

/// Runs a reload subscription loop until `shutdown` is cancelled.
///
/// The loop never carries configuration values over pub/sub. A matching
/// notification only tells this process to reload from its authoritative store.
/// Failed reload hints are coalesced and retried with bounded exponential
/// backoff, even if no further notification arrives. Callbacks must support
/// repeated authoritative reads; an empty key list means reload all config.
/// Shutdown drops in-flight callbacks, which must be safe to cancel and repeat.
///
/// The loop uses the transport supervision of [`run_config_reload_supervisor`],
/// but has no authoritative reconcile callback. It retries received hints;
/// repairing notifications lost during disconnection requires the supervisor
/// API with a full reconcile callback.
///
/// # Errors
///
/// Returns [`ConfigError`] when publishing, subscription, reconciliation, or reload handling fails.
pub async fn run_config_reload_worker<N, F, Fut>(
    notifier: Arc<N>,
    config: ConfigReloadWorkerConfig,
    shutdown: CancellationToken,
    reload: F,
) -> Result<()>
where
    N: ConfigChangeNotifier + ?Sized,
    F: FnMut(ConfigReloadMessage) -> Fut,
    Fut: Future<Output = Result<()>>,
{
    run_config_reload_worker_with_observer(
        notifier,
        config,
        shutdown,
        reload,
        None::<&dyn ConfigReloadObserver>,
    )
    .await
}

/// Runs a reload subscription loop and reports low-cardinality observations.
///
/// # Errors
///
/// Returns [`ConfigError`] when publishing, subscription, reconciliation, or reload handling fails.
pub async fn run_config_reload_worker_with_observer<N, F, Fut>(
    notifier: Arc<N>,
    config: ConfigReloadWorkerConfig,
    shutdown: CancellationToken,
    mut reload: F,
    observer: Option<&dyn ConfigReloadObserver>,
) -> Result<()>
where
    N: ConfigChangeNotifier + ?Sized,
    F: FnMut(ConfigReloadMessage) -> Fut,
    Fut: Future<Output = Result<()>>,
{
    run_config_reload_supervisor_inner(
        notifier,
        config,
        default_config_reload_reconnect_policy(),
        shutdown,
        None::<&mut fn() -> std::future::Ready<Result<()>>>,
        &mut reload,
        observer,
        None,
    )
    .await
}

/// Runs a reconnecting reload subscription with authoritative reconciliation.
///
/// `reconcile` runs after every successful subscription, including the initial
/// connection. This closes the race between the product's startup snapshot load
/// and the moment pub/sub begins receiving notifications. After a disconnect it
/// also repairs any changes missed while the transient transport was unavailable.
///
/// A disconnect is any of: subscribe failure, transport stream error or ending,
/// and local broadcast lag (the receiver fell behind and events were dropped).
/// Each one is observed, waited out with bounded exponential backoff (250 ms
/// initial, 30 s cap, jittered; the failure counter resets after a subscription
/// stays stable for 30 s), then followed by a fresh subscription and reconcile.
/// Failed reloads or reconciliations schedule one authoritative reconcile retry
/// with the same delay bounds, independent of transport health. Only successful
/// reconciliation clears this pending work; a later key-specific reload cannot
/// prove that all previously missed changes and derived caches are current.
/// Shutdown cancels both retry waits and in-flight callbacks. Callbacks must be
/// safe to drop and repeat after partially completed work.
/// The loop only returns when `shutdown` is cancelled.
///
/// # Errors
///
/// Returns [`ConfigError`] when publishing, subscription, reconciliation, or reload handling fails.
pub async fn run_config_reload_supervisor<N, R, RFut, F, Fut>(
    notifier: Arc<N>,
    config: ConfigReloadWorkerConfig,
    shutdown: CancellationToken,
    reconcile: R,
    reload: F,
) -> Result<()>
where
    N: ConfigChangeNotifier + ?Sized,
    R: FnMut() -> RFut,
    RFut: Future<Output = Result<()>>,
    F: FnMut(ConfigReloadMessage) -> Fut,
    Fut: Future<Output = Result<()>>,
{
    run_config_reload_supervisor_with_observers(
        notifier, config, shutdown, reconcile, reload, None, None,
    )
    .await
}

/// Runs a reconnecting reload subscription and reports reload and connection observations.
///
/// # Errors
///
/// Returns [`ConfigError`] when publishing, subscription, reconciliation, or reload handling fails.
pub async fn run_config_reload_supervisor_with_observers<N, R, RFut, F, Fut>(
    notifier: Arc<N>,
    config: ConfigReloadWorkerConfig,
    shutdown: CancellationToken,
    mut reconcile: R,
    mut reload: F,
    reload_observer: Option<&dyn ConfigReloadObserver>,
    connection_observer: Option<&dyn ConfigSyncConnectionObserver>,
) -> Result<()>
where
    N: ConfigChangeNotifier + ?Sized,
    R: FnMut() -> RFut,
    RFut: Future<Output = Result<()>>,
    F: FnMut(ConfigReloadMessage) -> Fut,
    Fut: Future<Output = Result<()>>,
{
    run_config_reload_supervisor_inner(
        notifier,
        config,
        default_config_reload_reconnect_policy(),
        shutdown,
        Some(&mut reconcile),
        &mut reload,
        reload_observer,
        connection_observer,
    )
    .await
}

#[expect(
    clippy::too_many_arguments,
    reason = "The supervisor accepts explicit notifier, retry, reload, cancellation, and observability dependencies."
)]
pub(super) async fn run_config_reload_supervisor_inner<N, R, RFut, F, Fut>(
    notifier: Arc<N>,
    config: ConfigReloadWorkerConfig,
    reconnect_policy: ConfigReloadReconnectPolicy,
    shutdown: CancellationToken,
    mut reconcile: Option<&mut R>,
    reload: &mut F,
    reload_observer: Option<&dyn ConfigReloadObserver>,
    connection_observer: Option<&dyn ConfigSyncConnectionObserver>,
) -> Result<()>
where
    N: ConfigChangeNotifier + ?Sized,
    R: FnMut() -> RFut,
    RFut: Future<Output = Result<()>>,
    F: FnMut(ConfigReloadMessage) -> Fut,
    Fut: Future<Output = Result<()>>,
{
    let source = Arc::new(ConfigNotifierSubscriptionSource { notifier });
    let (updates_tx, mut updates_rx) = mpsc::channel(1);
    let supervisor = aster_forge_events::supervise_event_subscription(
        source,
        reconnect_policy,
        shutdown.clone(),
        updates_tx,
    );
    tokio::pin!(supervisor);
    let mut retry = ConfigReloadRetry::default();

    loop {
        let update = tokio::select! {
            biased;
            () = shutdown.cancelled() => return Ok(()),
            () = retry.wait() => {
                retry.deadline = None;
                let Some(result) = await_config_callback(&shutdown, async {
                        if let Some(reconcile) = reconcile.as_mut() {
                            reconcile().await
                        } else if let Some(message) = retry.message.clone() {
                            process_config_reload_message(&config, message, reload, reload_observer)
                                .await.map(|_| ())
                        } else {
                            Ok(())
                        }
                    }).await else { return Ok(()); };
                if let Err(error) = result {
                    tracing::warn!(%error, "failed to retry runtime config reload");
                    retry.schedule(reconnect_policy);
                } else {
                    retry = ConfigReloadRetry::default();
                }
                continue;
            }
            () = &mut supervisor => return Ok(()),
            update = updates_rx.recv() => update,
        };
        match update {
            Some(aster_forge_events::EventSubscriptionUpdate::Connection(observation)) => {
                observe_config_sync_connection(
                    connection_observer,
                    observation.state,
                    observation.reconnect_attempt,
                    observation.backoff,
                );
                if matches!(
                    observation.state,
                    ConfigSyncConnectionState::Connected | ConfigSyncConnectionState::Recovered
                ) && let Some(reconcile) = reconcile.as_mut()
                {
                    let Some(result) = await_config_callback(&shutdown, reconcile()).await else {
                        return Ok(());
                    };
                    if let Err(error) = result {
                        tracing::warn!(
                            error = %error,
                            "failed to reconcile runtime config after subscription connected"
                        );
                        retry.schedule(reconnect_policy);
                    } else {
                        retry = ConfigReloadRetry::default();
                        tracing::debug!("runtime config reconciled after subscription connected");
                    }
                }
            }
            Some(aster_forge_events::EventSubscriptionUpdate::Item(ConfigChangeEvent::Reload(
                message,
            ))) => {
                let Some(result) = await_config_callback(
                    &shutdown,
                    process_config_reload_message(
                        &config,
                        message.clone(),
                        reload,
                        reload_observer,
                    ),
                )
                .await
                else {
                    return Ok(());
                };
                if result.is_err() {
                    if reconcile.is_none() {
                        retry.merge_message(message);
                    }
                    retry.schedule(reconnect_policy);
                }
            }
            None => return Ok(()),
        }
    }
}

async fn await_config_callback<T>(
    shutdown: &CancellationToken,
    callback: impl Future<Output = T>,
) -> Option<T> {
    tokio::select! {
        biased;
        () = shutdown.cancelled() => None,
        result = callback => Some(result),
    }
}

#[derive(Default)]
struct ConfigReloadRetry {
    deadline: Option<RetryInstant>,
    attempt: u32,
    message: Option<ConfigReloadMessage>,
}

impl ConfigReloadRetry {
    fn schedule(&mut self, policy: ConfigReloadReconnectPolicy) {
        // Further notifications must not postpone recovery or reset its backoff.
        if self.deadline.is_none() {
            self.attempt = self.attempt.saturating_add(1);
            let delay = policy
                .reconnect_delay(self.attempt)
                .max(Duration::from_millis(1));
            self.deadline = Some(RetryInstant::now() + delay);
        }
    }

    async fn wait(&self) {
        if let Some(deadline) = self.deadline {
            tokio::time::sleep_until(deadline).await;
        } else {
            std::future::pending::<()>().await;
        }
    }

    fn merge_message(&mut self, mut message: ConfigReloadMessage) {
        const MAX_PENDING_KEYS: usize = 1024;
        if let Some(pending) = &self.message {
            if pending.keys.is_empty() || message.keys.is_empty() {
                message.keys.clear();
            } else {
                message.keys.extend(pending.keys.iter().cloned());
                message.keys.sort();
                message.keys.dedup();
            }
        }
        // A long outage must not accumulate an unbounded list of failed hints.
        // Empty keys already mean an authoritative full reload.
        if message.keys.len() > MAX_PENDING_KEYS {
            message.keys.clear();
        }
        self.message = Some(message);
    }
}

async fn process_config_reload_message<F, Fut>(
    config: &ConfigReloadWorkerConfig,
    message: ConfigReloadMessage,
    reload: &mut F,
    observer: Option<&dyn ConfigReloadObserver>,
) -> Result<ConfigReloadDecision>
where
    F: FnMut(ConfigReloadMessage) -> Fut,
    Fut: Future<Output = Result<()>>,
{
    let changed_keys = message.keys.len();
    let started = Instant::now();
    let result = handle_config_reload_notification(config, message, reload).await;
    match &result {
        Ok(ConfigReloadDecision::Reloaded) => {
            observe_config_reload(
                observer,
                ConfigReloadDecision::Reloaded,
                "ok",
                changed_keys,
                started,
            );
            tracing::debug!("runtime config reloaded after remote notification");
        }
        Ok(
            decision @ (ConfigReloadDecision::IgnoredNamespace
            | ConfigReloadDecision::IgnoredOrigin),
        ) => {
            observe_config_reload(observer, *decision, "ok", changed_keys, started);
        }
        Err(error) => {
            observe_config_reload(
                observer,
                ConfigReloadDecision::Reloaded,
                "error",
                changed_keys,
                started,
            );
            tracing::warn!(
                error = %error,
                "failed to reload runtime config after remote notification"
            );
        }
    }
    result
}

#[cfg(test)]
pub(super) fn config_reload_reconnect_delay(
    policy: ConfigReloadReconnectPolicy,
    reconnect_attempt: u32,
) -> Duration {
    policy.reconnect_delay(reconnect_attempt)
}

pub(super) fn duration_millis_u64(duration: Duration) -> u64 {
    u64::try_from(duration.as_millis()).unwrap_or(u64::MAX)
}

fn observe_config_sync_connection(
    observer: Option<&dyn ConfigSyncConnectionObserver>,
    state: ConfigSyncConnectionState,
    reconnect_attempt: u32,
    backoff: Duration,
) {
    if let Some(observer) = observer {
        observer.observe_config_sync_connection(ConfigSyncConnectionObservation::new(
            state,
            reconnect_attempt,
            backoff,
        ));
    }
    match state {
        ConfigSyncConnectionState::Connected => {}
        ConfigSyncConnectionState::Recovered => {
            tracing::info!(reconnect_attempt, "config reload subscription recovered");
        }
        ConfigSyncConnectionState::Disconnected => {
            tracing::warn!(reconnect_attempt, "config reload subscription disconnected");
        }
        ConfigSyncConnectionState::Reconnecting => {
            tracing::warn!(
                reconnect_attempt,
                backoff_ms = duration_millis_u64(backoff),
                "waiting before config reload subscription reconnect"
            );
        }
    }
}

fn observe_config_reload(
    observer: Option<&dyn ConfigReloadObserver>,
    decision: ConfigReloadDecision,
    status: &'static str,
    changed_keys: usize,
    started: Instant,
) {
    if let Some(observer) = observer {
        observer.observe_config_reload(ConfigReloadObservation::new(
            "pubsub",
            decision,
            status,
            changed_keys,
            started.elapsed().as_secs_f64(),
        ));
    }
}
struct ConfigNotifierSubscriptionSource<N: ?Sized> {
    notifier: Arc<N>,
}

#[async_trait]
impl<N> aster_forge_events::EventSubscriptionSource for ConfigNotifierSubscriptionSource<N>
where
    N: ConfigChangeNotifier + ?Sized,
{
    type Item = ConfigChangeEvent;
    type Subscription = ConfigNotification;
    type Error = ConfigCoreError;

    async fn subscribe(&self) -> Result<Self::Subscription> {
        self.notifier.subscribe().await
    }

    async fn receive(&self, subscription: &mut Self::Subscription) -> Result<Self::Item> {
        subscription.recv().await
    }
}