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

risingwave_storage/hummock/compactor/iceberg_compaction/
mod.rs
