Optimize balance-related queries with a cache

CHANGELOG.md

@@ -29,6 +29,9 @@ and this project adheres to [Semantic Versioning](http://semver.org/).
 ### Changed
 - [2378](https://github.com/FuelLabs/fuel-core/pull/2378): Use cached hash of the topic instead of calculating it on each publishing gossip message.
 
+#### Breaking
+- [2383](https://github.com/FuelLabs/fuel-core/pull/2383): Asset balance queries now return U128 instead of U64. 
+
 ## [Version 0.40.0]
 
 ### Added

bin/e2e-test-client/src/test_context.rs

@@ -99,7 +99,7 @@ impl Wallet {
     }
 
     /// returns the balance associated with a wallet
-    pub async fn balance(&self, asset_id: Option<AssetId>) -> anyhow::Result<u64> {
+    pub async fn balance(&self, asset_id: Option<AssetId>) -> anyhow::Result<u128> {
         self.client
             .balance(&self.address, Some(&asset_id.unwrap_or_default()))
             .await

crates/client/assets/schema.sdl

@@ -4,7 +4,7 @@ scalar AssetId
 
 type Balance {
 	owner: Address!
-	amount: U64!
+	amount: U128!
 	assetId: AssetId!
 }
 
@@ -1259,6 +1259,8 @@ enum TxParametersVersion {
 
 scalar TxPointer
 
+scalar U128
+
 scalar U16
 
 scalar U32

crates/client/src/client.rs

@@ -1054,7 +1054,7 @@ impl FuelClient {
         &self,
         owner: &Address,
         asset_id: Option<&AssetId>,
-    ) -> io::Result<u64> {
+    ) -> io::Result<u128> {
         let owner: schema::Address = (*owner).into();
         let asset_id: schema::AssetId = match asset_id {
             Some(asset_id) => (*asset_id).into(),

crates/client/src/client/schema/balance.rs

@@ -4,7 +4,7 @@ use crate::client::{
         Address,
         AssetId,
         PageInfo,
-        U64,
+        U128,
     },
     PageDirection,
     PaginationRequest,
@@ -99,7 +99,7 @@ pub struct BalanceEdge {
 #[cynic(schema_path = "./assets/schema.sdl")]
 pub struct Balance {
     pub owner: Address,
-    pub amount: U64,
+    pub amount: U128,
     pub asset_id: AssetId,
 }
 

crates/client/src/client/schema/primitives.rs

@@ -272,6 +272,7 @@ macro_rules! number_scalar {
     };
 }
 
+number_scalar!(U128, u128);
 number_scalar!(U64, u64);
 number_scalar!(U32, u32);
 number_scalar!(U16, u16);

crates/client/src/client/types/balance.rs

@@ -10,7 +10,7 @@ use crate::client::{
 #[derive(Clone, Copy, Debug, PartialEq)]
 pub struct Balance {
     pub owner: Address,
-    pub amount: u64,
+    pub amount: u128,
     pub asset_id: AssetId,
 }
 

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,8 @@ use fuel_core_types::{
     blockchain::primitives::DaBlockHeight,
     fuel_types::BlockHeight,
 };
+use std::collections::HashSet;
+use strum::IntoEnumIterator;
 
 pub mod gas_price;
 pub mod off_chain;
@@ -67,24 +69,66 @@ 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,
+    strum::EnumIter,
+)]
+pub enum IndexationKind {
+    Balances,
+}
+
+impl IndexationKind {
+    pub fn all() -> impl Iterator<Item = Self> {
+        Self::iter()
+    }
+}
+
 /// 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/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/executor.rs

@@ -638,6 +638,7 @@ mod tests {
                 .next()
                 .unwrap()
                 .unwrap();
+
             assert_eq!(asset_id, AssetId::zeroed());
             assert_eq!(amount, expected_fee_amount_1 + expected_fee_amount_2);
         }

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,12 @@ use fuel_core_types::{
     },
     tai64::Tai64,
 };
-use std::sync::Arc;
+use std::{
+    collections::BTreeMap,
+    sync::Arc,
+};
+
+use super::storage::balances::TotalBalanceAmount;
 
 pub trait OffChainDatabase: Send + Sync {
     fn block_height(&self, block_id: &BlockId) -> StorageResult<BlockHeight>;
@@ -71,6 +76,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<TotalBalanceAmount>;
+
+    fn balances(
+        &self,
+        owner: &Address,
+        base_asset_id: &AssetId,
+    ) -> StorageResult<BTreeMap<AssetId, TotalBalanceAmount>>;
+
     fn owned_coins_ids(
         &self,
         owner: &Address,
@@ -273,6 +291,10 @@ pub mod worker {
             },
         },
         graphql_api::storage::{
+            balances::{
+                CoinBalances,
+                MessageBalances,
+            },
             da_compression::*,
             old::{
                 OldFuelBlockConsensus,
@@ -315,6 +337,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 +352,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 {