Optimize balance-related queries with a cache

crates/fuel-core/src/coins_query.rs

@@ -986,6 +986,7 @@ mod tests {
             let on_chain = self.database.on_chain().clone();
             let off_chain = self.database.off_chain().clone();
             ServiceDatabase::new(100, 0u32.into(), on_chain, off_chain)
+                .expect("should create service database")
         }
     }
 

crates/fuel-core/src/database.rs

@@ -508,7 +508,6 @@ where
 
     Ok(())
 }
-
 #[cfg(feature = "rocksdb")]
 pub fn convert_to_rocksdb_direction(direction: IterDirection) -> rocksdb::Direction {
     match direction {

crates/fuel-core/src/database/database_description.rs

@@ -4,6 +4,7 @@ use fuel_core_types::{
     blockchain::primitives::DaBlockHeight,
     fuel_types::BlockHeight,
 };
+use std::collections::HashSet;
 
 pub mod gas_price;
 pub mod off_chain;
@@ -67,24 +68,53 @@ pub trait DatabaseDescription: 'static + Copy + Debug + Send + Sync {
     fn prefix(column: &Self::Column) -> Option<usize>;
 }
 
+#[derive(
+    Copy, Clone, Debug, serde::Serialize, serde::Deserialize, Eq, PartialEq, Hash,
+)]
+pub enum IndexationKind {
+    Balances,
+    _CoinsToSpend,
+}
+
 /// The metadata of the database contains information about the version and its height.
-#[derive(Copy, Clone, Debug, serde::Serialize, serde::Deserialize)]
+#[derive(Clone, Debug, serde::Serialize, serde::Deserialize)]
 pub enum DatabaseMetadata<Height> {
-    V1 { version: u32, height: Height },
+    V1 {
+        version: u32,
+        height: Height,
+    },
+    V2 {
+        version: u32,
+        height: Height,
+        indexation_availability: HashSet<IndexationKind>,
+    },
 }
 
 impl<Height> DatabaseMetadata<Height> {
     /// Returns the version of the database.
     pub fn version(&self) -> u32 {
         match self {
             Self::V1 { version, .. } => *version,
+            Self::V2 { version, .. } => *version,
         }
     }
 
     /// Returns the height of the database.
     pub fn height(&self) -> &Height {
         match self {
             Self::V1 { height, .. } => height,
+            Self::V2 { height, .. } => height,
+        }
+    }
+
+    /// Returns true if the given indexation kind is available.
+    pub fn indexation_available(&self, kind: IndexationKind) -> bool {
+        match self {
+            Self::V1 { .. } => false,
+            Self::V2 {
+                indexation_availability,
+                ..
+            } => indexation_availability.contains(&kind),
         }
     }
 }

crates/fuel-core/src/database/database_description/off_chain.rs

@@ -11,6 +11,7 @@ impl DatabaseDescription for OffChain {
     type Column = fuel_core_graphql_api::storage::Column;
     type Height = BlockHeight;
 
+    // TODO[RC]: Do we bump this due to extended metadata?
     fn version() -> u32 {
         0
     }

crates/fuel-core/src/database/metadata.rs

@@ -17,6 +17,8 @@ use fuel_core_storage::{
     StorageInspect,
 };
 
+use super::database_description::IndexationKind;
+
 /// The table that stores all metadata about the database.
 pub struct MetadataTable<Description>(core::marker::PhantomData<Description>);
 
@@ -74,4 +76,12 @@ where
 
         Ok(metadata)
     }
+
+    pub fn indexation_available(&self, kind: IndexationKind) -> StorageResult<bool> {
+        let Some(metadata) = self.storage::<MetadataTable<Description>>().get(&())?
+        else {
+            return Ok(false)
+        };
+        Ok(metadata.indexation_available(kind))
+    }
 }

crates/fuel-core/src/graphql_api.rs

@@ -79,7 +79,7 @@ impl Default for Costs {
 }
 
 pub const DEFAULT_QUERY_COSTS: Costs = Costs {
-    balance_query: 40001,
+    balance_query: 40001, /* Cost will depend on whether balances index is available or not, but let's keep the default high to be on the safe side */
     coins_to_spend: 40001,
     get_peers: 40001,
     estimate_predicates: 40001,

crates/fuel-core/src/graphql_api/api_service.rs

@@ -90,6 +90,7 @@ use tower_http::{
 pub type Service = fuel_core_services::ServiceRunner<GraphqlService>;
 
 pub use super::database::ReadDatabase;
+use super::ports::worker;
 
 pub type BlockProducer = Box<dyn BlockProducerPort>;
 // In the future GraphQL should not be aware of `TxPool`. It should
@@ -229,7 +230,7 @@ pub fn new_service<OnChain, OffChain>(
 ) -> anyhow::Result<Service>
 where
     OnChain: AtomicView + 'static,
-    OffChain: AtomicView + 'static,
+    OffChain: AtomicView + worker::OffChainDatabase + 'static,
     OnChain::LatestView: OnChainDatabase,
     OffChain::LatestView: OffChainDatabase,
 {
@@ -241,7 +242,7 @@ where
         genesis_block_height,
         on_database,
         off_database,
-    );
+    )?;
     let request_timeout = config.config.api_request_timeout;
     let concurrency_limit = config.config.max_concurrent_queries;
     let body_limit = config.config.request_body_bytes_limit;

crates/fuel-core/src/graphql_api/database.rs

@@ -68,6 +68,8 @@ use std::{
     sync::Arc,
 };
 
+use super::ports::worker;
+
 mod arc_wrapper;
 
 /// The on-chain view of the database used by the [`ReadView`] to fetch on-chain data.
@@ -86,6 +88,8 @@ pub struct ReadDatabase {
     on_chain: Box<dyn AtomicView<LatestView = OnChainView>>,
     /// The off-chain database view provider.
     off_chain: Box<dyn AtomicView<LatestView = OffChainView>>,
+    /// The flag that indicates whether the Balances cache table is enabled.
+    balances_enabled: bool,
 }
 
 impl ReadDatabase {
@@ -95,19 +99,22 @@ impl ReadDatabase {
         genesis_height: BlockHeight,
         on_chain: OnChain,
         off_chain: OffChain,
-    ) -> Self
+    ) -> Result<Self, StorageError>
     where
         OnChain: AtomicView + 'static,
-        OffChain: AtomicView + 'static,
+        OffChain: AtomicView + worker::OffChainDatabase + 'static,
         OnChain::LatestView: OnChainDatabase,
         OffChain::LatestView: OffChainDatabase,
     {
-        Self {
+        let balances_enabled = off_chain.balances_enabled()?;
+
+        Ok(Self {
             batch_size,
             genesis_height,
             on_chain: Box::new(ArcWrapper::new(on_chain)),
             off_chain: Box::new(ArcWrapper::new(off_chain)),
-        }
+            balances_enabled,
+        })
     }
 
     /// Creates a consistent view of the database.
@@ -120,6 +127,7 @@ impl ReadDatabase {
             genesis_height: self.genesis_height,
             on_chain: self.on_chain.latest_view()?,
             off_chain: self.off_chain.latest_view()?,
+            balances_enabled: self.balances_enabled,
         })
     }
 
@@ -135,6 +143,7 @@ pub struct ReadView {
     pub(crate) genesis_height: BlockHeight,
     pub(crate) on_chain: OnChainView,
     pub(crate) off_chain: OffChainView,
+    pub(crate) balances_enabled: bool,
 }
 
 impl ReadView {

crates/fuel-core/src/graphql_api/ports.rs

@@ -62,7 +62,10 @@ use fuel_core_types::{
     },
     tai64::Tai64,
 };
-use std::sync::Arc;
+use std::{
+    collections::BTreeMap,
+    sync::Arc,
+};
 
 pub trait OffChainDatabase: Send + Sync {
     fn block_height(&self, block_id: &BlockId) -> StorageResult<BlockHeight>;
@@ -71,6 +74,19 @@ pub trait OffChainDatabase: Send + Sync {
 
     fn tx_status(&self, tx_id: &TxId) -> StorageResult<TransactionStatus>;
 
+    fn balance(
+        &self,
+        owner: &Address,
+        asset_id: &AssetId,
+        base_asset_id: &AssetId,
+    ) -> StorageResult<u64>;
+
+    fn balances(
+        &self,
+        owner: &Address,
+        base_asset_id: &AssetId,
+    ) -> StorageResult<BTreeMap<AssetId, u64>>;
+
     fn owned_coins_ids(
         &self,
         owner: &Address,
@@ -273,6 +289,10 @@ pub mod worker {
             },
         },
         graphql_api::storage::{
+            balances::{
+                CoinBalances,
+                MessageBalances,
+            },
             da_compression::*,
             old::{
                 OldFuelBlockConsensus,
@@ -315,6 +335,9 @@ pub mod worker {
 
         /// Creates a write database transaction.
         fn transaction(&mut self) -> Self::Transaction<'_>;
+
+        /// Checks if Balances cache functionality is available.
+        fn balances_enabled(&self) -> StorageResult<bool>;
     }
 
     pub trait OffChainDatabaseTransaction:
@@ -327,6 +350,8 @@ pub mod worker {
         + StorageMutate<OldTransactions, Error = StorageError>
         + StorageMutate<SpentMessages, Error = StorageError>
         + StorageMutate<RelayedTransactionStatuses, Error = StorageError>
+        + StorageMutate<CoinBalances, Error = StorageError>
+        + StorageMutate<MessageBalances, Error = StorageError>
         + StorageMutate<DaCompressedBlocks, Error = StorageError>
         + StorageMutate<DaCompressionTemporalRegistryAddress, Error = StorageError>
         + StorageMutate<DaCompressionTemporalRegistryAssetId, Error = StorageError>

crates/fuel-core/src/graphql_api/storage.rs

@@ -36,6 +36,7 @@ use fuel_core_types::{
 };
 use statistic::StatisticTable;
 
+pub mod balances;
 pub mod blocks;
 pub mod coins;
 pub mod contracts;
@@ -113,6 +114,10 @@ pub enum Column {
     DaCompressionTemporalRegistryScriptCode = 21,
     /// See [`DaCompressionTemporalRegistryPredicateCode`](da_compression::DaCompressionTemporalRegistryPredicateCode)
     DaCompressionTemporalRegistryPredicateCode = 22,
+    /// Coin balances per user and asset.
+    CoinBalances = 23,
+    /// Message balances per user.
+    MessageBalances = 24,
 }
 
 impl Column {

crates/fuel-core/src/graphql_api/storage/balances.rs

@@ -0,0 +1,70 @@
+use fuel_core_storage::{
+    blueprint::plain::Plain,
+    codec::{
+        postcard::Postcard,
+        raw::Raw,
+    },
+    structured_storage::TableWithBlueprint,
+    Mappable,
+};
+use fuel_core_types::{
+    fuel_tx::{
+        Address,
+        AssetId,
+    },
+    fuel_vm::double_key,
+};
+use rand::{
+    distributions::Standard,
+    prelude::Distribution,
+    Rng,
+};
+
+pub type Amount = u64;
+
+double_key!(BalancesKey, Address, address, AssetId, asset_id);
+impl Distribution<BalancesKey> for Standard {
+    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> BalancesKey {
+        let mut bytes = [0u8; BalancesKey::LEN];
+        rng.fill_bytes(bytes.as_mut());
+        BalancesKey::from_array(bytes)
+    }
+}
+
+/// This table stores the balances of coins per owner and asset id.
+pub struct CoinBalances;
+
+impl Mappable for CoinBalances {
+    type Key = BalancesKey;
+    type OwnedKey = Self::Key;
+    type Value = Amount;
+    type OwnedValue = Self::Value;
+}
+
+impl TableWithBlueprint for CoinBalances {
+    type Blueprint = Plain<Raw, Postcard>;
+    type Column = super::Column;
+
+    fn column() -> Self::Column {
+        Self::Column::CoinBalances
+    }
+}
+
+/// This table stores the balances of messages per owner.
+pub struct MessageBalances;
+
+impl Mappable for MessageBalances {
+    type Key = Address;
+    type OwnedKey = Self::Key;
+    type Value = Amount;
+    type OwnedValue = Self::Value;
+}
+
+impl TableWithBlueprint for MessageBalances {
+    type Blueprint = Plain<Raw, Postcard>;
+    type Column = super::Column;
+
+    fn column() -> Self::Column {
+        Self::Column::MessageBalances
+    }
+}