risingwave_storage/hummock/compactor/iceberg_compaction/
report.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::VecDeque;
16use std::time::SystemTime;
17
18use risingwave_pb::iceberg_compaction::{
19    SubscribeIcebergCompactionEventRequest, subscribe_iceberg_compaction_event_request,
20};
21use thiserror_ext::AsReport;
22use tokio::sync::mpsc;
23
24use super::TaskKey;
25
26#[derive(Debug)]
27pub(crate) struct IcebergPlanCompletion {
28    pub(crate) task_key: TaskKey,
29    pub(crate) error_message: Option<String>,
30}
31
32pub(crate) type IcebergTaskReport = subscribe_iceberg_compaction_event_request::ReportTask;
33
34pub(crate) enum ReportSendResult {
35    Sent,
36    RestartStream,
37}
38
39pub(crate) struct IcebergTaskTracker {
40    sink_id: u32,
41    total_plans: usize,
42    remaining_plans: usize,
43    successful_plans: usize,
44    failed_plans: usize,
45    first_error: Option<String>,
46}
47
48impl IcebergTaskTracker {
49    pub(crate) fn new(sink_id: u32, remaining_plans: usize) -> Self {
50        Self {
51            sink_id,
52            total_plans: remaining_plans,
53            remaining_plans,
54            successful_plans: 0,
55            failed_plans: 0,
56            first_error: None,
57        }
58    }
59
60    pub(crate) fn record_completion(&mut self, error_message: Option<String>) {
61        debug_assert!(self.remaining_plans > 0);
62        self.remaining_plans -= 1;
63        if let Some(error_message) = error_message {
64            self.failed_plans += 1;
65            if self.first_error.is_none() {
66                self.first_error = Some(error_message);
67            }
68        } else {
69            self.successful_plans += 1;
70        }
71    }
72
73    pub(crate) fn is_finished(&self) -> bool {
74        self.remaining_plans == 0
75    }
76
77    pub(crate) fn sink_id(&self) -> u32 {
78        self.sink_id
79    }
80
81    pub(crate) fn total_plans(&self) -> usize {
82        self.total_plans
83    }
84
85    pub(crate) fn successful_plans(&self) -> usize {
86        self.successful_plans
87    }
88
89    pub(crate) fn failed_plans(&self) -> usize {
90        self.failed_plans
91    }
92
93    pub(crate) fn into_report(self, task_id: u64) -> IcebergTaskReport {
94        let error_message = if self.successful_plans > 0 {
95            None
96        } else {
97            Some(
98                self.first_error
99                    .unwrap_or_else(|| "All admitted iceberg compaction plans failed".to_owned()),
100            )
101        };
102        build_iceberg_task_report(task_id, self.sink_id, error_message)
103    }
104}
105
106pub(crate) fn build_iceberg_task_report(
107    task_id: u64,
108    sink_id: u32,
109    error_message: Option<String>,
110) -> IcebergTaskReport {
111    subscribe_iceberg_compaction_event_request::ReportTask {
112        task_id,
113        sink_id,
114        status: if error_message.is_some() {
115            subscribe_iceberg_compaction_event_request::report_task::Status::Failed as i32
116        } else {
117            subscribe_iceberg_compaction_event_request::report_task::Status::Success as i32
118        },
119        error_message,
120    }
121}
122
123pub(crate) fn send_iceberg_task_report(
124    request_sender: &mpsc::UnboundedSender<SubscribeIcebergCompactionEventRequest>,
125    report_event: IcebergTaskReport,
126) -> Result<(), IcebergTaskReport> {
127    if let Err(e) = request_sender.send(SubscribeIcebergCompactionEventRequest {
128        event: Some(
129            subscribe_iceberg_compaction_event_request::Event::ReportTask(report_event.clone()),
130        ),
131        create_at: SystemTime::now()
132            .duration_since(std::time::UNIX_EPOCH)
133            .expect("Clock may have gone backwards")
134            .as_millis() as u64,
135    }) {
136        tracing::warn!(
137            iceberg_component = "compaction_worker",
138            iceberg_operation = "report_task",
139            error = %e.as_report(),
140            task_id = report_event.task_id,
141            sink_id = report_event.sink_id,
142            "iceberg_compaction_task_report_send_failed",
143        );
144        return Err(report_event);
145    }
146
147    Ok(())
148}
149
150pub(crate) fn send_or_buffer_iceberg_task_report(
151    request_sender: &mpsc::UnboundedSender<SubscribeIcebergCompactionEventRequest>,
152    pending_task_reports: &mut VecDeque<IcebergTaskReport>,
153    report: IcebergTaskReport,
154) -> ReportSendResult {
155    if let Err(report) = send_iceberg_task_report(request_sender, report) {
156        pending_task_reports.push_back(report);
157        return ReportSendResult::RestartStream;
158    }
159    ReportSendResult::Sent
160}
161
162pub(crate) fn flush_pending_iceberg_task_reports(
163    request_sender: &mpsc::UnboundedSender<SubscribeIcebergCompactionEventRequest>,
164    pending_task_reports: &mut VecDeque<IcebergTaskReport>,
165) -> ReportSendResult {
166    while let Some(report_event) = pending_task_reports.pop_front() {
167        if let Err(report_event) = send_iceberg_task_report(request_sender, report_event) {
168            pending_task_reports.push_front(report_event);
169            return ReportSendResult::RestartStream;
170        }
171    }
172    ReportSendResult::Sent
173}
174
175#[cfg(test)]
176mod tests {
177    use risingwave_pb::iceberg_compaction::subscribe_iceberg_compaction_event_request;
178
179    use super::*;
180
181    #[test]
182    fn test_send_iceberg_task_report_returns_payload_on_send_failure() {
183        let (tx, rx) = tokio::sync::mpsc::unbounded_channel();
184        drop(rx);
185
186        let report = build_iceberg_task_report(7, 9, Some("send failure".to_owned()));
187        let failed_report = send_iceberg_task_report(&tx, report.clone()).unwrap_err();
188
189        assert_eq!(failed_report.task_id, report.task_id);
190        assert_eq!(failed_report.sink_id, report.sink_id);
191        assert_eq!(failed_report.error_message, report.error_message);
192    }
193
194    #[test]
195    fn test_build_iceberg_task_result_partial_enqueue_is_success_if_admitted_plan_succeeds() {
196        let mut tracker = IcebergTaskTracker::new(9, 1);
197        tracker.record_completion(None);
198
199        let report = tracker.into_report(7);
200
201        assert_eq!(
202            report.status,
203            subscribe_iceberg_compaction_event_request::report_task::Status::Success as i32
204        );
205        assert!(report.error_message.is_none());
206    }
207
208    #[test]
209    fn test_build_iceberg_task_result_fails_if_all_admitted_plans_fail() {
210        let mut tracker = IcebergTaskTracker::new(9, 2);
211        tracker.record_completion(Some("first failure".to_owned()));
212        tracker.record_completion(Some("second failure".to_owned()));
213
214        let report = tracker.into_report(7);
215
216        assert_eq!(
217            report.status,
218            subscribe_iceberg_compaction_event_request::report_task::Status::Failed as i32
219        );
220        assert_eq!(report.error_message.as_deref(), Some("first failure"));
221    }
222}
risingwave_storage/hummock/compactor/iceberg_compaction/report.rs

risingwave_storage/hummock/compactor/iceberg_compaction/
report.rs
