spate_core/pipeline/

runtime.rs

//! Process assembly: threads, runtimes, observability, and the run loop.

use super::controller::{CLOSE_SHARE, ControllerContext, ControllerSignal, run_controller};
use super::driver::{DriverContext, DriverExit, DriverParams, run_driver};
use super::{DriverEvent, ExitReport, ExitState, FatalErrorReport, SinkRuntime, ThreadControl};
use crate::admin::{AdminServer, HealthState, HealthThresholds};
use crate::backpressure::{BackpressureParams, InflightBudget, WatermarkController};
use crate::checkpoint::Checkpointer;
use crate::clock::{Clock, SystemClock};
use crate::config::{MetricsExporter, PinningMode, PipelineConfig};
use crate::feed::memory::MemoryFeed;
use crate::feed::{
    DEAD_LETTER_READ_INTERVAL, DeadLetterSampler, Feed, FeedMetricsCtx, FeedOrigin, MeteredFeed,
    STOP_RESERVE,
};
use crate::metrics::{
    self, BackpressureMetrics, CheckpointMetrics, ComponentLabels, E2eBasis, Exporter, FeedMetrics,
    InflightBudgetMetrics, IoBusySampler, IoRuntimeMetrics, Meter, MetricRole, MetricsHandle,
    MetricsSettings, PipelineMetrics, PipelineState, SourceMetrics,
};
use crate::ops::RunnableChain;
use crate::sink::SinkProbeFn;
use crate::source::{DrainBarrier, Source};
use crate::task;
use std::net::SocketAddr;
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, Ordering};
use std::time::{Duration, Instant};
#[cfg(unix)]
use tokio::signal::unix::{SignalKind, signal};
use tokio::sync::watch;

/// Knobs that are not part of the user-facing YAML (loop granularities,
/// test hooks). The defaults suit production; tests shrink the timings.
#[derive(Clone, Debug)]
pub struct RuntimeOptions {
    /// Install SIGTERM/SIGINT handlers that trigger a graceful drain.
    /// Disable in tests and drive [`ShutdownHandle`] instead.
    pub handle_signals: bool,
    /// Max payloads per lane poll.
    pub max_records: usize,
    /// Lane poll timeout (also the paused/idle loop sleep).
    pub poll_timeout: Duration,
    /// Flush the chain after this long without new data.
    pub idle_flush: Duration,
    /// Sleep between retries of a blocked batch.
    pub blocked_retry: Duration,
    /// Controller `poll_events` timeout.
    pub event_poll_timeout: Duration,
    /// Version string published on `spate_pipeline_info`.
    pub version: String,
    /// How often to read the depth of the feed's dead-letter destination for
    /// `spate_feed_dead_letter_units`, starting at startup. A zero interval
    /// reads whenever no read is pending. `None` makes no such read, so every
    /// depth read the feed answers is its source's.
    pub dead_letter_read_interval: Option<Duration>,
}

impl Default for RuntimeOptions {
    fn default() -> Self {
        RuntimeOptions {
            handle_signals: true,
            max_records: 512,
            poll_timeout: Duration::from_millis(10),
            idle_flush: Duration::from_millis(100),
            blocked_retry: Duration::from_millis(2),
            event_poll_timeout: Duration::from_millis(50),
            version: env!("CARGO_PKG_VERSION").to_string(),
            dead_letter_read_interval: Some(DEAD_LETTER_READ_INTERVAL),
        }
    }
}

/// How often the main thread samples `spate_pipeline_io_busy_ratio`; the
/// upkeep cadence.
const IO_BUSY_INTERVAL: Duration = Duration::from_secs(5);

/// Triggers a graceful drain from anywhere (tests, custom signal wiring).
#[derive(Clone, Debug)]
pub struct ShutdownHandle(Arc<StopRequest>);

impl ShutdownHandle {
    /// Begin the graceful stop: a source that reads on what it holds
    /// finishes first, then the drain runs. The drain deadline counts from
    /// the first call. Idempotent.
    pub fn trigger(&self) {
        self.0.trigger();
    }
}

/// A graceful stop's request: the flag the controller and `StopSignal` read,
/// and the instant of the first trigger, read on the stop clock.
pub(crate) struct StopRequest {
    flag: Arc<AtomicBool>,
    at: std::sync::Mutex<Option<Instant>>,
    clock: std::sync::Mutex<Arc<dyn Clock + Send + Sync>>,
}

impl StopRequest {
    fn new(clock: Arc<dyn Clock + Send + Sync>) -> StopRequest {
        StopRequest {
            flag: Arc::new(AtomicBool::new(false)),
            at: std::sync::Mutex::new(None),
            clock: std::sync::Mutex::new(clock),
        }
    }

    fn trigger(&self) {
        {
            let mut at = self.at.lock().expect("stop request lock");
            if at.is_none() {
                *at = Some(self.clock.lock().expect("stop request lock").now());
            }
        }
        // Stored after the instant, so a reader that sees the flag finds it.
        self.flag.store(true, Ordering::Release);
    }

    /// Whether a stop was asked for.
    pub(crate) fn is_set(&self) -> bool {
        self.flag.load(Ordering::Acquire)
    }

    /// When the stop was first asked for, if it was.
    pub(crate) fn triggered_at(&self) -> Option<Instant> {
        *self.at.lock().expect("stop request lock")
    }

    pub(crate) fn flag(&self) -> Arc<AtomicBool> {
        Arc::clone(&self.flag)
    }
}

impl std::fmt::Debug for StopRequest {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("StopRequest")
            .field("set", &self.is_set())
            .field("at", &self.triggered_at())
            .finish_non_exhaustive()
    }
}

/// The pipeline could not start.
#[derive(Debug, thiserror::Error)]
#[non_exhaustive]
pub enum StartError {
    /// Invalid effective configuration.
    #[error("invalid runtime configuration: {0}")]
    Config(String),
    /// The metrics exporter could not be installed.
    #[error("metrics: {0}")]
    Metrics(String),
    /// The admin server could not take its address. Set `admin.listen` to a
    /// free address, or to `none` for a pipeline that needs no server.
    #[error("admin.listen: cannot bind {addr}: {source}")]
    AdminBind {
        /// The address `admin.listen` asked for.
        addr: SocketAddr,
        /// Why the kernel refused it.
        source: std::io::Error,
    },
    /// The I/O runtime or a pipeline thread could not start.
    #[error("io: {0}")]
    Io(#[from] std::io::Error),
}

/// One pipeline process: source, per-thread chains, and a sink.
///
/// The caller creates the shared [`InflightBudget`] first and wires it into
/// the chain terminals (which `add` on enqueue) and the sink workers (which
/// `sub` on durable write or abandonment) before handing everything here.
pub struct PipelineRuntime<S: Source> {
    config: PipelineConfig,
    source: S,
    chains: Box<dyn FnMut(usize) -> Box<dyn RunnableChain> + Send>,
    sink: SinkRuntime,
    budget: Arc<InflightBudget>,
    // Set by a stop request; the controller reads on to the reserve.
    stop: Arc<StopRequest>,
    // Set when the drain begins; the drivers read it.
    shutdown: Arc<AtomicBool>,
    stop_clock: Arc<dyn Clock + Send + Sync>,
    // Drives the commit tick, the fast-commit window, the stall check, the
    // idle flush and the health probe.
    clock: crate::clock::SharedClock,
    // How long a `CommitReady` partition keeps the tightened cadence; `None`
    // is one `checkpoint.interval`.
    chase_window: Option<Duration>,
    options: RuntimeOptions,
    io: Option<tokio::runtime::Runtime>,
    feed: Option<Arc<dyn Feed>>,
    io_busy_interval: Duration,
    io_busy_window: Duration,
}

impl<S: Source> std::fmt::Debug for PipelineRuntime<S> {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("PipelineRuntime")
            .field("pipeline", &self.config.pipeline.name)
            .field("options", &self.options)
            .finish_non_exhaustive()
    }
}

impl<S: Source + 'static> PipelineRuntime<S> {
    /// Assemble a runtime. `chains` builds one erased chain per pipeline
    /// thread (thread index in).
    pub fn new(
        config: PipelineConfig,
        source: S,
        chains: impl FnMut(usize) -> Box<dyn RunnableChain> + Send + 'static,
        sink: SinkRuntime,
        budget: Arc<InflightBudget>,
    ) -> Self {
        PipelineRuntime {
            config,
            source,
            chains: Box::new(chains),
            sink,
            budget,
            stop: Arc::new(StopRequest::new(Arc::new(SystemClock))),
            shutdown: Arc::new(AtomicBool::new(false)),
            stop_clock: Arc::new(SystemClock),
            clock: Arc::new(SystemClock),
            chase_window: None,
            options: RuntimeOptions::default(),
            io: None,
            feed: None,
            io_busy_interval: IO_BUSY_INTERVAL,
            io_busy_window: IoBusySampler::WINDOW,
        }
    }

    /// Sample the I/O runtime's busy ratio every `interval`, averaged over
    /// the trailing `window`.
    #[cfg(test)]
    #[must_use]
    pub(crate) fn with_io_busy_sampling(mut self, interval: Duration, window: Duration) -> Self {
        self.io_busy_interval = interval;
        self.io_busy_window = window;
        self
    }

    /// Override runtime options (returns `self` for chaining).
    #[must_use]
    pub fn with_options(mut self, options: RuntimeOptions) -> Self {
        self.options = options;
        self
    }

    /// Use a caller-owned tokio runtime as the I/O runtime instead of
    /// building one inside [`run`](Self::run). For assemblies whose
    /// connectors needed a handle before the runtime existed (sink workers
    /// spawned at construction, schema-registry fetchers, async pre-flight
    /// validation). `run` shuts it down on exit as it does the internally
    /// built one; connector tasks spawned on it earlier keep running until
    /// then. Without this, assemblies end up running a second runtime,
    /// doubling `pipeline.io_threads`.
    ///
    /// The runtime must be multi-thread, built with
    /// [`tokio::runtime::Builder::new_multi_thread`]. `run` waits on the
    /// controller outside `block_on`, and a current-thread runtime runs no
    /// spawned task then, so sink workers, the admin server and upkeep would
    /// stall while records flow. `run` returns [`StartError::Config`] for a
    /// current-thread runtime.
    #[must_use]
    pub fn with_io_runtime(mut self, io: tokio::runtime::Runtime) -> Self {
        self.io = Some(io);
        self
    }

    /// Hand the source a feed built in code, for the config's `feed:`
    /// section to name; see [`Pipeline::feed`](super::Pipeline::feed) for
    /// the rules `run` checks.
    #[must_use]
    pub fn with_feed(mut self, feed: Arc<dyn Feed>) -> Self {
        self.feed = Some(feed);
        self
    }

    /// A handle that triggers a graceful drain.
    #[must_use]
    pub fn shutdown_handle(&self) -> ShutdownHandle {
        ShutdownHandle(Arc::clone(&self.stop))
    }

    /// Read a graceful stop's reserve and deadlines, the revocation drain
    /// deadline, the startup-error grace and the sink's drain budget on
    /// `clock`, so a test moves them by hand.
    #[cfg(any(test, feature = "testing"))]
    #[doc(hidden)]
    #[must_use]
    pub fn with_stop_clock(mut self, clock: Arc<dyn Clock + Send + Sync>) -> Self {
        *self.stop.clock.lock().expect("stop request lock") = Arc::clone(&clock);
        self.stop_clock = clock;
        self
    }

    /// Read the commit tick, the fast-commit window, the stall check, the
    /// idle flush, the health probe, the feed's call timings and its
    /// dead-letter reads on `clock`, so a test moves them by hand. Every drain
    /// deadline is computed on the stop clock set by `with_stop_clock`, and the
    /// sink pool counts its share down on its own tokio clock.
    #[cfg(any(test, feature = "testing"))]
    #[doc(hidden)]
    #[must_use]
    pub fn with_clock(
        mut self,
        clock: Arc<dyn Clock + Send + Sync + std::panic::RefUnwindSafe>,
    ) -> Self {
        self.clock = clock;
        self
    }

    /// Close each `CommitReady` partition's tightened window after `window`,
    /// so a test reaches the end of a window without a clock.
    #[cfg(any(test, feature = "testing"))]
    #[doc(hidden)]
    #[must_use]
    pub fn with_chase_window(mut self, window: Duration) -> Self {
        self.chase_window = Some(window);
        self
    }

    /// Effective pipeline thread count: the config override, else
    /// [`default_thread_count`] over `available_parallelism`.
    /// `available_parallelism` respects cgroup CPU quotas, so Kubernetes
    /// limits size this correctly. Pods without limits see the node's cores;
    /// set `pipeline.threads` explicitly there.
    fn thread_count(&self) -> usize {
        self.config.pipeline.threads.unwrap_or_else(|| {
            let cores = std::thread::available_parallelism().map_or(2, usize::from);
            default_thread_count(cores, self.config.pipeline.io_threads)
        })
    }

    /// Run the pipeline to completion (blocking). Returns when the
    /// pipeline drained after a shutdown trigger/signal or failed.
    pub fn run(mut self) -> Result<ExitReport, StartError> {
        let threads = self.thread_count();
        if threads == 0 || self.config.pipeline.io_threads == 0 {
            return Err(StartError::Config("thread counts must be non-zero".into()));
        }
        if self.io.as_ref().is_some_and(|io| {
            io.handle().runtime_flavor() == tokio::runtime::RuntimeFlavor::CurrentThread
        }) {
            return Err(StartError::Config(
                "the I/O runtime must be multi-thread: a current-thread runtime runs no \
                 spawned task while `run` waits on the controller"
                    .into(),
            ));
        }
        self.config
            .reject_location_on_memory()
            .map_err(|e| StartError::Config(e.to_string()))?;
        self.config
            .reject_bad_settling_window()
            .map_err(|e| StartError::Config(e.to_string()))?;
        if let (Some(section), Some(feed)) = (&self.config.feed, &self.feed) {
            let named = section.backend().type_tag();
            if feed.backend() != named {
                return Err(StartError::Config(format!(
                    "the feed passed to `Pipeline::feed` is `{}`, but `feed:` names `{named}`",
                    feed.backend()
                )));
            }
        }
        let mut resolved_feed = None;
        match (&self.config.feed, self.feed.take()) {
            (Some(section), given) => {
                let (feed, origin) = match given {
                    Some(feed) => (feed, FeedOrigin::Given),
                    None if section.backend().type_tag() == "memory" => {
                        let feed = MemoryFeed::from_section(section)
                            .map_err(|e| StartError::Config(e.to_string()))?;
                        (Arc::new(feed) as Arc<dyn Feed>, FeedOrigin::Section)
                    }
                    None => {
                        return Err(StartError::Config(format!(
                            "the `{}` feed backend is built in code: build it from the `feed:` \
                             section and pass it to `Pipeline::feed`",
                            section.backend().type_tag()
                        )));
                    }
                };
                let drain_timeout = self.config.checkpoint.drain_timeout;
                if feed.hands_back() && drain_timeout <= STOP_RESERVE {
                    return Err(StartError::Config(format!(
                        "checkpoint.drain_timeout must be more than {STOP_RESERVE:?}, the stop \
                         reserve of the `{}` feed; found {drain_timeout:?}",
                        feed.backend()
                    )));
                }
                resolved_feed = Some((section, feed, origin));
            }
            (None, Some(_)) => {
                return Err(StartError::Config(
                    "a feed was given with `Pipeline::feed`, but the file has no `feed:` section \
                     naming it"
                        .into(),
                ));
            }
            (None, None) => {}
        }
        let pipeline_name = self.config.pipeline.name.clone();

        // The source's declared component_type feeds both its stage-metric
        // labels and the namespace of the custom-metrics Meter it receives at
        // `open`. Captured here, before `self.source` moves into the controller.
        let source_ct = self.source.component_type().to_string();
        let source_meter = Meter::for_component(
            &source_ct,
            MetricRole::Source,
            pipeline_name.clone(),
            "source",
        );

        // Observability first: everything after this records metrics.
        let handle = install_or_reuse(&metrics_settings(&self.config))?;

        let runtime_labels = ComponentLabels::new(pipeline_name.clone(), "runtime", "pipeline");
        // Fallible: a second live pipeline of this name in one process would
        // publish state and lifecycle gauges over this one's (see "Series
        // ownership" in `docs/METRICS.md`). Refuse to start instead.
        let pipeline_metrics = PipelineMetrics::try_new(&runtime_labels, &self.options.version)
            .map_err(|e| StartError::Metrics(e.to_string()))?;
        pipeline_metrics.set_state(PipelineState::Starting);
        pipeline_metrics.set_threads(threads);

        // The controller's handle sets, resolved here rather than at the
        // controller's own construction below. Both claim series, and a
        // duplicate must fail before any driver thread exists. Past that
        // point a bare `?` would leak the threads instead of stopping them.
        let checkpoint_metrics = CheckpointMetrics::try_new(
            &ComponentLabels::new(pipeline_name.clone(), "checkpoint", "checkpoint"),
            self.config.metrics.per_partition_detail,
        )
        .map_err(|e| StartError::Metrics(e.to_string()))?;
        // The owner of the budget gauge. One series per pipeline; the
        // drivers' `BackpressureMetrics` carry the pause series per thread.
        let budget_metrics = InflightBudgetMetrics::try_new(&runtime_labels)
            .map_err(|e| StartError::Metrics(e.to_string()))?;
        // The owner of the I/O busy gauge, which this thread samples while it
        // waits on the controller below.
        let io_busy_metrics = IoRuntimeMetrics::try_new(&runtime_labels)
            .map_err(|e| StartError::Metrics(e.to_string()))?;
        // The owner of the source series. It publishes active lanes, and its
        // clone goes to the source, which is the only thing that can measure
        // lag. The per-thread instances are shadows of this one.
        let controller_source_metrics = Arc::new(
            SourceMetrics::try_new(&ComponentLabels::new(
                pipeline_name.clone(),
                "source",
                source_ct.clone(),
            ))
            .map_err(|e| StartError::Metrics(e.to_string()))?,
        );

        // The feed's handles, and those its backend resolves from its `Meter`,
        // bind to the recorder installed above, so the feed is wrapped and
        // attached after it.
        let kept_feed = match resolved_feed {
            Some((section, feed, origin)) => {
                let metrics = FeedMetrics::try_new(&ComponentLabels::new(
                    pipeline_name.clone(),
                    "feed",
                    feed.backend().to_owned(),
                ))
                .map_err(|e| StartError::Metrics(e.to_string()))?;
                let feed = Arc::new(MeteredFeed::new(feed, metrics, Arc::clone(&self.clock)));
                feed.attach_metrics(
                    FeedMetricsCtx::new(Meter::for_component(
                        feed.backend(),
                        MetricRole::Feed,
                        pipeline_name.clone(),
                        "feed",
                    ))
                    .with_partition_detail(self.config.metrics.per_partition_detail),
                );
                self.source
                    .configure_feed(section, Arc::clone(&feed) as Arc<dyn Feed>, origin)
                    .map_err(|e| StartError::Config(e.to_string()))?;
                Some(feed)
            }
            None => None,
        };

        let health = HealthState::with_clock(
            threads,
            HealthThresholds::default(),
            Arc::clone(&self.clock),
        );

        // I/O runtime: sink workers (spawned by the caller-built SinkPool
        // onto this runtime via its own handle), admin server, upkeep,
        // signals, the feed's tasks. A caller-owned runtime
        // (`with_io_runtime`) is adopted instead of built; either way this
        // function shuts it down, after the source has dropped once the
        // controller has started.
        let io = match self.io.take() {
            Some(io) => io,
            None => tokio::runtime::Builder::new_multi_thread()
                .worker_threads(self.config.pipeline.io_threads)
                .thread_name("spate-io")
                .enable_all()
                .build()?,
        };

        // Admin bind, upkeep, and the controller thread are started *after*
        // the driver threads (below) so a failure in any of them can stop
        // the already-running drivers instead of leaking them.

        // The io tasks' handles live to the end of `run` and drop just before
        // the io runtime shuts down, since dropping one aborts its task.
        let signal_task = self.options.handle_signals.then(|| {
            let stop = Arc::clone(&self.stop);
            task::spawn_on(
                "shutdown-signal",
                async move {
                    wait_for_signal().await;
                    tracing::info!("shutdown signal received; draining");
                    stop.trigger();
                },
                io.handle(),
            )
        });

        // Sink readiness: probe at startup and periodically (tighter while
        // failing), driving the sinks-connected half of `/readyz`. No probe
        // hook means nothing to check, so report connected.
        let probe_task = match self.sink.probe.take() {
            Some(probe) => Some(task::spawn_on(
                "sink-probe",
                probe_sinks(probe, Arc::clone(&health)),
                io.handle(),
            )),
            None => {
                health.set_sinks_connected(true);
                None
            }
        };

        // Wiring.
        let (events_tx, events_rx) = crossbeam_channel::unbounded::<DriverEvent>();
        let (to_main_tx, to_main_rx) = crossbeam_channel::unbounded::<ControllerSignal>();
        let (sink_drained_tx, sink_drained_rx) = crossbeam_channel::unbounded::<()>();
        let mut checkpointer = Checkpointer::new();
        checkpointer.set_clock(Arc::clone(&self.clock));

        let bp_params = BackpressureParams::from_budget(
            usize::try_from(self.config.backpressure.max_inflight_bytes.as_u64())
                .unwrap_or(usize::MAX),
            self.config.backpressure.high_ratio,
            self.config.backpressure.low_ratio,
            self.config.backpressure.min_pause,
        );

        // Compact pinning: thread i on core i, low cores first, leaving the
        // remaining cores for the I/O runtime and librdkafka's threads.
        // Note for Kubernetes: exclusive cores require the kubelet static
        // CPU manager with Guaranteed QoS and integer CPU requests;
        // otherwise pinning only sets affinity within the shared cpuset.
        let core_ids: Vec<Option<core_affinity::CoreId>> =
            if self.config.pipeline.pinning == PinningMode::Compact {
                let mut ids = core_affinity::get_core_ids().unwrap_or_default();
                ids.sort_by_key(|c| c.id);
                if ids.len() < threads {
                    tracing::warn!(
                        cores = ids.len(),
                        threads,
                        "fewer cores than pipeline threads; surplus threads run unpinned"
                    );
                }
                (0..threads).map(|i| ids.get(i).copied()).collect()
            } else {
                vec![None; threads]
            };

        // Short grace for cleanup-time driver stops (startup errors, a
        // controller panic): the same budget the drain barrier uses.
        let drain_timeout = self.config.checkpoint.drain_timeout;

        let mut control_txs = Vec::with_capacity(threads);
        let mut driver_handles = Vec::with_capacity(threads);
        for i in 0..threads {
            let (control_tx, control_rx) = crossbeam_channel::unbounded::<ThreadControl<S::Lane>>();
            control_txs.push(control_tx);
            let ctx = DriverContext {
                params: DriverParams {
                    thread: i,
                    max_records: self.options.max_records,
                    poll_timeout: self.options.poll_timeout,
                    idle_flush: self.options.idle_flush,
                    blocked_retry: self.options.blocked_retry,
                    queue_low_ratio: self.config.backpressure.low_ratio,
                    max_pending_batches: self.config.checkpoint.max_pending_batches,
                },
                control: control_rx,
                events: events_tx.clone(),
                chain: (self.chains)(i),
                bp: WatermarkController::new(bp_params),
                budget: Arc::clone(&self.budget),
                queues: self.sink.queues.clone(),
                health: Arc::clone(&health),
                // Each driver's backpressure series is its own (`driver-{i}`),
                // so every thread owns what it publishes.
                bp_metrics: BackpressureMetrics::new(&ComponentLabels::new(
                    pipeline_name.clone(),
                    format!("driver-{i}"),
                    "driver",
                )),
                // A shadow. Every thread counts the records and polls it
                // performed (those sum), but the source *gauges* (lag and
                // active lanes) belong to the controller's instance below,
                // which is the only one that sees the assignment and the
                // only one the source itself is handed. A claiming
                // constructor here would take the series hostage from it.
                source_metrics: SourceMetrics::shadow(&ComponentLabels::new(
                    pipeline_name.clone(),
                    "source",
                    source_ct.clone(),
                )),
                shutdown: Arc::clone(&self.shutdown),
                clock: Arc::clone(&self.clock) as _,
                stop_clock: Arc::clone(&self.stop_clock),
            };
            let core = core_ids.get(i).copied().flatten();
            #[expect(
                clippy::disallowed_methods,
                reason = "`run` joins every driver thread, on each early return too"
            )]
            let spawned = std::thread::Builder::new()
                .name(format!("spate-pipeline-{i}"))
                .spawn(move || {
                    if let Some(core) = core
                        && !core_affinity::set_for_current(core)
                    {
                        tracing::warn!(core = core.id, "failed to pin pipeline thread");
                    }
                    run_driver(ctx)
                });
            match spawned {
                Ok(handle) => driver_handles.push(handle),
                Err(e) => {
                    // Stop the drivers already spawned before bailing out.
                    stop_drivers(
                        &self.shutdown,
                        &control_txs,
                        driver_handles,
                        drain_timeout,
                        &*self.stop_clock,
                    );
                    return Err(StartError::Io(e));
                }
            }
        }

        // Factories capture ShardQueues clones (their terminals need them),
        // and the sink only drains once every queue clone is gone. Holding
        // the factory through the drain would deadlock shutdown.
        drop(self.chains);

        // A cloned set of driver control senders kept by main, so it can stop
        // the drivers itself if a later startup step fails or the controller
        // thread dies (the originals are moved into the controller below).
        let control_txs_for_stop = control_txs.clone();

        // Admin bind now that the drivers are live: a bind failure (e.g. the
        // address is taken) stops them instead of leaking them.
        let (admin_stop_tx, admin_task) = match self.config.admin.listen {
            Some(addr) => {
                let render = handle.exports().then(|| handle.render_fn());
                let admin = match io.block_on(AdminServer::bind(addr, render, Arc::clone(&health)))
                {
                    Ok(admin) => admin,
                    Err(source) => {
                        stop_drivers(
                            &self.shutdown,
                            &control_txs_for_stop,
                            driver_handles,
                            drain_timeout,
                            &*self.stop_clock,
                        );
                        return Err(StartError::AdminBind { addr, source });
                    }
                };
                let (admin_stop_tx, admin_stop_rx) = watch::channel(false);
                tracing::info!(addr = %admin.local_addr(), "admin server listening");
                let admin_task =
                    task::spawn_on("admin-server", admin.run(admin_stop_rx), io.handle());
                (Some(admin_stop_tx), Some(admin_task))
            }
            None => {
                if handle.exports() {
                    tracing::warn!(
                        "an exporter is installed but `admin.listen` is `none`, so \
                         nothing serves /metrics; render the exposition through \
                         Pipeline::metrics(), or set `metrics.exporter: none`"
                    );
                }
                (None, None)
            }
        };
        let upkeep_task = {
            let _guard = io.enter();
            handle.spawn_upkeep(Duration::from_secs(5))
        };

        let controller_ctx = ControllerContext {
            source: self.source,
            checkpointer,
            control_txs,
            events_rx,
            to_main: to_main_tx,
            sink_drained_rx,
            stop: Arc::clone(&self.stop),
            shutdown: Arc::clone(&self.shutdown),
            stop_clock: Arc::clone(&self.stop_clock),
            clock: Arc::clone(&self.clock) as _,
            close_share: if kept_feed.as_ref().is_some_and(|f| f.hands_back()) {
                CLOSE_SHARE
            } else {
                Duration::ZERO
            },
            health: Arc::clone(&health),
            commit_interval: self.config.checkpoint.interval,
            chase_window: self.chase_window.unwrap_or(self.config.checkpoint.interval),
            drain_timeout: self.config.checkpoint.drain_timeout,
            event_poll_timeout: self.options.event_poll_timeout,
            stalled_fail_after: self.config.checkpoint.stalled_fail_after,
            sink_failures: self.sink.failures.clone(),
            checkpoint_metrics,
            budget: Arc::clone(&self.budget),
            budget_metrics,
            source_metrics: controller_source_metrics,
            source_meter,
            per_partition_detail: self.config.metrics.per_partition_detail,
            pipeline_metrics,
        };
        #[expect(
            clippy::disallowed_methods,
            reason = "`run` joins the controller after the driver threads, and on each early return"
        )]
        let controller_handle = match std::thread::Builder::new()
            .name("spate-controller".into())
            .spawn(move || run_controller(controller_ctx))
        {
            Ok(handle) => handle,
            Err(e) => {
                stop_drivers(
                    &self.shutdown,
                    &control_txs_for_stop,
                    driver_handles,
                    drain_timeout,
                    &*self.stop_clock,
                );
                return Err(StartError::Io(e));
            }
        };

        // Main: wait for the controller's choreography.
        let mut sink_drain = None;
        let mut driver_panic: Option<FatalErrorReport> = None;
        let sink_runtime = self.sink;
        let mut drain_fn = Some(sink_runtime.drain);
        drop(sink_runtime.queues);

        // Sampled here, off the I/O runtime, so the reading is neither
        // delayed by a saturated worker nor counted as its busy time. The
        // drain's `block_on` pauses it.
        let io_runtime_metrics = io.handle().metrics();
        let mut io_busy = IoBusySampler::new(
            self.io_busy_interval,
            self.io_busy_window,
            // Real time, since the busy time it is divided into is real.
            SystemClock.now(),
            metrics::io_busy_reading(&io_runtime_metrics).0,
        );
        let mut dead_letters = self
            .options
            .dead_letter_read_interval
            .map(|interval| DeadLetterSampler::new(interval, self.clock.now()));

        let (mut state, final_watermarks, close_by) = loop {
            let (busy, workers) = metrics::io_busy_reading(&io_runtime_metrics);
            sample_io_busy(
                &mut io_busy,
                &io_busy_metrics,
                SystemClock.now(),
                busy,
                workers,
                &pipeline_name,
            );
            if let (Some(feed), Some(sampler)) = (&kept_feed, &mut dead_letters) {
                sampler.sample(feed, self.clock.now());
            }
            match to_main_rx.recv_timeout(Duration::from_millis(100)) {
                Ok(ControllerSignal::LanesDrained { sink_deadline }) => {
                    for (i, h) in driver_handles.drain(..).enumerate() {
                        if h.join().is_err() {
                            driver_panic.get_or_insert(FatalErrorReport {
                                component: format!("driver-{i}"),
                                reason: "pipeline thread panicked outside the batch guard".into(),
                            });
                        }
                    }
                    if let Some(drain) = drain_fn.take() {
                        let budget = sink_deadline.saturating_duration_since(self.stop_clock.now());
                        sink_drain = Some(io.block_on(drain(budget)));
                    }
                    let _ = sink_drained_tx.send(());
                }
                Ok(ControllerSignal::Finished(report)) => {
                    break (report.state, report.final_watermarks, Some(report.close_by));
                }
                Err(crossbeam_channel::RecvTimeoutError::Timeout)
                    if !controller_handle.is_finished() =>
                {
                    // Controller still working; keep waiting on the 100ms tick.
                }
                Err(_) => {
                    // The controller thread ended without a Finished report
                    // (a timeout with a finished handle, or the signal
                    // channel disconnected), so it panicked. It never told
                    // the drivers to stop and never set the shutdown flag, so
                    // an untimed join here would wedge forever. Stop them
                    // ourselves, drain the sink, and fail the run.
                    stop_drivers(
                        &self.shutdown,
                        &control_txs_for_stop,
                        std::mem::take(&mut driver_handles),
                        drain_timeout,
                        &*self.stop_clock,
                    );
                    if let Some(drain) = drain_fn.take() {
                        sink_drain = Some(io.block_on(drain(drain_timeout)));
                    }
                    break (
                        ExitState::Failed(FatalErrorReport {
                            component: "controller".into(),
                            reason: "controller thread panicked".into(),
                        }),
                        Vec::new(),
                        None,
                    );
                }
            }
        };
        // Drivers are already joined on the drain path; on the
        // controller-died path make a best effort not to leak them.
        for h in driver_handles {
            let _ = h.join();
        }
        // A driver that panicked outside the batch guard is a bug worth
        // failing the run over, even if the drain otherwise completed.
        if let (ExitState::Completed, Some(report)) = (&state, driver_panic) {
            state = ExitState::Failed(report);
        }

        // After the final commit. This handle keeps the feed alive for its
        // close, whether or not the source has dropped its own.
        if let (Some(feed), Some(close_by)) = (&kept_feed, close_by) {
            close_feed(
                &io,
                feed.as_ref(),
                close_by.saturating_duration_since(self.stop_clock.now()),
            );
        }
        if let Some(tx) = &admin_stop_tx {
            let _ = tx.send(true);
        }
        // The controller drops the source on its way out, and a source's
        // `Drop` may still need the io runtime (the S3 source aborts its lister
        // there).
        let _ = controller_handle.join();
        drop((signal_task, probe_task, admin_task, upkeep_task));
        io.shutdown_timeout(Duration::from_secs(2));

        Ok(ExitReport {
            state,
            sink_drain,
            final_watermarks,
        })
    }
}

/// The pipeline thread count for `cores` when `pipeline.threads` is unset:
/// `cores` minus one core for every two I/O workers (rounded up) and one for
/// the controller, at least 1.
///
/// Row encoding runs on the pipeline threads, and so does compression when the
/// sink supplies a [`FrameCodec`](crate::sink::FrameCodec). The I/O workers
/// carry the socket writes and shard bookkeeping. While `cores` is at least
/// `ceil(io_threads / 2) + 2`, the runnable threads exceed the cores by
/// `floor(io_threads / 2)` when every thread is busy, and the I/O workers
/// share cores with pipeline threads. While `cores` is at least
/// `ceil(io_threads / 2) + 2`, compact pinning leaves the
/// `ceil(io_threads / 2) + 1` highest cores unpinned for the I/O runtime, the
/// controller and the source's own threads.
pub(super) fn default_thread_count(cores: usize, io_threads: usize) -> usize {
    cores.saturating_sub(io_threads.div_ceil(2) + 1).max(1)
}

/// Publish an I/O busy sample when one is due, and warn on the sample that
/// completes a saturated window.
///
/// `busy_total` is the I/O runtime's cumulative busy time over its
/// `io_workers` async workers.
pub(super) fn sample_io_busy(
    sampler: &mut IoBusySampler,
    gauge: &IoRuntimeMetrics,
    now: Instant,
    busy_total: Duration,
    io_workers: usize,
    pipeline: &str,
) {
    let Some(sample) = sampler.observe(now, busy_total, io_workers) else {
        return;
    };
    gauge.set_busy_ratio(sample.ratio);
    if sample.sustained {
        tracing::warn!(
            pipeline,
            io_workers,
            ratio = sample.ratio,
            "the I/O runtime's workers have been saturated for five minutes; if the \
             sink queues are filling, raise `pipeline.io_threads`"
        );
    }
}

/// Close `feed` and wait for it at most `within`. A close that fails or does
/// not settle in time leaves the units the feed held to its redelivery.
fn close_feed(io: &tokio::runtime::Runtime, feed: &dyn Feed, within: Duration) {
    let Some(confirmation) = feed.close(within) else {
        return;
    };
    // The timer is built inside `block_on`, where the runtime's reactor is.
    match io.block_on(async { tokio::time::timeout(within, confirmation).await }) {
        Ok(Ok(())) => {}
        Ok(Err(e)) => tracing::warn!(
            feed = feed.backend(),
            error = %e,
            "the feed did not close cleanly; the units it still held return through redelivery"
        ),
        Err(_) => tracing::warn!(
            feed = feed.backend(),
            ?within,
            "the feed did not close before the drain deadline; the units it still held return \
             through redelivery"
        ),
    }
}

/// Set the shutdown flag, tell every driver thread to stop within a bounded
/// drain barrier, and join them. Shared by the startup-error paths and the
/// controller-death path so an early failure or a controller panic never
/// leaves running pinned pipeline threads behind (`run` is a library API).
///
/// Joining bounds this. A driver observes the shutdown flag (abandoning any
/// blocked batch) and the `Shutdown` control message (flushing within
/// `grace`), then exits and drops its chain, closing the shard queues so the
/// sink can drain afterwards.
fn stop_drivers<L>(
    shutdown: &AtomicBool,
    control_txs: &[crossbeam_channel::Sender<ThreadControl<L>>],
    driver_handles: Vec<std::thread::JoinHandle<DriverExit>>,
    grace: Duration,
    stop_clock: &dyn Clock,
) {
    shutdown.store(true, Ordering::Relaxed);
    let deadline = stop_clock.now() + grace;
    let barrier = DrainBarrier::new(control_txs.len());
    for tx in control_txs {
        let _ = tx.send(ThreadControl::Shutdown {
            barrier: barrier.clone(),
            deadline,
        });
    }
    for handle in driver_handles {
        let _ = handle.join();
    }
}

/// Install the exporter, degrading gracefully when a foreign recorder
/// already owns the process. The pipeline keeps running against the
/// existing recorder with a detached (empty-rendering) handle for the
/// admin server. Shared by the runtime and the pipeline builder.
pub(crate) fn install_or_reuse(settings: &MetricsSettings) -> Result<MetricsHandle, StartError> {
    match metrics::install(settings) {
        Ok(h) => Ok(h),
        Err(metrics::MetricsError::AlreadyInstalled) => {
            tracing::warn!(
                "a metrics recorder is already installed; continuing \
                 with the existing one and a detached render handle"
            );
            metrics::install(&MetricsSettings {
                exporter: Exporter::None,
                ..settings.clone()
            })
            .map_err(|e| StartError::Metrics(e.to_string()))
        }
        Err(e) => Err(StartError::Metrics(e.to_string())),
    }
}

/// The [`MetricsSettings`] a pipeline configuration maps to.
///
/// Assemblies that pre-register metric handles (sink shard metrics, custom
/// metrics) should call
/// [`metrics::install`](crate::metrics::install)`(&metrics_settings(&config))`
/// **before** constructing them; the runtime's own install then reuses the
/// exporter. Handles built before any install bind to the no-op recorder
/// and render nothing.
#[must_use]
pub fn metrics_settings(config: &PipelineConfig) -> MetricsSettings {
    MetricsSettings {
        exporter: match config.metrics.exporter {
            MetricsExporter::Prometheus => Exporter::Prometheus,
            MetricsExporter::None => Exporter::None,
        },
        per_partition_detail: config.metrics.per_partition_detail,
        e2e_basis: match config.metrics.e2e_basis {
            crate::config::E2eBasis::Ingest => E2eBasis::Ingest,
            crate::config::E2eBasis::Event => E2eBasis::Event,
        },
    }
}

/// Probe the sinks at once and then periodically, keeping the
/// sinks-connected half of `/readyz` current. Runs until aborted.
///
/// A panic in `probe` marks the sinks not connected, so readiness does not
/// keep the last probe's answer once nothing updates it.
pub(crate) async fn probe_sinks(probe: SinkProbeFn, health: Arc<HealthState>) {
    struct NotConnectedOnPanic<'a>(&'a HealthState);
    impl Drop for NotConnectedOnPanic<'_> {
        fn drop(&mut self) {
            if std::thread::panicking() {
                self.0.set_sinks_connected(false);
            }
        }
    }
    let _guard = NotConnectedOnPanic(&health);
    loop {
        let connected = match probe().await {
            Ok(()) => true,
            Err(e) => {
                tracing::warn!(error = %e, "sink probe failed");
                false
            }
        };
        health.set_sinks_connected(connected);
        let recheck = if connected {
            Duration::from_secs(30)
        } else {
            Duration::from_secs(5)
        };
        tokio::time::sleep(recheck).await;
    }
}

async fn wait_for_signal() {
    #[cfg(unix)]
    {
        let mut term = match signal(SignalKind::terminate()) {
            Ok(s) => s,
            Err(e) => {
                tracing::error!(error = %e, "failed to install SIGTERM handler");
                std::future::pending::<()>().await;
                return;
            }
        };
        tokio::select! {
            _ = term.recv() => {}
            r = tokio::signal::ctrl_c() => {
                if let Err(e) = r {
                    tracing::error!(error = %e, "ctrl_c handler failed");
                    std::future::pending::<()>().await;
                }
            }
        }
    }
    #[cfg(not(unix))]
    {
        if let Err(e) = tokio::signal::ctrl_c().await {
            tracing::error!(error = %e, "ctrl_c handler failed");
            std::future::pending::<()>().await;
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::pipeline::fakes::{FakeSource, test_config};

    /// A runtime samples the I/O busy ratio every 5 s over the sampler's
    /// 5-minute window unless a test overrides both.
    #[test]
    fn the_io_busy_sampling_defaults_to_five_seconds_over_five_minutes() {
        let (source, _log, _script) = FakeSource::new();
        let (queues, _receivers) = crate::sink::shard_queues(1, 1);
        let sink = SinkRuntime {
            queues: vec![queues],
            drain: Box::new(|_budget| Box::pin(async { crate::sink::DrainReport::default() })),
            probe: None,
            failures: crate::sink::SinkFailures::new(),
        };
        let runtime = PipelineRuntime::new(
            test_config(1),
            source,
            |_thread| unreachable!("the runtime is never run"),
            sink,
            Arc::new(InflightBudget::new()),
        );

        assert_eq!(runtime.io_busy_interval, Duration::from_secs(5));
        assert_eq!(runtime.io_busy_window, IoBusySampler::WINDOW);
        assert_eq!(IoBusySampler::WINDOW, Duration::from_secs(300));
    }
}