iota_data_ingestion_core/reader/
v2.rs

1// Copyright (c) 2025 IOTA Stiftung
2// SPDX-License-Identifier: Apache-2.0
3
4use std::{
5    num::NonZeroUsize,
6    path::{Path, PathBuf},
7    sync::Arc,
8    time::Duration,
9};
10
11use backoff::backoff::Backoff;
12use futures::{StreamExt, TryStreamExt};
13use iota_config::object_storage_config::{ObjectStoreConfig, ObjectStoreType};
14use iota_grpc_client::Client as GrpcClient;
15use iota_metrics::spawn_monitored_task;
16use iota_types::{
17    full_checkpoint_content::CheckpointData, messages_checkpoint::CheckpointSequenceNumber,
18};
19use object_store::ObjectStore;
20use serde::{Deserialize, Serialize};
21use tap::Pipe;
22use tokio::{
23    sync::mpsc::{self},
24    task::JoinHandle,
25    time::timeout,
26};
27use tokio_util::sync::CancellationToken;
28use tracing::{debug, error, info};
29
30#[cfg(not(target_os = "macos"))]
31use crate::reader::fetch::init_watcher;
32use crate::{
33    IngestionError, IngestionResult, MAX_CHECKPOINTS_IN_PROGRESS,
34    config::CheckpointReaderConfigExt,
35    create_remote_store_client,
36    history::reader::{HistoricalReader, HistoricalReaderConfig},
37    reader::{
38        ReaderOptions,
39        common::DataLimiter,
40        fetch::{
41            GRPC_MAX_DECODING_MESSAGE_SIZE_BYTES, LocalRead, ReadSource, fetch_from_object_store,
42        },
43        filters::fullnode::TransactionFilter,
44    },
45};
46
47/// Available sources for checkpoint streams supported by the ingestion
48/// framework.
49///
50/// This enum represents the different types of remote sources from which
51/// checkpoint data can be fetched. Each variant corresponds to a supported
52/// backend or combination of backends for checkpoint retrieval.
53#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)]
54pub enum RemoteUrl {
55    /// The URL to the Fullnode server that exposes
56    /// checkpoint data streaming through gRPC.
57    ///
58    /// # Example
59    /// ```text
60    /// "http://127.0.0.1:50051"
61    /// ```
62    Fullnode(String),
63    /// A hybrid source combining historical object store and optional live
64    /// object store.
65    HybridHistoricalStore {
66        /// The URL path to the historical object store that contains `*.chk`,
67        /// `*.sum` & `MANIFEST` files.
68        ///
69        /// # Example
70        /// ```text
71        /// "https://checkpoints.mainnet.iota.cafe/ingestion/historical"
72        /// ```
73        historical_url: String,
74        /// The URL path to the live object store that contains `*.chk`
75        /// checkpoint files.
76        ///
77        /// # Example
78        /// ```text
79        /// "https://checkpoints.mainnet.iota.cafe/ingestion/live"
80        /// ```
81        live_url: Option<String>,
82    },
83}
84
85/// Represents a remote backend for checkpoint data retrieval.
86///
87/// This enum encapsulates the supported remote storage mechanisms that can be
88/// used by the ingestion framework to fetch checkpoint data. Each variant
89/// corresponds to a different type of remote source.
90enum RemoteStore {
91    Fullnode(GrpcClient),
92    HybridHistoricalStore {
93        historical: HistoricalReader,
94        live: Option<Box<dyn ObjectStore>>,
95    },
96}
97
98impl RemoteStore {
99    async fn new(
100        remote_url: RemoteUrl,
101        batch_size: usize,
102        timeout_secs: u64,
103    ) -> IngestionResult<Self> {
104        let store = match remote_url {
105            RemoteUrl::Fullnode(ref url) => {
106                let grpc_client = GrpcClient::new(url).map(|client| {
107                    client.with_max_decoding_message_size(GRPC_MAX_DECODING_MESSAGE_SIZE_BYTES)
108                })?;
109                RemoteStore::Fullnode(grpc_client)
110            }
111            RemoteUrl::HybridHistoricalStore {
112                historical_url,
113                live_url,
114            } => {
115                let remote_store_config = if let Some(dir) = historical_url.strip_prefix("file://")
116                {
117                    ObjectStoreConfig {
118                        object_store: Some(ObjectStoreType::File),
119                        directory: Some(PathBuf::from(dir)),
120                        ..Default::default()
121                    }
122                } else {
123                    ObjectStoreConfig {
124                        object_store: Some(ObjectStoreType::S3),
125                        object_store_connection_limit: 20,
126                        aws_endpoint: Some(historical_url),
127                        aws_virtual_hosted_style_request: true,
128                        no_sign_request: true,
129                        ..Default::default()
130                    }
131                };
132                let config = HistoricalReaderConfig {
133                    download_concurrency: NonZeroUsize::new(batch_size)
134                        .expect("batch size must be greater than zero"),
135                    remote_store_config,
136                };
137                let historical = HistoricalReader::new(config)
138                    .inspect_err(|e| error!("unable to instantiate historical reader: {e}"))?;
139
140                let live = live_url
141                    .map(|url| create_remote_store_client(url, Default::default(), timeout_secs))
142                    .transpose()?;
143
144                RemoteStore::HybridHistoricalStore { historical, live }
145            }
146        };
147        Ok(store)
148    }
149}
150
151/// Configuration options to control the behavior of a checkpoint
152/// reader.
153#[derive(Default, Clone)]
154pub struct CheckpointReaderConfig {
155    /// Config the checkpoint reader behavior for downloading new checkpoints.
156    pub reader_options: ReaderOptions,
157    /// Local path for checkpoint ingestion. If not provided, checkpoints will
158    /// be ingested from a temporary directory.
159    pub ingestion_path: Option<PathBuf>,
160    /// Remote source for checkpoint data stream.
161    pub remote_store_url: Option<RemoteUrl>,
162}
163
164/// Internal actor responsible for reading and streaming checkpoints.
165///
166/// `CheckpointReaderActor` is the core background task that manages the logic
167/// for fetching, batching, and streaming checkpoint data from local or remote
168/// sources. It handles checkpoint discovery, garbage collection signals, and
169/// coordinates with remote fetchers as needed.
170///
171/// This struct is intended to be run as an asynchronous task and is not
172/// typically interacted with directly. Instead, users should use
173/// [`CheckpointReader`], which provides a safe and ergonomic API for
174/// interacting with the running actor, such as receiving checkpoints, sending
175/// GC signals, or triggering shutdown.
176///
177/// # Responsibilities
178/// - Periodically scans for new checkpoints from configured sources.
179/// - Streams checkpoints to consumers via channels.
180/// - Handles garbage collection signals to prune processed checkpoints.
181/// - Coordinates with remote fetchers for batch downloads and retries.
182///
183/// # Usage
184/// Users should not construct or manage `CheckpointReader` directly. Instead,
185/// use [`CheckpointReader::new`] to spawn the actor and obtain a handle
186/// for interaction.
187struct CheckpointReaderActor {
188    /// Filesystem path to the local checkpoint directory.
189    path: PathBuf,
190    /// Start fetch from the current checkpoint sequence.
191    current_checkpoint_number: CheckpointSequenceNumber,
192    /// Keeps tracks the last processed checkpoint sequence number, used to
193    /// delete checkpoint files from ingestion path.
194    last_pruned_watermark: CheckpointSequenceNumber,
195    /// Channel for sending checkpoints to WorkerPools.
196    checkpoint_tx: mpsc::Sender<Arc<CheckpointData>>,
197    /// Sends a garbage collection (GC) signal to prune checkpoint files below
198    /// the specified watermark.
199    gc_signal_rx: mpsc::Receiver<CheckpointSequenceNumber>,
200    /// Remote checkpoint reader for fetching checkpoints from the network.
201    remote_store: Option<Arc<RemoteStore>>,
202    /// Shutdown signal for the actor.
203    token: CancellationToken,
204    /// Configures the behavior of the checkpoint reader.
205    reader_options: ReaderOptions,
206    /// Limit the amount of downloaded checkpoints held in memory to avoid OOM.
207    data_limiter: DataLimiter,
208    /// Filter applied to transactions within a checkpoint.
209    fullnode_transaction_filter: Option<TransactionFilter>,
210}
211
212impl LocalRead for CheckpointReaderActor {
213    fn exceeds_capacity(&self, checkpoint_number: CheckpointSequenceNumber) -> bool {
214        ((MAX_CHECKPOINTS_IN_PROGRESS as u64 + self.last_pruned_watermark) <= checkpoint_number)
215            || self.data_limiter.exceeds()
216    }
217
218    fn path(&self) -> &Path {
219        &self.path
220    }
221
222    fn current_checkpoint_number(&self) -> CheckpointSequenceNumber {
223        self.current_checkpoint_number
224    }
225
226    fn update_last_pruned_watermark(&mut self, watermark: CheckpointSequenceNumber) {
227        self.last_pruned_watermark = watermark;
228    }
229}
230
231impl CheckpointReaderActor {
232    fn should_fetch_from_remote(&self, checkpoints: &[Arc<CheckpointData>]) -> bool {
233        self.remote_store.is_some()
234            && (checkpoints.is_empty()
235                || self.is_checkpoint_ahead(&checkpoints[0], self.current_checkpoint_number))
236    }
237
238    /// Fetches checkpoints from the historical object store and streams them to
239    /// a channel.
240    async fn relay_from_historical(
241        &mut self,
242        historical_reader: &HistoricalReader,
243    ) -> IngestionResult<()> {
244        // Only sync the manifest when needed to avoid unnecessary network calls.
245        // If the requested checkpoint is beyond what's currently available in our
246        // cached manifest, we need to refresh it to check for newer checkpoints.
247        if self.current_checkpoint_number > historical_reader.latest_available_checkpoint().await? {
248            timeout(
249                Duration::from_secs(self.reader_options.timeout_secs),
250                historical_reader.sync_manifest_once(),
251            )
252            .await
253            .map_err(|_| {
254                IngestionError::HistoryRead("reading manifest exceeded the timeout".into())
255            })??;
256
257            // Verify the requested checkpoint is now available after the manifest refresh.
258            // If it's still not available, the checkpoint hasn't been published yet.
259            if self.current_checkpoint_number
260                > historical_reader.latest_available_checkpoint().await?
261            {
262                return Err(IngestionError::CheckpointNotAvailableYet);
263            }
264        }
265
266        let manifest = historical_reader.get_manifest().await;
267
268        let files = historical_reader.verify_and_get_manifest_files(manifest)?;
269
270        let start_index = match files.binary_search_by_key(&self.current_checkpoint_number, |s| {
271            s.checkpoint_seq_range.start
272        }) {
273            Ok(index) => index,
274            Err(index) => index - 1,
275        };
276
277        for metadata in files
278            .into_iter()
279            .enumerate()
280            .filter_map(|(index, metadata)| (index >= start_index).then_some(metadata))
281        {
282            let checkpoints = timeout(
283                Duration::from_secs(self.reader_options.timeout_secs),
284                historical_reader.iter_for_file(metadata.file_path()),
285            )
286            .await
287            .map_err(|_| {
288                IngestionError::HistoryRead(format!(
289                    "reading checkpoint {} exceeded the timeout",
290                    metadata.file_path()
291                ))
292            })??
293            .filter(|c| c.checkpoint_summary.sequence_number >= self.current_checkpoint_number)
294            .collect::<Vec<CheckpointData>>();
295
296            for checkpoint in checkpoints {
297                let size = bcs::serialized_size(&checkpoint)?;
298                self.send_remote_checkpoint_with_capacity_check(Arc::new(checkpoint), size)
299                    .await?;
300            }
301        }
302
303        Ok(())
304    }
305
306    /// Fetches checkpoints from the live object store and streams them to a
307    /// channel.
308    async fn relay_from_live(
309        &mut self,
310        batch_size: usize,
311        live: &dyn ObjectStore,
312    ) -> IngestionResult<()> {
313        let mut checkpoint_stream = (self.current_checkpoint_number..u64::MAX)
314            .map(|checkpoint_number| fetch_from_object_store(live, checkpoint_number))
315            .pipe(futures::stream::iter)
316            .buffered(batch_size);
317        while let Some((checkpoint, size)) = self
318            .token
319            .run_until_cancelled(checkpoint_stream.try_next())
320            .await
321            .transpose()?
322            .flatten()
323        {
324            self.send_remote_checkpoint_with_capacity_check(checkpoint, size)
325                .await?;
326        }
327        Ok(())
328    }
329
330    /// Fetches checkpoints from the fullnode through a gRPC streaming
331    /// connection and streams them to a channel.
332    async fn relay_from_fullnode(&mut self, client: &mut GrpcClient) -> IngestionResult<()> {
333        let mut checkpoints_stream = client
334            .stream_checkpoints(
335                Some(self.current_checkpoint_number),
336                None,
337                Some(iota_grpc_client::CHECKPOINT_RESPONSE_CHECKPOINT_DATA.into()),
338                self.fullnode_transaction_filter.clone().map(Into::into),
339                None,
340            )
341            .await
342            .map_err(|e| {
343                IngestionError::Grpc(format!("failed to initialize the checkpoint stream: {e}"))
344            })?
345            .into_inner();
346
347        while let Some(grpc_checkpoint) = self
348            .token
349            .run_until_cancelled(checkpoints_stream.try_next())
350            .await
351            .transpose()?
352            .flatten()
353        {
354            let checkpoint = grpc_checkpoint.checkpoint_data()?.try_into()?;
355            let size = bcs::serialized_size(&checkpoint)?;
356            self.send_remote_checkpoint_with_capacity_check(Arc::new(checkpoint), size)
357                .await?;
358        }
359
360        Ok(())
361    }
362
363    /// Fetches remote checkpoints from the remote store and streams them to the
364    /// channel.
365    ///
366    /// For every successfully fetched checkpoint, this function updates the
367    /// current checkpoint number and the data limiter. If an error occurs while
368    /// fetching a checkpoint, the function returns immediately with that error.
369    async fn fetch_and_send_to_channel(&mut self) -> IngestionResult<()> {
370        let Some(remote_store) = self.remote_store.as_ref().map(Arc::clone) else {
371            return Ok(());
372        };
373        let batch_size = self.reader_options.batch_size;
374        match remote_store.as_ref() {
375            RemoteStore::Fullnode(client) => {
376                self.relay_from_fullnode(&mut client.clone()).await?;
377            }
378            RemoteStore::HybridHistoricalStore { historical, live } => {
379                if let Some(Err(err)) = self
380                    .token
381                    .clone()
382                    .run_until_cancelled(self.relay_from_historical(historical))
383                    .await
384                {
385                    if matches!(err, IngestionError::CheckpointNotAvailableYet) {
386                        let live = live.as_ref().ok_or(err)?;
387                        return self.relay_from_live(batch_size, live).await;
388                    }
389                    return Err(err);
390                }
391            }
392        };
393        Ok(())
394    }
395
396    /// Fetches and sends checkpoints to the channel with retry logic.
397    ///
398    /// Uses an exponential backoff strategy to retry failed requests.
399    async fn fetch_and_send_to_channel_with_retry(&mut self) {
400        let mut backoff = backoff::ExponentialBackoff::default();
401        backoff.max_elapsed_time = Some(Duration::from_secs(60));
402        backoff.initial_interval = Duration::from_millis(100);
403        backoff.current_interval = backoff.initial_interval;
404        backoff.multiplier = 1.0;
405
406        loop {
407            match self.fetch_and_send_to_channel().await {
408                Ok(_) => break,
409                Err(IngestionError::MaxCheckpointsCapacityReached) => break,
410                Err(IngestionError::CheckpointNotAvailableYet) => {
411                    break info!("historical reader does not have the requested checkpoint yet");
412                }
413                Err(err) => match backoff.next_backoff() {
414                    Some(duration) => {
415                        if !err.to_string().to_lowercase().contains("not found") {
416                            debug!(
417                                "remote reader retry in {} ms. Error is {err:?}",
418                                duration.as_millis(),
419                            );
420                        }
421                        if self
422                            .token
423                            .run_until_cancelled(tokio::time::sleep(duration))
424                            .await
425                            .is_none()
426                        {
427                            break;
428                        }
429                    }
430                    None => {
431                        break error!("remote reader transient error {err:?}");
432                    }
433                },
434            }
435        }
436    }
437
438    /// Attempts to send a checkpoint from remote source to the channel if
439    /// capacity allows.
440    ///
441    /// If the checkpoint's sequence number would exceed the allowed capacity,
442    /// returns `IngestionError::MaxCheckpointsCapacityReached` and does not
443    /// send. Otherwise, adds the checkpoint to the data limiter and sends
444    /// it to the channel.
445    async fn send_remote_checkpoint_with_capacity_check(
446        &mut self,
447        checkpoint: Arc<CheckpointData>,
448        size: usize,
449    ) -> IngestionResult<()> {
450        if self.exceeds_capacity(checkpoint.checkpoint_summary.sequence_number) {
451            return Err(IngestionError::MaxCheckpointsCapacityReached);
452        }
453        self.data_limiter.add(&checkpoint, size);
454        self.send_checkpoint_to_channel(checkpoint).await
455    }
456
457    /// Sends a batch of local checkpoints to the channel in order.
458    ///
459    /// Each checkpoint is sent sequentially until a gap is detected (i.e., a
460    /// checkpoint with a sequence number greater than the current
461    /// checkpoint number). If a gap is found, the function breaks early. If
462    /// sending fails, returns the error immediately.
463    async fn send_local_checkpoints_to_channel(
464        &mut self,
465        checkpoints: Vec<Arc<CheckpointData>>,
466    ) -> IngestionResult<()> {
467        for checkpoint in checkpoints {
468            if self.is_checkpoint_ahead(&checkpoint, self.current_checkpoint_number) {
469                break;
470            }
471            self.send_checkpoint_to_channel(checkpoint).await?;
472        }
473        Ok(())
474    }
475
476    /// Sends a single checkpoint to the channel and advances the current
477    /// checkpoint number.
478    ///
479    /// Asserts that the checkpoint's sequence number matches the expected
480    /// current number. Increments the current checkpoint number after
481    /// sending.
482    async fn send_checkpoint_to_channel(
483        &mut self,
484        checkpoint: Arc<CheckpointData>,
485    ) -> IngestionResult<()> {
486        assert_eq!(
487            checkpoint.checkpoint_summary.sequence_number,
488            self.current_checkpoint_number
489        );
490        self.checkpoint_tx.send(checkpoint).await.map_err(|_| {
491            IngestionError::Channel(
492                "unable to send checkpoint to executor, receiver half closed".to_owned(),
493            )
494        })?;
495        self.current_checkpoint_number += 1;
496        Ok(())
497    }
498
499    /// Sync from either local or remote source new checkpoints to be processed
500    /// by the executor.
501    async fn sync(&mut self) -> IngestionResult<()> {
502        let mut remote_source = ReadSource::Local;
503        let checkpoints = self.read_local_files_with_retry().await?;
504        let should_fetch_from_remote = self.should_fetch_from_remote(&checkpoints);
505
506        if should_fetch_from_remote {
507            remote_source = ReadSource::Remote;
508            self.fetch_and_send_to_channel_with_retry().await;
509        } else {
510            self.send_local_checkpoints_to_channel(checkpoints).await?;
511        }
512
513        info!(
514            "Read from {remote_source}. Current checkpoint number: {}, pruning watermark: {}",
515            self.current_checkpoint_number, self.last_pruned_watermark,
516        );
517
518        Ok(())
519    }
520
521    /// Run the main loop of the checkpoint reader actor.
522    async fn run(mut self) {
523        let (_inotify_tx, mut inotify_rx) = mpsc::channel::<()>(1);
524        std::fs::create_dir_all(self.path()).expect("failed to create a directory");
525
526        #[cfg(not(target_os = "macos"))]
527        let _watcher = init_watcher(_inotify_tx, self.path());
528
529        self.data_limiter.gc(self.last_pruned_watermark);
530        self.gc_processed_files(self.last_pruned_watermark)
531            .expect("failed to clean the directory");
532
533        loop {
534            tokio::select! {
535                _ = self.token.cancelled() => break,
536                Some(watermark) = self.gc_signal_rx.recv() => {
537                    self.data_limiter.gc(watermark);
538                    self.gc_processed_files(watermark).expect("failed to clean the directory");
539                }
540                Ok(Some(_)) | Err(_) = timeout(Duration::from_millis(self.reader_options.tick_interval_ms), inotify_rx.recv())  => {
541                    self.sync().await.expect("failed to read checkpoint files");
542                }
543            }
544        }
545    }
546}
547
548/// Public API for interacting with the checkpoint reader actor.
549///
550/// It provides methods to receive streamed checkpoints, send garbage collection
551/// signals, and gracefully shut down the background checkpoint reading task.
552/// Internally, it communicates with a [`CheckpointReaderActor`], which manages
553/// the actual checkpoint fetching and streaming logic.
554pub(crate) struct CheckpointReader {
555    handle: JoinHandle<()>,
556    gc_signal_tx: mpsc::Sender<CheckpointSequenceNumber>,
557    checkpoint_rx: mpsc::Receiver<Arc<CheckpointData>>,
558    token: CancellationToken,
559}
560
561impl CheckpointReader {
562    pub(crate) async fn new(
563        starting_checkpoint_number: CheckpointSequenceNumber,
564        config: CheckpointReaderConfigExt,
565    ) -> IngestionResult<Self> {
566        if config.fullnode_transaction_filter.is_some()
567            && !matches!(config.base.remote_store_url, Some(RemoteUrl::Fullnode(_)))
568        {
569            return Err(IngestionError::Unsupported(
570                "filter is only supported on `RemoteUrl::Fullnode` connections".into(),
571            ));
572        }
573
574        let (checkpoint_tx, checkpoint_rx) = mpsc::channel(MAX_CHECKPOINTS_IN_PROGRESS);
575        let (gc_signal_tx, gc_signal_rx) = mpsc::channel(MAX_CHECKPOINTS_IN_PROGRESS);
576
577        let remote_store = if let Some(url) = config.base.remote_store_url {
578            Some(Arc::new(
579                RemoteStore::new(
580                    url,
581                    config.base.reader_options.batch_size,
582                    config.base.reader_options.timeout_secs,
583                )
584                .await?,
585            ))
586        } else {
587            None
588        };
589
590        let path = match config.base.ingestion_path {
591            Some(p) => p,
592            None => tempfile::tempdir()?.keep(),
593        };
594        let token = CancellationToken::new();
595        let reader = CheckpointReaderActor {
596            path,
597            current_checkpoint_number: starting_checkpoint_number,
598            last_pruned_watermark: starting_checkpoint_number,
599            checkpoint_tx,
600            gc_signal_rx,
601            remote_store,
602            token: token.clone(),
603            data_limiter: DataLimiter::new(config.base.reader_options.data_limit),
604            reader_options: config.base.reader_options,
605            fullnode_transaction_filter: config.fullnode_transaction_filter,
606        };
607
608        let handle = spawn_monitored_task!(reader.run());
609
610        Ok(Self {
611            handle,
612            gc_signal_tx,
613            checkpoint_rx,
614            token,
615        })
616    }
617
618    /// Read downloaded checkpoints from the queue.
619    pub(crate) async fn checkpoint(&mut self) -> Option<Arc<CheckpointData>> {
620        self.checkpoint_rx.recv().await
621    }
622
623    /// Sends a garbage collection (GC) signal to the checkpoint reader.
624    ///
625    /// Transmits a watermark to the checkpoint reader, indicating that all
626    /// checkpoints below this watermark can be safely pruned or cleaned up.
627    /// The signal is sent over an internal channel to the checkpoint reader
628    /// task.
629    pub(crate) async fn send_gc_signal(
630        &self,
631        watermark: CheckpointSequenceNumber,
632    ) -> IngestionResult<()> {
633        self.gc_signal_tx.send(watermark).await.map_err(|_| {
634            IngestionError::Channel(
635                "unable to send GC operation to checkpoint reader, receiver half closed".into(),
636            )
637        })
638    }
639
640    /// Gracefully shuts down the checkpoint reader task.
641    ///
642    /// It signals the background checkpoint reader actor to terminate, then
643    /// awaits its completion. Any in-progress checkpoint reading or streaming
644    /// operations will be stopped as part of the shutdown process.
645    pub(crate) async fn shutdown(self) -> IngestionResult<()> {
646        self.token.cancel();
647        self.handle.await.map_err(|err| IngestionError::Shutdown {
648            component: "checkpoint reader".into(),
649            msg: err.to_string(),
650        })
651    }
652}
iota_data_ingestion_core/reader/v2.rs

iota_data_ingestion_core/reader/
v2.rs
