feat(indexer): Introduce optimistic indexing

Author: tomxey

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

@@ -2,32 +2,358 @@
 // Modifications Copyright (c) 2024 IOTA Stiftung
 // SPDX-License-Identifier: Apache-2.0
 
+use std::collections::HashSet;
+
 use async_trait::async_trait;
-use fastcrypto::encoding::Base64;
+use diesel::{ExpressionMethods, QueryDsl, RunQueryDsl, SelectableHelper};
+use fastcrypto::{encoding::Base64, traits::ToFromBytes};
 use iota_json_rpc::IotaRpcModule;
 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::{ExecuteTransactionQueryParameters, client::TransactionExecutionResponse};
 use iota_types::{
-    base_types::IotaAddress, iota_serde::BigInt, quorum_driver_types::ExecuteTransactionRequestType,
+    base_types::IotaAddress,
+    digests::TransactionDigest,
+    effects::TransactionEffectsAPI,
+    full_checkpoint_content::CheckpointTransaction,
+    iota_serde::BigInt,
+    quorum_driver_types::ExecuteTransactionRequestType,
+    signature::GenericSignature,
+    transaction::{Transaction, TransactionData},
 };
 use jsonrpsee::{RpcModule, core::RpcResult, http_client::HttpClient};
 
-use crate::types::IotaTransactionBlockResponseWithOptions;
+use crate::{
+    errors::IndexerError,
+    handlers::{TransactionObjectChangesToCommit, checkpoint_handler::CheckpointHandler},
+    indexer_reader::IndexerReader,
+    metrics::IndexerMetrics,
+    models::{
+        event_indices::OptimisticEventIndices,
+        events::StoredEvent,
+        transactions::{StoredTransaction, TxInsertionOrder},
+        tx_indices::OptimisticTxIndices,
+    },
+    run_query_async,
+    schema::tx_insertion_order,
+    spawn_read_only_blocking,
+    store::{IndexerStore, PgIndexerStore},
+    types::{
+        EventIndex, IndexedDeletedObject, IndexedObject, IotaTransactionBlockResponseWithOptions,
+        TxIndex,
+    },
+};
 
 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) -> Self {
+    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,
         }
     }
+
+    async fn execute_and_optimistically_index_tx_effects(
+        &self,
+        tx_bytes: Base64,
+        signatures: Vec<Base64>,
+        options: Option<IotaTransactionBlockResponseOptions>,
+    ) -> RpcResult<IotaTransactionBlockResponse> {
+        let tx_data: TransactionData =
+            bcs::from_bytes(&tx_bytes.to_vec().map_err(IndexerError::FastCrypto)?)
+                .map_err(IndexerError::Bcs)?;
+        let mut sigs = Vec::new();
+        for sig in signatures {
+            sigs.push(
+                GenericSignature::from_bytes(&sig.to_vec().map_err(IndexerError::FastCrypto)?)
+                    .map_err(IndexerError::FastCrypto)?,
+            );
+        }
+        let transaction = Transaction::from_generic_sig_data(tx_data, sigs);
+
+        // TODO: shouldn't return type below be from rust-sdk types? Is this type
+        // correct?
+        let result = self
+            .fullnode_rest_client
+            .execute_transaction(
+                &ExecuteTransactionQueryParameters {
+                    events: true,
+                    balance_changes: false,
+                    input_objects: true,
+                    output_objects: true,
+                },
+                &transaction,
+            )
+            .await
+            .map_err(|e| IndexerError::Generic(e.to_string()))?;
+
+        // println!("RESULT: {:#?}", result);
+
+        let TransactionExecutionResponse {
+            effects,
+            finality: _,
+            events,
+            balance_changes: _,
+            input_objects,
+            output_objects,
+        } = result;
+
+        let tx_digest = match (input_objects, output_objects) {
+            (Some(input_objects), Some(output_objects)) => {
+                let full_tx_data = CheckpointTransaction {
+                    transaction,
+                    effects,
+                    events,
+                    input_objects,
+                    output_objects,
+                };
+                self.optimistically_index_transaction(&full_tx_data).await?;
+                *full_tx_data.effects.transaction_digest()
+            }
+            (_, _) => {
+                // todo: this could be simplified probably
+                let tx_digest = *effects.transaction_digest();
+                tx_digest
+            }
+        };
+
+        let tx_block_response = self
+            .wait_for_and_return_tx_block_response(tx_digest, options.clone())
+            .await?;
+
+        Ok(IotaTransactionBlockResponseWithOptions {
+            response: tx_block_response,
+            options: options.unwrap_or_default(),
+        }
+        .into())
+    }
+
+    async fn wait_for_and_return_tx_block_response(
+        &self,
+        tx_digest: TransactionDigest,
+        options: Option<IotaTransactionBlockResponseOptions>,
+    ) -> RpcResult<IotaTransactionBlockResponse> {
+        loop {
+            let tx_block_response = self
+                .inner
+                .multi_get_transaction_block_response_in_blocking_task(
+                    vec![tx_digest],
+                    options.clone().unwrap_or_default(),
+                )
+                .await? // the error returned here could also be a reason for retry I suppose (not fully
+                // indexed tx?)
+                .pop();
+
+            match tx_block_response {
+                Some(tx_block_response) => return Ok(tx_block_response),
+                None => continue, // TODO: implement some timeout, and delay between retries
+            }
+        }
+    }
+
+    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?;
+
+        // TODO: we write common tables here, we need to make sure we do not overwrite
+        // objects coming from checkpoint sync
+        // TODO: both insert and delete need to be handled
+        self.store.persist_objects(vec![object_changes]).await?; // this will insert, as well as remove from 'objects' table, is it ok?
+        // TODO: potential race conditions here also need to be handled
+        self.store.persist_displays(indexed_displays).await?;
+
+        self.store
+            .persist_optimistic_transaction(StoredTransaction::from(&indexed_tx).into())
+            .await?;
+        self.store
+            .persist_optimistic_events(
+                indexed_events
+                    .into_iter()
+                    .map(StoredEvent::from)
+                    .map(Into::into)
+                    .collect(),
+            )
+            .await?;
+        self.store
+            .persist_optimistic_event_indices(
+                Self::optimistic_event_indices_from_indexed_event_indices(events_indices),
+            )
+            .await?;
+        self.store
+            .persist_optimistic_tx_indices(Self::optimistic_tx_indices_from_indexed_tx_indices(
+                tx_indices,
+            ))
+            .await?;
+
+        Ok(())
+    }
+
+    fn optimistic_event_indices_from_indexed_event_indices(
+        event_indices: Vec<EventIndex>,
+    ) -> OptimisticEventIndices {
+        let (
+            event_emit_packages,
+            event_emit_modules,
+            event_senders,
+            event_struct_packages,
+            event_struct_modules,
+            event_struct_names,
+            event_struct_instantiations,
+        ) = event_indices.into_iter().map(|i| i.split()).fold(
+            (
+                Vec::new(),
+                Vec::new(),
+                Vec::new(),
+                Vec::new(),
+                Vec::new(),
+                Vec::new(),
+                Vec::new(),
+            ),
+            |(
+                mut event_emit_packages,
+                mut event_emit_modules,
+                mut event_senders,
+                mut event_struct_packages,
+                mut event_struct_modules,
+                mut event_struct_names,
+                mut event_struct_instantiations,
+            ),
+             index| {
+                event_emit_packages.push(index.0);
+                event_emit_modules.push(index.1);
+                event_senders.push(index.2);
+                event_struct_packages.push(index.3);
+                event_struct_modules.push(index.4);
+                event_struct_names.push(index.5);
+                event_struct_instantiations.push(index.6);
+                (
+                    event_emit_packages,
+                    event_emit_modules,
+                    event_senders,
+                    event_struct_packages,
+                    event_struct_modules,
+                    event_struct_names,
+                    event_struct_instantiations,
+                )
+            },
+        );
+
+        (
+            event_emit_packages.into_iter().map(Into::into).collect(),
+            event_emit_modules.into_iter().map(Into::into).collect(),
+            event_senders.into_iter().map(Into::into).collect(),
+            event_struct_packages.into_iter().map(Into::into).collect(),
+            event_struct_modules.into_iter().map(Into::into).collect(),
+            event_struct_names.into_iter().map(Into::into).collect(),
+            event_struct_instantiations
+                .into_iter()
+                .map(Into::into)
+                .collect(),
+        )
+    }
+
+    fn optimistic_tx_indices_from_indexed_tx_indices(tx_index: TxIndex) -> OptimisticTxIndices {
+        let (senders, recipients, input_objects, changed_objects, pkgs, mods, funs, _, kinds) =
+            tx_index.split();
+
+        (
+            senders.into_iter().map(Into::into).collect(),
+            recipients.into_iter().map(Into::into).collect(),
+            input_objects.into_iter().map(Into::into).collect(),
+            changed_objects.into_iter().map(Into::into).collect(),
+            pkgs.into_iter().map(Into::into).collect(),
+            mods.into_iter().map(Into::into).collect(),
+            funs.into_iter().map(Into::into).collect(),
+            kinds.into_iter().map(Into::into).collect(),
+        )
+    }
 }
 
 #[async_trait]
@@ -37,13 +363,26 @@ impl WriteApiServer for WriteApi {
         tx_bytes: Base64,
         signatures: Vec<Base64>,
         options: Option<IotaTransactionBlockResponseOptions>,
-        request_type: Option<ExecuteTransactionRequestType>,
+        _request_type: Option<ExecuteTransactionRequestType>,
     ) -> RpcResult<IotaTransactionBlockResponse> {
-        let iota_transaction_response = self
-            .fullnode
-            .execute_transaction_block(tx_bytes, signatures, options.clone(), request_type)
-            .await
-            .map_err(error_object_from_rpc)?;
+        let request_type = Some(ExecuteTransactionRequestType::WaitForLocalExecution); // force local execution, for testing purposes
+        let iota_transaction_response = match request_type {
+            None | Some(ExecuteTransactionRequestType::WaitForEffectsCert) => self
+                .fullnode
+                .execute_transaction_block(tx_bytes, signatures, options.clone(), request_type)
+                .await
+                .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,
+                    signatures,
+                    options.clone(),
+                )
+                .await?
+            }
+        };
         Ok(IotaTransactionBlockResponseWithOptions {
             response: iota_transaction_response,
             options: options.unwrap_or_default(),