fix: make the rollback atomic to the kind of database

Author: rymnc

crates/fuel-core/src/combined_database.rs

@@ -318,38 +318,26 @@ impl CombinedDatabase {
     /// we leave it to the caller to decide how to bring them up to date.
     /// We don't rollback the on-chain database as it is the source of truth.
     /// The target height of the rollback is the latest height of the on-chain database.
-    pub fn sync_aux_db_heights<S>(&self, shutdown_listener: &mut S) -> anyhow::Result<()>
-    where
-        S: ShutdownListener,
-    {
-        while !shutdown_listener.is_cancelled() {
-            let on_chain_height = match self.on_chain().latest_height_from_metadata()? {
-                Some(height) => height,
-                None => break, // Exit loop if on-chain height is None
-            };
-
-            let off_chain_height = self.off_chain().latest_height_from_metadata()?;
-            let gas_price_height = self.gas_price().latest_height_from_metadata()?;
+    pub fn sync_aux_db_heights(&self) -> anyhow::Result<()> {
+        let on_chain_height = match self.on_chain().latest_height_from_metadata()? {
+            Some(height) => height,
+            None => return Ok(()), // Exit loop if on-chain height is None
+        };
 
-            // Handle off-chain rollback if necessary
-            if let Some(off_height) = off_chain_height {
-                if off_height > on_chain_height {
-                    self.off_chain().rollback_last_block()?;
-                }
-            }
+        let off_chain_height = self.off_chain().latest_height_from_metadata()?;
+        let gas_price_height = self.gas_price().latest_height_from_metadata()?;
 
-            // Handle gas price rollback if necessary
-            if let Some(gas_height) = gas_price_height {
-                if gas_height > on_chain_height {
-                    self.gas_price().rollback_last_block()?;
-                }
+        // Handle off-chain rollback if necessary
+        if let Some(off_height) = off_chain_height {
+            if off_height > on_chain_height {
+                self.off_chain().rollback_to(on_chain_height)?;
             }
+        }
 
-            // If both off-chain and gas price heights are synced, break
-            if off_chain_height.map_or(true, |h| h <= on_chain_height)
-                && gas_price_height.map_or(true, |h| h <= on_chain_height)
-            {
-                break;
+        // Handle gas price rollback if necessary
+        if let Some(gas_height) = gas_price_height {
+            if gas_height > on_chain_height {
+                self.gas_price().rollback_to(on_chain_height)?;
             }
         }
 

crates/fuel-core/src/database.rs

@@ -272,19 +272,40 @@ where
 impl<Description> Database<Description>
 where
     Description: DatabaseDescription,
+    Description::Height: PartialOrd,
 {
     pub fn rollback_last_block(&self) -> StorageResult<()> {
+        self.perform_rollback(None)
+    }
+
+    pub fn rollback_to(&self, to_height: Description::Height) -> StorageResult<()> {
+        self.perform_rollback(Some(to_height))
+    }
+
+    fn perform_rollback(
+        &self,
+        to_height: Option<Description::Height>,
+    ) -> StorageResult<()> {
         let mut lock = self.inner_storage().stage.height.lock();
-        let height = *lock;
+        let current_height = *lock;
+
+        let current_height = current_height.ok_or_else(|| {
+            anyhow::anyhow!("Database doesn't have a height to rollback")
+        })?;
 
-        let Some(height) = height else {
+        let to_height = to_height.unwrap_or(current_height);
+
+        if to_height > current_height {
             return Err(
-                anyhow::anyhow!("Database doesn't have a height to rollback").into(),
+                anyhow::anyhow!("Can't rollback to the height {:?} because it's greater than the current height {:?}",
+                    to_height, current_height).into(),
             );
-        };
-        self.inner_storage().data.rollback_block_to(&height)?;
-        let new_height = height.rollback_height();
+        }
+
+        self.inner_storage().data.rollback_block_to(&to_height)?;
+        let new_height = to_height.rollback_height();
         *lock = new_height;
+
         tracing::info!(
             "Rollback of the {} to the height {:?} was successful",
             Description::name(),

crates/fuel-core/src/service.rs

@@ -134,7 +134,7 @@ impl FuelService {
 
         // initialize sub services
         tracing::info!("Initializing sub services");
-        database.sync_aux_db_heights(shutdown_listener)?;
+        database.sync_aux_db_heights()?;
         let (services, shared) = sub_services::init_sub_services(&config, database)?;
 
         let sub_services = Arc::new(services);