1mod block_stream;
16mod compaction_executor;
17mod compaction_filter;
18pub mod compaction_utils;
19mod iceberg_compaction;
20use risingwave_hummock_sdk::compact_task::{CompactTask, ValidationTask};
21use risingwave_pb::compactor::{DispatchCompactionTaskRequest, dispatch_compaction_task_request};
22use risingwave_pb::hummock::PbCompactTask;
23use risingwave_pb::hummock::report_compaction_task_request::{
24 Event as ReportCompactionTaskEvent, HeartBeat as SharedHeartBeat,
25 ReportTask as ReportSharedTask,
26};
27use risingwave_pb::iceberg_compaction::{
28 SubscribeIcebergCompactionEventRequest, SubscribeIcebergCompactionEventResponse,
29 subscribe_iceberg_compaction_event_request,
30};
31use risingwave_rpc_client::GrpcCompactorProxyClient;
32use thiserror_ext::AsReport;
33use tokio::sync::mpsc;
34use tonic::Request;
35
36pub mod compactor_runner;
37mod context;
38pub mod fast_compactor_runner;
39mod iterator;
40mod shared_buffer_compact;
41pub(super) mod task_progress;
42
43use std::collections::hash_map::Entry;
44use std::collections::{HashMap, VecDeque};
45use std::marker::PhantomData;
46use std::sync::atomic::{AtomicU32, AtomicU64, Ordering};
47use std::sync::{Arc, Mutex};
48use std::time::{Duration, SystemTime};
49
50use await_tree::{InstrumentAwait, SpanExt};
51pub use compaction_executor::CompactionExecutor;
52pub use compaction_filter::{
53 CompactionFilter, DummyCompactionFilter, MultiCompactionFilter, TtlCompactionFilter,
54};
55pub use context::{
56 CompactionAwaitTreeRegRef, CompactorContext, await_tree_key, new_compaction_await_tree_reg_ref,
57};
58use futures::{StreamExt, pin_mut};
59use iceberg_compaction::iceberg_compactor_runner::IcebergCompactorRunnerConfigBuilder;
61#[cfg(madsim)]
62pub use iceberg_compaction::set_simulated_pk_index_compaction_result;
63#[cfg(madsim)]
64use iceberg_compaction::simulated_pk_index_compaction_result;
65use iceberg_compaction::{
66 IcebergPlanCompletion, IcebergTaskExecution, IcebergTaskQueue, IcebergTaskReport,
67 IcebergTaskTracker, PushResult, ReportSendResult, build_drained_iceberg_task_report,
68 build_iceberg_task_report, create_task_execution, flush_pending_iceberg_task_reports,
69 send_or_buffer_iceberg_task_report,
70};
71pub use iterator::{ConcatSstableIterator, SstableStreamIterator};
72use more_asserts::assert_ge;
73use risingwave_hummock_sdk::table_stats::{TableStatsMap, to_prost_table_stats_map};
74use risingwave_hummock_sdk::{
75 HummockCompactionTaskId, HummockSstableObjectId, LocalSstableInfo, compact_task_to_string,
76};
77use risingwave_pb::hummock::compact_task::TaskStatus;
78use risingwave_pb::hummock::subscribe_compaction_event_request::{
79 Event as RequestEvent, HeartBeat, PullTask, ReportTask,
80};
81use risingwave_pb::hummock::subscribe_compaction_event_response::Event as ResponseEvent;
82use risingwave_pb::hummock::{
83 CompactTaskProgress, PbSstableFilterLayout, PbSstableFilterType, ReportCompactionTaskRequest,
84 SubscribeCompactionEventRequest, SubscribeCompactionEventResponse,
85};
86use risingwave_pb::id::IcebergCompactionTaskId;
87use risingwave_rpc_client::HummockMetaClient;
88pub use shared_buffer_compact::compact;
89use tokio::sync::oneshot::Sender;
90use tokio::task::JoinHandle;
91use tokio::time::Instant;
92
93pub use self::compaction_utils::{
94 CompactionStatistics, RemoteBuilderFactory, TaskConfig, check_compaction_result,
95 check_flush_result,
96};
97pub use self::task_progress::TaskProgress;
98use super::multi_builder::CapacitySplitTableBuilder;
99use super::{
100 GetObjectId, HummockError, HummockErrorInner, HummockResult, ObjectIdManager,
101 SstableBuilderOptions, Xor8FilterBuilder, Xor16FilterBuilder,
102};
103use crate::compaction_catalog_manager::{
104 CompactionCatalogAgentRef, CompactionCatalogManager, CompactionCatalogManagerRef,
105};
106use crate::hummock::compactor::compaction_utils::calculate_task_parallelism;
107use crate::hummock::compactor::compactor_runner::{compact_and_build_sst, compact_done};
108use crate::hummock::compactor::iceberg_compaction::TaskKey;
109use crate::hummock::iterator::{Forward, HummockIterator};
110use crate::hummock::{
111 BlockedXor8FilterBuilder, BlockedXor16FilterBuilder, FilterBuilder, NoneFilterBuilder,
112 SharedComapctorObjectIdManager, SstableWriterFactory, StreamingSstableWriterFactory,
113 validate_ssts,
114};
115use crate::monitor::CompactorMetrics;
116
117const COMPACTION_HEARTBEAT_LOG_INTERVAL: Duration = Duration::from_secs(60);
119
120#[derive(Debug, Clone, PartialEq, Eq)]
122struct CompactionLogState {
123 running_parallelism: u32,
124 pull_task_ack: bool,
125 pending_pull_task_count: u32,
126}
127
128#[derive(Debug, Clone, PartialEq, Eq)]
130struct IcebergCompactionLogState {
131 running_parallelism: u32,
132 waiting_parallelism: u32,
133 available_parallelism: u32,
134 pull_task_ack: bool,
135 pending_pull_task_count: u32,
136}
137
138struct LogThrottler<T: PartialEq> {
140 last_logged_state: Option<T>,
141 last_heartbeat: Instant,
142 heartbeat_interval: Duration,
143}
144
145impl<T: PartialEq> LogThrottler<T> {
146 fn new(heartbeat_interval: Duration) -> Self {
147 Self {
148 last_logged_state: None,
149 last_heartbeat: Instant::now(),
150 heartbeat_interval,
151 }
152 }
153
154 fn should_log(&self, current_state: &T) -> bool {
156 self.last_logged_state.as_ref() != Some(current_state)
157 || self.last_heartbeat.elapsed() >= self.heartbeat_interval
158 }
159
160 fn update(&mut self, current_state: T) {
162 self.last_logged_state = Some(current_state);
163 self.last_heartbeat = Instant::now();
164 }
165}
166
167pub struct Compactor {
169 context: CompactorContext,
171 object_id_getter: Arc<dyn GetObjectId>,
172 task_config: TaskConfig,
173 options: SstableBuilderOptions,
174 get_id_time: Arc<AtomicU64>,
175}
176
177pub type CompactOutput = (usize, Vec<LocalSstableInfo>, CompactionStatistics);
178
179impl Compactor {
180 pub fn new(
182 context: CompactorContext,
183 options: SstableBuilderOptions,
184 task_config: TaskConfig,
185 object_id_getter: Arc<dyn GetObjectId>,
186 ) -> Self {
187 Self {
188 context,
189 options,
190 task_config,
191 get_id_time: Arc::new(AtomicU64::new(0)),
192 object_id_getter,
193 }
194 }
195
196 async fn compact_key_range(
201 &self,
202 iter: impl HummockIterator<Direction = Forward>,
203 compaction_filter: impl CompactionFilter,
204 compaction_catalog_agent_ref: CompactionCatalogAgentRef,
205 task_progress: Option<Arc<TaskProgress>>,
206 task_id: Option<HummockCompactionTaskId>,
207 split_index: Option<usize>,
208 ) -> HummockResult<(Vec<LocalSstableInfo>, CompactionStatistics)> {
209 let compact_timer = if self.context.is_share_buffer_compact {
211 self.context
212 .compactor_metrics
213 .write_build_l0_sst_duration
214 .start_timer()
215 } else {
216 self.context
217 .compactor_metrics
218 .compact_sst_duration
219 .start_timer()
220 };
221
222 let (split_table_outputs, table_stats_map) = {
223 let factory = StreamingSstableWriterFactory::new(self.context.sstable_store.clone());
224 match (
225 self.task_config.sstable_filter_type,
226 self.task_config.sstable_filter_layout,
227 ) {
228 (PbSstableFilterType::SstableFilterNone, _) => {
229 self.compact_key_range_impl::<_, NoneFilterBuilder>(
230 factory,
231 iter,
232 compaction_filter,
233 compaction_catalog_agent_ref,
234 task_progress.clone(),
235 self.object_id_getter.clone(),
236 )
237 .instrument_await("compact".verbose())
238 .await?
239 }
240 (PbSstableFilterType::SstableFilterXor8, PbSstableFilterLayout::Blocked) => {
241 self.compact_key_range_impl::<_, BlockedXor8FilterBuilder>(
242 factory,
243 iter,
244 compaction_filter,
245 compaction_catalog_agent_ref,
246 task_progress.clone(),
247 self.object_id_getter.clone(),
248 )
249 .instrument_await("compact".verbose())
250 .await?
251 }
252 (PbSstableFilterType::SstableFilterXor8, PbSstableFilterLayout::Plain) => {
253 self.compact_key_range_impl::<_, Xor8FilterBuilder>(
254 factory,
255 iter,
256 compaction_filter,
257 compaction_catalog_agent_ref,
258 task_progress.clone(),
259 self.object_id_getter.clone(),
260 )
261 .instrument_await("compact".verbose())
262 .await?
263 }
264 (PbSstableFilterType::SstableFilterXor16, PbSstableFilterLayout::Blocked) => {
265 self.compact_key_range_impl::<_, BlockedXor16FilterBuilder>(
266 factory,
267 iter,
268 compaction_filter,
269 compaction_catalog_agent_ref,
270 task_progress.clone(),
271 self.object_id_getter.clone(),
272 )
273 .instrument_await("compact".verbose())
274 .await?
275 }
276 (PbSstableFilterType::SstableFilterXor16, PbSstableFilterLayout::Plain) => {
277 self.compact_key_range_impl::<_, Xor16FilterBuilder>(
278 factory,
279 iter,
280 compaction_filter,
281 compaction_catalog_agent_ref,
282 task_progress.clone(),
283 self.object_id_getter.clone(),
284 )
285 .instrument_await("compact".verbose())
286 .await?
287 }
288 (filter_type, layout) => {
289 return Err(HummockError::compaction_executor(format!(
290 "unresolved SST filter layout in task config: {:?}, {:?}",
291 filter_type, layout
292 )));
293 }
294 }
295 };
296
297 compact_timer.observe_duration();
298
299 Self::report_progress(
300 self.context.compactor_metrics.clone(),
301 task_progress,
302 &split_table_outputs,
303 self.context.is_share_buffer_compact,
304 );
305
306 self.context
307 .compactor_metrics
308 .get_table_id_total_time_duration
309 .observe(self.get_id_time.load(Ordering::Relaxed) as f64 / 1000.0 / 1000.0);
310
311 debug_assert!(
312 split_table_outputs
313 .iter()
314 .all(|table_info| table_info.sst_info.table_ids.is_sorted())
315 );
316
317 if task_id.is_some() {
318 tracing::info!(
320 "Finish Task {:?} split_index {:?} sst count {}",
321 task_id,
322 split_index,
323 split_table_outputs.len()
324 );
325 }
326 Ok((split_table_outputs, table_stats_map))
327 }
328
329 pub fn report_progress(
330 metrics: Arc<CompactorMetrics>,
331 task_progress: Option<Arc<TaskProgress>>,
332 ssts: &Vec<LocalSstableInfo>,
333 is_share_buffer_compact: bool,
334 ) {
335 for sst_info in ssts {
336 let sst_size = sst_info.file_size();
337 if let Some(tracker) = &task_progress {
338 tracker.inc_ssts_uploaded();
339 tracker.dec_num_pending_write_io();
340 }
341 if is_share_buffer_compact {
342 metrics.shared_buffer_to_sstable_size.observe(sst_size as _);
343 } else {
344 metrics.compaction_upload_sst_counts.inc();
345 }
346 }
347 }
348
349 async fn compact_key_range_impl<F: SstableWriterFactory, B: FilterBuilder>(
350 &self,
351 writer_factory: F,
352 iter: impl HummockIterator<Direction = Forward>,
353 compaction_filter: impl CompactionFilter,
354 compaction_catalog_agent_ref: CompactionCatalogAgentRef,
355 task_progress: Option<Arc<TaskProgress>>,
356 object_id_getter: Arc<dyn GetObjectId>,
357 ) -> HummockResult<(Vec<LocalSstableInfo>, CompactionStatistics)> {
358 let builder_factory = RemoteBuilderFactory::<F, B> {
359 object_id_getter,
360 limiter: self.context.memory_limiter.clone(),
361 options: self.options.clone(),
362 policy: self.task_config.cache_policy,
363 remote_rpc_cost: self.get_id_time.clone(),
364 compaction_catalog_agent_ref: compaction_catalog_agent_ref.clone(),
365 sstable_writer_factory: writer_factory,
366 _phantom: PhantomData,
367 };
368
369 let mut sst_builder = CapacitySplitTableBuilder::new(
370 builder_factory,
371 self.context.compactor_metrics.clone(),
372 task_progress.clone(),
373 self.task_config.table_vnode_partition.clone(),
374 self.context
375 .storage_opts
376 .compactor_concurrent_uploading_sst_count,
377 compaction_catalog_agent_ref,
378 );
379 let compaction_statistics = compact_and_build_sst(
380 &mut sst_builder,
381 &self.task_config,
382 self.context.compactor_metrics.clone(),
383 iter,
384 compaction_filter,
385 )
386 .instrument_await("compact_and_build_sst".verbose())
387 .await?;
388
389 let ssts = sst_builder
390 .finish()
391 .instrument_await("builder_finish".verbose())
392 .await?;
393
394 Ok((ssts, compaction_statistics))
395 }
396}
397
398#[must_use]
401pub fn start_iceberg_compactor(
402 compactor_context: CompactorContext,
403 hummock_meta_client: Arc<dyn HummockMetaClient>,
404 iceberg_compaction_memory_limit_bytes: usize,
405) -> (JoinHandle<()>, Sender<()>) {
406 let (shutdown_tx, mut shutdown_rx) = tokio::sync::oneshot::channel();
407 let stream_retry_interval = Duration::from_secs(30);
408 let periodic_event_update_interval = Duration::from_millis(
409 compactor_context
410 .storage_opts
411 .iceberg_compaction_pull_interval_ms,
412 );
413 let worker_num = compactor_context.compaction_executor.worker_num();
414
415 let max_task_parallelism: u32 = (worker_num as f32
416 * compactor_context.storage_opts.compactor_max_task_multiplier)
417 .ceil() as u32;
418
419 let max_pull_task_count = std::cmp::min(
420 max_task_parallelism,
421 compactor_context
422 .storage_opts
423 .iceberg_compaction_max_pull_task_count,
424 );
425
426 assert_ge!(
427 compactor_context.storage_opts.compactor_max_task_multiplier,
428 0.0
429 );
430
431 let join_handle = tokio::spawn(async move {
432 let pending_parallelism_budget = (max_task_parallelism as f32
434 * compactor_context
435 .storage_opts
436 .iceberg_compaction_pending_parallelism_budget_multiplier)
437 .ceil() as u32;
438 let mut task_queue = IcebergTaskQueue::new_with_metrics(
439 max_task_parallelism,
440 pending_parallelism_budget,
441 iceberg_compaction_memory_limit_bytes,
442 compactor_context.compactor_metrics.clone(),
443 );
444
445 let shutdown_map = Arc::new(Mutex::new(HashMap::<TaskKey, Sender<()>>::new()));
447
448 let (task_completion_tx, mut task_completion_rx) =
450 tokio::sync::mpsc::unbounded_channel::<IcebergPlanCompletion>();
451 let mut task_trackers = HashMap::<IcebergCompactionTaskId, IcebergTaskTracker>::new();
452 let mut pending_task_reports = VecDeque::<IcebergTaskReport>::new();
455
456 let mut min_interval = tokio::time::interval(stream_retry_interval);
457 let mut periodic_event_interval = tokio::time::interval(periodic_event_update_interval);
458
459 let mut log_throttler =
461 LogThrottler::<IcebergCompactionLogState>::new(COMPACTION_HEARTBEAT_LOG_INTERVAL);
462
463 'start_stream: loop {
465 let mut pull_task_ack = true;
468 tokio::select! {
469 _ = min_interval.tick() => {},
471 _ = &mut shutdown_rx => {
473 tracing::info!("Compactor is shutting down");
474 return;
475 }
476 }
477
478 let (request_sender, response_event_stream) = match hummock_meta_client
479 .subscribe_iceberg_compaction_event()
480 .await
481 {
482 Ok((request_sender, response_event_stream)) => {
483 tracing::debug!("Succeeded subscribe_iceberg_compaction_event.");
484 (request_sender, response_event_stream)
485 }
486
487 Err(e) => {
488 tracing::warn!(
489 error = %e.as_report(),
490 "Subscribing to iceberg compaction tasks failed with error. Will retry.",
491 );
492 continue 'start_stream;
493 }
494 };
495
496 if matches!(
497 flush_pending_iceberg_task_reports(&request_sender, &mut pending_task_reports),
498 ReportSendResult::RestartStream
499 ) {
500 continue 'start_stream;
501 }
502
503 pin_mut!(response_event_stream);
504
505 let _executor = compactor_context.compaction_executor.clone();
506
507 let mut event_loop_iteration_now = Instant::now();
509 'consume_stream: loop {
510 {
511 compactor_context
513 .compactor_metrics
514 .compaction_event_loop_iteration_latency
515 .observe(event_loop_iteration_now.elapsed().as_millis() as _);
516 event_loop_iteration_now = Instant::now();
517 }
518
519 let request_sender = request_sender.clone();
520 let event: Option<Result<SubscribeIcebergCompactionEventResponse, _>> = tokio::select! {
521 Some(plan_completion) = task_completion_rx.recv() => {
523 let task_key = plan_completion.task_key;
524 let error_message = plan_completion.error_message;
525 let pk_index_result = plan_completion.pk_index_result;
526 tracing::debug!(
527 task_id = %task_key.0,
528 plan_index = task_key.1,
529 success = error_message.is_none(),
530 "Plan completed, updating queue state"
531 );
532 task_queue.finish_running(task_key);
533
534 let completed_task_id = task_key.0;
535 let Entry::Occupied(mut tracker_entry) =
536 task_trackers.entry(completed_task_id)
537 else {
538 continue 'consume_stream;
539 };
540 tracker_entry
541 .get_mut()
542 .record_completion(error_message, pk_index_result);
543 if !tracker_entry.get().is_finished() {
544 continue 'consume_stream;
545 }
546
547 let tracker = tracker_entry.remove();
548 let sink_id = tracker.sink_id();
549 let admitted_plans = tracker.admitted_plans();
550 let successful_plans = tracker.successful_plans();
551 let failed_plans = tracker.failed_plans();
552 let report = tracker.into_report(completed_task_id);
553 let task_succeeded = report.error_message.is_none();
554 if matches!(
555 send_or_buffer_iceberg_task_report(
556 &request_sender,
557 &mut pending_task_reports,
558 report,
559 ),
560 ReportSendResult::RestartStream
561 ) {
562 continue 'start_stream;
563 }
564 if task_succeeded {
565 tracing::info!(
566 iceberg_component = "compaction_worker",
567 iceberg_operation = "report_task",
568 task_id = %completed_task_id,
569 sink_id = sink_id,
570 admitted_plans = admitted_plans,
571 successful_plans = successful_plans,
572 failed_plans = failed_plans,
573 "iceberg_compaction_task_succeeded",
574 );
575 }
576 continue 'consume_stream;
577 }
578
579 _ = task_queue.wait_schedulable() => {
581 schedule_queued_tasks(
582 &mut task_queue,
583 &compactor_context,
584 &shutdown_map,
585 &task_completion_tx,
586 );
587 continue 'consume_stream;
588 }
589
590 _ = periodic_event_interval.tick() => {
591 let should_restart_stream = handle_meta_task_pulling(
593 &mut pull_task_ack,
594 &task_queue,
595 max_task_parallelism,
596 max_pull_task_count,
597 &request_sender,
598 &mut log_throttler,
599 );
600
601 if should_restart_stream {
602 continue 'start_stream;
603 }
604 continue;
605 }
606 event = response_event_stream.next() => {
607 event
608 }
609
610 _ = &mut shutdown_rx => {
611 tracing::info!("Iceberg Compactor is shutting down");
612 return
613 }
614 };
615
616 match event {
617 Some(Ok(SubscribeIcebergCompactionEventResponse {
618 event,
619 create_at: _create_at,
620 })) => {
621 let event = match event {
622 Some(event) => event,
623 None => continue 'consume_stream,
624 };
625
626 match event {
627 risingwave_pb::iceberg_compaction::subscribe_iceberg_compaction_event_response::Event::CompactTask(iceberg_compaction_task) => {
628 let task_id = iceberg_compaction_task.task_id;
629 let sink_id = iceberg_compaction_task.sink_id;
630 let is_bounded_round = iceberg_compaction_task.max_file_sequence_number.is_some();
631 let pk_index_coordinated = iceberg_compaction_task.pk_index_coordinated;
632 #[cfg(madsim)]
637 if iceberg_compaction_task.pk_index_coordinated
638 && fail::eval(
639 "iceberg_pk_index_simulated_compaction_report",
640 |_| (),
641 )
642 .is_some()
643 && let Some(result) = simulated_pk_index_compaction_result()
644 {
645 let mut report = build_iceberg_task_report(task_id, sink_id, None);
646 report.pk_index_result = Some(result);
647 if matches!(
648 send_or_buffer_iceberg_task_report(
649 &request_sender,
650 &mut pending_task_reports,
651 report,
652 ),
653 ReportSendResult::RestartStream
654 ) {
655 continue 'start_stream;
656 }
657 continue 'consume_stream;
658 }
659 let compactor_runner_config = match IcebergCompactorRunnerConfigBuilder::default()
661 .max_parallelism((worker_num as f32 * compactor_context.storage_opts.iceberg_compaction_task_parallelism_ratio) as u32)
662 .min_size_per_partition(compactor_context.storage_opts.iceberg_compaction_min_size_per_partition_mb as u64 * 1024 * 1024)
663 .max_file_count_per_partition(compactor_context.storage_opts.iceberg_compaction_max_file_count_per_partition)
664 .enable_validate_compaction(compactor_context.storage_opts.iceberg_compaction_enable_validate)
665 .max_record_batch_rows(compactor_context.storage_opts.iceberg_compaction_max_record_batch_rows)
666 .enable_heuristic_output_parallelism(compactor_context.storage_opts.iceberg_compaction_enable_heuristic_output_parallelism)
667 .max_concurrent_closes(compactor_context.storage_opts.iceberg_compaction_max_concurrent_closes)
668 .enable_prefetch(compactor_context.storage_opts.iceberg_compaction_enable_prefetch)
669 .target_binpack_group_size_mb(
670 compactor_context.storage_opts.iceberg_compaction_target_binpack_group_size_mb
671 )
672 .min_group_size_mb(
673 compactor_context.storage_opts.iceberg_compaction_min_group_size_mb
674 )
675 .min_group_file_count(
676 compactor_context.storage_opts.iceberg_compaction_min_group_file_count
677 )
678 .build() {
679 Ok(config) => config,
680 Err(e) => {
681 tracing::warn!(
682 iceberg_component = "compaction_worker",
683 iceberg_operation = "build_runner_config",
684 error = %e.as_report(),
685 task_id = %task_id,
686 sink_id = sink_id,
687 "iceberg_compaction_runner_config_failed",
688 );
689 let report = build_iceberg_task_report(
690 task_id,
691 sink_id,
692 Some(format!(
693 "Failed to build iceberg compactor runner config: {}",
694 e.as_report()
695 )),
696 );
697 if matches!(
698 send_or_buffer_iceberg_task_report(
699 &request_sender,
700 &mut pending_task_reports,
701 report,
702 ),
703 ReportSendResult::RestartStream
704 ) {
705 continue 'start_stream;
706 }
707 continue 'consume_stream;
708 }
709 };
710
711 let task_execution = match create_task_execution(
713 iceberg_compaction_task,
714 compactor_runner_config,
715 compactor_context.compactor_metrics.clone(),
716 ).await {
717 Ok(task_execution) => task_execution,
718 Err(e) => {
719 tracing::warn!(
720 iceberg_component = "compaction_worker",
721 iceberg_operation = "plan_task",
722 error = %e.as_report(),
723 task_id = %task_id,
724 sink_id = sink_id,
725 "iceberg_compaction_task_plan_failed",
726 );
727 let report = build_iceberg_task_report(
728 task_id,
729 sink_id,
730 Some(format!(
731 "Failed to create iceberg compaction task execution: {}",
732 e.as_report()
733 )),
734 );
735 if matches!(
736 send_or_buffer_iceberg_task_report(
737 &request_sender,
738 &mut pending_task_reports,
739 report,
740 ),
741 ReportSendResult::RestartStream
742 ) {
743 continue 'start_stream;
744 }
745 continue 'consume_stream;
746 }
747 };
748
749 let (sink_id, plan_runners) = match task_execution {
750 IcebergTaskExecution::Ready { sink_id, plan_runners } => {
751 (sink_id, plan_runners)
752 }
753 IcebergTaskExecution::Blocked { sink_id, status, reason } => {
754 let mut report = build_iceberg_task_report(task_id, sink_id, Some(reason));
755 report.status = status as i32;
756 if matches!(
757 send_or_buffer_iceberg_task_report(&request_sender, &mut pending_task_reports, report),
758 ReportSendResult::RestartStream
759 ) {
760 continue 'start_stream;
761 }
762 continue 'consume_stream;
763 }
764 IcebergTaskExecution::NoWork { sink_id } => {
765 let report = if is_bounded_round {
767 build_drained_iceberg_task_report(task_id, sink_id)
768 } else {
769 build_iceberg_task_report(task_id, sink_id, None)
770 };
771 if matches!(
772 send_or_buffer_iceberg_task_report(
773 &request_sender,
774 &mut pending_task_reports,
775 report,
776 ),
777 ReportSendResult::RestartStream
778 ) {
779 continue 'start_stream;
780 }
781 continue 'consume_stream;
782 }
783 };
784
785 let planned_count = plan_runners.len();
787 let mut admitted_count = 0;
788 let mut rejection_reason = None;
789 let mut rejection_status = None;
790
791 for runner in plan_runners {
792 let meta = runner.to_meta();
793 let plan_index = meta.plan_index;
794 let required_parallelism = runner.required_parallelism();
795 let memory_reservation_bytes = meta.memory_reservation_bytes;
796 let runner_sink_id = runner.sink_id;
797 let runner_compaction_kind = runner.compaction_kind();
798 let runner_table = runner.table_ident.to_string();
799 let push_result = task_queue.push(meta, Some(runner));
800
801 match push_result {
802 PushResult::Added => {
803 admitted_count += 1;
804 tracing::debug!(
805 iceberg_component = "compaction_worker",
806 iceberg_operation = "enqueue_plan",
807 task_id = %task_id,
808 sink_id = runner_sink_id,
809 plan_index = plan_index,
810 task_type = ?runner_compaction_kind,
811 table = %runner_table,
812 required_parallelism = required_parallelism,
813 "iceberg_compaction_plan_enqueued",
814 );
815 },
816 PushResult::RejectedCapacity => {
817 rejection_status = Some(risingwave_pb::iceberg_compaction::subscribe_iceberg_compaction_event_request::report_task::Status::Busy);
818 tracing::warn!(
819 iceberg_component = "compaction_worker",
820 iceberg_operation = "enqueue_plan",
821 task_id = %task_id,
822 sink_id = runner_sink_id,
823 plan_index = plan_index,
824 task_type = ?runner_compaction_kind,
825 table = %runner_table,
826 required_parallelism = required_parallelism,
827 pending_budget = pending_parallelism_budget,
828 admitted_count = admitted_count,
829 planned_count = planned_count,
830 "iceberg_compaction_plan_rejected_capacity",
831 );
832 break;
834 },
835 PushResult::RejectedTooLarge => {
836 rejection_status = Some(risingwave_pb::iceberg_compaction::subscribe_iceberg_compaction_event_request::report_task::Status::BlockedCapacity);
837 rejection_reason.get_or_insert_with(|| {
838 format!(
839 "Iceberg compaction plan {plan_index} requires parallelism {required_parallelism} and an estimated {memory_reservation_bytes} bytes of memory, exceeding worker capacities of {max_task_parallelism} parallelism and {iceberg_compaction_memory_limit_bytes} bytes"
840 )
841 });
842 tracing::error!(
843 iceberg_component = "compaction_worker",
844 iceberg_operation = "enqueue_plan",
845 task_id = %task_id,
846 sink_id = runner_sink_id,
847 plan_index = plan_index,
848 task_type = ?runner_compaction_kind,
849 table = %runner_table,
850 required_parallelism = required_parallelism,
851 max_parallelism = max_task_parallelism,
852 memory_reservation_bytes,
853 memory_budget_bytes = iceberg_compaction_memory_limit_bytes,
854 "iceberg_compaction_plan_rejected_too_large",
855 );
856 },
857 PushResult::RejectedInvalidParallelism => {
858 tracing::error!(
859 iceberg_component = "compaction_worker",
860 iceberg_operation = "enqueue_plan",
861 task_id = %task_id,
862 sink_id = runner_sink_id,
863 plan_index = plan_index,
864 task_type = ?runner_compaction_kind,
865 table = %runner_table,
866 required_parallelism = required_parallelism,
867 "iceberg_compaction_plan_rejected_invalid_parallelism",
868 );
869 },
870 PushResult::RejectedDuplicate => {
871 tracing::error!(
872 iceberg_component = "compaction_worker",
873 iceberg_operation = "enqueue_plan",
874 task_id = %task_id,
875 sink_id = runner_sink_id,
876 plan_index = plan_index,
877 task_type = ?runner_compaction_kind,
878 table = %runner_table,
879 "iceberg_compaction_plan_rejected_duplicate",
880 );
881 }
882 }
883 }
884
885 if admitted_count == 0 {
886 let mut report = build_iceberg_task_report(
887 task_id,
888 sink_id,
889 Some(rejection_reason.unwrap_or_else(|| {
890 "Failed to enqueue all iceberg compaction plans".to_owned()
891 })),
892 );
893 if let Some(status) = rejection_status {
894 report.status = status as i32;
895 }
896 if matches!(
897 send_or_buffer_iceberg_task_report(
898 &request_sender,
899 &mut pending_task_reports,
900 report,
901 ),
902 ReportSendResult::RestartStream
903 ) {
904 continue 'start_stream;
905 }
906 } else {
907 let fully_admitted_bounded_round =
911 is_bounded_round
912 && admitted_count == planned_count;
913 task_trackers.insert(
914 task_id,
915 IcebergTaskTracker::new(
916 sink_id,
917 admitted_count,
918 pk_index_coordinated,
919 fully_admitted_bounded_round,
920 ),
921 );
922 }
923
924 tracing::info!(
925 iceberg_component = "compaction_worker",
926 iceberg_operation = "enqueue_plan",
927 task_id = %task_id,
928 sink_id = sink_id,
929 planned_count = planned_count,
930 admitted_count = admitted_count,
931 "iceberg_compaction_task_enqueue_finished",
932 );
933 },
934 risingwave_pb::iceberg_compaction::subscribe_iceberg_compaction_event_response::Event::PullTaskAck(_) => {
935 pull_task_ack = true;
937 },
938 risingwave_pb::iceberg_compaction::subscribe_iceberg_compaction_event_response::Event::CancelCompactTask(cancel_compact_task) => {
939 cancel_iceberg_task(
940 cancel_compact_task.task_id,
941 &mut task_queue,
942 &shutdown_map,
943 &mut task_trackers,
944 );
945 },
946 }
947 }
948 Some(Err(e)) => {
949 tracing::warn!("Failed to consume stream. {}", e.message());
950 continue 'start_stream;
951 }
952 _ => {
953 continue 'start_stream;
955 }
956 }
957 }
958 }
959 });
960
961 (join_handle, shutdown_tx)
962}
963
964#[must_use]
967pub fn start_compactor(
968 compactor_context: CompactorContext,
969 hummock_meta_client: Arc<dyn HummockMetaClient>,
970 object_id_manager: Arc<ObjectIdManager>,
971 compaction_catalog_manager_ref: CompactionCatalogManagerRef,
972) -> (JoinHandle<()>, Sender<()>) {
973 type CompactionShutdownMap = Arc<Mutex<HashMap<u64, Sender<()>>>>;
974 let (shutdown_tx, mut shutdown_rx) = tokio::sync::oneshot::channel();
975 let stream_retry_interval = Duration::from_secs(30);
976 let task_progress = compactor_context.task_progress_manager.clone();
977 let periodic_event_update_interval = Duration::from_millis(1000);
978
979 let max_task_parallelism: u32 = (compactor_context.compaction_executor.worker_num() as f32
980 * compactor_context.storage_opts.compactor_max_task_multiplier)
981 .ceil() as u32;
982 let running_task_parallelism = Arc::new(AtomicU32::new(0));
983
984 const MAX_PULL_TASK_COUNT: u32 = 4;
985 let max_pull_task_count = std::cmp::min(max_task_parallelism, MAX_PULL_TASK_COUNT);
986
987 assert_ge!(
988 compactor_context.storage_opts.compactor_max_task_multiplier,
989 0.0
990 );
991
992 let join_handle = tokio::spawn(async move {
993 let shutdown_map = CompactionShutdownMap::default();
994 let mut min_interval = tokio::time::interval(stream_retry_interval);
995 let mut periodic_event_interval = tokio::time::interval(periodic_event_update_interval);
996
997 let mut log_throttler =
999 LogThrottler::<CompactionLogState>::new(COMPACTION_HEARTBEAT_LOG_INTERVAL);
1000
1001 'start_stream: loop {
1003 let mut pull_task_ack = true;
1006 tokio::select! {
1007 _ = min_interval.tick() => {},
1009 _ = &mut shutdown_rx => {
1011 tracing::info!("Compactor is shutting down");
1012 return;
1013 }
1014 }
1015
1016 let (request_sender, response_event_stream) =
1017 match hummock_meta_client.subscribe_compaction_event().await {
1018 Ok((request_sender, response_event_stream)) => {
1019 tracing::debug!("Succeeded subscribe_compaction_event.");
1020 (request_sender, response_event_stream)
1021 }
1022
1023 Err(e) => {
1024 tracing::warn!(
1025 error = %e.as_report(),
1026 "Subscribing to compaction tasks failed with error. Will retry.",
1027 );
1028 continue 'start_stream;
1029 }
1030 };
1031
1032 pin_mut!(response_event_stream);
1033
1034 let executor = compactor_context.compaction_executor.clone();
1035 let object_id_manager = object_id_manager.clone();
1036
1037 let mut event_loop_iteration_now = Instant::now();
1039 'consume_stream: loop {
1040 {
1041 compactor_context
1043 .compactor_metrics
1044 .compaction_event_loop_iteration_latency
1045 .observe(event_loop_iteration_now.elapsed().as_millis() as _);
1046 event_loop_iteration_now = Instant::now();
1047 }
1048
1049 let running_task_parallelism = running_task_parallelism.clone();
1050 let request_sender = request_sender.clone();
1051 let event: Option<Result<SubscribeCompactionEventResponse, _>> = tokio::select! {
1052 _ = periodic_event_interval.tick() => {
1053 let progress_list = get_task_progress(task_progress.clone());
1054
1055 if let Err(e) = request_sender.send(SubscribeCompactionEventRequest {
1056 event: Some(RequestEvent::HeartBeat(
1057 HeartBeat {
1058 progress: progress_list
1059 }
1060 )),
1061 create_at: SystemTime::now()
1062 .duration_since(std::time::UNIX_EPOCH)
1063 .expect("Clock may have gone backwards")
1064 .as_millis() as u64,
1065 }) {
1066 tracing::warn!(error = %e.as_report(), "Failed to report task progress");
1067 continue 'start_stream;
1069 }
1070
1071
1072 let mut pending_pull_task_count = 0;
1073 if pull_task_ack {
1074 pending_pull_task_count = (max_task_parallelism - running_task_parallelism.load(Ordering::SeqCst)).min(max_pull_task_count);
1076
1077 if pending_pull_task_count > 0 {
1078 if let Err(e) = request_sender.send(SubscribeCompactionEventRequest {
1079 event: Some(RequestEvent::PullTask(
1080 PullTask {
1081 pull_task_count: pending_pull_task_count,
1082 }
1083 )),
1084 create_at: SystemTime::now()
1085 .duration_since(std::time::UNIX_EPOCH)
1086 .expect("Clock may have gone backwards")
1087 .as_millis() as u64,
1088 }) {
1089 tracing::warn!(error = %e.as_report(), "Failed to pull task");
1090
1091 continue 'start_stream;
1093 } else {
1094 pull_task_ack = false;
1095 }
1096 }
1097 }
1098
1099 let running_count = running_task_parallelism.load(Ordering::SeqCst);
1100 let current_state = CompactionLogState {
1101 running_parallelism: running_count,
1102 pull_task_ack,
1103 pending_pull_task_count,
1104 };
1105
1106 if log_throttler.should_log(¤t_state) {
1108 tracing::info!(
1109 running_parallelism_count = %current_state.running_parallelism,
1110 pull_task_ack = %current_state.pull_task_ack,
1111 pending_pull_task_count = %current_state.pending_pull_task_count
1112 );
1113 log_throttler.update(current_state);
1114 }
1115
1116 continue;
1117 }
1118 event = response_event_stream.next() => {
1119 event
1120 }
1121
1122 _ = &mut shutdown_rx => {
1123 tracing::info!("Compactor is shutting down");
1124 return
1125 }
1126 };
1127
1128 fn send_report_task_event(
1129 compact_task: &CompactTask,
1130 table_stats: TableStatsMap,
1131 object_timestamps: HashMap<HummockSstableObjectId, u64>,
1132 request_sender: &mpsc::UnboundedSender<SubscribeCompactionEventRequest>,
1133 ) {
1134 if let Err(e) = request_sender.send(SubscribeCompactionEventRequest {
1135 event: Some(RequestEvent::ReportTask(ReportTask {
1136 task_id: compact_task.task_id,
1137 task_status: compact_task.task_status.into(),
1138 sorted_output_ssts: compact_task
1139 .sorted_output_ssts
1140 .iter()
1141 .map(|sst| sst.into())
1142 .collect(),
1143 table_stats_change: to_prost_table_stats_map(table_stats),
1144 object_timestamps,
1145 })),
1146 create_at: SystemTime::now()
1147 .duration_since(std::time::UNIX_EPOCH)
1148 .expect("Clock may have gone backwards")
1149 .as_millis() as u64,
1150 }) {
1151 let task_id = compact_task.task_id;
1152 tracing::warn!(error = %e.as_report(), "Failed to report task {task_id:?}");
1153 }
1154 }
1155
1156 match event {
1157 Some(Ok(SubscribeCompactionEventResponse { event, create_at })) => {
1158 let event = match event {
1159 Some(event) => event,
1160 None => continue 'consume_stream,
1161 };
1162 let shutdown = shutdown_map.clone();
1163 let context = compactor_context.clone();
1164 let consumed_latency_ms = SystemTime::now()
1165 .duration_since(std::time::UNIX_EPOCH)
1166 .expect("Clock may have gone backwards")
1167 .as_millis() as u64
1168 - create_at;
1169 context
1170 .compactor_metrics
1171 .compaction_event_consumed_latency
1172 .observe(consumed_latency_ms as _);
1173
1174 let object_id_manager = object_id_manager.clone();
1175 let compaction_catalog_manager_ref = compaction_catalog_manager_ref.clone();
1176
1177 match event {
1178 ResponseEvent::CompactTask(compact_task) => {
1179 let compact_task = CompactTask::from(compact_task);
1180 let parallelism =
1181 calculate_task_parallelism(&compact_task, &context);
1182
1183 assert_ne!(parallelism, 0, "splits cannot be empty");
1184
1185 if (max_task_parallelism
1186 - running_task_parallelism.load(Ordering::SeqCst))
1187 < parallelism as u32
1188 {
1189 tracing::warn!(
1190 "Not enough core parallelism to serve the task {} task_parallelism {} running_task_parallelism {} max_task_parallelism {}",
1191 compact_task.task_id,
1192 parallelism,
1193 max_task_parallelism,
1194 running_task_parallelism.load(Ordering::Relaxed),
1195 );
1196 let (compact_task, table_stats, object_timestamps) =
1197 compact_done(
1198 compact_task,
1199 context.clone(),
1200 vec![],
1201 TaskStatus::NoAvailCpuResourceCanceled,
1202 );
1203
1204 send_report_task_event(
1205 &compact_task,
1206 table_stats,
1207 object_timestamps,
1208 &request_sender,
1209 );
1210
1211 continue 'consume_stream;
1212 }
1213
1214 running_task_parallelism
1215 .fetch_add(parallelism as u32, Ordering::SeqCst);
1216 executor.spawn(async move {
1217 let (tx, rx) = tokio::sync::oneshot::channel();
1218 let task_id = compact_task.task_id;
1219 shutdown.lock().unwrap().insert(task_id, tx);
1220
1221 let (
1222 (compact_task, table_stats, object_timestamps),
1223 _memory_tracker,
1224 ) = compactor_runner::compact(
1225 context.clone(),
1226 compact_task,
1227 rx,
1228 object_id_manager.clone(),
1229 compaction_catalog_manager_ref.clone(),
1230 )
1231 .await;
1232
1233 shutdown.lock().unwrap().remove(&task_id);
1234 running_task_parallelism
1235 .fetch_sub(parallelism as u32, Ordering::SeqCst);
1236
1237 send_report_task_event(
1238 &compact_task,
1239 table_stats,
1240 object_timestamps,
1241 &request_sender,
1242 );
1243
1244 let enable_check_compaction_result =
1245 context.storage_opts.check_compaction_result;
1246 let need_check_task =
1247 !compact_task.sorted_output_ssts.is_empty()
1248 && compact_task.task_status == TaskStatus::Success;
1249
1250 if enable_check_compaction_result && need_check_task {
1251 let read_table_ids = compact_task
1252 .get_table_ids_from_input_ssts()
1253 .collect::<Vec<_>>();
1254 match compaction_catalog_manager_ref.acquire(read_table_ids).await {
1255 Ok(compaction_catalog_agent_ref) => {
1256 match check_compaction_result(&compact_task, context.clone(), compaction_catalog_agent_ref).await
1257 {
1258 Err(e) => {
1259 tracing::warn!(error = %e.as_report(), "Failed to check compaction task {}",compact_task.task_id);
1260 }
1261 Ok(true) => (),
1262 Ok(false) => {
1263 panic!("Failed to pass consistency check for result of compaction task:\n{:?}", compact_task_to_string(&compact_task));
1264 }
1265 }
1266 },
1267 Err(e) => {
1268 tracing::warn!(error = %e.as_report(), "failed to acquire compaction catalog agent");
1269 }
1270 }
1271 }
1272 });
1273 }
1274 #[expect(deprecated)]
1275 ResponseEvent::VacuumTask(_) => {
1276 unreachable!("unexpected vacuum task");
1277 }
1278 #[expect(deprecated)]
1279 ResponseEvent::FullScanTask(_) => {
1280 unreachable!("unexpected scan task");
1281 }
1282 #[expect(deprecated)]
1283 ResponseEvent::ValidationTask(validation_task) => {
1284 let validation_task = ValidationTask::from(validation_task);
1285 executor.spawn(async move {
1286 validate_ssts(validation_task, context.sstable_store.clone())
1287 .await;
1288 });
1289 }
1290 ResponseEvent::CancelCompactTask(cancel_compact_task) => match shutdown
1291 .lock()
1292 .unwrap()
1293 .remove(&cancel_compact_task.task_id)
1294 {
1295 Some(tx) => {
1296 if tx.send(()).is_err() {
1297 tracing::warn!(
1298 "Cancellation of compaction task failed. task_id: {}",
1299 cancel_compact_task.task_id
1300 );
1301 }
1302 }
1303 _ => {
1304 tracing::warn!(
1305 "Attempting to cancel non-existent compaction task. task_id: {}",
1306 cancel_compact_task.task_id
1307 );
1308 }
1309 },
1310
1311 ResponseEvent::PullTaskAck(_pull_task_ack) => {
1312 pull_task_ack = true;
1314 }
1315 }
1316 }
1317 Some(Err(e)) => {
1318 tracing::warn!("Failed to consume stream. {}", e.message());
1319 continue 'start_stream;
1320 }
1321 _ => {
1322 continue 'start_stream;
1324 }
1325 }
1326 }
1327 }
1328 });
1329
1330 (join_handle, shutdown_tx)
1331}
1332
1333#[must_use]
1336pub fn start_shared_compactor(
1337 grpc_proxy_client: GrpcCompactorProxyClient,
1338 mut receiver: mpsc::UnboundedReceiver<Request<DispatchCompactionTaskRequest>>,
1339 context: CompactorContext,
1340) -> (JoinHandle<()>, Sender<()>) {
1341 type CompactionShutdownMap = Arc<Mutex<HashMap<u64, Sender<()>>>>;
1342 let task_progress = context.task_progress_manager.clone();
1343 let (shutdown_tx, mut shutdown_rx) = tokio::sync::oneshot::channel();
1344 let periodic_event_update_interval = Duration::from_millis(1000);
1345
1346 let join_handle = tokio::spawn(async move {
1347 let shutdown_map = CompactionShutdownMap::default();
1348
1349 let mut periodic_event_interval = tokio::time::interval(periodic_event_update_interval);
1350 let executor = context.compaction_executor.clone();
1351 let report_heartbeat_client = grpc_proxy_client.clone();
1352 'consume_stream: loop {
1353 let request: Option<Request<DispatchCompactionTaskRequest>> = tokio::select! {
1354 _ = periodic_event_interval.tick() => {
1355 let progress_list = get_task_progress(task_progress.clone());
1356 let report_compaction_task_request = ReportCompactionTaskRequest{
1357 event: Some(ReportCompactionTaskEvent::HeartBeat(
1358 SharedHeartBeat {
1359 progress: progress_list
1360 }
1361 )),
1362 };
1363 if let Err(e) = report_heartbeat_client.report_compaction_task(report_compaction_task_request).await{
1364 tracing::warn!(error = %e.as_report(), "Failed to report heartbeat");
1365 }
1366 continue
1367 }
1368
1369
1370 _ = &mut shutdown_rx => {
1371 tracing::info!("Compactor is shutting down");
1372 return
1373 }
1374
1375 request = receiver.recv() => {
1376 request
1377 }
1378
1379 };
1380 match request {
1381 Some(request) => {
1382 let context = context.clone();
1383 let shutdown = shutdown_map.clone();
1384
1385 let cloned_grpc_proxy_client = grpc_proxy_client.clone();
1386 executor.spawn(async move {
1387 let DispatchCompactionTaskRequest {
1388 tables,
1389 output_object_ids,
1390 task: dispatch_task,
1391 } = request.into_inner();
1392 let table_id_to_catalog = tables.into_iter().fold(HashMap::new(), |mut acc, table| {
1393 acc.insert(table.id, table);
1394 acc
1395 });
1396
1397 let mut output_object_ids_deque: VecDeque<_> = VecDeque::new();
1398 output_object_ids_deque.extend(output_object_ids.into_iter().map(Into::<HummockSstableObjectId>::into));
1399 let shared_compactor_object_id_manager =
1400 SharedComapctorObjectIdManager::new(output_object_ids_deque, cloned_grpc_proxy_client.clone(), context.storage_opts.sstable_id_remote_fetch_number);
1401 match dispatch_task.unwrap() {
1402 dispatch_compaction_task_request::Task::CompactTask(compact_task) => {
1403 let compact_task = CompactTask::from(&compact_task);
1404 let (tx, rx) = tokio::sync::oneshot::channel();
1405 let task_id = compact_task.task_id;
1406 shutdown.lock().unwrap().insert(task_id, tx);
1407
1408 let compaction_catalog_manager_ref =
1409 Arc::new(CompactionCatalogManager::new_preloaded(
1410 table_id_to_catalog,
1411 ));
1412 let ((compact_task, table_stats, object_timestamps), _memory_tracker)= compactor_runner::compact(
1413 context.clone(),
1414 compact_task,
1415 rx,
1416 shared_compactor_object_id_manager,
1417 compaction_catalog_manager_ref.clone(),
1418 )
1419 .await;
1420 shutdown.lock().unwrap().remove(&task_id);
1421 let report_compaction_task_request = ReportCompactionTaskRequest {
1422 event: Some(ReportCompactionTaskEvent::ReportTask(ReportSharedTask {
1423 compact_task: Some(PbCompactTask::from(&compact_task)),
1424 table_stats_change: to_prost_table_stats_map(table_stats),
1425 object_timestamps,
1426 })),
1427 };
1428
1429 match cloned_grpc_proxy_client
1430 .report_compaction_task(report_compaction_task_request)
1431 .await
1432 {
1433 Ok(_) => {
1434 let enable_check_compaction_result =
1436 context.storage_opts.check_compaction_result;
1437 let need_check_task =
1438 !compact_task.sorted_output_ssts.is_empty()
1439 && compact_task.task_status
1440 == TaskStatus::Success;
1441 if enable_check_compaction_result && need_check_task {
1442 let read_table_ids = compact_task
1443 .get_table_ids_from_input_ssts()
1444 .collect::<Vec<_>>();
1445 match compaction_catalog_manager_ref.acquire(read_table_ids).await {
1446 Ok(compaction_catalog_agent_ref) => {
1447 match check_compaction_result(&compact_task, context.clone(), compaction_catalog_agent_ref).await
1448 {
1449 Err(e) => {
1450 tracing::warn!(error = %e.as_report(), "Failed to check compaction task {}", task_id);
1451 }
1452 Ok(true) => (),
1453 Ok(false) => {
1454 panic!("Failed to pass consistency check for result of compaction task:\n{:?}", compact_task_to_string(&compact_task));
1455 }
1456 }
1457 }
1458 Err(e) => {
1459 tracing::warn!(error = %e.as_report(), "failed to acquire compaction catalog agent");
1460 }
1461 }
1462 }
1463 }
1464 Err(e) => tracing::warn!(error = %e.as_report(), "Failed to report task {task_id:?}"),
1465 }
1466
1467 }
1468 dispatch_compaction_task_request::Task::VacuumTask(_) => {
1469 unreachable!("unexpected vacuum task");
1470 }
1471 dispatch_compaction_task_request::Task::FullScanTask(_) => {
1472 unreachable!("unexpected scan task");
1473 }
1474 dispatch_compaction_task_request::Task::ValidationTask(validation_task) => {
1475 let validation_task = ValidationTask::from(validation_task);
1476 validate_ssts(validation_task, context.sstable_store.clone()).await;
1477 }
1478 dispatch_compaction_task_request::Task::CancelCompactTask(cancel_compact_task) => {
1479 match shutdown
1480 .lock()
1481 .unwrap()
1482 .remove(&cancel_compact_task.task_id)
1483 { Some(tx) => {
1484 if tx.send(()).is_err() {
1485 tracing::warn!(
1486 "Cancellation of compaction task failed. task_id: {}",
1487 cancel_compact_task.task_id
1488 );
1489 }
1490 } _ => {
1491 tracing::warn!(
1492 "Attempting to cancel non-existent compaction task. task_id: {}",
1493 cancel_compact_task.task_id
1494 );
1495 }}
1496 }
1497 }
1498 });
1499 }
1500 None => continue 'consume_stream,
1501 }
1502 }
1503 });
1504 (join_handle, shutdown_tx)
1505}
1506
1507fn get_task_progress(
1508 task_progress: Arc<
1509 parking_lot::lock_api::Mutex<parking_lot::RawMutex, HashMap<u64, Arc<TaskProgress>>>,
1510 >,
1511) -> Vec<CompactTaskProgress> {
1512 let mut progress_list = Vec::new();
1513 for (&task_id, progress) in &*task_progress.lock() {
1514 progress_list.push(progress.snapshot(task_id));
1515 }
1516 progress_list
1517}
1518
1519fn schedule_queued_tasks(
1521 task_queue: &mut IcebergTaskQueue,
1522 compactor_context: &CompactorContext,
1523 shutdown_map: &Arc<Mutex<HashMap<TaskKey, Sender<()>>>>,
1524 task_completion_tx: &tokio::sync::mpsc::UnboundedSender<IcebergPlanCompletion>,
1525) {
1526 while let Some(popped_task) = task_queue.pop() {
1527 let task_id = popped_task.meta.task_id;
1528 let plan_index = popped_task.meta.plan_index;
1529 let task_key = (task_id, plan_index);
1530
1531 let unique_ident = popped_task.runner.as_ref().map(|r| r.unique_ident());
1533
1534 let Some(runner) = popped_task.runner else {
1535 tracing::error!(
1536 iceberg_component = "compaction_worker",
1537 iceberg_operation = "schedule_plan",
1538 task_id = %task_id,
1539 plan_index = plan_index,
1540 "iceberg_compaction_plan_missing_runner",
1541 );
1542 task_queue.finish_running(task_key);
1543 continue;
1544 };
1545 let runner_sink_id = runner.sink_id;
1546 let runner_compaction_kind = runner.compaction_kind();
1547 let runner_table = runner.table_ident.to_string();
1548
1549 let executor = compactor_context.compaction_executor.clone();
1550 let shutdown_map_clone = shutdown_map.clone();
1551 let completion_tx_clone = task_completion_tx.clone();
1552 let (tx, rx) = tokio::sync::oneshot::channel();
1553
1554 {
1555 let mut shutdown_guard = shutdown_map.lock().unwrap();
1556 shutdown_guard.insert(task_key, tx);
1557 }
1558
1559 tracing::info!(
1560 iceberg_component = "compaction_worker",
1561 iceberg_operation = "schedule_plan",
1562 task_id = %task_id,
1563 sink_id = runner_sink_id,
1564 plan_index = plan_index,
1565 task_type = ?runner_compaction_kind,
1566 table = %runner_table,
1567 unique_ident = ?unique_ident,
1568 required_parallelism = popped_task.meta.required_parallelism,
1569 "iceberg_compaction_plan_started_from_queue",
1570 );
1571
1572 executor.spawn(async move {
1573 let _cleanup_guard = scopeguard::guard(shutdown_map_clone, move |shutdown_map| {
1574 let mut shutdown_guard = shutdown_map.lock().unwrap();
1575 shutdown_guard.remove(&task_key);
1576 });
1577
1578 let result = Box::pin(runner.compact(rx)).await;
1579
1580 let completion = match result {
1581 Ok((_stats, pk_index_result)) => IcebergPlanCompletion {
1582 task_key,
1583 error_message: None,
1584 pk_index_result,
1585 },
1586 Err(e) => {
1587 if is_cancelled_iceberg_compaction_error(&e) {
1588 tracing::info!(
1589 iceberg_component = "compaction_worker",
1590 iceberg_operation = "execute_plan",
1591 task_id = %task_key.0,
1592 sink_id = runner_sink_id,
1593 plan_index = task_key.1,
1594 task_type = ?runner_compaction_kind,
1595 table = %runner_table,
1596 "iceberg_compaction_plan_cancelled",
1597 );
1598 } else {
1599 tracing::warn!(
1600 iceberg_component = "compaction_worker",
1601 iceberg_operation = "execute_plan",
1602 error = %e.as_report(),
1603 task_id = %task_key.0,
1604 sink_id = runner_sink_id,
1605 plan_index = task_key.1,
1606 task_type = ?runner_compaction_kind,
1607 table = %runner_table,
1608 "iceberg_compaction_plan_failed",
1609 );
1610 }
1611 IcebergPlanCompletion {
1612 task_key,
1613 error_message: Some(e.to_report_string()),
1614 pk_index_result: None,
1615 }
1616 }
1617 };
1618
1619 if completion_tx_clone.send(completion).is_err() {
1620 tracing::warn!(
1621 iceberg_component = "compaction_worker",
1622 iceberg_operation = "notify_plan_completion",
1623 task_id = %task_key.0,
1624 sink_id = runner_sink_id,
1625 plan_index = task_key.1,
1626 task_type = ?runner_compaction_kind,
1627 table = %runner_table,
1628 "iceberg_compaction_plan_completion_send_failed",
1629 );
1630 }
1631 });
1632 }
1633}
1634
1635fn is_cancelled_iceberg_compaction_error(error: &crate::hummock::HummockError) -> bool {
1636 matches!(
1637 error.inner(),
1638 HummockErrorInner::CompactionExecutor(message) if message == "Plan cancelled"
1639 )
1640}
1641
1642fn cancel_iceberg_task(
1643 task_id: IcebergCompactionTaskId,
1644 task_queue: &mut IcebergTaskQueue,
1645 shutdown_map: &Arc<Mutex<HashMap<TaskKey, Sender<()>>>>,
1646 task_trackers: &mut HashMap<IcebergCompactionTaskId, IcebergTaskTracker>,
1647) {
1648 let cancelled_waiting = task_queue.cancel_waiting_task(task_id);
1653 let removed_tracker = task_trackers.remove(&task_id).is_some();
1654
1655 let cancelled_running = {
1656 let mut shutdown_guard = shutdown_map.lock().unwrap();
1657 let task_keys: Vec<_> = shutdown_guard
1658 .keys()
1659 .filter(|(running_task_id, _)| *running_task_id == task_id)
1660 .copied()
1661 .collect();
1662
1663 for task_key in &task_keys {
1664 if let Some(tx) = shutdown_guard.remove(task_key)
1665 && tx.send(()).is_err()
1666 {
1667 tracing::debug!(
1668 task_id = %task_key.0,
1669 plan_index = task_key.1,
1670 "Iceberg compaction plan shutdown receiver already closed during cancellation"
1671 );
1672 }
1673 }
1674
1675 task_keys.len()
1676 };
1677
1678 if cancelled_waiting == 0 && cancelled_running == 0 && !removed_tracker {
1679 tracing::warn!(
1680 task_id = %task_id,
1681 "Attempting to cancel non-existent iceberg compaction task"
1682 );
1683 } else {
1684 tracing::info!(
1685 task_id = %task_id,
1686 cancelled_waiting = cancelled_waiting,
1687 cancelled_running = cancelled_running,
1688 removed_tracker = removed_tracker,
1689 "Cancelled iceberg compaction task"
1690 );
1691 }
1692}
1693
1694fn handle_meta_task_pulling(
1697 pull_task_ack: &mut bool,
1698 task_queue: &IcebergTaskQueue,
1699 max_task_parallelism: u32,
1700 max_pull_task_count: u32,
1701 request_sender: &mpsc::UnboundedSender<SubscribeIcebergCompactionEventRequest>,
1702 log_throttler: &mut LogThrottler<IcebergCompactionLogState>,
1703) -> bool {
1704 let mut pending_pull_task_count = 0;
1705 if *pull_task_ack {
1706 let current_reserved_parallelism = task_queue
1711 .running_parallelism_sum()
1712 .saturating_add(task_queue.waiting_parallelism_sum());
1713 pending_pull_task_count = max_task_parallelism
1714 .saturating_sub(current_reserved_parallelism)
1715 .min(max_pull_task_count);
1716
1717 if pending_pull_task_count > 0 {
1718 if let Err(e) = request_sender.send(SubscribeIcebergCompactionEventRequest {
1719 event: Some(subscribe_iceberg_compaction_event_request::Event::PullTask(
1720 subscribe_iceberg_compaction_event_request::PullTask {
1721 pull_task_count: pending_pull_task_count,
1722 },
1723 )),
1724 create_at: SystemTime::now()
1725 .duration_since(std::time::UNIX_EPOCH)
1726 .expect("Clock may have gone backwards")
1727 .as_millis() as u64,
1728 }) {
1729 tracing::warn!(error = %e.as_report(), "Failed to pull task - will retry on stream restart");
1730 return true; } else {
1732 *pull_task_ack = false;
1733 }
1734 }
1735 }
1736
1737 let running_count = task_queue.running_parallelism_sum();
1738 let waiting_count = task_queue.waiting_parallelism_sum();
1739 let available_count = max_task_parallelism.saturating_sub(running_count);
1740 let current_state = IcebergCompactionLogState {
1741 running_parallelism: running_count,
1742 waiting_parallelism: waiting_count,
1743 available_parallelism: available_count,
1744 pull_task_ack: *pull_task_ack,
1745 pending_pull_task_count,
1746 };
1747
1748 if log_throttler.should_log(¤t_state) {
1750 tracing::info!(
1751 running_parallelism_count = %current_state.running_parallelism,
1752 waiting_parallelism_count = %current_state.waiting_parallelism,
1753 available_parallelism = %current_state.available_parallelism,
1754 pull_task_ack = %current_state.pull_task_ack,
1755 pending_pull_task_count = %current_state.pending_pull_task_count
1756 );
1757 log_throttler.update(current_state);
1758 }
1759
1760 false }
1762
1763#[cfg(test)]
1764mod tests {
1765 use super::*;
1766
1767 #[test]
1768 fn test_cancel_iceberg_task_removes_waiting_plans_and_tracker() {
1769 let task_id = IcebergCompactionTaskId::new(42);
1770 let mut task_queue = IcebergTaskQueue::new(10, 30, 1024);
1771 assert_eq!(
1772 task_queue.push(
1773 iceberg_compaction::IcebergTaskMeta {
1774 task_id,
1775 plan_index: 0,
1776 required_parallelism: 3,
1777 memory_reservation_bytes: 64,
1778 },
1779 None,
1780 ),
1781 PushResult::Added
1782 );
1783 assert_eq!(
1784 task_queue.push(
1785 iceberg_compaction::IcebergTaskMeta {
1786 task_id,
1787 plan_index: 1,
1788 required_parallelism: 4,
1789 memory_reservation_bytes: 64,
1790 },
1791 None,
1792 ),
1793 PushResult::Added
1794 );
1795 assert_eq!(task_queue.waiting_parallelism_sum(), 7);
1796
1797 let shutdown_map = Arc::new(Mutex::new(HashMap::new()));
1798 let mut task_trackers =
1799 HashMap::from([(task_id, IcebergTaskTracker::new(10, 2, false, false))]);
1800
1801 cancel_iceberg_task(task_id, &mut task_queue, &shutdown_map, &mut task_trackers);
1802
1803 assert_eq!(task_queue.waiting_parallelism_sum(), 0);
1804 assert!(!task_trackers.contains_key(&task_id));
1805 }
1806
1807 #[test]
1808 fn test_cancel_iceberg_task_shuts_down_running_plans_and_tracker() {
1809 let task_id = IcebergCompactionTaskId::new(43);
1810 let task_key = (task_id, 0);
1811 let mut task_queue = IcebergTaskQueue::new(10, 30, 1024);
1812 let shutdown_map = Arc::new(Mutex::new(HashMap::new()));
1813 let (shutdown_tx, mut shutdown_rx) = tokio::sync::oneshot::channel();
1814 shutdown_map.lock().unwrap().insert(task_key, shutdown_tx);
1815 let mut task_trackers =
1816 HashMap::from([(task_id, IcebergTaskTracker::new(10, 1, false, false))]);
1817
1818 cancel_iceberg_task(task_id, &mut task_queue, &shutdown_map, &mut task_trackers);
1819
1820 assert!(shutdown_rx.try_recv().is_ok());
1821 assert!(shutdown_map.lock().unwrap().is_empty());
1822 assert!(!task_trackers.contains_key(&task_id));
1823 }
1824
1825 #[test]
1826 fn test_log_state_equality() {
1827 let state1 = CompactionLogState {
1829 running_parallelism: 10,
1830 pull_task_ack: true,
1831 pending_pull_task_count: 2,
1832 };
1833 let state2 = CompactionLogState {
1834 running_parallelism: 10,
1835 pull_task_ack: true,
1836 pending_pull_task_count: 2,
1837 };
1838 let state3 = CompactionLogState {
1839 running_parallelism: 11,
1840 pull_task_ack: true,
1841 pending_pull_task_count: 2,
1842 };
1843 assert_eq!(state1, state2);
1844 assert_ne!(state1, state3);
1845
1846 let ice_state1 = IcebergCompactionLogState {
1848 running_parallelism: 10,
1849 waiting_parallelism: 5,
1850 available_parallelism: 15,
1851 pull_task_ack: true,
1852 pending_pull_task_count: 2,
1853 };
1854 let ice_state2 = IcebergCompactionLogState {
1855 running_parallelism: 10,
1856 waiting_parallelism: 6,
1857 available_parallelism: 15,
1858 pull_task_ack: true,
1859 pending_pull_task_count: 2,
1860 };
1861 assert_ne!(ice_state1, ice_state2);
1862 }
1863
1864 #[test]
1865 fn test_log_throttler_state_change_detection() {
1866 let mut throttler = LogThrottler::<CompactionLogState>::new(Duration::from_secs(60));
1867 let state1 = CompactionLogState {
1868 running_parallelism: 10,
1869 pull_task_ack: true,
1870 pending_pull_task_count: 2,
1871 };
1872 let state2 = CompactionLogState {
1873 running_parallelism: 11,
1874 pull_task_ack: true,
1875 pending_pull_task_count: 2,
1876 };
1877
1878 assert!(throttler.should_log(&state1));
1880 throttler.update(state1.clone());
1881
1882 assert!(!throttler.should_log(&state1));
1884
1885 assert!(throttler.should_log(&state2));
1887 throttler.update(state2.clone());
1888
1889 assert!(!throttler.should_log(&state2));
1891 }
1892
1893 #[test]
1894 fn test_log_throttler_heartbeat() {
1895 let mut throttler = LogThrottler::<CompactionLogState>::new(Duration::from_millis(10));
1896 let state = CompactionLogState {
1897 running_parallelism: 10,
1898 pull_task_ack: true,
1899 pending_pull_task_count: 2,
1900 };
1901
1902 assert!(throttler.should_log(&state));
1904 throttler.update(state.clone());
1905
1906 assert!(!throttler.should_log(&state));
1908
1909 std::thread::sleep(Duration::from_millis(15));
1911
1912 assert!(throttler.should_log(&state));
1914 }
1915}