laminar_db/pipeline/streaming_coordinator/
source_runtime.rs

1//! Mechanically extracted coordinator responsibility.
2
3use super::{
4    Arc, AtomicBool, CheckpointBarrier, ConnectorTaskTracker, DbError, Duration, Ordering,
5    OwnedConnectorTaskFences, RecordBatch, SourceBatchCursor, SourceCheckpoint, SourceMsgTx,
6};
7#[cfg(feature = "cluster")]
8use super::{
9    AssignmentDrainId, CheckpointParticipant, ClusterController, SourceDrainOutcome,
10    SourceDrainRequest, SourceDrainResolution,
11};
12
13#[cfg(feature = "cluster")]
14pub(super) struct SourceProcessAuthority {
15    controller: Arc<ClusterController>,
16    pub(super) lost: tokio_util::sync::CancellationToken,
17    pub(super) watcher_abort: Option<tokio::task::AbortHandle>,
18}
19
20#[cfg(feature = "cluster")]
21impl Drop for SourceProcessAuthority {
22    fn drop(&mut self) {
23        if let Some(watcher) = &self.watcher_abort {
24            watcher.abort();
25        }
26    }
27}
28
29/// Shared process-lease observation for every source task in one coordinator.
30///
31/// One watcher owns the deadline timer. Per-poll code selects only on the retained cancellation
32/// token and rechecks the controller's monotonic deadline at publication boundaries.
33#[cfg(feature = "cluster")]
34impl SourceProcessAuthority {
35    pub(super) fn new(controller: Arc<ClusterController>) -> Arc<Self> {
36        let lost = tokio_util::sync::CancellationToken::new();
37        if !controller.process_lease_is_live() {
38            lost.cancel();
39        }
40        let watcher_abort = if lost.is_cancelled() {
41            None
42        } else {
43            let task_controller = Arc::clone(&controller);
44            let task_lost = lost.clone();
45            let watcher = tokio::spawn(async move {
46                task_controller.wait_for_process_lease_loss().await;
47                task_lost.cancel();
48            });
49            Some(watcher.abort_handle())
50        };
51        Arc::new(Self {
52            controller,
53            lost,
54            watcher_abort,
55        })
56    }
57
58    #[inline]
59    pub(super) fn is_live(&self) -> bool {
60        !self.lost.is_cancelled() && self.controller.process_lease_is_live()
61    }
62
63    pub(super) async fn cancelled(&self) {
64        self.lost.cancelled().await;
65    }
66}
67
68#[cfg(feature = "cluster")]
69#[inline]
70pub(super) fn source_process_authority_is_live(authority: Option<&SourceProcessAuthority>) -> bool {
71    authority.is_none_or(SourceProcessAuthority::is_live)
72}
73
74/// Wait for coordinator batching or idle work without letting process-lease loss remain hidden
75/// behind the timer.
76pub(super) async fn wait_coordinator_delay(
77    duration: Duration,
78    #[cfg(feature = "cluster")] process_authority: Option<&SourceProcessAuthority>,
79) -> bool {
80    #[cfg(feature = "cluster")]
81    if let Some(authority) = process_authority {
82        return tokio::select! {
83            biased;
84            () = authority.cancelled() => true,
85            () = tokio::time::sleep(duration) => false,
86        };
87    }
88
89    tokio::time::sleep(duration).await;
90    false
91}
92
93/// Message from a source task to the coordinator with its production-time cursor.
94pub(super) enum SourceMsg {
95    Batch {
96        source_idx: usize,
97        batch: RecordBatch,
98        /// Committed to `committed_offsets` only after successful cycle publication.
99        cursor: SourceBatchCursor,
100    },
101    Barrier {
102        source_idx: usize,
103        barrier: CheckpointBarrier,
104        checkpoint: SourceCheckpoint,
105    },
106}
107
108/// Publish into the bounded source FIFO without allowing backpressure to hide terminal shutdown or
109/// process-lease loss. Process authority is revalidated after a successful send because its fence
110/// may cross while the queue wakes the producer.
111pub(super) async fn send_source_msg(
112    tx: &SourceMsgTx,
113    msg: SourceMsg,
114    shutdown: &tokio::sync::Notify,
115    #[cfg(feature = "cluster")] process_authority: Option<&SourceProcessAuthority>,
116) -> bool {
117    #[cfg(feature = "cluster")]
118    if let Some(authority) = process_authority {
119        if !authority.is_live() {
120            return false;
121        }
122        let sent = tokio::select! {
123            biased;
124            () = authority.cancelled() => false,
125            () = shutdown.notified() => false,
126            result = tx.send(msg) => result.is_ok(),
127        };
128        return sent && authority.is_live();
129    }
130
131    tokio::select! {
132        biased;
133        () = shutdown.notified() => false,
134        result = tx.send(msg) => result.is_ok(),
135    }
136}
137
138/// Fatal source-task control event. Faults use a dedicated unbounded side channel so they cannot
139/// sit behind a full data queue or be hidden while external-commit backpressure pauses data.
140pub(super) struct SourceFault {
141    pub(super) source: Arc<str>,
142    pub(super) error: String,
143}
144
145/// Reports every source-task exit that was not initiated by coordinator shutdown, including a
146/// panic while polling connector code. `UnboundedSender::send` is synchronous and allocation-only,
147/// so this remains reliable from `Drop` during unwinding without blocking the data path.
148pub(super) struct SourceTaskExitGuard {
149    pub(super) source: Arc<str>,
150    pub(super) expected_shutdown: Arc<AtomicBool>,
151    pub(super) fault_tx: tokio::sync::mpsc::UnboundedSender<SourceFault>,
152}
153
154impl Drop for SourceTaskExitGuard {
155    fn drop(&mut self) {
156        if !self.expected_shutdown.load(Ordering::Acquire) {
157            let _ = self.fault_tx.send(SourceFault {
158                source: Arc::clone(&self.source),
159                error: "source task exited without coordinator shutdown".into(),
160            });
161        }
162    }
163}
164
165pub(super) struct SourceActorTerminalState {
166    finished: AtomicBool,
167    notify: tokio::sync::Notify,
168}
169
170impl SourceActorTerminalState {
171    fn new() -> Self {
172        Self {
173            finished: AtomicBool::new(false),
174            notify: tokio::sync::Notify::new(),
175        }
176    }
177
178    fn finish(&self) {
179        if !self.finished.swap(true, Ordering::AcqRel) {
180            self.notify.notify_waiters();
181        }
182    }
183
184    pub(super) fn is_finished(&self) -> bool {
185        self.finished.load(Ordering::Acquire)
186    }
187
188    async fn wait_finished(&self) {
189        loop {
190            let notified = self.notify.notified();
191            tokio::pin!(notified);
192            notified.as_mut().enable();
193            if self.is_finished() {
194                return;
195            }
196            notified.await;
197        }
198    }
199}
200
201#[cfg(test)]
202pub(super) struct SourceActorLifetime(Arc<SourceActorTerminalState>);
203
204#[cfg(test)]
205impl Drop for SourceActorLifetime {
206    fn drop(&mut self) {
207        self.0.finish();
208    }
209}
210
211#[cfg(test)]
212pub(super) fn source_actor_terminal_guard() -> (SourceActorLifetime, Arc<SourceActorTerminalState>)
213{
214    let terminal = Arc::new(SourceActorTerminalState::new());
215    (SourceActorLifetime(Arc::clone(&terminal)), terminal)
216}
217
218pub(super) struct SourceActorFuture<F> {
219    actor: Option<std::pin::Pin<Box<F>>>,
220    terminal: Arc<SourceActorTerminalState>,
221}
222
223// Moving this wrapper never moves the separately pinned actor allocation.
224impl<F> Unpin for SourceActorFuture<F> {}
225
226impl<F> std::future::Future for SourceActorFuture<F>
227where
228    F: std::future::Future<Output = ()>,
229{
230    type Output = ();
231
232    fn poll(
233        mut self: std::pin::Pin<&mut Self>,
234        context: &mut std::task::Context<'_>,
235    ) -> std::task::Poll<Self::Output> {
236        let actor = self
237            .actor
238            .as_mut()
239            .expect("source actor polled after terminal completion");
240        if actor.as_mut().poll(context).is_pending() {
241            return std::task::Poll::Pending;
242        }
243        // Drop the complete actor future, including its connector, before publishing exit.
244        self.actor.take();
245        self.terminal.finish();
246        std::task::Poll::Ready(())
247    }
248}
249
250impl<F> Drop for SourceActorFuture<F> {
251    fn drop(&mut self) {
252        // Cancellation must drop the actor and its connector before another generation can observe
253        // terminal completion.
254        self.actor.take();
255        self.terminal.finish();
256    }
257}
258
259pub(super) fn spawn_source_actor<F>(
260    runtime: &tokio::runtime::Handle,
261    actor: F,
262) -> (tokio::task::JoinHandle<()>, Arc<SourceActorTerminalState>)
263where
264    F: std::future::Future<Output = ()> + Send + 'static,
265{
266    let terminal = Arc::new(SourceActorTerminalState::new());
267    let join = runtime.spawn(SourceActorFuture {
268        actor: Some(Box::pin(actor)),
269        terminal: Arc::clone(&terminal),
270    });
271    (join, terminal)
272}
273
274#[derive(Clone)]
275pub(super) enum SourceTaskOutcome {
276    Success,
277    Cancelled,
278    Failed(Arc<str>),
279}
280
281pub(super) struct SourceTaskState {
282    pub(super) name: Arc<str>,
283    shutdown: Arc<tokio::sync::Notify>,
284    expected_shutdown: Arc<AtomicBool>,
285    abort: tokio::task::AbortHandle,
286    actor_terminal: Arc<SourceActorTerminalState>,
287    terminal_tasks: Option<ConnectorTaskTracker>,
288    outcome: parking_lot::Mutex<Option<SourceTaskOutcome>>,
289    #[cfg(feature = "cluster")]
290    drain: parking_lot::Mutex<Option<SourceDrainLeaseControl>>,
291}
292
293#[cfg(feature = "cluster")]
294#[derive(Clone)]
295pub(super) enum SourceDrainCommand {
296    Begin {
297        request: SourceDrainRequest,
298        participant: CheckpointParticipant,
299        deadline: tokio::time::Instant,
300    },
301    Resolve {
302        resolution: SourceDrainResolution,
303        deadline: tokio::time::Instant,
304    },
305}
306
307#[cfg(feature = "cluster")]
308#[derive(Clone, Copy)]
309pub(super) enum SourceDrainCommandPolicy {
310    Any,
311    ResolveOnly,
312}
313
314#[cfg(feature = "cluster")]
315#[derive(Clone, Debug)]
316pub(super) enum SourceDrainTaskStatus {
317    Idle,
318    Pausing(AssignmentDrainId),
319    Ready(SourceDrainReceipt),
320    Resolved {
321        round: AssignmentDrainId,
322        outcome: SourceDrainOutcome,
323    },
324}
325
326#[cfg(feature = "cluster")]
327#[derive(Clone, Debug, PartialEq, Eq)]
328pub(super) struct SourceDrainReceipt {
329    pub(super) round: AssignmentDrainId,
330    pub(super) participant: CheckpointParticipant,
331    pub(super) source_task_incarnation: uuid::Uuid,
332}
333
334#[cfg(feature = "cluster")]
335impl SourceDrainReceipt {
336    pub(super) fn is_canonical(&self) -> bool {
337        self.round.is_canonical()
338            && self.participant.node_id != 0
339            && !self.participant.boot_incarnation.is_nil()
340            && !self.source_task_incarnation.is_nil()
341    }
342}
343
344#[cfg(feature = "cluster")]
345pub(super) fn validate_source_drain_receipts(
346    round: AssignmentDrainId,
347    participant: CheckpointParticipant,
348    receipts: &[SourceDrainReceipt],
349) -> Result<(), String> {
350    if receipts.iter().any(|receipt| {
351        !receipt.is_canonical() || receipt.round != round || receipt.participant != participant
352    }) {
353        return Err("source drain contains a stale or non-canonical task receipt".into());
354    }
355    let mut task_incarnations: Vec<uuid::Uuid> = receipts
356        .iter()
357        .map(|receipt| receipt.source_task_incarnation)
358        .collect();
359    task_incarnations.sort_unstable();
360    if task_incarnations.windows(2).any(|pair| pair[0] == pair[1]) {
361        return Err("source drain contains duplicate task receipts".into());
362    }
363    Ok(())
364}
365
366#[cfg(feature = "cluster")]
367#[derive(Clone)]
368pub(super) struct SourceDrainLeaseControl {
369    pub(super) task_incarnation: uuid::Uuid,
370    pub(super) command_tx: tokio::sync::watch::Sender<Option<SourceDrainCommand>>,
371    pub(super) status_tx: tokio::sync::watch::Sender<SourceDrainTaskStatus>,
372    pub(super) wake: Arc<tokio::sync::Notify>,
373}
374
375#[cfg(feature = "cluster")]
376pub(super) struct ActiveSourceDrain {
377    pub(super) request: SourceDrainRequest,
378    pub(super) participant: CheckpointParticipant,
379    pub(super) provider_drain: bool,
380    pub(super) prepare_deadline: tokio::time::Instant,
381    pub(super) ready: bool,
382    pub(super) pending_resolution: Option<PendingSourceDrainResolution>,
383}
384
385#[cfg(feature = "cluster")]
386#[derive(Clone, Copy)]
387pub(super) struct PendingSourceDrainResolution {
388    pub(super) resolution: SourceDrainResolution,
389    pub(super) deadline: tokio::time::Instant,
390}
391
392/// DB-owned proof that one source generation has reached terminal completion.
393///
394/// Actor exit is published by a wrapper that owns the connector future before spawn, and connector
395/// children are observed through the exact tracker captured at construction. Neither proof depends
396/// on the detached outcome supervisor continuing to run.
397#[derive(Clone)]
398pub(crate) struct SourceTaskLease {
399    pub(super) state: Arc<SourceTaskState>,
400}
401
402pub(crate) type OwnedSourceTasks = Arc<parking_lot::Mutex<Vec<SourceTaskLease>>>;
403
404impl SourceTaskLease {
405    #[cfg(all(test, feature = "cluster"))]
406    pub(crate) fn spawn_for_test<F>(
407        name: &str,
408        actor: F,
409        terminal_tasks: Option<ConnectorTaskTracker>,
410    ) -> Self
411    where
412        F: std::future::Future<Output = ()> + Send + 'static,
413    {
414        let runtime = tokio::runtime::Handle::current();
415        let (join, terminal) = spawn_source_actor(&runtime, actor);
416        Self::supervise(
417            Arc::from(name),
418            Arc::new(tokio::sync::Notify::new()),
419            Arc::new(AtomicBool::new(false)),
420            join,
421            terminal,
422            terminal_tasks,
423            &runtime,
424        )
425    }
426
427    pub(super) fn supervise(
428        name: Arc<str>,
429        shutdown: Arc<tokio::sync::Notify>,
430        expected_shutdown: Arc<AtomicBool>,
431        join: tokio::task::JoinHandle<()>,
432        actor_terminal: Arc<SourceActorTerminalState>,
433        terminal_tasks: Option<ConnectorTaskTracker>,
434        runtime: &tokio::runtime::Handle,
435    ) -> Self {
436        let state = Arc::new(SourceTaskState {
437            name,
438            shutdown,
439            expected_shutdown,
440            abort: join.abort_handle(),
441            actor_terminal,
442            terminal_tasks,
443            outcome: parking_lot::Mutex::new(None),
444            #[cfg(feature = "cluster")]
445            drain: parking_lot::Mutex::new(None),
446        });
447        let supervisor_state = Arc::clone(&state);
448        let supervisor = runtime.spawn(async move {
449            let outcome = match join.await {
450                Ok(()) => SourceTaskOutcome::Success,
451                Err(error) if error.is_cancelled() => SourceTaskOutcome::Cancelled,
452                Err(error) => SourceTaskOutcome::Failed(Arc::from(error.to_string())),
453            };
454            *supervisor_state.outcome.lock() = Some(outcome);
455        });
456        // Tokio owns the scheduled outcome observer. Terminal replacement fencing is independent
457        // of this detached task and comes from the actor wrapper plus connector tracker above.
458        drop(supervisor);
459        Self { state }
460    }
461
462    pub(crate) fn name(&self) -> &str {
463        &self.state.name
464    }
465
466    #[cfg(feature = "cluster")]
467    pub(super) fn install_drain_control(&self, control: SourceDrainLeaseControl) {
468        *self.state.drain.lock() = Some(control);
469    }
470
471    #[cfg(feature = "cluster")]
472    pub(super) fn drain_control(&self) -> Option<SourceDrainLeaseControl> {
473        self.state.drain.lock().clone()
474    }
475
476    pub(crate) fn is_finished(&self) -> bool {
477        self.state.actor_terminal.is_finished()
478            && self
479                .state
480                .terminal_tasks
481                .as_ref()
482                .is_none_or(ConnectorTaskTracker::is_terminated)
483    }
484
485    #[cfg(feature = "cluster")]
486    pub(crate) fn is_running(&self) -> bool {
487        !self.state.actor_terminal.is_finished()
488            && !self.state.abort.is_finished()
489            && !self.state.expected_shutdown.load(Ordering::Acquire)
490    }
491
492    pub(crate) fn request_shutdown(&self) {
493        self.mark_expected_shutdown();
494        self.notify_shutdown();
495    }
496
497    pub(super) fn mark_expected_shutdown(&self) {
498        self.state.expected_shutdown.store(true, Ordering::Release);
499    }
500
501    pub(super) fn notify_shutdown(&self) {
502        self.state.shutdown.notify_one();
503    }
504
505    pub(crate) fn abort(&self) {
506        self.state.abort.abort();
507    }
508
509    pub(super) async fn wait_finished(&self) {
510        let actor = self.state.actor_terminal.wait_finished();
511        let connector_tasks = async {
512            if let Some(tasks) = self.state.terminal_tasks.as_ref() {
513                tasks.wait_terminated().await;
514            }
515        };
516        tokio::join!(actor, connector_tasks);
517    }
518
519    pub(crate) async fn wait_until(&self, deadline: tokio::time::Instant) -> bool {
520        if self.is_finished() {
521            return true;
522        }
523        tokio::time::timeout_at(deadline, self.wait_finished())
524            .await
525            .is_ok()
526            || self.is_finished()
527    }
528
529    pub(crate) fn log_terminal_outcome(&self) {
530        match self.state.outcome.lock().clone() {
531            Some(SourceTaskOutcome::Failed(error)) => {
532                tracing::warn!(source = %self.state.name, %error, "source task panicked");
533            }
534            Some(SourceTaskOutcome::Success | SourceTaskOutcome::Cancelled) | None => {}
535        }
536    }
537}
538
539/// Stable owner for source generations constructed through
540/// [`crate::pipeline::StreamingCoordinator::new`].
541#[derive(Clone)]
542pub struct StreamingCoordinatorRuntime {
543    pub(super) owned_source_tasks: OwnedSourceTasks,
544    pub(super) owned_connector_task_fences: OwnedConnectorTaskFences,
545    pub(super) construction: Arc<tokio::sync::Mutex<()>>,
546    pub(super) active_coordinator: Arc<AtomicBool>,
547}
548
549pub(super) struct StreamingCoordinatorGeneration {
550    runtime: StreamingCoordinatorRuntime,
551}
552
553impl Drop for StreamingCoordinatorGeneration {
554    fn drop(&mut self) {
555        self.runtime
556            .active_coordinator
557            .store(false, Ordering::Release);
558    }
559}
560
561impl StreamingCoordinatorRuntime {
562    /// Creates an empty runtime owner with no active source generation.
563    #[must_use]
564    pub fn new() -> Self {
565        Self {
566            owned_source_tasks: Arc::new(parking_lot::Mutex::new(Vec::new())),
567            owned_connector_task_fences: Arc::new(parking_lot::Mutex::new(Vec::new())),
568            construction: Arc::new(tokio::sync::Mutex::new(())),
569            active_coordinator: Arc::new(AtomicBool::new(false)),
570        }
571    }
572
573    pub(super) fn claim_generation(&self) -> Result<StreamingCoordinatorGeneration, DbError> {
574        self.active_coordinator
575            .compare_exchange(false, true, Ordering::AcqRel, Ordering::Acquire)
576            .map_err(|_| {
577                DbError::Pipeline(
578                    "cannot construct a replacement streaming coordinator while the prior \
579                     coordinator generation is still active"
580                        .into(),
581                )
582            })?;
583        Ok(StreamingCoordinatorGeneration {
584            runtime: self.clone(),
585        })
586    }
587
588    pub(super) fn prune_and_require_idle(&self) -> Result<(), DbError> {
589        let source_names = {
590            let mut sources = self.owned_source_tasks.lock();
591            sources.retain(|source| !source.is_finished());
592            sources
593                .iter()
594                .map(|source| source.name().to_owned())
595                .collect::<Vec<_>>()
596        };
597        let connector_names = {
598            let mut connectors = self.owned_connector_task_fences.lock();
599            connectors.retain(|connector| !connector.is_finished());
600            connectors
601                .iter()
602                .map(|connector| connector.name().to_owned())
603                .collect::<Vec<_>>()
604        };
605        if source_names.is_empty() && connector_names.is_empty() {
606            return Ok(());
607        }
608        Err(DbError::Pipeline(format!(
609            "cannot construct a replacement streaming coordinator while prior connector \
610             generations remain unresolved: sources=[{}], connector_tasks=[{}]",
611            source_names.join(", "),
612            connector_names.join(", ")
613        )))
614    }
615}
616
617impl Default for StreamingCoordinatorRuntime {
618    fn default() -> Self {
619        Self::new()
620    }
621}
laminar_db/pipeline/streaming_coordinator/source_runtime.rs

laminar_db/pipeline/streaming_coordinator/
source_runtime.rs
