Skip to main content

risingwave_meta/barrier/
rpc.rs

1// Copyright 2024 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::hash_map::Entry;
16use std::collections::{HashMap, HashSet};
17use std::error::Error;
18use std::fmt::{Debug, Formatter};
19use std::future::poll_fn;
20use std::sync::Arc;
21use std::task::{Context, Poll};
22use std::time::Duration;
23
24use anyhow::anyhow;
25use fail::fail_point;
26use futures::future::{BoxFuture, join_all};
27use futures::{FutureExt, StreamExt};
28use itertools::Itertools;
29use risingwave_common::bail;
30use risingwave_common::catalog::{DatabaseId, FragmentTypeFlag, TableId};
31use risingwave_common::id::JobId;
32use risingwave_common::util::epoch::Epoch;
33use risingwave_common::util::retry::exponential_backoff;
34use risingwave_common::util::stream_graph_visitor::visit_stream_node_cont;
35use risingwave_common::util::tracing::TracingContext;
36use risingwave_connector::source::SplitImpl;
37use risingwave_meta_model::WorkerId;
38use risingwave_pb::common::{HostAddress, WorkerNode};
39use risingwave_pb::hummock::HummockVersionStats;
40use risingwave_pb::id::PartialGraphId;
41use risingwave_pb::source::{PbCdcTableSnapshotSplits, PbCdcTableSnapshotSplitsWithGeneration};
42use risingwave_pb::stream_plan::barrier_mutation::Mutation;
43use risingwave_pb::stream_plan::stream_node::NodeBody;
44use risingwave_pb::stream_plan::{AddMutation, Barrier, BarrierMutation};
45use risingwave_pb::stream_service::inject_barrier_request::build_actor_info::UpstreamActors;
46use risingwave_pb::stream_service::inject_barrier_request::{
47    BuildActorInfo, FragmentBuildActorInfo,
48};
49use risingwave_pb::stream_service::streaming_control_stream_request::{
50    CreatePartialGraphRequest, PbCreatePartialGraphRequest, PbInitRequest,
51    RemovePartialGraphRequest, ResetPartialGraphsRequest,
52};
53use risingwave_pb::stream_service::{
54    InjectBarrierRequest, StreamingControlStreamRequest, streaming_control_stream_request,
55    streaming_control_stream_response,
56};
57use risingwave_rpc_client::StreamingControlHandle;
58use thiserror_ext::AsReport;
59use tokio::time::{Instant, sleep};
60use tracing::{debug, error, info, warn};
61use uuid::Uuid;
62
63use super::{BarrierKind, TracedEpoch};
64use crate::barrier::BackfillOrderState;
65use crate::barrier::backfill_order_control::get_nodes_with_backfill_dependencies;
66use crate::barrier::cdc_progress::CdcTableBackfillTracker;
67use crate::barrier::checkpoint::independent_job::{
68    IcebergV3JobCheckpointControl, IcebergV3RenderResult, IndependentJobControl, IndependentJobInfo,
69};
70use crate::barrier::checkpoint::{
71    BarrierWorkerState, BatchRefreshJobCheckpointControl, BatchRefreshRenderResult,
72    CreatingStreamingJobControl, DatabaseCheckpointControl, DatabaseCheckpointControlMetrics,
73    IndependentCheckpointJob, IndependentCheckpointJobControl, IndependentCheckpointJobStatus,
74};
75use crate::barrier::context::{GlobalBarrierWorkerContext, GlobalBarrierWorkerContextImpl};
76use crate::barrier::edge_builder::FragmentEdgeBuilder;
77use crate::barrier::info::{
78    BarrierInfo, CreateStreamingJobStatus, InflightDatabaseInfo, InflightStreamingJobInfo,
79    SubscriberType,
80};
81use crate::barrier::partial_graph::PartialGraphRecoverer;
82use crate::barrier::progress::CreateMviewProgressTracker;
83use crate::barrier::utils::NodeToCollect;
84use crate::controller::fragment::InflightFragmentInfo;
85use crate::controller::utils::StreamingJobExtraInfo;
86use crate::manager::MetaSrvEnv;
87use crate::model::{
88    ActorId, FragmentDownstreamRelation, FragmentId, StreamActor, StreamJobActorsToCreate,
89    SubscriptionId,
90};
91use crate::stream::cdc::{
92    CdcTableSnapshotSplits, is_parallelized_backfill_enabled_cdc_scan_fragment,
93};
94use crate::stream::{
95    ExtendedFragmentBackfillOrder, StreamFragmentGraph, UserDefinedFragmentBackfillOrder,
96    build_actor_connector_splits,
97};
98use crate::{MetaError, MetaResult};
99
100pub(crate) fn to_partial_graph_id(
101    database_id: DatabaseId,
102    creating_job_id: Option<JobId>,
103) -> PartialGraphId {
104    let raw_job_id = creating_job_id
105        .map(|job_id| {
106            assert_ne!(job_id, u32::MAX);
107            job_id.as_raw_id()
108        })
109        .unwrap_or(u32::MAX);
110    (((database_id.as_raw_id() as u64) << 32) | (raw_job_id as u64)).into()
111}
112
113pub(super) fn from_partial_graph_id(
114    partial_graph_id: PartialGraphId,
115) -> (DatabaseId, Option<JobId>) {
116    let id = partial_graph_id.as_raw_id();
117    let database_id = (id >> 32) as u32;
118    let raw_creating_job_id = (id & ((1 << 32) - 1)) as u32;
119    let creating_job_id = if raw_creating_job_id == u32::MAX {
120        None
121    } else {
122        Some(JobId::new(raw_creating_job_id))
123    };
124    (database_id.into(), creating_job_id)
125}
126
127pub(super) fn build_locality_fragment_state_table_mapping(
128    fragment_infos: &HashMap<FragmentId, InflightFragmentInfo>,
129) -> HashMap<FragmentId, Vec<TableId>> {
130    let mut mapping = HashMap::new();
131
132    for (fragment_id, fragment_info) in fragment_infos {
133        let mut state_table_ids = Vec::new();
134        visit_stream_node_cont(&fragment_info.nodes, |stream_node| {
135            if let Some(NodeBody::LocalityProvider(locality_provider)) =
136                stream_node.node_body.as_ref()
137            {
138                let state_table_id = locality_provider
139                    .state_table
140                    .as_ref()
141                    .expect("must have state table")
142                    .id;
143                state_table_ids.push(state_table_id);
144                false
145            } else {
146                true
147            }
148        });
149        if !state_table_ids.is_empty() {
150            mapping.insert(*fragment_id, state_table_ids);
151        }
152    }
153
154    mapping
155}
156
157pub(super) fn database_partial_graphs<'a>(
158    database_id: DatabaseId,
159    creating_jobs: impl Iterator<Item = JobId> + Sized + 'a,
160) -> impl Iterator<Item = PartialGraphId> + 'a {
161    creating_jobs
162        .map(Some)
163        .chain([None])
164        .map(move |creating_job_id| to_partial_graph_id(database_id, creating_job_id))
165}
166
167struct ControlStreamNode {
168    worker_id: WorkerId,
169    host: HostAddress,
170    handle: StreamingControlHandle,
171}
172
173enum WorkerNodeState {
174    Connected {
175        control_stream: ControlStreamNode,
176        removed: bool,
177    },
178    Reconnecting(BoxFuture<'static, StreamingControlHandle>),
179}
180
181pub(super) struct ControlStreamManager {
182    workers: HashMap<WorkerId, (WorkerNode, WorkerNodeState)>,
183    pub env: MetaSrvEnv,
184}
185
186impl ControlStreamManager {
187    pub(super) fn new(env: MetaSrvEnv) -> Self {
188        Self {
189            workers: Default::default(),
190            env,
191        }
192    }
193
194    pub(super) fn host_addr(&self, worker_id: WorkerId) -> HostAddress {
195        self.workers[&worker_id].0.host.clone().unwrap()
196    }
197
198    pub(super) async fn add_worker<'a>(
199        &mut self,
200        node: WorkerNode,
201        partial_graphs: impl Iterator<Item = (PartialGraphId, &'a str)>,
202        context: Arc<impl GlobalBarrierWorkerContext>,
203    ) {
204        let node_id = node.id;
205        if let Entry::Occupied(entry) = self.workers.entry(node_id) {
206            let (existing_node, worker_state) = entry.get();
207            assert_eq!(existing_node.host, node.host);
208            warn!(id = %node.id, host = ?node.host, "node already exists");
209            match worker_state {
210                WorkerNodeState::Connected { .. } => {
211                    warn!(id = %node.id, host = ?node.host, "new node already connected");
212                    return;
213                }
214                WorkerNodeState::Reconnecting(_) => {
215                    warn!(id = %node.id, host = ?node.host, "remove previous pending worker connect request and reconnect");
216                    entry.remove();
217                }
218            }
219        }
220        let node_host = node.host.clone().unwrap();
221        let mut backoff =
222            exponential_backoff(Duration::from_millis(100), 5, Duration::from_secs(3));
223        const MAX_RETRY: usize = 5;
224        for i in 1..=MAX_RETRY {
225            match context.new_control_stream(&node).await {
226                Ok(mut handle) => {
227                    WorkerNodeConnected {
228                        handle: &mut handle,
229                        node: &node,
230                    }
231                    .initialize(partial_graphs);
232                    info!(?node_host, "add control stream worker");
233                    assert!(
234                        self.workers
235                            .insert(
236                                node_id,
237                                (
238                                    node,
239                                    WorkerNodeState::Connected {
240                                        control_stream: ControlStreamNode {
241                                            worker_id: node_id as _,
242                                            host: node_host,
243                                            handle,
244                                        },
245                                        removed: false
246                                    }
247                                )
248                            )
249                            .is_none()
250                    );
251                    return;
252                }
253                Err(e) => {
254                    // It may happen that the dns information of newly registered worker node
255                    // has not been propagated to the meta node and cause error. Wait for a while and retry
256                    let delay = backoff.next().unwrap();
257                    error!(
258                        attempt = i,
259                        backoff_delay = ?delay,
260                        err = %e.as_report(),
261                        ?node_host,
262                        "failed to resolve the worker node address",
263                    );
264                    sleep(delay).await;
265                }
266            }
267        }
268        error!(?node_host, "failed to create the worker node after retries");
269        assert!(
270            self.workers
271                .insert(
272                    node_id,
273                    (
274                        node.clone(),
275                        WorkerNodeState::Reconnecting(ControlStreamManager::retry_connect(
276                            node, context,
277                        ))
278                    )
279                )
280                .is_none()
281        );
282    }
283
284    pub(super) fn remove_worker(&mut self, node: WorkerNode) {
285        if let Entry::Occupied(mut entry) = self.workers.entry(node.id) {
286            let (_, worker_state) = entry.get_mut();
287            match worker_state {
288                WorkerNodeState::Connected { removed, .. } => {
289                    info!(worker_id = %node.id, "mark connected worker as removed");
290                    *removed = true;
291                }
292                WorkerNodeState::Reconnecting(_) => {
293                    info!(worker_id = %node.id, "remove worker");
294                    entry.remove();
295                }
296            }
297        }
298    }
299
300    fn retry_connect(
301        node: WorkerNode,
302        context: Arc<impl GlobalBarrierWorkerContext>,
303    ) -> BoxFuture<'static, StreamingControlHandle> {
304        async move {
305            let mut attempt = 0;
306            let backoff = exponential_backoff(
307                Duration::from_millis(100),
308                5,
309                Duration::from_mins(1),
310            );
311            for delay in backoff {
312                attempt += 1;
313                sleep(delay).await;
314                match context.new_control_stream(&node).await {
315                    Ok(handle) => {
316                        return handle;
317                    }
318                    Err(e) => {
319                        warn!(e = %e.as_report(), ?node, attempt, "failed to create the control stream worker");
320                    }
321                }
322            }
323            unreachable!("end of retry backoff")
324        }.boxed()
325    }
326
327    pub(super) async fn recover(
328        env: MetaSrvEnv,
329        nodes: &HashMap<WorkerId, WorkerNode>,
330        context: Arc<impl GlobalBarrierWorkerContext>,
331    ) -> Self {
332        let reset_start_time = Instant::now();
333        let nodes = join_all(nodes.iter().map(|(worker_id, node)| async {
334            let result = context.new_control_stream(node).await;
335            (*worker_id, node.clone(), result)
336        }))
337        .await;
338        let mut unconnected_workers = HashSet::new();
339        let mut workers = HashMap::new();
340        for (worker_id, node, result) in nodes {
341            match result {
342                Ok(handle) => {
343                    let control_stream = ControlStreamNode {
344                        worker_id: node.id,
345                        host: node.host.clone().unwrap(),
346                        handle,
347                    };
348                    assert!(
349                        workers
350                            .insert(
351                                worker_id,
352                                (
353                                    node,
354                                    WorkerNodeState::Connected {
355                                        control_stream,
356                                        removed: false
357                                    }
358                                )
359                            )
360                            .is_none()
361                    );
362                }
363                Err(e) => {
364                    unconnected_workers.insert(worker_id);
365                    warn!(
366                        e = %e.as_report(),
367                        %worker_id,
368                        ?node,
369                        "failed to connect to node"
370                    );
371                    assert!(
372                        workers
373                            .insert(
374                                worker_id,
375                                (
376                                    node.clone(),
377                                    WorkerNodeState::Reconnecting(Self::retry_connect(
378                                        node,
379                                        context.clone()
380                                    ))
381                                )
382                            )
383                            .is_none()
384                    );
385                }
386            }
387        }
388
389        info!(elapsed=?reset_start_time.elapsed(), ?unconnected_workers, "control stream reset");
390
391        Self { workers, env }
392    }
393
394    /// Clear all nodes and response streams in the manager.
395    pub(super) fn clear(&mut self) {
396        *self = Self::new(self.env.clone());
397    }
398}
399
400pub(super) struct WorkerNodeConnected<'a> {
401    node: &'a WorkerNode,
402    handle: &'a mut StreamingControlHandle,
403}
404
405impl<'a> WorkerNodeConnected<'a> {
406    pub(super) fn initialize<'b>(
407        self,
408        partial_graphs: impl Iterator<Item = (PartialGraphId, &'b str)>,
409    ) {
410        for (partial_graph_id, term_id) in partial_graphs {
411            if let Err(e) = self.handle.send_request(StreamingControlStreamRequest {
412                request: Some(
413                    streaming_control_stream_request::Request::CreatePartialGraph(
414                        PbCreatePartialGraphRequest {
415                            partial_graph_id,
416                            term_id: term_id.to_owned(),
417                        },
418                    ),
419                ),
420            }) {
421                warn!(e = %e.as_report(), node = ?self.node, "failed to send initial partial graph request");
422            }
423        }
424    }
425}
426
427pub(super) enum WorkerNodeEvent<'a> {
428    Response(MetaResult<streaming_control_stream_response::Response>),
429    Connected(WorkerNodeConnected<'a>),
430}
431
432impl ControlStreamManager {
433    fn poll_next_event<'a>(
434        this_opt: &mut Option<&'a mut Self>,
435        cx: &mut Context<'_>,
436        context: &Arc<impl GlobalBarrierWorkerContext>,
437        poll_reconnect: bool,
438    ) -> Poll<(WorkerId, WorkerNodeEvent<'a>)> {
439        let this = this_opt.as_mut().expect("Future polled after completion");
440        if this.workers.is_empty() {
441            return Poll::Pending;
442        }
443        {
444            for (&worker_id, (node, worker_state)) in &mut this.workers {
445                let control_stream = match worker_state {
446                    WorkerNodeState::Connected { control_stream, .. } => control_stream,
447                    WorkerNodeState::Reconnecting(_) if !poll_reconnect => {
448                        continue;
449                    }
450                    WorkerNodeState::Reconnecting(join_handle) => {
451                        match join_handle.poll_unpin(cx) {
452                            Poll::Ready(handle) => {
453                                info!(id=%node.id, host=?node.host, "reconnected to worker");
454                                *worker_state = WorkerNodeState::Connected {
455                                    control_stream: ControlStreamNode {
456                                        worker_id: node.id,
457                                        host: node.host.clone().unwrap(),
458                                        handle,
459                                    },
460                                    removed: false,
461                                };
462                                let this = this_opt.take().expect("should exist");
463                                let (node, worker_state) =
464                                    this.workers.get_mut(&worker_id).expect("should exist");
465                                let WorkerNodeState::Connected { control_stream, .. } =
466                                    worker_state
467                                else {
468                                    unreachable!()
469                                };
470                                return Poll::Ready((
471                                    worker_id,
472                                    WorkerNodeEvent::Connected(WorkerNodeConnected {
473                                        node,
474                                        handle: &mut control_stream.handle,
475                                    }),
476                                ));
477                            }
478                            Poll::Pending => {
479                                continue;
480                            }
481                        }
482                    }
483                };
484                match control_stream.handle.response_stream.poll_next_unpin(cx) {
485                    Poll::Ready(result) => {
486                        {
487                            let result = result
488                                .ok_or_else(|| (false, anyhow!("end of stream").into()))
489                                .and_then(|result| {
490                                    result.map_err(|err| -> (bool, MetaError) { (false, err.into()) }).and_then(|resp| {
491                                        match resp
492                                            .response
493                                            .ok_or_else(|| (false, anyhow!("empty response").into()))?
494                                        {
495                                            streaming_control_stream_response::Response::Shutdown(_) => Err((true, anyhow!(
496                                                "worker node {worker_id} is shutting down"
497                                            )
498                                                .into())),
499                                            streaming_control_stream_response::Response::Init(_) => {
500                                                // This arm should be unreachable.
501                                                Err((false, anyhow!("get unexpected init response").into()))
502                                            }
503                                            resp => {
504                                                if let streaming_control_stream_response::Response::CompleteBarrier(barrier_resp) = &resp {
505                                                    assert_eq!(worker_id, barrier_resp.worker_id);
506                                                }
507                                                Ok(resp)
508                                            }
509                                        }
510                                    })
511                                });
512                            let result = match result {
513                                Ok(resp) => Ok(resp),
514                                Err((shutdown, err)) => {
515                                    warn!(worker_id = %node.id, host = ?node.host, err = %err.as_report(), "get error from response stream");
516                                    let WorkerNodeState::Connected { removed, .. } = worker_state
517                                    else {
518                                        unreachable!("checked connected")
519                                    };
520                                    if *removed || shutdown {
521                                        this.workers.remove(&worker_id);
522                                    } else {
523                                        *worker_state = WorkerNodeState::Reconnecting(
524                                            ControlStreamManager::retry_connect(
525                                                node.clone(),
526                                                context.clone(),
527                                            ),
528                                        );
529                                    }
530                                    Err(err)
531                                }
532                            };
533                            return Poll::Ready((worker_id, WorkerNodeEvent::Response(result)));
534                        }
535                    }
536                    Poll::Pending => {
537                        continue;
538                    }
539                }
540            }
541        };
542
543        Poll::Pending
544    }
545
546    #[await_tree::instrument("control_stream_next_event")]
547    pub(super) async fn next_event<'a>(
548        &'a mut self,
549        context: &Arc<impl GlobalBarrierWorkerContext>,
550    ) -> (WorkerId, WorkerNodeEvent<'a>) {
551        let mut this = Some(self);
552        poll_fn(|cx| Self::poll_next_event(&mut this, cx, context, true)).await
553    }
554
555    #[await_tree::instrument("control_stream_next_response")]
556    pub(super) async fn next_response(
557        &mut self,
558        context: &Arc<impl GlobalBarrierWorkerContext>,
559    ) -> (
560        WorkerId,
561        MetaResult<streaming_control_stream_response::Response>,
562    ) {
563        let mut this = Some(self);
564        let (worker_id, event) =
565            poll_fn(|cx| Self::poll_next_event(&mut this, cx, context, false)).await;
566        match event {
567            WorkerNodeEvent::Response(result) => (worker_id, result),
568            WorkerNodeEvent::Connected(_) => {
569                unreachable!("set poll_reconnect=false")
570            }
571        }
572    }
573}
574
575pub(super) struct DatabaseInitialBarrierCollector {
576    pub(super) database_id: DatabaseId,
577    pub(super) initializing_partial_graphs: HashSet<PartialGraphId>,
578    pub(super) database: DatabaseCheckpointControl,
579}
580
581impl Debug for DatabaseInitialBarrierCollector {
582    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
583        f.debug_struct("DatabaseInitialBarrierCollector")
584            .field("database_id", &self.database_id)
585            .field("initializing_graphs", &self.initializing_partial_graphs)
586            .finish()
587    }
588}
589
590impl DatabaseInitialBarrierCollector {
591    pub(super) fn is_collected(&self) -> bool {
592        self.initializing_partial_graphs.is_empty()
593    }
594
595    pub(super) fn partial_graph_initialized(&mut self, partial_graph_id: PartialGraphId) {
596        assert!(self.initializing_partial_graphs.remove(&partial_graph_id));
597    }
598
599    pub(super) fn all_partial_graphs(&self) -> impl Iterator<Item = PartialGraphId> + '_ {
600        database_partial_graphs(
601            self.database_id,
602            self.database
603                .independent_checkpoint_job_controls
604                .keys()
605                .copied(),
606        )
607    }
608
609    pub(super) fn finish(self) -> DatabaseCheckpointControl {
610        assert!(self.is_collected());
611        self.database
612    }
613
614    pub(super) fn is_valid_after_worker_err(&self, worker_id: WorkerId) -> bool {
615        self.database.is_valid_after_worker_err(worker_id)
616    }
617}
618
619impl PartialGraphRecoverer<'_> {
620    /// Extract information from the loaded runtime barrier worker snapshot info, and inject the initial barrier.
621    #[expect(clippy::too_many_arguments)]
622    pub(super) fn inject_database_initial_barrier(
623        &mut self,
624        database_id: DatabaseId,
625        jobs: HashMap<JobId, HashMap<FragmentId, InflightFragmentInfo>>,
626        job_extra_info: &HashMap<JobId, StreamingJobExtraInfo>,
627        state_table_committed_epochs: &mut HashMap<TableId, u64>,
628        state_table_log_epochs: &mut HashMap<TableId, Vec<(Vec<u64>, u64)>>,
629        fragment_relations: &FragmentDownstreamRelation,
630        stream_actors: &HashMap<ActorId, StreamActor>,
631        source_splits: &mut HashMap<ActorId, Vec<SplitImpl>>,
632        creating_jobs: &mut HashSet<JobId>,
633        mv_depended_subscriptions: &mut HashMap<TableId, HashMap<SubscriptionId, u64>>,
634        is_paused: bool,
635        hummock_version_stats: &HummockVersionStats,
636        cdc_table_snapshot_splits: &mut HashMap<JobId, CdcTableSnapshotSplits>,
637        batch_refresh: HashMap<JobId, BatchRefreshRenderResult>,
638        iceberg_v3: HashMap<JobId, IcebergV3RenderResult>,
639    ) -> MetaResult<DatabaseCheckpointControl> {
640        let term_id = Uuid::new_v4().to_string();
641
642        fn collect_source_splits(
643            fragment_infos: impl Iterator<Item = &InflightFragmentInfo>,
644            source_splits: &mut HashMap<ActorId, Vec<SplitImpl>>,
645        ) -> HashMap<ActorId, Vec<SplitImpl>> {
646            fragment_infos
647                .flat_map(|info| info.actors.keys())
648                .filter_map(|actor_id| {
649                    let actor_id = *actor_id as ActorId;
650                    source_splits
651                        .remove(&actor_id)
652                        .map(|splits| (actor_id, splits))
653                })
654                .collect()
655        }
656        fn build_mutation(
657            splits: &HashMap<ActorId, Vec<SplitImpl>>,
658            cdc_table_snapshot_split_assignment: HashMap<ActorId, PbCdcTableSnapshotSplits>,
659            backfill_orders: &ExtendedFragmentBackfillOrder,
660            is_paused: bool,
661        ) -> Mutation {
662            let backfill_nodes_to_pause = get_nodes_with_backfill_dependencies(backfill_orders)
663                .into_iter()
664                .collect();
665            Mutation::Add(AddMutation {
666                // Actors built during recovery is not treated as newly added actors.
667                actor_dispatchers: Default::default(),
668                added_actors: Default::default(),
669                actor_splits: build_actor_connector_splits(splits),
670                actor_cdc_table_snapshot_splits: Some(PbCdcTableSnapshotSplitsWithGeneration {
671                    splits: cdc_table_snapshot_split_assignment,
672                }),
673                pause: is_paused,
674                subscriptions_to_add: Default::default(),
675                backfill_nodes_to_pause,
676                new_upstream_sinks: Default::default(),
677                dropped_actors: Default::default(),
678                sink_log_store_flush: Default::default(),
679            })
680        }
681
682        fn resolve_jobs_committed_epoch(
683            state_table_committed_epochs: &mut HashMap<TableId, u64>,
684            table_ids: impl Iterator<Item = TableId>,
685        ) -> u64 {
686            let mut epochs = table_ids.map(|table_id| {
687                (
688                    table_id,
689                    state_table_committed_epochs
690                        .remove(&table_id)
691                        .expect("should exist"),
692                )
693            });
694            let (first_table_id, prev_epoch) = epochs.next().expect("non-empty");
695            for (table_id, epoch) in epochs {
696                assert_eq!(
697                    prev_epoch, epoch,
698                    "{} has different committed epoch to {}",
699                    first_table_id, table_id
700                );
701            }
702            prev_epoch
703        }
704        fn job_backfill_orders(
705            job_extra_info: &HashMap<JobId, StreamingJobExtraInfo>,
706            job_id: JobId,
707        ) -> UserDefinedFragmentBackfillOrder {
708            UserDefinedFragmentBackfillOrder::new(
709                job_extra_info
710                    .get(&job_id)
711                    .and_then(|info| info.backfill_orders.clone())
712                    .map_or_else(HashMap::new, |orders| orders.0),
713            )
714        }
715
716        fn recover_job_backfill_order(
717            job_extra_info: &HashMap<JobId, StreamingJobExtraInfo>,
718            job_id: JobId,
719            downstreams: &FragmentDownstreamRelation,
720            fragment_infos: &HashMap<FragmentId, InflightFragmentInfo>,
721        ) -> ExtendedFragmentBackfillOrder {
722            let backfill_order = job_backfill_orders(job_extra_info, job_id);
723            StreamFragmentGraph::extend_fragment_backfill_ordering_with_locality_backfill(
724                backfill_order,
725                downstreams,
726                || {
727                    fragment_infos.iter().map(|(fragment_id, fragment)| {
728                        (*fragment_id, fragment.fragment_type_mask, &fragment.nodes)
729                    })
730                },
731            )
732        }
733
734        let mut subscribers: HashMap<_, HashMap<_, _>> = jobs
735            .keys()
736            .filter_map(|job_id| {
737                mv_depended_subscriptions
738                    .remove(&job_id.as_mv_table_id())
739                    .map(|subscriptions| {
740                        (
741                            job_id.as_mv_table_id(),
742                            subscriptions
743                                .into_iter()
744                                .map(|(subscription_id, retention)| {
745                                    (
746                                        subscription_id.as_subscriber_id(),
747                                        SubscriberType::Subscription(retention),
748                                    )
749                                })
750                                .collect(),
751                        )
752                    })
753            })
754            .collect();
755
756        // Batch-refresh jobs are rendered outside `jobs`, but their upstream tables
757        // must still start with log-store-enabled subscribers after recovery.
758        for (job_id, render_result) in &batch_refresh {
759            let snapshot_backfill_info = StreamFragmentGraph::collect_snapshot_backfill_info_impl(
760                render_result
761                    .fragment_infos
762                    .values()
763                    .map(|fragment| (&fragment.nodes, fragment.fragment_type_mask)),
764            )?
765            .0
766            .ok_or_else(|| anyhow!("batch refresh job {} has no snapshot backfill info", job_id))?;
767
768            for upstream_table_id in snapshot_backfill_info
769                .upstream_mv_table_id_to_backfill_epoch
770                .keys()
771            {
772                subscribers
773                    .entry(*upstream_table_id)
774                    .or_default()
775                    .try_insert(job_id.as_subscriber_id(), SubscriberType::SnapshotBackfill)
776                    .expect("non-duplicate");
777            }
778        }
779
780        for (job_id, render_result) in &iceberg_v3 {
781            let snapshot_backfill_info = StreamFragmentGraph::collect_snapshot_backfill_info_impl(
782                render_result
783                    .active_fragment_infos()
784                    .values()
785                    .map(|fragment| (&fragment.nodes, fragment.fragment_type_mask)),
786            )?
787            .0
788            .ok_or_else(|| anyhow!("Iceberg V3 job {} has no snapshot backfill info", job_id))?;
789
790            for upstream_table_id in snapshot_backfill_info
791                .upstream_mv_table_id_to_backfill_epoch
792                .keys()
793            {
794                subscribers
795                    .entry(*upstream_table_id)
796                    .or_default()
797                    .try_insert(job_id.as_subscriber_id(), SubscriberType::SnapshotBackfill)
798                    .expect("non-duplicate");
799            }
800        }
801
802        let mut database_jobs = HashMap::new();
803        let mut snapshot_backfill_jobs = HashMap::new();
804
805        for (job_id, job_fragments) in jobs {
806            if creating_jobs.remove(&job_id) {
807                if job_fragments.values().any(|fragment| {
808                    fragment
809                        .fragment_type_mask
810                        .contains(FragmentTypeFlag::SnapshotBackfillStreamScan)
811                }) {
812                    debug!(%job_id, "recovered snapshot backfill job");
813                    snapshot_backfill_jobs.insert(job_id, job_fragments);
814                } else {
815                    database_jobs.insert(job_id, (job_fragments, true));
816                }
817            } else {
818                database_jobs.insert(job_id, (job_fragments, false));
819            }
820        }
821
822        let database_job_log_epochs: HashMap<_, _> = database_jobs
823            .keys()
824            .filter_map(|job_id| {
825                state_table_log_epochs
826                    .remove(&job_id.as_mv_table_id())
827                    .map(|epochs| (job_id.as_mv_table_id(), epochs))
828            })
829            .collect();
830
831        let prev_epoch = resolve_jobs_committed_epoch(
832            state_table_committed_epochs,
833            InflightFragmentInfo::existing_table_ids(
834                database_jobs.values().flat_map(|(job, _)| job.values()),
835            ),
836        );
837        let prev_epoch = TracedEpoch::new(Epoch(prev_epoch));
838        // Use a different `curr_epoch` for each recovery attempt.
839        let curr_epoch = prev_epoch.next();
840        let barrier_info = BarrierInfo {
841            prev_epoch,
842            curr_epoch,
843            kind: BarrierKind::Initial,
844        };
845
846        let mut ongoing_snapshot_backfill_jobs: HashMap<JobId, _> = HashMap::new();
847        for (job_id, fragment_infos) in snapshot_backfill_jobs {
848            let committed_epoch = resolve_jobs_committed_epoch(
849                state_table_committed_epochs,
850                InflightFragmentInfo::existing_table_ids(fragment_infos.values()),
851            );
852            if committed_epoch == barrier_info.prev_epoch() {
853                info!(
854                    "recovered creating snapshot backfill job {} catch up with upstream already",
855                    job_id
856                );
857                database_jobs
858                    .try_insert(job_id, (fragment_infos, true))
859                    .expect("non-duplicate");
860                continue;
861            }
862            let snapshot_backfill_info = StreamFragmentGraph::collect_snapshot_backfill_info_impl(
863                fragment_infos
864                    .values()
865                    .map(|fragment| (&fragment.nodes, fragment.fragment_type_mask)),
866            )?
867            .0
868            .ok_or_else(|| {
869                anyhow!(
870                    "recovered snapshot backfill job {} has no snapshot backfill info",
871                    job_id
872                )
873            })?;
874            let mut snapshot_epoch = None;
875            let upstream_table_ids: HashSet<_> = snapshot_backfill_info
876                .upstream_mv_table_id_to_backfill_epoch
877                .keys()
878                .cloned()
879                .collect();
880            for (upstream_table_id, epoch) in
881                snapshot_backfill_info.upstream_mv_table_id_to_backfill_epoch
882            {
883                let epoch = epoch.ok_or_else(|| anyhow!("recovered snapshot backfill job {} to upstream {} has not set snapshot epoch", job_id, upstream_table_id))?;
884                let snapshot_epoch = snapshot_epoch.get_or_insert(epoch);
885                if *snapshot_epoch != epoch {
886                    return Err(anyhow!("snapshot epoch {} to upstream {} different to snapshot epoch {} to previous upstream", epoch, upstream_table_id, snapshot_epoch).into());
887                }
888            }
889            let snapshot_epoch = snapshot_epoch.ok_or_else(|| {
890                anyhow!(
891                    "snapshot backfill job {} has not set snapshot epoch",
892                    job_id
893                )
894            })?;
895            for upstream_table_id in &upstream_table_ids {
896                subscribers
897                    .entry(*upstream_table_id)
898                    .or_default()
899                    .try_insert(job_id.as_subscriber_id(), SubscriberType::SnapshotBackfill)
900                    .expect("non-duplicate");
901            }
902            ongoing_snapshot_backfill_jobs
903                .try_insert(
904                    job_id,
905                    (
906                        fragment_infos,
907                        upstream_table_ids,
908                        committed_epoch,
909                        snapshot_epoch,
910                    ),
911                )
912                .expect("non-duplicated");
913        }
914
915        let mut cdc_table_snapshot_split_assignment: HashMap<ActorId, PbCdcTableSnapshotSplits> =
916            HashMap::new();
917
918        let database_jobs: HashMap<JobId, InflightStreamingJobInfo> = {
919            database_jobs
920                .into_iter()
921                .map(|(job_id, (fragment_infos, is_creating))| {
922                    let status = if is_creating {
923                        let backfill_ordering = job_backfill_orders(job_extra_info, job_id);
924                        let backfill_ordering = StreamFragmentGraph::extend_fragment_backfill_ordering_with_locality_backfill(
925                            backfill_ordering,
926                            fragment_relations,
927                            || fragment_infos.iter().map(|(fragment_id, fragment)| {
928                            (*fragment_id, fragment.fragment_type_mask, &fragment.nodes)
929                        }));
930                        let backfill_order_state = BackfillOrderState::recover_from_fragment_infos(
931                            &backfill_ordering,
932                            &fragment_infos,
933                        );
934                        CreateStreamingJobStatus::Creating {
935                            tracker: CreateMviewProgressTracker::recover(
936                                job_id,
937                                &fragment_infos,
938                                backfill_order_state,
939                                hummock_version_stats,
940                            ),
941                        }
942                    } else {
943                        CreateStreamingJobStatus::Created
944                    };
945                    let cdc_table_backfill_tracker =
946                        if let Some(splits) = cdc_table_snapshot_splits.remove(&job_id) {
947                            let cdc_fragment = fragment_infos
948                                .values()
949                                .find(|fragment| {
950                                    is_parallelized_backfill_enabled_cdc_scan_fragment(
951                                        fragment.fragment_type_mask,
952                                        &fragment.nodes,
953                                    )
954                                    .is_some()
955                                })
956                                .expect("should have parallel cdc fragment");
957                            let cdc_actors = cdc_fragment.actors.keys().copied().collect();
958                            let mut tracker =
959                                CdcTableBackfillTracker::restore(cdc_fragment.fragment_id, splits);
960                            cdc_table_snapshot_split_assignment
961                                .extend(tracker.reassign_splits(cdc_actors)?);
962                            Some(tracker)
963                        } else {
964                            None
965                        };
966                    Ok((
967                        job_id,
968                        InflightStreamingJobInfo {
969                            job_id,
970                            fragment_infos,
971                            subscribers: subscribers
972                                .remove(&job_id.as_mv_table_id())
973                                .unwrap_or_default(),
974                            status,
975                            cdc_table_backfill_tracker,
976                        },
977                    ))
978                })
979                .try_collect::<_, _, MetaError>()
980        }?;
981
982        let (mut edges, _) = FragmentEdgeBuilder::new()
983            .add_new_fragments(
984                database_jobs.values().flat_map(|job| job.fragment_infos()),
985                to_partial_graph_id(database_id, None),
986                self.control_stream_manager(),
987            )
988            .finish_fragments()
989            .add_relations(fragment_relations)?
990            .build();
991
992        {
993            let new_actors =
994                edges.collect_actors_to_create(database_jobs.values().flat_map(move |job| {
995                    job.fragment_infos.values().map(move |fragment_infos| {
996                        (
997                            fragment_infos.fragment_id,
998                            &fragment_infos.nodes,
999                            fragment_infos.actors.iter().map(move |(actor_id, actor)| {
1000                                (
1001                                    stream_actors.get(actor_id).expect("should exist"),
1002                                    actor.worker_id,
1003                                )
1004                            }),
1005                            job.subscribers.keys().copied(),
1006                        )
1007                    })
1008                }));
1009
1010            let nodes_actors =
1011                InflightFragmentInfo::actor_ids_to_collect(database_jobs.values().flatten());
1012            let database_job_source_splits =
1013                collect_source_splits(database_jobs.values().flatten(), source_splits);
1014            let database_backfill_orders =
1015                UserDefinedFragmentBackfillOrder::merge(database_jobs.values().map(|job| {
1016                    if matches!(job.status, CreateStreamingJobStatus::Creating { .. }) {
1017                        job_backfill_orders(job_extra_info, job.job_id)
1018                    } else {
1019                        UserDefinedFragmentBackfillOrder::default()
1020                    }
1021                }));
1022            let database_backfill_orders =
1023                StreamFragmentGraph::extend_fragment_backfill_ordering_with_locality_backfill(
1024                    database_backfill_orders,
1025                    fragment_relations,
1026                    || {
1027                        database_jobs.values().flat_map(|job_fragments| {
1028                            job_fragments
1029                                .fragment_infos
1030                                .iter()
1031                                .map(|(fragment_id, fragment)| {
1032                                    (*fragment_id, fragment.fragment_type_mask, &fragment.nodes)
1033                                })
1034                        })
1035                    },
1036                );
1037            let mutation = build_mutation(
1038                &database_job_source_splits,
1039                cdc_table_snapshot_split_assignment,
1040                &database_backfill_orders,
1041                is_paused,
1042            );
1043
1044            let partial_graph_id = to_partial_graph_id(database_id, None);
1045            self.recover_graph(
1046                partial_graph_id,
1047                &term_id,
1048                mutation,
1049                &barrier_info,
1050                &nodes_actors,
1051                InflightFragmentInfo::existing_table_ids(database_jobs.values().flatten()),
1052                new_actors,
1053                DatabaseCheckpointControlMetrics::new(database_id),
1054            )?;
1055            debug!(
1056                %database_id,
1057                "inject initial barrier"
1058            );
1059        };
1060
1061        let mut independent_checkpoint_job_controls: HashMap<
1062            JobId,
1063            IndependentCheckpointJobControl,
1064        > = HashMap::new();
1065        for (job_id, (info, upstream_table_ids, committed_epoch, snapshot_epoch)) in
1066            ongoing_snapshot_backfill_jobs
1067        {
1068            let partial_graph_id = to_partial_graph_id(database_id, Some(job_id));
1069            let (mut job_edges, _) = FragmentEdgeBuilder::new()
1070                .add_new_fragments(
1071                    info.values(),
1072                    partial_graph_id,
1073                    self.control_stream_manager(),
1074                )
1075                .finish_fragments()
1076                .add_relations(fragment_relations)?
1077                .build();
1078            let new_actors =
1079                job_edges.collect_actors_to_create(info.values().map(|fragment_infos| {
1080                    (
1081                        fragment_infos.fragment_id,
1082                        &fragment_infos.nodes,
1083                        fragment_infos.actors.iter().map(move |(actor_id, actor)| {
1084                            (
1085                                stream_actors.get(actor_id).expect("should exist"),
1086                                actor.worker_id,
1087                            )
1088                        }),
1089                        vec![], // no subscribers for backfilling jobs
1090                    )
1091                }));
1092
1093            let job_source_splits = collect_source_splits(info.values(), source_splits);
1094            assert!(
1095                !cdc_table_snapshot_splits.contains_key(&job_id),
1096                "snapshot backfill job {job_id} should not have cdc backfill"
1097            );
1098            if is_paused {
1099                bail!("should not pause when having snapshot backfill job {job_id}");
1100            }
1101            let job_backfill_orders =
1102                recover_job_backfill_order(job_extra_info, job_id, fragment_relations, &info);
1103            let mutation = build_mutation(
1104                &job_source_splits,
1105                Default::default(), // no cdc backfill job for
1106                &job_backfill_orders,
1107                false,
1108            );
1109            let control = IndependentJobControl::recovered(
1110                IndependentJobInfo::from_fragment_infos(
1111                    database_id,
1112                    job_id,
1113                    snapshot_epoch,
1114                    upstream_table_ids,
1115                    info.values(),
1116                ),
1117                committed_epoch,
1118            );
1119
1120            let job = CreatingStreamingJobControl::recover(
1121                control,
1122                &database_job_log_epochs,
1123                &barrier_info,
1124                info,
1125                job_backfill_orders,
1126                fragment_relations,
1127                hummock_version_stats,
1128                new_actors,
1129                mutation.clone(),
1130                &term_id,
1131                self,
1132            )?;
1133            let job = IndependentCheckpointJobControl::creating_streaming_job(
1134                job_id,
1135                partial_graph_id,
1136                IndependentCheckpointJobStatus::Ready,
1137                job,
1138            );
1139            independent_checkpoint_job_controls.insert(job_id, job);
1140        }
1141
1142        // Recover Iceberg V3 jobs from their separately rendered closed graphs. The resolver
1143        // partition has already been removed from the active runtime graph.
1144        for (job_id, render_result) in iceberg_v3 {
1145            creating_jobs.remove(&job_id);
1146            debug!(%job_id, "recovered Iceberg V3 job");
1147
1148            let committed_epoch = resolve_jobs_committed_epoch(
1149                state_table_committed_epochs,
1150                InflightFragmentInfo::existing_table_ids(render_result.all_fragment_infos()),
1151            );
1152            let snapshot_backfill_info = StreamFragmentGraph::collect_snapshot_backfill_info_impl(
1153                render_result
1154                    .active_fragment_infos()
1155                    .values()
1156                    .map(|fragment| (&fragment.nodes, fragment.fragment_type_mask)),
1157            )?
1158            .0
1159            .ok_or_else(|| anyhow!("Iceberg V3 job {} has no snapshot backfill info", job_id))?;
1160            let upstream_table_ids: HashSet<_> = snapshot_backfill_info
1161                .upstream_mv_table_id_to_backfill_epoch
1162                .keys()
1163                .copied()
1164                .collect();
1165            let mut snapshot_epochs = snapshot_backfill_info
1166                .upstream_mv_table_id_to_backfill_epoch
1167                .values()
1168                .copied();
1169            let snapshot_epoch = snapshot_epochs
1170                .next()
1171                .flatten()
1172                .ok_or_else(|| anyhow!("Iceberg V3 job {} has no snapshot epoch", job_id))?;
1173            if snapshot_epochs.any(|epoch| epoch != Some(snapshot_epoch)) {
1174                return Err(
1175                    anyhow!("Iceberg V3 job {} has inconsistent snapshot epochs", job_id).into(),
1176                );
1177            }
1178
1179            let job_backfill_orders = recover_job_backfill_order(
1180                job_extra_info,
1181                job_id,
1182                render_result.active_downstreams(),
1183                render_result.active_fragment_infos(),
1184            );
1185            let mutation = build_mutation(
1186                &Default::default(),
1187                Default::default(),
1188                &job_backfill_orders,
1189                false,
1190            );
1191            let partial_graph_id = to_partial_graph_id(database_id, Some(job_id));
1192            let control = IndependentJobControl::recovered(
1193                IndependentJobInfo::from_fragment_infos(
1194                    database_id,
1195                    job_id,
1196                    snapshot_epoch,
1197                    upstream_table_ids,
1198                    render_result.all_fragment_infos(),
1199                ),
1200                committed_epoch,
1201            );
1202            let iceberg_job = IcebergV3JobCheckpointControl::recover(
1203                control,
1204                &database_job_log_epochs,
1205                &barrier_info,
1206                job_backfill_orders,
1207                hummock_version_stats,
1208                mutation,
1209                &term_id,
1210                self,
1211                render_result,
1212            )?;
1213            independent_checkpoint_job_controls.insert(
1214                job_id,
1215                IndependentCheckpointJobControl::iceberg_v3(
1216                    job_id,
1217                    partial_graph_id,
1218                    IndependentCheckpointJobStatus::Ready,
1219                    iceberg_job,
1220                ),
1221            );
1222        }
1223
1224        // Recover batch refresh jobs (both idle and consuming snapshot).
1225        // Actors were already rendered by `render_runtime_info()`.
1226        for (job_id, render_result) in batch_refresh {
1227            creating_jobs.remove(&job_id);
1228            debug!(%job_id, "recovered batch refresh job");
1229
1230            // Resolve committed epoch from state tables.
1231            let committed_epoch = resolve_jobs_committed_epoch(
1232                state_table_committed_epochs,
1233                InflightFragmentInfo::existing_table_ids(render_result.fragment_infos.values()),
1234            );
1235
1236            let snapshot_backfill_info = StreamFragmentGraph::collect_snapshot_backfill_info_impl(
1237                render_result
1238                    .fragment_infos
1239                    .values()
1240                    .map(|fragment| (&fragment.nodes, fragment.fragment_type_mask)),
1241            )?
1242            .0
1243            .ok_or_else(|| anyhow!("batch refresh job {} has no snapshot backfill info", job_id))?;
1244
1245            let upstream_table_ids: HashSet<TableId> = snapshot_backfill_info
1246                .upstream_mv_table_id_to_backfill_epoch
1247                .keys()
1248                .copied()
1249                .collect();
1250            let snapshot_epoch = snapshot_backfill_info
1251                .upstream_mv_table_id_to_backfill_epoch
1252                .values()
1253                .find_map(|e| *e)
1254                .unwrap_or(committed_epoch);
1255
1256            let job_backfill_orders = recover_job_backfill_order(
1257                job_extra_info,
1258                job_id,
1259                fragment_relations,
1260                &render_result.fragment_infos,
1261            );
1262            let mutation = build_mutation(
1263                &Default::default(), // batch refresh has no source splits
1264                Default::default(),
1265                &job_backfill_orders,
1266                false,
1267            );
1268
1269            let refresh_interval_sec = job_extra_info
1270                .get(&job_id)
1271                .and_then(|info| info.refresh_interval_sec)
1272                .expect("batch refresh job should have refresh_interval_sec in job extra info");
1273            let partial_graph_id = to_partial_graph_id(database_id, Some(job_id));
1274            let control = IndependentJobControl::recovered(
1275                IndependentJobInfo::from_fragment_infos(
1276                    database_id,
1277                    job_id,
1278                    snapshot_epoch,
1279                    upstream_table_ids,
1280                    render_result.fragment_infos.values(),
1281                ),
1282                committed_epoch,
1283            );
1284
1285            let job = BatchRefreshJobCheckpointControl::recover(
1286                control,
1287                job_backfill_orders,
1288                hummock_version_stats,
1289                mutation,
1290                render_result,
1291                &term_id,
1292                self,
1293                refresh_interval_sec,
1294            )?;
1295            independent_checkpoint_job_controls.insert(
1296                job_id,
1297                IndependentCheckpointJobControl::batch_refresh(
1298                    job_id,
1299                    partial_graph_id,
1300                    IndependentCheckpointJobStatus::Ready,
1301                    job,
1302                ),
1303            );
1304        }
1305
1306        self.control_stream_manager()
1307            .env
1308            .shared_actor_infos()
1309            .recover_database(
1310                database_id,
1311                database_jobs
1312                    .values()
1313                    .flat_map(|info| {
1314                        info.fragment_infos()
1315                            .map(move |fragment| (fragment, info.job_id))
1316                    })
1317                    .chain(
1318                        independent_checkpoint_job_controls
1319                            .iter()
1320                            .flat_map(|(job_id, job)| {
1321                                let job_id = *job_id;
1322                                job.fragment_infos()
1323                                    .into_iter()
1324                                    .flat_map(move |infos| infos.values().map(move |f| (f, job_id)))
1325                            }),
1326                    )
1327                    .chain(independent_checkpoint_job_controls.iter().filter_map(
1328                        |(job_id, job)| match job.running() {
1329                            Some(IndependentCheckpointJob::IcebergV3(job)) => {
1330                                Some((job.resolver_fragment_info(), *job_id))
1331                            }
1332                            _ => None,
1333                        },
1334                    )),
1335            );
1336
1337        let committed_epoch = barrier_info.prev_epoch();
1338        let new_epoch = barrier_info.curr_epoch;
1339        let database_info = InflightDatabaseInfo::recover(
1340            database_id,
1341            database_jobs.into_values(),
1342            self.control_stream_manager()
1343                .env
1344                .shared_actor_infos()
1345                .clone(),
1346        );
1347        let database_state = BarrierWorkerState::recovery(new_epoch, is_paused);
1348        Ok(DatabaseCheckpointControl::recovery(
1349            database_id,
1350            term_id,
1351            database_state,
1352            committed_epoch,
1353            database_info,
1354            independent_checkpoint_job_controls,
1355        ))
1356    }
1357}
1358
1359impl ControlStreamManager {
1360    fn connected_workers(&self) -> impl Iterator<Item = (WorkerId, &ControlStreamNode)> + '_ {
1361        self.workers
1362            .iter()
1363            .filter_map(|(worker_id, (_, worker_state))| match worker_state {
1364                WorkerNodeState::Connected { control_stream, .. } => {
1365                    Some((*worker_id, control_stream))
1366                }
1367                WorkerNodeState::Reconnecting(_) => None,
1368            })
1369    }
1370
1371    pub(super) fn inject_barrier(
1372        &mut self,
1373        partial_graph_id: PartialGraphId,
1374        mutation: Option<Mutation>,
1375        barrier_info: &BarrierInfo,
1376        node_actors: &HashMap<WorkerId, HashSet<ActorId>>,
1377        table_ids_to_sync: impl Iterator<Item = TableId>,
1378        nodes_to_sync_table: impl Iterator<Item = WorkerId>,
1379        mut new_actors: Option<StreamJobActorsToCreate>,
1380    ) -> MetaResult<NodeToCollect> {
1381        fail_point!("inject_barrier_err", |_| risingwave_common::bail!(
1382            "inject_barrier_err"
1383        ));
1384
1385        let nodes_to_sync_table: HashSet<_> = nodes_to_sync_table.collect();
1386
1387        nodes_to_sync_table.iter().for_each(|worker_id| {
1388            assert!(node_actors.contains_key(worker_id), "worker_id {worker_id} in nodes_to_sync_table {nodes_to_sync_table:?} but not in node_actors {node_actors:?}");
1389        });
1390
1391        let mut node_need_collect = NodeToCollect::new();
1392        let table_ids_to_sync = table_ids_to_sync.collect_vec();
1393
1394        node_actors.iter()
1395            .try_for_each(|(worker_id, actor_ids_to_collect)| {
1396                assert!(!actor_ids_to_collect.is_empty(), "empty actor_ids_to_collect on worker {worker_id} in node_actors {node_actors:?}");
1397                let table_ids_to_sync = if nodes_to_sync_table.contains(worker_id) {
1398                    table_ids_to_sync.clone()
1399                } else {
1400                    vec![]
1401                };
1402
1403                let node = if let Some((_, worker_state)) = self.workers.get(worker_id)
1404                    &&
1405                    let WorkerNodeState::Connected { control_stream, .. } = worker_state
1406                {
1407                    control_stream
1408                } else {
1409                    return Err(anyhow!("unconnected worker node {}", worker_id).into());
1410                };
1411
1412                {
1413                    let mutation = mutation.clone();
1414                    let barrier = Barrier {
1415                        epoch: Some(risingwave_pb::data::Epoch {
1416                            curr: barrier_info.curr_epoch(),
1417                            prev: barrier_info.prev_epoch(),
1418                        }),
1419                        mutation: mutation.clone().map(|_| BarrierMutation { mutation }),
1420                        tracing_context: TracingContext::from_span(barrier_info.curr_epoch.span())
1421                            .to_protobuf(),
1422                        kind: barrier_info.kind.to_protobuf() as i32,
1423                    };
1424
1425                    node.handle
1426                        .request_sender
1427                        .send(StreamingControlStreamRequest {
1428                            request: Some(
1429                                streaming_control_stream_request::Request::InjectBarrier(
1430                                    InjectBarrierRequest {
1431                                        request_id: Uuid::new_v4().to_string(),
1432                                        barrier: Some(barrier),
1433                                        actor_ids_to_collect: actor_ids_to_collect.iter().copied().collect(),
1434                                        table_ids_to_sync,
1435                                        partial_graph_id,
1436                                        actors_to_build: new_actors
1437                                            .as_mut()
1438                                            .map(|new_actors| new_actors.remove(worker_id))
1439                                            .into_iter()
1440                                            .flatten()
1441                                            .flatten()
1442                                            .map(|(fragment_id, (node, actors, initial_subscriber_ids))| {
1443                                                FragmentBuildActorInfo {
1444                                                    fragment_id,
1445                                                    node: Some(node),
1446                                                    actors: actors
1447                                                        .into_iter()
1448                                                        .map(|(actor, upstreams, dispatchers)| {
1449                                                            BuildActorInfo {
1450                                                                actor_id: actor.actor_id,
1451                                                                fragment_upstreams: upstreams
1452                                                                    .into_iter()
1453                                                                    .map(|(fragment_id, upstreams)| {
1454                                                                        (
1455                                                                            fragment_id,
1456                                                                            UpstreamActors {
1457                                                                                actors: upstreams
1458                                                                                    .into_values()
1459                                                                                    .collect(),
1460                                                                            },
1461                                                                        )
1462                                                                    })
1463                                                                    .collect(),
1464                                                                dispatchers,
1465                                                                vnode_bitmap: actor.vnode_bitmap.map(|bitmap| bitmap.to_protobuf()),
1466                                                                mview_definition: actor.mview_definition,
1467                                                                expr_context: actor.expr_context,
1468                                                                config_override: actor.config_override.to_string(),
1469                                                                initial_subscriber_ids: initial_subscriber_ids.iter().copied().collect(),
1470                                                            }
1471                                                        })
1472                                                        .collect(),
1473                                                }
1474                                            })
1475                                            .collect(),
1476                                    },
1477                                ),
1478                            ),
1479                        })
1480                        .map_err(|_| {
1481                            MetaError::from(anyhow!(
1482                                "failed to send request to {} {:?}",
1483                                node.worker_id,
1484                                node.host
1485                            ))
1486                        })?;
1487
1488                    node_need_collect.insert(*worker_id);
1489                    Result::<_, MetaError>::Ok(())
1490                }
1491            })
1492            .inspect_err(|e| {
1493                // Record failure in event log.
1494                use risingwave_pb::meta::event_log;
1495                let event = event_log::EventInjectBarrierFail {
1496                    prev_epoch: barrier_info.prev_epoch(),
1497                    cur_epoch: barrier_info.curr_epoch(),
1498                    error: e.to_report_string(),
1499                };
1500                self.env
1501                    .event_log_manager_ref()
1502                    .add_event_logs(vec![event_log::Event::InjectBarrierFail(event)]);
1503            })?;
1504        Ok(node_need_collect)
1505    }
1506
1507    pub(super) fn add_partial_graph(&mut self, partial_graph_id: PartialGraphId, term_id: &str) {
1508        self.connected_workers().for_each(|(_, node)| {
1509            if node
1510                .handle
1511                .request_sender
1512                .send(StreamingControlStreamRequest {
1513                    request: Some(
1514                        streaming_control_stream_request::Request::CreatePartialGraph(
1515                            CreatePartialGraphRequest {
1516                                partial_graph_id,
1517                                term_id: term_id.to_owned(),
1518                            },
1519                        ),
1520                    ),
1521                }).is_err() {
1522                let (database_id, creating_job_id) = from_partial_graph_id(partial_graph_id);
1523                warn!(%database_id, ?creating_job_id, worker_id = %node.worker_id, "failed to add the partial graph to the worker")
1524            }
1525        });
1526    }
1527
1528    pub(super) fn remove_partial_graphs(&mut self, partial_graph_ids: Vec<PartialGraphId>) {
1529        self.connected_workers().for_each(|(_, node)| {
1530            if node.handle
1531                .request_sender
1532                .send(StreamingControlStreamRequest {
1533                    request: Some(
1534                        streaming_control_stream_request::Request::RemovePartialGraph(
1535                            RemovePartialGraphRequest {
1536                                partial_graph_ids: partial_graph_ids.clone(),
1537                            },
1538                        ),
1539                    ),
1540                })
1541                .is_err()
1542            {
1543                warn!(worker_id = %node.worker_id,node = ?node.host,"failed to send remove partial graph request");
1544            }
1545        })
1546    }
1547
1548    pub(super) fn reset_partial_graphs(
1549        &mut self,
1550        partial_graph_ids: Vec<PartialGraphId>,
1551    ) -> HashSet<WorkerId> {
1552        self.connected_workers()
1553            .filter_map(|(worker_id, node)| {
1554                if node
1555                    .handle
1556                    .request_sender
1557                    .send(StreamingControlStreamRequest {
1558                        request: Some(
1559                            streaming_control_stream_request::Request::ResetPartialGraphs(
1560                                ResetPartialGraphsRequest {
1561                                    partial_graph_ids: partial_graph_ids.clone(),
1562                                },
1563                            ),
1564                        ),
1565                    })
1566                    .is_err()
1567                {
1568                    warn!(%worker_id, node = ?node.host,"failed to send reset database request");
1569                    None
1570                } else {
1571                    Some(worker_id)
1572                }
1573            })
1574            .collect()
1575    }
1576}
1577
1578impl GlobalBarrierWorkerContextImpl {
1579    pub(super) async fn new_control_stream_impl(
1580        &self,
1581        node: &WorkerNode,
1582    ) -> MetaResult<StreamingControlHandle> {
1583        let handle = self
1584            .env
1585            .stream_client_pool()
1586            .get(node)
1587            .await?
1588            .start_streaming_control(PbInitRequest::default())
1589            .await?;
1590        Ok(handle)
1591    }
1592}
1593
1594pub(super) fn merge_node_rpc_errors<E: Error + Send + Sync + 'static>(
1595    message: &str,
1596    errors: impl IntoIterator<Item = (WorkerId, E)>,
1597) -> MetaError {
1598    use std::fmt::Write;
1599
1600    use risingwave_common::error::error_request_copy;
1601    use risingwave_common::error::tonic::extra::Score;
1602
1603    let errors = errors.into_iter().collect_vec();
1604
1605    if errors.is_empty() {
1606        return anyhow!(message.to_owned()).into();
1607    }
1608
1609    // Create the error from the single error.
1610    let single_error = |(worker_id, e)| {
1611        anyhow::Error::from(e)
1612            .context(format!("{message}, in worker node {worker_id}"))
1613            .into()
1614    };
1615
1616    if errors.len() == 1 {
1617        return single_error(errors.into_iter().next().unwrap());
1618    }
1619
1620    // Find the error with the highest score.
1621    let max_score = errors
1622        .iter()
1623        .filter_map(|(_, e)| error_request_copy::<Score>(e))
1624        .max();
1625
1626    if let Some(max_score) = max_score {
1627        let mut errors = errors;
1628        let max_scored = errors
1629            .extract_if(.., |(_, e)| {
1630                error_request_copy::<Score>(e) == Some(max_score)
1631            })
1632            .next()
1633            .unwrap();
1634
1635        return single_error(max_scored);
1636    }
1637
1638    // The errors do not have scores, so simply concatenate them.
1639    let concat: String = errors
1640        .into_iter()
1641        .fold(format!("{message}: "), |mut s, (w, e)| {
1642            write!(&mut s, " in worker node {}, {};", w, e.as_report()).unwrap();
1643            s
1644        });
1645    anyhow!(concat).into()
1646}
1647
1648#[cfg(test)]
1649mod test_partial_graph_id {
1650    use crate::barrier::rpc::{from_partial_graph_id, to_partial_graph_id};
1651
1652    #[test]
1653    fn test_partial_graph_id_conversion() {
1654        let database_id = 233.into();
1655        let job_id = 233.into();
1656        assert_eq!(
1657            (database_id, None),
1658            from_partial_graph_id(to_partial_graph_id(database_id, None))
1659        );
1660        assert_eq!(
1661            (database_id, Some(job_id)),
1662            from_partial_graph_id(to_partial_graph_id(database_id, Some(job_id)))
1663        );
1664    }
1665}