From 9b36c1f9e78295cd81dc62168b9aa9a75e40f9a3 Mon Sep 17 00:00:00 2001
From: Aaryamann Challani <43716372+rymnc@users.noreply.github.com>
Date: Wed, 29 Jan 2025 17:45:58 +0530
Subject: [PATCH] port(storage): use the cached block height from the database
 at the time of creation of the view

---
 CHANGELOG.md                                  |  5 +-
 crates/fuel-core/src/database.rs              | 89 +++++++++++++------
 crates/fuel-core/src/database/storage.rs      |  6 +-
 .../service/genesis/importer/import_task.rs   |  6 +-
 crates/fuel-core/src/state.rs                 | 28 +++---
 .../fuel-core/src/state/generic_database.rs   | 43 +++++----
 .../fuel-core/src/state/historical_rocksdb.rs | 16 +++-
 .../src/state/in_memory/memory_store.rs       |  9 +-
 8 files changed, 133 insertions(+), 69 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index ef58b7360af..886c905cb3e 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -8,16 +8,15 @@ and this project adheres to [Semantic Versioning](http://semver.org/).
 
 ### Added
 - [2635](https://github.com/FuelLabs/fuel-core/pull/2635): Add metrics to gas price service
+- [2553](https://github.com/FuelLabs/fuel-core/pull/2553): Scaffold global merkle root storage crate.
 
 ### Changed
 - [2387](https://github.com/FuelLabs/fuel-core/pull/2387): Update description `tx-max-depth` flag.
 - [2630](https://github.com/FuelLabs/fuel-core/pull/2630): Removed some noisy `tracing::info!` logs
 
-### Added
-- [2553](https://github.com/FuelLabs/fuel-core/pull/2553): Scaffold global merkle root storage crate.
-
 ### Fixed
 - [2632](https://github.com/FuelLabs/fuel-core/pull/2632): Improved performance of certain async trait impls in the gas price service.
+- [2646](https://github.com/FuelLabs/fuel-core/pull/2646): Improved performance of fetching block height by caching it when the view is created.
 
 ## [Version 0.41.4]
 
diff --git a/crates/fuel-core/src/database.rs b/crates/fuel-core/src/database.rs
index 1f3abfb59a9..5bdeecf35d6 100644
--- a/crates/fuel-core/src/database.rs
+++ b/crates/fuel-core/src/database.rs
@@ -25,6 +25,7 @@ use crate::{
     },
 };
 use fuel_core_chain_config::TableEntry;
+pub use fuel_core_database::Error;
 use fuel_core_gas_price_service::common::fuel_core_storage_adapter::storage::GasPriceMetadata;
 use fuel_core_services::SharedMutex;
 use fuel_core_storage::{
@@ -60,17 +61,12 @@ use itertools::Itertools;
 use std::{
     borrow::Cow,
     fmt::Debug,
+    io::Empty,
     sync::Arc,
 };
-
-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,
-    indexation_availability,
-};
 #[cfg(feature = "rocksdb")]
 use crate::state::{
     historical_rocksdb::{
@@ -84,6 +80,13 @@ use crate::state::{
         RocksDb,
     },
 };
+use crate::{
+    database::database_description::{
+        gas_price::GasPriceDatabase,
+        indexation_availability,
+    },
+    state::HeightType,
+};
 #[cfg(feature = "rocksdb")]
 use std::path::Path;
 
@@ -126,10 +129,13 @@ where
 }
 
 pub type Database<Description = OnChain, Stage = RegularStage<Description>> =
-    GenericDatabase<DataSource<Description, Stage>>;
-pub type OnChainIterableKeyValueView = IterableKeyValueView<ColumnType<OnChain>>;
-pub type OffChainIterableKeyValueView = IterableKeyValueView<ColumnType<OffChain>>;
-pub type RelayerIterableKeyValueView = IterableKeyValueView<ColumnType<Relayer>>;
+    GenericDatabase<DataSource<Description, Stage>, Empty>;
+pub type OnChainIterableKeyValueView =
+    IterableKeyValueView<ColumnType<OnChain>, HeightType<OnChain>>;
+pub type OffChainIterableKeyValueView =
+    IterableKeyValueView<ColumnType<OffChain>, HeightType<OffChain>>;
+pub type RelayerIterableKeyValueView =
+    IterableKeyValueView<ColumnType<Relayer>, HeightType<Relayer>>;
 
 pub type GenesisDatabase<Description = OnChain> = Database<Description, GenesisStage>;
 
@@ -141,7 +147,9 @@ impl OnChainIterableKeyValueView {
     }
 
     pub fn latest_height(&self) -> StorageResult<BlockHeight> {
-        self.maybe_latest_height()?.ok_or(not_found!("BlockHeight"))
+        self.metadata()
+            .cloned()
+            .ok_or_else(|| not_found!("Metadata"))
     }
 
     pub fn latest_block(&self) -> StorageResult<CompressedBlock> {
@@ -176,7 +184,10 @@ where
     Description: DatabaseDescription,
 {
     pub fn new(data_source: DataSourceType<Description>) -> Self {
-        GenesisDatabase::from_storage(DataSource::new(data_source, GenesisStage))
+        GenesisDatabase::from_storage_and_metadata(
+            DataSource::new(data_source, GenesisStage),
+            None,
+        )
     }
 }
 
@@ -187,12 +198,15 @@ where
         StorageInspect<MetadataTable<Description>, Error = StorageError>,
 {
     pub fn new(data_source: DataSourceType<Description>) -> Self {
-        let mut database = Self::from_storage(DataSource::new(
-            data_source,
-            RegularStage {
-                height: SharedMutex::new(None),
-            },
-        ));
+        let mut database = Self::from_storage_and_metadata(
+            DataSource::new(
+                data_source,
+                RegularStage {
+                    height: SharedMutex::new(None),
+                },
+            ),
+            Some(Empty::default()),
+        );
         let height = database
             .latest_height_from_metadata()
             .expect("Failed to get latest height during creation of the database");
@@ -235,14 +249,14 @@ where
     ) -> core::result::Result<GenesisDatabase<Description>, GenesisDatabase<Description>>
     {
         if !self.stage.height.lock().is_some() {
-            Ok(GenesisDatabase::new(self.into_inner().data))
+            Ok(GenesisDatabase::new(self.into_inner().0.data))
         } else {
             tracing::warn!(
                 "Converting regular database into genesis, \
                 while height is already set for `{}`",
                 Description::name()
             );
-            Err(GenesisDatabase::new(self.into_inner().data))
+            Err(GenesisDatabase::new(self.into_inner().0.data))
         }
     }
 }
@@ -254,7 +268,10 @@ where
 {
     pub fn in_memory() -> Self {
         let data = Arc::<MemoryStore<Description>>::new(MemoryStore::default());
-        Self::from_storage(DataSource::new(data, Stage::default()))
+        Self::from_storage_and_metadata(
+            DataSource::new(data, Stage::default()),
+            Some(Empty::default()),
+        )
     }
 
     #[cfg(feature = "rocksdb")]
@@ -323,16 +340,30 @@ where
 
         Ok(())
     }
+
+    fn latest_view_with_height(
+        &self,
+        height: Option<Description::Height>,
+    ) -> StorageResult<IterableKeyValueView<ColumnType<Description>, Description::Height>>
+    {
+        let view = self.inner_storage().data.latest_view()?;
+
+        let (view, _) = view.into_inner();
+        Ok(IterableKeyValueView::from_storage_and_metadata(
+            view, height,
+        ))
+    }
 }
 
 impl<Description> AtomicView for Database<Description>
 where
     Description: DatabaseDescription,
 {
-    type LatestView = IterableKeyValueView<ColumnType<Description>>;
+    type LatestView = IterableKeyValueView<ColumnType<Description>, Description::Height>;
 
     fn latest_view(&self) -> StorageResult<Self::LatestView> {
-        self.inner_storage().data.latest_view()
+        let height = *self.inner_storage().stage.height.lock();
+        self.latest_view_with_height(height)
     }
 }
 
@@ -341,7 +372,7 @@ where
     Description: DatabaseDescription,
 {
     type Height = Description::Height;
-    type ViewAtHeight = KeyValueView<ColumnType<Description>>;
+    type ViewAtHeight = KeyValueView<ColumnType<Description>, Description::Height>;
 
     fn latest_height(&self) -> Option<Self::Height> {
         *self.inner_storage().stage.height.lock()
@@ -351,9 +382,15 @@ where
         let lock = self.inner_storage().stage.height.lock();
 
         match *lock {
-            None => return self.latest_view().map(|view| view.into_key_value_view()),
+            None => {
+                return self
+                    .latest_view_with_height(None)
+                    .map(|view| view.into_key_value_view())
+            }
             Some(current_height) if &current_height == height => {
-                return self.latest_view().map(|view| view.into_key_value_view())
+                return self
+                    .latest_view_with_height(Some(current_height))
+                    .map(|view| view.into_key_value_view())
             }
             _ => {}
         };
diff --git a/crates/fuel-core/src/database/storage.rs b/crates/fuel-core/src/database/storage.rs
index d00e292963f..4d5a408af48 100644
--- a/crates/fuel-core/src/database/storage.rs
+++ b/crates/fuel-core/src/database/storage.rs
@@ -16,7 +16,7 @@ use fuel_core_storage::{
     StorageWrite,
 };
 
-impl<Storage, M> StorageMutate<M> for GenericDatabase<Storage>
+impl<Storage, M, Metadata> StorageMutate<M> for GenericDatabase<Storage, Metadata>
 where
     M: Mappable,
     Self: Modifiable,
@@ -50,7 +50,7 @@ where
     }
 }
 
-impl<Storage, M> StorageWrite<M> for GenericDatabase<Storage>
+impl<Storage, M, Metadata> StorageWrite<M> for GenericDatabase<Storage, Metadata>
 where
     M: Mappable,
     StructuredStorage<Storage>: StorageInspect<M, Error = StorageError>,
@@ -95,7 +95,7 @@ where
     }
 }
 
-impl<Storage, M> StorageBatchMutate<M> for GenericDatabase<Storage>
+impl<Storage, M, Metadata> StorageBatchMutate<M> for GenericDatabase<Storage, Metadata>
 where
     M: Mappable,
     StructuredStorage<Storage>: StorageInspect<M, Error = StorageError>,
diff --git a/crates/fuel-core/src/service/genesis/importer/import_task.rs b/crates/fuel-core/src/service/genesis/importer/import_task.rs
index 5cdcf636c8c..a7936ed3e25 100644
--- a/crates/fuel-core/src/service/genesis/importer/import_task.rs
+++ b/crates/fuel-core/src/service/genesis/importer/import_task.rs
@@ -578,11 +578,13 @@ mod tests {
         fn view_at_height(
             &self,
             _: &BlockHeight,
-        ) -> StorageResult<KeyValueView<Self::Column>> {
+        ) -> StorageResult<KeyValueView<Self::Column, BlockHeight>> {
             Err(anyhow::anyhow!("I refuse to work!").into())
         }
 
-        fn latest_view(&self) -> StorageResult<IterableKeyValueView<Self::Column>> {
+        fn latest_view(
+            &self,
+        ) -> StorageResult<IterableKeyValueView<Self::Column, BlockHeight>> {
             Err(anyhow::anyhow!("I refuse to work!").into())
         }
 
diff --git a/crates/fuel-core/src/state.rs b/crates/fuel-core/src/state.rs
index 06ee176c652..7a2549eb054 100644
--- a/crates/fuel-core/src/state.rs
+++ b/crates/fuel-core/src/state.rs
@@ -30,23 +30,25 @@ pub mod rocks_db;
 pub mod rocks_db_key_iterator;
 
 pub type ColumnType<Description> = <Description as DatabaseDescription>::Column;
+pub type HeightType<Description> = <Description as DatabaseDescription>::Height;
 
 /// A type extends the `KeyValueView`, allowing iteration over the storage.
-pub type IterableKeyValueView<Column> =
-    GenericDatabase<IterableKeyValueViewWrapper<Column>>;
+pub type IterableKeyValueView<Column, BlockHeight> =
+    GenericDatabase<IterableKeyValueViewWrapper<Column>, BlockHeight>;
 
 /// The basic view available for the key value storage.
-pub type KeyValueView<Column> = GenericDatabase<KeyValueViewWrapper<Column>>;
+pub type KeyValueView<Column, BlockHeight> =
+    GenericDatabase<KeyValueViewWrapper<Column>, BlockHeight>;
 
-impl<Column> IterableKeyValueView<Column>
+impl<Column, Height> IterableKeyValueView<Column, Height>
 where
     Column: StorageColumn + 'static,
 {
     /// Downgrades the `IterableKeyValueView` into the `KeyValueView`.
-    pub fn into_key_value_view(self) -> KeyValueView<Column> {
-        let iterable = self.into_inner();
+    pub fn into_key_value_view(self) -> KeyValueView<Column, Height> {
+        let (iterable, metadata) = self.into_inner();
         let storage = KeyValueViewWrapper::new(iterable);
-        KeyValueView::from_storage(storage)
+        KeyValueView::from_storage_and_metadata(storage, metadata)
     }
 }
 
@@ -61,9 +63,9 @@ pub trait TransactableStorage<Height>: IterableStore + Debug + Send + Sync {
     fn view_at_height(
         &self,
         height: &Height,
-    ) -> StorageResult<KeyValueView<Self::Column>>;
+    ) -> StorageResult<KeyValueView<Self::Column, Height>>;
 
-    fn latest_view(&self) -> StorageResult<IterableKeyValueView<Self::Column>>;
+    fn latest_view(&self) -> StorageResult<IterableKeyValueView<Self::Column, Height>>;
 
     fn rollback_block_to(&self, height: &Height) -> StorageResult<()>;
 }
@@ -79,14 +81,16 @@ where
         unimplemented!()
     }
 
-    fn view_at_height(&self, _: &Height) -> StorageResult<KeyValueView<Self::Column>> {
+    fn view_at_height(
+        &self,
+        _: &Height,
+    ) -> StorageResult<KeyValueView<Self::Column, Height>> {
         unimplemented!()
     }
 
-    fn latest_view(&self) -> StorageResult<IterableKeyValueView<Self::Column>> {
+    fn latest_view(&self) -> StorageResult<IterableKeyValueView<Self::Column, Height>> {
         unimplemented!()
     }
-
     fn rollback_block_to(&self, _: &Height) -> StorageResult<()> {
         unimplemented!()
     }
diff --git a/crates/fuel-core/src/state/generic_database.rs b/crates/fuel-core/src/state/generic_database.rs
index 5a31591268f..a99b81bdf12 100644
--- a/crates/fuel-core/src/state/generic_database.rs
+++ b/crates/fuel-core/src/state/generic_database.rs
@@ -28,27 +28,36 @@ use std::{
 };
 
 #[derive(Debug, Clone)]
-pub struct GenericDatabase<Storage> {
+pub struct GenericDatabase<Storage, Metadata> {
     storage: StructuredStorage<Storage>,
+    metadata: Option<Metadata>,
 }
 
-impl<Storage> GenericDatabase<Storage> {
+impl<Storage, Metadata> GenericDatabase<Storage, Metadata> {
     pub fn inner_storage(&self) -> &Storage {
         self.storage.as_ref()
     }
 
-    pub fn from_storage(storage: Storage) -> Self {
+    pub fn metadata(&self) -> Option<&Metadata> {
+        self.metadata.as_ref()
+    }
+
+    pub fn from_storage_and_metadata(
+        storage: Storage,
+        metadata: Option<Metadata>,
+    ) -> Self {
         Self {
             storage: StructuredStorage::new(storage),
+            metadata,
         }
     }
 
-    pub fn into_inner(self) -> Storage {
-        self.storage.into_storage()
+    pub fn into_inner(self) -> (Storage, Option<Metadata>) {
+        (self.storage.into_inner(), self.metadata)
     }
 }
 
-impl<M, Storage> StorageInspect<M> for GenericDatabase<Storage>
+impl<M, Storage, Metadata> StorageInspect<M> for GenericDatabase<Storage, Metadata>
 where
     M: Mappable,
     StructuredStorage<Storage>: StorageInspect<M, Error = StorageError>,
@@ -64,7 +73,7 @@ where
     }
 }
 
-impl<Storage, M> StorageSize<M> for GenericDatabase<Storage>
+impl<M, Storage, Metadata> StorageSize<M> for GenericDatabase<Storage, Metadata>
 where
     M: Mappable,
     StructuredStorage<Storage>: StorageSize<M, Error = StorageError>,
@@ -74,7 +83,7 @@ where
     }
 }
 
-impl<Storage, M> StorageRead<M> for GenericDatabase<Storage>
+impl<M, Storage, Metadata> StorageRead<M> for GenericDatabase<Storage, Metadata>
 where
     M: Mappable,
     StructuredStorage<Storage>: StorageRead<M, Error = StorageError>,
@@ -93,7 +102,8 @@ where
     }
 }
 
-impl<Key, M, Storage> MerkleRootStorage<Key, M> for GenericDatabase<Storage>
+impl<Key, M, Storage, Metadata> MerkleRootStorage<Key, M>
+    for GenericDatabase<Storage, Metadata>
 where
     M: Mappable,
     StructuredStorage<Storage>: MerkleRootStorage<Key, M, Error = StorageError>,
@@ -103,7 +113,7 @@ where
     }
 }
 
-impl<Storage> KeyValueInspect for GenericDatabase<Storage>
+impl<Storage, Metadata> KeyValueInspect for GenericDatabase<Storage, Metadata>
 where
     Storage: KeyValueInspect,
 {
@@ -136,7 +146,7 @@ where
     }
 }
 
-impl<Storage> IterableStore for GenericDatabase<Storage>
+impl<Storage, Metadata> IterableStore for GenericDatabase<Storage, Metadata>
 where
     Storage: IterableStore,
 {
@@ -162,19 +172,19 @@ where
     }
 }
 
-impl<Storage> AsRef<Storage> for GenericDatabase<Storage> {
+impl<Storage, Metadata> AsRef<Storage> for GenericDatabase<Storage, Metadata> {
     fn as_ref(&self) -> &Storage {
         self.storage.as_ref()
     }
 }
 
-impl<Storage> AsMut<Storage> for GenericDatabase<Storage> {
+impl<Storage, Metadata> AsMut<Storage> for GenericDatabase<Storage, Metadata> {
     fn as_mut(&mut self) -> &mut Storage {
         self.storage.as_mut()
     }
 }
 
-impl<Storage> core::ops::Deref for GenericDatabase<Storage> {
+impl<Storage, Metadata> core::ops::Deref for GenericDatabase<Storage, Metadata> {
     type Target = Storage;
 
     fn deref(&self) -> &Self::Target {
@@ -182,13 +192,14 @@ impl<Storage> core::ops::Deref for GenericDatabase<Storage> {
     }
 }
 
-impl<Storage> core::ops::DerefMut for GenericDatabase<Storage> {
+impl<Storage, Metadata> core::ops::DerefMut for GenericDatabase<Storage, Metadata> {
     fn deref_mut(&mut self) -> &mut Self::Target {
         self.as_mut()
     }
 }
 
-impl<Storage, Error> PredicateStorageRequirements for GenericDatabase<Storage>
+impl<Storage, Error, Metadata> PredicateStorageRequirements
+    for GenericDatabase<Storage, Metadata>
 where
     Self: StorageRead<BlobData, Error = Error>,
     Error: Debug,
diff --git a/crates/fuel-core/src/state/historical_rocksdb.rs b/crates/fuel-core/src/state/historical_rocksdb.rs
index 45ce0d797de..a231335736a 100644
--- a/crates/fuel-core/src/state/historical_rocksdb.rs
+++ b/crates/fuel-core/src/state/historical_rocksdb.rs
@@ -20,6 +20,7 @@ use crate::{
         key_value_view::KeyValueViewWrapper,
         rocks_db::RocksDb,
         ColumnType,
+        HeightType,
         IterableKeyValueView,
         KeyValueView,
         TransactableStorage,
@@ -595,17 +596,24 @@ where
     fn view_at_height(
         &self,
         height: &Description::Height,
-    ) -> StorageResult<KeyValueView<ColumnType<Description>>> {
+    ) -> StorageResult<KeyValueView<ColumnType<Description>, HeightType<Description>>>
+    {
         let view = self.create_view_at(height)?;
-        Ok(KeyValueView::from_storage(KeyValueViewWrapper::new(view)))
+        Ok(KeyValueView::from_storage_and_metadata(
+            KeyValueViewWrapper::new(view),
+            Some(*height),
+        ))
     }
 
     fn latest_view(
         &self,
-    ) -> StorageResult<IterableKeyValueView<ColumnType<Description>>> {
+    ) -> StorageResult<
+        IterableKeyValueView<ColumnType<Description>, HeightType<Description>>,
+    > {
         let view = self.latest_view();
-        Ok(IterableKeyValueView::from_storage(
+        Ok(IterableKeyValueView::from_storage_and_metadata(
             IterableKeyValueViewWrapper::new(view),
+            None,
         ))
     }
 
diff --git a/crates/fuel-core/src/state/in_memory/memory_store.rs b/crates/fuel-core/src/state/in_memory/memory_store.rs
index 17ee7f9ea5b..dd4aa94ecfe 100644
--- a/crates/fuel-core/src/state/in_memory/memory_store.rs
+++ b/crates/fuel-core/src/state/in_memory/memory_store.rs
@@ -191,7 +191,7 @@ where
     fn view_at_height(
         &self,
         _: &Description::Height,
-    ) -> StorageResult<KeyValueView<Self::Column>> {
+    ) -> StorageResult<KeyValueView<Self::Column, Description::Height>> {
         // TODO: https://github.com/FuelLabs/fuel-core/issues/1995
         Err(
             anyhow::anyhow!("The historical view is not implemented for `MemoryStore`")
@@ -199,10 +199,13 @@ where
         )
     }
 
-    fn latest_view(&self) -> StorageResult<IterableKeyValueView<Self::Column>> {
+    fn latest_view(
+        &self,
+    ) -> StorageResult<IterableKeyValueView<Self::Column, Description::Height>> {
         let view = self.create_view();
-        Ok(IterableKeyValueView::from_storage(
+        Ok(IterableKeyValueView::from_storage_and_metadata(
             IterableKeyValueViewWrapper::new(view),
+            None,
         ))
     }