Optimize balance-related queries with a cache (#2383)

## Linked Issues/PRs
Closes #1965

## Description
1. The client uses the off-chain database metadata to tell if it can use
the optimized queries or fall-back to the legacy ones.
2. When client starts the metadata is either:
1. Kept intact in-case genesis exists (meaning - no optimized queries
are available)
* this stems from the fact that the pre-existing DB contains some coins
already and we'll **not** be re-indexing them in order to create the
lookup indexes
* in this mode of operation the "old" way of getting balances is used as
confirmed in logs:
          * `worker_service: 607: Balances cache available: false`
* `query::balance: 100: Querying balances without balances cache
owner=53a9c6a...`
    2. Initialized with `V2` if genesis is missing
* this means that while re-syncing the DB the balance indexes are going
to be created and could be used to satisfy the requests
          * `worker_service: 607: Balances cache available: true`
* `query::balance: 151: Querying balances using balances cache
owner=53a9c6a...`

Technical considerations:
* Metadata is extended with `V2` which now contains
`indexation_availability`. This is a set that holds the available
indexations (currently only `Balances`)
* New databases are added: `CoinBalances` and `MessageBalances` that
keep the coin and message balances respectively
   * `CoinBalances` stores the balances per `owner` and `asset_id`
* `MessageBalances` stores balances per `owner`, separately for
retryable and non-retryable messages
* In case balances indexation is available (see "Description" above):
* these databases are updated in the GrapQL API worker service when
processing one of these events: `MessageImported`, `MessageConsumed`,
`CoinCreated`, `CoinConsumed`
* `balance()` and `balances()` queries use these databases to satisfy
the query

## Checklist
- [X] New behavior is reflected in tests

### Before requesting review
- [X] I have reviewed the code myself

---------

Co-authored-by: Green Baneling <XgreenX9999@gmail.com>

CHANGELOG.md

@@ -43,6 +43,7 @@ and this project adheres to [Semantic Versioning](http://semver.org/).
 
 #### Breaking
 - [2389](https://github.com/FuelLabs/fuel-core/pull/2258): Updated the `messageProof` GraphQL schema to return a non-nullable `MessageProof`.
+- [2383](https://github.com/FuelLabs/fuel-core/pull/2383): Asset balance queries now return U128 instead of U64. 
 - [2154](https://github.com/FuelLabs/fuel-core/pull/2154): Transaction graphql endpoints use `TransactionType` instead of `fuel_tx::Transaction`.
 - [2446](https://github.com/FuelLabs/fuel-core/pull/2446): Use graphiql instead of graphql-playground due to known vulnerability and stale development.
 - [2379](https://github.com/FuelLabs/fuel-core/issues/2379): Change `kv_store::Value` to be `Arc<[u8]>` instead of `Arc<Vec<u8>>`.

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

@@ -1069,7 +1069,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

@@ -58,6 +58,7 @@ use fuel_core_types::{
 };
 use itertools::Itertools;
 use std::{
+    borrow::Cow,
     fmt::Debug,
     sync::Arc,
 };
@@ -66,7 +67,10 @@ pub use fuel_core_database::Error;
 pub type Result<T> = core::result::Result<T, Error>;
 
 // TODO: Extract `Database` and all belongs into `fuel-core-database`.
-use crate::database::database_description::gas_price::GasPriceDatabase;
+use crate::database::database_description::{
+    gas_price::GasPriceDatabase,
+    indexation_availability,
+};
 #[cfg(feature = "rocksdb")]
 use crate::state::{
     historical_rocksdb::{
@@ -482,15 +486,13 @@ where
             ConflictPolicy::Overwrite,
             changes,
         );
+        let maybe_current_metadata = transaction
+            .storage_as_mut::<MetadataTable<Description>>()
+            .get(&())?;
+        let metadata = update_metadata::<Description>(maybe_current_metadata, new_height);
         transaction
             .storage_as_mut::<MetadataTable<Description>>()
-            .insert(
-                &(),
-                &DatabaseMetadata::V1 {
-                    version: Description::version(),
-                    height: new_height,
-                },
-            )?;
+            .insert(&(), &metadata)?;
 
         transaction.into_changes()
     } else {
@@ -509,6 +511,39 @@ where
     Ok(())
 }
 
+fn update_metadata<Description>(
+    maybe_current_metadata: Option<
+        Cow<DatabaseMetadata<<Description as DatabaseDescription>::Height>>,
+    >,
+    new_height: <Description as DatabaseDescription>::Height,
+) -> DatabaseMetadata<<Description as DatabaseDescription>::Height>
+where
+    Description: DatabaseDescription,
+{
+    let updated_metadata = match maybe_current_metadata.as_ref() {
+        Some(metadata) => match metadata.as_ref() {
+            DatabaseMetadata::V1 { .. } => DatabaseMetadata::V1 {
+                version: Description::version(),
+                height: new_height,
+            },
+            DatabaseMetadata::V2 {
+                indexation_availability,
+                ..
+            } => DatabaseMetadata::V2 {
+                version: Description::version(),
+                height: new_height,
+                indexation_availability: indexation_availability.clone(),
+            },
+        },
+        None => DatabaseMetadata::V2 {
+            version: Description::version(),
+            height: new_height,
+            indexation_availability: indexation_availability::<Description>(None),
+        },
+    };
+    updated_metadata
+}
+
 #[cfg(feature = "rocksdb")]
 pub fn convert_to_rocksdb_direction(direction: IterDirection) -> rocksdb::Direction {
     match direction {
@@ -524,10 +559,6 @@ mod tests {
         database_description::DatabaseDescription,
         Database,
     };
-    use fuel_core_storage::{
-        tables::FuelBlocks,
-        StorageAsMut,
-    };
 
     fn column_keys_not_exceed_count<Description>()
     where
@@ -1084,4 +1115,145 @@ mod tests {
         // rocks db fails
         test(db);
     }
+
+    mod metadata {
+        use crate::database::database_description::IndexationKind;
+        use fuel_core_storage::kv_store::StorageColumn;
+        use std::{
+            borrow::Cow,
+            collections::HashSet,
+        };
+        use strum::EnumCount;
+
+        use super::{
+            database_description::DatabaseDescription,
+            update_metadata,
+            DatabaseHeight,
+            DatabaseMetadata,
+        };
+
+        #[derive(Debug, Copy, Clone, Default, PartialEq, Eq)]
+        struct HeightMock(u64);
+        impl DatabaseHeight for HeightMock {
+            fn as_u64(&self) -> u64 {
+                1
+            }
+
+            fn advance_height(&self) -> Option<Self> {
+                None
+            }
+
+            fn rollback_height(&self) -> Option<Self> {
+                None
+            }
+        }
+
+        const MOCK_VERSION: u32 = 0;
+
+        #[derive(EnumCount, enum_iterator::Sequence, Debug, Clone, Copy)]
+        enum ColumnMock {
+            Column1,
+        }
+
+        impl StorageColumn for ColumnMock {
+            fn name(&self) -> String {
+                "column".to_string()
+            }
+
+            fn id(&self) -> u32 {
+                42
+            }
+        }
+
+        #[derive(Debug, Clone, Copy)]
+        struct DatabaseDescriptionMock;
+        impl DatabaseDescription for DatabaseDescriptionMock {
+            type Column = ColumnMock;
+
+            type Height = HeightMock;
+
+            fn version() -> u32 {
+                MOCK_VERSION
+            }
+
+            fn name() -> String {
+                "mock".to_string()
+            }
+
+            fn metadata_column() -> Self::Column {
+                Self::Column::Column1
+            }
+
+            fn prefix(_: &Self::Column) -> Option<usize> {
+                None
+            }
+        }
+
+        #[test]
+        fn update_metadata_preserves_v1() {
+            let current_metadata: DatabaseMetadata<HeightMock> = DatabaseMetadata::V1 {
+                version: MOCK_VERSION,
+                height: HeightMock(1),
+            };
+            let new_metadata = update_metadata::<DatabaseDescriptionMock>(
+                Some(Cow::Borrowed(&current_metadata)),
+                HeightMock(2),
+            );
+
+            match new_metadata {
+                DatabaseMetadata::V1 { version, height } => {
+                    assert_eq!(version, current_metadata.version());
+                    assert_eq!(height, HeightMock(2));
+                }
+                DatabaseMetadata::V2 { .. } => panic!("should be V1"),
+            }
+        }
+
+        #[test]
+        fn update_metadata_preserves_v2() {
+            let available_indexation = HashSet::new();
+
+            let current_metadata: DatabaseMetadata<HeightMock> = DatabaseMetadata::V2 {
+                version: MOCK_VERSION,
+                height: HeightMock(1),
+                indexation_availability: available_indexation.clone(),
+            };
+            let new_metadata = update_metadata::<DatabaseDescriptionMock>(
+                Some(Cow::Borrowed(&current_metadata)),
+                HeightMock(2),
+            );
+
+            match new_metadata {
+                DatabaseMetadata::V1 { .. } => panic!("should be V2"),
+                DatabaseMetadata::V2 {
+                    version,
+                    height,
+                    indexation_availability,
+                } => {
+                    assert_eq!(version, current_metadata.version());
+                    assert_eq!(height, HeightMock(2));
+                    assert_eq!(indexation_availability, available_indexation);
+                }
+            }
+        }
+
+        #[test]
+        fn update_metadata_none_becomes_v2() {
+            let new_metadata =
+                update_metadata::<DatabaseDescriptionMock>(None, HeightMock(2));
+
+            match new_metadata {
+                DatabaseMetadata::V1 { .. } => panic!("should be V2"),
+                DatabaseMetadata::V2 {
+                    version,
+                    height,
+                    indexation_availability,
+                } => {
+                    assert_eq!(version, MOCK_VERSION);
+                    assert_eq!(height, HeightMock(2));
+                    assert_eq!(indexation_availability, IndexationKind::all().collect());
+                }
+            }
+        }
+    }
 }