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

risingwave_meta/barrier/
complete_task.rs
