feat(indexer): Introduce optimistic indexing

Author: tomxey

crates/iota-indexer/src/apis/write_api.rs

@@ -2,48 +2,68 @@
 // Modifications Copyright (c) 2024 IOTA Stiftung
 // SPDX-License-Identifier: Apache-2.0
 
+use std::{collections::HashSet, time::Duration};
+
+use anyhow::Context;
 use async_trait::async_trait;
+use diesel::{ExpressionMethods, QueryDsl, RunQueryDsl, SelectableHelper};
+use downcast::Any;
 use fastcrypto::{encoding::Base64, traits::ToFromBytes};
 use iota_json_rpc::IotaRpcModule;
-use iota_json_rpc_api::{error_object_from_rpc, WriteApiClient, WriteApiServer};
+use iota_json_rpc_api::{WriteApiClient, WriteApiServer, error_object_from_rpc};
 use iota_json_rpc_types::{
     DevInspectArgs, DevInspectResults, DryRunTransactionBlockResponse,
     IotaTransactionBlockResponse, IotaTransactionBlockResponseOptions,
 };
 use iota_open_rpc::Module;
-use iota_rest_api::{client::TransactionExecutionResponse, ExecuteTransactionQueryParameters};
+use iota_rest_api::{ExecuteTransactionQueryParameters, client::TransactionExecutionResponse};
 use iota_types::{
     base_types::IotaAddress,
     effects::TransactionEffectsAPI,
+    execution_status::ExecutionStatus,
     full_checkpoint_content::CheckpointTransaction,
     iota_serde::BigInt,
     quorum_driver_types::ExecuteTransactionRequestType,
     signature::GenericSignature,
     transaction::{Transaction, TransactionData},
 };
-use jsonrpsee::{core::RpcResult, http_client::HttpClient, RpcModule};
+use jsonrpsee::{RpcModule, core::RpcResult, http_client::HttpClient};
 
 use crate::{
-    errors::IndexerError, indexer_reader::IndexerReader,
-    types::IotaTransactionBlockResponseWithOptions,
+    errors::IndexerError,
+    handlers::{TransactionObjectChangesToCommit, checkpoint_handler::CheckpointHandler},
+    indexer_reader::IndexerReader,
+    metrics::IndexerMetrics,
+    models::transactions::{StoredTransaction, TxInsertionOrder},
+    run_query_async,
+    schema::tx_insertion_order,
+    spawn_read_only_blocking,
+    store::{IndexerStore, PgIndexerStore},
+    types::{IndexedDeletedObject, IndexedObject, IotaTransactionBlockResponseWithOptions},
 };
 
 pub(crate) struct WriteApi {
     fullnode: HttpClient,
     fullnode_rest_client: iota_rest_api::Client,
     inner: IndexerReader,
+    store: PgIndexerStore,
+    metrics: IndexerMetrics,
 }
 
 impl WriteApi {
     pub fn new(
         fullnode_client: HttpClient,
         fullnode_rest_client: iota_rest_api::Client,
         inner: IndexerReader,
+        store: PgIndexerStore,
+        metrics: IndexerMetrics,
     ) -> Self {
         Self {
             fullnode: fullnode_client,
             fullnode_rest_client,
             inner,
+            store,
+            metrics,
         }
     }
 
@@ -101,6 +121,7 @@ impl WriteApi {
         };
 
         println!("{:#?}", full_tx_data);
+        self.optimistically_index_transaction(&full_tx_data).await?;
 
         // TODO: optimistically index the TX, so we can read it in the lines below right
         // away, without looping
@@ -127,6 +148,84 @@ impl WriteApi {
         }
         .into())
     }
+
+    async fn optimistically_index_transaction(
+        &self,
+        full_tx_data: &CheckpointTransaction,
+    ) -> Result<(), IndexerError> {
+        let assigned_insertion_order: TxInsertionOrder = {
+            let tx_digest_bytes = full_tx_data.transaction.digest().inner().to_vec();
+
+            // // how do we return info about failed txs if we don't index them?
+            // if *full_tx_data.effects.status() == ExecutionStatus::Success {
+
+            self.store
+                .persist_tx_insertion_order(vec![TxInsertionOrder {
+                    insertion_order: -1, // ignored value
+                    tx_digest: tx_digest_bytes.clone(),
+                }])
+                .await?;
+
+            let pool = self.inner.get_pool();
+            run_query_async!(&pool, |conn| {
+                tx_insertion_order::table
+                    .select(TxInsertionOrder::as_select())
+                    .filter(tx_insertion_order::tx_digest.eq(tx_digest_bytes))
+                    .first::<TxInsertionOrder>(conn)
+            })?
+        };
+
+        let object_changes = {
+            let indexed_eventually_removed_objects = full_tx_data
+                .removed_object_refs_post_version()
+                .map(|obj_ref| IndexedDeletedObject {
+                    object_id: obj_ref.0,
+                    object_version: obj_ref.1.into(),
+                    checkpoint_sequence_number: 0,
+                })
+                .collect::<Vec<_>>();
+
+            let changed_objects = full_tx_data
+                .output_objects
+                .iter()
+                .map(|o| {
+                    crate::handlers::checkpoint_handler::try_extract_df_kind(o)
+                        .map(|df_kind| IndexedObject::from_object(0, o.clone(), df_kind))
+                })
+                .collect::<Result<Vec<_>, _>>()?;
+
+            let output_objects_ids = changed_objects
+                .iter()
+                .map(|o| o.object.id())
+                .collect::<HashSet<_>>();
+            assert!(
+                indexed_eventually_removed_objects
+                    .iter()
+                    .all(|o| !output_objects_ids.contains(&o.object_id))
+            );
+
+            TransactionObjectChangesToCommit {
+                changed_objects,
+                deleted_objects: indexed_eventually_removed_objects,
+            }
+        };
+
+        let (indexed_tx, tx_indices, indexed_events, events_indices, indexed_displays) =
+            CheckpointHandler::index_transaction(
+                &full_tx_data,
+                assigned_insertion_order.insertion_order.try_into().unwrap(),
+                0,
+                0,
+                &self.metrics,
+            )
+            .await?;
+
+        self.store.persist_objects(vec![object_changes]).await?; // this will insert, as well as remove from 'objects' table, is it ok?
+        self.store.persist_displays(indexed_displays).await?;
+        self.store
+            .persist_optimistic_transaction(StoredTransaction::from(&indexed_tx).into())
+            .await
+    }
 }
 
 #[async_trait]
@@ -143,7 +242,9 @@ impl WriteApiServer for WriteApi {
                 .fullnode
                 .execute_transaction_block(tx_bytes, signatures, options.clone(), request_type)
                 .await
-                .map_err(error_object_from_rpc)?,
+                .map_err(error_object_from_rpc)?, // should this be the default option? What if
+            // fullnode returns that the tx is locally
+            // executed?
             Some(ExecuteTransactionRequestType::WaitForLocalExecution) => {
                 self.execute_and_optimistically_index_tx_effects(
                     tx_bytes,

crates/iota-indexer/src/db.rs

@@ -307,7 +307,15 @@ pub mod setup_postgres {
             let store = PgIndexerStore::new(blocking_cp, indexer_metrics.clone());
             return Indexer::start_writer(&indexer_config, store, indexer_metrics).await;
         } else if indexer_config.rpc_server_worker {
-            return Indexer::start_reader(&indexer_config, &registry, db_url.to_string()).await;
+            let store = PgIndexerStore::new(blocking_cp, indexer_metrics.clone());
+            return Indexer::start_reader(
+                &indexer_config,
+                store,
+                &registry,
+                db_url.to_string(),
+                indexer_metrics,
+            )
+            .await;
         } else if indexer_config.analytical_worker {
             let store = PgIndexerAnalyticalStore::new(blocking_cp);
             return Indexer::start_analytical_worker(store, indexer_metrics.clone()).await;