iota_indexer/store/
pg_partition_manager.rs

1// Copyright (c) Mysten Labs, Inc.
2// Modifications Copyright (c) 2024 IOTA Stiftung
3// SPDX-License-Identifier: Apache-2.0
4
5use std::{
6    collections::{BTreeMap, HashMap},
7    time::Duration,
8};
9
10use diesel::{
11    QueryableByName, RunQueryDsl,
12    sql_types::{BigInt, VarChar},
13};
14use downcast::Any;
15use tracing::{error, info};
16
17use crate::{
18    db::ConnectionPool, errors::IndexerError, ingestion::primary::persist::EpochToCommit,
19    models::epoch::StoredEpochInfo, store::diesel_macro::*,
20};
21
22const GET_PARTITION_SQL: &str = {
23    r"
24SELECT parent.relname                                            AS table_name,
25       MIN(CAST(SUBSTRING(child.relname FROM '\d+$') AS BIGINT)) AS first_partition,
26       MAX(CAST(SUBSTRING(child.relname FROM '\d+$') AS BIGINT)) AS last_partition
27FROM pg_inherits
28         JOIN pg_class parent ON pg_inherits.inhparent = parent.oid
29         JOIN pg_class child ON pg_inherits.inhrelid = child.oid
30         JOIN pg_namespace nmsp_parent ON nmsp_parent.oid = parent.relnamespace
31         JOIN pg_namespace nmsp_child ON nmsp_child.oid = child.relnamespace
32WHERE parent.relkind = 'p'
33GROUP BY table_name;
34"
35};
36
37pub struct PgPartitionManager {
38    cp: ConnectionPool,
39    partition_strategies: HashMap<&'static str, PgPartitionStrategy>,
40}
41
42impl Clone for PgPartitionManager {
43    fn clone(&self) -> PgPartitionManager {
44        Self {
45            cp: self.cp.clone(),
46            partition_strategies: self.partition_strategies.clone(),
47        }
48    }
49}
50
51#[derive(Clone, Copy)]
52pub enum PgPartitionStrategy {
53    CheckpointSequenceNumber,
54    TxSequenceNumber,
55    ObjectId,
56}
57
58impl PgPartitionStrategy {
59    pub fn is_epoch_partitioned(&self) -> bool {
60        matches!(
61            self,
62            Self::CheckpointSequenceNumber | Self::TxSequenceNumber
63        )
64    }
65}
66
67#[derive(Clone, Debug)]
68pub struct EpochPartitionData {
69    last_epoch: u64,
70    next_epoch: u64,
71    last_epoch_start_cp: u64,
72    next_epoch_start_cp: u64,
73    last_epoch_start_tx: u64,
74    next_epoch_start_tx: u64,
75}
76
77impl EpochPartitionData {
78    pub fn compose_data(epoch: EpochToCommit, last_db_epoch: StoredEpochInfo) -> Self {
79        let last_epoch = last_db_epoch.epoch as u64;
80        let last_epoch_start_cp = last_db_epoch.first_checkpoint_id as u64;
81        let next_epoch = epoch.new_epoch.epoch as u64;
82        let next_epoch_start_cp = epoch.new_epoch.first_checkpoint_id as u64;
83
84        let next_epoch_start_tx = epoch.new_epoch.first_tx_sequence_number as u64;
85        let last_epoch_start_tx = last_db_epoch.first_tx_sequence_number as u64;
86
87        Self {
88            last_epoch,
89            next_epoch,
90            last_epoch_start_cp,
91            next_epoch_start_cp,
92            last_epoch_start_tx,
93            next_epoch_start_tx,
94        }
95    }
96}
97
98impl PgPartitionManager {
99    pub fn new(cp: ConnectionPool) -> Self {
100        let mut partition_strategies = HashMap::new();
101        partition_strategies.insert("events", PgPartitionStrategy::TxSequenceNumber);
102        partition_strategies.insert("transactions", PgPartitionStrategy::TxSequenceNumber);
103        partition_strategies.insert("objects_version", PgPartitionStrategy::ObjectId);
104        Self {
105            cp,
106            partition_strategies,
107        }
108    }
109
110    pub fn get_table_partitions(&self) -> Result<BTreeMap<String, (u64, u64)>, IndexerError> {
111        #[derive(QueryableByName, Debug, Clone)]
112        struct PartitionedTable {
113            #[diesel(sql_type = VarChar)]
114            table_name: String,
115            #[diesel(sql_type = BigInt)]
116            first_partition: i64,
117            #[diesel(sql_type = BigInt)]
118            last_partition: i64,
119        }
120
121        Ok(
122            read_only_blocking!(&self.cp, |conn| diesel::RunQueryDsl::load(
123                diesel::sql_query(GET_PARTITION_SQL),
124                conn
125            ))?
126            .into_iter()
127            .map(|table: PartitionedTable| {
128                (
129                    table.table_name,
130                    (table.first_partition as u64, table.last_partition as u64),
131                )
132            })
133            .collect(),
134        )
135    }
136
137    /// Tries to fetch the partitioning strategy for the given partitioned
138    /// table. Defaults to `CheckpointSequenceNumber` as the majority of our
139    /// tables are partitioned on an epoch's checkpoints today.
140    pub fn get_strategy(&self, table_name: &str) -> PgPartitionStrategy {
141        self.partition_strategies
142            .get(table_name)
143            .copied()
144            .unwrap_or(PgPartitionStrategy::CheckpointSequenceNumber)
145    }
146
147    pub fn determine_epoch_partition_range(
148        &self,
149        table_name: &str,
150        data: &EpochPartitionData,
151    ) -> Option<(u64, u64)> {
152        match self.get_strategy(table_name) {
153            PgPartitionStrategy::CheckpointSequenceNumber => {
154                Some((data.last_epoch_start_cp, data.next_epoch_start_cp))
155            }
156            PgPartitionStrategy::TxSequenceNumber => {
157                Some((data.last_epoch_start_tx, data.next_epoch_start_tx))
158            }
159            PgPartitionStrategy::ObjectId => None,
160        }
161    }
162
163    fn advance_partition(
164        &self,
165        table: &str,
166        last_partition: u64,
167        next_partition: u64,
168        last_partition_start: u64,
169        next_partition_start: u64,
170    ) -> Result<(), IndexerError> {
171        transactional_blocking_with_retry!(
172            &self.cp,
173            |conn| {
174                RunQueryDsl::execute(
175                    diesel::sql_query("CALL advance_partition($1, $2, $3, $4, $5)")
176                        .bind::<diesel::sql_types::Text, _>(table)
177                        .bind::<diesel::sql_types::BigInt, _>(last_partition as i64)
178                        .bind::<diesel::sql_types::BigInt, _>(next_partition as i64)
179                        .bind::<diesel::sql_types::BigInt, _>(last_partition_start as i64)
180                        .bind::<diesel::sql_types::BigInt, _>(next_partition_start as i64),
181                    conn,
182                )
183            },
184            Duration::from_secs(10)
185        )?;
186        Ok(())
187    }
188
189    pub fn advance_epoch(
190        &self,
191        table: String,
192        last_partition: u64,
193        data: &EpochPartitionData,
194    ) -> Result<(), IndexerError> {
195        let Some((last_epoch_start, next_epoch_start)) =
196            self.determine_epoch_partition_range(&table, data)
197        else {
198            return Ok(());
199        };
200        if data.next_epoch == 0 {
201            info!("Epoch 0 partition has been created in the initial setup.");
202            return Ok(());
203        }
204        if last_partition == data.last_epoch {
205            self.advance_partition(
206                &table,
207                data.last_epoch,
208                data.next_epoch,
209                last_epoch_start,
210                next_epoch_start,
211            )?;
212            info!(
213                "Advanced epoch partition for table {} from {} to {}, prev partition upper bound {}",
214                table, last_partition, data.next_epoch, next_epoch_start
215            );
216        } else if last_partition == 0 {
217            // this applies to instances that were restored from a formal snapshot before
218            // v1.31
219            info!("Healing epoch-based partitions. This might take some time.");
220            self.advance_partition(&table, last_partition, data.next_epoch, 0, next_epoch_start)?;
221            info!(
222                "Healed epoch partition for table {} from {} to {}, prev partition upper bound {}",
223                table, last_partition, data.next_epoch, next_epoch_start
224            );
225        } else if last_partition != data.next_epoch {
226            error!(
227                "Advancing to epoch {} failed. Corrupted partitions for table {table}",
228                data.last_epoch
229            );
230            return Err(IndexerError::PostgresWrite);
231        } else {
232            info!(
233                "Epoch has been advanced to {} already, skipping.",
234                data.next_epoch
235            );
236        }
237        Ok(())
238    }
239
240    pub fn drop_table_partition(&self, table: String, partition: u64) -> Result<(), IndexerError> {
241        transactional_blocking_with_retry!(
242            &self.cp,
243            |conn| {
244                RunQueryDsl::execute(
245                    diesel::sql_query("CALL drop_partition($1, $2)")
246                        .bind::<diesel::sql_types::Text, _>(table.clone())
247                        .bind::<diesel::sql_types::BigInt, _>(partition as i64),
248                    conn,
249                )
250            },
251            Duration::from_secs(10)
252        )?;
253        Ok(())
254    }
255}
iota_indexer/store/pg_partition_manager.rs

iota_indexer/store/
pg_partition_manager.rs
