iota_data_ingestion_core/reducer.rs
1// Copyright (c) Mysten Labs, Inc.
2// Modifications Copyright (c) 2025 IOTA Stiftung
3// SPDX-License-Identifier: Apache-2.0
4
5use std::{collections::HashMap, sync::Arc};
6
7use async_trait::async_trait;
8use backoff::{ExponentialBackoff, backoff::Backoff};
9use futures::StreamExt;
10use iota_types::messages_checkpoint::CheckpointSequenceNumber;
11use tokio::sync::mpsc;
12use tokio_stream::wrappers::ReceiverStream;
13use tokio_util::sync::CancellationToken;
14
15use crate::{
16 IngestionError, IngestionResult, MAX_CHECKPOINTS_IN_PROGRESS, Worker, util::reset_backoff,
17 worker_pool::WorkerPoolStatus,
18};
19
20/// Represents the outcome of a commit operation with retry logic.
21pub enum CommitStatus {
22 /// Commit succeeded, continue processing.
23 Success,
24 /// Processing should stop due to shutdown signal.
25 Shutdown,
26}
27
28/// Processes and commits batches of messages produced by workers.
29///
30/// The Reducer trait provides batch processing capabilities for messages
31/// generated by [`Worker`]s. It allows for custom batching logic and efficient
32/// processing of worker results.
33///
34/// # Batch Processing
35///
36/// [`Messages`](Worker::Message) are accumulated in batches based on the
37/// [`should_close_batch`](Reducer::should_close_batch) policy.
38///
39/// When a batch is
40/// ready to be committed (as determined by `should_close_batch`), the
41/// [`commit`](Reducer::commit) method is called with the collected messages.
42#[async_trait]
43pub trait Reducer<Mapper: Worker>: Send + Sync {
44 /// Commits a batch of messages.
45 ///
46 /// This method is called when a batch is ready to be processed, as
47 /// determined by [`should_close_batch`](Reducer::should_close_batch). The
48 /// implementation should handle the batch processing logic and return
49 /// an error if the operation fails.
50 ///
51 /// # Note
52 /// Messages within each batch are guaranteed to be ordered by checkpoint
53 /// sequence number.
54 async fn commit(&self, batch: &[Mapper::Message]) -> Result<(), Mapper::Error>;
55
56 /// Attempts to commit a batch of messages with exponential backoff retries
57 /// on failure.
58 ///
59 /// This function repeatedly calls the [`commit`](Reducer::commit) method of
60 /// the provided [`Reducer`] until either:
61 /// - The commit succeeds, returning `CommitStatus::Success`.
62 /// - A cancellation signal is received via the [`CancellationToken`], returning
63 /// `CommitStatus::Shutdown`.
64 /// - All retry attempts are exhausted within backoff's maximum elapsed time: the pool is
65 /// cancelled and `CommitStatus::Shutdown` is returned.
66 ///
67 /// # Retry Mechanism:
68 /// - Uses [`ExponentialBackoff`](backoff::ExponentialBackoff) to introduce increasing delays
69 /// between retry attempts.
70 /// - Checks for cancellation both before and after each commit attempt.
71 /// - If a cancellation signal is received during a backoff delay, the function exits
72 /// immediately with `CommitStatus::Shutdown`.
73 async fn commit_with_retry(
74 &self,
75 batch: &[Mapper::Message],
76 mut backoff: ExponentialBackoff,
77 token: &CancellationToken,
78 ) -> CommitStatus {
79 loop {
80 // check for cancellation before attempting commit.
81 if token.is_cancelled() {
82 return CommitStatus::Shutdown;
83 }
84 // attempt to commit.
85 match self.commit(batch).await {
86 Ok(_) => return CommitStatus::Success,
87 Err(err) => {
88 let err = IngestionError::Reducer(format!("failed to commit batch: {err}"));
89 tracing::warn!("transient reducer commit error {err:?}");
90 // check for cancellation after failed commit.
91 if token.is_cancelled() {
92 return CommitStatus::Shutdown;
93 }
94 }
95 }
96 // Retrying the same batch any further cannot help, and a panic
97 // would abort the whole process in release builds, so stop the pool.
98 let Some(duration) = backoff.next_backoff() else {
99 tracing::error!(
100 "reducer commit failed and the retries are exhausted, shutting down"
101 );
102 token.cancel();
103 return CommitStatus::Shutdown;
104 };
105 // if cancellation occurs during backoff wait, exit early with Shutdown.
106 // Otherwise (if timeout expires), continue with the next retry attempt
107 if tokio::time::timeout(duration, token.cancelled())
108 .await
109 .is_ok()
110 {
111 return CommitStatus::Shutdown;
112 }
113 }
114 }
115
116 /// Determines if the current batch should be closed and committed.
117 ///
118 /// This method is called by the framework in two situations:
119 ///
120 /// - **`next_item = Some(msg)`**: a new message is about to be added to the batch. Returning
121 /// `true` triggers a commit of the existing batch first. The `msg` is added to a new batch.
122 /// Returning `false` appends `msg` to the current batch without committing.
123 ///
124 /// - **`next_item = None`**: the current chunk from the upstream stream has been fully consumed
125 /// and there are no more messages to add in this iteration (the stream itself may still
126 /// produce more chunks later). Returning `true` flushes the partial batch. Returning `false`
127 /// leaves the batch in memory until the next chunk arrives.
128 ///
129 /// # Default Implementation
130 ///
131 /// Returns `true` only when `next_item` is `None`, so each chunk's
132 /// remaining batch is flushed at the chunk boundary.
133 fn should_close_batch(
134 &self,
135 _batch: &[Mapper::Message],
136 next_item: Option<&Mapper::Message>,
137 ) -> bool {
138 next_item.is_none()
139 }
140}
141
142/// Processes worker messages and manages batching through a reducer.
143///
144/// This function is the core of the reduction pipeline, handling message
145/// batching, watermark tracking, and progress reporting. It maintains message
146/// order by checkpoint sequence number and applies batching logic through the
147/// provided reducer.
148///
149/// # Message Processing Flow
150///
151/// 1. Receives messages in chunks up to [`MAX_CHECKPOINTS_IN_PROGRESS`].
152/// 2. Maintains message order using checkpoint sequence numbers.
153/// 3. Batches messages according to reducer's [`Reducer::should_close_batch`] policy and after that
154/// its committed.
155/// 4. Progress is updated after each commit.
156/// 5. Reports progress back to the executor.
157///
158/// # Shutdown Behavior
159///
160/// The function will gracefully exit when receiving a shutdown signal,
161/// ensuring no data loss for processed messages.
162pub(crate) async fn reduce<W: Worker>(
163 task_name: String,
164 mut current_checkpoint_number: CheckpointSequenceNumber,
165 watermark_receiver: mpsc::Receiver<(CheckpointSequenceNumber, Option<W::Message>)>,
166 executor_progress_sender: mpsc::Sender<WorkerPoolStatus>,
167 reducer: Box<dyn Reducer<W>>,
168 backoff: Arc<ExponentialBackoff>,
169 token: CancellationToken,
170) -> IngestionResult<()> {
171 // convert to a stream of MAX_CHECKPOINTS_IN_PROGRESS size. This way, each
172 // iteration of the loop will process all ready messages.
173 let mut stream =
174 ReceiverStream::new(watermark_receiver).ready_chunks(MAX_CHECKPOINTS_IN_PROGRESS);
175 // store unprocessed progress messages from workers.
176 let mut unprocessed = HashMap::new();
177 // buffer to accumulate results before passing them to the reducer.
178 // The size of this batch is dynamically determined by the reducer's
179 // `should_close_batch` method.
180 let mut batch = vec![];
181 // track the next unprocessed checkpoint number for reporting progress
182 // after each chunk of messages is received from the stream.
183 let mut progress_update = None;
184 // flag to indicate a shutdown has been triggered.
185 let mut trigger_shutdown = false;
186
187 while let Some(update_batch) = stream.next().await {
188 unprocessed.extend(update_batch);
189 // Process messages sequentially based on checkpoint sequence number.
190 // This ensures in-order processing and maintains progress integrity.
191 while let Some(maybe_worker_message) = unprocessed.remove(¤t_checkpoint_number) {
192 let Some(message) = maybe_worker_message else {
193 current_checkpoint_number += 1;
194 if batch.is_empty() {
195 progress_update = Some(current_checkpoint_number);
196 }
197 continue;
198 };
199
200 // reducer is configured, collect messages in batch based on
201 // `reducer.should_close_batch` policy, once a batch is collected it gets
202 // committed and a new batch is created with the current message.
203 if reducer.should_close_batch(&batch, Some(&message)) {
204 match reducer
205 .commit_with_retry(&std::mem::take(&mut batch), reset_backoff(&backoff), &token)
206 .await
207 {
208 CommitStatus::Success => {
209 batch = vec![message];
210 progress_update = Some(current_checkpoint_number);
211 }
212 CommitStatus::Shutdown => {
213 trigger_shutdown = true;
214 break;
215 }
216 };
217 } else {
218 // Add message to existing batch since no commit needed.
219 batch.push(message);
220 }
221 current_checkpoint_number += 1;
222 }
223 // Handle final batch processing.
224 // Check if the final batch should be committed.
225 // None parameter indicates no more messages available.
226 // Only check if batch is non-empty to avoid unnecessary calls.
227 if !batch.is_empty() && reducer.should_close_batch(&batch, None) && !trigger_shutdown {
228 match reducer
229 .commit_with_retry(&std::mem::take(&mut batch), reset_backoff(&backoff), &token)
230 .await
231 {
232 CommitStatus::Success => {
233 progress_update = Some(current_checkpoint_number);
234 }
235 CommitStatus::Shutdown => trigger_shutdown = true,
236 }
237 }
238 // report progress update to executor.
239 if let Some(watermark) = progress_update {
240 executor_progress_sender
241 .send(WorkerPoolStatus::Running((task_name.clone(), watermark)))
242 .await
243 .map_err(|_| IngestionError::Channel("unable to send worker pool progress updates to executor, receiver half closed".into()))?;
244 progress_update = None;
245 }
246
247 // Check for shutdown signal after progress update to ensure progress
248 // is reported even during shutdown.
249 if trigger_shutdown {
250 break;
251 }
252 }
253 Ok(())
254}