1use std::collections::hash_map::Entry;
16use std::sync::Arc;
17use std::time::{Duration, Instant, SystemTime};
18
19use futures::StreamExt;
20use itertools::Itertools;
21use parking_lot::RwLock;
22use risingwave_common::license::Feature;
23use risingwave_connector::connector_common::{
24 IcebergCommittedSnapshot, IcebergSinkCompactionUpdate,
25};
26use risingwave_connector::sink::SinkParam;
27use risingwave_connector::sink::catalog::{SinkCatalog, SinkId};
28use risingwave_connector::sink::iceberg::{
29 CompactionType, IcebergConfig, IcebergWriteMode, commit_branch,
30};
31use risingwave_hummock_sdk::HummockContextId;
32use risingwave_pb::catalog::PbStreamJobStatus;
33use risingwave_pb::iceberg_compaction::IcebergCompactionTask;
34use risingwave_pb::iceberg_compaction::iceberg_compaction_task::TaskType;
35use risingwave_pb::iceberg_compaction::subscribe_iceberg_compaction_event_request::ReportTask as IcebergReportTask;
36use risingwave_pb::iceberg_compaction::subscribe_iceberg_compaction_event_request::report_task::Status as IcebergReportTaskStatus;
37use risingwave_pb::id::IcebergCompactionTaskId;
38use thiserror_ext::AsReport;
39use tokio::sync::oneshot;
40use tokio::task::JoinHandle;
41
42use super::*;
43
44const COMPACTION_RETRY_BACKOFF: Duration = Duration::from_secs(1);
45
46#[derive(Debug, Clone)]
48enum CompactionTrackState {
49 Idle {
54 next_compaction_time: Instant,
55 manual_task_type: Option<TaskType>,
57 },
58 PendingDispatch { attempt: Arc<CompactionAttempt> },
60 InFlight {
63 task_id: IcebergCompactionTaskId,
64 compactor_context_id: HummockContextId,
65 attempt: Arc<CompactionAttempt>,
66 report_deadline: Instant,
67 },
68}
69
70#[derive(Debug)]
75struct CompactionAttempt {
76 task_type: TaskType,
77 max_file_sequence_number: Option<i64>,
78 pending_commit_count_at_start: usize,
79 gc_watermark_snapshot: Option<IcebergCommittedSnapshot>,
80}
81
82#[derive(Debug, Clone, Copy)]
83struct ScheduledCompactionTask {
84 task_id: IcebergCompactionTaskId,
85 compactor_context_id: HummockContextId,
86}
87
88#[derive(Debug, Clone, Default)]
93struct CompactionHealth {
94 last_success_at: Option<SystemTime>,
95 last_failure_at: Option<SystemTime>,
96 last_error: Option<String>,
97 consecutive_failures: u32,
98 compacted_snapshot: Option<IcebergCommittedSnapshot>,
101 round_start_snapshot: Option<IcebergCommittedSnapshot>,
104}
105
106impl CompactionHealth {
107 fn record_progress(&mut self) {
108 self.last_success_at = Some(SystemTime::now());
109 self.consecutive_failures = 0;
110 }
111
112 fn record_completion(&mut self, compacted_snapshot: Option<IcebergCommittedSnapshot>) {
113 self.record_progress();
114 if compacted_snapshot.is_some() {
115 self.compacted_snapshot = compacted_snapshot;
116 }
117 }
118
119 fn record_failure(&mut self, error: String) {
120 self.last_failure_at = Some(SystemTime::now());
121 self.last_error = Some(error);
122 self.consecutive_failures = self.consecutive_failures.saturating_add(1);
123 }
124}
125
126#[derive(Debug, Clone, Copy, PartialEq, Eq)]
127enum CompactionTrackFinishAction {
128 KeepTrack,
129 RemoveTrack,
130}
131
132#[derive(Debug, Clone)]
133pub(super) struct CompactionTrack {
134 configured_task_type: TaskType,
136 write_mode: IcebergWriteMode,
137 trigger_interval_sec: u64,
138 trigger_snapshot_count: usize,
141 report_timeout: Duration,
142 last_config_refresh_at: Instant,
143 pending_commit_count: usize,
144 latest_observed_snapshot: Option<IcebergCommittedSnapshot>,
145 round_max_file_sequence_number: Option<i64>,
147 finish_action: CompactionTrackFinishAction,
151 state: CompactionTrackState,
152 health: CompactionHealth,
153}
154
155impl CompactionTrack {
156 fn new(
157 configured_task_type: TaskType,
158 write_mode: IcebergWriteMode,
159 trigger_interval_sec: u64,
160 trigger_snapshot_count: usize,
161 report_timeout: Duration,
162 now: Instant,
163 ) -> Self {
164 Self {
165 configured_task_type,
166 write_mode,
167 trigger_interval_sec,
168 trigger_snapshot_count,
169 report_timeout,
170 last_config_refresh_at: now,
171 pending_commit_count: 0,
172 latest_observed_snapshot: None,
173 round_max_file_sequence_number: None,
174 finish_action: CompactionTrackFinishAction::KeepTrack,
175 state: CompactionTrackState::Idle {
176 next_compaction_time: now + Duration::from_secs(trigger_interval_sec),
177 manual_task_type: None,
178 },
179 health: CompactionHealth::default(),
180 }
181 }
182
183 fn should_trigger(&self, now: Instant) -> bool {
196 let next_compaction_time = match &self.state {
197 CompactionTrackState::Idle {
198 next_compaction_time,
199 ..
200 } => *next_compaction_time,
201 CompactionTrackState::PendingDispatch { .. }
202 | CompactionTrackState::InFlight { .. } => return false,
203 };
204
205 let time_ready = now >= next_compaction_time;
206 if self.round_max_file_sequence_number.is_some() {
207 return time_ready;
212 }
213
214 let commit_ready = self.pending_commit_count >= self.trigger_snapshot_count;
215 let has_commits = self.pending_commit_count > 0;
216 commit_ready || (time_ready && has_commits)
217 }
218
219 fn record_observed_snapshot(&mut self, observed_snapshot: IcebergCommittedSnapshot) {
220 self.latest_observed_snapshot = Some(observed_snapshot);
221 }
222
223 fn record_commit(&mut self) {
224 self.pending_commit_count = self.pending_commit_count.saturating_add(1);
225 }
226
227 fn record_force_compaction(&mut self, now: Instant, forced_task_type: Option<TaskType>) {
228 if let CompactionTrackState::Idle {
229 next_compaction_time,
230 manual_task_type,
231 } = &mut self.state
232 {
233 if let Some(task_type) = forced_task_type {
234 *manual_task_type = Some(task_type);
235 }
236 *next_compaction_time = now;
237 if self.round_max_file_sequence_number.is_none() {
240 self.pending_commit_count = self.pending_commit_count.max(1);
241 }
242 }
243 }
244
245 fn needs_config_refresh(&self, now: Instant, refresh_interval: Duration) -> bool {
246 now.saturating_duration_since(self.last_config_refresh_at) >= refresh_interval
247 }
248
249 fn mark_config_refreshed(&mut self, now: Instant) {
250 self.last_config_refresh_at = now;
251 }
252
253 fn effective_task_type(&self) -> TaskType {
254 match &self.state {
255 CompactionTrackState::Idle {
256 manual_task_type, ..
257 } => manual_task_type.unwrap_or(self.configured_task_type),
258 CompactionTrackState::PendingDispatch { attempt }
259 | CompactionTrackState::InFlight { attempt, .. } => attempt.task_type,
260 }
261 }
262
263 fn start_attempt(&mut self) -> Arc<CompactionAttempt> {
264 let CompactionTrackState::Idle {
265 manual_task_type, ..
266 } = &mut self.state
267 else {
268 unreachable!("Cannot start an attempt while another attempt is active")
269 };
270
271 let manual_task_type = manual_task_type.take();
272 let task_type = manual_task_type.unwrap_or(self.configured_task_type);
273 let new_round_boundary = if manual_task_type.is_none()
274 && self.write_mode == IcebergWriteMode::MergeOnRead
275 && self.round_max_file_sequence_number.is_none()
276 {
277 self.latest_observed_snapshot
278 .as_ref()
279 .expect("automatic merge-on-read compaction must start from an observed snapshot")
280 .max_file_sequence_number
281 } else {
282 None
283 };
284
285 if let Some(max_file_sequence_number) = new_round_boundary {
290 self.round_max_file_sequence_number = Some(max_file_sequence_number);
291 self.pending_commit_count = 0;
292 self.health.round_start_snapshot = self.latest_observed_snapshot.clone();
293 }
294 let attempt = Arc::new(CompactionAttempt {
295 task_type,
296 max_file_sequence_number: self.round_max_file_sequence_number,
297 pending_commit_count_at_start: self.pending_commit_count,
298 gc_watermark_snapshot: self.latest_observed_snapshot.clone(),
299 });
300 self.state = CompactionTrackState::PendingDispatch {
301 attempt: attempt.clone(),
302 };
303 attempt
304 }
305
306 fn mark_dispatched(
307 &mut self,
308 task_id: IcebergCompactionTaskId,
309 compactor_context_id: HummockContextId,
310 now: Instant,
311 ) {
312 let CompactionTrackState::PendingDispatch { attempt } = &self.state else {
313 unreachable!("Only a pending attempt can be marked as dispatched")
314 };
315 self.state = CompactionTrackState::InFlight {
316 task_id,
317 compactor_context_id,
318 attempt: attempt.clone(),
319 report_deadline: now + self.report_timeout,
320 };
321 }
322
323 pub(super) fn active_attempt_gc_watermark_snapshot(
324 &self,
325 ) -> Option<Option<&IcebergCommittedSnapshot>> {
326 match &self.state {
327 CompactionTrackState::PendingDispatch { attempt }
328 | CompactionTrackState::InFlight { attempt, .. } => {
329 Some(attempt.gc_watermark_snapshot.as_ref())
330 }
331 CompactionTrackState::Idle { .. } => None,
332 }
333 }
334
335 fn is_pending_dispatch(&self) -> bool {
336 matches!(self.state, CompactionTrackState::PendingDispatch { .. })
337 }
338
339 fn removes_track_after_finish(&self) -> bool {
340 self.finish_action == CompactionTrackFinishAction::RemoveTrack
341 }
342
343 fn has_queued_task_or_active_attempt(&self) -> bool {
344 matches!(
345 &self.state,
346 CompactionTrackState::Idle {
347 manual_task_type: Some(_),
348 ..
349 } | CompactionTrackState::PendingDispatch { .. }
350 | CompactionTrackState::InFlight { .. }
351 )
352 }
353
354 pub(super) fn is_in_flight_task(&self, task_id: IcebergCompactionTaskId) -> bool {
355 matches!(
356 &self.state,
357 CompactionTrackState::InFlight {
358 task_id: current_task_id,
359 ..
360 } if *current_task_id == task_id
361 )
362 }
363
364 fn scheduled_task(&self) -> Option<ScheduledCompactionTask> {
365 match &self.state {
366 CompactionTrackState::InFlight {
367 task_id,
368 compactor_context_id,
369 ..
370 } => Some(ScheduledCompactionTask {
371 task_id: *task_id,
372 compactor_context_id: *compactor_context_id,
373 }),
374 CompactionTrackState::Idle { .. } | CompactionTrackState::PendingDispatch { .. } => {
375 None
376 }
377 }
378 }
379
380 fn is_report_timed_out(&self, now: Instant) -> bool {
381 matches!(
382 &self.state,
383 CompactionTrackState::InFlight {
384 report_deadline,
385 ..
386 } if now >= *report_deadline
387 )
388 }
389
390 fn finish_failed(&mut self, now: Instant, error: String) -> CompactionTrackFinishAction {
391 if !matches!(self.state, CompactionTrackState::InFlight { .. }) {
392 unreachable!("Only an in-flight attempt can finish")
393 }
394 self.health.record_failure(error);
395 self.state = CompactionTrackState::Idle {
396 next_compaction_time: now + COMPACTION_RETRY_BACKOFF,
397 manual_task_type: None,
398 };
399 self.finish_action
400 }
401
402 fn revert_pre_dispatch_failure(&mut self, now: Instant) -> CompactionTrackFinishAction {
413 if !self.is_pending_dispatch() {
414 unreachable!("Only a pending attempt can be reverted")
415 }
416 self.health
417 .record_failure("compaction task failed before dispatch to a compactor".to_owned());
418 self.state = CompactionTrackState::Idle {
419 next_compaction_time: now + COMPACTION_RETRY_BACKOFF,
420 manual_task_type: None,
421 };
422 self.finish_action
423 }
424
425 fn update_interval(&mut self, new_interval_sec: u64, now: Instant) {
426 if self.trigger_interval_sec == new_interval_sec {
427 return;
428 }
429
430 self.trigger_interval_sec = new_interval_sec;
431 if self.round_max_file_sequence_number.is_some() {
432 return;
435 }
436
437 match &mut self.state {
438 CompactionTrackState::Idle {
439 next_compaction_time,
440 ..
441 } => {
442 *next_compaction_time = now + Duration::from_secs(new_interval_sec);
443 }
444 CompactionTrackState::PendingDispatch { .. }
445 | CompactionTrackState::InFlight { .. } => {}
446 }
447 }
448
449 fn finish_success(&mut self, now: Instant) -> CompactionTrackFinishAction {
450 let CompactionTrackState::InFlight { attempt, .. } = &self.state else {
451 unreachable!("Only an in-flight attempt can finish")
452 };
453 if attempt.max_file_sequence_number.is_some() {
454 self.health.record_progress();
457 self.state = CompactionTrackState::Idle {
458 next_compaction_time: now,
459 manual_task_type: None,
460 };
461 } else {
462 self.pending_commit_count = self
463 .pending_commit_count
464 .saturating_sub(attempt.pending_commit_count_at_start);
465 self.health
466 .record_completion(attempt.gc_watermark_snapshot.clone());
467 self.state = CompactionTrackState::Idle {
468 next_compaction_time: now + Duration::from_secs(self.trigger_interval_sec),
469 manual_task_type: None,
470 };
471 }
472 self.finish_action
473 }
474
475 fn finish_drained(&mut self, now: Instant) -> CompactionTrackFinishAction {
478 let CompactionTrackState::InFlight { attempt, .. } = &self.state else {
479 unreachable!("Only an in-flight attempt can finish")
480 };
481 debug_assert_eq!(
482 attempt.max_file_sequence_number,
483 self.round_max_file_sequence_number
484 );
485 debug_assert!(attempt.max_file_sequence_number.is_some());
486 self.round_max_file_sequence_number = None;
487 let round_start_snapshot = self.health.round_start_snapshot.take();
488 self.health.record_completion(round_start_snapshot);
489 self.state = CompactionTrackState::Idle {
490 next_compaction_time: now + Duration::from_secs(self.trigger_interval_sec),
491 manual_task_type: None,
492 };
493 self.finish_action
494 }
495
496 fn compaction_lag(&self) -> Option<Duration> {
497 let latest = self.latest_observed_snapshot.as_ref()?;
498 let compacted = self.health.compacted_snapshot.as_ref()?;
499 let lag_ms = latest.timestamp_ms.saturating_sub(compacted.timestamp_ms);
500 Some(Duration::from_millis(lag_ms.max(0) as u64))
501 }
502
503 fn is_in_flight_bounded_attempt(&self) -> bool {
504 matches!(
505 &self.state,
506 CompactionTrackState::InFlight { attempt, .. }
507 if attempt.max_file_sequence_number.is_some()
508 )
509 }
510}
511
512pub(crate) struct IcebergCompactionHandle {
513 sink_id: SinkId,
514 attempt: Arc<CompactionAttempt>,
515 inner: Arc<RwLock<IcebergCompactionManagerInner>>,
516 metadata_manager: MetadataManager,
517 dispatched: bool,
518}
519
520impl IcebergCompactionHandle {
521 fn new(
522 sink_id: SinkId,
523 attempt: Arc<CompactionAttempt>,
524 inner: Arc<RwLock<IcebergCompactionManagerInner>>,
525 metadata_manager: MetadataManager,
526 ) -> Self {
527 Self {
528 sink_id,
529 attempt,
530 inner,
531 metadata_manager,
532 dispatched: false,
533 }
534 }
535
536 pub async fn send_compact_task(
537 mut self,
538 compactor: Arc<crate::hummock::IcebergCompactor>,
539 task_id: IcebergCompactionTaskId,
540 ) -> MetaResult<()> {
541 let Some(prost_sink_catalog) = self
542 .metadata_manager
543 .catalog_controller
544 .get_sink_by_id(self.sink_id)
545 .await?
546 else {
547 tracing::warn!(
548 iceberg_component = "compaction_scheduler",
549 iceberg_operation = "dispatch_task",
550 sink_id = %self.sink_id,
551 task_id = %task_id,
552 "iceberg_compaction_dispatch_sink_not_found",
553 );
554 return Ok(());
555 };
556 let sink_catalog = SinkCatalog::from(prost_sink_catalog);
557 let param = SinkParam::try_from_sink_catalog(sink_catalog)?;
558
559 self.try_dispatch_task(
560 &compactor,
561 IcebergCompactionTask {
562 task_id,
563 sink_id: self.sink_id.as_raw_id(),
564 props: param.properties,
565 task_type: self.attempt.task_type as i32,
566 pk_index_coordinated: false,
567 max_file_sequence_number: self.attempt.max_file_sequence_number,
568 },
569 )
570 }
571
572 fn try_dispatch_task(
573 &mut self,
574 compactor: &crate::hummock::IcebergCompactor,
575 task: IcebergCompactionTask,
576 ) -> MetaResult<()> {
577 use risingwave_pb::iceberg_compaction::subscribe_iceberg_compaction_event_response::Event as IcebergResponseEvent;
578
579 let task_id = task.task_id;
580 let mut guard = self.inner.write();
583 let Some(track) = guard
584 .sink_schedules
585 .get_mut(&self.sink_id)
586 .filter(|track| track.is_pending_dispatch())
587 else {
588 tracing::warn!(
589 iceberg_component = "compaction_scheduler",
590 iceberg_operation = "dispatch_task",
591 sink_id = %self.sink_id,
592 task_id = %task_id,
593 "iceberg_compaction_dispatch_track_not_pending",
594 );
595 return Ok(());
596 };
597
598 let result = compactor.send_event(IcebergResponseEvent::CompactTask(task));
599 if result.is_ok() {
600 track.mark_dispatched(task_id, compactor.context_id(), Instant::now());
601 self.dispatched = true;
602 }
603 result
604 }
605}
606
607impl Drop for IcebergCompactionHandle {
608 fn drop(&mut self) {
609 let waiter = {
610 let mut guard = self.inner.write();
611 let finish_action = if !self.dispatched
612 && let Some(track) = guard.sink_schedules.get_mut(&self.sink_id)
613 && track.is_pending_dispatch()
614 {
615 Some(track.revert_pre_dispatch_failure(Instant::now()))
616 } else {
617 None
618 };
619 if let Some(finish_action) = finish_action {
620 let waiter = guard.manual_compaction_waiters.remove(&self.sink_id);
621 IcebergCompactionManager::apply_track_finish_action(
622 &mut guard,
623 self.sink_id,
624 finish_action,
625 );
626 waiter
627 } else {
628 None
629 }
630 };
631
632 if let Some(waiter) = waiter {
633 let _ = waiter.send(Err(anyhow!(
634 "Iceberg compaction task failed before dispatch for sink {}",
635 self.sink_id
636 )
637 .into()));
638 }
639 }
640}
641
642#[derive(Debug, Clone)]
643enum SinkUpdateKind {
644 Commit {
646 observed_snapshot: IcebergCommittedSnapshot,
647 },
648 ForceCompaction {
651 observed_snapshot: IcebergCommittedSnapshot,
652 },
653 ManualForceCompaction { task_type: TaskType },
656}
657
658impl SinkUpdateKind {
659 fn apply_to_track(self, track: &mut CompactionTrack, now: Instant) {
660 match self {
661 SinkUpdateKind::Commit { observed_snapshot } => {
662 track.record_observed_snapshot(observed_snapshot);
663 track.record_commit();
664 }
665 SinkUpdateKind::ForceCompaction { observed_snapshot } => {
666 track.record_observed_snapshot(observed_snapshot);
667 track.record_force_compaction(now, None);
668 }
669 SinkUpdateKind::ManualForceCompaction { task_type } => {
670 track.record_force_compaction(now, Some(task_type))
671 }
672 }
673 }
674
675 fn allows_disabled_compaction(&self) -> bool {
676 matches!(self, SinkUpdateKind::ManualForceCompaction { .. })
677 }
678}
679
680struct PreparedSinkUpdate {
690 sink_id: SinkId,
691 kind: SinkUpdateKind,
692 now: Instant,
693 allow_track_initialization: bool,
694 loaded_config: Option<IcebergConfig>,
695}
696
697#[derive(Debug, Clone)]
698pub struct IcebergCompactionScheduleStatus {
699 pub sink_id: SinkId,
700 pub task_type: String,
701 pub trigger_interval_sec: u64,
702 pub trigger_snapshot_count: usize,
703 pub schedule_state: String,
704 pub next_compaction_after_sec: Option<u64>,
705 pub pending_snapshot_count: Option<usize>,
706 pub is_triggerable: bool,
707 pub last_success_at: Option<SystemTime>,
708 pub last_failure_at: Option<SystemTime>,
709 pub last_error: Option<String>,
710 pub consecutive_failures: u32,
711 pub compaction_lag: Option<Duration>,
715}
716
717impl IcebergCompactionManager {
718 fn apply_track_finish_action(
719 guard: &mut IcebergCompactionManagerInner,
720 sink_id: SinkId,
721 finish_action: CompactionTrackFinishAction,
722 ) {
723 match finish_action {
724 CompactionTrackFinishAction::KeepTrack => {}
725 CompactionTrackFinishAction::RemoveTrack => {
726 guard.sink_schedules.remove(&sink_id);
727 }
728 }
729 }
730
731 pub(super) fn refresh_schedule_config(
732 &self,
733 track: &mut CompactionTrack,
734 iceberg_config: &IcebergConfig,
735 now: Instant,
736 ) {
737 let (configured_task_type, write_mode, trigger_interval_sec, trigger_snapshot_count) =
738 self.resolve_schedule_values(iceberg_config);
739 debug_assert_eq!(
740 track.write_mode, write_mode,
741 "Iceberg write mode cannot change while a schedule track exists"
742 );
743 track.configured_task_type = configured_task_type;
744 track.trigger_snapshot_count = trigger_snapshot_count;
745 track.update_interval(trigger_interval_sec, now);
746 track.mark_config_refreshed(now);
747 }
748
749 pub fn schedule_on_startup_task(manager: Arc<Self>) -> (JoinHandle<()>, oneshot::Sender<()>) {
752 let (shutdown_tx, shutdown_rx) = oneshot::channel();
753 let join_handle = tokio::spawn(async move {
754 tokio::select! {
755 _ = manager.schedule_existing_sinks() => {}
756 _ = shutdown_rx => {}
757 }
758 });
759 (join_handle, shutdown_tx)
760 }
761
762 async fn schedule_existing_sinks(&self) {
763 let sinks = match self.metadata_manager.catalog_controller.list_sinks().await {
764 Ok(sinks) => sinks,
765 Err(e) => {
766 tracing::warn!(
767 error = %e.as_report(),
768 "failed to list sinks to schedule iceberg compaction on startup"
769 );
770 return;
771 }
772 };
773 let sink_ids = sinks
774 .into_iter()
775 .filter(|sink| {
776 sink.stream_job_status() == PbStreamJobStatus::Created
777 && is_iceberg_sink(&sink.properties)
778 })
779 .map(|sink| sink.id);
780 futures::stream::iter(sink_ids)
781 .for_each_concurrent(16, |sink_id| async move {
782 if let Err(e) = self.schedule_sink_on_startup(sink_id).await {
783 tracing::warn!(
784 %sink_id,
785 error = %e.as_report(),
786 "failed to schedule iceberg compaction on startup"
787 );
788 }
789 })
790 .await;
791 }
792
793 async fn schedule_sink_on_startup(&self, sink_id: SinkId) -> MetaResult<()> {
796 if self.inner.read().sink_schedules.contains_key(&sink_id) {
798 return Ok(());
799 }
800 let config = self.load_iceberg_config(sink_id).await?;
801 if !config.enable_compaction {
802 return Ok(());
803 }
804 let table = config.load_table().await?;
805 let branch = commit_branch(config.r#type.as_str(), config.write_mode);
806 let Some(observed_snapshot) =
808 IcebergCommittedSnapshot::from_branch_head(table.metadata(), &branch)
809 else {
810 return Ok(());
811 };
812 let prepared_update = PreparedSinkUpdate {
813 sink_id,
814 kind: SinkUpdateKind::Commit { observed_snapshot },
815 now: Instant::now(),
816 allow_track_initialization: true,
817 loaded_config: Some(config),
818 };
819 let mut guard = self.inner.write();
820 if !guard.sink_schedules.contains_key(&sink_id) {
823 self.apply_sink_update(&mut guard, prepared_update);
824 }
825 Ok(())
826 }
827
828 pub async fn update_iceberg_commit_info(&self, msg: IcebergSinkCompactionUpdate) {
829 let IcebergSinkCompactionUpdate {
830 sink_id,
831 force_compaction,
832 observed_snapshot,
833 } = msg;
834 let kind = if force_compaction {
835 SinkUpdateKind::ForceCompaction { observed_snapshot }
836 } else {
837 SinkUpdateKind::Commit { observed_snapshot }
838 };
839 let prepared_update = self
840 .prepare_sink_update(sink_id, kind, Instant::now())
841 .await;
842
843 let mut guard = self.inner.write();
844 self.apply_sink_update(&mut guard, prepared_update);
845 }
846
847 async fn prepare_sink_update(
848 &self,
849 sink_id: SinkId,
850 kind: SinkUpdateKind,
851 now: Instant,
852 ) -> PreparedSinkUpdate {
853 let refresh_interval = self.config_refresh_interval();
854 let (allow_track_initialization, should_refresh_config) = {
855 let guard = self.inner.read();
856 match guard.sink_schedules.get(&sink_id) {
857 Some(track) => (false, track.needs_config_refresh(now, refresh_interval)),
858 None => (true, true),
859 }
860 };
861
862 let loaded_config = if should_refresh_config {
863 match self.load_iceberg_config(sink_id).await {
864 Ok(config) => Some(config),
865 Err(e) => {
866 tracing::warn!(
867 error = ?e.as_report(),
868 "Failed to load iceberg config for sink {}",
869 sink_id
870 );
871 None
872 }
873 }
874 } else {
875 None
876 };
877
878 PreparedSinkUpdate {
879 sink_id,
880 kind,
881 now,
882 allow_track_initialization,
883 loaded_config,
884 }
885 }
886
887 fn apply_sink_update(
888 &self,
889 guard: &mut IcebergCompactionManagerInner,
890 prepared_update: PreparedSinkUpdate,
891 ) -> bool {
892 let PreparedSinkUpdate {
893 sink_id,
894 kind,
895 now,
896 allow_track_initialization,
897 loaded_config,
898 } = prepared_update;
899 let refresh_interval = self.config_refresh_interval();
900
901 if let Some(config) = loaded_config.as_ref() {
902 if config.enable_snapshot_expiration {
903 guard.snapshot_expiration_sink_ids.insert(sink_id);
904 } else {
905 guard.snapshot_expiration_sink_ids.remove(&sink_id);
906 }
907 if config.enable_manifest_rewrite {
908 guard.manifest_rewrite_sink_ids.insert(sink_id);
909 } else {
910 guard.manifest_rewrite_sink_ids.remove(&sink_id);
911 }
912
913 if !config.enable_compaction && !kind.allows_disabled_compaction() {
914 let keep_until_task_finishes =
915 guard.sink_schedules.get_mut(&sink_id).is_some_and(|track| {
916 let keep = track.has_queued_task_or_active_attempt();
917 if keep {
918 track.finish_action = CompactionTrackFinishAction::RemoveTrack;
922 }
923 keep
924 });
925 if !keep_until_task_finishes {
926 guard.sink_schedules.remove(&sink_id);
927 }
928 return false;
929 }
930 }
931
932 match guard.sink_schedules.entry(sink_id) {
933 Entry::Occupied(entry) => {
934 let track = entry.into_mut();
935 if track.removes_track_after_finish()
936 && !kind.allows_disabled_compaction()
937 && !loaded_config
938 .as_ref()
939 .is_some_and(|config| config.enable_compaction)
940 {
941 return false;
942 }
943 if track.needs_config_refresh(now, refresh_interval)
944 && let Some(config) = loaded_config.as_ref()
945 {
946 self.refresh_schedule_config(track, config, now);
947 }
948 if let Some(config) = loaded_config.as_ref() {
949 track.finish_action =
950 if kind.allows_disabled_compaction() && !config.enable_compaction {
951 CompactionTrackFinishAction::RemoveTrack
952 } else {
953 CompactionTrackFinishAction::KeepTrack
954 };
955 }
956
957 kind.apply_to_track(track, now);
958 true
959 }
960 Entry::Vacant(entry) => {
961 if !allow_track_initialization {
962 tracing::warn!(
963 iceberg_component = "compaction_scheduler",
964 iceberg_operation = "apply_sink_update",
965 sink_id = %sink_id,
966 "iceberg_compaction_update_ignored_track_missing",
967 );
968 return false;
969 }
970
971 let Some(config) = loaded_config.as_ref() else {
972 tracing::warn!(
973 iceberg_component = "compaction_scheduler",
974 iceberg_operation = "apply_sink_update",
975 sink_id = %sink_id,
976 "iceberg_compaction_update_ignored_config_unavailable",
977 );
978 return false;
979 };
980
981 let track = entry.insert(self.create_compaction_track(config, now));
982 track.finish_action =
983 if kind.allows_disabled_compaction() && !config.enable_compaction {
984 CompactionTrackFinishAction::RemoveTrack
985 } else {
986 CompactionTrackFinishAction::KeepTrack
987 };
988 kind.apply_to_track(track, now);
989 true
990 }
991 }
992 }
993
994 pub(super) fn create_compaction_track(
995 &self,
996 iceberg_config: &IcebergConfig,
997 now: Instant,
998 ) -> CompactionTrack {
999 let (configured_task_type, write_mode, trigger_interval_sec, trigger_snapshot_count) =
1000 self.resolve_schedule_values(iceberg_config);
1001
1002 CompactionTrack::new(
1003 configured_task_type,
1004 write_mode,
1005 trigger_interval_sec,
1006 trigger_snapshot_count,
1007 self.report_timeout(),
1008 now,
1009 )
1010 }
1011
1012 fn resolve_schedule_values(
1013 &self,
1014 iceberg_config: &IcebergConfig,
1015 ) -> (TaskType, IcebergWriteMode, u64, usize) {
1016 let configured_type = match iceberg_config.write_mode {
1019 IcebergWriteMode::CopyOnWrite => None,
1020 IcebergWriteMode::MergeOnRead => iceberg_config.compaction_type,
1021 };
1022 let compaction_type = match configured_type {
1023 Some(compaction_type) => compaction_type,
1024 None if Feature::IcebergCompaction.check_available().is_ok() => CompactionType::Auto,
1025 None => CompactionType::Full,
1026 };
1027
1028 (
1029 match compaction_type {
1030 CompactionType::Auto => TaskType::Auto,
1031 CompactionType::Full => TaskType::Full,
1032 CompactionType::SmallFiles => TaskType::SmallFiles,
1033 CompactionType::FilesWithDelete => TaskType::FilesWithDelete,
1034 },
1035 iceberg_config.write_mode,
1036 iceberg_config.compaction_interval_sec(),
1037 iceberg_config.trigger_snapshot_count(),
1038 )
1039 }
1040
1041 pub(super) async fn start_manual_compaction(
1042 &self,
1043 sink_id: SinkId,
1044 ) -> MetaResult<oneshot::Receiver<MetaResult<IcebergCompactionTaskId>>> {
1045 let prepared_update = self
1046 .prepare_sink_update(
1047 sink_id,
1048 SinkUpdateKind::ManualForceCompaction {
1049 task_type: TaskType::Full,
1050 },
1051 Instant::now(),
1052 )
1053 .await;
1054 let mut guard = self.inner.write();
1055 let now = Instant::now();
1056 if guard.manual_compaction_waiters.contains_key(&sink_id) {
1057 return Err(anyhow!(
1058 "manual iceberg compaction is already waiting for sink {}",
1059 sink_id
1060 )
1061 .into());
1062 }
1063
1064 if let Some(track) = guard.sink_schedules.get(&sink_id) {
1065 if track.round_max_file_sequence_number.is_some() {
1066 return Err(anyhow!(
1067 "manual Full compaction is rejected while an automatic round is active for sink {}",
1068 sink_id
1069 )
1070 .into());
1071 }
1072 match &track.state {
1073 CompactionTrackState::PendingDispatch { attempt } => {
1074 return Err(anyhow!(
1075 "iceberg compaction task is already running for sink {} \
1076 (state=pending_dispatch, pending_commit_count_at_start={}, \
1077 pending_commit_count={})",
1078 sink_id,
1079 attempt.pending_commit_count_at_start,
1080 track.pending_commit_count
1081 )
1082 .into());
1083 }
1084 CompactionTrackState::InFlight {
1085 task_id,
1086 attempt,
1087 report_deadline,
1088 ..
1089 } => {
1090 return Err(anyhow!(
1091 "iceberg compaction task is already running for sink {} \
1092 (state=in_flight, task_id={}, pending_commit_count_at_start={}, \
1093 pending_commit_count={}, report_timeout_after_sec={})",
1094 sink_id,
1095 task_id,
1096 attempt.pending_commit_count_at_start,
1097 track.pending_commit_count,
1098 report_deadline.saturating_duration_since(now).as_secs()
1099 )
1100 .into());
1101 }
1102 CompactionTrackState::Idle { .. } => {}
1103 }
1104 }
1105
1106 if self.apply_sink_update(&mut guard, prepared_update) {
1107 let (tx, rx) = oneshot::channel();
1108 guard.manual_compaction_waiters.insert(sink_id, tx);
1109 Ok(rx)
1110 } else {
1111 Err(anyhow!(
1112 "failed to trigger manual iceberg compaction for sink {}",
1113 sink_id
1114 )
1115 .into())
1116 }
1117 }
1118
1119 pub(super) fn cancel_manual_compaction_waiter(&self, sink_id: SinkId) {
1120 self.inner
1121 .write()
1122 .manual_compaction_waiters
1123 .remove(&sink_id);
1124 }
1125
1126 fn finish_timed_out_compaction_tasks(
1127 guard: &mut IcebergCompactionManagerInner,
1128 now: Instant,
1129 ) -> Vec<(SinkId, ManualCompactionWaiter)> {
1130 let mut timed_out_tasks = Vec::new();
1131 for (&sink_id, track) in &mut guard.sink_schedules {
1132 if track.is_report_timed_out(now) {
1133 tracing::warn!(
1134 iceberg_component = "compaction_scheduler",
1135 iceberg_operation = "report_timeout",
1136 sink_id = %sink_id,
1137 "iceberg_compaction_task_report_timed_out",
1138 );
1139 let error = format!(
1140 "compaction task report timed out after {}s",
1141 track.report_timeout.as_secs()
1142 );
1143 timed_out_tasks.push((sink_id, track.finish_failed(now, error)));
1144 }
1145 }
1146
1147 let mut timed_out_waiters = Vec::new();
1148 for (sink_id, finish_action) in timed_out_tasks {
1149 if let Some(waiter) = guard.manual_compaction_waiters.remove(&sink_id) {
1150 timed_out_waiters.push((sink_id, waiter));
1151 }
1152 Self::apply_track_finish_action(guard, sink_id, finish_action);
1153 }
1154 timed_out_waiters
1155 }
1156
1157 pub(crate) fn get_top_n_iceberg_commit_sink_ids(
1158 &self,
1159 n: usize,
1160 ) -> Vec<IcebergCompactionHandle> {
1161 let now = Instant::now();
1162 let (handles, timed_out_waiters) = {
1163 let mut guard = self.inner.write();
1164 let timed_out_waiters = Self::finish_timed_out_compaction_tasks(&mut guard, now);
1165
1166 let mut candidates = Vec::new();
1167 for (sink_id, track) in &guard.sink_schedules {
1168 if track.should_trigger(now)
1169 && let CompactionTrackState::Idle {
1170 next_compaction_time,
1171 ..
1172 } = &track.state
1173 {
1174 candidates.push((*sink_id, *next_compaction_time));
1175 }
1176 }
1177
1178 candidates.sort_by_key(|c| c.1);
1179
1180 let handles = candidates
1181 .into_iter()
1182 .take(n)
1183 .filter_map(|(sink_id, _)| {
1184 let track = guard.sink_schedules.get_mut(&sink_id)?;
1185 let attempt = track.start_attempt();
1186
1187 Some(IcebergCompactionHandle::new(
1188 sink_id,
1189 attempt,
1190 self.inner.clone(),
1191 self.metadata_manager.clone(),
1192 ))
1193 })
1194 .collect();
1195
1196 (handles, timed_out_waiters)
1197 };
1198
1199 for (sink_id, waiter) in timed_out_waiters {
1200 let _ = waiter.send(Err(anyhow!(
1201 "Iceberg compaction task report timed out for sink {}",
1202 sink_id
1203 )
1204 .into()));
1205 }
1206
1207 handles
1208 }
1209
1210 pub fn clear_iceberg_maintenance_by_sink_id(&self, sink_id: SinkId) {
1211 let (task_to_cancel, waiter) = {
1212 let mut guard = self.inner.write();
1213 let task_to_cancel = Self::remove_sink_schedule(&mut guard, sink_id);
1214 guard.snapshot_expiration_sink_ids.remove(&sink_id);
1215 guard.manifest_rewrite_sink_ids.remove(&sink_id);
1216 let waiter = guard.manual_compaction_waiters.remove(&sink_id);
1217 (task_to_cancel, waiter)
1218 };
1219 self.cancel_scheduled_task_if_any(sink_id, task_to_cancel);
1220
1221 if let Some(waiter) = waiter {
1222 let _ = waiter.send(Err(anyhow!(
1223 "Iceberg compaction maintenance was cleared for sink {}",
1224 sink_id
1225 )
1226 .into()));
1227 }
1228 }
1229
1230 fn remove_sink_schedule(
1231 guard: &mut IcebergCompactionManagerInner,
1232 sink_id: SinkId,
1233 ) -> Option<ScheduledCompactionTask> {
1234 guard
1235 .sink_schedules
1236 .remove(&sink_id)
1237 .and_then(|track| track.scheduled_task())
1238 }
1239
1240 fn cancel_scheduled_task_if_any(&self, sink_id: SinkId, task: Option<ScheduledCompactionTask>) {
1241 let Some(ScheduledCompactionTask {
1242 task_id,
1243 compactor_context_id,
1244 }) = task
1245 else {
1246 return;
1247 };
1248
1249 let Some(compactor) = self
1250 .iceberg_compactor_manager
1251 .get_compactor(compactor_context_id)
1252 else {
1253 tracing::warn!(
1254 sink_id = %sink_id,
1255 task_id = %task_id,
1256 compactor_context_id = %compactor_context_id,
1257 "Unable to cancel iceberg compaction task because compactor is no longer registered",
1258 );
1259 return;
1260 };
1261
1262 tracing::info!(
1263 sink_id = %sink_id,
1264 task_id = %task_id,
1265 compactor_context_id = %compactor_context_id,
1266 "Cancelling iceberg compaction task for removed schedule",
1267 );
1268
1269 if let Err(e) = compactor.cancel_task(task_id) {
1270 tracing::warn!(
1271 error = %e.as_report(),
1272 sink_id = %sink_id,
1273 task_id = %task_id,
1274 compactor_context_id = %compactor_context_id,
1275 "Failed to cancel iceberg compaction task for removed schedule",
1276 );
1277 }
1278 }
1279
1280 pub fn list_compaction_statuses(&self) -> Vec<IcebergCompactionScheduleStatus> {
1281 let now = Instant::now();
1282 let schedules = {
1283 let guard = self.inner.read();
1284 guard
1285 .sink_schedules
1286 .iter()
1287 .map(|(&sink_id, track)| (sink_id, track.clone()))
1288 .collect_vec()
1289 };
1290
1291 let mut statuses = schedules
1292 .into_iter()
1293 .map(|(sink_id, track)| {
1294 let next_compaction_after_sec = match &track.state {
1295 CompactionTrackState::Idle {
1296 next_compaction_time,
1297 ..
1298 } => Some(
1299 next_compaction_time
1300 .saturating_duration_since(now)
1301 .as_secs(),
1302 ),
1303 CompactionTrackState::PendingDispatch { .. }
1304 | CompactionTrackState::InFlight { .. } => None,
1305 };
1306 let is_triggerable = track.should_trigger(now);
1307 let compaction_lag = track.compaction_lag();
1308
1309 IcebergCompactionScheduleStatus {
1310 sink_id,
1311 task_type: track
1312 .effective_task_type()
1313 .as_str_name()
1314 .to_ascii_lowercase(),
1315 trigger_interval_sec: track.trigger_interval_sec,
1316 trigger_snapshot_count: track.trigger_snapshot_count,
1317 schedule_state: match track.state {
1318 CompactionTrackState::Idle { .. } => "idle".to_owned(),
1319 CompactionTrackState::PendingDispatch { .. }
1320 | CompactionTrackState::InFlight { .. } => "processing".to_owned(),
1321 },
1322 next_compaction_after_sec,
1323 pending_snapshot_count: Some(track.pending_commit_count),
1324 is_triggerable,
1325 last_success_at: track.health.last_success_at,
1326 last_failure_at: track.health.last_failure_at,
1327 last_error: track.health.last_error,
1328 consecutive_failures: track.health.consecutive_failures,
1329 compaction_lag,
1330 }
1331 })
1332 .collect_vec();
1333
1334 statuses.sort_by_key(|status| status.sink_id);
1335 statuses
1336 }
1337
1338 pub fn handle_report_task(&self, report: IcebergReportTask) {
1339 let sink_id = SinkId::from(report.sink_id);
1340 let task_id = report.task_id;
1341 let status = IcebergReportTaskStatus::try_from(report.status)
1342 .unwrap_or(IcebergReportTaskStatus::Unspecified);
1343 let now = Instant::now();
1344
1345 let waiter = {
1346 let mut guard = self.inner.write();
1347 let mut waiter = None;
1348
1349 match guard.sink_schedules.get_mut(&sink_id) {
1350 Some(track) if track.is_in_flight_task(task_id) => {
1351 let finish_action = match status {
1352 IcebergReportTaskStatus::Success => track.finish_success(now),
1353 IcebergReportTaskStatus::Drained
1354 if track.is_in_flight_bounded_attempt() =>
1355 {
1356 track.finish_drained(now)
1357 }
1358 IcebergReportTaskStatus::BlockedCapacity
1359 | IcebergReportTaskStatus::BlockedStatistics
1360 | IcebergReportTaskStatus::Busy
1361 | IcebergReportTaskStatus::Drained
1362 | IcebergReportTaskStatus::Failed
1363 | IcebergReportTaskStatus::Unspecified => {
1364 tracing::warn!(
1365 iceberg_component = "compaction_scheduler",
1366 iceberg_operation = "handle_report",
1367 sink_id = %sink_id,
1368 task_id = %task_id,
1369 status = ?status,
1370 error_message = report.error_message.as_deref().unwrap_or_default(),
1371 "iceberg_compaction_task_reported_failure",
1372 );
1373 let error = report.error_message.clone().unwrap_or_else(|| {
1374 format!("compaction task reported {}", status.as_str_name())
1375 });
1376 track.finish_failed(now, error)
1377 }
1378 };
1379
1380 Self::apply_track_finish_action(&mut guard, sink_id, finish_action);
1381 waiter = guard.manual_compaction_waiters.remove(&sink_id);
1382 }
1383 Some(_) => {
1384 tracing::warn!(
1385 iceberg_component = "compaction_scheduler",
1386 iceberg_operation = "handle_report",
1387 sink_id = %sink_id,
1388 task_id = %task_id,
1389 status = ?status,
1390 "iceberg_compaction_report_ignored_stale",
1391 );
1392 }
1393 None => {
1394 tracing::warn!(
1395 iceberg_component = "compaction_scheduler",
1396 iceberg_operation = "handle_report",
1397 sink_id = %sink_id,
1398 task_id = %task_id,
1399 status = ?status,
1400 "iceberg_compaction_report_unknown_sink",
1401 );
1402 }
1403 }
1404
1405 waiter
1406 };
1407
1408 if let Some(waiter) = waiter {
1409 Self::complete_manual_task_waiter(waiter, &report);
1410 }
1411 }
1412}
1413
1414#[cfg(test)]
1415mod tests;