1use super::controller::{ControllerContext, ControllerSignal, run_controller};
4use super::driver::{DriverContext, DriverExit, DriverParams, run_driver};
5use super::{DriverEvent, ExitReport, ExitState, FatalErrorReport, SinkRuntime, ThreadControl};
6use crate::admin::{AdminServer, HealthState, HealthThresholds};
7use crate::backpressure::{BackpressureParams, InflightBudget, WatermarkController};
8use crate::checkpoint::Checkpointer;
9use crate::config::{MetricsExporter, PinningMode, PipelineConfig};
10use crate::metrics::{
11 self, BackpressureMetrics, CheckpointMetrics, ComponentLabels, E2eBasis, Exporter,
12 InflightBudgetMetrics, Meter, MetricRole, MetricsHandle, MetricsSettings, PipelineMetrics,
13 PipelineState, SourceMetrics,
14};
15use crate::ops::RunnableChain;
16use crate::source::{DrainBarrier, Source};
17use std::net::SocketAddr;
18use std::sync::Arc;
19use std::sync::atomic::{AtomicBool, Ordering};
20use std::time::{Duration, Instant};
21
22#[derive(Clone, Debug)]
25pub struct RuntimeOptions {
26 pub handle_signals: bool,
29 pub max_records: usize,
31 pub poll_timeout: Duration,
33 pub idle_flush: Duration,
35 pub blocked_retry: Duration,
37 pub event_poll_timeout: Duration,
39 pub version: String,
41}
42
43impl Default for RuntimeOptions {
44 fn default() -> Self {
45 RuntimeOptions {
46 handle_signals: true,
47 max_records: 512,
48 poll_timeout: Duration::from_millis(10),
49 idle_flush: Duration::from_millis(100),
50 blocked_retry: Duration::from_millis(2),
51 event_poll_timeout: Duration::from_millis(50),
52 version: env!("CARGO_PKG_VERSION").to_string(),
53 }
54 }
55}
56
57#[derive(Clone, Debug)]
59pub struct ShutdownHandle(Arc<AtomicBool>);
60
61impl ShutdownHandle {
62 pub fn trigger(&self) {
64 self.0.store(true, Ordering::Relaxed);
65 }
66}
67
68#[derive(Debug, thiserror::Error)]
70#[non_exhaustive]
71pub enum StartError {
72 #[error("invalid runtime configuration: {0}")]
74 Config(String),
75 #[error("metrics: {0}")]
77 Metrics(String),
78 #[error("admin.listen: cannot bind {addr}: {source}")]
81 AdminBind {
82 addr: SocketAddr,
84 source: std::io::Error,
86 },
87 #[error("io: {0}")]
89 Io(#[from] std::io::Error),
90}
91
92pub struct PipelineRuntime<S: Source> {
98 config: PipelineConfig,
99 source: S,
100 chains: Box<dyn FnMut(usize) -> Box<dyn RunnableChain> + Send>,
101 sink: SinkRuntime,
102 budget: Arc<InflightBudget>,
103 shutdown: Arc<AtomicBool>,
104 options: RuntimeOptions,
105 io: Option<tokio::runtime::Runtime>,
106}
107
108impl<S: Source> std::fmt::Debug for PipelineRuntime<S> {
109 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
110 f.debug_struct("PipelineRuntime")
111 .field("pipeline", &self.config.pipeline.name)
112 .field("options", &self.options)
113 .finish_non_exhaustive()
114 }
115}
116
117impl<S: Source + 'static> PipelineRuntime<S> {
118 pub fn new(
121 config: PipelineConfig,
122 source: S,
123 chains: impl FnMut(usize) -> Box<dyn RunnableChain> + Send + 'static,
124 sink: SinkRuntime,
125 budget: Arc<InflightBudget>,
126 ) -> Self {
127 PipelineRuntime {
128 config,
129 source,
130 chains: Box::new(chains),
131 sink,
132 budget,
133 shutdown: Arc::new(AtomicBool::new(false)),
134 options: RuntimeOptions::default(),
135 io: None,
136 }
137 }
138
139 #[must_use]
141 pub fn with_options(mut self, options: RuntimeOptions) -> Self {
142 self.options = options;
143 self
144 }
145
146 #[must_use]
155 pub fn with_io_runtime(mut self, io: tokio::runtime::Runtime) -> Self {
156 self.io = Some(io);
157 self
158 }
159
160 #[must_use]
162 pub fn shutdown_handle(&self) -> ShutdownHandle {
163 ShutdownHandle(Arc::clone(&self.shutdown))
164 }
165
166 fn thread_count(&self) -> usize {
172 self.config.pipeline.threads.unwrap_or_else(|| {
173 let cores = std::thread::available_parallelism().map_or(2, usize::from);
174 cores
175 .saturating_sub(self.config.pipeline.io_threads + 1)
176 .max(1)
177 })
178 }
179
180 pub fn run(mut self) -> Result<ExitReport, StartError> {
183 let threads = self.thread_count();
184 if threads == 0 || self.config.pipeline.io_threads == 0 {
185 return Err(StartError::Config("thread counts must be non-zero".into()));
186 }
187 if let Some(section) = &self.config.coordination {
188 tracing::warn!(
189 "the `coordination:` section is deprecated, and v0.4.0 removes it; a work feed \
190 replaces split coordination (https://github.com/spate-etl/spate/issues/1088)"
191 );
192 #[expect(deprecated, reason = "v0.4.0 removes this with spate-coordination")]
193 self.source
194 .configure_coordination(section)
195 .map_err(|e| StartError::Config(e.to_string()))?;
196 }
197 let pipeline_name = self.config.pipeline.name.clone();
198
199 let source_ct = self.source.component_type().to_string();
203 let source_meter = Meter::for_component(
204 &source_ct,
205 MetricRole::Source,
206 pipeline_name.clone(),
207 "source",
208 );
209
210 let handle = install_or_reuse(&metrics_settings(&self.config))?;
212
213 let runtime_labels = ComponentLabels::new(pipeline_name.clone(), "runtime", "pipeline");
214 let pipeline_metrics = PipelineMetrics::try_new(&runtime_labels, &self.options.version)
218 .map_err(|e| StartError::Metrics(e.to_string()))?;
219 pipeline_metrics.set_state(PipelineState::Starting);
220 pipeline_metrics.set_threads(threads);
221
222 let checkpoint_metrics = CheckpointMetrics::try_new(
227 &ComponentLabels::new(pipeline_name.clone(), "checkpoint", "checkpoint"),
228 self.config.metrics.per_partition_detail,
229 )
230 .map_err(|e| StartError::Metrics(e.to_string()))?;
231 let budget_metrics = InflightBudgetMetrics::try_new(&runtime_labels)
234 .map_err(|e| StartError::Metrics(e.to_string()))?;
235 let controller_source_metrics = Arc::new(
239 SourceMetrics::try_new(&ComponentLabels::new(
240 pipeline_name.clone(),
241 "source",
242 source_ct.clone(),
243 ))
244 .map_err(|e| StartError::Metrics(e.to_string()))?,
245 );
246
247 let health = HealthState::new(threads, HealthThresholds::default());
248
249 let io = match self.io.take() {
256 Some(io) => io,
257 None => tokio::runtime::Builder::new_multi_thread()
258 .worker_threads(self.config.pipeline.io_threads)
259 .thread_name("spate-io")
260 .enable_all()
261 .build()?,
262 };
263
264 if self.options.handle_signals {
269 let shutdown = Arc::clone(&self.shutdown);
270 io.spawn(async move {
271 wait_for_signal().await;
272 tracing::info!("shutdown signal received; draining");
273 shutdown.store(true, Ordering::Relaxed);
274 });
275 }
276
277 match self.sink.probe.take() {
281 Some(probe) => {
282 let health_probe = Arc::clone(&health);
283 io.spawn(async move {
284 loop {
285 let connected = match probe().await {
286 Ok(()) => true,
287 Err(e) => {
288 tracing::warn!(error = %e, "sink probe failed");
289 false
290 }
291 };
292 health_probe.set_sinks_connected(connected);
293 let recheck = if connected {
294 Duration::from_secs(30)
295 } else {
296 Duration::from_secs(5)
297 };
298 tokio::time::sleep(recheck).await;
299 }
300 });
301 }
302 None => health.set_sinks_connected(true),
303 }
304
305 let (events_tx, events_rx) = crossbeam_channel::unbounded::<DriverEvent>();
307 let (to_main_tx, to_main_rx) = crossbeam_channel::unbounded::<ControllerSignal>();
308 let (sink_drained_tx, sink_drained_rx) = crossbeam_channel::unbounded::<()>();
309 let checkpointer = Checkpointer::new();
310
311 let bp_params = BackpressureParams::from_budget(
312 usize::try_from(self.config.backpressure.max_inflight_bytes.as_u64())
313 .unwrap_or(usize::MAX),
314 self.config.backpressure.high_ratio,
315 self.config.backpressure.low_ratio,
316 self.config.backpressure.min_pause,
317 );
318
319 let core_ids: Vec<Option<core_affinity::CoreId>> =
325 if self.config.pipeline.pinning == PinningMode::Compact {
326 let mut ids = core_affinity::get_core_ids().unwrap_or_default();
327 ids.sort_by_key(|c| c.id);
328 if ids.len() < threads {
329 tracing::warn!(
330 cores = ids.len(),
331 threads,
332 "fewer cores than pipeline threads; surplus threads run unpinned"
333 );
334 }
335 (0..threads).map(|i| ids.get(i).copied()).collect()
336 } else {
337 vec![None; threads]
338 };
339
340 let drain_timeout = self.config.checkpoint.drain_timeout;
343
344 let mut control_txs = Vec::with_capacity(threads);
345 let mut driver_handles = Vec::with_capacity(threads);
346 for i in 0..threads {
347 let (control_tx, control_rx) = crossbeam_channel::unbounded::<ThreadControl<S::Lane>>();
348 control_txs.push(control_tx);
349 let ctx = DriverContext {
350 params: DriverParams {
351 thread: i,
352 max_records: self.options.max_records,
353 poll_timeout: self.options.poll_timeout,
354 idle_flush: self.options.idle_flush,
355 blocked_retry: self.options.blocked_retry,
356 queue_low_ratio: self.config.backpressure.low_ratio,
357 max_pending_batches: self.config.checkpoint.max_pending_batches,
358 },
359 control: control_rx,
360 events: events_tx.clone(),
361 chain: (self.chains)(i),
362 bp: WatermarkController::new(bp_params),
363 budget: Arc::clone(&self.budget),
364 queues: self.sink.queues.clone(),
365 health: Arc::clone(&health),
366 bp_metrics: BackpressureMetrics::new(&ComponentLabels::new(
369 pipeline_name.clone(),
370 format!("driver-{i}"),
371 "driver",
372 )),
373 source_metrics: SourceMetrics::shadow(&ComponentLabels::new(
380 pipeline_name.clone(),
381 "source",
382 source_ct.clone(),
383 )),
384 shutdown: Arc::clone(&self.shutdown),
385 };
386 let core = core_ids.get(i).copied().flatten();
387 let spawned = std::thread::Builder::new()
388 .name(format!("spate-pipeline-{i}"))
389 .spawn(move || {
390 if let Some(core) = core
391 && !core_affinity::set_for_current(core)
392 {
393 tracing::warn!(core = core.id, "failed to pin pipeline thread");
394 }
395 run_driver(ctx)
396 });
397 match spawned {
398 Ok(handle) => driver_handles.push(handle),
399 Err(e) => {
400 stop_drivers(&self.shutdown, &control_txs, driver_handles, drain_timeout);
402 return Err(StartError::Io(e));
403 }
404 }
405 }
406
407 drop(self.chains);
411
412 let control_txs_for_stop = control_txs.clone();
416
417 let admin_stop_tx = match self.config.admin.listen {
420 Some(addr) => {
421 let render = handle.exports().then(|| handle.render_fn());
422 let admin = match io.block_on(AdminServer::bind(addr, render, Arc::clone(&health)))
423 {
424 Ok(admin) => admin,
425 Err(source) => {
426 stop_drivers(
427 &self.shutdown,
428 &control_txs_for_stop,
429 driver_handles,
430 drain_timeout,
431 );
432 return Err(StartError::AdminBind { addr, source });
433 }
434 };
435 let (admin_stop_tx, admin_stop_rx) = tokio::sync::watch::channel(false);
436 tracing::info!(addr = %admin.local_addr(), "admin server listening");
437 io.spawn(admin.run(admin_stop_rx));
438 Some(admin_stop_tx)
439 }
440 None => {
441 if handle.exports() {
442 tracing::warn!(
443 "an exporter is installed but `admin.listen` is `none`, so \
444 nothing serves /metrics; render the exposition through \
445 Pipeline::metrics(), or set `metrics.exporter: none`"
446 );
447 }
448 None
449 }
450 };
451 {
452 let _guard = io.enter();
454 let _upkeep = handle.spawn_upkeep(Duration::from_secs(5));
455 }
456
457 let controller_ctx = ControllerContext {
458 source: self.source,
459 checkpointer,
460 control_txs,
461 events_rx,
462 to_main: to_main_tx,
463 sink_drained_rx,
464 shutdown: Arc::clone(&self.shutdown),
465 health: Arc::clone(&health),
466 commit_interval: self.config.checkpoint.interval,
467 drain_timeout: self.config.checkpoint.drain_timeout,
468 event_poll_timeout: self.options.event_poll_timeout,
469 stalled_fail_after: self.config.checkpoint.stalled_fail_after,
470 sink_failures: self.sink.failures.clone(),
471 checkpoint_metrics,
472 budget: Arc::clone(&self.budget),
473 budget_metrics,
474 source_metrics: controller_source_metrics,
475 source_meter,
476 per_partition_detail: self.config.metrics.per_partition_detail,
477 pipeline_metrics,
478 };
479 let controller_handle = match std::thread::Builder::new()
480 .name("spate-controller".into())
481 .spawn(move || run_controller(controller_ctx))
482 {
483 Ok(handle) => handle,
484 Err(e) => {
485 stop_drivers(
486 &self.shutdown,
487 &control_txs_for_stop,
488 driver_handles,
489 drain_timeout,
490 );
491 return Err(StartError::Io(e));
492 }
493 };
494
495 let mut sink_drain = None;
497 let mut driver_panic: Option<FatalErrorReport> = None;
498 let sink_runtime = self.sink;
499 let mut drain_fn = Some(sink_runtime.drain);
500 drop(sink_runtime.queues);
501
502 let (mut state, final_watermarks) = loop {
503 match to_main_rx.recv_timeout(Duration::from_millis(100)) {
504 Ok(ControllerSignal::LanesDrained { sink_deadline }) => {
505 for (i, h) in driver_handles.drain(..).enumerate() {
506 if h.join().is_err() {
507 driver_panic.get_or_insert(FatalErrorReport {
508 component: format!("driver-{i}"),
509 reason: "pipeline thread panicked outside the batch guard".into(),
510 });
511 }
512 }
513 if let Some(drain) = drain_fn.take() {
514 let budget = sink_deadline.saturating_duration_since(Instant::now());
515 sink_drain = Some(io.block_on(drain(budget)));
516 }
517 let _ = sink_drained_tx.send(());
518 }
519 Ok(ControllerSignal::Finished(report)) => {
520 break (report.state, report.final_watermarks);
521 }
522 Err(crossbeam_channel::RecvTimeoutError::Timeout)
523 if !controller_handle.is_finished() =>
524 {
525 }
527 Err(_) => {
528 stop_drivers(
535 &self.shutdown,
536 &control_txs_for_stop,
537 std::mem::take(&mut driver_handles),
538 drain_timeout,
539 );
540 if let Some(drain) = drain_fn.take() {
541 sink_drain = Some(io.block_on(drain(drain_timeout)));
542 }
543 break (
544 ExitState::Failed(FatalErrorReport {
545 component: "controller".into(),
546 reason: "controller thread panicked".into(),
547 }),
548 Vec::new(),
549 );
550 }
551 }
552 };
553 for h in driver_handles {
556 let _ = h.join();
557 }
558 if let (ExitState::Completed, Some(report)) = (&state, driver_panic) {
561 state = ExitState::Failed(report);
562 }
563
564 if let Some(tx) = &admin_stop_tx {
565 let _ = tx.send(true);
566 }
567 let _ = controller_handle.join();
570 io.shutdown_timeout(Duration::from_secs(2));
571
572 Ok(ExitReport {
573 state,
574 sink_drain,
575 final_watermarks,
576 })
577 }
578}
579
580fn stop_drivers<L>(
590 shutdown: &AtomicBool,
591 control_txs: &[crossbeam_channel::Sender<ThreadControl<L>>],
592 driver_handles: Vec<std::thread::JoinHandle<DriverExit>>,
593 grace: Duration,
594) {
595 shutdown.store(true, Ordering::Relaxed);
596 let deadline = Instant::now() + grace;
597 let barrier = DrainBarrier::new(control_txs.len());
598 for tx in control_txs {
599 let _ = tx.send(ThreadControl::Shutdown {
600 barrier: barrier.clone(),
601 deadline,
602 });
603 }
604 for handle in driver_handles {
605 let _ = handle.join();
606 }
607}
608
609pub(crate) fn install_or_reuse(settings: &MetricsSettings) -> Result<MetricsHandle, StartError> {
614 match metrics::install(settings) {
615 Ok(h) => Ok(h),
616 Err(metrics::MetricsError::AlreadyInstalled) => {
617 tracing::warn!(
618 "a metrics recorder is already installed; continuing \
619 with the existing one and a detached render handle"
620 );
621 metrics::install(&MetricsSettings {
622 exporter: Exporter::None,
623 ..settings.clone()
624 })
625 .map_err(|e| StartError::Metrics(e.to_string()))
626 }
627 Err(e) => Err(StartError::Metrics(e.to_string())),
628 }
629}
630
631#[must_use]
640pub fn metrics_settings(config: &PipelineConfig) -> MetricsSettings {
641 MetricsSettings {
642 exporter: match config.metrics.exporter {
643 MetricsExporter::Prometheus => Exporter::Prometheus,
644 MetricsExporter::None => Exporter::None,
645 },
646 per_partition_detail: config.metrics.per_partition_detail,
647 e2e_basis: match config.metrics.e2e_basis {
648 crate::config::E2eBasis::Ingest => E2eBasis::Ingest,
649 crate::config::E2eBasis::Event => E2eBasis::Event,
650 },
651 }
652}
653
654async fn wait_for_signal() {
655 #[cfg(unix)]
656 {
657 let mut term =
658 match tokio::signal::unix::signal(tokio::signal::unix::SignalKind::terminate()) {
659 Ok(s) => s,
660 Err(e) => {
661 tracing::error!(error = %e, "failed to install SIGTERM handler");
662 std::future::pending::<()>().await;
663 return;
664 }
665 };
666 tokio::select! {
667 _ = term.recv() => {}
668 r = tokio::signal::ctrl_c() => {
669 if let Err(e) = r {
670 tracing::error!(error = %e, "ctrl_c handler failed");
671 std::future::pending::<()>().await;
672 }
673 }
674 }
675 }
676 #[cfg(not(unix))]
677 {
678 if let Err(e) = tokio::signal::ctrl_c().await {
679 tracing::error!(error = %e, "ctrl_c handler failed");
680 std::future::pending::<()>().await;
681 }
682 }
683}