sui_indexer_alt_framework/

lib.rs

// Copyright (c) Mysten Labs, Inc.
// SPDX-License-Identifier: Apache-2.0

use std::collections::BTreeSet;
use std::ops::Bound;
use std::sync::Arc;
use std::time::Duration;

use anyhow::Context;
use anyhow::bail;
use anyhow::ensure;
use ingestion::ClientArgs;
use ingestion::IngestionConfig;
use ingestion::IngestionService;
use ingestion::ingestion_client::IngestionClient;
use metrics::IndexerMetrics;
use prometheus::Registry;
use sui_indexer_alt_framework_store_traits::ConcurrentStore;
use sui_indexer_alt_framework_store_traits::Connection;
use sui_indexer_alt_framework_store_traits::SequentialStore;
use sui_indexer_alt_framework_store_traits::Store;
use sui_indexer_alt_framework_store_traits::pipeline_task;
use tracing::info;

use crate::metrics::IngestionMetrics;
use crate::pipeline::Processor;
use crate::pipeline::concurrent::ConcurrentConfig;
use crate::pipeline::concurrent::{self};
use crate::pipeline::sequential::SequentialConfig;
use crate::pipeline::sequential::{self};
use crate::service::Service;

pub use anyhow::Result;
pub use sui_field_count::FieldCount;
pub use sui_futures::service;
/// External users access the store trait through framework::store
pub use sui_indexer_alt_framework_store_traits as store;
pub use sui_types as types;

#[cfg(feature = "cluster")]
pub mod cluster;
pub mod config;
pub mod ingestion;
pub mod metrics;
pub mod pipeline;
#[cfg(feature = "postgres")]
pub mod postgres;

#[cfg(test)]
pub mod mocks;

/// Command-line arguments for the indexer
#[derive(clap::Args, Default, Debug, Clone)]
pub struct IndexerArgs {
    /// Override the next checkpoint for all pipelines without a committer watermark to start
    /// processing from, which is 0 by default. Pipelines with existing watermarks will ignore this
    /// setting and always resume from their committer watermark + 1.
    ///
    /// Setting this value indirectly affects ingestion, as the checkpoint to start ingesting from
    /// is the minimum across all pipelines' next checkpoints.
    #[arg(long)]
    pub first_checkpoint: Option<u64>,

    /// Override for the checkpoint to end ingestion at (inclusive) -- useful for backfills. By
    /// default, ingestion will not stop, and will continue to poll for new checkpoints.
    #[arg(long)]
    pub last_checkpoint: Option<u64>,

    /// Only run the following pipelines. If not provided, all pipelines found in the
    /// configuration file will be run.
    #[arg(long, action = clap::ArgAction::Append)]
    pub pipeline: Vec<String>,

    /// Additional configurations for running a tasked indexer.
    #[clap(flatten)]
    pub task: TaskArgs,
}

/// Command-line arguments for configuring a tasked indexer.
#[derive(clap::Parser, Default, Debug, Clone)]
pub struct TaskArgs {
    /// An optional task name for this indexer. When set, pipelines will record watermarks using the
    /// delimiter defined on the store. This allows the same pipelines to run under multiple
    /// indexers (e.g. for backfills or temporary workflows) while maintaining separate watermark
    /// entries in the database.
    ///
    /// By default there is no task name, and watermarks are keyed only by `pipeline`.
    ///
    /// Sequential pipelines cannot be attached to a tasked indexer.
    ///
    /// The framework ensures that tasked pipelines never commit checkpoints below the main
    /// pipeline’s pruner watermark. Requires `--reader-interval-ms`.
    #[arg(long, requires = "reader_interval_ms")]
    task: Option<String>,

    /// The interval in milliseconds at which each of the pipelines on a tasked indexer should
    /// refetch its main pipeline's reader watermark.
    ///
    /// This is required when `--task` is set and should should ideally be set to a value that is
    /// an order of magnitude smaller than the main pipeline's pruning interval, to ensure this
    /// task pipeline can pick up the new reader watermark before the main pipeline prunes up to
    /// it.
    ///
    /// If the main pipeline does not have pruning enabled, this value can be set to some high
    /// value, as the tasked pipeline will never see an updated reader watermark.
    #[arg(long, requires = "task")]
    reader_interval_ms: Option<u64>,
}

pub struct Indexer<S: Store> {
    /// The storage backend that the indexer uses to write and query indexed data. This
    /// generic implementation allows for plugging in different storage solutions that implement the
    /// `Store` trait.
    store: S,

    /// Prometheus Metrics.
    metrics: Arc<IndexerMetrics>,

    /// Service for downloading and disseminating checkpoint data.
    ingestion_service: IngestionService,

    /// Optional override of the checkpoint lowerbound. When set, pipelines without a committer
    /// watermark will start processing at this checkpoint.
    first_checkpoint: Option<u64>,

    /// Optional override of the checkpoint upperbound. When set, the indexer will stop ingestion at
    /// this checkpoint.
    last_checkpoint: Option<u64>,

    /// The minimum `next_checkpoint` across all pipelines. This is the checkpoint for the indexer
    /// to start ingesting from.
    next_checkpoint: u64,

    /// The minimum `next_checkpoint` across all sequential pipelines. This is used to initialize
    /// the regulator to prevent ingestion from running too far ahead of sequential pipelines.
    next_sequential_checkpoint: Option<u64>,

    /// An optional task name for this indexer. When set, pipelines will record watermarks using the
    /// delimiter defined on the store. This allows the same pipelines to run under multiple
    /// indexers (e.g. for backfills or temporary workflows) while maintaining separate watermark
    /// entries in the database.
    ///
    /// By default there is no task name, and watermarks are keyed only by `pipeline`.
    ///
    /// Sequential pipelines cannot be attached to a tasked indexer.
    ///
    /// The framework ensures that tasked pipelines never commit checkpoints below the main
    /// pipeline’s pruner watermark.
    task: Option<Task>,

    /// Optional filter for pipelines to run. If `None`, all pipelines added to the indexer will
    /// run. Any pipelines that are present in this filter but not added to the indexer will yield
    /// a warning when the indexer is run.
    enabled_pipelines: Option<BTreeSet<String>>,

    /// Pipelines that have already been registered with the indexer. Used to make sure a pipeline
    /// with the same name isn't added twice.
    added_pipelines: BTreeSet<&'static str>,

    /// The service handles for every pipeline, used to manage lifetimes and graceful shutdown.
    pipelines: Vec<Service>,
}

/// Configuration for a tasked indexer.
#[derive(Clone)]
pub(crate) struct Task {
    /// Name of the tasked indexer, to be used with the delimiter defined on the indexer's store to
    /// record pipeline watermarks.
    task: String,
    /// The interval at which each of the pipelines on a tasked indexer should refecth its main
    /// pipeline's reader watermark.
    reader_interval: Duration,
}

impl TaskArgs {
    pub fn tasked(task: String, reader_interval_ms: u64) -> Self {
        Self {
            task: Some(task),
            reader_interval_ms: Some(reader_interval_ms),
        }
    }

    fn into_task(self) -> Option<Task> {
        Some(Task {
            task: self.task?,
            reader_interval: Duration::from_millis(self.reader_interval_ms?),
        })
    }
}

impl<S: Store> Indexer<S> {
    /// Create a new instance of the indexer framework from a store that implements the `Store`
    /// trait, along with `indexer_args`, `client_args`, and `ingestion_config`. Together, these
    /// arguments configure the following:
    ///
    /// - What is indexed (which checkpoints, which pipelines, whether to track and update
    ///   watermarks) and where to serve metrics from,
    /// - Where to download checkpoints from,
    /// - Concurrency and buffering parameters for downloading checkpoints.
    ///
    /// After initialization, at least one pipeline must be added using [Self::concurrent_pipeline]
    /// or [Self::sequential_pipeline], before the indexer is started using [Self::run].
    pub async fn new(
        store: S,
        indexer_args: IndexerArgs,
        client_args: ClientArgs,
        ingestion_config: IngestionConfig,
        metrics_prefix: Option<&str>,
        registry: &Registry,
    ) -> Result<Self> {
        let IndexerArgs {
            first_checkpoint,
            last_checkpoint,
            pipeline,
            task,
        } = indexer_args;

        let metrics = IndexerMetrics::new(metrics_prefix, registry);

        let ingestion_service =
            IngestionService::new(client_args, ingestion_config, metrics_prefix, registry)?;

        Ok(Self {
            store,
            metrics,
            ingestion_service,
            first_checkpoint,
            last_checkpoint,
            next_checkpoint: u64::MAX,
            next_sequential_checkpoint: None,
            task: task.into_task(),
            enabled_pipelines: if pipeline.is_empty() {
                None
            } else {
                Some(pipeline.into_iter().collect())
            },
            added_pipelines: BTreeSet::new(),
            pipelines: vec![],
        })
    }

    /// The store used by the indexer.
    pub fn store(&self) -> &S {
        &self.store
    }

    /// The ingestion client used by the indexer to fetch checkpoints.
    pub fn ingestion_client(&self) -> &IngestionClient {
        self.ingestion_service.ingestion_client()
    }

    /// The indexer's metrics.
    pub fn indexer_metrics(&self) -> &Arc<IndexerMetrics> {
        &self.metrics
    }

    /// The ingestion service's metrics.
    pub fn ingestion_metrics(&self) -> &Arc<IngestionMetrics> {
        self.ingestion_service.metrics()
    }

    /// The pipelines that this indexer will run.
    pub fn pipelines(&self) -> impl Iterator<Item = &'static str> + '_ {
        self.added_pipelines.iter().copied().filter(|p| {
            self.enabled_pipelines
                .as_ref()
                .is_none_or(|e| e.contains(*p))
        })
    }

    /// The minimum next checkpoint across all sequential pipelines. This value is used to
    /// initialize the ingestion regulator's high watermark to prevent ingestion from running
    /// too far ahead of sequential pipelines.
    pub fn next_sequential_checkpoint(&self) -> Option<u64> {
        self.next_sequential_checkpoint
    }

    /// Start ingesting checkpoints from `next_checkpoint`. Individual pipelines
    /// will start processing and committing once the ingestion service has caught up to their
    /// respective watermarks.
    ///
    /// Ingestion will stop after consuming the configured `last_checkpoint` if one is provided.
    pub async fn run(self) -> anyhow::Result<Service> {
        if let Some(enabled_pipelines) = self.enabled_pipelines {
            ensure!(
                enabled_pipelines.is_empty(),
                "Tried to enable pipelines that this indexer does not know about: \
                {enabled_pipelines:#?}",
            );
        }

        let start = self.next_checkpoint;
        let end = self.last_checkpoint;
        info!(start, end, "Ingestion range");

        let mut service = self
            .ingestion_service
            .run(
                (
                    Bound::Included(start),
                    end.map_or(Bound::Unbounded, Bound::Included),
                ),
                self.next_sequential_checkpoint,
            )
            .await
            .context("Failed to start ingestion service")?;

        for pipeline in self.pipelines {
            service = service.merge(pipeline);
        }

        Ok(service)
    }

    /// Determine the checkpoint for the pipeline to resume processing from. This is either the
    /// checkpoint after its watermark, or if that doesn't exist, then the provided
    /// [Self::first_checkpoint], and if that is not set, then 0 (genesis).
    ///
    /// Update the starting ingestion checkpoint as the minimum across all the next checkpoints
    /// calculated above.
    ///
    /// Returns `Ok(None)` if the pipeline is disabled.
    async fn add_pipeline<P: Processor + 'static>(
        &mut self,
        pipeline_task: String,
    ) -> Result<Option<u64>> {
        ensure!(
            self.added_pipelines.insert(P::NAME),
            "Pipeline {:?} already added",
            P::NAME,
        );

        if let Some(enabled_pipelines) = &mut self.enabled_pipelines
            && !enabled_pipelines.remove(P::NAME)
        {
            info!(pipeline = P::NAME, "Skipping");
            return Ok(None);
        }

        // Create a new record based on `proposed_next_checkpoint` if one does not exist.
        // Otherwise, use the existing record and disregard the proposed value.
        let proposed_next_checkpoint = self.first_checkpoint.unwrap_or(0);
        let mut conn = self.store.connect().await?;
        let init_watermark = conn
            .init_watermark(&pipeline_task, proposed_next_checkpoint.checked_sub(1))
            .await
            .with_context(|| format!("Failed to init watermark for {pipeline_task}"))?;

        let next_checkpoint = if let Some(init_watermark) = init_watermark {
            if let Some(checkpoint_hi_inclusive) = init_watermark.checkpoint_hi_inclusive {
                checkpoint_hi_inclusive + 1
            } else {
                0
            }
        } else {
            proposed_next_checkpoint
        };

        self.next_checkpoint = next_checkpoint.min(self.next_checkpoint);

        Ok(Some(next_checkpoint))
    }
}

impl<S: ConcurrentStore> Indexer<S> {
    /// Adds a new pipeline to this indexer and starts it up. Although their tasks have started,
    /// they will be idle until the ingestion service starts, and serves it checkpoint data.
    ///
    /// Concurrent pipelines commit checkpoint data out-of-order to maximise throughput, and they
    /// keep the watermark table up-to-date with the highest point they can guarantee all data
    /// exists for, for their pipeline.
    pub async fn concurrent_pipeline<H: concurrent::Handler<Store = S>>(
        &mut self,
        handler: H,
        config: ConcurrentConfig,
    ) -> Result<()> {
        let pipeline_task =
            pipeline_task::<S>(H::NAME, self.task.as_ref().map(|t| t.task.as_str()))?;
        let Some(next_checkpoint) = self.add_pipeline::<H>(pipeline_task).await? else {
            return Ok(());
        };

        self.pipelines.push(concurrent::pipeline::<H>(
            handler,
            next_checkpoint,
            config,
            self.store.clone(),
            self.task.clone(),
            self.ingestion_service.subscribe().0,
            self.metrics.clone(),
        ));

        Ok(())
    }
}

impl<T: SequentialStore> Indexer<T> {
    /// Adds a new pipeline to this indexer and starts it up. Although their tasks have started,
    /// they will be idle until the ingestion service starts, and serves it checkpoint data.
    ///
    /// Sequential pipelines commit checkpoint data in-order which sacrifices throughput, but may be
    /// required to handle pipelines that modify data in-place (where each update is not an insert,
    /// but could be a modification of an existing row, where ordering between updates is
    /// important).
    ///
    /// The pipeline can optionally be configured to lag behind the ingestion service by a fixed
    /// number of checkpoints (configured by `checkpoint_lag`).
    pub async fn sequential_pipeline<H: sequential::Handler<Store = T>>(
        &mut self,
        handler: H,
        config: SequentialConfig,
    ) -> Result<()> {
        if self.task.is_some() {
            bail!(
                "Sequential pipelines do not support pipeline tasks. \
                These pipelines guarantee that each checkpoint is committed exactly once and in order. \
                Running the same pipeline under a different task would violate these guarantees."
            );
        }

        let Some(next_checkpoint) = self.add_pipeline::<H>(H::NAME.to_owned()).await? else {
            return Ok(());
        };

        // Track the minimum next_checkpoint across all sequential pipelines
        self.next_sequential_checkpoint = Some(
            self.next_sequential_checkpoint
                .map_or(next_checkpoint, |n| n.min(next_checkpoint)),
        );

        let (checkpoint_rx, commit_hi_tx) = self.ingestion_service.subscribe();

        self.pipelines.push(sequential::pipeline::<H>(
            handler,
            next_checkpoint,
            config,
            self.store.clone(),
            checkpoint_rx,
            commit_hi_tx,
            self.metrics.clone(),
        ));

        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use std::sync::Arc;

    use async_trait::async_trait;
    use clap::Parser;
    use sui_indexer_alt_framework_store_traits::PrunerWatermark;
    use sui_synthetic_ingestion::synthetic_ingestion;
    use tokio::sync::watch;

    use crate::FieldCount;
    use crate::config::ConcurrencyConfig;
    use crate::ingestion::ingestion_client::IngestionClientArgs;
    use crate::mocks::store::MockStore;
    use crate::pipeline::CommitterConfig;
    use crate::pipeline::Processor;
    use crate::pipeline::concurrent::ConcurrentConfig;
    use crate::store::CommitterWatermark;
    use crate::store::ConcurrentConnection as _;
    use crate::store::Connection as _;

    use super::*;

    #[allow(dead_code)]
    #[derive(Clone, FieldCount)]
    struct MockValue(u64);

    /// A handler that can be controlled externally to block checkpoint processing.
    struct ControllableHandler {
        /// Process checkpoints less than or equal to this value.
        process_below: watch::Receiver<u64>,
    }

    impl ControllableHandler {
        fn with_limit(limit: u64) -> (Self, watch::Sender<u64>) {
            let (tx, rx) = watch::channel(limit);
            (Self { process_below: rx }, tx)
        }
    }

    #[async_trait]
    impl Processor for ControllableHandler {
        const NAME: &'static str = "controllable";
        type Value = MockValue;

        async fn process(
            &self,
            checkpoint: &Arc<sui_types::full_checkpoint_content::Checkpoint>,
        ) -> anyhow::Result<Vec<Self::Value>> {
            let cp_num = checkpoint.summary.sequence_number;

            // Wait until the checkpoint is allowed to be processed
            self.process_below
                .clone()
                .wait_for(|&limit| cp_num <= limit)
                .await
                .ok();

            Ok(vec![MockValue(cp_num)])
        }
    }

    #[async_trait]
    impl concurrent::Handler for ControllableHandler {
        type Store = MockStore;
        type Batch = Vec<MockValue>;

        fn batch(
            &self,
            batch: &mut Self::Batch,
            values: &mut std::vec::IntoIter<Self::Value>,
        ) -> concurrent::BatchStatus {
            batch.extend(values);
            concurrent::BatchStatus::Ready
        }

        async fn commit<'a>(
            &self,
            batch: &Self::Batch,
            conn: &mut <Self::Store as Store>::Connection<'a>,
        ) -> anyhow::Result<usize> {
            for value in batch {
                conn.0
                    .commit_data(Self::NAME, value.0, vec![value.0])
                    .await?;
            }
            Ok(batch.len())
        }
    }

    macro_rules! test_pipeline {
        ($handler:ident, $name:literal) => {
            struct $handler;

            #[async_trait]
            impl Processor for $handler {
                const NAME: &'static str = $name;
                type Value = MockValue;
                async fn process(
                    &self,
                    checkpoint: &Arc<sui_types::full_checkpoint_content::Checkpoint>,
                ) -> anyhow::Result<Vec<Self::Value>> {
                    Ok(vec![MockValue(checkpoint.summary.sequence_number)])
                }
            }

            #[async_trait]
            impl crate::pipeline::concurrent::Handler for $handler {
                type Store = MockStore;
                type Batch = Vec<Self::Value>;

                fn batch(
                    &self,
                    batch: &mut Self::Batch,
                    values: &mut std::vec::IntoIter<Self::Value>,
                ) -> crate::pipeline::concurrent::BatchStatus {
                    batch.extend(values);
                    crate::pipeline::concurrent::BatchStatus::Pending
                }

                async fn commit<'a>(
                    &self,
                    batch: &Self::Batch,
                    conn: &mut <Self::Store as Store>::Connection<'a>,
                ) -> anyhow::Result<usize> {
                    for value in batch {
                        conn.0
                            .commit_data(Self::NAME, value.0, vec![value.0])
                            .await?;
                    }
                    Ok(batch.len())
                }
            }

            #[async_trait]
            impl crate::pipeline::sequential::Handler for $handler {
                type Store = MockStore;
                type Batch = Vec<Self::Value>;

                fn batch(&self, batch: &mut Self::Batch, values: std::vec::IntoIter<Self::Value>) {
                    batch.extend(values);
                }

                async fn commit<'a>(
                    &self,
                    _batch: &Self::Batch,
                    _conn: &mut <Self::Store as Store>::Connection<'a>,
                ) -> anyhow::Result<usize> {
                    Ok(1)
                }
            }
        };
    }

    test_pipeline!(MockHandler, "test_processor");
    test_pipeline!(SequentialHandler, "sequential_handler");
    test_pipeline!(MockCheckpointSequenceNumberHandler, "test");

    /// If `ingestion_data` is `Some((num_checkpoints, checkpoint_size))`, synthetic ingestion
    /// data will be generated in the temp directory before creating the indexer.
    async fn create_test_indexer(
        store: MockStore,
        indexer_args: IndexerArgs,
        registry: &Registry,
        ingestion_data: Option<(u64, u64)>,
    ) -> (Indexer<MockStore>, tempfile::TempDir) {
        let temp_dir = tempfile::tempdir().unwrap();
        if let Some((num_checkpoints, checkpoint_size)) = ingestion_data {
            synthetic_ingestion::generate_ingestion(synthetic_ingestion::Config {
                ingestion_dir: temp_dir.path().to_owned(),
                starting_checkpoint: 0,
                num_checkpoints,
                checkpoint_size,
            })
            .await;
        }
        let client_args = ClientArgs {
            ingestion: IngestionClientArgs {
                local_ingestion_path: Some(temp_dir.path().to_owned()),
                ..Default::default()
            },
            ..Default::default()
        };
        let indexer = Indexer::new(
            store,
            indexer_args,
            client_args,
            IngestionConfig::default(),
            None,
            registry,
        )
        .await
        .unwrap();
        (indexer, temp_dir)
    }

    async fn set_committer_watermark(
        conn: &mut <MockStore as Store>::Connection<'_>,
        name: &str,
        hi: u64,
    ) {
        conn.set_committer_watermark(
            name,
            CommitterWatermark {
                checkpoint_hi_inclusive: hi,
                ..Default::default()
            },
        )
        .await
        .unwrap();
    }

    async fn add_concurrent<H: concurrent::Handler<Store = MockStore>>(
        indexer: &mut Indexer<MockStore>,
        handler: H,
    ) {
        indexer
            .concurrent_pipeline(handler, ConcurrentConfig::default())
            .await
            .unwrap();
    }

    async fn add_sequential<H: sequential::Handler<Store = MockStore>>(
        indexer: &mut Indexer<MockStore>,
        handler: H,
    ) {
        indexer
            .sequential_pipeline(handler, SequentialConfig::default())
            .await
            .unwrap();
    }

    macro_rules! assert_out_of_order {
        ($metrics:expr, $pipeline:expr, $expected:expr) => {
            assert_eq!(
                $metrics
                    .total_watermarks_out_of_order
                    .get_metric_with_label_values(&[$pipeline])
                    .unwrap()
                    .get(),
                $expected,
            );
        };
    }

    async fn test_init_watermark(
        first_checkpoint: Option<u64>,
        is_concurrent: bool,
    ) -> (Option<CommitterWatermark>, Option<PrunerWatermark>) {
        let registry = Registry::new();
        let store = MockStore::default();

        test_pipeline!(A, "pipeline_name");

        let mut conn = store.connect().await.unwrap();

        let indexer_args = IndexerArgs {
            first_checkpoint,
            ..IndexerArgs::default()
        };
        let (mut indexer, _temp_dir) =
            create_test_indexer(store.clone(), indexer_args, &registry, None).await;

        if is_concurrent {
            add_concurrent(&mut indexer, A).await;
        } else {
            add_sequential(&mut indexer, A).await;
        }

        (
            conn.committer_watermark(A::NAME).await.unwrap(),
            conn.pruner_watermark(A::NAME, Duration::ZERO)
                .await
                .unwrap(),
        )
    }

    #[test]
    fn test_arg_parsing() {
        #[derive(Parser)]
        struct Args {
            #[clap(flatten)]
            indexer: IndexerArgs,
        }

        let args = Args::try_parse_from([
            "cmd",
            "--first-checkpoint",
            "10",
            "--last-checkpoint",
            "100",
            "--pipeline",
            "a",
            "--pipeline",
            "b",
            "--task",
            "t",
            "--reader-interval-ms",
            "5000",
        ])
        .unwrap();

        assert_eq!(args.indexer.first_checkpoint, Some(10));
        assert_eq!(args.indexer.last_checkpoint, Some(100));
        assert_eq!(args.indexer.pipeline, vec!["a", "b"]);
        assert_eq!(args.indexer.task.task, Some("t".to_owned()));
        assert_eq!(args.indexer.task.reader_interval_ms, Some(5000));
    }

    /// next_checkpoint is smallest among existing watermarks + 1.
    #[tokio::test]
    async fn test_next_checkpoint_all_pipelines_have_watermarks() {
        let registry = Registry::new();
        let store = MockStore::default();

        test_pipeline!(A, "concurrent_a");
        test_pipeline!(B, "concurrent_b");
        test_pipeline!(C, "sequential_c");
        test_pipeline!(D, "sequential_d");

        let mut conn = store.connect().await.unwrap();

        conn.init_watermark(A::NAME, Some(0)).await.unwrap();
        set_committer_watermark(&mut conn, A::NAME, 100).await;

        conn.init_watermark(B::NAME, Some(0)).await.unwrap();
        set_committer_watermark(&mut conn, B::NAME, 10).await;

        conn.init_watermark(C::NAME, Some(0)).await.unwrap();
        set_committer_watermark(&mut conn, C::NAME, 1).await;

        conn.init_watermark(D::NAME, Some(0)).await.unwrap();
        set_committer_watermark(&mut conn, D::NAME, 50).await;

        let (mut indexer, _temp_dir) =
            create_test_indexer(store, IndexerArgs::default(), &registry, None).await;

        add_concurrent(&mut indexer, A).await;
        add_concurrent(&mut indexer, B).await;
        add_sequential(&mut indexer, C).await;
        add_sequential(&mut indexer, D).await;

        assert_eq!(indexer.first_checkpoint, None);
        assert_eq!(indexer.last_checkpoint, None);
        assert_eq!(indexer.next_checkpoint, 2);
        assert_eq!(indexer.next_sequential_checkpoint, Some(2));
    }

    /// next_checkpoint is 0 when at least one pipeline has no watermark.
    #[tokio::test]
    async fn test_next_checkpoint_not_all_pipelines_have_watermarks() {
        let registry = Registry::new();
        let store = MockStore::default();

        test_pipeline!(A, "concurrent_a");
        test_pipeline!(B, "concurrent_b");
        test_pipeline!(C, "sequential_c");
        test_pipeline!(D, "sequential_d");

        let mut conn = store.connect().await.unwrap();
        set_committer_watermark(&mut conn, B::NAME, 10).await;
        set_committer_watermark(&mut conn, C::NAME, 1).await;

        let (mut indexer, _temp_dir) =
            create_test_indexer(store, IndexerArgs::default(), &registry, None).await;

        add_concurrent(&mut indexer, A).await;
        add_concurrent(&mut indexer, B).await;
        add_sequential(&mut indexer, C).await;
        add_sequential(&mut indexer, D).await;

        assert_eq!(indexer.first_checkpoint, None);
        assert_eq!(indexer.last_checkpoint, None);
        assert_eq!(indexer.next_checkpoint, 0);
        assert_eq!(indexer.next_sequential_checkpoint, Some(0));
    }

    /// next_checkpoint is 1 when smallest committer watermark is 0.
    #[tokio::test]
    async fn test_next_checkpoint_smallest_is_0() {
        let registry = Registry::new();
        let store = MockStore::default();

        test_pipeline!(A, "concurrent_a");
        test_pipeline!(B, "concurrent_b");
        test_pipeline!(C, "sequential_c");
        test_pipeline!(D, "sequential_d");

        let mut conn = store.connect().await.unwrap();
        set_committer_watermark(&mut conn, A::NAME, 100).await;
        set_committer_watermark(&mut conn, B::NAME, 10).await;
        set_committer_watermark(&mut conn, C::NAME, 1).await;
        set_committer_watermark(&mut conn, D::NAME, 0).await;

        let (mut indexer, _temp_dir) =
            create_test_indexer(store, IndexerArgs::default(), &registry, None).await;

        add_concurrent(&mut indexer, A).await;
        add_concurrent(&mut indexer, B).await;
        add_sequential(&mut indexer, C).await;
        add_sequential(&mut indexer, D).await;

        assert_eq!(indexer.next_checkpoint, 1);
    }

    /// next_checkpoint is first_checkpoint when at least one pipeline has no
    /// watermark, and first_checkpoint is smallest.
    #[tokio::test]
    async fn test_next_checkpoint_first_checkpoint_and_no_watermark() {
        let registry = Registry::new();
        let store = MockStore::default();

        test_pipeline!(A, "concurrent_a");
        test_pipeline!(B, "concurrent_b");
        test_pipeline!(C, "sequential_c");
        test_pipeline!(D, "sequential_d");

        let mut conn = store.connect().await.unwrap();
        set_committer_watermark(&mut conn, B::NAME, 50).await;
        set_committer_watermark(&mut conn, C::NAME, 10).await;

        let indexer_args = IndexerArgs {
            first_checkpoint: Some(5),
            ..Default::default()
        };
        let (mut indexer, _temp_dir) =
            create_test_indexer(store, indexer_args, &registry, None).await;

        add_concurrent(&mut indexer, A).await;
        add_concurrent(&mut indexer, B).await;
        add_sequential(&mut indexer, C).await;
        add_sequential(&mut indexer, D).await;

        assert_eq!(indexer.first_checkpoint, Some(5));
        assert_eq!(indexer.last_checkpoint, None);
        assert_eq!(indexer.next_checkpoint, 5);
        assert_eq!(indexer.next_sequential_checkpoint, Some(5));
    }

    /// next_checkpoint is smallest among existing watermarks + 1 if
    /// all pipelines have watermarks (ignores first_checkpoint).
    #[tokio::test]
    async fn test_next_checkpoint_ignore_first_checkpoint() {
        let registry = Registry::new();
        let store = MockStore::default();

        test_pipeline!(B, "concurrent_b");
        test_pipeline!(C, "sequential_c");

        let mut conn = store.connect().await.unwrap();
        set_committer_watermark(&mut conn, B::NAME, 50).await;
        set_committer_watermark(&mut conn, C::NAME, 10).await;

        let indexer_args = IndexerArgs {
            first_checkpoint: Some(5),
            ..Default::default()
        };
        let (mut indexer, _temp_dir) =
            create_test_indexer(store, indexer_args, &registry, None).await;

        add_concurrent(&mut indexer, B).await;
        add_sequential(&mut indexer, C).await;

        assert_eq!(indexer.first_checkpoint, Some(5));
        assert_eq!(indexer.last_checkpoint, None);
        assert_eq!(indexer.next_checkpoint, 11);
        assert_eq!(indexer.next_sequential_checkpoint, Some(11));
    }

    /// If the first_checkpoint is being considered, because pipelines are missing watermarks, it
    /// will not be used as the starting point if it is not the smallest valid committer watermark
    /// to resume ingesting from.
    #[tokio::test]
    async fn test_next_checkpoint_large_first_checkpoint() {
        let registry = Registry::new();
        let store = MockStore::default();

        test_pipeline!(A, "concurrent_a");
        test_pipeline!(B, "concurrent_b");
        test_pipeline!(C, "sequential_c");

        let mut conn = store.connect().await.unwrap();
        set_committer_watermark(&mut conn, B::NAME, 50).await;
        set_committer_watermark(&mut conn, C::NAME, 10).await;

        let indexer_args = IndexerArgs {
            first_checkpoint: Some(24),
            ..Default::default()
        };
        let (mut indexer, _temp_dir) =
            create_test_indexer(store, indexer_args, &registry, None).await;

        add_concurrent(&mut indexer, A).await;
        add_concurrent(&mut indexer, B).await;
        add_sequential(&mut indexer, C).await;

        assert_eq!(indexer.first_checkpoint, Some(24));
        assert_eq!(indexer.last_checkpoint, None);
        assert_eq!(indexer.next_checkpoint, 11);
        assert_eq!(indexer.next_sequential_checkpoint, Some(11));
    }

    // test ingestion, all pipelines have watermarks, no first_checkpoint provided
    #[tokio::test]
    async fn test_indexer_ingestion_existing_watermarks_no_first_checkpoint() {
        let registry = Registry::new();
        let store = MockStore::default();

        test_pipeline!(A, "concurrent_a");
        test_pipeline!(B, "concurrent_b");
        test_pipeline!(C, "sequential_c");
        test_pipeline!(D, "sequential_d");

        let mut conn = store.connect().await.unwrap();
        set_committer_watermark(&mut conn, A::NAME, 5).await;
        set_committer_watermark(&mut conn, B::NAME, 10).await;
        set_committer_watermark(&mut conn, C::NAME, 15).await;
        set_committer_watermark(&mut conn, D::NAME, 20).await;

        let indexer_args = IndexerArgs {
            last_checkpoint: Some(29),
            ..Default::default()
        };
        let (mut indexer, _temp_dir) =
            create_test_indexer(store.clone(), indexer_args, &registry, Some((30, 1))).await;

        add_concurrent(&mut indexer, A).await;
        add_concurrent(&mut indexer, B).await;
        add_sequential(&mut indexer, C).await;
        add_sequential(&mut indexer, D).await;

        let ingestion_metrics = indexer.ingestion_metrics().clone();
        let indexer_metrics = indexer.indexer_metrics().clone();

        indexer.run().await.unwrap().join().await.unwrap();

        assert_eq!(ingestion_metrics.total_ingested_checkpoints.get(), 24);
        assert_out_of_order!(indexer_metrics, A::NAME, 0);
        assert_out_of_order!(indexer_metrics, B::NAME, 5);
        assert_out_of_order!(indexer_metrics, C::NAME, 10);
        assert_out_of_order!(indexer_metrics, D::NAME, 15);
    }

    // test ingestion, some pipelines missing watermarks, no first_checkpoint provided
    #[tokio::test]
    async fn test_indexer_ingestion_missing_watermarks_no_first_checkpoint() {
        let registry = Registry::new();
        let store = MockStore::default();

        test_pipeline!(A, "concurrent_a");
        test_pipeline!(B, "concurrent_b");
        test_pipeline!(C, "sequential_c");
        test_pipeline!(D, "sequential_d");

        let mut conn = store.connect().await.unwrap();
        set_committer_watermark(&mut conn, B::NAME, 10).await;
        set_committer_watermark(&mut conn, C::NAME, 15).await;
        set_committer_watermark(&mut conn, D::NAME, 20).await;

        let indexer_args = IndexerArgs {
            last_checkpoint: Some(29),
            ..Default::default()
        };
        let (mut indexer, _temp_dir) =
            create_test_indexer(store.clone(), indexer_args, &registry, Some((30, 1))).await;

        add_concurrent(&mut indexer, A).await;
        add_concurrent(&mut indexer, B).await;
        add_sequential(&mut indexer, C).await;
        add_sequential(&mut indexer, D).await;

        let ingestion_metrics = indexer.ingestion_metrics().clone();
        let metrics = indexer.indexer_metrics().clone();
        indexer.run().await.unwrap().join().await.unwrap();

        assert_eq!(ingestion_metrics.total_ingested_checkpoints.get(), 30);
        assert_out_of_order!(metrics, A::NAME, 0);
        assert_out_of_order!(metrics, B::NAME, 11);
        assert_out_of_order!(metrics, C::NAME, 16);
        assert_out_of_order!(metrics, D::NAME, 21);
    }

    // test ingestion, some pipelines missing watermarks, use first_checkpoint
    #[tokio::test]
    async fn test_indexer_ingestion_use_first_checkpoint() {
        let registry = Registry::new();
        let store = MockStore::default();

        test_pipeline!(A, "concurrent_a");
        test_pipeline!(B, "concurrent_b");
        test_pipeline!(C, "sequential_c");
        test_pipeline!(D, "sequential_d");

        let mut conn = store.connect().await.unwrap();
        set_committer_watermark(&mut conn, B::NAME, 10).await;
        set_committer_watermark(&mut conn, C::NAME, 15).await;
        set_committer_watermark(&mut conn, D::NAME, 20).await;

        let indexer_args = IndexerArgs {
            first_checkpoint: Some(10),
            last_checkpoint: Some(29),
            ..Default::default()
        };
        let (mut indexer, _temp_dir) =
            create_test_indexer(store.clone(), indexer_args, &registry, Some((30, 1))).await;

        add_concurrent(&mut indexer, A).await;
        add_concurrent(&mut indexer, B).await;
        add_sequential(&mut indexer, C).await;
        add_sequential(&mut indexer, D).await;

        let ingestion_metrics = indexer.ingestion_metrics().clone();
        let metrics = indexer.indexer_metrics().clone();
        indexer.run().await.unwrap().join().await.unwrap();

        assert_eq!(ingestion_metrics.total_ingested_checkpoints.get(), 20);
        assert_out_of_order!(metrics, A::NAME, 0);
        assert_out_of_order!(metrics, B::NAME, 1);
        assert_out_of_order!(metrics, C::NAME, 6);
        assert_out_of_order!(metrics, D::NAME, 11);
    }

    #[tokio::test]
    async fn test_init_watermark_concurrent_no_first_checkpoint() {
        let (committer_watermark, pruner_watermark) = test_init_watermark(None, true).await;
        assert_eq!(committer_watermark, None);
        assert_eq!(pruner_watermark, None);
    }

    #[tokio::test]
    async fn test_init_watermark_concurrent_first_checkpoint_0() {
        let (committer_watermark, pruner_watermark) = test_init_watermark(Some(0), true).await;
        assert_eq!(committer_watermark, None);
        assert_eq!(pruner_watermark, None);
    }

    #[tokio::test]
    async fn test_init_watermark_concurrent_first_checkpoint_1() {
        let (committer_watermark, pruner_watermark) = test_init_watermark(Some(1), true).await;

        let committer_watermark = committer_watermark.unwrap();
        assert_eq!(committer_watermark.checkpoint_hi_inclusive, 0);

        let pruner_watermark = pruner_watermark.unwrap();
        assert_eq!(pruner_watermark.reader_lo, 1);
        assert_eq!(pruner_watermark.pruner_hi, 1);
    }

    #[tokio::test]
    async fn test_init_watermark_sequential() {
        let (committer_watermark, pruner_watermark) = test_init_watermark(Some(1), false).await;

        let committer_watermark = committer_watermark.unwrap();
        assert_eq!(committer_watermark.checkpoint_hi_inclusive, 0);

        let pruner_watermark = pruner_watermark.unwrap();
        assert_eq!(pruner_watermark.reader_lo, 1);
        assert_eq!(pruner_watermark.pruner_hi, 1);
    }

    #[tokio::test]
    async fn test_multiple_sequential_pipelines_next_checkpoint() {
        let registry = Registry::new();
        let store = MockStore::default();

        let mut conn = store.connect().await.unwrap();
        set_committer_watermark(&mut conn, MockHandler::NAME, 10).await;
        set_committer_watermark(&mut conn, SequentialHandler::NAME, 5).await;

        let indexer_args = IndexerArgs {
            first_checkpoint: None,
            last_checkpoint: Some(19),
            pipeline: vec![],
            ..Default::default()
        };
        let (mut indexer, _temp_dir) =
            create_test_indexer(store.clone(), indexer_args, &registry, Some((20, 2))).await;

        // Add first sequential pipeline
        add_sequential(&mut indexer, MockHandler).await;

        // Verify next_sequential_checkpoint is set correctly (10 + 1 = 11)
        assert_eq!(
            indexer.next_sequential_checkpoint(),
            Some(11),
            "next_sequential_checkpoint should be 11"
        );

        // Add second sequential pipeline
        add_sequential(&mut indexer, SequentialHandler).await;

        // Should change to 6 (minimum of 6 and 11)
        assert_eq!(
            indexer.next_sequential_checkpoint(),
            Some(6),
            "next_sequential_checkpoint should still be 6"
        );

        // Run indexer to verify it can make progress past the initial hi and finish ingesting.
        indexer.run().await.unwrap().join().await.unwrap();

        // Verify each pipeline made some progress independently
        let watermark1 = conn.committer_watermark(MockHandler::NAME).await.unwrap();
        let watermark2 = conn
            .committer_watermark(SequentialHandler::NAME)
            .await
            .unwrap();

        assert_eq!(watermark1.unwrap().checkpoint_hi_inclusive, 19);
        assert_eq!(watermark2.unwrap().checkpoint_hi_inclusive, 19);
    }

    /// When a tasked indexer is initialized such that a tasked pipeline is run with a
    /// `first_checkpoint` less than the main pipeline's reader_lo, the indexer will correctly skip
    /// committing checkpoints less than the main pipeline's reader watermark.
    #[tokio::test]
    async fn test_tasked_pipelines_ignore_below_main_reader_lo() {
        let registry = Registry::new();
        let store = MockStore::default();

        // Mock the store as if we have a main pipeline with a committer watermark at `10` and a
        // reader watermark at `7`.
        let mut conn = store.connect().await.unwrap();
        set_committer_watermark(&mut conn, MockCheckpointSequenceNumberHandler::NAME, 10).await;
        conn.set_reader_watermark(MockCheckpointSequenceNumberHandler::NAME, 7)
            .await
            .unwrap();

        // Start a tasked indexer that will ingest from checkpoint 0. Checkpoints 0 through 6 should
        // be ignored by the tasked indexer.
        let indexer_args = IndexerArgs {
            first_checkpoint: Some(0),
            last_checkpoint: Some(15),
            task: TaskArgs::tasked("task".to_string(), 10),
            ..Default::default()
        };
        let (mut tasked_indexer, _temp_dir) =
            create_test_indexer(store.clone(), indexer_args, &registry, Some((16, 2))).await;

        add_concurrent(&mut tasked_indexer, MockCheckpointSequenceNumberHandler).await;

        let ingestion_metrics = tasked_indexer.ingestion_metrics().clone();
        let metrics = tasked_indexer.indexer_metrics().clone();

        tasked_indexer.run().await.unwrap().join().await.unwrap();

        assert_eq!(ingestion_metrics.total_ingested_checkpoints.get(), 16);
        assert_eq!(
            metrics
                .total_collector_skipped_checkpoints
                .get_metric_with_label_values(&[MockCheckpointSequenceNumberHandler::NAME])
                .unwrap()
                .get(),
            7
        );
        let data = store
            .data
            .get(MockCheckpointSequenceNumberHandler::NAME)
            .unwrap();
        assert_eq!(data.len(), 9);
        for i in 0..7 {
            assert!(data.get(&i).is_none());
        }
        for i in 7..16 {
            assert!(data.get(&i).is_some());
        }
    }

    /// Tasked pipelines can run ahead of the main pipeline's committer watermark.
    #[tokio::test]
    async fn test_tasked_pipelines_surpass_main_pipeline_committer_hi() {
        let registry = Registry::new();
        let store = MockStore::default();

        let mut conn = store.connect().await.unwrap();
        set_committer_watermark(&mut conn, "test", 10).await;
        conn.set_reader_watermark("test", 5).await.unwrap();

        // Start a tasked indexer that will ingest from checkpoint 9 and go past the main pipeline's
        // watermarks.
        let indexer_args = IndexerArgs {
            first_checkpoint: Some(9),
            last_checkpoint: Some(25),
            task: TaskArgs::tasked("task".to_string(), 10),
            ..Default::default()
        };
        let (mut tasked_indexer, _temp_dir) =
            create_test_indexer(store.clone(), indexer_args, &registry, Some((26, 2))).await;

        add_concurrent(&mut tasked_indexer, MockCheckpointSequenceNumberHandler).await;

        let ingestion_metrics = tasked_indexer.ingestion_metrics().clone();
        let metrics = tasked_indexer.indexer_metrics().clone();

        tasked_indexer.run().await.unwrap().join().await.unwrap();

        assert_eq!(ingestion_metrics.total_ingested_checkpoints.get(), 17);
        assert_out_of_order!(metrics, "test", 0);
        assert_eq!(
            metrics
                .total_collector_skipped_checkpoints
                .get_metric_with_label_values(&[MockCheckpointSequenceNumberHandler::NAME])
                .unwrap()
                .get(),
            0
        );

        let data = store.data.get("test").unwrap();
        assert_eq!(data.len(), 17);
        for i in 0..9 {
            assert!(data.get(&i).is_none());
        }
        for i in 9..26 {
            assert!(data.get(&i).is_some());
        }
        let main_pipeline_watermark = store.watermark("test").unwrap();
        // assert that the main pipeline's watermarks are not updated
        assert_eq!(main_pipeline_watermark.checkpoint_hi_inclusive, Some(10));
        assert_eq!(main_pipeline_watermark.reader_lo, 5);
        let tasked_pipeline_watermark = store.watermark("test@task").unwrap();
        assert_eq!(tasked_pipeline_watermark.checkpoint_hi_inclusive, Some(25));
        assert_eq!(tasked_pipeline_watermark.reader_lo, 9);
    }

    /// Test that when the collector observes `reader_lo = X`, that all checkpoints >= X will be
    /// committed, and any checkpoints inflight < X will be skipped.
    #[tokio::test]
    async fn test_tasked_pipelines_skip_checkpoints_trailing_main_reader_lo() {
        let registry = Registry::new();
        let store = MockStore::default();
        let mut conn = store.connect().await.unwrap();
        // Set the main pipeline watermark.
        set_committer_watermark(&mut conn, ControllableHandler::NAME, 11).await;

        // Generate 500 checkpoints upfront, for the indexer to process all at once.
        let indexer_args = IndexerArgs {
            first_checkpoint: Some(0),
            last_checkpoint: Some(500),
            task: TaskArgs::tasked("task".to_string(), 10 /* reader_interval_ms */),
            ..Default::default()
        };
        let (mut tasked_indexer, _temp_dir) =
            create_test_indexer(store.clone(), indexer_args, &registry, Some((501, 2))).await;
        let mut allow_process = 10;
        // Limit the pipeline to process only checkpoints `[0, 10]`.
        let (controllable_handler, process_below) = ControllableHandler::with_limit(allow_process);
        let _ = tasked_indexer
            .concurrent_pipeline(
                controllable_handler,
                ConcurrentConfig {
                    committer: CommitterConfig {
                        collect_interval_ms: 10,
                        watermark_interval_ms: 10,
                        ..Default::default()
                    },
                    // High fixed concurrency so all checkpoints can be processed
                    // concurrently despite out-of-order arrival.
                    fanout: Some(ConcurrencyConfig::Fixed { value: 501 }),
                    ..Default::default()
                },
            )
            .await;
        let metrics = tasked_indexer.indexer_metrics().clone();

        let mut s_indexer = tasked_indexer.run().await.unwrap();

        // Wait for pipeline to commit up to configured checkpoint 10 inclusive. With the main
        // pipeline `reader_lo` currently unset, all checkpoints are allowed and should be
        // committed.
        store
            .wait_for_watermark(
                &pipeline_task::<MockStore>(ControllableHandler::NAME, Some("task")).unwrap(),
                10,
                Duration::from_secs(10),
            )
            .await;

        // Set the reader_lo to 250, simulating the main pipeline getting ahead. The
        // track_main_reader_lo task will eventually pick this up and update the atomic. The
        // collector reads from the atomic when it receives checkpoints, so we release checkpoints
        // one at a time until the collector_reader_lo metric shows the new value.
        conn.set_reader_watermark(ControllableHandler::NAME, 250)
            .await
            .unwrap();

        let reader_lo = metrics
            .collector_reader_lo
            .with_label_values(&[ControllableHandler::NAME]);

        // Send checkpoints one at a time at 10ms intervals. The tasked indexer has a reader refresh
        // interval of 10ms as well, so the collector should pick up the new reader_lo after a few
        // checkpoints have been processed.
        let mut interval = tokio::time::interval(Duration::from_millis(10));
        while reader_lo.get() != 250 {
            interval.tick().await;
            // allow_process is initialized to 11, bump to 11 for the next checkpoint
            allow_process += 1;
            assert!(
                allow_process <= 500,
                "Released all checkpoints but collector never observed new reader_lo"
            );
            process_below.send(allow_process).ok();
        }

        // At this point, the collector has observed reader_lo = 250. Release all remaining
        // checkpoints. Guarantees:
        // - [0, 10]: committed (before reader_lo was set)
        // - [11, allow_process]: some committed, some skipped (timing-dependent during detection)
        // - (allow_process, 250): skipped (in-flight, filtered by collector)
        // - [250, 500]: committed (>= reader_lo)
        process_below.send(500).ok();

        s_indexer.join().await.unwrap();

        let data = store.data.get(ControllableHandler::NAME).unwrap();

        // Checkpoints (allow_process, 250) must be skipped.
        for chkpt in (allow_process + 1)..250 {
            assert!(
                data.get(&chkpt).is_none(),
                "Checkpoint {chkpt} should have been skipped"
            );
        }

        // Checkpoints >= reader_lo must be committed.
        for chkpt in 250..=500 {
            assert!(
                data.get(&chkpt).is_some(),
                "Checkpoint {chkpt} should have been committed (>= reader_lo)"
            );
        }

        // Baseline: checkpoints [0, 10] were committed before reader_lo was set.
        for chkpt in 0..=10 {
            assert!(
                data.get(&chkpt).is_some(),
                "Checkpoint {chkpt} should have been committed (baseline)"
            );
        }
    }
}