1use 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 epochs_to_ack: HashMap<DatabaseId, (Option<u64>, Vec<(JobId, u64)>)>,
51
52 join_handle: JoinHandle<MetaResult<HummockVersionStats>>,
56 },
57 #[expect(dead_code)]
58 Err(MetaError),
59}
60
61#[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 pub(super) epoch_infos: HashMap<PartialGraphId, PartialGraphBarrierInfo>,
69 pub(super) list_finished_source_ids: Vec<PbListFinishedSource>,
71 pub(super) load_finished_source_ids: Vec<PbLoadFinishedSource>,
73 pub(super) refresh_finished_actors: Vec<PbRefreshFinishedActor>,
75 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 let mut iceberg_pk_index_commit_sink_ids = Vec::new();
119 if !self.iceberg_pk_index_pre_commit_metadata.is_empty() {
120 let res = context
121 .pre_commit_iceberg_pk_index_sink_metadata(
122 self.iceberg_pk_index_pre_commit_metadata,
123 )
124 .await?;
125 iceberg_pk_index_commit_sink_ids = res;
126 }
127
128 let version_stats = context.commit_epoch(self.commit_info).await?;
129
130 if !iceberg_pk_index_commit_sink_ids.is_empty() {
131 context
132 .commit_iceberg_pk_index_sink_metadata(iceberg_pk_index_commit_sink_ids)
133 .await?;
134 }
135 let epochs = self
136 .epoch_infos
137 .iter()
138 .map(|(id, info)| (*id, info.barrier_info.prev_epoch()));
139 context.advance_iceberg_pk_index_sink_committed_epochs(epochs);
140
141 for (partial_graph_id, info) in self.epoch_infos {
142 let (database_id, job_id) = from_partial_graph_id(partial_graph_id);
143 let command_name = info.post_collect_command.command_name().to_owned();
144 let elapsed_secs = info.elapsed_secs();
145 notifiers.extend(info.notifier);
146 context
147 .post_collect_command(info.post_collect_command)
148 .await?;
149 if job_id.is_none() {
150 Self::report_complete_event(
151 &env,
152 database_id,
153 elapsed_secs,
154 &info.barrier_info,
155 command_name,
156 );
157 }
158 }
159
160 if !self.list_finished_source_ids.is_empty() {
163 context
164 .handle_list_finished_source_ids(self.list_finished_source_ids)
165 .await?;
166 }
167 if !self.load_finished_source_ids.is_empty() {
168 context
169 .handle_load_finished_source_ids(self.load_finished_source_ids)
170 .await?;
171 }
172 if !self.refresh_finished_actors.is_empty() {
173 context
174 .handle_refresh_finished_actors(self.refresh_finished_actors)
175 .await?;
176 }
177
178 wait_commit_timer.observe_duration();
179 version_stats
180 };
181
182 let version_stats = {
183 let version_stats = match result {
184 Ok(version_stats) => version_stats,
185 Err(e) => {
186 for notifier in notifiers {
187 notifier.notify_collection_failed(e.clone());
188 }
189 return Err(e);
190 }
191 };
192 notifiers.into_iter().for_each(|notifier| {
193 notifier.notify_collected();
194 });
195 try_join_all(
196 self.finished_jobs
197 .into_iter()
198 .map(|finished_job| context.finish_creating_job(finished_job)),
199 )
200 .await?;
201 try_join_all(
202 self.finished_cdc_table_backfill
203 .into_iter()
204 .map(|job_id| context.finish_cdc_table_backfill(job_id)),
205 )
206 .await?;
207 version_stats
208 };
209
210 Ok(version_stats)
211 }
212}
213
214impl CompleteBarrierTask {
215 fn report_complete_event(
216 env: &MetaSrvEnv,
217 database_id: DatabaseId,
218 duration_sec: f64,
219 barrier_info: &BarrierInfo,
220 command: String,
221 ) {
222 use risingwave_pb::meta::event_log;
224 let event = event_log::EventBarrierComplete {
225 prev_epoch: barrier_info.prev_epoch(),
226 cur_epoch: barrier_info.curr_epoch(),
227 duration_sec,
228 command,
229 barrier_kind: barrier_info.kind.as_str_name().to_owned(),
230 database_id: database_id.as_raw_id(),
231 };
232 if cfg!(debug_assertions) || Deployment::current().is_ci() {
233 if duration_sec > 5.0 {
235 tracing::warn!(event = ?event,"high barrier latency observed!")
236 }
237 }
238 env.event_log_manager_ref()
239 .add_event_logs(vec![event_log::Event::BarrierComplete(event)]);
240 }
241}
242
243pub(super) struct BarrierCompleteOutput {
244 #[expect(clippy::type_complexity)]
245 pub epochs_to_ack: HashMap<DatabaseId, (Option<u64>, Vec<(JobId, u64)>)>,
247 pub hummock_version_stats: HummockVersionStats,
248}
249
250impl CompletingTask {
251 pub(super) fn next_completed_barrier<'a>(
252 &'a mut self,
253 periodic_barriers: &mut PeriodicBarriers,
254 checkpoint_control: &mut CheckpointControl,
255 partial_graph_manager: &mut PartialGraphManager,
256 context: &Arc<impl GlobalBarrierWorkerContext>,
257 env: &MetaSrvEnv,
258 ) -> impl Future<Output = MetaResult<BarrierCompleteOutput>> + 'a {
259 if let CompletingTask::None = self
262 && let Some(task) = checkpoint_control
263 .next_complete_barrier_task(periodic_barriers, partial_graph_manager)
264 {
265 {
266 let epochs_to_ack = task.epochs_to_ack();
267 let context = context.clone();
268 let await_tree_reg = env.await_tree_reg().clone();
269 let env = env.clone();
270
271 let fut = async move { task.complete_barrier(&*context, env).await };
272 let fut = await_tree_reg
273 .register_derived_root("Barrier Completion Task")
274 .instrument(fut);
275 let join_handle = tokio::spawn(fut);
276
277 *self = CompletingTask::Completing {
278 epochs_to_ack,
279 join_handle,
280 };
281 }
282 }
283
284 async move {
285 if !matches!(self, CompletingTask::Completing { .. }) {
286 return pending().await;
287 };
288 self.next_completed_barrier_inner().await
289 }
290 }
291
292 #[await_tree::instrument]
293 pub(super) async fn wait_completing_task(
294 &mut self,
295 ) -> MetaResult<Option<BarrierCompleteOutput>> {
296 match self {
297 CompletingTask::None => Ok(None),
298 CompletingTask::Completing { .. } => {
299 self.next_completed_barrier_inner().await.map(Some)
300 }
301 CompletingTask::Err(_) => {
302 unreachable!("should not be called on previous err")
303 }
304 }
305 }
306
307 async fn next_completed_barrier_inner(&mut self) -> MetaResult<BarrierCompleteOutput> {
308 let CompletingTask::Completing { join_handle, .. } = self else {
309 unreachable!()
310 };
311
312 {
313 {
314 let join_result: MetaResult<_> = try {
315 join_handle
316 .await
317 .context("failed to join completing command")
318 .map_err(MetaError::from)??
319 };
320 let next_completing_command_status = if let Err(e) = &join_result {
323 CompletingTask::Err(e.clone())
324 } else {
325 CompletingTask::None
326 };
327 let completed_command = replace(self, next_completing_command_status);
328 let hummock_version_stats = join_result?;
329
330 must_match!(completed_command, CompletingTask::Completing {
331 epochs_to_ack,
332 ..
333 } => {
334 Ok(BarrierCompleteOutput {
335 epochs_to_ack,
336 hummock_version_stats,
337 })
338 })
339 }
340 }
341 }
342}