risingwave_meta/barrier/
partial_graph.rs

1// Copyright 2026 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::Bound::Unbounded;
16use std::collections::hash_map::Entry;
17use std::collections::{HashMap, HashSet};
18use std::mem::take;
19use std::ops::{Bound, RangeBounds};
20use std::sync::Arc;
21use std::time::Instant;
22
23use educe::Educe;
24use itertools::Itertools;
25use risingwave_common::util::epoch::EpochPair;
26use risingwave_pb::common::WorkerNode;
27use risingwave_pb::id::{ActorId, PartialGraphId, TableId, WorkerId};
28use risingwave_pb::stream_plan::barrier_mutation::Mutation;
29use risingwave_pb::stream_service::BarrierCompleteResponse;
30use risingwave_pb::stream_service::streaming_control_stream_response::{
31    ResetPartialGraphResponse, Response,
32};
33use tracing::{debug, warn};
34use uuid::Uuid;
35
36use crate::barrier::BarrierKind;
37use crate::barrier::command::PostCollectCommand;
38use crate::barrier::context::GlobalBarrierWorkerContext;
39use crate::barrier::info::BarrierInfo;
40use crate::barrier::notifier::Notifier;
41use crate::barrier::rpc::{ControlStreamManager, WorkerNodeEvent};
42use crate::barrier::utils::{BarrierItemCollector, NodeToCollect, is_valid_after_worker_err};
43use crate::manager::MetaSrvEnv;
44use crate::model::StreamJobActorsToCreate;
45use crate::{MetaError, MetaResult};
46
47#[derive(Debug)]
48pub(super) struct PartialGraphBarrierInfo {
49    enqueue_time: Instant,
50    pub(super) post_collect_command: PostCollectCommand,
51    pub(super) barrier_info: BarrierInfo,
52    pub(super) notifiers: Vec<Notifier>,
53    pub(super) table_ids_to_commit: HashSet<TableId>,
54}
55
56impl PartialGraphBarrierInfo {
57    pub(super) fn new(
58        post_collect_command: PostCollectCommand,
59        barrier_info: BarrierInfo,
60        notifiers: Vec<Notifier>,
61        table_ids_to_commit: HashSet<TableId>,
62    ) -> Self {
63        Self {
64            enqueue_time: Instant::now(),
65            post_collect_command,
66            barrier_info,
67            notifiers,
68            table_ids_to_commit,
69        }
70    }
71
72    pub(super) fn elapsed_secs(&self) -> f64 {
73        self.enqueue_time.elapsed().as_secs_f64()
74    }
75}
76
77pub(super) trait PartialGraphStat: Send + Sync + 'static {
78    fn observe_barrier_latency(&self, epoch: EpochPair, barrier_latency_secs: f64);
79    fn observe_barrier_num(&self, inflight_barrier_num: usize, collected_barrier_num: usize);
80}
81
82#[derive(Educe)]
83#[educe(Debug)]
84struct PartialGraphRunningState {
85    barrier_item_collector:
86        BarrierItemCollector<WorkerId, BarrierCompleteResponse, PartialGraphBarrierInfo>,
87    completing_epoch: Option<u64>,
88    #[educe(Debug(ignore))]
89    stat: Box<dyn PartialGraphStat>,
90}
91
92impl PartialGraphRunningState {
93    fn new(stat: Box<dyn PartialGraphStat>) -> Self {
94        Self {
95            barrier_item_collector: BarrierItemCollector::new(true),
96            completing_epoch: None,
97            stat,
98        }
99    }
100
101    fn is_empty(&self) -> bool {
102        self.barrier_item_collector.is_empty() && self.completing_epoch.is_none()
103    }
104
105    fn pending_barrier_num(&self) -> usize {
106        self.barrier_item_collector.inflight_barrier_num()
107            + self.barrier_item_collector.collected_barrier_num()
108            + usize::from(self.completing_epoch.is_some())
109    }
110
111    fn enqueue(&mut self, node_to_collect: NodeToCollect, mut info: PartialGraphBarrierInfo) {
112        let epoch = info.barrier_info.epoch();
113        assert_ne!(info.barrier_info.kind, BarrierKind::Initial);
114        if info.post_collect_command.should_checkpoint() {
115            assert!(info.barrier_info.kind.is_checkpoint());
116        }
117        info.notifiers.iter_mut().for_each(|n| n.notify_started());
118        self.barrier_item_collector
119            .enqueue(epoch, node_to_collect, info);
120        self.stat.observe_barrier_num(
121            self.barrier_item_collector.inflight_barrier_num(),
122            self.barrier_item_collector.collected_barrier_num(),
123        );
124    }
125
126    fn collect(&mut self, resp: BarrierCompleteResponse) {
127        debug!(
128            epoch = resp.epoch,
129            worker_id = %resp.worker_id,
130            partial_graph_id = %resp.partial_graph_id,
131            "collect barrier from worker"
132        );
133        self.barrier_item_collector
134            .collect(resp.epoch, resp.worker_id, resp);
135    }
136
137    fn barrier_collected<'a>(
138        &mut self,
139        temp_ref: &'a CollectedBarrierTempRef,
140    ) -> Option<CollectedBarrier<'a>> {
141        if let Some((epoch, info)) = self.barrier_item_collector.barrier_collected() {
142            self.stat
143                .observe_barrier_latency(epoch, info.elapsed_secs());
144            self.stat.observe_barrier_num(
145                self.barrier_item_collector.inflight_barrier_num(),
146                self.barrier_item_collector.collected_barrier_num(),
147            );
148            Some(temp_ref.collected_barrier(epoch, self.pending_barrier_num()))
149        } else {
150            None
151        }
152    }
153}
154
155#[derive(Debug)]
156struct ResetPartialGraphCollector {
157    remaining_workers: HashSet<WorkerId>,
158    reset_resps: HashMap<WorkerId, ResetPartialGraphResponse>,
159}
160
161impl ResetPartialGraphCollector {
162    fn collect(&mut self, worker_id: WorkerId, resp: ResetPartialGraphResponse) -> bool {
163        assert!(self.remaining_workers.remove(&worker_id));
164        self.reset_resps
165            .try_insert(worker_id, resp)
166            .expect("non-duplicate");
167        self.remaining_workers.is_empty()
168    }
169}
170
171#[derive(Educe)]
172#[educe(Debug)]
173enum PartialGraphStatus {
174    Running(PartialGraphRunningState),
175    Resetting(ResetPartialGraphCollector),
176    Initializing {
177        epoch: EpochPair,
178        node_to_collect: NodeToCollect,
179        #[educe(Debug(ignore))]
180        stat: Option<Box<dyn PartialGraphStat>>,
181    },
182}
183
184impl PartialGraphStatus {
185    fn collect<'a>(
186        &mut self,
187        worker_id: WorkerId,
188        resp: BarrierCompleteResponse,
189        temp_ref: &'a CollectedBarrierTempRef,
190    ) -> Option<PartialGraphEvent<'a>> {
191        assert_eq!(worker_id, resp.worker_id);
192        match self {
193            PartialGraphStatus::Running(state) => {
194                state.collect(resp);
195                state
196                    .barrier_collected(temp_ref)
197                    .map(PartialGraphEvent::BarrierCollected)
198            }
199            PartialGraphStatus::Resetting(_) => None,
200            PartialGraphStatus::Initializing {
201                epoch,
202                node_to_collect,
203                stat,
204            } => {
205                assert_eq!(epoch.prev, resp.epoch);
206                assert!(node_to_collect.remove(&worker_id));
207                if node_to_collect.is_empty() {
208                    *self = PartialGraphStatus::Running(PartialGraphRunningState::new(
209                        stat.take().expect("should be taken for once"),
210                    ));
211                    Some(PartialGraphEvent::Initialized)
212                } else {
213                    None
214                }
215            }
216        }
217    }
218}
219
220/// Holding the reference to a static empty map as a temporary workaround for borrow checker limitation on `CollectedBarrier`.
221/// Mod local method can create a `CollectedBarrierTempRef` locally, and the created `CollectedBarrier` will temporarily reference
222/// to the `CollectedBarrierTempRef`. The method must fill in the correct reference before returning in a mod-pub method.
223///
224/// This struct *must* be private to ensure that the invalid lifetime won't be leaked outside.
225struct CollectedBarrierTempRef {
226    resps: &'static HashMap<WorkerId, BarrierCompleteResponse>,
227}
228
229impl CollectedBarrierTempRef {
230    fn new() -> Self {
231        static EMPTY: std::sync::LazyLock<HashMap<WorkerId, BarrierCompleteResponse>> =
232            std::sync::LazyLock::new(HashMap::new);
233        CollectedBarrierTempRef { resps: &EMPTY }
234    }
235
236    fn collected_barrier(
237        &self,
238        epoch: EpochPair,
239        pending_barrier_num: usize,
240    ) -> CollectedBarrier<'_> {
241        CollectedBarrier {
242            epoch,
243            resps: self.resps,
244            pending_barrier_num,
245        }
246    }
247
248    fn correct_lifetime<'a>(
249        &self,
250        event: PartialGraphManagerEvent<'_>,
251        manager: &'a PartialGraphManager,
252    ) -> PartialGraphManagerEvent<'a> {
253        match event {
254            PartialGraphManagerEvent::PartialGraph(partial_graph_id, event) => {
255                let event = match event {
256                    PartialGraphEvent::BarrierCollected(collected) => {
257                        let pending_barrier_num = collected.pending_barrier_num;
258                        let state = manager.running_graph(partial_graph_id);
259                        let (epoch, resps, _) = state
260                            .barrier_item_collector
261                            .last_collected()
262                            .expect("should exist");
263                        assert_eq!(epoch, collected.epoch);
264                        PartialGraphEvent::BarrierCollected(CollectedBarrier {
265                            epoch,
266                            resps,
267                            pending_barrier_num,
268                        })
269                    }
270                    PartialGraphEvent::Reset(resps) => PartialGraphEvent::Reset(resps),
271                    PartialGraphEvent::Initialized => PartialGraphEvent::Initialized,
272                    PartialGraphEvent::Error(worker_id) => PartialGraphEvent::Error(worker_id),
273                };
274                PartialGraphManagerEvent::PartialGraph(partial_graph_id, event)
275            }
276            PartialGraphManagerEvent::Worker(worker_id, event) => {
277                PartialGraphManagerEvent::Worker(worker_id, event)
278            }
279        }
280    }
281}
282
283#[derive(Debug)]
284pub(super) struct CollectedBarrier<'a> {
285    pub epoch: EpochPair,
286    pub resps: &'a HashMap<WorkerId, BarrierCompleteResponse>,
287    pub pending_barrier_num: usize,
288}
289
290pub(super) enum PartialGraphEvent<'a> {
291    BarrierCollected(CollectedBarrier<'a>),
292    Reset(HashMap<WorkerId, ResetPartialGraphResponse>),
293    Initialized,
294    Error(WorkerId),
295}
296
297pub(super) enum WorkerEvent {
298    WorkerError {
299        err: MetaError,
300        affected_partial_graphs: HashSet<PartialGraphId>,
301    },
302    WorkerConnected,
303}
304
305fn existing_graphs(
306    graphs: &HashMap<PartialGraphId, PartialGraphStatus>,
307) -> impl Iterator<Item = PartialGraphId> + '_ {
308    graphs
309        .iter()
310        .filter_map(|(partial_graph_id, status)| match status {
311            PartialGraphStatus::Running(_) | PartialGraphStatus::Initializing { .. } => {
312                Some(*partial_graph_id)
313            }
314            PartialGraphStatus::Resetting(_) => None,
315        })
316}
317
318pub(super) struct PartialGraphManager {
319    control_stream_manager: ControlStreamManager,
320    term_id: String,
321    graphs: HashMap<PartialGraphId, PartialGraphStatus>,
322}
323
324impl PartialGraphManager {
325    pub(super) fn uninitialized(env: MetaSrvEnv) -> Self {
326        Self {
327            control_stream_manager: ControlStreamManager::new(env),
328            term_id: "uninitialized".to_owned(),
329            graphs: HashMap::new(),
330        }
331    }
332
333    pub(super) async fn recover(
334        env: MetaSrvEnv,
335        nodes: &HashMap<WorkerId, WorkerNode>,
336        context: Arc<impl GlobalBarrierWorkerContext>,
337    ) -> Self {
338        let term_id = Uuid::new_v4().to_string();
339        let control_stream_manager =
340            ControlStreamManager::recover(env, nodes, &term_id, context).await;
341        Self {
342            control_stream_manager,
343            term_id,
344            graphs: Default::default(),
345        }
346    }
347
348    pub(super) fn control_stream_manager(&self) -> &ControlStreamManager {
349        &self.control_stream_manager
350    }
351
352    pub(super) async fn add_worker(
353        &mut self,
354        node: WorkerNode,
355        context: Arc<impl GlobalBarrierWorkerContext>,
356    ) {
357        self.control_stream_manager
358            .add_worker(node, existing_graphs(&self.graphs), &self.term_id, context)
359            .await
360    }
361
362    pub(super) fn remove_worker(&mut self, node: WorkerNode) {
363        self.control_stream_manager.remove_worker(node);
364    }
365
366    pub(super) fn clear_worker(&mut self) {
367        self.control_stream_manager.clear();
368    }
369
370    pub(crate) fn notify_all_err(&mut self, err: &MetaError) {
371        for (_, graph) in self.graphs.drain() {
372            if let PartialGraphStatus::Running(graph) = graph {
373                for info in graph.barrier_item_collector.into_infos() {
374                    for notifier in info.notifiers {
375                        notifier.notify_collection_failed(err.clone());
376                    }
377                }
378            }
379        }
380    }
381}
382
383#[must_use]
384pub(super) struct PartialGraphAdder<'a> {
385    partial_graph_id: PartialGraphId,
386    manager: &'a mut PartialGraphManager,
387    consumed: bool,
388}
389
390impl PartialGraphAdder<'_> {
391    pub(super) fn added(mut self) {
392        self.consumed = true;
393    }
394
395    pub(super) fn failed(mut self) {
396        self.manager.reset_partial_graphs([self.partial_graph_id]);
397        self.consumed = true;
398    }
399
400    pub(super) fn manager(&mut self) -> &mut PartialGraphManager {
401        self.manager
402    }
403}
404
405impl Drop for PartialGraphAdder<'_> {
406    fn drop(&mut self) {
407        debug_assert!(self.consumed, "unconsumed graph adder");
408        if !self.consumed {
409            warn!(partial_graph_id = %self.partial_graph_id, "unconsumed graph adder");
410        }
411    }
412}
413
414impl PartialGraphManager {
415    pub(super) fn add_partial_graph(
416        &mut self,
417        partial_graph_id: PartialGraphId,
418        stat: impl PartialGraphStat,
419    ) -> PartialGraphAdder<'_> {
420        self.graphs
421            .try_insert(
422                partial_graph_id,
423                PartialGraphStatus::Running(PartialGraphRunningState::new(Box::new(stat))),
424            )
425            .expect("non-duplicated");
426        self.control_stream_manager
427            .add_partial_graph(partial_graph_id);
428        PartialGraphAdder {
429            partial_graph_id,
430            manager: self,
431            consumed: false,
432        }
433    }
434
435    pub(super) fn remove_partial_graphs(&mut self, partial_graphs: Vec<PartialGraphId>) {
436        for partial_graph_id in &partial_graphs {
437            let graph = self.graphs.remove(partial_graph_id).expect("should exist");
438            let PartialGraphStatus::Running(state) = graph else {
439                panic!("graph to be explicitly removed should be running");
440            };
441            assert!(state.is_empty());
442        }
443        self.control_stream_manager
444            .remove_partial_graphs(partial_graphs);
445    }
446
447    pub(super) fn reset_partial_graphs(
448        &mut self,
449        partial_graph_ids: impl IntoIterator<Item = PartialGraphId>,
450    ) {
451        let partial_graph_ids = partial_graph_ids.into_iter().collect_vec();
452        let remaining_workers = self
453            .control_stream_manager
454            .reset_partial_graphs(partial_graph_ids.clone());
455        let new_collector = || ResetPartialGraphCollector {
456            remaining_workers: remaining_workers.clone(),
457            reset_resps: Default::default(),
458        };
459        for partial_graph_id in partial_graph_ids {
460            match self.graphs.entry(partial_graph_id) {
461                Entry::Vacant(entry) => {
462                    entry.insert(PartialGraphStatus::Resetting(new_collector()));
463                }
464                Entry::Occupied(mut entry) => {
465                    let graph = entry.get_mut();
466                    match graph {
467                        PartialGraphStatus::Resetting(_) => {
468                            unreachable!("should not reset again")
469                        }
470                        PartialGraphStatus::Running(_)
471                        | PartialGraphStatus::Initializing { .. } => {
472                            *graph = PartialGraphStatus::Resetting(new_collector());
473                        }
474                    }
475                }
476            }
477        }
478    }
479
480    pub(super) fn assert_resetting(&self, partial_graph_id: PartialGraphId) {
481        let graph = self.graphs.get(&partial_graph_id).expect("should exist");
482        let PartialGraphStatus::Resetting(..) = graph else {
483            panic!("should be at resetting but at {:?}", graph);
484        };
485    }
486
487    pub(super) fn inject_barrier(
488        &mut self,
489        partial_graph_id: PartialGraphId,
490        mutation: Option<Mutation>,
491        node_actors: &HashMap<WorkerId, HashSet<ActorId>>,
492        table_ids_to_sync: impl Iterator<Item = TableId>,
493        nodes_to_sync_table: impl Iterator<Item = WorkerId>,
494        new_actors: Option<StreamJobActorsToCreate>,
495        info: PartialGraphBarrierInfo,
496    ) -> MetaResult<()> {
497        let graph = self
498            .graphs
499            .get_mut(&partial_graph_id)
500            .expect("should exist");
501        let node_to_collect = self.control_stream_manager.inject_barrier(
502            partial_graph_id,
503            mutation,
504            &info.barrier_info,
505            node_actors,
506            table_ids_to_sync,
507            nodes_to_sync_table,
508            new_actors,
509        )?;
510        let PartialGraphStatus::Running(state) = graph else {
511            panic!("should not inject barrier on non-running status: {graph:?}")
512        };
513        state.enqueue(node_to_collect, info);
514        Ok(())
515    }
516
517    fn running_graph(&self, partial_graph_id: PartialGraphId) -> &PartialGraphRunningState {
518        let PartialGraphStatus::Running(graph) = &self.graphs[&partial_graph_id] else {
519            unreachable!("should be running")
520        };
521        graph
522    }
523
524    fn running_graph_mut(
525        &mut self,
526        partial_graph_id: PartialGraphId,
527    ) -> &mut PartialGraphRunningState {
528        let PartialGraphStatus::Running(graph) = self
529            .graphs
530            .get_mut(&partial_graph_id)
531            .expect("should exist")
532        else {
533            unreachable!("should be running")
534        };
535        graph
536    }
537
538    pub(super) fn pending_barrier_num(&self, partial_graph_id: PartialGraphId) -> usize {
539        self.running_graph(partial_graph_id).pending_barrier_num()
540    }
541
542    pub(super) fn first_inflight_barrier(
543        &self,
544        partial_graph_id: PartialGraphId,
545    ) -> Option<EpochPair> {
546        self.running_graph(partial_graph_id)
547            .barrier_item_collector
548            .first_inflight_epoch()
549    }
550
551    pub(super) fn pending_barrier_infos(
552        &self,
553        partial_graph_id: PartialGraphId,
554    ) -> impl Iterator<Item = &BarrierInfo> {
555        self.running_graph(partial_graph_id)
556            .barrier_item_collector
557            .iter_infos()
558            .map(|info| &info.barrier_info)
559    }
560
561    pub(super) fn start_completing(
562        &mut self,
563        partial_graph_id: PartialGraphId,
564        epoch_end_bound: Bound<u64>,
565        mut on_non_checkpoint_epoch: impl FnMut(
566            EpochPair,
567            HashMap<WorkerId, BarrierCompleteResponse>,
568            PostCollectCommand,
569        ),
570    ) -> Option<(
571        u64,
572        HashMap<WorkerId, BarrierCompleteResponse>,
573        PartialGraphBarrierInfo,
574    )> {
575        let graph = self.running_graph_mut(partial_graph_id);
576        assert!(graph.completing_epoch.is_none());
577        let epoch_range: (Bound<u64>, Bound<u64>) = (Unbounded, epoch_end_bound);
578        while let Some((epoch, resps, info)) = graph
579            .barrier_item_collector
580            .take_collected_if(|epoch| epoch_range.contains(&epoch.prev))
581        {
582            if info.post_collect_command.should_checkpoint() {
583                assert!(info.barrier_info.kind.is_checkpoint());
584            } else if !info.barrier_info.kind.is_checkpoint() {
585                info.notifiers
586                    .into_iter()
587                    .for_each(Notifier::notify_collected);
588                on_non_checkpoint_epoch(epoch, resps, info.post_collect_command);
589                continue;
590            }
591            let prev_epoch = info.barrier_info.prev_epoch();
592            graph.completing_epoch = Some(prev_epoch);
593            return Some((prev_epoch, resps, info));
594        }
595        None
596    }
597
598    pub(super) fn ack_completed(&mut self, partial_graph_id: PartialGraphId, prev_epoch: u64) {
599        assert_eq!(
600            self.running_graph_mut(partial_graph_id)
601                .completing_epoch
602                .take(),
603            Some(prev_epoch)
604        );
605    }
606
607    pub(super) fn has_pending_checkpoint_barrier(&self, partial_graph_id: PartialGraphId) -> bool {
608        self.running_graph(partial_graph_id)
609            .barrier_item_collector
610            .iter_infos()
611            .any(|info| info.barrier_info.kind.is_checkpoint())
612    }
613
614    pub(super) fn start_recover(&mut self) -> PartialGraphRecoverer<'_> {
615        PartialGraphRecoverer {
616            added_partial_graphs: Default::default(),
617            manager: self,
618            consumed: false,
619        }
620    }
621}
622
623#[must_use]
624pub(super) struct PartialGraphRecoverer<'a> {
625    added_partial_graphs: HashSet<PartialGraphId>,
626    manager: &'a mut PartialGraphManager,
627    consumed: bool,
628}
629
630impl PartialGraphRecoverer<'_> {
631    pub(super) fn recover_graph(
632        &mut self,
633        partial_graph_id: PartialGraphId,
634        mutation: Mutation,
635        barrier_info: &BarrierInfo,
636        node_actors: &HashMap<WorkerId, HashSet<ActorId>>,
637        table_ids_to_sync: impl Iterator<Item = TableId>,
638        new_actors: StreamJobActorsToCreate,
639        stat: impl PartialGraphStat,
640    ) -> MetaResult<()> {
641        assert!(
642            self.added_partial_graphs.insert(partial_graph_id),
643            "duplicated recover graph {partial_graph_id}"
644        );
645        self.manager
646            .control_stream_manager
647            .add_partial_graph(partial_graph_id);
648        assert!(barrier_info.kind.is_initial());
649        let node_to_collect = self.manager.control_stream_manager.inject_barrier(
650            partial_graph_id,
651            Some(mutation),
652            barrier_info,
653            node_actors,
654            table_ids_to_sync,
655            node_actors.keys().copied(),
656            Some(new_actors),
657        )?;
658        self.manager
659            .graphs
660            .try_insert(
661                partial_graph_id,
662                PartialGraphStatus::Initializing {
663                    epoch: barrier_info.epoch(),
664                    node_to_collect,
665                    stat: Some(Box::new(stat)),
666                },
667            )
668            .expect("non-duplicated");
669        Ok(())
670    }
671
672    pub(super) fn control_stream_manager(&self) -> &ControlStreamManager {
673        &self.manager.control_stream_manager
674    }
675
676    pub(super) fn all_initializing(mut self) -> HashSet<PartialGraphId> {
677        self.consumed = true;
678        take(&mut self.added_partial_graphs)
679    }
680
681    pub(super) fn failed(mut self) -> HashSet<PartialGraphId> {
682        self.manager
683            .reset_partial_graphs(self.added_partial_graphs.iter().copied());
684        self.consumed = true;
685        take(&mut self.added_partial_graphs)
686    }
687}
688
689impl Drop for PartialGraphRecoverer<'_> {
690    fn drop(&mut self) {
691        debug_assert!(self.consumed, "unconsumed graph recoverer");
692        if !self.consumed {
693            warn!(partial_graph_ids = ?self.added_partial_graphs, "unconsumed graph recoverer");
694        }
695    }
696}
697
698#[must_use]
699pub(super) enum PartialGraphManagerEvent<'a> {
700    PartialGraph(PartialGraphId, PartialGraphEvent<'a>),
701    Worker(WorkerId, WorkerEvent),
702}
703
704impl PartialGraphManager {
705    pub(super) async fn next_event<'a>(
706        &'a mut self,
707        context: &Arc<impl GlobalBarrierWorkerContext>,
708    ) -> PartialGraphManagerEvent<'a> {
709        let temp_ref = CollectedBarrierTempRef::new();
710        let event = self.next_event_inner(context, &temp_ref).await;
711        temp_ref.correct_lifetime(event, self)
712    }
713
714    async fn next_event_inner<'a>(
715        &mut self,
716        context: &Arc<impl GlobalBarrierWorkerContext>,
717        temp_ref: &'a CollectedBarrierTempRef,
718    ) -> PartialGraphManagerEvent<'a> {
719        for (&partial_graph_id, graph) in &mut self.graphs {
720            match graph {
721                PartialGraphStatus::Running(state) => {
722                    if let Some(collected) = state.barrier_collected(temp_ref) {
723                        return PartialGraphManagerEvent::PartialGraph(
724                            partial_graph_id,
725                            PartialGraphEvent::BarrierCollected(collected),
726                        );
727                    }
728                }
729                PartialGraphStatus::Resetting(collector) => {
730                    if collector.remaining_workers.is_empty() {
731                        let resps = take(&mut collector.reset_resps);
732                        self.graphs.remove(&partial_graph_id);
733                        return PartialGraphManagerEvent::PartialGraph(
734                            partial_graph_id,
735                            PartialGraphEvent::Reset(resps),
736                        );
737                    }
738                }
739                PartialGraphStatus::Initializing {
740                    node_to_collect,
741                    stat,
742                    ..
743                } => {
744                    if node_to_collect.is_empty() {
745                        *graph = PartialGraphStatus::Running(PartialGraphRunningState::new(
746                            stat.take().expect("should be taken once"),
747                        ));
748                        return PartialGraphManagerEvent::PartialGraph(
749                            partial_graph_id,
750                            PartialGraphEvent::Initialized,
751                        );
752                    }
753                }
754            }
755        }
756        loop {
757            let (worker_id, event) = self
758                .control_stream_manager
759                .next_event(&self.term_id, context)
760                .await;
761            match event {
762                WorkerNodeEvent::Response(result) => match result {
763                    Ok(resp) => match resp {
764                        Response::CompleteBarrier(resp) => {
765                            let partial_graph_id = resp.partial_graph_id;
766                            if let Some(event) = self
767                                .graphs
768                                .get_mut(&partial_graph_id)
769                                .expect("should exist")
770                                .collect(worker_id, resp, temp_ref)
771                            {
772                                return PartialGraphManagerEvent::PartialGraph(
773                                    partial_graph_id,
774                                    event,
775                                );
776                            }
777                        }
778                        Response::ReportPartialGraphFailure(resp) => {
779                            let partial_graph_id = resp.partial_graph_id;
780                            let graph = self
781                                .graphs
782                                .get_mut(&partial_graph_id)
783                                .expect("should exist");
784                            match graph {
785                                PartialGraphStatus::Resetting(_) => {
786                                    // ignore reported error when resetting
787                                }
788                                PartialGraphStatus::Running(_)
789                                | PartialGraphStatus::Initializing { .. } => {
790                                    return PartialGraphManagerEvent::PartialGraph(
791                                        partial_graph_id,
792                                        PartialGraphEvent::Error(worker_id),
793                                    );
794                                }
795                            }
796                        }
797                        Response::ResetPartialGraph(resp) => {
798                            let partial_graph_id = resp.partial_graph_id;
799                            let graph = self
800                                .graphs
801                                .get_mut(&partial_graph_id)
802                                .expect("should exist");
803                            match graph {
804                                PartialGraphStatus::Running(_)
805                                | PartialGraphStatus::Initializing { .. } => {
806                                    if cfg!(debug_assertions) {
807                                        unreachable!(
808                                            "should not have reset request when not in resetting state"
809                                        )
810                                    } else {
811                                        warn!(
812                                            ?resp,
813                                            "ignore reset resp when not in Resetting state"
814                                        );
815                                    }
816                                }
817                                PartialGraphStatus::Resetting(collector) => {
818                                    if collector.collect(worker_id, resp) {
819                                        let resps = take(&mut collector.reset_resps);
820                                        self.graphs.remove(&partial_graph_id);
821                                        return PartialGraphManagerEvent::PartialGraph(
822                                            partial_graph_id,
823                                            PartialGraphEvent::Reset(resps),
824                                        );
825                                    }
826                                }
827                            }
828                        }
829                        Response::Init(_) | Response::Shutdown(_) => {
830                            unreachable!("should be handled in control stream manager")
831                        }
832                    },
833                    Err(error) => {
834                        let affected_partial_graphs = self
835                            .graphs
836                            .iter_mut()
837                            .filter_map(|(partial_graph_id, graph)| match graph {
838                                PartialGraphStatus::Running(state) => state
839                                    .barrier_item_collector
840                                    .iter_to_collect()
841                                    .any(|to_collect| {
842                                        !is_valid_after_worker_err(to_collect, worker_id)
843                                    })
844                                    .then_some(*partial_graph_id),
845                                PartialGraphStatus::Resetting(collector) => {
846                                    collector.remaining_workers.remove(&worker_id);
847                                    None
848                                }
849                                PartialGraphStatus::Initializing {
850                                    node_to_collect, ..
851                                } => (!is_valid_after_worker_err(node_to_collect, worker_id))
852                                    .then_some(*partial_graph_id),
853                            })
854                            .collect();
855                        return PartialGraphManagerEvent::Worker(
856                            worker_id,
857                            WorkerEvent::WorkerError {
858                                err: error,
859                                affected_partial_graphs,
860                            },
861                        );
862                    }
863                },
864                WorkerNodeEvent::Connected(connected) => {
865                    connected.initialize(existing_graphs(&self.graphs));
866                    return PartialGraphManagerEvent::Worker(
867                        worker_id,
868                        WorkerEvent::WorkerConnected,
869                    );
870                }
871            }
872        }
873    }
874}
875
876#[cfg(test)]
877mod tests {
878    use risingwave_common::util::epoch::Epoch;
879
880    use super::*;
881    use crate::barrier::TracedEpoch;
882
883    struct TestPartialGraphStat;
884
885    impl PartialGraphStat for TestPartialGraphStat {
886        fn observe_barrier_latency(&self, _epoch: EpochPair, _barrier_latency_secs: f64) {}
887
888        fn observe_barrier_num(&self, _inflight_barrier_num: usize, _collected_barrier_num: usize) {
889        }
890    }
891
892    fn barrier_info(prev_epoch: u64, curr_epoch: u64) -> PartialGraphBarrierInfo {
893        PartialGraphBarrierInfo::new(
894            PostCollectCommand::barrier(),
895            BarrierInfo {
896                prev_epoch: TracedEpoch::new(Epoch(prev_epoch)),
897                curr_epoch: TracedEpoch::new(Epoch(curr_epoch)),
898                kind: BarrierKind::Barrier,
899            },
900            vec![],
901            HashSet::new(),
902        )
903    }
904
905    #[test]
906    fn test_pending_barrier_num_includes_collected_and_completing() {
907        let mut state = PartialGraphRunningState::new(Box::new(TestPartialGraphStat));
908        let worker_id: WorkerId = 1.into();
909        state.barrier_item_collector.enqueue(
910            EpochPair::new(2, 1),
911            HashSet::from([worker_id]),
912            barrier_info(1, 2),
913        );
914        state.barrier_item_collector.enqueue(
915            EpochPair::new(3, 2),
916            HashSet::from([worker_id]),
917            barrier_info(2, 3),
918        );
919        assert_eq!(state.pending_barrier_num(), 2);
920
921        state
922            .barrier_item_collector
923            .collect(1, worker_id, BarrierCompleteResponse::default());
924        state.barrier_item_collector.barrier_collected();
925        assert_eq!(state.pending_barrier_num(), 2);
926
927        state
928            .barrier_item_collector
929            .take_collected_if(|epoch| epoch.prev == 1)
930            .expect("the first barrier should be collected");
931        state.completing_epoch = Some(1);
932        assert_eq!(state.pending_barrier_num(), 2);
933    }
934}
risingwave_meta/barrier/partial_graph.rs

risingwave_meta/barrier/
partial_graph.rs
