risingwave_storage/hummock/compactor/iceberg_compaction/
mod.rs

1// Copyright 2025 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
15pub use self::iceberg_compactor_runner::create_task_execution;
16#[cfg(madsim)]
17pub use self::report::set_simulated_pk_index_compaction_result;
18#[cfg(madsim)]
19pub(crate) use self::report::simulated_pk_index_compaction_result;
20pub(crate) use self::report::{
21    IcebergPlanCompletion, IcebergTaskReport, IcebergTaskTracker, PkIndexCompactionResult,
22    ReportSendResult, build_iceberg_task_report, build_pk_index_compaction_result,
23    flush_pending_iceberg_task_reports, send_or_buffer_iceberg_task_report,
24};
25use crate::hummock::compactor::iceberg_compaction::iceberg_compactor_runner::IcebergCompactionPlanRunner;
26
27pub(crate) mod iceberg_compactor_runner;
28pub(crate) mod report;
29
30use std::collections::{HashMap, VecDeque};
31use std::sync::Arc;
32
33use risingwave_pb::id::IcebergCompactionTaskId;
34use tokio::sync::Notify;
35
36/// Unique key combining `(task_id, plan_index)` since one task can have multiple plans.
37pub(crate) type TaskKey = (IcebergCompactionTaskId, usize);
38
39/// Task metadata for queue operations.
40#[derive(Debug, Clone)]
41pub struct IcebergTaskMeta {
42    pub task_id: IcebergCompactionTaskId,
43    pub plan_index: usize,
44    /// Must be in range `1..=max_parallelism`
45    pub required_parallelism: u32,
46}
47
48#[derive(Debug)]
49pub struct PoppedIcebergTask {
50    pub meta: IcebergTaskMeta,
51    pub runner: Option<IcebergCompactionPlanRunner>,
52}
53
54impl IcebergTaskMeta {
55    fn key(&self) -> TaskKey {
56        (self.task_id, self.plan_index)
57    }
58}
59
60/// Internal storage for the task queue.
61struct IcebergTaskQueueInner {
62    /// FIFO queue of waiting task metadata
63    deque: VecDeque<IcebergTaskMeta>,
64    /// Maps `(task_id, plan_index)` to `required_parallelism` for tracking
65    id_map: HashMap<TaskKey, u32>,
66    /// Sum of `required_parallelism` for all waiting tasks
67    waiting_parallelism_sum: u32,
68    /// Sum of `required_parallelism` for all running tasks
69    running_parallelism_sum: u32,
70    /// Optional runner payloads indexed by `(task_id, plan_index)`
71    runners: HashMap<TaskKey, IcebergCompactionPlanRunner>,
72}
73
74/// FIFO task queue with parallelism-based scheduling for Iceberg compaction.
75///
76/// Tasks execute in submission order when sufficient parallelism is available.
77/// The queue tracks waiting and running tasks to prevent over-commitment of resources.
78///
79/// Constraints:
80/// - Each task requires `1..=max_parallelism` units
81/// - Total waiting parallelism cannot exceed `pending_parallelism_budget`
82/// - Total running parallelism cannot exceed `max_parallelism`
83/// - Tasks block until enough parallelism is available
84///
85/// Note: The queue does NOT deduplicate or reorder tasks. Task management
86/// (deduplication, merging, cancellation) is Meta's responsibility.
87pub struct IcebergTaskQueue {
88    inner: IcebergTaskQueueInner,
89    /// Maximum concurrent parallelism for running tasks
90    max_parallelism: u32,
91    /// Maximum total parallelism for waiting tasks (backpressure limit)
92    pending_parallelism_budget: u32,
93    /// Notification for event-driven scheduling
94    schedule_notify: Arc<Notify>,
95}
96
97#[derive(Debug, PartialEq, Eq)]
98pub enum PushResult {
99    Added,
100    /// Would exceed `pending_parallelism_budget`
101    RejectedCapacity,
102    /// `required_parallelism` > `max_parallelism`
103    RejectedTooLarge,
104    /// `required_parallelism` == 0
105    RejectedInvalidParallelism,
106    /// Task with same `(task_id, plan_index)` already exists
107    RejectedDuplicate,
108}
109
110impl IcebergTaskQueue {
111    pub fn new(max_parallelism: u32, pending_parallelism_budget: u32) -> Self {
112        assert!(max_parallelism > 0, "max_parallelism must be > 0");
113        assert!(
114            pending_parallelism_budget >= max_parallelism,
115            "pending budget should allow at least one task"
116        );
117        Self {
118            inner: IcebergTaskQueueInner {
119                deque: VecDeque::new(),
120                id_map: HashMap::new(),
121                waiting_parallelism_sum: 0,
122                running_parallelism_sum: 0,
123                runners: HashMap::new(),
124            },
125            max_parallelism,
126            pending_parallelism_budget,
127            schedule_notify: Arc::new(Notify::new()),
128        }
129    }
130
131    /// Waits until there are tasks that can be scheduled.
132    ///
133    /// Returns `true` if there are schedulable tasks, `false` otherwise.
134    /// Use this in a `tokio::select!` to wake up when tasks become schedulable.
135    pub async fn wait_schedulable(&self) -> bool {
136        // Check if we have tasks that can be scheduled right now
137        if self.has_schedulable_tasks() {
138            return true;
139        }
140        // Otherwise wait for notification
141        self.schedule_notify.notified().await;
142        self.has_schedulable_tasks()
143    }
144
145    fn has_schedulable_tasks(&self) -> bool {
146        if let Some(front_task) = self.inner.deque.front() {
147            let available_parallelism = self
148                .max_parallelism
149                .saturating_sub(self.inner.running_parallelism_sum);
150            available_parallelism >= front_task.required_parallelism
151        } else {
152            false
153        }
154    }
155
156    fn notify_schedulable(&self) {
157        if self.has_schedulable_tasks() {
158            self.schedule_notify.notify_one();
159        }
160    }
161
162    pub fn running_parallelism_sum(&self) -> u32 {
163        self.inner.running_parallelism_sum
164    }
165
166    pub fn waiting_parallelism_sum(&self) -> u32 {
167        self.inner.waiting_parallelism_sum
168    }
169
170    fn available_parallelism(&self) -> u32 {
171        self.max_parallelism
172            .saturating_sub(self.inner.running_parallelism_sum)
173    }
174
175    /// Push a task into the queue.
176    ///
177    /// The task is validated and added to the end of the FIFO queue if constraints are met.
178    pub fn push(
179        &mut self,
180        meta: IcebergTaskMeta,
181        runner: Option<IcebergCompactionPlanRunner>,
182    ) -> PushResult {
183        if meta.required_parallelism == 0 {
184            return PushResult::RejectedInvalidParallelism;
185        }
186        if meta.required_parallelism > self.max_parallelism {
187            return PushResult::RejectedTooLarge;
188        }
189
190        let key = meta.key();
191
192        // Reject duplicate keys to prevent inconsistent state between id_map and deque
193        if self.inner.id_map.contains_key(&key) {
194            return PushResult::RejectedDuplicate;
195        }
196
197        let new_total = self.inner.waiting_parallelism_sum + meta.required_parallelism;
198        if new_total > self.pending_parallelism_budget {
199            return PushResult::RejectedCapacity;
200        }
201
202        self.inner.id_map.insert(key, meta.required_parallelism);
203        self.inner.waiting_parallelism_sum = new_total;
204
205        if let Some(r) = runner {
206            self.inner.runners.insert(key, r);
207        }
208
209        self.inner.deque.push_back(meta);
210
211        self.notify_schedulable();
212        PushResult::Added
213    }
214
215    /// Pop the next task if sufficient parallelism is available.
216    ///
217    /// Returns `None` if the queue is empty or the front task cannot fit
218    /// within the available parallelism budget.
219    pub fn pop(&mut self) -> Option<PoppedIcebergTask> {
220        let front = self.inner.deque.front()?;
221        if front.required_parallelism > self.available_parallelism() {
222            return None;
223        }
224
225        let meta = self.inner.deque.pop_front()?;
226        self.inner.waiting_parallelism_sum = self
227            .inner
228            .waiting_parallelism_sum
229            .saturating_sub(meta.required_parallelism);
230        self.inner.running_parallelism_sum = self
231            .inner
232            .running_parallelism_sum
233            .saturating_add(meta.required_parallelism);
234
235        let runner = self.inner.runners.remove(&meta.key());
236        Some(PoppedIcebergTask { meta, runner })
237    }
238
239    /// Mark a task as finished, freeing its parallelism for other tasks.
240    ///
241    /// Returns `true` if the task was found and removed, `false` otherwise.
242    pub fn finish_running(&mut self, task_key: TaskKey) -> bool {
243        let Some(required) = self.inner.id_map.remove(&task_key) else {
244            tracing::warn!(
245                task_id = %task_key.0,
246                plan_index = task_key.1,
247                "finish_running called for unknown task key, possible bug: double-finish or invalid key"
248            );
249            return false;
250        };
251
252        self.inner.running_parallelism_sum =
253            self.inner.running_parallelism_sum.saturating_sub(required);
254        self.inner.runners.remove(&task_key);
255        self.notify_schedulable();
256        true
257    }
258
259    /// Cancel all waiting plans belonging to the given task.
260    ///
261    /// Returns the number of waiting plans removed from the queue.
262    pub fn cancel_waiting_task(&mut self, task_id: IcebergCompactionTaskId) -> usize {
263        let mut retained = VecDeque::with_capacity(self.inner.deque.len());
264        let mut cancelled_parallelism = 0;
265        let mut cancelled_count = 0;
266
267        while let Some(meta) = self.inner.deque.pop_front() {
268            if meta.task_id == task_id {
269                cancelled_parallelism += meta.required_parallelism;
270                cancelled_count += 1;
271                self.inner.id_map.remove(&meta.key());
272                self.inner.runners.remove(&meta.key());
273            } else {
274                retained.push_back(meta);
275            }
276        }
277
278        self.inner.deque = retained;
279        self.inner.waiting_parallelism_sum = self
280            .inner
281            .waiting_parallelism_sum
282            .saturating_sub(cancelled_parallelism);
283
284        if cancelled_count > 0 {
285            self.notify_schedulable();
286        }
287
288        cancelled_count
289    }
290}
291
292#[cfg(test)]
293mod tests {
294    use super::*;
295
296    fn mk_meta(id: u64, plan_index: usize, p: u32) -> IcebergTaskMeta {
297        IcebergTaskMeta {
298            task_id: id.into(),
299            plan_index,
300            required_parallelism: p,
301        }
302    }
303
304    #[test]
305    fn test_basic_push_pop() {
306        let mut q = IcebergTaskQueue::new(8, 32);
307        assert_eq!(q.push(mk_meta(1, 0, 4), None), PushResult::Added);
308        assert_eq!(q.waiting_parallelism_sum(), 4);
309
310        let popped = q.pop().expect("should pop");
311        assert_eq!(popped.meta.task_id.as_raw_id(), 1);
312        assert_eq!(q.waiting_parallelism_sum(), 0);
313        assert_eq!(q.running_parallelism_sum(), 4);
314
315        assert!(q.finish_running((1.into(), 0)));
316        assert_eq!(q.running_parallelism_sum(), 0);
317    }
318
319    #[test]
320    fn test_fifo_ordering() {
321        let mut q = IcebergTaskQueue::new(8, 32);
322        assert_eq!(q.push(mk_meta(1, 0, 2), None), PushResult::Added);
323        assert_eq!(q.push(mk_meta(2, 0, 2), None), PushResult::Added);
324        assert_eq!(q.push(mk_meta(3, 0, 2), None), PushResult::Added);
325
326        assert_eq!(q.pop().unwrap().meta.task_id.as_raw_id(), 1);
327        assert_eq!(q.pop().unwrap().meta.task_id.as_raw_id(), 2);
328        assert_eq!(q.pop().unwrap().meta.task_id.as_raw_id(), 3);
329    }
330
331    #[test]
332    fn test_capacity_reject() {
333        let mut q = IcebergTaskQueue::new(4, 6);
334        assert_eq!(q.push(mk_meta(1, 0, 3), None), PushResult::Added);
335        assert_eq!(q.push(mk_meta(2, 0, 3), None), PushResult::Added); // sum=6
336        assert_eq!(q.push(mk_meta(3, 0, 1), None), PushResult::RejectedCapacity); // would exceed
337    }
338
339    #[test]
340    fn test_invalid_parallelism() {
341        let mut q = IcebergTaskQueue::new(4, 10);
342        assert_eq!(
343            q.push(mk_meta(1, 0, 0), None),
344            PushResult::RejectedInvalidParallelism
345        );
346        assert_eq!(q.push(mk_meta(2, 0, 5), None), PushResult::RejectedTooLarge); // > max
347    }
348
349    #[test]
350    fn test_duplicate_key_rejected() {
351        let mut q = IcebergTaskQueue::new(8, 32);
352        assert_eq!(q.push(mk_meta(1, 0, 3), None), PushResult::Added);
353        // Same (task_id, plan_index) should be rejected
354        assert_eq!(
355            q.push(mk_meta(1, 0, 5), None),
356            PushResult::RejectedDuplicate
357        );
358        // Parallelism sum should not have changed
359        assert_eq!(q.waiting_parallelism_sum(), 3);
360
361        // Different plan_index is allowed
362        assert_eq!(q.push(mk_meta(1, 1, 2), None), PushResult::Added);
363        assert_eq!(q.waiting_parallelism_sum(), 5);
364
365        // After pop and finish, the key can be reused
366        let p = q.pop().unwrap();
367        assert_eq!(p.meta.task_id.as_raw_id(), 1);
368        assert_eq!(p.meta.plan_index, 0);
369        q.finish_running((1.into(), 0));
370
371        // Now the same key can be pushed again
372        assert_eq!(q.push(mk_meta(1, 0, 4), None), PushResult::Added);
373    }
374
375    #[test]
376    fn test_pop_insufficient_parallelism() {
377        let mut q = IcebergTaskQueue::new(8, 32);
378        assert_eq!(q.push(mk_meta(1, 0, 6), None), PushResult::Added);
379        assert_eq!(q.push(mk_meta(2, 0, 4), None), PushResult::Added);
380
381        let p1 = q.pop().unwrap();
382        assert_eq!(p1.meta.task_id.as_raw_id(), 1);
383        // Not enough remaining parallelism (only 2 left)
384        assert!(q.pop().is_none());
385
386        // Finish first, then second becomes schedulable
387        assert!(q.finish_running((1.into(), 0)));
388        let p2 = q.pop().unwrap();
389        assert_eq!(p2.meta.task_id.as_raw_id(), 2);
390    }
391
392    #[test]
393    fn test_finish_nonexistent_task() {
394        let mut q = IcebergTaskQueue::new(4, 16);
395        assert!(!q.finish_running((999.into(), 0)));
396        assert_eq!(q.running_parallelism_sum(), 0);
397    }
398
399    #[test]
400    fn test_double_finish() {
401        let mut q = IcebergTaskQueue::new(8, 32);
402        assert_eq!(q.push(mk_meta(1, 0, 4), None), PushResult::Added);
403        q.pop().unwrap();
404        assert_eq!(q.running_parallelism_sum(), 4);
405
406        // First finish succeeds
407        assert!(q.finish_running((1.into(), 0)));
408        assert_eq!(q.running_parallelism_sum(), 0);
409
410        // Second finish on same key returns false (triggers warn log)
411        assert!(!q.finish_running((1.into(), 0)));
412        assert_eq!(q.running_parallelism_sum(), 0);
413    }
414
415    #[test]
416    fn test_max_parallelism_boundary() {
417        // pending_budget == max_parallelism: minimal valid configuration
418        let mut q = IcebergTaskQueue::new(4, 4);
419
420        // Task with required_parallelism == max_parallelism should be accepted
421        assert_eq!(q.push(mk_meta(1, 0, 4), None), PushResult::Added);
422        // Budget exhausted
423        assert_eq!(q.push(mk_meta(2, 0, 1), None), PushResult::RejectedCapacity);
424
425        // Can pop and run at full parallelism
426        let p = q.pop().unwrap();
427        assert_eq!(p.meta.required_parallelism, 4);
428        assert_eq!(q.running_parallelism_sum(), 4);
429
430        // No room for any new running task
431        assert_eq!(q.push(mk_meta(3, 0, 1), None), PushResult::Added);
432        assert!(q.pop().is_none());
433    }
434
435    #[test]
436    fn test_same_task_id_multiple_plans() {
437        let mut q = IcebergTaskQueue::new(10, 30);
438        let task_id = 1u64;
439
440        // Same task_id with different plan_index are independent
441        assert_eq!(q.push(mk_meta(task_id, 0, 3), None), PushResult::Added);
442        assert_eq!(q.push(mk_meta(task_id, 1, 4), None), PushResult::Added);
443        assert_eq!(q.push(mk_meta(task_id, 2, 2), None), PushResult::Added);
444        assert_eq!(q.waiting_parallelism_sum(), 9);
445
446        // Pop all
447        for i in 0..3 {
448            let p = q.pop().unwrap();
449            assert_eq!(p.meta.task_id.as_raw_id(), task_id);
450            assert_eq!(p.meta.plan_index, i);
451        }
452        assert_eq!(q.running_parallelism_sum(), 9);
453
454        // Finish out of order
455        assert!(q.finish_running((task_id.into(), 1)));
456        assert_eq!(q.running_parallelism_sum(), 5);
457        assert!(q.finish_running((task_id.into(), 0)));
458        assert!(q.finish_running((task_id.into(), 2)));
459        assert_eq!(q.running_parallelism_sum(), 0);
460    }
461
462    #[test]
463    fn test_cancel_waiting_task_only_removes_waiting_plans() {
464        let mut q = IcebergTaskQueue::new(10, 30);
465        let task_id = 1u64;
466
467        assert_eq!(q.push(mk_meta(task_id, 0, 3), None), PushResult::Added);
468        assert_eq!(q.push(mk_meta(task_id, 1, 4), None), PushResult::Added);
469        assert_eq!(q.push(mk_meta(2, 0, 2), None), PushResult::Added);
470
471        let popped = q.pop().unwrap();
472        assert_eq!(popped.meta.task_id.as_raw_id(), task_id);
473        assert_eq!(popped.meta.plan_index, 0);
474        assert_eq!(q.running_parallelism_sum(), 3);
475        assert_eq!(q.waiting_parallelism_sum(), 6);
476
477        assert_eq!(q.cancel_waiting_task(task_id.into()), 1);
478        assert_eq!(q.running_parallelism_sum(), 3);
479        assert_eq!(q.waiting_parallelism_sum(), 2);
480
481        assert!(q.finish_running((task_id.into(), 0)));
482        let next = q.pop().unwrap();
483        assert_eq!(next.meta.task_id.as_raw_id(), 2);
484        assert_eq!(next.meta.plan_index, 0);
485    }
486
487    #[test]
488    fn test_empty_queue_behavior() {
489        let mut q = IcebergTaskQueue::new(8, 32);
490        assert!(q.pop().is_none());
491        assert!(!q.finish_running((1.into(), 0)));
492        assert_eq!(q.waiting_parallelism_sum(), 0);
493        assert_eq!(q.running_parallelism_sum(), 0);
494    }
495}
risingwave_storage/hummock/compactor/iceberg_compaction/mod.rs

risingwave_storage/hummock/compactor/iceberg_compaction/
mod.rs
