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.