sui-indexer-alt-framework: separate trailing_init_watermark from init_watermark

Author: evan-wall-mysten

crates/sui-analytics-indexer/src/store/mod.rs

@@ -23,7 +23,6 @@ use std::sync::RwLock;
 use std::time::Duration;
 use std::time::Instant;
 
-use anyhow::Result;
 use async_trait::async_trait;
 use object_store::PutPayload;
 use object_store::path::Path as ObjectPath;
@@ -33,8 +32,6 @@ use sui_indexer_alt_framework::store::Connection;
 use sui_indexer_alt_framework::store::Store;
 use sui_indexer_alt_framework::store::TransactionalStore;
 use sui_indexer_alt_framework_store_traits::CommitterWatermark;
-use sui_indexer_alt_framework_store_traits::InitWatermark;
-use sui_indexer_alt_framework_store_traits::init_with_committer_watermark;
 use sui_types::base_types::EpochId;
 use tokio::sync::mpsc;
 use tracing::debug;
@@ -287,7 +284,7 @@ impl AnalyticsStore {
         &self,
         first_checkpoint: Option<u64>,
         last_checkpoint: Option<u64>,
-    ) -> Result<(Option<u64>, Option<u64>)> {
+    ) -> anyhow::Result<(Option<u64>, Option<u64>)> {
         match &self.mode {
             StoreMode::Live(_) => Ok((first_checkpoint, last_checkpoint)),
             StoreMode::Migration(store) => {
@@ -533,7 +530,7 @@ impl<'a> AnalyticsConnection<'a> {
     pub async fn commit_batch<P: Processor>(
         &mut self,
         batch_from_framework: &[CheckpointRows],
-    ) -> Result<usize> {
+    ) -> anyhow::Result<usize> {
         let pipeline = P::NAME;
         let pipeline_config = self.pipeline_config(pipeline);
 
@@ -623,18 +620,6 @@ impl TransactionalStore for AnalyticsStore {
 
 #[async_trait]
 impl Connection for AnalyticsConnection<'_> {
-    /// Initialize watermark.
-    ///
-    /// In live mode: Watermarks are derived from file names, so just delegates to `committer_watermark`.
-    /// In migration mode: Delegates to `MigrationStore::init_watermark`.
-    async fn init_watermark(
-        &mut self,
-        pipeline_task: &str,
-        init_watermark: InitWatermark,
-    ) -> anyhow::Result<InitWatermark> {
-        init_with_committer_watermark(self, pipeline_task, init_watermark).await
-    }
-
     /// Determine the watermark.
     ///
     /// In live mode: scans file names in the object store.

crates/sui-indexer-alt-consistent-store/src/store/mod.rs

@@ -13,9 +13,7 @@ use prometheus::Registry;
 use scoped_futures::ScopedBoxFuture;
 use sui_indexer_alt_framework::service::Service;
 use sui_indexer_alt_framework::store::CommitterWatermark;
-use sui_indexer_alt_framework::store::InitWatermark;
 use sui_indexer_alt_framework::store::Store as _;
-use sui_indexer_alt_framework::store::init_with_committer_watermark;
 use sui_indexer_alt_framework::store::{self};
 
 use crate::db::Db;
@@ -167,14 +165,6 @@ impl<S: Send + Sync + 'static> store::TransactionalStore for Store<S> {
 
 #[async_trait]
 impl<S: Send + Sync> store::Connection for Connection<'_, S> {
-    async fn init_watermark(
-        &mut self,
-        pipeline_task: &str,
-        init_watermark: InitWatermark,
-    ) -> anyhow::Result<InitWatermark> {
-        init_with_committer_watermark(self, pipeline_task, init_watermark).await
-    }
-
     async fn committer_watermark(
         &mut self,
         pipeline_task: &str,

crates/sui-indexer-alt-framework-store-traits/src/lib.rs

@@ -15,11 +15,21 @@ pub trait Connection: Send {
     /// Returns the `InitWatermark` based on the existing watermark if it exists.
     /// Otherwise, initializes a new watermark record with `InitWatermark` and returns
     /// the value passed in.
+    ///
+    /// The default implementation does not perform initialization and delegates to `committer_watermark`.
     async fn init_watermark(
         &mut self,
         pipeline_task: &str,
-        init_watermark: InitWatermark,
-    ) -> anyhow::Result<InitWatermark>;
+        _init_watermark: InitWatermark,
+    ) -> anyhow::Result<InitWatermark> {
+        let checkpoint_hi_inclusive = self
+            .committer_watermark(pipeline_task)
+            .await?
+            .map(|w| w.checkpoint_hi_inclusive);
+        Ok(InitWatermark {
+            checkpoint_hi_inclusive,
+        })
+    }
 
     /// Given a `pipeline_task` representing either a pipeline name or a pipeline with an associated
     /// task (formatted as `{pipeline}{Store::DELIMITER}{task}`), return the committer watermark
@@ -45,6 +55,15 @@ pub trait Connection: Send {
 /// the pruning subsystem.
 #[async_trait]
 pub trait TrailingConnection: Connection {
+    /// Returns the `InitWatermark` based on the existing watermark if it exists.
+    /// Otherwise, initializes a new watermark record with `InitWatermark` and returns
+    /// the value passed in.
+    async fn trailing_init_watermark(
+        &mut self,
+        pipeline_task: &str,
+        trailing_init_watermark: TrailingInitWatermark,
+    ) -> anyhow::Result<TrailingInitWatermark>;
+
     /// Given a pipeline, return the reader watermark from the database. This is used by the indexer
     /// to determine the new `reader_lo` or inclusive lower bound of available data.
     async fn reader_watermark(
@@ -120,12 +139,18 @@ pub trait TransactionalStore: Store {
             &'r mut Self::Connection<'_>,
         ) -> ScopedBoxFuture<'a, 'r, anyhow::Result<R>>;
 }
-
 /// Used during watermark initialization to set and return state.
 #[derive(Default, Debug, Clone, Copy, PartialEq)]
 pub struct InitWatermark {
     /// Calculated by the framework as `default_next_checkpoint.checked_sub(1)`.
     pub checkpoint_hi_inclusive: Option<u64>,
+}
+
+/// Used during watermark initialization to set and return state.
+#[derive(Default, Debug, Clone, Copy, PartialEq)]
+pub struct TrailingInitWatermark {
+    /// Calculated by the framework as `default_next_checkpoint.checked_sub(1)`.
+    pub checkpoint_hi_inclusive: Option<u64>,
     /// Calculated by the framework as `default_next_checkpoint`.
     pub reader_lo: u64,
 }
@@ -212,23 +237,6 @@ impl PrunerWatermark {
     }
 }
 
-/// A utility function for connections that do not have special initialization logic. These
-/// connections delegate initialization to `Connection::committer_watermark`.
-pub async fn init_with_committer_watermark(
-    connection: &mut impl Connection,
-    pipeline_task: &str,
-    init_watermark: InitWatermark,
-) -> anyhow::Result<InitWatermark> {
-    let checkpoint_hi_inclusive = connection
-        .committer_watermark(pipeline_task)
-        .await?
-        .map(|w| w.checkpoint_hi_inclusive);
-    Ok(InitWatermark {
-        checkpoint_hi_inclusive,
-        ..init_watermark
-    })
-}
-
 /// Check that the pipeline name does not contain the store's delimiter, and construct the string
 /// used for tracking a pipeline's watermarks in the store. This is either the pipeline name itself,
 /// or `{pipeline}{Store::DELIMITER}{task}` if a task name is provided.

crates/sui-indexer-alt-framework/src/lib.rs

@@ -14,12 +14,12 @@ use ingestion::IngestionService;
 use ingestion::ingestion_client::IngestionClient;
 use metrics::IndexerMetrics;
 use prometheus::Registry;
-use sui_indexer_alt_framework_store_traits::Connection;
-use sui_indexer_alt_framework_store_traits::InitWatermark;
 use sui_indexer_alt_framework_store_traits::Store;
 use sui_indexer_alt_framework_store_traits::TrailingConnection;
+use sui_indexer_alt_framework_store_traits::TrailingInitWatermark;
 use sui_indexer_alt_framework_store_traits::TransactionalStore;
 use sui_indexer_alt_framework_store_traits::pipeline_task;
+use sui_indexer_alt_framework_store_traits::{Connection, InitWatermark};
 use tracing::info;
 
 use crate::metrics::IngestionMetrics;
@@ -294,7 +294,13 @@ impl<S: Store> Indexer<S> {
     where
         H: concurrent::Handler<Store = S> + Send + Sync + 'static,
     {
-        let Some(next_checkpoint) = self.add_pipeline::<H>().await? else {
+        let default_next_checkpoint = self.default_next_checkpoint;
+        let Some(next_checkpoint) = self
+            .add_pipeline::<H>(async |pt, conn| {
+                Self::next_checkpoint(default_next_checkpoint, pt, conn).await
+            })
+            .await?
+        else {
             return Ok(());
         };
 
@@ -345,6 +351,26 @@ impl<S: Store> Indexer<S> {
         Ok(service)
     }
 
+    fn next_checkpoint<'c>(
+        default_next_checkpoint: u64,
+        pipeline_task: String,
+        mut conn: S::Connection<'c>,
+    ) -> impl Future<Output = anyhow::Result<u64>> + 'c {
+        async move {
+            let InitWatermark {
+                checkpoint_hi_inclusive,
+            } = conn
+                .init_watermark(
+                    &pipeline_task,
+                    InitWatermark {
+                        checkpoint_hi_inclusive: default_next_checkpoint.checked_sub(1),
+                    },
+                )
+                .await?;
+            Ok(checkpoint_hi_inclusive.map_or(0, |c| c + 1))
+        }
+    }
+
     /// 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).
@@ -353,7 +379,10 @@ impl<S: Store> Indexer<S> {
     /// calculated above.
     ///
     /// Returns `Ok(None)` if the pipeline is disabled.
-    async fn add_pipeline<P: Processor + 'static>(&mut self) -> Result<Option<u64>> {
+    async fn add_pipeline<P: Processor + 'static>(
+        &mut self,
+        next_checkpoint_fn: impl for<'c> AsyncFnOnce(String, S::Connection<'c>) -> anyhow::Result<u64>,
+    ) -> Result<Option<u64>> {
         ensure!(
             self.added_pipelines.insert(P::NAME),
             "Pipeline {:?} already added",
@@ -367,7 +396,7 @@ impl<S: Store> Indexer<S> {
             return Ok(None);
         }
 
-        let mut conn = self
+        let conn = self
             .store
             .connect()
             .await
@@ -376,21 +405,9 @@ impl<S: Store> Indexer<S> {
         let pipeline_task =
             pipeline_task::<S>(P::NAME, self.task.as_ref().map(|t| t.task.as_str()))?;
 
-        let InitWatermark {
-            checkpoint_hi_inclusive,
-            reader_lo,
-        } = conn
-            .init_watermark(
-                &pipeline_task,
-                InitWatermark {
-                    checkpoint_hi_inclusive: self.default_next_checkpoint.checked_sub(1),
-                    reader_lo: self.default_next_checkpoint,
-                },
-            )
+        let next_checkpoint = next_checkpoint_fn(pipeline_task.clone(), conn)
             .await
-            .with_context(|| format!("Failed to init watermark for {pipeline_task}"))?;
-
-        let next_checkpoint = checkpoint_hi_inclusive.map_or(reader_lo, |c| c + 1);
+            .with_context(|| format!("Failed to calculate next_checkpoint for {pipeline_task}"))?;
 
         self.first_ingestion_checkpoint = next_checkpoint.min(self.first_ingestion_checkpoint);
 
@@ -412,7 +429,25 @@ where
     where
         H: concurrent::Handler<Store = S> + Send + Sync + 'static,
     {
-        let Some(next_checkpoint) = self.add_pipeline::<H>().await? else {
+        let default_next_checkpoint = self.default_next_checkpoint;
+        let Some(next_checkpoint) = self
+            .add_pipeline::<H>(async |pipeline_task, mut conn| {
+                let TrailingInitWatermark {
+                    checkpoint_hi_inclusive,
+                    reader_lo,
+                } = conn
+                    .trailing_init_watermark(
+                        &pipeline_task,
+                        TrailingInitWatermark {
+                            checkpoint_hi_inclusive: default_next_checkpoint.checked_sub(1),
+                            reader_lo: default_next_checkpoint,
+                        },
+                    )
+                    .await?;
+                Ok(checkpoint_hi_inclusive.map_or(reader_lo, |c| c + 1))
+            })
+            .await?
+        else {
             return Ok(());
         };
 
@@ -449,10 +484,6 @@ impl<T: TransactionalStore> Indexer<T> {
     where
         H: sequential::Handler<Store = T> + Send + Sync + 'static,
     {
-        let Some(next_checkpoint) = self.add_pipeline::<H>().await? else {
-            return Ok(());
-        };
-
         if self.task.is_some() {
             bail!(
                 "Sequential pipelines do not support pipeline tasks. \
@@ -461,6 +492,16 @@ impl<T: TransactionalStore> Indexer<T> {
             );
         }
 
+        let default_next_checkpoint = self.default_next_checkpoint;
+        let Some(next_checkpoint) = self
+            .add_pipeline::<H>(async |pt, conn| {
+                Self::next_checkpoint(default_next_checkpoint, pt, conn).await
+            })
+            .await?
+        else {
+            return Ok(());
+        };
+
         // Track the minimum next_checkpoint across all sequential pipelines
         self.next_sequential_checkpoint = Some(
             self.next_sequential_checkpoint
@@ -740,7 +781,7 @@ mod tests {
 
         let mut conn = store.connect().await.unwrap();
 
-        conn.init_watermark(A::NAME, InitWatermark::default())
+        conn.trailing_init_watermark(A::NAME, TrailingInitWatermark::default())
             .await
             .unwrap();
         conn.set_committer_watermark(
@@ -753,7 +794,7 @@ mod tests {
         .await
         .unwrap();
 
-        conn.init_watermark(B::NAME, InitWatermark::default())
+        conn.trailing_init_watermark(B::NAME, TrailingInitWatermark::default())
             .await
             .unwrap();
         conn.set_committer_watermark(
@@ -766,7 +807,7 @@ mod tests {
         .await
         .unwrap();
 
-        conn.init_watermark(C::NAME, InitWatermark::default())
+        conn.trailing_init_watermark(C::NAME, TrailingInitWatermark::default())
             .await
             .unwrap();
         conn.set_committer_watermark(
@@ -779,7 +820,7 @@ mod tests {
         .await
         .unwrap();
 
-        conn.init_watermark(D::NAME, InitWatermark::default())
+        conn.trailing_init_watermark(D::NAME, TrailingInitWatermark::default())
             .await
             .unwrap();
         conn.set_committer_watermark(