Skip to main content

sui_indexer_alt_framework/postgres/
handler.rs

1// Copyright (c) Mysten Labs, Inc.
2// SPDX-License-Identifier: Apache-2.0
3
4//! Postgres-specific handler trait for concurrent indexing pipelines.
5//!
6//! This module provides an interface for handlers that need to respect
7//! PostgreSQL's bind parameter limit (32,767 parameters per query). When inserting multiple rows,
8//! each field becomes a bind parameter, so the maximum number of rows per batch is:
9//!
10//! ```text
11//! max_rows = 32,767 / fields_per_row
12//! ```
13//!
14//! The `Handler` trait in this module extends the framework's concurrent Handler trait with
15//! Postgres-specific batching logic that automatically handles this limitation.
16//!
17//! # TODO: Consider moving FieldCount to sui-pg-db
18//!
19//! Currently, FieldCount lives in this framework crate but is Postgres-specific. Ideally it should
20//! live in sui-pg-db. However, this creates a circular dependency:
21//! - sui-indexer-alt-framework depends on sui-pg-db (for IndexerCluster and other utilities)
22//! - This blanket impl needs FieldCount to implement concurrent::Handler for postgres indexers
23//! - Moving FieldCount to sui-pg-db would require framework to depend on sui-pg-db (circular!)
24//!
25//! To fully decouple, we'd need to move all postgres-specific code (including IndexerCluster) to
26//! sui-pg-db, which would be a much larger breaking change. Consider this for a future refactor.
27//!
28//! See: <https://github.com/MystenLabs/sui/pull/24055#issuecomment-3471278182>
29
30use async_trait::async_trait;
31
32use crate::pipeline::Processor;
33use crate::pipeline::concurrent;
34use crate::postgres::Connection;
35use crate::postgres::Db;
36use crate::postgres::FieldCount;
37
38/// Postgres-specific handler trait for concurrent indexing pipelines.
39///
40/// The trait automatically implements the framework's Handler trait with a PgBatch that respects
41/// the 32,767 bind parameter limit.
42#[async_trait]
43pub trait Handler: Processor<Value: FieldCount> {
44    /// If at least this many rows are pending, the committer will commit them eagerly.
45    const MIN_EAGER_ROWS: usize = 50;
46
47    /// If there are more than this many rows pending, the committer applies backpressure.
48    const MAX_PENDING_ROWS: usize = 5000;
49
50    /// The maximum number of watermarks that can show up in a single batch.
51    const MAX_WATERMARK_UPDATES: usize = 10_000;
52
53    /// Take a chunk of values and commit them to the database, returning the number of rows
54    /// affected.
55    async fn commit<'a>(values: &[Self::Value], conn: &mut Connection<'a>)
56    -> anyhow::Result<usize>;
57
58    /// Clean up data between checkpoints `_from` and `_to_exclusive` (exclusive) in the database,
59    /// returning the number of rows affected. This function is optional, and defaults to not
60    /// pruning at all.
61    async fn prune<'a>(
62        &self,
63        _from: u64,
64        _to_exclusive: u64,
65        _conn: &mut Connection<'a>,
66    ) -> anyhow::Result<usize> {
67        Ok(0)
68    }
69}
70
71/// Calculate the maximum number of rows that can be inserted in a single batch,
72/// given the number of fields per row.
73const fn max_chunk_rows<T: FieldCount>() -> usize {
74    match (i16::MAX as usize).checked_div(T::FIELD_COUNT) {
75        Some(rows) => rows,
76        None => i16::MAX as usize,
77    }
78}
79
80/// Blanket implementation of the framework's Handler trait for any type implementing the
81/// Postgres-specific Handler trait.
82#[async_trait]
83impl<H> concurrent::Handler for H
84where
85    H: Handler,
86    H::Value: FieldCount + Send + Sync,
87{
88    type Store = Db;
89    type Batch = Vec<H::Value>;
90
91    const MIN_EAGER_ROWS: usize = H::MIN_EAGER_ROWS;
92    const MAX_PENDING_ROWS: usize = H::MAX_PENDING_ROWS;
93    const MAX_WATERMARK_UPDATES: usize = H::MAX_WATERMARK_UPDATES;
94
95    fn batch(
96        &self,
97        batch: &mut Self::Batch,
98        values: &mut std::vec::IntoIter<Self::Value>,
99    ) -> crate::pipeline::concurrent::BatchStatus {
100        let max_chunk_rows = max_chunk_rows::<H::Value>();
101        let current_len = batch.len();
102
103        if current_len + values.len() > max_chunk_rows {
104            // Batch would exceed the limit, take only what fits
105            let remaining_capacity = max_chunk_rows - current_len;
106            batch.extend(values.take(remaining_capacity));
107            crate::pipeline::concurrent::BatchStatus::Ready
108        } else {
109            // All values fit, take them all
110            batch.extend(values);
111            crate::pipeline::concurrent::BatchStatus::Pending
112        }
113    }
114
115    async fn commit<'a>(
116        &self,
117        batch: &Self::Batch,
118        conn: &mut <Self::Store as crate::store::Store>::Connection<'a>,
119    ) -> anyhow::Result<usize> {
120        H::commit(batch, conn).await
121    }
122
123    async fn prune<'a>(
124        &self,
125        from: u64,
126        to_exclusive: u64,
127        conn: &mut <Self::Store as crate::store::Store>::Connection<'a>,
128    ) -> anyhow::Result<usize> {
129        <H as Handler>::prune(self, from, to_exclusive, conn).await
130    }
131}