Skip to main content

risingwave_meta/manager/iceberg_compaction/
schedule.rs

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