fix tx_global_index consistency with transactions table

Author: OBorce

api-server/api-server-common/src/storage/impls/in_memory/mod.rs

@@ -15,12 +15,13 @@
 
 pub mod transactional;
 
-use crate::storage::storage_api::{
-    block_aux_data::{BlockAuxData, BlockWithExtraData},
-    AmountWithDecimals, ApiServerStorageError, BlockInfo, CoinOrTokenStatistic, Delegation,
-    FungibleTokenData, LockedUtxo, NftWithOwner, Order, PoolBlockStats, PoolDataWithExtraInfo,
-    TokenTransaction, TransactionInfo, TransactionWithBlockInfo, Utxo, UtxoLock, UtxoWithExtraInfo,
+use std::{
+    cmp::{Ordering, Reverse},
+    collections::{BTreeMap, BTreeSet},
+    ops::Bound::{Excluded, Unbounded},
+    sync::Arc,
 };
+
 use common::{
     address::Address,
     chain::{
@@ -31,24 +32,42 @@ use common::{
     },
     primitives::{id::WithId, Amount, BlockHeight, CoinOrTokenId, Id, Idable},
 };
-use itertools::Itertools as _;
-use std::{
-    cmp::Reverse,
-    collections::{BTreeMap, BTreeSet},
-    ops::Bound::{Excluded, Unbounded},
-    sync::Arc,
+
+use crate::storage::storage_api::{
+    block_aux_data::{BlockAuxData, BlockWithExtraData},
+    AmountWithDecimals, ApiServerStorageError, BlockInfo, CoinOrTokenStatistic, Delegation,
+    FungibleTokenData, LockedUtxo, NftWithOwner, Order, PoolBlockStats, PoolDataWithExtraInfo,
+    TokenTransaction, TransactionInfo, TransactionWithBlockInfo, Utxo, UtxoLock, UtxoWithExtraInfo,
 };
 
+use itertools::Itertools as _;
+
 use super::CURRENT_STORAGE_VERSION;
 
+#[derive(Debug, Clone, PartialEq, Eq)]
+struct TokenTransactionOrderedByTxId(TokenTransaction);
+
+impl PartialOrd for TokenTransactionOrderedByTxId {
+    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
+        Some(self.cmp(other))
+    }
+}
+
+impl Ord for TokenTransactionOrderedByTxId {
+    fn cmp(&self, other: &Self) -> Ordering {
+        self.0.tx_id.cmp(&other.0.tx_id)
+    }
+}
+
 #[derive(Debug, Clone)]
 struct ApiServerInMemoryStorage {
     block_table: BTreeMap<Id<Block>, BlockWithExtraData>,
     block_aux_data_table: BTreeMap<Id<Block>, BlockAuxData>,
     address_balance_table: BTreeMap<String, BTreeMap<CoinOrTokenId, BTreeMap<BlockHeight, Amount>>>,
     address_locked_balance_table: BTreeMap<String, BTreeMap<(CoinOrTokenId, BlockHeight), Amount>>,
     address_transactions_table: BTreeMap<String, BTreeMap<BlockHeight, Vec<Id<Transaction>>>>,
-    token_transactions_table: BTreeMap<TokenId, BTreeMap<BlockHeight, BTreeSet<TokenTransaction>>>,
+    token_transactions_table:
+        BTreeMap<TokenId, BTreeMap<BlockHeight, BTreeSet<TokenTransactionOrderedByTxId>>>,
     delegation_table: BTreeMap<DelegationId, BTreeMap<BlockHeight, Delegation>>,
     main_chain_blocks_table: BTreeMap<BlockHeight, Id<Block>>,
     pool_data_table: BTreeMap<PoolId, BTreeMap<BlockHeight, PoolDataWithExtraInfo>>,
@@ -188,9 +207,9 @@ impl ApiServerInMemoryStorage {
                 transactions
                     .iter()
                     .rev()
-                    .flat_map(|(_, txs)| txs.iter())
-                    .cloned()
+                    .flat_map(|(_, txs)| txs.iter().map(|tx| &tx.0))
                     .flat_map(|tx| (tx.tx_global_index < tx_global_index).then_some(tx))
+                    .cloned()
                     .take(len as usize)
                     .collect()
             }))
@@ -979,36 +998,22 @@ impl ApiServerInMemoryStorage {
         Ok(())
     }
 
-    fn set_token_transactions_at_height(
+    fn set_token_transaction_at_height(
         &mut self,
         token_id: TokenId,
-        transaction_ids: BTreeSet<Id<Transaction>>,
+        tx_id: Id<Transaction>,
         block_height: BlockHeight,
+        tx_global_index: u64,
     ) -> Result<(), ApiServerStorageError> {
-        if transaction_ids.is_empty() {
-            return Ok(());
-        }
-
-        let next_tx_idx = self
-            .token_transactions_table
-            .values()
-            .flat_map(|by_height| by_height.values())
-            .flat_map(|tx_set| tx_set.iter())
-            .map(|tx| tx.tx_global_index + 1)
-            .max()
-            .unwrap_or(0);
-
         self.token_transactions_table
             .entry(token_id)
             .or_default()
             .entry(block_height)
             .or_default()
-            .extend(
-                transaction_ids.into_iter().enumerate().map(|(idx, tx_id)| TokenTransaction {
-                    tx_global_index: next_tx_idx + idx as u64,
-                    tx_id,
-                }),
-            );
+            .replace(TokenTransactionOrderedByTxId(TokenTransaction {
+                tx_global_index,
+                tx_id,
+            }));
 
         Ok(())
     }

api-server/api-server-common/src/storage/impls/in_memory/transactional/write.rs

@@ -111,14 +111,19 @@ impl ApiServerStorageWrite for ApiServerInMemoryStorageTransactionalRw<'_> {
             .set_address_transactions_at_height(address, transactions, block_height)
     }
 
-    async fn set_token_transactions_at_height(
+    async fn set_token_transaction_at_height(
         &mut self,
         token_id: TokenId,
-        transactions: BTreeSet<Id<Transaction>>,
+        tx_id: Id<Transaction>,
         block_height: BlockHeight,
+        tx_global_index: u64,
     ) -> Result<(), ApiServerStorageError> {
-        self.transaction
-            .set_token_transactions_at_height(token_id, transactions, block_height)
+        self.transaction.set_token_transaction_at_height(
+            token_id,
+            tx_id,
+            block_height,
+            tx_global_index,
+        )
     }
 
     async fn set_mainchain_block(

api-server/api-server-common/src/storage/impls/postgres/queries.rs

@@ -604,25 +604,25 @@ impl<'a, 'b> QueryFromConnection<'a, 'b> {
     pub async fn set_token_transactions_at_height(
         &mut self,
         token_id: TokenId,
-        transaction_ids: BTreeSet<Id<Transaction>>,
+        transaction_id: Id<Transaction>,
         block_height: BlockHeight,
+        tx_global_index: u64,
     ) -> Result<(), ApiServerStorageError> {
         let height = Self::block_height_to_postgres_friendly(block_height);
+        let tx_global_index = Self::tx_global_index_to_postgres_friendly(tx_global_index)?;
 
-        for transaction_id in transaction_ids {
-            self.tx
-                .execute(
-                    r#"
-                        INSERT INTO ml.token_transactions (token_id, block_height, transaction_id)
-                        VALUES ($1, $2, $3)
-                        ON CONFLICT (token_id, block_height, transaction_id)
-                        DO NOTHING;
-                    "#,
-                    &[&token_id.encode(), &height, &transaction_id.encode()],
-                )
-                .await
-                .map_err(|e| ApiServerStorageError::LowLevelStorageError(e.to_string()))?;
-        }
+        self.tx
+            .execute(
+                r#"
+                    INSERT INTO ml.token_transactions (token_id, block_height, transaction_id, tx_global_index)
+                    VALUES ($1, $2, $3, $4)
+                    ON CONFLICT (token_id, transaction_id, block_height)
+                    DO NOTHING;
+                "#,
+                &[&token_id.encode(), &height, &transaction_id.encode(), &tx_global_index],
+            )
+            .await
+            .map_err(|e| ApiServerStorageError::LowLevelStorageError(e.to_string()))?;
 
         Ok(())
     }
@@ -823,11 +823,11 @@ impl<'a, 'b> QueryFromConnection<'a, 'b> {
 
         self.just_execute(
             "CREATE TABLE ml.token_transactions (
-                    tx_global_index bigint PRIMARY KEY GENERATED ALWAYS AS IDENTITY (START WITH 0 MINVALUE 0),
+                    tx_global_index bigint PRIMARY KEY,
                     token_id bytea NOT NULL,
                     block_height bigint NOT NULL,
                     transaction_id bytea NOT NULL,
-                    UNIQUE (token_id, block_height, transaction_id)
+                    UNIQUE (token_id, transaction_id, block_height)
                 );",
         )
         .await?;

api-server/api-server-common/src/storage/impls/postgres/transactional/write.rs

@@ -131,15 +131,21 @@ impl ApiServerStorageWrite for ApiServerPostgresTransactionalRw<'_> {
         Ok(())
     }
 
-    async fn set_token_transactions_at_height(
+    async fn set_token_transaction_at_height(
         &mut self,
         token_id: TokenId,
-        transaction_ids: BTreeSet<Id<Transaction>>,
+        transaction_id: Id<Transaction>,
         block_height: BlockHeight,
+        tx_global_index: u64,
     ) -> Result<(), ApiServerStorageError> {
         let mut conn = QueryFromConnection::new(self.connection.as_ref().expect(CONN_ERR));
-        conn.set_token_transactions_at_height(token_id, transaction_ids, block_height)
-            .await?;
+        conn.set_token_transactions_at_height(
+            token_id,
+            transaction_id,
+            block_height,
+            tx_global_index,
+        )
+        .await?;
 
         Ok(())
     }

api-server/api-server-common/src/storage/storage_api/mod.rs

@@ -14,7 +14,6 @@
 // limitations under the License.
 
 use std::{
-    cmp::Ordering,
     collections::{BTreeMap, BTreeSet},
     fmt::Display,
     str::FromStr,
@@ -69,7 +68,7 @@ pub enum ApiServerStorageError {
     AddressableError,
     #[error("Block timestamp too high: {0}")]
     TimestampTooHigh(BlockTimestamp),
-    #[error("Tx global index to hight: {0}")]
+    #[error("Tx global index too hight: {0}")]
     TxGlobalIndexTooHigh(u64),
     #[error("Id creation error: {0}")]
     IdCreationError(#[from] IdCreationError),
@@ -590,18 +589,6 @@ pub struct TokenTransaction {
     pub tx_id: Id<Transaction>,
 }
 
-impl PartialOrd for TokenTransaction {
-    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
-        Some(self.cmp(other))
-    }
-}
-
-impl Ord for TokenTransaction {
-    fn cmp(&self, other: &Self) -> Ordering {
-        self.tx_id.cmp(&other.tx_id)
-    }
-}
-
 #[async_trait::async_trait]
 pub trait ApiServerStorageRead: Sync {
     async fn is_initialized(&self) -> Result<bool, ApiServerStorageError>;
@@ -630,10 +617,9 @@ pub trait ApiServerStorageRead: Sync {
         address: &str,
     ) -> Result<Vec<Id<Transaction>>, ApiServerStorageError>;
 
-    /// Return a page of TX IDs that reference this token_id, with a limit of len and older
-    /// tx_global_index than the specified.
-    /// The tx_global_index is only ordered by block height and are not continuous for a specific
-    /// token_id.
+    /// Returns a page of transaction IDs that reference this `token_id`, limited to `len` entries
+    /// and with a `tx_global_index` older than the specified value.
+    /// The `tx_global_index` and is not continuous for a specific `token_id`.
     async fn get_token_transactions(
         &self,
         token_id: TokenId,
@@ -866,12 +852,15 @@ pub trait ApiServerStorageWrite: ApiServerStorageRead {
         block_height: BlockHeight,
     ) -> Result<(), ApiServerStorageError>;
 
-    /// Append new token transactions with increasing tx_global_index at this block height
-    async fn set_token_transactions_at_height(
+    /// Sets the `token_id`–`transaction_id` pair at the specified `block_height` along with the
+    /// `tx_global_index`.
+    /// If the pair already exists at that `block_height`, the `tx_global_index` is updated.
+    async fn set_token_transaction_at_height(
         &mut self,
         token_id: TokenId,
-        transaction_ids: BTreeSet<Id<Transaction>>,
+        transaction_id: Id<Transaction>,
         block_height: BlockHeight,
+        tx_global_index: u64,
     ) -> Result<(), ApiServerStorageError>;
 
     async fn set_mainchain_block(