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    pub(super) fn supervise(
406        name: Arc<str>,
407        shutdown: Arc<tokio::sync::Notify>,
408        expected_shutdown: Arc<AtomicBool>,
409        join: tokio::task::JoinHandle<()>,
410        actor_terminal: Arc<SourceActorTerminalState>,
411        terminal_tasks: Option<ConnectorTaskTracker>,
412        runtime: &tokio::runtime::Handle,
413    ) -> Self {
414        let state = Arc::new(SourceTaskState {
415            name,
416            shutdown,
417            expected_shutdown,
418            abort: join.abort_handle(),
419            actor_terminal,
420            terminal_tasks,
421            outcome: parking_lot::Mutex::new(None),
422            #[cfg(feature = "cluster")]
423            drain: parking_lot::Mutex::new(None),
424        });
425        let supervisor_state = Arc::clone(&state);
426        let supervisor = runtime.spawn(async move {
427            let outcome = match join.await {
428                Ok(()) => SourceTaskOutcome::Success,
429                Err(error) if error.is_cancelled() => SourceTaskOutcome::Cancelled,
430                Err(error) => SourceTaskOutcome::Failed(Arc::from(error.to_string())),
431            };
432            *supervisor_state.outcome.lock() = Some(outcome);
433        });
434        // Tokio owns the scheduled outcome observer. Terminal replacement fencing is independent
435        // of this detached task and comes from the actor wrapper plus connector tracker above.
436        drop(supervisor);
437        Self { state }
438    }
439
440    pub(crate) fn name(&self) -> &str {
441        &self.state.name
442    }
443
444    #[cfg(feature = "cluster")]
445    pub(super) fn install_drain_control(&self, control: SourceDrainLeaseControl) {
446        *self.state.drain.lock() = Some(control);
447    }
448
449    #[cfg(feature = "cluster")]
450    pub(super) fn drain_control(&self) -> Option<SourceDrainLeaseControl> {
451        self.state.drain.lock().clone()
452    }
453
454    pub(crate) fn is_finished(&self) -> bool {
455        self.state.actor_terminal.is_finished()
456            && self
457                .state
458                .terminal_tasks
459                .as_ref()
460                .is_none_or(ConnectorTaskTracker::is_terminated)
461    }
462
463    pub(crate) fn request_shutdown(&self) {
464        self.mark_expected_shutdown();
465        self.notify_shutdown();
466    }
467
468    pub(super) fn mark_expected_shutdown(&self) {
469        self.state.expected_shutdown.store(true, Ordering::Release);
470    }
471
472    pub(super) fn notify_shutdown(&self) {
473        self.state.shutdown.notify_one();
474    }
475
476    pub(crate) fn abort(&self) {
477        self.state.abort.abort();
478    }
479
480    pub(super) async fn wait_finished(&self) {
481        let actor = self.state.actor_terminal.wait_finished();
482        let connector_tasks = async {
483            if let Some(tasks) = self.state.terminal_tasks.as_ref() {
484                tasks.wait_terminated().await;
485            }
486        };
487        tokio::join!(actor, connector_tasks);
488    }
489
490    pub(crate) async fn wait_until(&self, deadline: tokio::time::Instant) -> bool {
491        if self.is_finished() {
492            return true;
493        }
494        tokio::time::timeout_at(deadline, self.wait_finished())
495            .await
496            .is_ok()
497            || self.is_finished()
498    }
499
500    pub(crate) fn log_terminal_outcome(&self) {
501        match self.state.outcome.lock().clone() {
502            Some(SourceTaskOutcome::Failed(error)) => {
503                tracing::warn!(source = %self.state.name, %error, "source task panicked");
504            }
505            Some(SourceTaskOutcome::Success | SourceTaskOutcome::Cancelled) | None => {}
506        }
507    }
508}
509
510/// Stable owner for source generations constructed through
511/// [`crate::pipeline::StreamingCoordinator::new`].
512#[derive(Clone)]
513pub struct StreamingCoordinatorRuntime {
514    pub(super) owned_source_tasks: OwnedSourceTasks,
515    pub(super) owned_connector_task_fences: OwnedConnectorTaskFences,
516    pub(super) construction: Arc<tokio::sync::Mutex<()>>,
517    pub(super) active_coordinator: Arc<AtomicBool>,
518}
519
520pub(super) struct StreamingCoordinatorGeneration {
521    runtime: StreamingCoordinatorRuntime,
522}
523
524impl Drop for StreamingCoordinatorGeneration {
525    fn drop(&mut self) {
526        self.runtime
527            .active_coordinator
528            .store(false, Ordering::Release);
529    }
530}
531
532impl StreamingCoordinatorRuntime {
533    /// Creates an empty runtime owner with no active source generation.
534    #[must_use]
535    pub fn new() -> Self {
536        Self {
537            owned_source_tasks: Arc::new(parking_lot::Mutex::new(Vec::new())),
538            owned_connector_task_fences: Arc::new(parking_lot::Mutex::new(Vec::new())),
539            construction: Arc::new(tokio::sync::Mutex::new(())),
540            active_coordinator: Arc::new(AtomicBool::new(false)),
541        }
542    }
543
544    pub(super) fn claim_generation(&self) -> Result<StreamingCoordinatorGeneration, DbError> {
545        self.active_coordinator
546            .compare_exchange(false, true, Ordering::AcqRel, Ordering::Acquire)
547            .map_err(|_| {
548                DbError::Pipeline(
549                    "cannot construct a replacement streaming coordinator while the prior \
550                     coordinator generation is still active"
551                        .into(),
552                )
553            })?;
554        Ok(StreamingCoordinatorGeneration {
555            runtime: self.clone(),
556        })
557    }
558
559    pub(super) fn prune_and_require_idle(&self) -> Result<(), DbError> {
560        let source_names = {
561            let mut sources = self.owned_source_tasks.lock();
562            sources.retain(|source| !source.is_finished());
563            sources
564                .iter()
565                .map(|source| source.name().to_owned())
566                .collect::<Vec<_>>()
567        };
568        let connector_names = {
569            let mut connectors = self.owned_connector_task_fences.lock();
570            connectors.retain(|connector| !connector.is_finished());
571            connectors
572                .iter()
573                .map(|connector| connector.name().to_owned())
574                .collect::<Vec<_>>()
575        };
576        if source_names.is_empty() && connector_names.is_empty() {
577            return Ok(());
578        }
579        Err(DbError::Pipeline(format!(
580            "cannot construct a replacement streaming coordinator while prior connector \
581             generations remain unresolved: sources=[{}], connector_tasks=[{}]",
582            source_names.join(", "),
583            connector_names.join(", ")
584        )))
585    }
586}
587
588impl Default for StreamingCoordinatorRuntime {
589    fn default() -> Self {
590        Self::new()
591    }
592}
laminar_db/pipeline/streaming_coordinator/source_runtime.rs

laminar_db/pipeline/streaming_coordinator/
source_runtime.rs
