risingwave_stream/task/barrier_worker/
managed_state.rs

1// Copyright 2025 RisingWave Labs
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use std::collections::{BTreeMap, BTreeSet, HashMap, HashSet, VecDeque};
16use std::fmt::{Debug, Display, Formatter};
17use std::future::{Future, pending, poll_fn};
18use std::mem::replace;
19use std::sync::Arc;
20use std::task::{Context, Poll};
21use std::time::{Duration, Instant};
22
23use anyhow::anyhow;
24use futures::future::BoxFuture;
25use futures::stream::{FuturesOrdered, FuturesUnordered};
26use futures::{FutureExt, StreamExt};
27use prometheus::HistogramTimer;
28use risingwave_common::catalog::TableId;
29use risingwave_common::id::{SinkId, SourceId};
30use risingwave_common::util::epoch::EpochPair;
31use risingwave_pb::connector_service::SinkMetadata;
32use risingwave_pb::stream_plan::barrier::BarrierKind;
33use risingwave_pb::stream_service::barrier_complete_response::{
34    IcebergPkIndexSinkMetadata as PbIcebergPkIndexSinkMetadata, PbCdcSourceOffsetUpdated,
35    PbCdcTableBackfillProgress, PbCreateMviewProgress, PbListFinishedSource, PbLoadFinishedSource,
36};
37use risingwave_storage::StateStoreImpl;
38use tokio::sync::mpsc;
39use tokio::sync::mpsc::{UnboundedReceiver, UnboundedSender, unbounded_channel};
40use tokio::task::JoinHandle;
41
42use crate::error::{StreamError, StreamResult};
43use crate::executor::Barrier;
44use crate::executor::monitor::StreamingMetrics;
45use crate::task::progress::BackfillState;
46use crate::task::{
47    ActorId, LocalBarrierEvent, LocalBarrierManager, NewOutputRequest, PartialGraphId,
48    StreamActorManager, UpDownActorIds,
49};
50
51struct IssuedState {
52    /// Actor ids remaining to be collected.
53    pub remaining_actors: BTreeSet<ActorId>,
54
55    pub barrier_inflight_latency: HistogramTimer,
56}
57
58impl Debug for IssuedState {
59    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
60        f.debug_struct("IssuedState")
61            .field("remaining_actors", &self.remaining_actors)
62            .finish()
63    }
64}
65
66/// The state machine of local barrier manager.
67#[derive(Debug)]
68enum ManagedBarrierStateInner {
69    /// Meta service has issued a `send_barrier` request. We're collecting barriers now.
70    Issued(IssuedState),
71
72    /// The barrier has been collected by all remaining actors
73    AllCollected {
74        create_mview_progress: Vec<PbCreateMviewProgress>,
75        list_finished_source_ids: Vec<PbListFinishedSource>,
76        load_finished_source_ids: Vec<PbLoadFinishedSource>,
77        cdc_table_backfill_progress: Vec<PbCdcTableBackfillProgress>,
78        cdc_source_offset_updated: Vec<PbCdcSourceOffsetUpdated>,
79        iceberg_pk_index_sink_metadata: Vec<PbIcebergPkIndexSinkMetadata>,
80        truncate_tables: Vec<TableId>,
81        refresh_finished_tables: Vec<TableId>,
82    },
83}
84
85#[derive(Debug)]
86struct BarrierState {
87    barrier: Barrier,
88    /// Only be `Some(_)` when `barrier.kind` is `Checkpoint`
89    table_ids: Option<HashSet<TableId>>,
90    inner: ManagedBarrierStateInner,
91}
92
93use risingwave_common::must_match;
94use risingwave_pb::id::FragmentId;
95use risingwave_pb::stream_service::{InjectBarrierRequest, PbIcebergPkIndexSinkRole};
96
97use crate::executor::exchange::permit;
98use crate::task::barrier_worker::await_epoch_completed_future::AwaitEpochCompletedFuture;
99use crate::task::barrier_worker::{ScoredStreamError, TakeReceiverRequest};
100use crate::task::cdc_progress::CdcTableBackfillState;
101
102pub(super) struct ManagedBarrierStateDebugInfo<'a> {
103    running_actors: BTreeSet<ActorId>,
104    graph_state: &'a PartialGraphManagedBarrierState,
105}
106
107impl Display for ManagedBarrierStateDebugInfo<'_> {
108    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
109        write!(f, "running_actors: ")?;
110        for actor_id in &self.running_actors {
111            write!(f, "{}, ", actor_id)?;
112        }
113        {
114            writeln!(f, "graph states")?;
115            write!(f, "{}", self.graph_state)?;
116        }
117        Ok(())
118    }
119}
120
121impl Display for &'_ PartialGraphManagedBarrierState {
122    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
123        let mut prev_epoch = 0u64;
124        for (epoch, barrier_state) in &self.epoch_barrier_state_map {
125            write!(f, "> Epoch {}: ", epoch)?;
126            match &barrier_state.inner {
127                ManagedBarrierStateInner::Issued(state) => {
128                    write!(
129                        f,
130                        "Issued [{:?}]. Remaining actors: [",
131                        barrier_state.barrier.kind
132                    )?;
133                    let mut is_prev_epoch_issued = false;
134                    if prev_epoch != 0 {
135                        let bs = &self.epoch_barrier_state_map[&prev_epoch];
136                        if let ManagedBarrierStateInner::Issued(IssuedState {
137                            remaining_actors: remaining_actors_prev,
138                            ..
139                        }) = &bs.inner
140                        {
141                            // Only show the actors that are not in the previous epoch.
142                            is_prev_epoch_issued = true;
143                            let mut duplicates = 0usize;
144                            for actor_id in &state.remaining_actors {
145                                if !remaining_actors_prev.contains(actor_id) {
146                                    write!(f, "{}, ", actor_id)?;
147                                } else {
148                                    duplicates += 1;
149                                }
150                            }
151                            if duplicates > 0 {
152                                write!(f, "...and {} actors in prev epoch", duplicates)?;
153                            }
154                        }
155                    }
156                    if !is_prev_epoch_issued {
157                        for actor_id in &state.remaining_actors {
158                            write!(f, "{}, ", actor_id)?;
159                        }
160                    }
161                    write!(f, "]")?;
162                }
163                ManagedBarrierStateInner::AllCollected { .. } => {
164                    write!(f, "AllCollected")?;
165                }
166            }
167            prev_epoch = *epoch;
168            writeln!(f)?;
169        }
170
171        if !self.create_mview_progress.is_empty() {
172            writeln!(f, "Create MView Progress:")?;
173            for (epoch, progress) in &self.create_mview_progress {
174                write!(f, "> Epoch {}:", epoch)?;
175                for (actor_id, (_, state)) in progress {
176                    write!(f, ">> Actor {}: {}, ", actor_id, state)?;
177                }
178            }
179        }
180
181        Ok(())
182    }
183}
184
185enum InflightActorStatus {
186    /// The actor has been issued some barriers, but has not collected the first barrier
187    IssuedFirst(Vec<Barrier>),
188    /// The actor has been issued some barriers, and has collected the first barrier
189    Running(u64),
190}
191
192impl InflightActorStatus {
193    fn max_issued_epoch(&self) -> u64 {
194        match self {
195            InflightActorStatus::Running(epoch) => *epoch,
196            InflightActorStatus::IssuedFirst(issued_barriers) => {
197                issued_barriers.last().expect("non-empty").epoch.prev
198            }
199        }
200    }
201}
202
203pub(crate) struct InflightActorState {
204    actor_id: ActorId,
205    barrier_senders: Vec<mpsc::UnboundedSender<Barrier>>,
206    /// `prev_epoch`. `push_back` and `pop_front`
207    pub(in crate::task) inflight_barriers: VecDeque<u64>,
208    status: InflightActorStatus,
209    /// Whether the actor has been issued a stop barrier
210    is_stopping: bool,
211
212    new_output_request_tx: UnboundedSender<(ActorId, NewOutputRequest)>,
213    join_handle: JoinHandle<()>,
214    monitor_task_handle: Option<JoinHandle<()>>,
215}
216
217impl InflightActorState {
218    pub(super) fn start(
219        actor_id: ActorId,
220        initial_barrier: &Barrier,
221        new_output_request_tx: UnboundedSender<(ActorId, NewOutputRequest)>,
222        join_handle: JoinHandle<()>,
223        monitor_task_handle: Option<JoinHandle<()>>,
224    ) -> Self {
225        Self {
226            actor_id,
227            barrier_senders: vec![],
228            inflight_barriers: VecDeque::from_iter([initial_barrier.epoch.prev]),
229            status: InflightActorStatus::IssuedFirst(vec![initial_barrier.clone()]),
230            is_stopping: false,
231            new_output_request_tx,
232            join_handle,
233            monitor_task_handle,
234        }
235    }
236
237    pub(super) fn issue_barrier(&mut self, barrier: &Barrier, is_stop: bool) -> StreamResult<()> {
238        assert!(barrier.epoch.prev > self.status.max_issued_epoch());
239
240        for barrier_sender in &self.barrier_senders {
241            barrier_sender.send(barrier.clone()).map_err(|_| {
242                StreamError::barrier_send(
243                    barrier.clone(),
244                    self.actor_id,
245                    "failed to send to registered sender",
246                )
247            })?;
248        }
249
250        if let Some(prev_epoch) = self.inflight_barriers.back() {
251            assert!(*prev_epoch < barrier.epoch.prev);
252        }
253        self.inflight_barriers.push_back(barrier.epoch.prev);
254
255        match &mut self.status {
256            InflightActorStatus::IssuedFirst(pending_barriers) => {
257                pending_barriers.push(barrier.clone());
258            }
259            InflightActorStatus::Running(prev_epoch) => {
260                *prev_epoch = barrier.epoch.prev;
261            }
262        };
263
264        if is_stop {
265            assert!(!self.is_stopping, "stopped actor should not issue barrier");
266            self.is_stopping = true;
267        }
268        Ok(())
269    }
270
271    pub(super) fn collect(&mut self, epoch: EpochPair) -> bool {
272        let prev_epoch = self.inflight_barriers.pop_front().expect("should exist");
273        assert_eq!(prev_epoch, epoch.prev);
274        match &self.status {
275            InflightActorStatus::IssuedFirst(pending_barriers) => {
276                assert_eq!(
277                    prev_epoch,
278                    pending_barriers.first().expect("non-empty").epoch.prev
279                );
280                self.status = InflightActorStatus::Running(
281                    pending_barriers.last().expect("non-empty").epoch.prev,
282                );
283            }
284            InflightActorStatus::Running(_) => {}
285        }
286
287        self.inflight_barriers.is_empty() && self.is_stopping
288    }
289}
290
291/// Part of [`PartialGraphState`]
292pub(crate) struct PartialGraphManagedBarrierState {
293    /// Record barrier state for each epoch of concurrent checkpoints.
294    ///
295    /// The key is `prev_epoch`, and the first value is `curr_epoch`
296    epoch_barrier_state_map: BTreeMap<u64, BarrierState>,
297
298    prev_barrier_table_ids: Option<(EpochPair, HashSet<TableId>)>,
299
300    /// Record the progress updates of creating mviews for each epoch of concurrent checkpoints.
301    ///
302    /// The process of progress reporting is as follows:
303    /// 1. updated by [`crate::task::barrier_manager::CreateMviewProgressReporter::update`]
304    /// 2. converted to [`ManagedBarrierStateInner`] in [`Self::may_have_collected_all`]
305    /// 3. handled by [`Self::pop_barrier_to_complete`]
306    /// 4. put in [`crate::task::barrier_worker::BarrierCompleteResult`] and reported to meta.
307    pub(crate) create_mview_progress: HashMap<u64, HashMap<ActorId, (FragmentId, BackfillState)>>,
308
309    /// Record the source list finished reports for each epoch of concurrent checkpoints.
310    /// Used for refreshable batch source. The map key is epoch and the value is
311    /// a list of pb messages reported by actors.
312    pub(crate) list_finished_source_ids: HashMap<u64, Vec<PbListFinishedSource>>,
313
314    /// Record the source load finished reports for each epoch of concurrent checkpoints.
315    /// Used for refreshable batch source. The map key is epoch and the value is
316    /// a list of pb messages reported by actors.
317    pub(crate) load_finished_source_ids: HashMap<u64, Vec<PbLoadFinishedSource>>,
318
319    pub(crate) cdc_table_backfill_progress: HashMap<u64, HashMap<ActorId, CdcTableBackfillState>>,
320
321    /// Record CDC source offset updated reports for each epoch of concurrent checkpoints.
322    /// Used to track when CDC sources have updated their offset at least once.
323    pub(crate) cdc_source_offset_updated: HashMap<u64, Vec<PbCdcSourceOffsetUpdated>>,
324
325    /// Record Iceberg pk-index sink metadata reports per epoch for concurrent checkpoints.
326    pub(crate) iceberg_pk_index_sink_metadata: HashMap<u64, Vec<PbIcebergPkIndexSinkMetadata>>,
327
328    /// Record the tables to truncate for each epoch of concurrent checkpoints.
329    pub(crate) truncate_tables: HashMap<u64, HashSet<TableId>>,
330    /// Record the tables that have finished refresh for each epoch of concurrent checkpoints.
331    /// Used for materialized view refresh completion reporting.
332    pub(crate) refresh_finished_tables: HashMap<u64, HashSet<TableId>>,
333
334    state_store: StateStoreImpl,
335
336    streaming_metrics: Arc<StreamingMetrics>,
337}
338
339impl PartialGraphManagedBarrierState {
340    pub(super) fn new(actor_manager: &StreamActorManager) -> Self {
341        Self::new_inner(
342            actor_manager.env.state_store(),
343            actor_manager.streaming_metrics.clone(),
344        )
345    }
346
347    fn new_inner(state_store: StateStoreImpl, streaming_metrics: Arc<StreamingMetrics>) -> Self {
348        Self {
349            epoch_barrier_state_map: Default::default(),
350            prev_barrier_table_ids: None,
351            create_mview_progress: Default::default(),
352            list_finished_source_ids: Default::default(),
353            load_finished_source_ids: Default::default(),
354            cdc_table_backfill_progress: Default::default(),
355            cdc_source_offset_updated: Default::default(),
356            iceberg_pk_index_sink_metadata: Default::default(),
357            truncate_tables: Default::default(),
358            refresh_finished_tables: Default::default(),
359            state_store,
360            streaming_metrics,
361        }
362    }
363
364    #[cfg(test)]
365    pub(crate) fn for_test() -> Self {
366        Self::new_inner(
367            StateStoreImpl::for_test(),
368            Arc::new(StreamingMetrics::unused()),
369        )
370    }
371
372    pub(super) fn is_empty(&self) -> bool {
373        self.epoch_barrier_state_map.is_empty()
374    }
375}
376
377pub(crate) struct SuspendedPartialGraphState {
378    pub(super) suspend_time: Instant,
379    inner: PartialGraphState,
380    failure: Option<(Option<ActorId>, StreamError)>,
381}
382
383impl SuspendedPartialGraphState {
384    fn new(
385        state: PartialGraphState,
386        failure: Option<(Option<ActorId>, StreamError)>,
387        _completing_futures: Option<FuturesOrdered<AwaitEpochCompletedFuture>>, /* discard the completing futures */
388    ) -> Self {
389        Self {
390            suspend_time: Instant::now(),
391            inner: state,
392            failure,
393        }
394    }
395
396    async fn reset(mut self) -> ResetPartialGraphOutput {
397        let root_err = self.inner.try_find_root_actor_failure(self.failure).await;
398        self.inner.abort_and_wait_actors().await;
399        ResetPartialGraphOutput { root_err }
400    }
401}
402
403pub(crate) struct ResetPartialGraphOutput {
404    pub(crate) root_err: Option<ScoredStreamError>,
405}
406
407pub(in crate::task) enum PartialGraphStatus {
408    ReceivedExchangeRequest(Vec<(UpDownActorIds, TakeReceiverRequest)>),
409    Running(PartialGraphState),
410    Suspended(SuspendedPartialGraphState),
411    Resetting,
412    /// temporary place holder
413    Unspecified,
414}
415
416impl PartialGraphStatus {
417    pub(crate) async fn abort(&mut self) {
418        match self {
419            PartialGraphStatus::ReceivedExchangeRequest(pending_requests) => {
420                for (_, request) in pending_requests.drain(..) {
421                    if let TakeReceiverRequest::Remote { result_sender, .. } = request {
422                        let _ = result_sender.send(Err(anyhow!("partial graph aborted").into()));
423                    }
424                }
425            }
426            PartialGraphStatus::Running(state) => {
427                state.abort_and_wait_actors().await;
428            }
429            PartialGraphStatus::Suspended(SuspendedPartialGraphState { inner: state, .. }) => {
430                state.abort_and_wait_actors().await;
431            }
432            PartialGraphStatus::Resetting => {}
433            PartialGraphStatus::Unspecified => {
434                unreachable!()
435            }
436        }
437    }
438
439    pub(crate) fn state_for_request(&mut self) -> Option<&mut PartialGraphState> {
440        match self {
441            PartialGraphStatus::ReceivedExchangeRequest(_) => {
442                unreachable!("should not handle request")
443            }
444            PartialGraphStatus::Running(state) => Some(state),
445            PartialGraphStatus::Suspended(_) => None,
446            PartialGraphStatus::Resetting => {
447                unreachable!("should not receive further request during cleaning")
448            }
449            PartialGraphStatus::Unspecified => {
450                unreachable!()
451            }
452        }
453    }
454
455    pub(super) fn poll_next_event(
456        &mut self,
457        cx: &mut Context<'_>,
458    ) -> Poll<ManagedBarrierStateEvent> {
459        match self {
460            PartialGraphStatus::ReceivedExchangeRequest(_) => Poll::Pending,
461            PartialGraphStatus::Running(state) => state.poll_next_event(cx),
462            PartialGraphStatus::Suspended(_) | PartialGraphStatus::Resetting => Poll::Pending,
463            PartialGraphStatus::Unspecified => {
464                unreachable!()
465            }
466        }
467    }
468
469    pub(super) fn suspend(
470        &mut self,
471        failed_actor: Option<ActorId>,
472        err: StreamError,
473        completing_futures: Option<FuturesOrdered<AwaitEpochCompletedFuture>>,
474    ) {
475        let state = must_match!(replace(self, PartialGraphStatus::Unspecified), PartialGraphStatus::Running(state) => state);
476        *self = PartialGraphStatus::Suspended(SuspendedPartialGraphState::new(
477            state,
478            Some((failed_actor, err)),
479            completing_futures,
480        ));
481    }
482
483    pub(super) fn start_reset(
484        &mut self,
485        partial_graph_id: PartialGraphId,
486        completing_futures: Option<FuturesOrdered<AwaitEpochCompletedFuture>>,
487        table_ids_to_clear: &mut HashSet<TableId>,
488    ) -> BoxFuture<'static, ResetPartialGraphOutput> {
489        match replace(self, PartialGraphStatus::Resetting) {
490            PartialGraphStatus::ReceivedExchangeRequest(pending_requests) => {
491                for (_, request) in pending_requests {
492                    if let TakeReceiverRequest::Remote { result_sender, .. } = request {
493                        let _ = result_sender.send(Err(anyhow!("partial graph reset").into()));
494                    }
495                }
496                async move { ResetPartialGraphOutput { root_err: None } }.boxed()
497            }
498            PartialGraphStatus::Running(state) => {
499                assert_eq!(partial_graph_id, state.partial_graph_id);
500                info!(
501                    %partial_graph_id,
502                    "start partial graph reset from Running"
503                );
504                table_ids_to_clear.extend(state.table_ids.iter().copied());
505                SuspendedPartialGraphState::new(state, None, completing_futures)
506                    .reset()
507                    .boxed()
508            }
509            PartialGraphStatus::Suspended(state) => {
510                assert!(
511                    completing_futures.is_none(),
512                    "should have been clear when suspended"
513                );
514                assert_eq!(partial_graph_id, state.inner.partial_graph_id);
515                info!(
516                    %partial_graph_id,
517                    suspend_elapsed = ?state.suspend_time.elapsed(),
518                    "start partial graph reset after suspended"
519                );
520                table_ids_to_clear.extend(state.inner.table_ids.iter().copied());
521                state.reset().boxed()
522            }
523            PartialGraphStatus::Resetting => {
524                unreachable!("should not reset for twice");
525            }
526            PartialGraphStatus::Unspecified => {
527                unreachable!()
528            }
529        }
530    }
531}
532
533#[derive(Default)]
534pub(in crate::task) struct ManagedBarrierState {
535    pub(super) partial_graphs: HashMap<PartialGraphId, PartialGraphStatus>,
536    pub(super) resetting_graphs:
537        FuturesUnordered<JoinHandle<Vec<(PartialGraphId, ResetPartialGraphOutput)>>>,
538}
539
540pub(super) enum ManagedBarrierStateEvent {
541    BarrierCollected {
542        partial_graph_id: PartialGraphId,
543        barrier: Barrier,
544    },
545    ActorError {
546        partial_graph_id: PartialGraphId,
547        actor_id: ActorId,
548        err: StreamError,
549    },
550    PartialGraphsReset(Vec<(PartialGraphId, ResetPartialGraphOutput)>),
551    RegisterLocalUpstreamOutput {
552        actor_id: ActorId,
553        upstream_actor_id: ActorId,
554        upstream_partial_graph_id: PartialGraphId,
555        tx: permit::Sender,
556    },
557}
558
559impl ManagedBarrierState {
560    pub(super) fn next_event(&mut self) -> impl Future<Output = ManagedBarrierStateEvent> + '_ {
561        poll_fn(|cx| {
562            for graph in self.partial_graphs.values_mut() {
563                if let Poll::Ready(event) = graph.poll_next_event(cx) {
564                    return Poll::Ready(event);
565                }
566            }
567            if let Poll::Ready(Some(result)) = self.resetting_graphs.poll_next_unpin(cx) {
568                let outputs = result.expect("failed to join resetting future");
569                for (partial_graph_id, _) in &outputs {
570                    let PartialGraphStatus::Resetting = self
571                        .partial_graphs
572                        .remove(partial_graph_id)
573                        .expect("should exist")
574                    else {
575                        panic!("should be resetting")
576                    };
577                }
578                return Poll::Ready(ManagedBarrierStateEvent::PartialGraphsReset(outputs));
579            }
580            Poll::Pending
581        })
582    }
583}
584
585/// Per-partial-graph barrier state manager. Handles barriers for one specific partial graph.
586/// Part of [`ManagedBarrierState`] in [`super::LocalBarrierWorker`].
587///
588/// See [`crate::task`] for architecture overview.
589pub(crate) struct PartialGraphState {
590    partial_graph_id: PartialGraphId,
591    pub(crate) actor_states: HashMap<ActorId, InflightActorState>,
592    pub(super) actor_pending_new_output_requests:
593        HashMap<ActorId, Vec<(ActorId, NewOutputRequest)>>,
594
595    pub(crate) graph_state: PartialGraphManagedBarrierState,
596
597    table_ids: HashSet<TableId>,
598
599    actor_manager: Arc<StreamActorManager>,
600
601    pub(super) local_barrier_manager: LocalBarrierManager,
602
603    barrier_event_rx: UnboundedReceiver<LocalBarrierEvent>,
604    pub(super) actor_failure_rx: UnboundedReceiver<(ActorId, StreamError)>,
605}
606
607impl PartialGraphState {
608    /// Create a barrier manager state. This will be called only once.
609    pub(super) fn new(
610        partial_graph_id: PartialGraphId,
611        term_id: String,
612        actor_manager: Arc<StreamActorManager>,
613    ) -> Self {
614        let (local_barrier_manager, barrier_event_rx, actor_failure_rx) =
615            LocalBarrierManager::new(term_id, actor_manager.env.clone());
616        Self {
617            partial_graph_id,
618            actor_states: Default::default(),
619            actor_pending_new_output_requests: Default::default(),
620            graph_state: PartialGraphManagedBarrierState::new(&actor_manager),
621            table_ids: Default::default(),
622            actor_manager,
623            local_barrier_manager,
624            barrier_event_rx,
625            actor_failure_rx,
626        }
627    }
628
629    pub(super) fn to_debug_info(&self) -> ManagedBarrierStateDebugInfo<'_> {
630        ManagedBarrierStateDebugInfo {
631            running_actors: self.actor_states.keys().cloned().collect(),
632            graph_state: &self.graph_state,
633        }
634    }
635
636    async fn abort_and_wait_actors(&mut self) {
637        for (actor_id, state) in &self.actor_states {
638            tracing::debug!("force stopping actor {}", actor_id);
639            state.join_handle.abort();
640            if let Some(monitor_task_handle) = &state.monitor_task_handle {
641                monitor_task_handle.abort();
642            }
643        }
644
645        for (actor_id, state) in self.actor_states.drain() {
646            tracing::debug!("join actor {}", actor_id);
647            let result = state.join_handle.await;
648            assert!(result.is_ok() || result.unwrap_err().is_cancelled());
649        }
650    }
651}
652
653impl InflightActorState {
654    pub(super) fn register_barrier_sender(
655        &mut self,
656        tx: mpsc::UnboundedSender<Barrier>,
657    ) -> StreamResult<()> {
658        match &self.status {
659            InflightActorStatus::IssuedFirst(pending_barriers) => {
660                for barrier in pending_barriers {
661                    tx.send(barrier.clone()).map_err(|_| {
662                        StreamError::barrier_send(
663                            barrier.clone(),
664                            self.actor_id,
665                            "failed to send pending barriers to newly registered sender",
666                        )
667                    })?;
668                }
669                self.barrier_senders.push(tx);
670            }
671            InflightActorStatus::Running(_) => {
672                unreachable!("should not register barrier sender when entering Running status")
673            }
674        }
675        Ok(())
676    }
677}
678
679impl PartialGraphState {
680    pub(super) fn register_barrier_sender(
681        &mut self,
682        actor_id: ActorId,
683        tx: mpsc::UnboundedSender<Barrier>,
684    ) -> StreamResult<()> {
685        self.actor_states
686            .get_mut(&actor_id)
687            .expect("should exist")
688            .register_barrier_sender(tx)
689    }
690}
691
692impl PartialGraphState {
693    pub(super) fn transform_to_issued(
694        &mut self,
695        barrier: &Barrier,
696        request: InjectBarrierRequest,
697    ) -> StreamResult<()> {
698        assert_eq!(self.partial_graph_id, request.partial_graph_id);
699        let actor_to_stop = barrier.all_stop_actors();
700        let is_stop_actor = |actor_id| {
701            actor_to_stop
702                .map(|actors| actors.contains(&actor_id))
703                .unwrap_or(false)
704        };
705
706        let table_ids = HashSet::from_iter(request.table_ids_to_sync);
707        self.table_ids.extend(table_ids.iter().cloned());
708
709        self.graph_state.transform_to_issued(
710            barrier,
711            request.actor_ids_to_collect.iter().copied(),
712            table_ids,
713        );
714
715        let mut new_actors = HashSet::new();
716        for (node, fragment_id, actor) in
717            request
718                .actors_to_build
719                .into_iter()
720                .flat_map(|fragment_actors| {
721                    let node = Arc::new(fragment_actors.node.unwrap());
722                    fragment_actors
723                        .actors
724                        .into_iter()
725                        .map(move |actor| (node.clone(), fragment_actors.fragment_id, actor))
726                })
727        {
728            let actor_id = actor.actor_id;
729            assert!(!is_stop_actor(actor_id));
730            assert!(new_actors.insert(actor_id));
731            assert!(request.actor_ids_to_collect.contains(&actor_id));
732            let (new_output_request_tx, new_output_request_rx) = unbounded_channel();
733            if let Some(pending_requests) = self.actor_pending_new_output_requests.remove(&actor_id)
734            {
735                for request in pending_requests {
736                    let _ = new_output_request_tx.send(request);
737                }
738            }
739            let (join_handle, monitor_join_handle) = self.actor_manager.spawn_actor(
740                actor,
741                fragment_id,
742                node,
743                self.local_barrier_manager.clone(),
744                new_output_request_rx,
745            );
746            assert!(
747                self.actor_states
748                    .try_insert(
749                        actor_id,
750                        InflightActorState::start(
751                            actor_id,
752                            barrier,
753                            new_output_request_tx,
754                            join_handle,
755                            monitor_join_handle
756                        )
757                    )
758                    .is_ok()
759            );
760        }
761
762        // Spawn a trivial join handle to be compatible with the unit test. In the unit tests that involve local barrier manager,
763        // actors are spawned in the local test logic, but we assume that there is an entry for each spawned actor in ·actor_states`,
764        // so under cfg!(test) we add a dummy entry for each new actor.
765        if cfg!(test) {
766            for &actor_id in &request.actor_ids_to_collect {
767                if !self.actor_states.contains_key(&actor_id) {
768                    let (tx, rx) = unbounded_channel();
769                    let join_handle = self.actor_manager.runtime.spawn(async move {
770                        // The rx is spawned so that tx.send() will not fail.
771                        let _ = rx;
772                        pending().await
773                    });
774                    assert!(
775                        self.actor_states
776                            .try_insert(
777                                actor_id,
778                                InflightActorState::start(actor_id, barrier, tx, join_handle, None,)
779                            )
780                            .is_ok()
781                    );
782                    new_actors.insert(actor_id);
783                }
784            }
785        }
786
787        // Note: it's important to issue barrier to actor after issuing to graph to ensure that
788        // we call `start_epoch` on the graph before the actors receive the barrier
789        for &actor_id in &request.actor_ids_to_collect {
790            if new_actors.contains(&actor_id) {
791                continue;
792            }
793            self.actor_states
794                .get_mut(&actor_id)
795                .unwrap_or_else(|| {
796                    panic!(
797                        "should exist: {} {:?}",
798                        actor_id, request.actor_ids_to_collect
799                    );
800                })
801                .issue_barrier(barrier, is_stop_actor(actor_id))?;
802        }
803
804        Ok(())
805    }
806
807    pub(super) fn new_actor_output_request(
808        &mut self,
809        actor_id: ActorId,
810        upstream_actor_id: ActorId,
811        request: TakeReceiverRequest,
812    ) {
813        let request = match request {
814            TakeReceiverRequest::Remote {
815                result_sender,
816                upstream_fragment_id,
817            } => {
818                let upstream_fragment_id_str = upstream_fragment_id.to_string();
819                let fragment_channel_buffered_bytes = self
820                    .actor_manager
821                    .streaming_metrics
822                    .fragment_channel_buffered_bytes
823                    .with_guarded_label_values(&[&upstream_fragment_id_str]);
824                let (tx, rx) = permit::channel_from_config_with_metrics(
825                    self.local_barrier_manager.env.global_config(),
826                    permit::ChannelMetrics {
827                        sender_actor_channel_buffered_bytes: fragment_channel_buffered_bytes
828                            .clone(),
829                        receiver_actor_channel_buffered_bytes: fragment_channel_buffered_bytes,
830                    },
831                );
832                let _ = result_sender.send(Ok(rx));
833                NewOutputRequest::Remote(tx)
834            }
835            TakeReceiverRequest::Local(tx) => NewOutputRequest::Local(tx),
836        };
837        if let Some(actor) = self.actor_states.get_mut(&upstream_actor_id) {
838            let _ = actor.new_output_request_tx.send((actor_id, request));
839        } else {
840            self.actor_pending_new_output_requests
841                .entry(upstream_actor_id)
842                .or_default()
843                .push((actor_id, request));
844        }
845    }
846
847    /// Handles [`LocalBarrierEvent`] from [`crate::task::barrier_manager::LocalBarrierManager`].
848    pub(super) fn poll_next_event(
849        &mut self,
850        cx: &mut Context<'_>,
851    ) -> Poll<ManagedBarrierStateEvent> {
852        if let Poll::Ready(option) = self.actor_failure_rx.poll_recv(cx) {
853            let (actor_id, err) = option.expect("non-empty when tx in local_barrier_manager");
854            return Poll::Ready(ManagedBarrierStateEvent::ActorError {
855                actor_id,
856                err,
857                partial_graph_id: self.partial_graph_id,
858            });
859        }
860        // yield some pending collected epochs
861        {
862            if let Some(barrier) = self.graph_state.may_have_collected_all() {
863                return Poll::Ready(ManagedBarrierStateEvent::BarrierCollected {
864                    barrier,
865                    partial_graph_id: self.partial_graph_id,
866                });
867            }
868        }
869        while let Poll::Ready(event) = self.barrier_event_rx.poll_recv(cx) {
870            match event.expect("non-empty when tx in local_barrier_manager") {
871                LocalBarrierEvent::ReportActorCollected { actor_id, epoch } => {
872                    if let Some(barrier) = self.collect(actor_id, epoch) {
873                        return Poll::Ready(ManagedBarrierStateEvent::BarrierCollected {
874                            barrier,
875                            partial_graph_id: self.partial_graph_id,
876                        });
877                    }
878                }
879                LocalBarrierEvent::ReportCreateProgress {
880                    epoch,
881                    fragment_id,
882                    actor,
883                    state,
884                } => {
885                    self.update_create_mview_progress(epoch, fragment_id, actor, state);
886                }
887                LocalBarrierEvent::ReportSourceListFinished {
888                    epoch,
889                    actor_id,
890                    table_id,
891                    associated_source_id,
892                } => {
893                    self.report_source_list_finished(
894                        epoch,
895                        actor_id,
896                        table_id,
897                        associated_source_id,
898                    );
899                }
900                LocalBarrierEvent::ReportSourceLoadFinished {
901                    epoch,
902                    actor_id,
903                    table_id,
904                    associated_source_id,
905                } => {
906                    self.report_source_load_finished(
907                        epoch,
908                        actor_id,
909                        table_id,
910                        associated_source_id,
911                    );
912                }
913                LocalBarrierEvent::RefreshFinished {
914                    epoch,
915                    actor_id,
916                    table_id,
917                    staging_table_id,
918                } => {
919                    self.report_refresh_finished(epoch, actor_id, table_id, staging_table_id);
920                }
921                LocalBarrierEvent::RegisterBarrierSender {
922                    actor_id,
923                    barrier_sender,
924                } => {
925                    if let Err(err) = self.register_barrier_sender(actor_id, barrier_sender) {
926                        return Poll::Ready(ManagedBarrierStateEvent::ActorError {
927                            actor_id,
928                            err,
929                            partial_graph_id: self.partial_graph_id,
930                        });
931                    }
932                }
933                LocalBarrierEvent::RegisterLocalUpstreamOutput {
934                    actor_id,
935                    upstream_actor_id,
936                    upstream_partial_graph_id,
937                    tx,
938                } => {
939                    return Poll::Ready(ManagedBarrierStateEvent::RegisterLocalUpstreamOutput {
940                        actor_id,
941                        upstream_actor_id,
942                        upstream_partial_graph_id,
943                        tx,
944                    });
945                }
946                LocalBarrierEvent::ReportCdcTableBackfillProgress {
947                    actor_id,
948                    epoch,
949                    state,
950                } => {
951                    self.update_cdc_table_backfill_progress(epoch, actor_id, state);
952                }
953                LocalBarrierEvent::ReportCdcSourceOffsetUpdated {
954                    epoch,
955                    actor_id,
956                    source_id,
957                } => {
958                    self.report_cdc_source_offset_updated(epoch, actor_id, source_id);
959                }
960                LocalBarrierEvent::ReportIcebergPkIndexSinkMetadata {
961                    epoch,
962                    sink_id,
963                    actor_id,
964                    role,
965                    metadata,
966                } => {
967                    self.report_iceberg_pk_index_sink_metadata(
968                        epoch, sink_id, actor_id, role, metadata,
969                    );
970                }
971            }
972        }
973
974        debug_assert!(self.graph_state.may_have_collected_all().is_none());
975        Poll::Pending
976    }
977}
978
979impl PartialGraphState {
980    #[must_use]
981    pub(super) fn collect(&mut self, actor_id: ActorId, epoch: EpochPair) -> Option<Barrier> {
982        let is_finished = self
983            .actor_states
984            .get_mut(&actor_id)
985            .expect("should exist")
986            .collect(epoch);
987        if is_finished {
988            let state = self.actor_states.remove(&actor_id).expect("should exist");
989            if let Some(monitor_task_handle) = state.monitor_task_handle {
990                monitor_task_handle.abort();
991            }
992        }
993        self.graph_state.collect(actor_id, epoch);
994        self.graph_state.may_have_collected_all()
995    }
996
997    pub(super) fn pop_barrier_to_complete(&mut self, prev_epoch: u64) -> BarrierToComplete {
998        self.graph_state.pop_barrier_to_complete(prev_epoch)
999    }
1000
1001    /// Collect actor errors for a while and find the one that might be the root cause.
1002    ///
1003    /// Returns `None` if there's no actor error received.
1004    async fn try_find_root_actor_failure(
1005        &mut self,
1006        first_failure: Option<(Option<ActorId>, StreamError)>,
1007    ) -> Option<ScoredStreamError> {
1008        let mut later_errs = vec![];
1009        // fetch more actor errors within a timeout
1010        let _ = tokio::time::timeout(Duration::from_secs(3), async {
1011            let mut uncollected_actors: HashSet<_> = self.actor_states.keys().cloned().collect();
1012            if let Some((Some(failed_actor), _)) = &first_failure {
1013                uncollected_actors.remove(failed_actor);
1014            }
1015            while !uncollected_actors.is_empty()
1016                && let Some((actor_id, error)) = self.actor_failure_rx.recv().await
1017            {
1018                uncollected_actors.remove(&actor_id);
1019                later_errs.push(error);
1020            }
1021        })
1022        .await;
1023
1024        first_failure
1025            .into_iter()
1026            .map(|(_, err)| err)
1027            .chain(later_errs)
1028            .map(|e| e.with_score())
1029            .max_by_key(|e| e.score)
1030    }
1031
1032    /// Report that a source has finished listing for a specific epoch
1033    pub(super) fn report_source_list_finished(
1034        &mut self,
1035        epoch: EpochPair,
1036        actor_id: ActorId,
1037        table_id: TableId,
1038        associated_source_id: SourceId,
1039    ) {
1040        // Find the correct partial graph state by matching the actor's partial graph id
1041        if let Some(actor_state) = self.actor_states.get(&actor_id)
1042            && actor_state.inflight_barriers.contains(&epoch.prev)
1043        {
1044            self.graph_state
1045                .list_finished_source_ids
1046                .entry(epoch.curr)
1047                .or_default()
1048                .push(PbListFinishedSource {
1049                    reporter_actor_id: actor_id,
1050                    table_id,
1051                    associated_source_id,
1052                });
1053        } else {
1054            warn!(
1055                ?epoch,
1056                %actor_id, %table_id, %associated_source_id, "ignore source list finished"
1057            );
1058        }
1059    }
1060
1061    /// Report that a source has finished loading for a specific epoch
1062    pub(super) fn report_source_load_finished(
1063        &mut self,
1064        epoch: EpochPair,
1065        actor_id: ActorId,
1066        table_id: TableId,
1067        associated_source_id: SourceId,
1068    ) {
1069        // Find the correct partial graph state by matching the actor's partial graph id
1070        if let Some(actor_state) = self.actor_states.get(&actor_id)
1071            && actor_state.inflight_barriers.contains(&epoch.prev)
1072        {
1073            self.graph_state
1074                .load_finished_source_ids
1075                .entry(epoch.curr)
1076                .or_default()
1077                .push(PbLoadFinishedSource {
1078                    reporter_actor_id: actor_id,
1079                    table_id,
1080                    associated_source_id,
1081                });
1082        } else {
1083            warn!(
1084                ?epoch,
1085                %actor_id, %table_id, %associated_source_id, "ignore source load finished"
1086            );
1087        }
1088    }
1089
1090    /// Report that a CDC source has updated its offset at least once
1091    pub(super) fn report_cdc_source_offset_updated(
1092        &mut self,
1093        epoch: EpochPair,
1094        actor_id: ActorId,
1095        source_id: SourceId,
1096    ) {
1097        if let Some(actor_state) = self.actor_states.get(&actor_id)
1098            && actor_state.inflight_barriers.contains(&epoch.prev)
1099        {
1100            self.graph_state
1101                .cdc_source_offset_updated
1102                .entry(epoch.curr)
1103                .or_default()
1104                .push(PbCdcSourceOffsetUpdated {
1105                    reporter_actor_id: actor_id,
1106                    source_id,
1107                });
1108        } else {
1109            warn!(
1110                ?epoch,
1111                %actor_id, %source_id, "ignore cdc source offset updated"
1112            );
1113        }
1114    }
1115
1116    /// Record a Iceberg pk-index sink metadata report.
1117    pub(super) fn report_iceberg_pk_index_sink_metadata(
1118        &mut self,
1119        epoch: EpochPair,
1120        sink_id: SinkId,
1121        actor_id: ActorId,
1122        role: PbIcebergPkIndexSinkRole,
1123        metadata: Option<SinkMetadata>,
1124    ) {
1125        if let Some(actor_state) = self.actor_states.get(&actor_id)
1126            && actor_state.inflight_barriers.contains(&epoch.prev)
1127        {
1128            self.graph_state
1129                .iceberg_pk_index_sink_metadata
1130                .entry(epoch.curr)
1131                .or_default()
1132                .push(PbIcebergPkIndexSinkMetadata {
1133                    sink_id,
1134                    reporter_actor_id: actor_id,
1135                    prev_epoch: epoch.prev,
1136                    role: role as i32,
1137                    metadata,
1138                });
1139        } else {
1140            tracing::warn!(
1141                ?epoch,
1142                %actor_id,
1143                %sink_id,
1144                ?role,
1145                "ignore iceberg v3 sink metadata report from non-inflight actor"
1146            );
1147        }
1148    }
1149
1150    /// Report that a table has finished refreshing for a specific epoch
1151    pub(super) fn report_refresh_finished(
1152        &mut self,
1153        epoch: EpochPair,
1154        actor_id: ActorId,
1155        table_id: TableId,
1156        staging_table_id: TableId,
1157    ) {
1158        // Find the correct partial graph state by matching the actor's partial graph id
1159        let Some(actor_state) = self.actor_states.get(&actor_id) else {
1160            warn!(
1161                ?epoch,
1162                %actor_id, %table_id, "ignore refresh finished table: actor_state not found"
1163            );
1164            return;
1165        };
1166        if !actor_state.inflight_barriers.contains(&epoch.prev) {
1167            warn!(
1168                ?epoch,
1169                %actor_id,
1170                %table_id,
1171                inflight_barriers = ?actor_state.inflight_barriers,
1172                "ignore refresh finished table: partial_graph_id not found in inflight_barriers"
1173            );
1174            return;
1175        };
1176        self.graph_state
1177            .refresh_finished_tables
1178            .entry(epoch.curr)
1179            .or_default()
1180            .insert(table_id);
1181        self.graph_state
1182            .truncate_tables
1183            .entry(epoch.curr)
1184            .or_default()
1185            .insert(staging_table_id);
1186    }
1187}
1188
1189impl PartialGraphManagedBarrierState {
1190    /// Check whether any `Issued` barrier has been collected by all actors and, if so,
1191    /// transition it to `AllCollected`.
1192    fn may_have_collected_all(&mut self) -> Option<Barrier> {
1193        for barrier_state in self.epoch_barrier_state_map.values_mut() {
1194            match &barrier_state.inner {
1195                ManagedBarrierStateInner::Issued(IssuedState {
1196                    remaining_actors, ..
1197                }) if remaining_actors.is_empty() => {}
1198                ManagedBarrierStateInner::AllCollected { .. } => {
1199                    continue;
1200                }
1201                ManagedBarrierStateInner::Issued(_) => {
1202                    break;
1203                }
1204            }
1205
1206            self.streaming_metrics.barrier_manager_progress.inc();
1207
1208            let create_mview_progress = self
1209                .create_mview_progress
1210                .remove(&barrier_state.barrier.epoch.curr)
1211                .unwrap_or_default()
1212                .into_iter()
1213                .map(|(actor, (fragment_id, state))| state.to_pb(fragment_id, actor))
1214                .collect();
1215
1216            let list_finished_source_ids = self
1217                .list_finished_source_ids
1218                .remove(&barrier_state.barrier.epoch.curr)
1219                .unwrap_or_default();
1220
1221            let load_finished_source_ids = self
1222                .load_finished_source_ids
1223                .remove(&barrier_state.barrier.epoch.curr)
1224                .unwrap_or_default();
1225
1226            let cdc_table_backfill_progress = self
1227                .cdc_table_backfill_progress
1228                .remove(&barrier_state.barrier.epoch.curr)
1229                .unwrap_or_default()
1230                .into_iter()
1231                .map(|(actor, state)| state.to_pb(actor, barrier_state.barrier.epoch.curr))
1232                .collect();
1233
1234            let cdc_source_offset_updated = self
1235                .cdc_source_offset_updated
1236                .remove(&barrier_state.barrier.epoch.curr)
1237                .unwrap_or_default();
1238
1239            let iceberg_pk_index_sink_metadata = self
1240                .iceberg_pk_index_sink_metadata
1241                .remove(&barrier_state.barrier.epoch.curr)
1242                .unwrap_or_default();
1243
1244            let truncate_tables = self
1245                .truncate_tables
1246                .remove(&barrier_state.barrier.epoch.curr)
1247                .unwrap_or_default()
1248                .into_iter()
1249                .collect();
1250
1251            let refresh_finished_tables = self
1252                .refresh_finished_tables
1253                .remove(&barrier_state.barrier.epoch.curr)
1254                .unwrap_or_default()
1255                .into_iter()
1256                .collect();
1257            let prev_state = replace(
1258                &mut barrier_state.inner,
1259                ManagedBarrierStateInner::AllCollected {
1260                    create_mview_progress,
1261                    list_finished_source_ids,
1262                    load_finished_source_ids,
1263                    truncate_tables,
1264                    refresh_finished_tables,
1265                    cdc_table_backfill_progress,
1266                    cdc_source_offset_updated,
1267                    iceberg_pk_index_sink_metadata,
1268                },
1269            );
1270
1271            must_match!(prev_state, ManagedBarrierStateInner::Issued(IssuedState {
1272                barrier_inflight_latency: timer,
1273                ..
1274            }) => {
1275                timer.observe_duration();
1276            });
1277
1278            return Some(barrier_state.barrier.clone());
1279        }
1280        None
1281    }
1282
1283    fn pop_barrier_to_complete(&mut self, prev_epoch: u64) -> BarrierToComplete {
1284        let (popped_prev_epoch, barrier_state) = self
1285            .epoch_barrier_state_map
1286            .pop_first()
1287            .expect("should exist");
1288
1289        assert_eq!(prev_epoch, popped_prev_epoch);
1290
1291        let (
1292            create_mview_progress,
1293            list_finished_source_ids,
1294            load_finished_source_ids,
1295            cdc_table_backfill_progress,
1296            cdc_source_offset_updated,
1297            iceberg_pk_index_sink_metadata,
1298            truncate_tables,
1299            refresh_finished_tables,
1300        ) = must_match!(barrier_state.inner, ManagedBarrierStateInner::AllCollected {
1301            create_mview_progress,
1302            list_finished_source_ids,
1303            load_finished_source_ids,
1304            truncate_tables,
1305            refresh_finished_tables,
1306            cdc_table_backfill_progress,
1307            cdc_source_offset_updated,
1308            iceberg_pk_index_sink_metadata,
1309        } => {
1310            (create_mview_progress, list_finished_source_ids, load_finished_source_ids, cdc_table_backfill_progress, cdc_source_offset_updated, iceberg_pk_index_sink_metadata, truncate_tables, refresh_finished_tables)
1311        });
1312        BarrierToComplete {
1313            barrier: barrier_state.barrier,
1314            table_ids: barrier_state.table_ids,
1315            create_mview_progress,
1316            list_finished_source_ids,
1317            load_finished_source_ids,
1318            truncate_tables,
1319            refresh_finished_tables,
1320            cdc_table_backfill_progress,
1321            cdc_source_offset_updated,
1322            iceberg_pk_index_sink_metadata,
1323        }
1324    }
1325}
1326
1327pub(crate) struct BarrierToComplete {
1328    pub barrier: Barrier,
1329    pub table_ids: Option<HashSet<TableId>>,
1330    pub create_mview_progress: Vec<PbCreateMviewProgress>,
1331    pub list_finished_source_ids: Vec<PbListFinishedSource>,
1332    pub load_finished_source_ids: Vec<PbLoadFinishedSource>,
1333    pub truncate_tables: Vec<TableId>,
1334    pub refresh_finished_tables: Vec<TableId>,
1335    pub cdc_table_backfill_progress: Vec<PbCdcTableBackfillProgress>,
1336    pub cdc_source_offset_updated: Vec<PbCdcSourceOffsetUpdated>,
1337    pub iceberg_pk_index_sink_metadata: Vec<PbIcebergPkIndexSinkMetadata>,
1338}
1339
1340impl PartialGraphManagedBarrierState {
1341    /// Collect a `barrier` from the actor with `actor_id`.
1342    pub(super) fn collect(&mut self, actor_id: impl Into<ActorId>, epoch: EpochPair) {
1343        let actor_id = actor_id.into();
1344        tracing::debug!(
1345            target: "events::stream::barrier::manager::collect",
1346            ?epoch, %actor_id, state = ?self.epoch_barrier_state_map,
1347            "collect_barrier",
1348        );
1349
1350        match self.epoch_barrier_state_map.get_mut(&epoch.prev) {
1351            None => {
1352                // If the barrier's state is stashed, this occurs exclusively in scenarios where the barrier has not been
1353                // injected by the barrier manager, or the barrier message is blocked at the `RemoteInput` side waiting for injection.
1354                // Given these conditions, it's inconceivable for an actor to attempt collect at this point.
1355                panic!(
1356                    "cannot collect new actor barrier {:?} at current state: None",
1357                    epoch,
1358                )
1359            }
1360            Some(&mut BarrierState {
1361                ref barrier,
1362                inner:
1363                    ManagedBarrierStateInner::Issued(IssuedState {
1364                        ref mut remaining_actors,
1365                        ..
1366                    }),
1367                ..
1368            }) => {
1369                let exist = remaining_actors.remove(&actor_id);
1370                assert!(
1371                    exist,
1372                    "the actor doesn't exist. actor_id: {:?}, curr_epoch: {:?}",
1373                    actor_id, epoch.curr
1374                );
1375                assert_eq!(barrier.epoch.curr, epoch.curr);
1376            }
1377            Some(BarrierState { inner, .. }) => {
1378                panic!(
1379                    "cannot collect new actor barrier {:?} at current state: {:?}",
1380                    epoch, inner
1381                )
1382            }
1383        }
1384    }
1385
1386    /// When the meta service issues a `send_barrier` request, call this function to transform to
1387    /// `Issued` and start to collect or to notify.
1388    pub(super) fn transform_to_issued(
1389        &mut self,
1390        barrier: &Barrier,
1391        actor_ids_to_collect: impl IntoIterator<Item = ActorId>,
1392        table_ids: HashSet<TableId>,
1393    ) {
1394        let timer = self
1395            .streaming_metrics
1396            .barrier_inflight_latency
1397            .start_timer();
1398
1399        if let Some(hummock) = self.state_store.as_hummock() {
1400            hummock.start_epoch(barrier.epoch.curr, table_ids.clone());
1401        }
1402
1403        let table_ids = match barrier.kind {
1404            BarrierKind::Unspecified => {
1405                unreachable!()
1406            }
1407            BarrierKind::Initial => {
1408                assert!(
1409                    self.prev_barrier_table_ids.is_none(),
1410                    "non empty table_ids at initial barrier: {:?}",
1411                    self.prev_barrier_table_ids
1412                );
1413                info!(epoch = ?barrier.epoch, "initialize at Initial barrier");
1414                self.prev_barrier_table_ids = Some((barrier.epoch, table_ids));
1415                None
1416            }
1417            BarrierKind::Barrier => {
1418                if let Some((prev_epoch, prev_table_ids)) = self.prev_barrier_table_ids.as_mut() {
1419                    assert_eq!(prev_epoch.curr, barrier.epoch.prev);
1420                    assert_eq!(prev_table_ids, &table_ids);
1421                    *prev_epoch = barrier.epoch;
1422                } else {
1423                    info!(epoch = ?barrier.epoch, "initialize at non-checkpoint barrier");
1424                    self.prev_barrier_table_ids = Some((barrier.epoch, table_ids));
1425                }
1426                None
1427            }
1428            BarrierKind::Checkpoint => Some(
1429                if let Some((prev_epoch, prev_table_ids)) = self
1430                    .prev_barrier_table_ids
1431                    .replace((barrier.epoch, table_ids))
1432                    && prev_epoch.curr == barrier.epoch.prev
1433                {
1434                    prev_table_ids
1435                } else {
1436                    debug!(epoch = ?barrier.epoch, "reinitialize at Checkpoint barrier");
1437                    HashSet::new()
1438                },
1439            ),
1440        };
1441
1442        if let Some(&mut BarrierState { ref inner, .. }) =
1443            self.epoch_barrier_state_map.get_mut(&barrier.epoch.prev)
1444        {
1445            {
1446                panic!(
1447                    "barrier epochs{:?} state has already been `Issued`. Current state: {:?}",
1448                    barrier.epoch, inner
1449                );
1450            }
1451        };
1452
1453        self.epoch_barrier_state_map.insert(
1454            barrier.epoch.prev,
1455            BarrierState {
1456                barrier: barrier.clone(),
1457                inner: ManagedBarrierStateInner::Issued(IssuedState {
1458                    remaining_actors: BTreeSet::from_iter(actor_ids_to_collect),
1459                    barrier_inflight_latency: timer,
1460                }),
1461                table_ids,
1462            },
1463        );
1464    }
1465
1466    #[cfg(test)]
1467    async fn pop_next_completed_epoch(&mut self) -> u64 {
1468        if let Some(barrier) = self.may_have_collected_all() {
1469            self.pop_barrier_to_complete(barrier.epoch.prev);
1470            return barrier.epoch.prev;
1471        }
1472        pending().await
1473    }
1474}
1475
1476#[cfg(test)]
1477mod tests {
1478    use std::collections::HashSet;
1479
1480    use risingwave_common::util::epoch::test_epoch;
1481
1482    use crate::executor::Barrier;
1483    use crate::task::barrier_worker::managed_state::PartialGraphManagedBarrierState;
1484
1485    #[tokio::test]
1486    async fn test_managed_state_add_actor() {
1487        let mut managed_barrier_state = PartialGraphManagedBarrierState::for_test();
1488        let barrier1 = Barrier::new_test_barrier(test_epoch(1));
1489        let barrier2 = Barrier::new_test_barrier(test_epoch(2));
1490        let barrier3 = Barrier::new_test_barrier(test_epoch(3));
1491        let actor_ids_to_collect1 = HashSet::from([1.into(), 2.into()]);
1492        let actor_ids_to_collect2 = HashSet::from([1.into(), 2.into()]);
1493        let actor_ids_to_collect3 = HashSet::from([1.into(), 2.into(), 3.into()]);
1494        managed_barrier_state.transform_to_issued(&barrier1, actor_ids_to_collect1, HashSet::new());
1495        managed_barrier_state.transform_to_issued(&barrier2, actor_ids_to_collect2, HashSet::new());
1496        managed_barrier_state.transform_to_issued(&barrier3, actor_ids_to_collect3, HashSet::new());
1497        managed_barrier_state.collect(1, barrier1.epoch);
1498        managed_barrier_state.collect(2, barrier1.epoch);
1499        assert_eq!(
1500            managed_barrier_state.pop_next_completed_epoch().await,
1501            test_epoch(0)
1502        );
1503        assert_eq!(
1504            managed_barrier_state
1505                .epoch_barrier_state_map
1506                .first_key_value()
1507                .unwrap()
1508                .0,
1509            &test_epoch(1)
1510        );
1511        managed_barrier_state.collect(1, barrier2.epoch);
1512        managed_barrier_state.collect(1, barrier3.epoch);
1513        managed_barrier_state.collect(2, barrier2.epoch);
1514        assert_eq!(
1515            managed_barrier_state.pop_next_completed_epoch().await,
1516            test_epoch(1)
1517        );
1518        assert_eq!(
1519            managed_barrier_state
1520                .epoch_barrier_state_map
1521                .first_key_value()
1522                .unwrap()
1523                .0,
1524            &test_epoch(2)
1525        );
1526        managed_barrier_state.collect(2, barrier3.epoch);
1527        managed_barrier_state.collect(3, barrier3.epoch);
1528        assert_eq!(
1529            managed_barrier_state.pop_next_completed_epoch().await,
1530            test_epoch(2)
1531        );
1532        assert!(managed_barrier_state.epoch_barrier_state_map.is_empty());
1533    }
1534
1535    #[tokio::test]
1536    async fn test_managed_state_stop_actor() {
1537        let mut managed_barrier_state = PartialGraphManagedBarrierState::for_test();
1538        let barrier1 = Barrier::new_test_barrier(test_epoch(1));
1539        let barrier2 = Barrier::new_test_barrier(test_epoch(2));
1540        let barrier3 = Barrier::new_test_barrier(test_epoch(3));
1541        let actor_ids_to_collect1 = HashSet::from([1.into(), 2.into(), 3.into(), 4.into()]);
1542        let actor_ids_to_collect2 = HashSet::from([1.into(), 2.into(), 3.into()]);
1543        let actor_ids_to_collect3 = HashSet::from([1.into(), 2.into()]);
1544        managed_barrier_state.transform_to_issued(&barrier1, actor_ids_to_collect1, HashSet::new());
1545        managed_barrier_state.transform_to_issued(&barrier2, actor_ids_to_collect2, HashSet::new());
1546        managed_barrier_state.transform_to_issued(&barrier3, actor_ids_to_collect3, HashSet::new());
1547
1548        managed_barrier_state.collect(1, barrier1.epoch);
1549        managed_barrier_state.collect(1, barrier2.epoch);
1550        managed_barrier_state.collect(1, barrier3.epoch);
1551        managed_barrier_state.collect(2, barrier1.epoch);
1552        managed_barrier_state.collect(2, barrier2.epoch);
1553        managed_barrier_state.collect(2, barrier3.epoch);
1554        assert_eq!(
1555            managed_barrier_state
1556                .epoch_barrier_state_map
1557                .first_key_value()
1558                .unwrap()
1559                .0,
1560            &0
1561        );
1562        managed_barrier_state.collect(3, barrier1.epoch);
1563        managed_barrier_state.collect(3, barrier2.epoch);
1564        assert_eq!(
1565            managed_barrier_state
1566                .epoch_barrier_state_map
1567                .first_key_value()
1568                .unwrap()
1569                .0,
1570            &0
1571        );
1572        managed_barrier_state.collect(4, barrier1.epoch);
1573        assert_eq!(
1574            managed_barrier_state.pop_next_completed_epoch().await,
1575            test_epoch(0)
1576        );
1577        assert_eq!(
1578            managed_barrier_state.pop_next_completed_epoch().await,
1579            test_epoch(1)
1580        );
1581        assert_eq!(
1582            managed_barrier_state.pop_next_completed_epoch().await,
1583            test_epoch(2)
1584        );
1585        assert!(managed_barrier_state.epoch_barrier_state_map.is_empty());
1586    }
1587}
risingwave_stream/task/barrier_worker/managed_state.rs

risingwave_stream/task/barrier_worker/
managed_state.rs
