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::IcebergPkIndexCompactionContext;
29use risingwave_pb::stream_plan::barrier_mutation::Mutation;
30use risingwave_pb::stream_service::BarrierCompleteResponse;
31use risingwave_pb::stream_service::streaming_control_stream_response::{
32    ResetPartialGraphResponse, Response,
33};
34use tracing::{debug, warn};
35use uuid::Uuid;
36
37use crate::barrier::BarrierKind;
38use crate::barrier::command::PostCollectCommand;
39use crate::barrier::context::GlobalBarrierWorkerContext;
40use crate::barrier::info::BarrierInfo;
41use crate::barrier::notifier::{CollectionNotifier, NotifierStarter};
42use crate::barrier::rpc::{ControlStreamManager, WorkerNodeEvent};
43use crate::barrier::utils::{BarrierItemCollector, NodeToCollect, is_valid_after_worker_err};
44use crate::manager::MetaSrvEnv;
45use crate::model::StreamJobActorsToCreate;
46use crate::{MetaError, MetaResult};
47
48#[derive(Debug)]
49pub(super) struct PartialGraphBarrierInfo {
50    enqueue_time: Instant,
51    pub(super) post_collect_command: PostCollectCommand,
52    pub(super) barrier_info: BarrierInfo,
53    pub(super) notifier: Option<CollectionNotifier>,
54    pub(super) table_ids_to_commit: HashSet<TableId>,
55}
56
57impl PartialGraphBarrierInfo {
58    pub(super) fn new(
59        post_collect_command: PostCollectCommand,
60        barrier_info: BarrierInfo,
61        notifier: Option<&mut NotifierStarter>,
62        table_ids_to_commit: HashSet<TableId>,
63    ) -> Self {
64        Self {
65            enqueue_time: Instant::now(),
66            post_collect_command,
67            barrier_info,
68            notifier: notifier.map(NotifierStarter::add_notify),
69            table_ids_to_commit,
70        }
71    }
72
73    pub(super) fn elapsed_secs(&self) -> f64 {
74        self.enqueue_time.elapsed().as_secs_f64()
75    }
76}
77
78pub(super) trait PartialGraphStat: Send + Sync + 'static {
79    fn observe_barrier_latency(&self, epoch: EpochPair, barrier_latency_secs: f64);
80    fn observe_barrier_num(&self, inflight_barrier_num: usize, collected_barrier_num: usize);
81}
82
83#[derive(Educe)]
84#[educe(Debug)]
85struct PartialGraphRunningState {
86    barrier_item_collector:
87        BarrierItemCollector<WorkerId, BarrierCompleteResponse, PartialGraphBarrierInfo>,
88    completing_epoch: Option<u64>,
89    #[educe(Debug(ignore))]
90    stat: Box<dyn PartialGraphStat>,
91}
92
93impl PartialGraphRunningState {
94    fn new(stat: Box<dyn PartialGraphStat>) -> Self {
95        Self {
96            barrier_item_collector: BarrierItemCollector::new(true),
97            completing_epoch: None,
98            stat,
99        }
100    }
101
102    fn is_empty(&self) -> bool {
103        self.barrier_item_collector.is_empty() && self.completing_epoch.is_none()
104    }
105
106    fn pending_barrier_num(&self) -> usize {
107        self.barrier_item_collector.inflight_barrier_num()
108            + self.barrier_item_collector.collected_barrier_num()
109            + usize::from(self.completing_epoch.is_some())
110    }
111
112    fn enqueue(&mut self, node_to_collect: NodeToCollect, info: PartialGraphBarrierInfo) {
113        let epoch = info.barrier_info.epoch();
114        assert_ne!(info.barrier_info.kind, BarrierKind::Initial);
115        if info.post_collect_command.should_checkpoint() {
116            assert!(info.barrier_info.kind.is_checkpoint());
117        }
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                    if let Some(notifier) = info.notifier {
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        compaction_context: Option<IcebergPkIndexCompactionContext>,
492        node_actors: &HashMap<WorkerId, HashSet<ActorId>>,
493        table_ids_to_sync: impl Iterator<Item = TableId>,
494        nodes_to_sync_table: impl Iterator<Item = WorkerId>,
495        new_actors: Option<StreamJobActorsToCreate>,
496        info: PartialGraphBarrierInfo,
497    ) -> MetaResult<()> {
498        let graph = self
499            .graphs
500            .get_mut(&partial_graph_id)
501            .expect("should exist");
502        let node_to_collect = self.control_stream_manager.inject_barrier(
503            partial_graph_id,
504            mutation,
505            compaction_context,
506            &info.barrier_info,
507            node_actors,
508            table_ids_to_sync,
509            nodes_to_sync_table,
510            new_actors,
511        )?;
512        let PartialGraphStatus::Running(state) = graph else {
513            panic!("should not inject barrier on non-running status: {graph:?}")
514        };
515        state.enqueue(node_to_collect, info);
516        Ok(())
517    }
518
519    fn running_graph(&self, partial_graph_id: PartialGraphId) -> &PartialGraphRunningState {
520        let PartialGraphStatus::Running(graph) = &self.graphs[&partial_graph_id] else {
521            unreachable!("should be running")
522        };
523        graph
524    }
525
526    fn running_graph_mut(
527        &mut self,
528        partial_graph_id: PartialGraphId,
529    ) -> &mut PartialGraphRunningState {
530        let PartialGraphStatus::Running(graph) = self
531            .graphs
532            .get_mut(&partial_graph_id)
533            .expect("should exist")
534        else {
535            unreachable!("should be running")
536        };
537        graph
538    }
539
540    pub(super) fn pending_barrier_num(&self, partial_graph_id: PartialGraphId) -> usize {
541        self.running_graph(partial_graph_id).pending_barrier_num()
542    }
543
544    pub(super) fn first_inflight_barrier(
545        &self,
546        partial_graph_id: PartialGraphId,
547    ) -> Option<EpochPair> {
548        self.running_graph(partial_graph_id)
549            .barrier_item_collector
550            .first_inflight_epoch()
551    }
552
553    pub(super) fn pending_barrier_infos(
554        &self,
555        partial_graph_id: PartialGraphId,
556    ) -> impl Iterator<Item = &BarrierInfo> {
557        self.running_graph(partial_graph_id)
558            .barrier_item_collector
559            .iter_infos()
560            .map(|info| &info.barrier_info)
561    }
562
563    pub(super) fn start_completing(
564        &mut self,
565        partial_graph_id: PartialGraphId,
566        epoch_end_bound: Bound<u64>,
567        mut on_non_checkpoint_epoch: impl FnMut(
568            EpochPair,
569            HashMap<WorkerId, BarrierCompleteResponse>,
570            PostCollectCommand,
571        ),
572    ) -> Option<(
573        u64,
574        HashMap<WorkerId, BarrierCompleteResponse>,
575        PartialGraphBarrierInfo,
576    )> {
577        let graph = self.running_graph_mut(partial_graph_id);
578        assert!(graph.completing_epoch.is_none());
579        let epoch_range: (Bound<u64>, Bound<u64>) = (Unbounded, epoch_end_bound);
580        while let Some((epoch, resps, info)) = graph
581            .barrier_item_collector
582            .take_collected_if(|epoch| epoch_range.contains(&epoch.prev))
583        {
584            if info.post_collect_command.should_checkpoint() {
585                assert!(info.barrier_info.kind.is_checkpoint());
586            } else if !info.barrier_info.kind.is_checkpoint() {
587                if let Some(notifier) = info.notifier {
588                    notifier.notify_collected();
589                }
590                on_non_checkpoint_epoch(epoch, resps, info.post_collect_command);
591                continue;
592            }
593            let prev_epoch = info.barrier_info.prev_epoch();
594            graph.completing_epoch = Some(prev_epoch);
595            return Some((prev_epoch, resps, info));
596        }
597        None
598    }
599
600    pub(super) fn ack_completed(&mut self, partial_graph_id: PartialGraphId, prev_epoch: u64) {
601        assert_eq!(
602            self.running_graph_mut(partial_graph_id)
603                .completing_epoch
604                .take(),
605            Some(prev_epoch)
606        );
607    }
608
609    pub(super) fn has_pending_checkpoint_barrier(&self, partial_graph_id: PartialGraphId) -> bool {
610        self.running_graph(partial_graph_id)
611            .barrier_item_collector
612            .iter_infos()
613            .any(|info| info.barrier_info.kind.is_checkpoint())
614    }
615
616    pub(super) fn start_recover(&mut self) -> PartialGraphRecoverer<'_> {
617        PartialGraphRecoverer {
618            added_partial_graphs: Default::default(),
619            manager: self,
620            consumed: false,
621        }
622    }
623}
624
625#[must_use]
626pub(super) struct PartialGraphRecoverer<'a> {
627    added_partial_graphs: HashSet<PartialGraphId>,
628    manager: &'a mut PartialGraphManager,
629    consumed: bool,
630}
631
632impl PartialGraphRecoverer<'_> {
633    pub(super) fn recover_graph(
634        &mut self,
635        partial_graph_id: PartialGraphId,
636        mutation: Mutation,
637        barrier_info: &BarrierInfo,
638        node_actors: &HashMap<WorkerId, HashSet<ActorId>>,
639        table_ids_to_sync: impl Iterator<Item = TableId>,
640        new_actors: StreamJobActorsToCreate,
641        stat: impl PartialGraphStat,
642    ) -> MetaResult<()> {
643        assert!(
644            self.added_partial_graphs.insert(partial_graph_id),
645            "duplicated recover graph {partial_graph_id}"
646        );
647        self.manager
648            .control_stream_manager
649            .add_partial_graph(partial_graph_id);
650        assert!(barrier_info.kind.is_initial());
651        let node_to_collect = self.manager.control_stream_manager.inject_barrier(
652            partial_graph_id,
653            Some(mutation),
654            None,
655            barrier_info,
656            node_actors,
657            table_ids_to_sync,
658            node_actors.keys().copied(),
659            Some(new_actors),
660        )?;
661        self.manager
662            .graphs
663            .try_insert(
664                partial_graph_id,
665                PartialGraphStatus::Initializing {
666                    epoch: barrier_info.epoch(),
667                    node_to_collect,
668                    stat: Some(Box::new(stat)),
669                },
670            )
671            .expect("non-duplicated");
672        Ok(())
673    }
674
675    pub(super) fn control_stream_manager(&self) -> &ControlStreamManager {
676        &self.manager.control_stream_manager
677    }
678
679    pub(super) fn all_initializing(mut self) -> HashSet<PartialGraphId> {
680        self.consumed = true;
681        take(&mut self.added_partial_graphs)
682    }
683
684    pub(super) fn failed(mut self) -> HashSet<PartialGraphId> {
685        self.manager
686            .reset_partial_graphs(self.added_partial_graphs.iter().copied());
687        self.consumed = true;
688        take(&mut self.added_partial_graphs)
689    }
690}
691
692impl Drop for PartialGraphRecoverer<'_> {
693    fn drop(&mut self) {
694        debug_assert!(self.consumed, "unconsumed graph recoverer");
695        if !self.consumed {
696            warn!(partial_graph_ids = ?self.added_partial_graphs, "unconsumed graph recoverer");
697        }
698    }
699}
700
701#[must_use]
702pub(super) enum PartialGraphManagerEvent<'a> {
703    PartialGraph(PartialGraphId, PartialGraphEvent<'a>),
704    Worker(WorkerId, WorkerEvent),
705}
706
707impl PartialGraphManager {
708    pub(super) async fn next_event<'a>(
709        &'a mut self,
710        context: &Arc<impl GlobalBarrierWorkerContext>,
711    ) -> PartialGraphManagerEvent<'a> {
712        let temp_ref = CollectedBarrierTempRef::new();
713        let event = self.next_event_inner(context, &temp_ref).await;
714        temp_ref.correct_lifetime(event, self)
715    }
716
717    async fn next_event_inner<'a>(
718        &mut self,
719        context: &Arc<impl GlobalBarrierWorkerContext>,
720        temp_ref: &'a CollectedBarrierTempRef,
721    ) -> PartialGraphManagerEvent<'a> {
722        for (&partial_graph_id, graph) in &mut self.graphs {
723            match graph {
724                PartialGraphStatus::Running(state) => {
725                    if let Some(collected) = state.barrier_collected(temp_ref) {
726                        return PartialGraphManagerEvent::PartialGraph(
727                            partial_graph_id,
728                            PartialGraphEvent::BarrierCollected(collected),
729                        );
730                    }
731                }
732                PartialGraphStatus::Resetting(collector) => {
733                    if collector.remaining_workers.is_empty() {
734                        let resps = take(&mut collector.reset_resps);
735                        self.graphs.remove(&partial_graph_id);
736                        return PartialGraphManagerEvent::PartialGraph(
737                            partial_graph_id,
738                            PartialGraphEvent::Reset(resps),
739                        );
740                    }
741                }
742                PartialGraphStatus::Initializing {
743                    node_to_collect,
744                    stat,
745                    ..
746                } => {
747                    if node_to_collect.is_empty() {
748                        *graph = PartialGraphStatus::Running(PartialGraphRunningState::new(
749                            stat.take().expect("should be taken once"),
750                        ));
751                        return PartialGraphManagerEvent::PartialGraph(
752                            partial_graph_id,
753                            PartialGraphEvent::Initialized,
754                        );
755                    }
756                }
757            }
758        }
759        loop {
760            let (worker_id, event) = self
761                .control_stream_manager
762                .next_event(&self.term_id, context)
763                .await;
764            match event {
765                WorkerNodeEvent::Response(result) => match result {
766                    Ok(resp) => match resp {
767                        Response::CompleteBarrier(resp) => {
768                            let partial_graph_id = resp.partial_graph_id;
769                            if let Some(event) = self
770                                .graphs
771                                .get_mut(&partial_graph_id)
772                                .expect("should exist")
773                                .collect(worker_id, resp, temp_ref)
774                            {
775                                return PartialGraphManagerEvent::PartialGraph(
776                                    partial_graph_id,
777                                    event,
778                                );
779                            }
780                        }
781                        Response::ReportPartialGraphFailure(resp) => {
782                            let partial_graph_id = resp.partial_graph_id;
783                            let graph = self
784                                .graphs
785                                .get_mut(&partial_graph_id)
786                                .expect("should exist");
787                            match graph {
788                                PartialGraphStatus::Resetting(_) => {
789                                    // ignore reported error when resetting
790                                }
791                                PartialGraphStatus::Running(_)
792                                | PartialGraphStatus::Initializing { .. } => {
793                                    return PartialGraphManagerEvent::PartialGraph(
794                                        partial_graph_id,
795                                        PartialGraphEvent::Error(worker_id),
796                                    );
797                                }
798                            }
799                        }
800                        Response::ResetPartialGraph(resp) => {
801                            let partial_graph_id = resp.partial_graph_id;
802                            let graph = self
803                                .graphs
804                                .get_mut(&partial_graph_id)
805                                .expect("should exist");
806                            match graph {
807                                PartialGraphStatus::Running(_)
808                                | PartialGraphStatus::Initializing { .. } => {
809                                    if cfg!(debug_assertions) {
810                                        unreachable!(
811                                            "should not have reset request when not in resetting state"
812                                        )
813                                    } else {
814                                        warn!(
815                                            ?resp,
816                                            "ignore reset resp when not in Resetting state"
817                                        );
818                                    }
819                                }
820                                PartialGraphStatus::Resetting(collector) => {
821                                    if collector.collect(worker_id, resp) {
822                                        let resps = take(&mut collector.reset_resps);
823                                        self.graphs.remove(&partial_graph_id);
824                                        return PartialGraphManagerEvent::PartialGraph(
825                                            partial_graph_id,
826                                            PartialGraphEvent::Reset(resps),
827                                        );
828                                    }
829                                }
830                            }
831                        }
832                        Response::Init(_) | Response::Shutdown(_) => {
833                            unreachable!("should be handled in control stream manager")
834                        }
835                    },
836                    Err(error) => {
837                        let affected_partial_graphs = self
838                            .graphs
839                            .iter_mut()
840                            .filter_map(|(partial_graph_id, graph)| match graph {
841                                PartialGraphStatus::Running(state) => state
842                                    .barrier_item_collector
843                                    .iter_to_collect()
844                                    .any(|to_collect| {
845                                        !is_valid_after_worker_err(to_collect, worker_id)
846                                    })
847                                    .then_some(*partial_graph_id),
848                                PartialGraphStatus::Resetting(collector) => {
849                                    collector.remaining_workers.remove(&worker_id);
850                                    None
851                                }
852                                PartialGraphStatus::Initializing {
853                                    node_to_collect, ..
854                                } => (!is_valid_after_worker_err(node_to_collect, worker_id))
855                                    .then_some(*partial_graph_id),
856                            })
857                            .collect();
858                        return PartialGraphManagerEvent::Worker(
859                            worker_id,
860                            WorkerEvent::WorkerError {
861                                err: error,
862                                affected_partial_graphs,
863                            },
864                        );
865                    }
866                },
867                WorkerNodeEvent::Connected(connected) => {
868                    connected.initialize(existing_graphs(&self.graphs));
869                    return PartialGraphManagerEvent::Worker(
870                        worker_id,
871                        WorkerEvent::WorkerConnected,
872                    );
873                }
874            }
875        }
876    }
877}
878
879#[cfg(test)]
880mod tests {
881    use risingwave_common::util::epoch::Epoch;
882
883    use super::*;
884    use crate::barrier::TracedEpoch;
885
886    struct TestPartialGraphStat;
887
888    impl PartialGraphStat for TestPartialGraphStat {
889        fn observe_barrier_latency(&self, _epoch: EpochPair, _barrier_latency_secs: f64) {}
890
891        fn observe_barrier_num(&self, _inflight_barrier_num: usize, _collected_barrier_num: usize) {
892        }
893    }
894
895    fn barrier_info(prev_epoch: u64, curr_epoch: u64) -> PartialGraphBarrierInfo {
896        PartialGraphBarrierInfo::new(
897            PostCollectCommand::barrier(),
898            BarrierInfo {
899                prev_epoch: TracedEpoch::new(Epoch(prev_epoch)),
900                curr_epoch: TracedEpoch::new(Epoch(curr_epoch)),
901                kind: BarrierKind::Barrier,
902            },
903            None,
904            HashSet::new(),
905        )
906    }
907
908    #[test]
909    fn test_pending_barrier_num_includes_collected_and_completing() {
910        let mut state = PartialGraphRunningState::new(Box::new(TestPartialGraphStat));
911        let worker_id: WorkerId = 1.into();
912        state.barrier_item_collector.enqueue(
913            EpochPair::new(2, 1),
914            HashSet::from([worker_id]),
915            barrier_info(1, 2),
916        );
917        state.barrier_item_collector.enqueue(
918            EpochPair::new(3, 2),
919            HashSet::from([worker_id]),
920            barrier_info(2, 3),
921        );
922        assert_eq!(state.pending_barrier_num(), 2);
923
924        state
925            .barrier_item_collector
926            .collect(1, worker_id, BarrierCompleteResponse::default());
927        state.barrier_item_collector.barrier_collected();
928        assert_eq!(state.pending_barrier_num(), 2);
929
930        state
931            .barrier_item_collector
932            .take_collected_if(|epoch| epoch.prev == 1)
933            .expect("the first barrier should be collected");
934        state.completing_epoch = Some(1);
935        assert_eq!(state.pending_barrier_num(), 2);
936    }
937}
risingwave_meta/barrier/partial_graph.rs

risingwave_meta/barrier/
partial_graph.rs
