feat(indexer): Make concurrent modifications to objects and disaply table safe

tomxey · tomxey · commit a2690faea596 · 2025-04-11T16:53:18.000+02:00
diff --git a/crates/iota-indexer/migrations/pg/2025-04-10-103714_deleted_object_versions_table/down.sql b/crates/iota-indexer/migrations/pg/2025-04-10-103714_deleted_object_versions_table/down.sql
@@ -0,0 +1 @@
+DROP TABLE deleted_objects_versions;
diff --git a/crates/iota-indexer/migrations/pg/2025-04-10-103714_deleted_object_versions_table/up.sql b/crates/iota-indexer/migrations/pg/2025-04-10-103714_deleted_object_versions_table/up.sql
@@ -0,0 +1,8 @@
+-- Version of objects that already got deleted.
+-- This is to handle race conditions between optimistic indexing and checkpoint indexing.
+-- If object version is present in this table,
+-- any inserts of lower versions for this object should be skipped.
+CREATE TABLE deleted_objects_versions (
+    object_id      bytea   PRIMARY KEY,
+    object_version bigint  NOT NULL
+);
diff --git a/crates/iota-indexer/src/schema.rs b/crates/iota-indexer/src/schema.rs
@@ -64,6 +64,13 @@ diesel::table! {
     }
 }
 
+diesel::table! {
+    deleted_objects_versions (object_id) {
+        object_id -> Bytea,
+        object_version -> Int8,
+    }
+}
+
 diesel::table! {
     display (object_type) {
         object_type -> Text,
@@ -601,6 +608,7 @@ macro_rules! for_all_tables {
             addresses,
             chain_identifier,
             checkpoints,
+            deleted_objects_versions,
             display,
             epoch_peak_tps,
             epochs,
diff --git a/crates/iota-indexer/src/store/mod.rs b/crates/iota-indexer/src/store/mod.rs
@@ -166,6 +166,21 @@ pub mod diesel_macro {
         }};
     }
 
+    #[macro_export]
+    macro_rules! on_conflict_do_update_with_condition {
+        ($table:expr, $values:expr, $target:expr, $pg_columns:expr, $condition:expr, $conn:expr) => {{
+            use diesel::{ExpressionMethods, RunQueryDsl, query_dsl::methods::FilterDsl};
+
+            diesel::insert_into($table)
+                .values($values)
+                .on_conflict($target)
+                .do_update()
+                .set($pg_columns)
+                .filter($condition)
+                .execute($conn)?;
+        }};
+    }
+
     #[macro_export]
     macro_rules! run_query {
         ($pool:expr, $query:expr) => {{
diff --git a/crates/iota-indexer/src/store/pg_indexer_store.rs b/crates/iota-indexer/src/store/pg_indexer_store.rs
@@ -13,6 +13,7 @@ use async_trait::async_trait;
 use diesel::{
     ExpressionMethods, OptionalExtension, QueryDsl, RunQueryDsl,
     dsl::{max, min},
+    sql_types,
     upsert::excluded,
 };
 use downcast::Any;
@@ -48,11 +49,12 @@ use crate::{
         transactions::{OptimisticTransaction, StoredTransaction, TxInsertionOrder},
         tx_indices::OptimisticTxIndices,
     },
-    on_conflict_do_update, persist_chunk_into_table, read_only_blocking,
+    on_conflict_do_update, on_conflict_do_update_with_condition, persist_chunk_into_table,
+    read_only_blocking,
     schema::{
         chain_identifier, checkpoints, display, epochs, event_emit_module, event_emit_package,
         event_senders, event_struct_instantiation, event_struct_module, event_struct_name,
-        event_struct_package, events, feature_flags, objects, objects_history, objects_snapshot,
+        event_struct_package, events, feature_flags, objects_history, objects_snapshot,
         objects_version, optimistic_event_emit_module, optimistic_event_emit_package,
         optimistic_event_senders, optimistic_event_struct_instantiation,
         optimistic_event_struct_module, optimistic_event_struct_name,
@@ -339,7 +341,7 @@ impl PgIndexerStore {
         transactional_blocking_with_retry!(
             &self.blocking_cp,
             |conn| {
-                on_conflict_do_update!(
+                on_conflict_do_update_with_condition!(
                     display::table,
                     display_updates.values().collect::<Vec<_>>(),
                     display::object_type,
@@ -348,6 +350,7 @@ impl PgIndexerStore {
                         display::version.eq(excluded(display::version)),
                         display::bcs.eq(excluded(display::bcs)),
                     ),
+                    excluded(display::version).gt(display::version),
                     conn
                 );
                 Ok::<(), IndexerError>(())
@@ -362,6 +365,100 @@ impl PgIndexerStore {
         &self,
         mutated_object_mutation_chunk: Vec<StoredObject>,
     ) -> Result<(), IndexerError> {
+        let mut object_id_vec = Vec::new();
+        let mut object_version_vec = Vec::new();
+        let mut object_digest_vec = Vec::new();
+        let mut owner_type_vec = Vec::new();
+        let mut owner_id_vec = Vec::new();
+        let mut object_type_vec = Vec::new();
+        let mut serialized_object_vec = Vec::new();
+        let mut coin_type_vec = Vec::new();
+        let mut coin_balance_vec = Vec::new();
+        let mut df_kind_vec = Vec::new();
+
+        for obj in mutated_object_mutation_chunk.iter().cloned() {
+            object_id_vec.push(obj.object_id);
+            object_version_vec.push(obj.object_version);
+            object_digest_vec.push(obj.object_digest);
+            owner_type_vec.push(obj.owner_type);
+            owner_id_vec.push(obj.owner_id);
+            object_type_vec.push(obj.object_type);
+            serialized_object_vec.push(obj.serialized_object);
+            coin_type_vec.push(obj.coin_type);
+            coin_balance_vec.push(obj.coin_balance);
+            df_kind_vec.push(obj.df_kind);
+        }
+
+        let query = diesel::sql_query(
+            r#"
+        WITH new_data AS (
+            SELECT
+                unnest($1::bytea[]) AS object_id,
+                unnest($2::bigint[]) AS object_version,
+                unnest($3::bytea[]) AS object_digest,
+                unnest($4::smallint[]) AS owner_type,
+                unnest($5::bytea[]  ) AS owner_id,
+                unnest($6::text[]   ) AS object_type,
+                unnest($7::bytea[]) AS serialized_object,
+                unnest($8::text[] ) AS coin_type,
+                unnest($9::bigint[]) AS coin_balance,
+                unnest($10::smallint[]) AS df_kind
+        ),
+        locked_objects AS (
+            SELECT o.*
+            FROM objects o
+            JOIN new_data nd ON o.object_id = nd.object_id
+            FOR UPDATE
+        ),
+        locked_deletes AS (
+            SELECT del.*
+            FROM deleted_objects_versions del
+            JOIN new_data nd ON del.object_id = nd.object_id
+            FOR SHARE
+        )
+        INSERT INTO objects (
+            object_id,
+            object_version,
+            object_digest,
+            owner_type,
+            owner_id,
+            object_type,
+            serialized_object,
+            coin_type,
+            coin_balance,
+            df_kind
+        )
+        SELECT nd.*
+        FROM new_data nd
+        LEFT JOIN deleted_objects_versions del
+          ON del.object_id = nd.object_id
+        WHERE COALESCE(del.object_version, -1) < nd.object_version
+        ON CONFLICT (object_id)
+        DO UPDATE SET
+            object_version = EXCLUDED.object_version,
+            object_digest = EXCLUDED.object_digest,
+            owner_type = EXCLUDED.owner_type,
+            owner_id = EXCLUDED.owner_id,
+            object_type = EXCLUDED.object_type,
+            serialized_object = EXCLUDED.serialized_object,
+            coin_type = EXCLUDED.coin_type,
+            coin_balance = EXCLUDED.coin_balance,
+            df_kind = EXCLUDED.df_kind
+        WHERE
+            EXCLUDED.object_version > objects.object_version;
+    "#,
+        )
+        .bind::<sql_types::Array<sql_types::Binary>, _>(object_id_vec)
+        .bind::<sql_types::Array<sql_types::BigInt>, _>(object_version_vec)
+        .bind::<sql_types::Array<sql_types::Binary>, _>(object_digest_vec)
+        .bind::<sql_types::Array<sql_types::SmallInt>, _>(owner_type_vec)
+        .bind::<sql_types::Array<sql_types::Nullable<sql_types::Binary>>, _>(owner_id_vec)
+        .bind::<sql_types::Array<sql_types::Nullable<sql_types::Text>>, _>(object_type_vec)
+        .bind::<sql_types::Array<sql_types::Binary>, _>(serialized_object_vec)
+        .bind::<sql_types::Array<sql_types::Nullable<sql_types::Text>>, _>(coin_type_vec)
+        .bind::<sql_types::Array<sql_types::Nullable<sql_types::BigInt>>, _>(coin_balance_vec)
+        .bind::<sql_types::Array<sql_types::Nullable<sql_types::SmallInt>>, _>(df_kind_vec);
+
         let guard = self
             .metrics
             .checkpoint_db_commit_latency_objects_chunks
@@ -370,24 +467,7 @@ impl PgIndexerStore {
         transactional_blocking_with_retry!(
             &self.blocking_cp,
             |conn| {
-                on_conflict_do_update!(
-                    objects::table,
-                    mutated_object_mutation_chunk.clone(),
-                    objects::object_id,
-                    (
-                        objects::object_id.eq(excluded(objects::object_id)),
-                        objects::object_version.eq(excluded(objects::object_version)),
-                        objects::object_digest.eq(excluded(objects::object_digest)),
-                        objects::owner_type.eq(excluded(objects::owner_type)),
-                        objects::owner_id.eq(excluded(objects::owner_id)),
-                        objects::object_type.eq(excluded(objects::object_type)),
-                        objects::serialized_object.eq(excluded(objects::serialized_object)),
-                        objects::coin_type.eq(excluded(objects::coin_type)),
-                        objects::coin_balance.eq(excluded(objects::coin_balance)),
-                        objects::df_kind.eq(excluded(objects::df_kind)),
-                    ),
-                    conn
-                );
+                query.clone().execute(conn)?;
                 Ok::<(), IndexerError>(())
             },
             PG_DB_COMMIT_SLEEP_DURATION
@@ -410,23 +490,53 @@ impl PgIndexerStore {
             .checkpoint_db_commit_latency_objects_chunks
             .start_timer();
         let len = deleted_objects_chunk.len();
+
+        let (object_id_vec, object_version_vec): (Vec<_>, Vec<_>) = deleted_objects_chunk
+            .into_iter()
+            .map(|obj| (obj.object_id, obj.object_version))
+            .unzip();
+
+        let query = diesel::sql_query(
+            r#"
+        WITH new_data AS (
+            SELECT
+                unnest($1::bytea[])  AS object_id,
+                unnest($2::bigint[]) AS object_version
+        ),
+        locked_objects AS (
+            SELECT o.*
+            FROM objects o
+            JOIN new_data nd ON o.object_id = nd.object_id
+            FOR UPDATE
+        ),
+        locked_deletes AS (
+            SELECT del.*
+            FROM deleted_objects_versions del
+            JOIN new_data nd ON del.object_id = nd.object_id
+            FOR UPDATE
+        ),
+        deleted AS (
+            DELETE FROM objects o
+            USING new_data nd
+            WHERE o.object_id = nd.object_id
+              AND nd.object_version > o.object_version
+        )
+        INSERT INTO deleted_objects_versions (object_id, object_version)
+        SELECT object_id, object_version
+        FROM new_data
+        ON CONFLICT (object_id)
+        DO UPDATE
+        SET object_version = EXCLUDED.object_version
+        WHERE EXCLUDED.object_version > deleted_objects_versions.object_version;
+    "#,
+        )
+        .bind::<sql_types::Array<sql_types::Bytea>, _>(object_id_vec)
+        .bind::<sql_types::Array<sql_types::BigInt>, _>(object_version_vec);
+
         transactional_blocking_with_retry!(
             &self.blocking_cp,
             |conn| {
-                diesel::delete(
-                    objects::table.filter(
-                        objects::object_id.eq_any(
-                            deleted_objects_chunk
-                                .iter()
-                                .map(|o| o.object_id.clone())
-                                .collect::<Vec<_>>(),
-                        ),
-                    ),
-                )
-                .execute(conn)
-                .map_err(IndexerError::from)
-                .context("Failed to write object deletion to PostgresDB")?;
-
+                query.clone().execute(conn)?;
                 Ok::<(), IndexerError>(())
             },
             PG_DB_COMMIT_SLEEP_DURATION
@@ -677,17 +787,17 @@ impl PgIndexerStore {
             },
             PG_DB_COMMIT_SLEEP_DURATION
         )
-        .tap_ok(|_| {
-            let elapsed = guard.stop_and_record();
-            info!(
+            .tap_ok(|_| {
+                let elapsed = guard.stop_and_record();
+                info!(
                 elapsed,
                 "Persisted {} checkpoints",
                 stored_checkpoints.len()
             );
-        })
-        .tap_err(|e| {
-            tracing::error!("Failed to persist checkpoints with error: {}", e);
-        })
+            })
+            .tap_err(|e| {
+                tracing::error!("Failed to persist checkpoints with error: {}", e);
+            })
     }
 
     fn persist_transactions_chunk(
