1use std::collections::hash_map::Entry;
16use std::sync::Arc;
17use std::time::{Duration, Instant};
18
19use itertools::Itertools;
20use parking_lot::RwLock;
21use risingwave_connector::connector_common::{
22 IcebergCommittedSnapshot, IcebergSinkCompactionUpdate,
23};
24use risingwave_connector::sink::SinkParam;
25use risingwave_connector::sink::catalog::{SinkCatalog, SinkId};
26use risingwave_connector::sink::iceberg::{
27 CompactionType, IcebergConfig, should_enable_iceberg_cow,
28};
29use risingwave_hummock_sdk::HummockContextId;
30use risingwave_pb::iceberg_compaction::IcebergCompactionTask;
31use risingwave_pb::iceberg_compaction::iceberg_compaction_task::TaskType;
32use risingwave_pb::iceberg_compaction::subscribe_iceberg_compaction_event_request::ReportTask as IcebergReportTask;
33use risingwave_pb::iceberg_compaction::subscribe_iceberg_compaction_event_request::report_task::Status as IcebergReportTaskStatus;
34use thiserror_ext::AsReport;
35use tokio::sync::oneshot;
36
37use super::*;
38
39#[derive(Debug, Clone)]
41enum CompactionTrackState {
42 Idle {
47 next_compaction_time: Instant,
48 next_task_type_override: Option<TaskType>,
51 },
52 PendingDispatch {
54 task_type: TaskType,
55 pending_commit_count_at_dispatch: usize,
56 gc_watermark_snapshot: Option<IcebergCommittedSnapshot>,
57 },
58 InFlight {
61 task_id: u64,
62 compactor_context_id: HummockContextId,
63 task_type: TaskType,
64 pending_commit_count_at_dispatch: usize,
65 report_deadline: Instant,
66 gc_watermark_snapshot: Option<IcebergCommittedSnapshot>,
67 },
68}
69
70#[derive(Debug, Clone, Copy)]
71struct ScheduledCompactionTask {
72 task_id: u64,
73 compactor_context_id: HummockContextId,
74}
75
76#[derive(Debug, Clone, Copy, PartialEq, Eq)]
77enum CompactionTrackFinishAction {
78 KeepTrack,
79 RemoveTrack,
80}
81
82#[derive(Debug, Clone)]
83pub(super) struct CompactionTrack {
84 task_type: TaskType,
85 trigger_interval_sec: u64,
86 trigger_snapshot_count: usize,
89 report_timeout: Duration,
90 last_config_refresh_at: Instant,
91 pending_commit_count: usize,
92 latest_observed_snapshot: Option<IcebergCommittedSnapshot>,
93 finish_action: CompactionTrackFinishAction,
97 state: CompactionTrackState,
98}
99
100impl CompactionTrack {
101 fn new(
102 task_type: TaskType,
103 trigger_interval_sec: u64,
104 trigger_snapshot_count: usize,
105 report_timeout: Duration,
106 now: Instant,
107 ) -> Self {
108 Self {
109 task_type,
110 trigger_interval_sec,
111 trigger_snapshot_count,
112 report_timeout,
113 last_config_refresh_at: now,
114 pending_commit_count: 0,
115 latest_observed_snapshot: None,
116 finish_action: CompactionTrackFinishAction::KeepTrack,
117 state: CompactionTrackState::Idle {
118 next_compaction_time: now + Duration::from_secs(trigger_interval_sec),
119 next_task_type_override: None,
120 },
121 }
122 }
123
124 fn should_trigger(&self, now: Instant) -> bool {
136 let next_compaction_time = match &self.state {
137 CompactionTrackState::Idle {
138 next_compaction_time,
139 ..
140 } => *next_compaction_time,
141 CompactionTrackState::PendingDispatch { .. }
142 | CompactionTrackState::InFlight { .. } => return false,
143 };
144
145 let time_ready = now >= next_compaction_time;
146 let commit_ready = self.pending_commit_count >= self.trigger_snapshot_count;
147 let has_commits = self.pending_commit_count > 0;
148
149 commit_ready || (time_ready && has_commits)
150 }
151
152 fn record_observed_snapshot(&mut self, observed_snapshot: IcebergCommittedSnapshot) {
153 self.latest_observed_snapshot = Some(observed_snapshot);
154 }
155
156 fn record_commit(&mut self) {
157 self.pending_commit_count = self.pending_commit_count.saturating_add(1);
158 }
159
160 fn record_force_compaction(&mut self, now: Instant, forced_task_type: Option<TaskType>) {
161 if let CompactionTrackState::Idle {
162 next_compaction_time,
163 next_task_type_override,
164 } = &mut self.state
165 {
166 if let Some(task_type) = forced_task_type {
167 *next_task_type_override = Some(task_type);
168 }
169 *next_compaction_time = now;
170 self.pending_commit_count = self.pending_commit_count.max(1);
171 }
172 }
173
174 fn needs_config_refresh(&self, now: Instant, refresh_interval: Duration) -> bool {
175 now.saturating_duration_since(self.last_config_refresh_at) >= refresh_interval
176 }
177
178 fn should_refresh_config(&self, now: Instant, refresh_interval: Duration) -> bool {
179 matches!(self.state, CompactionTrackState::Idle { .. })
180 && self.needs_config_refresh(now, refresh_interval)
181 }
182
183 fn mark_config_refreshed(&mut self, now: Instant) {
184 self.last_config_refresh_at = now;
185 }
186
187 fn current_task_type(&self) -> TaskType {
188 match &self.state {
189 CompactionTrackState::Idle {
190 next_task_type_override,
191 ..
192 } => next_task_type_override.unwrap_or(self.task_type),
193 CompactionTrackState::PendingDispatch { task_type, .. }
194 | CompactionTrackState::InFlight { task_type, .. } => *task_type,
195 }
196 }
197
198 fn start_processing(&mut self) -> TaskType {
199 match &mut self.state {
200 CompactionTrackState::Idle {
201 next_task_type_override,
202 ..
203 } => {
204 let task_type = next_task_type_override.take().unwrap_or(self.task_type);
205 self.state = CompactionTrackState::PendingDispatch {
206 task_type,
207 pending_commit_count_at_dispatch: self.pending_commit_count,
208 gc_watermark_snapshot: self.latest_observed_snapshot.clone(),
209 };
210 task_type
211 }
212 CompactionTrackState::PendingDispatch { .. }
213 | CompactionTrackState::InFlight { .. } => {
214 unreachable!("Cannot start processing when already processing")
215 }
216 }
217 }
218
219 fn mark_dispatched(
220 &mut self,
221 task_id: u64,
222 compactor_context_id: HummockContextId,
223 now: Instant,
224 ) {
225 let (task_type, pending_commit_count_at_dispatch, gc_watermark_snapshot) = match &mut self
226 .state
227 {
228 CompactionTrackState::PendingDispatch {
229 task_type,
230 pending_commit_count_at_dispatch,
231 gc_watermark_snapshot,
232 ..
233 } => {
234 let gc_watermark_snapshot = gc_watermark_snapshot.take();
235 (
236 *task_type,
237 *pending_commit_count_at_dispatch,
238 gc_watermark_snapshot,
239 )
240 }
241 CompactionTrackState::Idle { .. } => unreachable!("Cannot mark dispatched when idle"),
242 CompactionTrackState::InFlight { .. } => {
243 unreachable!("Cannot mark dispatched when already in flight")
244 }
245 };
246 self.state = CompactionTrackState::InFlight {
247 task_id,
248 compactor_context_id,
249 task_type,
250 pending_commit_count_at_dispatch,
251 report_deadline: now + self.report_timeout,
252 gc_watermark_snapshot,
253 };
254 }
255
256 pub(super) fn processing_gc_watermark_snapshot(
257 &self,
258 ) -> Option<Option<&IcebergCommittedSnapshot>> {
259 match &self.state {
260 CompactionTrackState::PendingDispatch {
261 gc_watermark_snapshot,
262 ..
263 }
264 | CompactionTrackState::InFlight {
265 gc_watermark_snapshot,
266 ..
267 } => Some(gc_watermark_snapshot.as_ref()),
268 CompactionTrackState::Idle { .. } => None,
269 }
270 }
271
272 fn is_pending_dispatch(&self) -> bool {
273 matches!(self.state, CompactionTrackState::PendingDispatch { .. })
274 }
275
276 fn removes_track_after_finish(&self) -> bool {
277 self.finish_action == CompactionTrackFinishAction::RemoveTrack
278 }
279
280 pub(super) fn is_processing_task(&self, task_id: u64) -> bool {
281 matches!(
282 &self.state,
283 CompactionTrackState::InFlight {
284 task_id: current_task_id,
285 ..
286 } if *current_task_id == task_id
287 )
288 }
289
290 fn scheduled_task(&self) -> Option<ScheduledCompactionTask> {
291 match &self.state {
292 CompactionTrackState::InFlight {
293 task_id,
294 compactor_context_id,
295 ..
296 } => Some(ScheduledCompactionTask {
297 task_id: *task_id,
298 compactor_context_id: *compactor_context_id,
299 }),
300 CompactionTrackState::Idle { .. } | CompactionTrackState::PendingDispatch { .. } => {
301 None
302 }
303 }
304 }
305
306 fn is_report_timed_out(&self, now: Instant) -> bool {
307 matches!(
308 &self.state,
309 CompactionTrackState::InFlight {
310 report_deadline,
311 ..
312 } if now >= *report_deadline
313 )
314 }
315
316 fn finish_success(&mut self, now: Instant) -> CompactionTrackFinishAction {
317 match &self.state {
318 CompactionTrackState::InFlight {
319 pending_commit_count_at_dispatch,
320 ..
321 } => {
322 self.pending_commit_count = self
323 .pending_commit_count
324 .saturating_sub(*pending_commit_count_at_dispatch);
325 self.state = CompactionTrackState::Idle {
326 next_compaction_time: now + Duration::from_secs(self.trigger_interval_sec),
327 next_task_type_override: None,
328 };
329 self.finish_action
330 }
331 CompactionTrackState::Idle { .. } => unreachable!("Cannot finish success when idle"),
332 CompactionTrackState::PendingDispatch { .. } => {
333 unreachable!("Cannot finish success before task dispatch")
334 }
335 }
336 }
337
338 fn finish_failed(&mut self, now: Instant) -> CompactionTrackFinishAction {
339 match &self.state {
340 CompactionTrackState::InFlight { .. } => {
341 self.state = CompactionTrackState::Idle {
342 next_compaction_time: now,
343 next_task_type_override: None,
344 };
345 self.finish_action
346 }
347 CompactionTrackState::Idle { .. } => unreachable!("Cannot finish failed when idle"),
348 CompactionTrackState::PendingDispatch { .. } => {
349 unreachable!("Cannot finish failed before task dispatch")
350 }
351 }
352 }
353
354 fn revert_pre_dispatch_failure(&mut self, now: Instant) -> CompactionTrackFinishAction {
366 match &self.state {
367 CompactionTrackState::PendingDispatch { .. } => {
368 self.state = CompactionTrackState::Idle {
369 next_compaction_time: now,
370 next_task_type_override: None,
371 };
372 self.finish_action
373 }
374 CompactionTrackState::Idle { .. } => {
375 unreachable!("Cannot revert a pre-dispatch failure when idle")
376 }
377 CompactionTrackState::InFlight { .. } => {
378 unreachable!("Cannot revert a pre-dispatch failure after dispatch")
379 }
380 }
381 }
382
383 fn update_interval(&mut self, new_interval_sec: u64, now: Instant) {
384 if self.trigger_interval_sec == new_interval_sec {
385 return;
386 }
387
388 self.trigger_interval_sec = new_interval_sec;
389
390 match &mut self.state {
391 CompactionTrackState::Idle {
392 next_compaction_time,
393 ..
394 } => {
395 *next_compaction_time = now + Duration::from_secs(new_interval_sec);
396 }
397 CompactionTrackState::PendingDispatch { .. }
398 | CompactionTrackState::InFlight { .. } => {}
399 }
400 }
401}
402
403pub(crate) struct IcebergCompactionHandle {
404 sink_id: SinkId,
405 task_type: TaskType,
406 inner: Arc<RwLock<IcebergCompactionManagerInner>>,
407 metadata_manager: MetadataManager,
408 handle_success: bool,
409}
410
411impl IcebergCompactionHandle {
412 fn new(
413 sink_id: SinkId,
414 task_type: TaskType,
415 inner: Arc<RwLock<IcebergCompactionManagerInner>>,
416 metadata_manager: MetadataManager,
417 ) -> Self {
418 Self {
419 sink_id,
420 task_type,
421 inner,
422 metadata_manager,
423 handle_success: false,
424 }
425 }
426
427 pub async fn send_compact_task(
428 mut self,
429 compactor: Arc<crate::hummock::IcebergCompactor>,
430 task_id: u64,
431 ) -> MetaResult<()> {
432 use risingwave_pb::iceberg_compaction::subscribe_iceberg_compaction_event_response::Event as IcebergResponseEvent;
433
434 let Some(prost_sink_catalog) = self
435 .metadata_manager
436 .catalog_controller
437 .get_sink_by_id(self.sink_id)
438 .await?
439 else {
440 tracing::warn!(
441 iceberg_component = "compaction_scheduler",
442 iceberg_operation = "dispatch_task",
443 sink_id = %self.sink_id,
444 task_id,
445 "iceberg_compaction_dispatch_sink_not_found",
446 );
447 return Ok(());
448 };
449 let sink_catalog = SinkCatalog::from(prost_sink_catalog);
450 let param = SinkParam::try_from_sink_catalog(sink_catalog)?;
451
452 let result =
453 compactor.send_event(IcebergResponseEvent::CompactTask(IcebergCompactionTask {
454 task_id,
455 sink_id: self.sink_id.as_raw_id(),
456 props: param.properties,
457 task_type: self.task_type as i32,
458 }));
459
460 if result.is_ok() {
461 let mut should_cancel_sent_task = false;
462 let mut guard = self.inner.write();
463 let mut dispatched = false;
464 if let Some(track) = guard.sink_schedules.get_mut(&self.sink_id)
465 && track.is_pending_dispatch()
466 {
467 track.mark_dispatched(task_id, compactor.context_id(), Instant::now());
468 dispatched = true;
469 }
470 self.handle_success = dispatched;
471 if !dispatched {
472 should_cancel_sent_task = true;
473 tracing::warn!(
474 iceberg_component = "compaction_scheduler",
475 iceberg_operation = "dispatch_task",
476 sink_id = %self.sink_id,
477 task_id,
478 "iceberg_compaction_dispatch_track_not_pending",
479 );
480 }
481 drop(guard);
482
483 if should_cancel_sent_task {
484 self.cancel_sent_task(&compactor, task_id);
485 }
486 }
487
488 result
489 }
490
491 fn cancel_sent_task(&self, compactor: &crate::hummock::IcebergCompactor, task_id: u64) {
492 if let Err(e) = compactor.cancel_task(task_id) {
493 tracing::warn!(
494 iceberg_component = "compaction_scheduler",
495 iceberg_operation = "cancel_task",
496 error = %e.as_report(),
497 sink_id = %self.sink_id,
498 task_id,
499 "iceberg_compaction_cancel_after_schedule_removal_failed",
500 );
501 }
502 }
503}
504
505impl Drop for IcebergCompactionHandle {
506 fn drop(&mut self) {
507 let waiter = {
508 let mut guard = self.inner.write();
509 if !self.handle_success
510 && let Some(track) = guard.sink_schedules.get_mut(&self.sink_id)
511 && track.is_pending_dispatch()
512 {
513 let finish_action = track.revert_pre_dispatch_failure(Instant::now());
514 let waiter = guard.manual_compaction_waiters.remove(&self.sink_id);
515 IcebergCompactionManager::apply_track_finish_action(
516 &mut guard,
517 self.sink_id,
518 finish_action,
519 );
520 waiter
521 } else {
522 None
523 }
524 };
525
526 if let Some(waiter) = waiter {
527 let _ = waiter.send(Err(anyhow!(
528 "Iceberg compaction task failed before dispatch for sink {}",
529 self.sink_id
530 )
531 .into()));
532 }
533 }
534}
535
536#[derive(Debug, Clone)]
537enum SinkUpdateKind {
538 Commit {
540 observed_snapshot: IcebergCommittedSnapshot,
541 },
542 ForceCompaction {
545 observed_snapshot: IcebergCommittedSnapshot,
546 },
547 ManualForceCompaction { task_type: TaskType },
550}
551
552impl SinkUpdateKind {
553 fn apply_to_track(self, track: &mut CompactionTrack, now: Instant) {
554 match self {
555 SinkUpdateKind::Commit { observed_snapshot } => {
556 track.record_observed_snapshot(observed_snapshot);
557 track.record_commit();
558 }
559 SinkUpdateKind::ForceCompaction { observed_snapshot } => {
560 if matches!(track.state, CompactionTrackState::Idle { .. }) {
561 track.record_observed_snapshot(observed_snapshot);
562 }
563 track.record_force_compaction(now, None);
564 }
565 SinkUpdateKind::ManualForceCompaction { task_type } => {
566 track.record_force_compaction(now, Some(task_type))
567 }
568 }
569 }
570
571 fn allows_disabled_compaction(&self) -> bool {
572 matches!(self, SinkUpdateKind::ManualForceCompaction { .. })
573 }
574}
575
576struct PreparedSinkUpdate {
586 sink_id: SinkId,
587 kind: SinkUpdateKind,
588 now: Instant,
589 allow_track_initialization: bool,
590 loaded_config: Option<IcebergConfig>,
591}
592
593#[derive(Debug, Clone)]
594pub struct IcebergCompactionScheduleStatus {
595 pub sink_id: SinkId,
596 pub task_type: String,
597 pub trigger_interval_sec: u64,
598 pub trigger_snapshot_count: usize,
599 pub schedule_state: String,
600 pub next_compaction_after_sec: Option<u64>,
601 pub pending_snapshot_count: Option<usize>,
602 pub is_triggerable: bool,
603}
604
605impl IcebergCompactionManager {
606 fn apply_track_finish_action(
607 guard: &mut IcebergCompactionManagerInner,
608 sink_id: SinkId,
609 finish_action: CompactionTrackFinishAction,
610 ) {
611 match finish_action {
612 CompactionTrackFinishAction::KeepTrack => {}
613 CompactionTrackFinishAction::RemoveTrack => {
614 guard.sink_schedules.remove(&sink_id);
615 }
616 }
617 }
618
619 pub(super) fn refresh_schedule_config(
620 &self,
621 track: &mut CompactionTrack,
622 iceberg_config: &IcebergConfig,
623 now: Instant,
624 ) {
625 let (task_type, trigger_interval_sec, trigger_snapshot_count) =
626 self.resolve_schedule_values(iceberg_config);
627 track.task_type = task_type;
628 track.trigger_snapshot_count = trigger_snapshot_count;
629 track.update_interval(trigger_interval_sec, now);
630 track.mark_config_refreshed(now);
631 }
632
633 pub async fn update_iceberg_commit_info(&self, msg: IcebergSinkCompactionUpdate) {
634 let IcebergSinkCompactionUpdate {
635 sink_id,
636 force_compaction,
637 observed_snapshot,
638 } = msg;
639 let kind = if force_compaction {
640 SinkUpdateKind::ForceCompaction { observed_snapshot }
641 } else {
642 SinkUpdateKind::Commit { observed_snapshot }
643 };
644 let prepared_update = self
645 .prepare_sink_update(sink_id, kind, Instant::now())
646 .await;
647
648 let mut guard = self.inner.write();
649 self.apply_sink_update(&mut guard, prepared_update);
650 }
651
652 async fn prepare_sink_update(
653 &self,
654 sink_id: SinkId,
655 kind: SinkUpdateKind,
656 now: Instant,
657 ) -> PreparedSinkUpdate {
658 let refresh_interval = self.config_refresh_interval();
659 let (allow_track_initialization, should_refresh_config) = {
660 let guard = self.inner.read();
661 match guard.sink_schedules.get(&sink_id) {
662 Some(track) => (false, track.should_refresh_config(now, refresh_interval)),
663 None => (true, true),
664 }
665 };
666
667 let loaded_config = if should_refresh_config {
668 match self.load_iceberg_config(sink_id).await {
669 Ok(config) => Some(config),
670 Err(e) => {
671 tracing::warn!(
672 error = ?e.as_report(),
673 "Failed to load iceberg config for sink {}",
674 sink_id
675 );
676 None
677 }
678 }
679 } else {
680 None
681 };
682
683 PreparedSinkUpdate {
684 sink_id,
685 kind,
686 now,
687 allow_track_initialization,
688 loaded_config,
689 }
690 }
691
692 fn apply_sink_update(
693 &self,
694 guard: &mut IcebergCompactionManagerInner,
695 prepared_update: PreparedSinkUpdate,
696 ) -> bool {
697 let PreparedSinkUpdate {
698 sink_id,
699 kind,
700 now,
701 allow_track_initialization,
702 loaded_config,
703 } = prepared_update;
704 let refresh_interval = self.config_refresh_interval();
705
706 if let Some(config) = loaded_config.as_ref() {
707 if config.enable_snapshot_expiration {
708 guard.snapshot_expiration_sink_ids.insert(sink_id);
709 } else {
710 guard.snapshot_expiration_sink_ids.remove(&sink_id);
711 }
712
713 if !config.enable_compaction && !kind.allows_disabled_compaction() {
714 if !guard.sink_schedules.get(&sink_id).is_some_and(|track| {
715 matches!(
716 &track.state,
717 CompactionTrackState::PendingDispatch { .. }
718 | CompactionTrackState::InFlight { .. }
719 ) || track.removes_track_after_finish()
720 }) {
721 guard.sink_schedules.remove(&sink_id);
722 }
723 return false;
724 }
725 }
726
727 match guard.sink_schedules.entry(sink_id) {
728 Entry::Occupied(entry) => {
729 let track = entry.into_mut();
730 if track.removes_track_after_finish()
731 && !kind.allows_disabled_compaction()
732 && !loaded_config
733 .as_ref()
734 .is_some_and(|config| config.enable_compaction)
735 {
736 return false;
737 }
738 if track.should_refresh_config(now, refresh_interval)
739 && let Some(config) = loaded_config.as_ref()
740 {
741 self.refresh_schedule_config(track, config, now);
742 }
743 if let Some(config) = loaded_config.as_ref() {
744 track.finish_action =
745 if kind.allows_disabled_compaction() && !config.enable_compaction {
746 CompactionTrackFinishAction::RemoveTrack
747 } else {
748 CompactionTrackFinishAction::KeepTrack
749 };
750 }
751
752 kind.apply_to_track(track, now);
753 true
754 }
755 Entry::Vacant(entry) => {
756 if !allow_track_initialization {
757 tracing::warn!(
758 iceberg_component = "compaction_scheduler",
759 iceberg_operation = "apply_sink_update",
760 sink_id = %sink_id,
761 "iceberg_compaction_update_ignored_track_missing",
762 );
763 return false;
764 }
765
766 let Some(config) = loaded_config.as_ref() else {
767 tracing::warn!(
768 iceberg_component = "compaction_scheduler",
769 iceberg_operation = "apply_sink_update",
770 sink_id = %sink_id,
771 "iceberg_compaction_update_ignored_config_unavailable",
772 );
773 return false;
774 };
775
776 let track = entry.insert(self.create_compaction_track(config, now));
777 track.finish_action =
778 if kind.allows_disabled_compaction() && !config.enable_compaction {
779 CompactionTrackFinishAction::RemoveTrack
780 } else {
781 CompactionTrackFinishAction::KeepTrack
782 };
783 kind.apply_to_track(track, now);
784 true
785 }
786 }
787 }
788
789 pub(super) fn create_compaction_track(
790 &self,
791 iceberg_config: &IcebergConfig,
792 now: Instant,
793 ) -> CompactionTrack {
794 let (task_type, trigger_interval_sec, trigger_snapshot_count) =
795 self.resolve_schedule_values(iceberg_config);
796
797 CompactionTrack::new(
798 task_type,
799 trigger_interval_sec,
800 trigger_snapshot_count,
801 self.report_timeout(),
802 now,
803 )
804 }
805
806 fn resolve_schedule_values(&self, iceberg_config: &IcebergConfig) -> (TaskType, u64, usize) {
807 (
808 if should_enable_iceberg_cow(iceberg_config.r#type.as_str(), iceberg_config.write_mode)
809 {
810 TaskType::Full
811 } else {
812 match iceberg_config.compaction_type() {
813 CompactionType::Full => TaskType::Full,
814 CompactionType::SmallFiles => TaskType::SmallFiles,
815 CompactionType::FilesWithDelete => TaskType::FilesWithDelete,
816 }
817 },
818 iceberg_config.compaction_interval_sec(),
819 iceberg_config.trigger_snapshot_count(),
820 )
821 }
822
823 pub(super) async fn start_manual_compaction(
824 &self,
825 sink_id: SinkId,
826 ) -> MetaResult<oneshot::Receiver<MetaResult<u64>>> {
827 let prepared_update = self
828 .prepare_sink_update(
829 sink_id,
830 SinkUpdateKind::ManualForceCompaction {
831 task_type: TaskType::Full,
832 },
833 Instant::now(),
834 )
835 .await;
836 let mut guard = self.inner.write();
837 let now = Instant::now();
838 if guard.manual_compaction_waiters.contains_key(&sink_id) {
839 return Err(anyhow!(
840 "manual iceberg compaction is already waiting for sink {}",
841 sink_id
842 )
843 .into());
844 }
845
846 if let Some(track) = guard.sink_schedules.get(&sink_id) {
847 match &track.state {
848 CompactionTrackState::PendingDispatch {
849 pending_commit_count_at_dispatch,
850 ..
851 } => {
852 return Err(anyhow!(
853 "iceberg compaction task is already running for sink {} \
854 (state=pending_dispatch, pending_commit_count_at_dispatch={}, \
855 pending_commit_count={})",
856 sink_id,
857 pending_commit_count_at_dispatch,
858 track.pending_commit_count
859 )
860 .into());
861 }
862 CompactionTrackState::InFlight {
863 task_id,
864 pending_commit_count_at_dispatch,
865 report_deadline,
866 ..
867 } => {
868 return Err(anyhow!(
869 "iceberg compaction task is already running for sink {} \
870 (state=in_flight, task_id={}, pending_commit_count_at_dispatch={}, \
871 pending_commit_count={}, report_timeout_after_sec={})",
872 sink_id,
873 task_id,
874 pending_commit_count_at_dispatch,
875 track.pending_commit_count,
876 report_deadline.saturating_duration_since(now).as_secs()
877 )
878 .into());
879 }
880 CompactionTrackState::Idle { .. } => {}
881 }
882 }
883
884 if self.apply_sink_update(&mut guard, prepared_update) {
885 let (tx, rx) = oneshot::channel();
886 guard.manual_compaction_waiters.insert(sink_id, tx);
887 Ok(rx)
888 } else {
889 Err(anyhow!(
890 "failed to trigger manual iceberg compaction for sink {}",
891 sink_id
892 )
893 .into())
894 }
895 }
896
897 pub(super) fn cancel_manual_compaction_waiter(&self, sink_id: SinkId) {
898 self.inner
899 .write()
900 .manual_compaction_waiters
901 .remove(&sink_id);
902 }
903
904 fn finish_timed_out_compaction_tasks(
905 guard: &mut IcebergCompactionManagerInner,
906 now: Instant,
907 ) -> Vec<(SinkId, ManualCompactionWaiter)> {
908 let mut timed_out_tasks = Vec::new();
909 for (&sink_id, track) in &mut guard.sink_schedules {
910 if track.is_report_timed_out(now) {
911 tracing::warn!(
912 iceberg_component = "compaction_scheduler",
913 iceberg_operation = "report_timeout",
914 sink_id = %sink_id,
915 "iceberg_compaction_task_report_timed_out",
916 );
917 timed_out_tasks.push((sink_id, track.finish_failed(now)));
918 }
919 }
920
921 let mut timed_out_waiters = Vec::new();
922 for (sink_id, finish_action) in timed_out_tasks {
923 if let Some(waiter) = guard.manual_compaction_waiters.remove(&sink_id) {
924 timed_out_waiters.push((sink_id, waiter));
925 }
926 Self::apply_track_finish_action(guard, sink_id, finish_action);
927 }
928 timed_out_waiters
929 }
930
931 pub(crate) fn get_top_n_iceberg_commit_sink_ids(
932 &self,
933 n: usize,
934 ) -> Vec<IcebergCompactionHandle> {
935 let now = Instant::now();
936 let (handles, timed_out_waiters) = {
937 let mut guard = self.inner.write();
938 let timed_out_waiters = Self::finish_timed_out_compaction_tasks(&mut guard, now);
939
940 let mut candidates = Vec::new();
941 for (sink_id, track) in &guard.sink_schedules {
942 if track.should_trigger(now)
943 && let CompactionTrackState::Idle {
944 next_compaction_time,
945 ..
946 } = &track.state
947 {
948 candidates.push((*sink_id, *next_compaction_time));
949 }
950 }
951
952 candidates.sort_by_key(|c| c.1);
953
954 let handles = candidates
955 .into_iter()
956 .take(n)
957 .filter_map(|(sink_id, _)| {
958 let track = guard.sink_schedules.get_mut(&sink_id)?;
959 let task_type = track.start_processing();
960
961 Some(IcebergCompactionHandle::new(
962 sink_id,
963 task_type,
964 self.inner.clone(),
965 self.metadata_manager.clone(),
966 ))
967 })
968 .collect();
969
970 (handles, timed_out_waiters)
971 };
972
973 for (sink_id, waiter) in timed_out_waiters {
974 let _ = waiter.send(Err(anyhow!(
975 "Iceberg compaction task report timed out for sink {}",
976 sink_id
977 )
978 .into()));
979 }
980
981 handles
982 }
983
984 pub fn clear_iceberg_maintenance_by_sink_id(&self, sink_id: SinkId) {
985 let (task_to_cancel, waiter) = {
986 let mut guard = self.inner.write();
987 let task_to_cancel = Self::remove_sink_schedule(&mut guard, sink_id);
988 guard.snapshot_expiration_sink_ids.remove(&sink_id);
989 let waiter = guard.manual_compaction_waiters.remove(&sink_id);
990 (task_to_cancel, waiter)
991 };
992 self.cancel_scheduled_task_if_any(sink_id, task_to_cancel);
993
994 if let Some(waiter) = waiter {
995 let _ = waiter.send(Err(anyhow!(
996 "Iceberg compaction maintenance was cleared for sink {}",
997 sink_id
998 )
999 .into()));
1000 }
1001 }
1002
1003 fn remove_sink_schedule(
1004 guard: &mut IcebergCompactionManagerInner,
1005 sink_id: SinkId,
1006 ) -> Option<ScheduledCompactionTask> {
1007 guard
1008 .sink_schedules
1009 .remove(&sink_id)
1010 .and_then(|track| track.scheduled_task())
1011 }
1012
1013 fn cancel_scheduled_task_if_any(&self, sink_id: SinkId, task: Option<ScheduledCompactionTask>) {
1014 let Some(ScheduledCompactionTask {
1015 task_id,
1016 compactor_context_id,
1017 }) = task
1018 else {
1019 return;
1020 };
1021
1022 let Some(compactor) = self
1023 .iceberg_compactor_manager
1024 .get_compactor(compactor_context_id)
1025 else {
1026 tracing::warn!(
1027 sink_id = %sink_id,
1028 task_id,
1029 compactor_context_id = %compactor_context_id,
1030 "Unable to cancel iceberg compaction task because compactor is no longer registered",
1031 );
1032 return;
1033 };
1034
1035 tracing::info!(
1036 sink_id = %sink_id,
1037 task_id,
1038 compactor_context_id = %compactor_context_id,
1039 "Cancelling iceberg compaction task for removed schedule",
1040 );
1041
1042 if let Err(e) = compactor.cancel_task(task_id) {
1043 tracing::warn!(
1044 error = %e.as_report(),
1045 sink_id = %sink_id,
1046 task_id,
1047 compactor_context_id = %compactor_context_id,
1048 "Failed to cancel iceberg compaction task for removed schedule",
1049 );
1050 }
1051 }
1052
1053 pub fn list_compaction_statuses(&self) -> Vec<IcebergCompactionScheduleStatus> {
1054 let now = Instant::now();
1055 let schedules = {
1056 let guard = self.inner.read();
1057 guard
1058 .sink_schedules
1059 .iter()
1060 .map(|(&sink_id, track)| (sink_id, track.clone()))
1061 .collect_vec()
1062 };
1063
1064 let mut statuses = schedules
1065 .into_iter()
1066 .map(|(sink_id, track)| {
1067 let next_compaction_after_sec = match &track.state {
1068 CompactionTrackState::Idle {
1069 next_compaction_time,
1070 ..
1071 } => Some(
1072 next_compaction_time
1073 .saturating_duration_since(now)
1074 .as_secs(),
1075 ),
1076 CompactionTrackState::PendingDispatch { .. }
1077 | CompactionTrackState::InFlight { .. } => None,
1078 };
1079 let is_triggerable = track.should_trigger(now);
1080
1081 IcebergCompactionScheduleStatus {
1082 sink_id,
1083 task_type: track.current_task_type().as_str_name().to_ascii_lowercase(),
1084 trigger_interval_sec: track.trigger_interval_sec,
1085 trigger_snapshot_count: track.trigger_snapshot_count,
1086 schedule_state: match track.state {
1087 CompactionTrackState::Idle { .. } => "idle".to_owned(),
1088 CompactionTrackState::PendingDispatch { .. }
1089 | CompactionTrackState::InFlight { .. } => "processing".to_owned(),
1090 },
1091 next_compaction_after_sec,
1092 pending_snapshot_count: Some(track.pending_commit_count),
1093 is_triggerable,
1094 }
1095 })
1096 .collect_vec();
1097
1098 statuses.sort_by_key(|status| status.sink_id);
1099 statuses
1100 }
1101
1102 pub fn handle_report_task(&self, report: IcebergReportTask) {
1103 let sink_id = SinkId::from(report.sink_id);
1104 let task_id = report.task_id;
1105 let status = IcebergReportTaskStatus::try_from(report.status)
1106 .unwrap_or(IcebergReportTaskStatus::Unspecified);
1107 let now = Instant::now();
1108
1109 let waiter = {
1110 let mut guard = self.inner.write();
1111 let mut waiter = None;
1112
1113 match guard.sink_schedules.get_mut(&sink_id) {
1114 Some(track) if track.is_processing_task(task_id) => {
1115 let finish_action = match status {
1116 IcebergReportTaskStatus::Success => track.finish_success(now),
1117 IcebergReportTaskStatus::Failed | IcebergReportTaskStatus::Unspecified => {
1118 tracing::warn!(
1119 iceberg_component = "compaction_scheduler",
1120 iceberg_operation = "handle_report",
1121 sink_id = %sink_id,
1122 task_id,
1123 status = ?status,
1124 error_message = report.error_message.as_deref().unwrap_or_default(),
1125 "iceberg_compaction_task_reported_failure",
1126 );
1127 track.finish_failed(now)
1128 }
1129 };
1130
1131 Self::apply_track_finish_action(&mut guard, sink_id, finish_action);
1132 waiter = guard.manual_compaction_waiters.remove(&sink_id);
1133 }
1134 Some(_) => {
1135 tracing::warn!(
1136 iceberg_component = "compaction_scheduler",
1137 iceberg_operation = "handle_report",
1138 sink_id = %sink_id,
1139 task_id,
1140 status = ?status,
1141 "iceberg_compaction_report_ignored_stale",
1142 );
1143 }
1144 None => {
1145 tracing::warn!(
1146 iceberg_component = "compaction_scheduler",
1147 iceberg_operation = "handle_report",
1148 sink_id = %sink_id,
1149 task_id,
1150 status = ?status,
1151 "iceberg_compaction_report_unknown_sink",
1152 );
1153 }
1154 }
1155
1156 waiter
1157 };
1158
1159 if let Some(waiter) = waiter {
1160 Self::complete_manual_task_waiter(waiter, &report);
1161 }
1162 }
1163}
1164
1165#[cfg(test)]
1166mod tests;