Update pruner to start a separate task that will update reader lower …

…bound watermarks, without doing any actual pruning.

crates/sui-indexer/src/handlers/pruner.rs

@@ -1,6 +1,7 @@
 // Copyright (c) Mysten Labs, Inc.
 // SPDX-License-Identifier: Apache-2.0
 
+use mysten_metrics::spawn_monitored_task;
 use serde::{Deserialize, Serialize};
 use std::collections::HashMap;
 use std::time::Duration;
@@ -10,6 +11,7 @@ use tracing::{error, info};
 
 use crate::config::RetentionConfig;
 use crate::errors::IndexerError;
+use crate::models::watermarks::{PrunableWatermark, StoredWatermark};
 use crate::store::pg_partition_manager::PgPartitionManager;
 use crate::store::PgIndexerStore;
 use crate::{metrics::IndexerMetrics, store::IndexerStore, types::IndexerResult};
@@ -69,6 +71,39 @@ pub enum PrunableTable {
     PrunerCpWatermark,
 }
 
+impl PrunableTable {
+    pub fn select_lower_bound(&self, cp: u64, tx: u64) -> u64 {
+        match self {
+            PrunableTable::ObjectsHistory => cp,
+            PrunableTable::Transactions => tx,
+            PrunableTable::Events => tx,
+
+            PrunableTable::EventEmitPackage => tx,
+            PrunableTable::EventEmitModule => tx,
+            PrunableTable::EventSenders => tx,
+            PrunableTable::EventStructInstantiation => tx,
+            PrunableTable::EventStructModule => tx,
+            PrunableTable::EventStructName => tx,
+            PrunableTable::EventStructPackage => tx,
+
+            PrunableTable::TxAffectedAddresses => tx,
+            PrunableTable::TxAffectedObjects => tx,
+            PrunableTable::TxCallsPkg => tx,
+            PrunableTable::TxCallsMod => tx,
+            PrunableTable::TxCallsFun => tx,
+            PrunableTable::TxChangedObjects => tx,
+            PrunableTable::TxDigests => tx,
+            PrunableTable::TxInputObjects => tx,
+            PrunableTable::TxKinds => tx,
+            PrunableTable::TxRecipients => tx,
+            PrunableTable::TxSenders => tx,
+
+            PrunableTable::Checkpoints => cp,
+            PrunableTable::PrunerCpWatermark => cp,
+        }
+    }
+}
+
 impl Pruner {
     /// Instantiates a pruner with default retention and overrides. Pruner will finalize the
     /// retention policies so there is a value for every prunable table.
@@ -101,6 +136,15 @@ impl Pruner {
     }
 
     pub async fn start(&self, cancel: CancellationToken) -> IndexerResult<()> {
+        let store_clone = self.store.clone();
+        let retention_policies = self.retention_policies.clone();
+        let cancel_clone = cancel.clone();
+        spawn_monitored_task!(update_watermarks_lower_bounds_task(
+            store_clone,
+            retention_policies,
+            cancel_clone
+        ));
+
         let mut last_seen_max_epoch = 0;
         // The first epoch that has not yet been pruned.
         let mut next_prune_epoch = None;
@@ -177,3 +221,77 @@ impl Pruner {
         Ok(())
     }
 }
+
+/// Task to periodically query the `watermarks` table and update the lower bounds for all watermarks
+/// if the entry exceeds epoch-level retention policy.
+async fn update_watermarks_lower_bounds_task(
+    store: PgIndexerStore,
+    retention_policies: HashMap<PrunableTable, u64>,
+    cancel: CancellationToken,
+) -> IndexerResult<()> {
+    let mut interval = tokio::time::interval(Duration::from_secs(5));
+    loop {
+        tokio::select! {
+            _ = cancel.cancelled() => {
+                info!("Pruner watermark lower bound update task cancelled.");
+                return Ok(());
+            }
+            _ = interval.tick() => {
+                update_watermarks_lower_bounds(&store, &retention_policies, &cancel).await?;
+            }
+        }
+    }
+}
+
+/// Fetches all entries from the `watermarks` table, and updates the lower bounds for all watermarks
+/// if the entry's epoch range exceeds the respective retention policy.
+async fn update_watermarks_lower_bounds(
+    store: &PgIndexerStore,
+    retention_policies: &HashMap<PrunableTable, u64>,
+    cancel: &CancellationToken,
+) -> IndexerResult<()> {
+    let (watermarks, latest_db_timestamp) = store.get_watermarks().await?;
+    let mut lower_bound_updates = vec![];
+
+    for watermark in watermarks.iter() {
+        if cancel.is_cancelled() {
+            info!("Pruner watermark lower bound update task cancelled.");
+            return Ok(());
+        }
+
+        let Some(watermark) = PrunableWatermark::new(watermark.clone(), latest_db_timestamp) else {
+            continue;
+        };
+
+        let Some(epochs_to_keep) = retention_policies.get(&watermark.entity) else {
+            continue;
+        };
+
+        if watermark.epoch_lo + epochs_to_keep <= watermark.epoch_hi {
+            let new_inclusive_epoch_lower_bound =
+                watermark.epoch_hi.saturating_sub(epochs_to_keep - 1);
+
+            // TODO: (wlmyng) now that epochs table is not pruned, we can add `first_tx_seq_num` or
+            // something and use it as a lookup table.
+            let (min_cp, _) = store
+                .get_checkpoint_range_for_epoch(new_inclusive_epoch_lower_bound)
+                .await?;
+            let (min_tx, _) = store.get_transaction_range_for_checkpoint(min_cp).await?;
+
+            lower_bound_updates.push(StoredWatermark::from_lower_bound_update(
+                watermark.entity.as_ref(),
+                new_inclusive_epoch_lower_bound,
+                watermark.entity.select_lower_bound(min_cp, min_tx),
+            ))
+        }
+    }
+
+    if !lower_bound_updates.is_empty() {
+        store
+            .update_watermarks_lower_bound(lower_bound_updates)
+            .await?;
+        info!("Finished updating lower bounds for watermarks");
+    }
+
+    Ok(())
+}

crates/sui-indexer/src/models/watermarks.rs

@@ -1,14 +1,17 @@
 // Copyright (c) Mysten Labs, Inc.
 // SPDX-License-Identifier: Apache-2.0
 
+use std::str::FromStr;
+
+use diesel::prelude::*;
+
 use crate::{
-    handlers::CommitterWatermark,
+    handlers::{pruner::PrunableTable, CommitterWatermark},
     schema::watermarks::{self},
 };
-use diesel::prelude::*;
 
 /// Represents a row in the `watermarks` table.
-#[derive(Queryable, Insertable, Default, QueryableByName)]
+#[derive(Queryable, Insertable, Default, QueryableByName, Clone)]
 #[diesel(table_name = watermarks, primary_key(entity))]
 pub struct StoredWatermark {
     /// The table governed by this watermark, i.e `epochs`, `checkpoints`, `transactions`.
@@ -39,6 +42,48 @@ pub struct StoredWatermark {
     pub pruned_lo: Option<i64>,
 }
 
+#[derive(Debug)]
+pub struct PrunableWatermark {
+    pub entity: PrunableTable,
+    pub epoch_hi: u64,
+    pub epoch_lo: u64,
+    pub checkpoint_hi: u64,
+    pub tx_hi: u64,
+    pub reader_lo: u64,
+    /// Timestamp when the watermark's lower bound was last updated.
+    pub timestamp_ms: i64,
+    /// Latest timestamp read from db.
+    pub current_timestamp_ms: i64,
+    /// Data at and below `pruned_lo` is considered pruned by the pruner.
+    pub pruned_lo: Option<u64>,
+}
+
+impl PrunableWatermark {
+    pub fn new(stored: StoredWatermark, latest_db_timestamp: i64) -> Option<Self> {
+        let Some(entity) = PrunableTable::from_str(&stored.entity).ok() else {
+            return None;
+        };
+
+        Some(PrunableWatermark {
+            entity,
+            epoch_hi: stored.epoch_hi as u64,
+            epoch_lo: stored.epoch_lo as u64,
+            checkpoint_hi: stored.checkpoint_hi as u64,
+            tx_hi: stored.tx_hi as u64,
+            reader_lo: stored.reader_lo as u64,
+            timestamp_ms: stored.timestamp_ms,
+            current_timestamp_ms: latest_db_timestamp,
+            pruned_lo: stored.pruned_lo.map(|lo| lo as u64),
+        })
+    }
+
+    /// Represents the first `unit` (checkpoint, tx, epoch) that has not yet been pruned. If
+    /// `pruned_lo` is not set in db, default to 0. Otherwise, this is `pruned_lo + `.
+    pub fn pruner_lo(&self) -> u64 {
+        self.pruned_lo.map_or(0, |lo| lo.saturating_add(1))
+    }
+}
+
 impl StoredWatermark {
     pub fn from_upper_bound_update(entity: &str, watermark: CommitterWatermark) -> Self {
         StoredWatermark {
@@ -49,4 +94,13 @@ impl StoredWatermark {
             ..StoredWatermark::default()
         }
     }
+
+    pub fn from_lower_bound_update(entity: &str, epoch_lo: u64, reader_lo: u64) -> Self {
+        StoredWatermark {
+            entity: entity.to_string(),
+            epoch_lo: epoch_lo as i64,
+            reader_lo: reader_lo as i64,
+            ..StoredWatermark::default()
+        }
+    }
 }

crates/sui-indexer/src/store/indexer_store.rs

@@ -12,6 +12,7 @@ use crate::models::display::StoredDisplay;
 use crate::models::obj_indices::StoredObjectVersion;
 use crate::models::objects::{StoredDeletedObject, StoredObject};
 use crate::models::raw_checkpoints::StoredRawCheckpoint;
+use crate::models::watermarks::StoredWatermark;
 use crate::types::{
     EventIndex, IndexedCheckpoint, IndexedEvent, IndexedPackage, IndexedTransaction, TxIndex,
 };
@@ -123,4 +124,12 @@ pub trait IndexerStore: Clone + Sync + Send + 'static {
     ) -> Result<(), IndexerError>
     where
         E::Iterator: Iterator<Item: AsRef<str>>;
+
+    async fn update_watermarks_lower_bound(
+        &self,
+        watermarks: Vec<StoredWatermark>,
+    ) -> Result<(), IndexerError>;
+
+    /// Load all watermark entries from the store, and the latest timestamp from the db.
+    async fn get_watermarks(&self) -> Result<(Vec<StoredWatermark>, i64), IndexerError>;
 }

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

@@ -52,3 +52,42 @@ where
     })
     .await
 }
+
+pub async fn read_with_retry<'a, Q, T>(
+    pool: &ConnectionPool,
+    timeout: Duration,
+    query: Q,
+) -> Result<T, IndexerError>
+where
+    Q: for<'r> FnOnce(
+            &'r mut AsyncPgConnection,
+        ) -> ScopedBoxFuture<'a, 'r, Result<T, IndexerError>>
+        + Send,
+    Q: Clone,
+    T: 'a,
+{
+    let backoff = backoff::ExponentialBackoff {
+        max_elapsed_time: Some(timeout),
+        ..Default::default()
+    };
+    backoff::future::retry(backoff, || async {
+        let mut connection = pool.get().await.map_err(|e| backoff::Error::Transient {
+            err: IndexerError::PostgresWriteError(e.to_string()),
+            retry_after: None,
+        })?;
+
+        connection
+            .build_transaction()
+            .read_only()
+            .run(query.clone())
+            .await
+            .map_err(|e| {
+                tracing::error!("Error with reading data from DB: {:?}, retrying...", e);
+                backoff::Error::Transient {
+                    err: IndexerError::PostgresWriteError(e.to_string()),
+                    retry_after: None,
+                }
+            })
+    })
+    .await
+}