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

@@ -3,9 +3,11 @@
 
 use crate::config::RetentionPolicies;
 use crate::errors::IndexerError;
+use crate::models::watermarks::StoredWatermark;
 use crate::store::pg_partition_manager::PgPartitionManager;
 use crate::store::PgIndexerStore;
 use crate::{metrics::IndexerMetrics, store::IndexerStore, types::IndexerResult};
+use mysten_metrics::spawn_monitored_task;
 use serde::{Deserialize, Serialize};
 use std::collections::HashMap;
 use std::time::Duration;
@@ -108,6 +110,15 @@ impl Pruner {
  }
 
  pub async fn start(&self, cancel: CancellationToken) -> IndexerResult<()> {
+ let store_clone = self.store.clone();
+ let retention_policies = self.retention_policies.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;
@@ -185,3 +196,73 @@ 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 = 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(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.map_to_reader_unit(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,11 +1,17 @@
 // Copyright (c) Mysten Labs, Inc.
 // SPDX-License-Identifier: Apache-2.0
 
-use crate::schema::watermarks::{self};
+use std::str::FromStr;
+
+use crate::{
+ errors::IndexerError,
+ handlers::pruner::PrunableTable,
+ 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`.
@@ -38,6 +44,51 @@ pub struct StoredWatermark {
  pub pruned_lo: Option<i64>,
 }
 
+#[derive(Debug)]
+pub struct WatermarkRead {
+ pub entity: PrunableTable,
+ pub epoch_hi: u64,
+ pub epoch_lo: u64,
+ pub checkpoint_hi: u64,
+ pub reader_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 WatermarkRead {
+ pub fn new(stored: StoredWatermark, latest_db_timestamp: i64) -> Result<Self, IndexerError> {
+ let entity = PrunableTable::from_str(&stored.entity).map_err(|e| {
+ IndexerError::PersistentStorageDataCorruptionError(format!(
+ "Unexpected entity in watermarks table: {}",
+ e
+ ))
+ })?;
+
+ Ok(WatermarkRead {
+ entity,
+ epoch_hi: stored.epoch_hi as u64,
+ epoch_lo: stored.epoch_lo as u64,
+ checkpoint_hi: stored.checkpoint_hi as u64,
+ reader_hi: stored.reader_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,
@@ -53,4 +104,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, WatermarkRead};
 use crate::types::{
  EventIndex, IndexedCheckpoint, IndexedEvent, IndexedPackage, IndexedTransaction, TxIndex,
 };
@@ -123,4 +124,13 @@ pub trait IndexerStore: Clone + Sync + Send + 'static {
  cp: u64,
  tx: u64,
  ) -> Result<(), IndexerError>;
+
+ async fn update_watermarks_lower_bound(
+ &self,
+ watermarks: Vec<StoredWatermark>,
+ ) -> Result<(), IndexerError>;
+
+ /// Load watermarks from the store. Watermarks that are not recognized by `PrunableTable` are
+ /// filtered out.
+ async fn get_watermarks(&self) -> Result<Vec<WatermarkRead>, 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
+}