1use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
10use std::sync::Arc;
11use std::time::{Duration, Instant};
12
13use arrow_array::RecordBatch;
14use crossfire::{mpsc, AsyncRx, MAsyncTx};
15use laminar_connectors::checkpoint::SourceCheckpoint;
16#[cfg(test)]
17use laminar_connectors::checkpoint::SourceCheckpointDelta;
18use laminar_connectors::connector::{
19 schema_with_source_mutations_and_row_positions, schema_with_source_row_positions,
20 strip_source_row_positions, ConnectorCancellationPolicy, ConnectorTaskTracker,
21 DeliveryGuarantee, SourceBatch, SourceBatchCursor, SourceCheckpointUnavailablePolicy,
22 SourceConnector, SourceConsistency, SourceContract, SourceInputMode, SourcePosition,
23 SourceRowPositionCapability, SourceStart, SOURCE_MUTATION_COLUMN,
24};
25#[cfg(feature = "cluster")]
26use laminar_connectors::connector::{
27 SourceDrainOutcome, SourceDrainRequest, SourceDrainResolution,
28};
29use laminar_connectors::error::ConnectorError;
30#[cfg(feature = "cluster")]
31use laminar_core::checkpoint::{
32 AssignmentDrainId, AssignmentDrainTransition, CheckpointParticipant,
33};
34use laminar_core::checkpoint::{CheckpointAttempt, CheckpointAttemptRelation};
35use laminar_core::checkpoint::{CheckpointBarrier, CheckpointBarrierInjector};
36#[cfg(feature = "cluster")]
37use laminar_core::cluster::control::ClusterController;
38use rustc_hash::{FxHashMap, FxHashSet};
39
40use super::callback::{
41 BarrierOutcome, CheckpointAssignmentAdmission, CheckpointCompletion, CheckpointControlOutcome,
42 CheckpointControlWake, CycleError, CycleOutcome, PipelineCallback, SourceBarrierControl,
43 SourceBarrierSignal, SourceRegistration,
44};
45#[cfg(test)]
46use super::config::CheckpointSchedule;
47use super::config::PipelineConfig;
48use crate::catalog::{schema_has_reserved_mutation_columns, validate_source_batch};
49use crate::connector_task_fence::{ConnectorTaskFenceRegistration, OwnedConnectorTaskFences};
50use crate::error::DbError;
51
52type SourceMsgRx = AsyncRx<mpsc::Array<SourceMsg>>;
53type SourceMsgTx = MAsyncTx<mpsc::Array<SourceMsg>>;
54type ControlMsgRx = AsyncRx<mpsc::Array<super::ControlMsg>>;
55type ForceCheckpointRequest = crate::db::ForceCheckpointRequest;
56
57#[derive(Clone, Copy, Debug)]
58struct CyclePublicationDurations {
59 output_store_ns: u64,
60 sink_enqueue_ns: u64,
61}
62
63mod barrier_handling;
64mod checkpoint_admission;
65mod checkpoint_lifecycle;
66mod construction;
67mod cycle;
68mod drain_ownership;
69mod execution;
70mod intake_backpressure;
71mod shutdown;
72mod source_actor;
73mod source_drain;
74mod source_lifecycle;
75mod source_runtime;
76
77use source_drain::{acknowledge_latest_source_commit, take_batch_cursor, try_source_checkpoint};
78#[cfg(feature = "cluster")]
79use source_drain::{
80 apply_latest_source_drain_command_fenced, publish_source_drain_ready_fenced,
81 resolve_pending_source_drain_fenced, source_drain_flushing, source_drain_held,
82};
83#[cfg(feature = "cluster")]
84use source_lifecycle::{
85 check_source_sync_fence, source_operation_authority_error, wait_source_drain_hold,
86};
87use source_lifecycle::{
88 poll_source_once, run_source_operation, source_operation_deadline_error, start_source_once,
89 wait_source_barrier_release, wait_source_idle, SourceConnectorLifecycle,
90 SourceOperationOutcome, SourcePollOutcome, SourceStartFailure, SourceStartOutcome,
91};
92use source_runtime::{
93 send_source_msg, spawn_source_actor, wait_coordinator_delay, SourceFault, SourceMsg,
94 SourceTaskExitGuard, StreamingCoordinatorGeneration,
95};
96#[cfg(feature = "cluster")]
97use source_runtime::{
98 source_process_authority_is_live, validate_source_drain_receipts, ActiveSourceDrain,
99 PendingSourceDrainResolution, SourceDrainCommand, SourceDrainCommandPolicy,
100 SourceDrainLeaseControl, SourceDrainReceipt, SourceDrainTaskStatus, SourceProcessAuthority,
101};
102
103#[cfg(all(test, feature = "cluster"))]
104pub(crate) use drain_ownership::install_replacement_source_drain_task_for_test;
105#[cfg(feature = "cluster")]
106pub(crate) use drain_ownership::{
107 owned_source_drain_resolved, prepare_owned_source_drain, resolve_owned_source_drain,
108};
109pub use source_runtime::StreamingCoordinatorRuntime;
110pub(crate) use source_runtime::{OwnedSourceTasks, SourceTaskLease};
111
112struct SourceHandle {
113 recovery_cursor: bool,
116 task: SourceTaskLease,
117 startup_activation: Option<crossfire::oneshot::TxOneshot<()>>,
120 barrier_injector: CheckpointBarrierInjector,
121 barrier_release_tx: tokio::sync::watch::Sender<Option<SourceBarrierSignal>>,
123 epoch_committed_tx: tokio::sync::watch::Sender<Option<(u64, SourceCheckpoint)>>,
127}
128
129pub(crate) struct TrackedSourceRegistration {
130 source: SourceRegistration,
131 contract: SourceContract,
132 expected_schema: arrow_schema::SchemaRef,
133 positioned_schema: arrow_schema::SchemaRef,
134 mutation_schema: arrow_schema::SchemaRef,
135 primary_key: Vec<String>,
136 primary_key_indices: Vec<usize>,
137 schema_admitted: bool,
138 admitted_non_append_mode: Option<SourceInputMode>,
139 task_fence: ConnectorTaskFenceRegistration,
140}
141
142pub(crate) const MUTATION_SOURCE_NOT_ADMITTED: &str =
143 "[LDB-5039] mutation sources require an exclusively admitted stateful operator route";
144
145pub(crate) fn admit_append_only_source(
146 contract: SourceContract,
147 has_reserved_mutation_columns: bool,
148) -> Result<(), &'static str> {
149 if contract.input_mode == SourceInputMode::AppendOnly && !has_reserved_mutation_columns {
150 Ok(())
151 } else {
152 Err(MUTATION_SOURCE_NOT_ADMITTED)
153 }
154}
155
156impl TrackedSourceRegistration {
157 fn metadata_schemas(
158 source_name: &str,
159 contract: SourceContract,
160 expected_schema: &arrow_schema::SchemaRef,
161 ) -> Result<(arrow_schema::SchemaRef, arrow_schema::SchemaRef), DbError> {
162 let map_error = |error| {
163 DbError::Config(format!(
164 "source '{source_name}' has an invalid source-metadata schema: {error}"
165 ))
166 };
167 let positioned = schema_with_source_row_positions(expected_schema).map_err(map_error)?;
168 let mutations =
169 schema_with_source_mutations_and_row_positions(expected_schema).map_err(map_error)?;
170 if contract.row_positions == SourceRowPositionCapability::OrderedDeterministic {
171 Ok((positioned, mutations))
172 } else {
173 Ok((Arc::clone(expected_schema), Arc::clone(expected_schema)))
174 }
175 }
176
177 fn resolve_contract(source: &SourceRegistration) -> Result<SourceContract, DbError> {
178 let contract = source.connector.contract(&source.config).map_err(|error| {
179 DbError::Config(format!(
180 "source '{}' (type '{}') has an invalid contract: {error}",
181 source.name,
182 source.config.connector_type()
183 ))
184 })?;
185 Ok(contract)
186 }
187
188 pub(crate) fn capture(
189 source: SourceRegistration,
190 owned: &OwnedConnectorTaskFences,
191 ) -> Result<Self, DbError> {
192 let contract = Self::resolve_contract(&source)?;
193 let expected_schema = source.connector.schema();
194 let (positioned_schema, mutation_schema) =
195 Self::metadata_schemas(&source.name, contract, &expected_schema)?;
196 let task_fence = ConnectorTaskFenceRegistration::capture_registered(
197 Arc::<str>::from(format!("source:{}", source.name)),
198 source.connector.terminal_task_tracker(),
199 owned,
200 );
201 Ok(Self {
202 source,
203 contract,
204 expected_schema,
205 positioned_schema,
206 mutation_schema,
207 primary_key: Vec::new(),
208 primary_key_indices: Vec::new(),
209 schema_admitted: false,
210 admitted_non_append_mode: None,
211 task_fence,
212 })
213 }
214
215 pub(crate) fn from_captured(
216 source: SourceRegistration,
217 task_fence: ConnectorTaskFenceRegistration,
218 ) -> Result<Self, DbError> {
219 let contract = Self::resolve_contract(&source)?;
220 let expected_schema = source.connector.schema();
221 let (positioned_schema, mutation_schema) =
222 Self::metadata_schemas(&source.name, contract, &expected_schema)?;
223 Ok(Self {
224 source,
225 contract,
226 expected_schema,
227 positioned_schema,
228 mutation_schema,
229 primary_key: Vec::new(),
230 primary_key_indices: Vec::new(),
231 schema_admitted: false,
232 admitted_non_append_mode: None,
233 task_fence,
234 })
235 }
236
237 pub(crate) fn with_admitted_schema(
238 mut self,
239 expected_schema: arrow_schema::SchemaRef,
240 primary_key: Vec<String>,
241 ) -> Result<Self, DbError> {
242 let primary_key_indices = primary_key
243 .iter()
244 .map(|column| {
245 expected_schema.index_of(column).map_err(|_| {
246 DbError::Config(format!(
247 "source '{}' primary-key column '{column}' is absent from its admitted schema",
248 self.name
249 ))
250 })
251 })
252 .collect::<Result<Vec<_>, _>>()?;
253 self.expected_schema = expected_schema;
254 (self.positioned_schema, self.mutation_schema) =
255 Self::metadata_schemas(&self.name, self.contract, &self.expected_schema)?;
256 self.primary_key = primary_key;
257 self.primary_key_indices = primary_key_indices;
258 self.schema_admitted = true;
259 Ok(self)
260 }
261
262 pub(crate) const fn contract(&self) -> SourceContract {
263 self.contract
264 }
265
266 pub(crate) fn with_temporal_right_mutations(mut self) -> Self {
267 debug_assert_eq!(self.contract.input_mode, SourceInputMode::KeyedUpsert);
268 self.admitted_non_append_mode = Some(SourceInputMode::KeyedUpsert);
269 self
270 }
271
272 pub(crate) fn with_ordered_interval_input_mode(
273 mut self,
274 mode: SourceInputMode,
275 ) -> Result<Self, DbError> {
276 if mode == SourceInputMode::AppendOnly || self.contract.input_mode != mode {
277 return Err(DbError::Config(format!(
278 "source '{}' lost its admitted bounded-interval input mode",
279 self.name
280 )));
281 }
282 if self
283 .admitted_non_append_mode
284 .is_some_and(|admitted| admitted != mode)
285 {
286 return Err(DbError::Config(format!(
287 "source '{}' has conflicting stateful mutation routes",
288 self.name
289 )));
290 }
291 self.admitted_non_append_mode = Some(mode);
292 Ok(self)
293 }
294
295 fn has_reserved_mutation_columns(&self) -> bool {
296 schema_has_reserved_mutation_columns(self.expected_schema.as_ref())
297 }
298}
299
300fn prepare_encoded_source_batch(
301 source_name: &str,
302 expected_schema: &arrow_schema::SchemaRef,
303 positioned_schema: &arrow_schema::SchemaRef,
304 mutation_schema: &arrow_schema::SchemaRef,
305 primary_key: &[String],
306 primary_key_indices: &[usize],
307 capability: SourceRowPositionCapability,
308 batch: SourceBatch,
309) -> Result<RecordBatch, laminar_core::streaming::StreamingError> {
310 validate_source_batch(
311 source_name,
312 expected_schema,
313 primary_key,
314 primary_key_indices,
315 &batch.records,
316 )?;
317 batch
318 .into_records_with_metadata(capability, positioned_schema, mutation_schema)
319 .map_err(|error| {
320 laminar_core::streaming::StreamingError::InvalidConfig(format!(
321 "source '{source_name}' emitted invalid source metadata: {error}"
322 ))
323 })
324}
325
326impl std::ops::Deref for TrackedSourceRegistration {
327 type Target = SourceRegistration;
328
329 fn deref(&self) -> &Self::Target {
330 &self.source
331 }
332}
333
334impl std::ops::DerefMut for TrackedSourceRegistration {
335 fn deref_mut(&mut self) -> &mut Self::Target {
336 &mut self.source
337 }
338}
339
340struct PreparedSourceGeneration {
341 registration: TrackedSourceRegistration,
342}
343
344impl SourceHandle {
345 fn barrier_control(&self) -> SourceBarrierControl {
346 SourceBarrierControl::new(
347 self.barrier_injector.clone(),
348 self.barrier_release_tx.clone(),
349 )
350 }
351}
352
353#[derive(Debug)]
355pub enum ExitReason {
356 Shutdown,
358 Halt(String),
361 Fault(String),
363}
364
365struct PendingWatermarkBatch {
366 source_name: Arc<str>,
367 batch: RecordBatch,
368 admission_floor: i64,
369 input_channels: Option<Arc<[Vec<u8>]>>,
370}
371
372pub struct StreamingCoordinator {
374 config: PipelineConfig,
375 rx: SourceMsgRx,
376 source_fault_rx: tokio::sync::mpsc::UnboundedReceiver<SourceFault>,
377 source_handles: Vec<SourceHandle>,
378 source_names: Vec<Arc<str>>,
379 source_input_modes: Vec<Option<SourceInputMode>>,
380 shutdown: Arc<tokio::sync::Notify>,
381 terminal_shutdown: tokio_util::sync::CancellationToken,
382 pending_barrier: PendingBarrier,
383 last_checkpoint: Instant,
384 checkpoint_retry_not_before: Option<Instant>,
385 checkpoint_retry_backoff: Duration,
386 source_batches_buf: FxHashMap<Arc<str>, Vec<RecordBatch>>,
387 parked_source_msg: Option<SourceMsg>,
390 pending_watermark_batches: Vec<PendingWatermarkBatch>,
391 barrier_seen: FxHashSet<usize>,
394 committed_offsets: Vec<Option<SourceCheckpoint>>,
396 pending_offsets: Vec<Option<SourceBatchCursor>>,
399 replay_pending: bool,
402 control_rx: ControlMsgRx,
403 checkpoint_complete_rx:
404 Option<crossfire::AsyncRx<crossfire::mpsc::Array<CheckpointCompletion>>>,
405 force_ckpt_rx: Option<crate::db::ForceCheckpointRx>,
407 manual_waiting: Vec<ForceCheckpointRequest>,
408 manual_handoff_required: bool,
410 manual_active: Option<ManualCheckpointAttempt>,
412 checkpoint_in_flight: Arc<AtomicU64>,
414 last_published_checkpoint: Option<CheckpointAttempt>,
417 #[cfg(feature = "cluster")]
418 process_authority: Option<Arc<SourceProcessAuthority>>,
419 public_generation: Option<StreamingCoordinatorGeneration>,
420}
421
422#[allow(clippy::struct_excessive_bools)]
425struct CoordinatorRunState {
426 batch_window: Duration,
427 checkpoint_control_wake: Option<CheckpointControlWake>,
428 shuffle_work_wake: Option<Arc<tokio::sync::Notify>>,
429 checkpoint_control_poll_at: tokio::time::Instant,
430 checkpoint_control_pending: bool,
431 intake_was_paused: bool,
434 barriers: Vec<(usize, CheckpointBarrier, SourceCheckpoint)>,
435 fault: Option<String>,
436 halted: bool,
437 halt_reason: Option<String>,
438 source_channel_expected: bool,
439}
440
441impl CoordinatorRunState {
442 fn halt(&mut self, reason: impl Into<String>) {
443 self.halted = true;
444 self.halt_reason.get_or_insert_with(|| reason.into());
445 self.fault = None;
446 }
447}
448
449#[derive(Default)]
450struct CoordinatorWake {
451 message: Option<SourceMsg>,
452 retrying_replay: bool,
453 checkpoint_control_due: bool,
454 gates: CoordinatorGates,
455}
456
457#[derive(Default)]
458struct CoordinatorGates {
459 intake_paused: bool,
460 external_commit_backpressured: bool,
461}
462
463enum CoordinatorWaitAction {
464 Cycle,
465 Continue,
466 Stop,
467}
468
469struct CoordinatorWait {
470 action: CoordinatorWaitAction,
471 wake: CoordinatorWake,
472}
473
474impl CoordinatorWait {
475 fn cycle(wake: CoordinatorWake) -> Self {
476 Self {
477 action: CoordinatorWaitAction::Cycle,
478 wake,
479 }
480 }
481
482 fn continue_loop() -> Self {
483 Self {
484 action: CoordinatorWaitAction::Continue,
485 wake: CoordinatorWake::default(),
486 }
487 }
488
489 fn stop() -> Self {
490 Self {
491 action: CoordinatorWaitAction::Stop,
492 wake: CoordinatorWake::default(),
493 }
494 }
495}
496
497impl Drop for StreamingCoordinator {
498 fn drop(&mut self) {
499 for handle in &self.source_handles {
503 handle.task.request_shutdown();
504 }
505 drop(self.public_generation.take());
507 }
508}
509
510struct ManualCheckpointAttempt {
512 attempt: CheckpointAttempt,
513 flags: u64,
514 requests: Vec<ForceCheckpointRequest>,
515}
516
517struct CheckpointAdmission {
518 manual: bool,
519 flags: u64,
520 assignment_fence: Option<laminar_core::cluster::control::CheckpointAssignmentFence>,
521 assignment_guard: Option<tokio::sync::OwnedMutexGuard<()>>,
525 deadline: tokio::time::Instant,
526}
527
528#[derive(Clone, Copy, Debug, PartialEq, Eq)]
529enum CheckpointCleanupOwner {
530 Originator,
532 Follower,
534}
535
536struct AlignedCheckpointContext {
537 cleanup_owner: CheckpointCleanupOwner,
538 attempt: CheckpointAttempt,
539 started_at: Instant,
540 flags: u64,
541 assignment_fence: Option<laminar_core::cluster::control::CheckpointAssignmentFence>,
542}
543
544struct PendingBarrier {
546 attempt: Option<CheckpointAttempt>,
547 sources_total: usize,
548 sources_aligned: FxHashSet<usize>,
549 source_checkpoints: FxHashMap<String, SourceCheckpoint>,
550 started_at: Instant,
551 deadline: Option<tokio::time::Instant>,
552 active: bool,
553 flags: u64,
554 assignment_fence: Option<laminar_core::cluster::control::CheckpointAssignmentFence>,
555 cleanup_owner: CheckpointCleanupOwner,
556}
557
558impl PendingBarrier {
559 fn new() -> Self {
560 Self {
561 attempt: None,
562 sources_total: 0,
563 sources_aligned: FxHashSet::default(),
564 source_checkpoints: FxHashMap::default(),
565 started_at: Instant::now(),
566 deadline: None,
567 active: false,
568 flags: laminar_core::checkpoint::flags::NONE,
569 assignment_fence: None,
570 cleanup_owner: CheckpointCleanupOwner::Originator,
571 }
572 }
573
574 #[cfg(test)]
575 fn reset(&mut self, attempt: CheckpointAttempt, sources_total: usize) {
576 self.reset_with_assignment(
577 attempt,
578 sources_total,
579 laminar_core::checkpoint::flags::NONE,
580 None,
581 None,
582 );
583 }
584
585 fn reset_follower(&mut self, attempt: CheckpointAttempt, sources_total: usize, flags: u64) {
586 self.reset_inner(
587 attempt,
588 sources_total,
589 flags,
590 None,
591 CheckpointCleanupOwner::Follower,
592 );
593 }
594
595 fn reset_with_assignment(
596 &mut self,
597 attempt: CheckpointAttempt,
598 sources_total: usize,
599 flags: u64,
600 assignment_fence: Option<laminar_core::cluster::control::CheckpointAssignmentFence>,
601 deadline: Option<tokio::time::Instant>,
602 ) {
603 self.reset_inner(
604 attempt,
605 sources_total,
606 flags,
607 assignment_fence,
608 CheckpointCleanupOwner::Originator,
609 );
610 self.deadline = deadline;
611 }
612
613 fn attempt_deadline(&self, checkpoint_timeout: Duration) -> tokio::time::Instant {
614 self.deadline
615 .unwrap_or_else(|| tokio::time::Instant::from_std(self.started_at) + checkpoint_timeout)
616 }
617
618 fn reset_inner(
619 &mut self,
620 attempt: CheckpointAttempt,
621 sources_total: usize,
622 flags: u64,
623 assignment_fence: Option<laminar_core::cluster::control::CheckpointAssignmentFence>,
624 cleanup_owner: CheckpointCleanupOwner,
625 ) {
626 self.attempt = Some(attempt);
627 self.sources_total = sources_total;
628 self.sources_aligned.clear();
629 self.source_checkpoints.clear();
630 self.started_at = Instant::now();
631 self.deadline = None;
632 self.active = true;
633 self.flags = flags;
634 self.assignment_fence = assignment_fence;
635 self.cleanup_owner = cleanup_owner;
636 }
637
638 fn take_active_attempt(&mut self) -> Option<(CheckpointAttempt, CheckpointCleanupOwner)> {
640 if !self.active {
641 return None;
642 }
643 let cleanup_owner = self.cleanup_owner;
644 self.active = false;
645 self.sources_total = 0;
646 self.sources_aligned.clear();
647 self.source_checkpoints.clear();
648 self.flags = laminar_core::checkpoint::flags::NONE;
649 self.assignment_fence = None;
650 self.deadline = None;
651 self.cleanup_owner = CheckpointCleanupOwner::Originator;
652 self.attempt.take().map(|attempt| (attempt, cleanup_owner))
653 }
654
655 fn clear(&mut self) {
656 self.active = false;
657 self.attempt = None;
658 self.sources_total = 0;
659 self.sources_aligned.clear();
660 self.source_checkpoints.clear();
661 self.flags = laminar_core::checkpoint::flags::NONE;
662 self.assignment_fence = None;
663 self.deadline = None;
664 self.cleanup_owner = CheckpointCleanupOwner::Originator;
665 }
666}
667
668pub(crate) const IDLE_TIMEOUT: Duration = Duration::from_millis(100);
670
671const CHECKPOINT_RETRY_BASE: Duration = Duration::from_millis(100);
673const CHECKPOINT_RETRY_MAX: Duration = Duration::from_secs(5);
674
675const SHUTDOWN_DRAIN_BUDGET: Duration = Duration::from_secs(2);
677
678const SHUTDOWN_JOIN_TIMEOUT: Duration = Duration::from_secs(8);
680
681const SHUTDOWN_COMPLETION_TICK: Duration = Duration::from_millis(10);
685
686#[cfg(test)]
687mod tests;