risingwave_meta/barrier/
complete_task.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::HashMap;
16use std::future::{Future, pending};
17use std::mem::replace;
18use std::sync::Arc;
19
20use anyhow::Context;
21use futures::future::try_join_all;
22use risingwave_common::id::JobId;
23use risingwave_common::must_match;
24use risingwave_common::util::deployment::Deployment;
25use risingwave_pb::hummock::HummockVersionStats;
26use risingwave_pb::id::{DatabaseId, PartialGraphId};
27use risingwave_pb::stream_service::barrier_complete_response::{
28    PbIcebergPkIndexSinkMetadata, PbListFinishedSource, PbLoadFinishedSource,
29};
30use tokio::task::JoinHandle;
31
32use crate::barrier::checkpoint::CheckpointControl;
33use crate::barrier::context::GlobalBarrierWorkerContext;
34use crate::barrier::info::BarrierInfo;
35use crate::barrier::partial_graph::{PartialGraphBarrierInfo, PartialGraphManager};
36use crate::barrier::progress::TrackingJob;
37use crate::barrier::rpc::from_partial_graph_id;
38use crate::barrier::schedule::PeriodicBarriers;
39use crate::hummock::CommitEpochInfo;
40use crate::manager::MetaSrvEnv;
41use crate::rpc::metrics::GLOBAL_META_METRICS;
42use crate::{MetaError, MetaResult};
43
44pub(super) enum CompletingTask {
45    None,
46    Completing {
47        #[expect(clippy::type_complexity)]
48        /// `database_id` -> (`Some(database_graph_committed_epoch)`, [(`creating_job_id`, `creating_job_committed_epoch`)])
49        epochs_to_ack: HashMap<DatabaseId, (Option<u64>, Vec<(JobId, u64)>)>,
50
51        // The join handle of a spawned task that completes the barrier.
52        // The return value indicate whether there is some create streaming job command
53        // that has finished but not checkpointed. If there is any, we will force checkpoint on the next barrier
54        join_handle: JoinHandle<MetaResult<HummockVersionStats>>,
55    },
56    #[expect(dead_code)]
57    Err(MetaError),
58}
59
60/// Only for checkpoint barrier. For normal barrier, there won't be a task.
61#[derive(Default)]
62pub(super) struct CompleteBarrierTask {
63    pub(super) commit_info: CommitEpochInfo,
64    pub(super) finished_jobs: Vec<TrackingJob>,
65    pub(super) finished_cdc_table_backfill: Vec<JobId>,
66    /// `partial_graph_id` -> barrier info for post-collect processing
67    pub(super) epoch_infos: HashMap<PartialGraphId, PartialGraphBarrierInfo>,
68    /// Source listing completion events that need `ListFinish` commands
69    pub(super) list_finished_source_ids: Vec<PbListFinishedSource>,
70    /// Source load completion events that need `LoadFinish` commands
71    pub(super) load_finished_source_ids: Vec<PbLoadFinishedSource>,
72    /// Table IDs that have finished materialize refresh and need completion signaling
73    pub(super) refresh_finished_table_job_ids: Vec<JobId>,
74    /// Iceberg pk-index sink reports collected during this barrier
75    pub(super) iceberg_pk_index_sink_metadata: Vec<PbIcebergPkIndexSinkMetadata>,
76}
77
78impl CompleteBarrierTask {
79    #[expect(clippy::type_complexity)]
80    pub(super) fn epochs_to_ack(&self) -> HashMap<DatabaseId, (Option<u64>, Vec<(JobId, u64)>)> {
81        let mut epochs_to_ack: HashMap<DatabaseId, (Option<u64>, Vec<(JobId, u64)>)> =
82            HashMap::new();
83        for (partial_graph_id, info) in &self.epoch_infos {
84            let (database_id, creating_job_id) = from_partial_graph_id(*partial_graph_id);
85            let epoch = info.barrier_info.prev_epoch();
86            let (database, jobs) = epochs_to_ack.entry(database_id).or_default();
87            if let Some(job_id) = creating_job_id {
88                jobs.push((job_id, epoch));
89            } else {
90                *database = Some(epoch);
91            }
92        }
93        epochs_to_ack
94    }
95}
96
97impl CompleteBarrierTask {
98    pub(super) async fn complete_barrier(
99        self,
100        context: &impl GlobalBarrierWorkerContext,
101        env: MetaSrvEnv,
102    ) -> MetaResult<HummockVersionStats> {
103        let mut notifiers = Vec::new();
104        let result: MetaResult<HummockVersionStats> = try {
105            let wait_commit_timer = GLOBAL_META_METRICS
106                .barrier_wait_commit_latency
107                .start_timer();
108
109            // Iceberg pk-index sink metadata reports are handled in three steps around hummock `commit_epoch`:
110            //   1. pre_commit: persist pending rows under `pending_sink_state` (no iceberg I/O).
111            //   2. commit_epoch: advance hummock.
112            //   3. commit: drive iceberg overwrite_files for queued epochs.
113            //   4. advance_committed_epochs: advance the per-partial-graph committed epoch cursor.
114            let mut iceberg_pk_index_commit_sink_ids = Vec::new();
115            if !self.iceberg_pk_index_sink_metadata.is_empty() {
116                let res = context
117                    .pre_commit_iceberg_pk_index_sink_metadata(self.iceberg_pk_index_sink_metadata)
118                    .await?;
119                iceberg_pk_index_commit_sink_ids = res;
120            }
121
122            let version_stats = context.commit_epoch(self.commit_info).await?;
123
124            if !iceberg_pk_index_commit_sink_ids.is_empty() {
125                context
126                    .commit_iceberg_pk_index_sink_metadata(iceberg_pk_index_commit_sink_ids)
127                    .await?;
128            }
129            let epochs = self
130                .epoch_infos
131                .iter()
132                .map(|(id, info)| (*id, info.barrier_info.prev_epoch()));
133            context.advance_iceberg_pk_index_sink_committed_epochs(epochs);
134
135            // Handle list finished source IDs for refreshable batch sources
136            // Spawn this asynchronously to avoid deadlock during barrier collection
137            //
138            // This step is for fs-like refreshable-batch sources, which need to list the data first finishing loading. It guarantees finishing listing before loading.
139            // The other sources can skip this step.
140
141            if !self.list_finished_source_ids.is_empty() {
142                context
143                    .handle_list_finished_source_ids(self.list_finished_source_ids.clone())
144                    .await?;
145            }
146
147            // Handle load finished source IDs for refreshable batch sources
148            // Spawn this asynchronously to avoid deadlock during barrier collection
149            if !self.load_finished_source_ids.is_empty() {
150                context
151                    .handle_load_finished_source_ids(self.load_finished_source_ids.clone())
152                    .await?;
153            }
154
155            // Handle refresh finished table IDs for materialized view refresh completion
156            if !self.refresh_finished_table_job_ids.is_empty() {
157                context
158                    .handle_refresh_finished_table_ids(self.refresh_finished_table_job_ids.clone())
159                    .await?;
160            }
161
162            for (partial_graph_id, info) in self.epoch_infos {
163                let (database_id, job_id) = from_partial_graph_id(partial_graph_id);
164                let command_name = info.post_collect_command.command_name().to_owned();
165                let elapsed_secs = info.elapsed_secs();
166                notifiers.extend(info.notifiers);
167                context
168                    .post_collect_command(info.post_collect_command)
169                    .await?;
170                if job_id.is_none() {
171                    Self::report_complete_event(
172                        &env,
173                        database_id,
174                        elapsed_secs,
175                        &info.barrier_info,
176                        command_name,
177                    );
178                }
179            }
180
181            wait_commit_timer.observe_duration();
182            version_stats
183        };
184
185        let version_stats = {
186            let version_stats = match result {
187                Ok(version_stats) => version_stats,
188                Err(e) => {
189                    for notifier in notifiers {
190                        notifier.notify_collection_failed(e.clone());
191                    }
192                    return Err(e);
193                }
194            };
195            notifiers.into_iter().for_each(|notifier| {
196                notifier.notify_collected();
197            });
198            try_join_all(
199                self.finished_jobs
200                    .into_iter()
201                    .map(|finished_job| context.finish_creating_job(finished_job)),
202            )
203            .await?;
204            try_join_all(
205                self.finished_cdc_table_backfill
206                    .into_iter()
207                    .map(|job_id| context.finish_cdc_table_backfill(job_id)),
208            )
209            .await?;
210            version_stats
211        };
212
213        Ok(version_stats)
214    }
215}
216
217impl CompleteBarrierTask {
218    fn report_complete_event(
219        env: &MetaSrvEnv,
220        database_id: DatabaseId,
221        duration_sec: f64,
222        barrier_info: &BarrierInfo,
223        command: String,
224    ) {
225        // Record barrier latency in event log.
226        use risingwave_pb::meta::event_log;
227        let event = event_log::EventBarrierComplete {
228            prev_epoch: barrier_info.prev_epoch(),
229            cur_epoch: barrier_info.curr_epoch(),
230            duration_sec,
231            command,
232            barrier_kind: barrier_info.kind.as_str_name().to_owned(),
233            database_id: database_id.as_raw_id(),
234        };
235        if cfg!(debug_assertions) || Deployment::current().is_ci() {
236            // Add a warning log so that debug mode / CI can observe it
237            if duration_sec > 5.0 {
238                tracing::warn!(event = ?event,"high barrier latency observed!")
239            }
240        }
241        env.event_log_manager_ref()
242            .add_event_logs(vec![event_log::Event::BarrierComplete(event)]);
243    }
244}
245
246pub(super) struct BarrierCompleteOutput {
247    #[expect(clippy::type_complexity)]
248    /// `database_id` -> (`Some(database_graph_committed_epoch)`, [(`creating_job_id`, `creating_job_committed_epoch`)])
249    pub epochs_to_ack: HashMap<DatabaseId, (Option<u64>, Vec<(JobId, u64)>)>,
250    pub hummock_version_stats: HummockVersionStats,
251}
252
253impl CompletingTask {
254    pub(super) fn next_completed_barrier<'a>(
255        &'a mut self,
256        periodic_barriers: &mut PeriodicBarriers,
257        checkpoint_control: &mut CheckpointControl,
258        partial_graph_manager: &mut PartialGraphManager,
259        context: &Arc<impl GlobalBarrierWorkerContext>,
260        env: &MetaSrvEnv,
261    ) -> impl Future<Output = MetaResult<BarrierCompleteOutput>> + 'a {
262        // If there is no completing barrier, try to start completing the earliest barrier if
263        // it has been collected.
264        if let CompletingTask::None = self
265            && let Some(task) = checkpoint_control
266                .next_complete_barrier_task(periodic_barriers, partial_graph_manager)
267        {
268            {
269                let epochs_to_ack = task.epochs_to_ack();
270                let context = context.clone();
271                let await_tree_reg = env.await_tree_reg().clone();
272                let env = env.clone();
273
274                let fut = async move { task.complete_barrier(&*context, env).await };
275                let fut = await_tree_reg
276                    .register_derived_root("Barrier Completion Task")
277                    .instrument(fut);
278                let join_handle = tokio::spawn(fut);
279
280                *self = CompletingTask::Completing {
281                    epochs_to_ack,
282                    join_handle,
283                };
284            }
285        }
286
287        async move {
288            if !matches!(self, CompletingTask::Completing { .. }) {
289                return pending().await;
290            };
291            self.next_completed_barrier_inner().await
292        }
293    }
294
295    #[await_tree::instrument]
296    pub(super) async fn wait_completing_task(
297        &mut self,
298    ) -> MetaResult<Option<BarrierCompleteOutput>> {
299        match self {
300            CompletingTask::None => Ok(None),
301            CompletingTask::Completing { .. } => {
302                self.next_completed_barrier_inner().await.map(Some)
303            }
304            CompletingTask::Err(_) => {
305                unreachable!("should not be called on previous err")
306            }
307        }
308    }
309
310    async fn next_completed_barrier_inner(&mut self) -> MetaResult<BarrierCompleteOutput> {
311        let CompletingTask::Completing { join_handle, .. } = self else {
312            unreachable!()
313        };
314
315        {
316            {
317                let join_result: MetaResult<_> = try {
318                    join_handle
319                        .await
320                        .context("failed to join completing command")
321                        .map_err(MetaError::from)??
322                };
323                // It's important to reset the completing_command after await no matter the result is err
324                // or not, and otherwise the join handle will be polled again after ready.
325                let next_completing_command_status = if let Err(e) = &join_result {
326                    CompletingTask::Err(e.clone())
327                } else {
328                    CompletingTask::None
329                };
330                let completed_command = replace(self, next_completing_command_status);
331                let hummock_version_stats = join_result?;
332
333                must_match!(completed_command, CompletingTask::Completing {
334                    epochs_to_ack,
335                    ..
336                } => {
337                    Ok(BarrierCompleteOutput {
338                        epochs_to_ack,
339                        hummock_version_stats,
340                    })
341                })
342            }
343        }
344    }
345}
risingwave_meta/barrier/complete_task.rs

risingwave_meta/barrier/
complete_task.rs
