orion-server 1.8.1

//! Engine reload and the Kafka-consumer lifecycle that rides along with it.
//!
//! R25: this lived in `server/routes/mod.rs` — 226 lines of engine and Kafka
//! lifecycle with no HTTP in them, in a file whose `api_routes()` is fourteen
//! lines and whose name says *route table*. The admin handlers that trigger a
//! reload and the cluster epoch watcher both call in here; neither is a route.

use std::sync::Arc;

use crate::server::state::AppState;

/// Options for [`reload_engine_with_opts`].
#[derive(Clone, Copy, Default)]
pub struct ReloadOpts {
    /// Add 0–5 s of per-node jitter before a full Kafka consumer restart.
    /// Epoch-driven reloads fire on every node near-simultaneously; without
    /// jitter they would all leave and rejoin the consumer group at once.
    pub kafka_restart_jitter: bool,
}

/// Reload the engine with all active channels and workflows from the database.
pub async fn reload_engine(state: &AppState) -> Result<(), crate::errors::OrionError> {
    reload_engine_with_opts(state, ReloadOpts::default()).await
}

/// Epoch-driven resync from the database, run by the watcher when another
/// node's mutation advanced the config epoch.
///
/// `scope` is what the bumping node said it changed. It used to be nothing at
/// all — the epoch was a bare counter — so this always did the widest resync
/// there is: reload every connector and evict every cached SQL, MongoDB and
/// cache pool. One workflow activation was therefore a fleet-wide reconnect
/// storm, every node dropping every pooled connection for a change that
/// touched no connector.
///
/// [`EpochScope::Definitions`](crate::cluster::EpochScope::Definitions) skips the connector half. That is sound in the
/// direction that matters: the channel loader keys its per-channel reuse on
/// the connector registry's generation token
/// (`channel::registry::DepsFingerprint`), and not reloading connectors leaves
/// that token where it is, so a channel whose connectors did not change is
/// carried over with its limiters and semaphores intact — which is exactly
/// what should happen. It also honours that module's standing rule, *never
/// evict a pool without loading connectors*, by doing neither.
///
/// An unknown scope reads as [`EpochScope::All`](crate::cluster::EpochScope::All), so a bump from a newer node
/// this one does not understand costs the storm rather than a missed change.
pub async fn resync_from_db(
    state: &AppState,
    scope: crate::cluster::EpochScope,
) -> Result<(), crate::errors::OrionError> {
    if scope.touches_connectors() {
        state
            .connector_registry
            .reload(state.repos.connectors.as_ref())
            .await?;
        // Paired with the reload above, always: a pool dropped without its
        // connector being re-read is a pool rebuilt from config this node has
        // not re-checked.
        state.caches.sql_pool_cache.evict_all().await;
        state.caches.mongo_pool_cache.evict_all().await;
        state.caches.cache_pool.evict_all_pools().await;
    }
    reload_engine_with_opts(
        state,
        ReloadOpts {
            kafka_restart_jitter: true,
        },
    )
    .await
}

#[tracing::instrument(skip(state, opts))]
pub async fn reload_engine_with_opts(
    state: &AppState,
    opts: ReloadOpts,
) -> Result<(), crate::errors::OrionError> {
    let start = std::time::Instant::now();

    // One reload at a time, process-wide. Everything below is a
    // read-modify-write over the database rows and the published generation;
    // overlapping reloads read the same rows and the loser's publish wins.
    // See `AppStateInner::reload_lock`.
    let _reload_guard = state.reload_lock.lock().await;

    let result = async {
        // Two independent reads, overlapped: this runs on every admin mutation
        // and, in cluster mode, on every node whose epoch watcher sees a bump,
        // so the sequential form paid two round trips where one suffices.
        let (channels, active_workflows) = tokio::try_join!(
            state.repos.channels.list_active(),
            state.repos.workflows.list_active(),
        )?;
        let channels = crate::engine::filter_channels(channels, &state.config.channel_filter);
        // The outgoing generation, read once: it is the screen for the new
        // workflows, the base the new engine is built from, and the reuse
        // cache the new channel estate is built against. Reading it once is
        // what makes those three agree.
        //
        // The running engine is the screen because `with_new_workflows`
        // carries the handler registry across, so the handlers that will run
        // these workflows are exactly the ones already registered there. A
        // workflow they cannot dispatch is quarantined per channel rather than
        // failing the reload — which would take down every channel on every
        // node over one bad stored row.
        let current = state.runtime.load();

        // The plugin half decides how the engine is built. An unchanged
        // active plugin set keeps the cheap path: `with_new_workflows` carries
        // the running engine's handler map across, so the running engine is
        // the screen and the handlers that will run these workflows are
        // exactly the ones already registered there. A changed set — an
        // activation, an archive, a new digest — needs handlers the old
        // engine does not have, so the engine is assembled afresh from the
        // node's live components plus the loaded plugins and screened against
        // that builder. Either way a workflow the handlers cannot dispatch is
        // quarantined per channel rather than failing the reload, both halves
        // are built before anything is published, and a failure here leaves
        // the previous generation serving whole. Back when the channel estate
        // was published first, this `?` could leave the node *permanently*
        // mixed: new guards over old workflows, until some later reload
        // happened to succeed.
        let plugin_rows = state.repos.plugins.list_active().await?;
        let plugins_changed =
            crate::plugin::PluginSet::fingerprint_of(&plugin_rows) != current.plugins.fingerprint();
        let (new_engine, functions, plugins, engine_issues) = if plugins_changed {
            let plugins = Arc::new(
                crate::plugin::load_active(
                    plugin_rows,
                    state.repos.plugins.as_ref(),
                    state.plugins.as_ref(),
                    &state.config.plugins,
                )
                .await,
            );
            let functions = Arc::new(
                crate::engine::FunctionRegistry::builtin()
                    .with_entries(plugins.entries())
                    .map_err(|e| crate::errors::OrionError::Config {
                        message: format!("plugin function registry: {e}"),
                    })?,
            );
            let mut handlers = super::build_handlers(state);
            for (name, handler) in plugins.handlers() {
                handlers.insert(name, handler);
            }
            let builder = crate::engine::operators::with_ops_budget(
                crate::engine::operators::with_orion_engine_defaults(
                    dataflow_rs::Engine::builder(),
                    &state.secrets,
                ),
                state.config.engine.ops_budget,
            )
            .with_handlers(handlers);
            let (workflows, mut engine_issues) =
                crate::engine::build_engine_workflows(&channels, &active_workflows, &builder);
            for issue in &mut engine_issues {
                plugins.annotate(&mut issue.reason);
            }
            let engine = Arc::new(
                builder
                    .with_workflows(workflows)
                    .build()
                    .map_err(crate::errors::OrionError::Engine)?
                    .with_observer(Arc::new(crate::engine::MetricsObserver)),
            );
            (engine, functions, plugins, engine_issues)
        } else {
            let (workflows, mut engine_issues) = crate::engine::build_engine_workflows(
                &channels,
                &active_workflows,
                &*current.engine,
            );
            for issue in &mut engine_issues {
                current.plugins.annotate(&mut issue.reason);
            }
            let engine = Arc::new(
                current
                    .engine
                    .with_new_workflows(workflows)
                    .map_err(crate::errors::OrionError::Engine)?,
            );
            (
                engine,
                current.functions.clone(),
                current.plugins.clone(),
                engine_issues,
            )
        };

        // The model half, compiled on the engine just built — never carried
        // across, because a `datalogic` program is bound to the engine that
        // compiled it and both paths above produced a fresh one, including
        // the cheap one: `with_new_workflows` compiles on a fresh expression
        // engine too. So there is no model equivalent of the plugin
        // fingerprint comparison above, and a set is rebuilt whether or not
        // the rows moved; `ModelSet`'s own docs carry the reason, and
        // `a_reload_hands_the_model_set_a_new_expression_engine` pins the
        // upstream fact it rests on. A workflow naming a model the set does
        // not serve is quarantined with the reason, the way a workflow
        // naming an unavailable plugin function is.
        let model_rows = state.repos.models.list_active().await?;
        let models = Arc::new(crate::model::ModelSet::load_active(
            &model_rows,
            &state.config.models,
            state.models.is_some(),
            new_engine.datalogic(),
        ));
        let mut engine_issues = engine_issues;
        engine_issues.extend(super::models::load_issues(
            &channels,
            &active_workflows,
            &models,
        ));

        // Channels that fail to load are quarantined — refused at every
        // ingress — and the reload proceeds (F35). It used to abort here,
        // which meant one unparseable `config_json` failed every activate,
        // archive, delete and rollout with a 500, and stopped the cluster
        // epoch watcher resyncing all nodes.
        // The quarantine set is recorded on the generation and read back with
        // `ChannelSnapshot::quarantined` (reported via `/health`), so it is
        // deliberately not propagated as an error here.
        let new_channels = state
            .channel_loader
            .build(
                &current.channels,
                &channels,
                crate::channel::ReloadDeps {
                    vars: state.vars.as_ref(),
                    connector_registry: &state.connector_registry,
                    cache_pool: &state.caches.cache_pool,
                    datalogic: &state.datalogic,
                    jwks: &state.jwks,
                    http_client: &state.http_client,
                    allow_private_token_urls: state.config.oauth2_login.allow_private_token_urls,
                    global_trace_storage: &state.config.trace_storage,
                    cron_enabled: state.config.cron.enabled,
                },
                engine_issues,
            )
            .await;

        // One store, both halves. There is no window in which a reader is held
        // off, so this needs neither a timeout nor a carefully-scoped drop
        // before the Kafka restart below — and no window in which a request is
        // admitted by one generation and executed by another, which is what
        // two stores here could not avoid however they were ordered.
        let generation = state.runtime.publish(
            new_engine,
            Arc::new(new_channels),
            functions,
            plugins,
            models,
        );

        // Warm what `models.preload` selects, off the publish's critical
        // path: the generation serves now, and a request that arrives before
        // its model is resident shares the load in flight.
        super::models::spawn_preload(
            super::models::PreloadDeps::from_state(state),
            state.runtime.load(),
            &active_workflows,
        );

        // Update active workflows gauge
        crate::metrics::set_active_workflows(active_workflows.len() as f64);

        // Restart Kafka consumer if async channel topics changed
        if state.config.kafka.enabled {
            restart_kafka_consumer_if_needed(state, &channels, opts).await;
        }

        tracing::info!(
            generation,
            workflow_count = active_workflows.len(),
            channel_count = channels.len(),
            "Runtime generation published"
        );
        Ok(())
    }
    .await;

    let duration = start.elapsed().as_secs_f64();
    crate::metrics::record_engine_reload_duration(duration);

    // The degraded flag is set here rather than at the call sites, so it holds
    // for every way a reload is triggered — an admin mutation, an explicit
    // `POST /engine/reload`, and the cluster epoch watcher's resync. A failure
    // leaves this node serving the *previous* generation: correct, but not what
    // the database says, and nothing else would report that.
    match &result {
        Ok(()) => {
            crate::metrics::record_engine_reload("success");
            state
                .reload_degraded
                .store(false, std::sync::atomic::Ordering::Release);
        }
        Err(e) => {
            crate::metrics::record_engine_reload("failure");
            state
                .reload_degraded
                .store(true, std::sync::atomic::Ordering::Release);
            tracing::error!(
                error = %e,
                "Engine reload failed: this node is still serving the previous \
                 generation, so its running config no longer matches the database"
            );
        }
    }

    result
}

/// Restart the Kafka consumer when async channel topic mappings have changed.
///
/// Merges config-file topics with DB-driven async channel topics. If the set
/// of topics differs from what the current consumer is subscribed to, the old
/// consumer is shut down and a new one is started.
async fn restart_kafka_consumer_if_needed(
    state: &AppState,
    channels: &[crate::storage::models::Channel],
    opts: ReloadOpts,
) {
    use std::collections::HashSet;

    let all_topics = crate::kafka::merge_kafka_topics(&state.config.kafka, channels);

    let new_topic_set: HashSet<String> = all_topics.iter().map(|t| t.topic.clone()).collect();

    let mut handle_guard = state.kafka.consumer_handle.lock().await;

    // Optimisation: if topics haven't changed, pause/resume instead of full restart
    if let Some(ref existing_handle) = *handle_guard
        && *existing_handle.topics() == new_topic_set
    {
        tracing::info!("Kafka topics unchanged, pausing consumer during engine swap");
        if let Err(e) = existing_handle.pause() {
            tracing::warn!(error = %e, "Failed to pause Kafka consumer, falling back to full restart");
        } else {
            // Brief sleep to allow in-flight messages to finish processing
            tokio::time::sleep(std::time::Duration::from_millis(100)).await;
            if let Err(e) = existing_handle.resume() {
                tracing::error!(error = %e, "Failed to resume Kafka consumer after engine reload");
                // Fall through to full restart below
            } else {
                tracing::info!("Kafka consumer resumed after engine reload");
                return;
            }
        }
    }

    // Full restart path: pause first to minimize gap, then shutdown and restart
    if opts.kafka_restart_jitter {
        let jitter_ms = rand::random_range(0..=5000u64);
        tracing::info!(
            jitter_ms,
            "Jittering Kafka consumer restart (epoch-driven reload)"
        );
        tokio::time::sleep(std::time::Duration::from_millis(jitter_ms)).await;
    }
    if let Some(ref existing_handle) = *handle_guard {
        let _ = existing_handle.pause(); // Best-effort pause before shutdown
        tokio::time::sleep(std::time::Duration::from_millis(50)).await;
    }
    if let Some(old_handle) = handle_guard.take() {
        tracing::info!("Shutting down Kafka consumer for topic refresh...");
        old_handle.shutdown().await;
    }

    // K7: the old handle is gone. A start failure below must not leave that
    // as the permanent end state with every probe green — flag ingestion as
    // degraded and hand recovery to the supervisor.
    match try_start_ingest(state, channels) {
        Ok(Some(new_handle)) => {
            tracing::info!(
                topics = ?new_topic_set,
                "Kafka consumer restarted with updated topics"
            );
            *handle_guard = Some(new_handle);
            state.kafka.ingest_status.set_degraded(false);
        }
        Ok(None) => {
            tracing::info!("No Kafka topics configured or from DB, consumer not started");
            state.kafka.ingest_status.set_degraded(false);
        }
        Err(e) => {
            state.kafka.ingest_status.set_degraded(true);
            crate::metrics::record_error("kafka_restart");
            tracing::error!(
                error = %e,
                "Failed to restart Kafka consumer; ingestion is down — retrying with backoff"
            );
            drop(handle_guard);
            spawn_kafka_restart_supervisor(state);
        }
    }
}

/// Start an ingest consumer for the current channel set through the same
/// builder startup uses ([`crate::bootstrap::start_kafka_ingest`]), so the
/// boot, reload and supervisor paths cannot drift. `Ok(None)` when the
/// merged topic list is empty. The error is stringified because the caller
/// holds it across a spawned task boundary.
fn try_start_ingest(
    state: &AppState,
    channels: &[crate::storage::models::Channel],
) -> Result<Option<crate::kafka::consumer::ConsumerHandle>, String> {
    crate::bootstrap::start_kafka_ingest(
        &state.config.kafka,
        channels,
        crate::bootstrap::IngestDeps {
            runtime: state.runtime.clone(),
            datalogic: state.datalogic.clone(),
            vars: state.vars.clone(),
            kafka_producer: state.kafka.producer.clone(),
            instance_id: state
                .cluster
                .enabled
                .then(|| state.cluster.instance_id.clone()),
            trace_repo: state.repos.traces.clone(),
            persistence_queue: state.trace_persistence_queue.clone(),
            max_result_size_bytes: state.config.trace_queue.max_result_size_bytes,
        },
    )
    .map_err(|e| e.to_string())
}

/// Supervise a Kafka consumer that failed to (re)start: retry with capped
/// exponential backoff (1 s doubling to 60 s) until a consumer is running
/// again, then clear the degraded flag (K7). The channel list is re-read
/// from the database on every attempt, so topic changes made while
/// ingestion was down are honoured. At most one supervisor runs per
/// process; it stands down when another reload restores the consumer first,
/// or when the node starts draining.
///
/// Slot invariant: every exit path releases the supervisor slot while the
/// handle mutex is still held. A failed reload runs its start attempt and
/// `set_degraded(true)` inside that same mutex, so by the time it drops the
/// mutex and calls `claim_supervisor` the slot is either free (this
/// supervisor's final critical section already ran — the claim succeeds and
/// a fresh supervisor spawns) or claimed by a supervisor that is still
/// looping and will retry the new failure itself. Releasing *after* the
/// mutex was a TOCTOU: a reload failing in the unlock→release gap saw the
/// slot occupied, spawned nothing, and left the process degraded with no
/// supervisor until the next reload.
pub fn spawn_kafka_restart_supervisor(state: &AppState) {
    if !state.kafka.ingest_status.claim_supervisor() {
        return;
    }
    let state = state.clone();
    tokio::spawn(async move {
        let mut backoff_ms = crate::kafka::consumer::INITIAL_RETRY_BACKOFF_MS;
        loop {
            tokio::time::sleep(std::time::Duration::from_millis(backoff_ms)).await;
            let mut handle_guard = state.kafka.consumer_handle.lock().await;
            // Draining: do not resurrect a consumer mid-shutdown. (Checked
            // under the mutex like every other exit, so even this release
            // cannot race a failing reload's claim — and a spawn lost to a
            // drain-window race would only ever supervise a node that is
            // shutting down.)
            if !state.ready.load(std::sync::atomic::Ordering::Acquire) {
                state.kafka.ingest_status.release_supervisor();
                break;
            }
            if handle_guard.is_some() {
                // Another reload already restarted the consumer.
                state.kafka.ingest_status.release_supervisor();
                break;
            }
            let channels = match state.repos.channels.list_active().await {
                Ok(channels) => {
                    crate::engine::filter_channels(channels, &state.config.channel_filter)
                }
                Err(e) => {
                    tracing::warn!(
                        error = %e,
                        backoff_ms,
                        "Kafka restart supervisor could not list channels; will retry"
                    );
                    drop(handle_guard);
                    backoff_ms = crate::kafka::consumer::next_backoff_ms(backoff_ms);
                    continue;
                }
            };
            let started = try_start_ingest(&state, &channels);
            match started {
                Ok(Some(handle)) => {
                    *handle_guard = Some(handle);
                    // Slot before flag: both are Release stores, so an
                    // observer that sees the degraded flag cleared also
                    // sees the slot free.
                    state.kafka.ingest_status.release_supervisor();
                    state.kafka.ingest_status.set_degraded(false);
                    tracing::info!("Kafka consumer restored by restart supervisor");
                    break;
                }
                Ok(None) => {
                    // Nothing to ingest any more — idle, not degraded.
                    state.kafka.ingest_status.release_supervisor();
                    state.kafka.ingest_status.set_degraded(false);
                    tracing::info!("No Kafka topics remain; restart supervisor standing down");
                    break;
                }
                Err(e) => {
                    crate::metrics::record_error("kafka_restart");
                    tracing::error!(
                        error = %e,
                        backoff_ms,
                        "Kafka consumer restart failed; ingestion still down"
                    );
                    drop(handle_guard);
                    backoff_ms = crate::kafka::consumer::next_backoff_ms(backoff_ms);
                }
            }
        }
    });
}