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