If a format change is cancelled, don't mark the database as upgraded, and don't check the upgrade (#7442)

zebra-state/src/service/finalized_state/disk_format/upgrade.rs

@@ -66,7 +66,8 @@ pub struct DbFormatChangeThreadHandle {
     /// A handle that can wait for the running format change thread to finish.
     ///
     /// Panics from this thread are propagated into Zebra's state service.
-    update_task: Option<Arc<JoinHandle<()>>>,
+    /// The task returns an error if the upgrade was cancelled by a shutdown.
+    update_task: Option<Arc<JoinHandle<Result<(), CancelFormatChange>>>>,
 
     /// A channel that tells the running format thread to finish early.
     /// A channel that tells the running format thread to finish early.
@@ -159,7 +160,7 @@ impl DbFormatChange {
                     initial_tip_height,
                     upgrade_db,
                     cancel_receiver,
-                );
+                )
             })
         });
 
@@ -176,7 +177,7 @@ impl DbFormatChange {
     /// Apply this format change to the database.
     ///
     /// Format changes are launched in an independent `std::thread` by `apply_format_upgrade()`.
-    /// This thread runs until the upgrade is finished.
+    /// This thread runs until the upgrade is finished or cancelled.
     ///
     /// See `apply_format_upgrade()` for details.
     fn apply_format_change(
@@ -186,7 +187,7 @@ impl DbFormatChange {
         initial_tip_height: Option<Height>,
         upgrade_db: ZebraDb,
         cancel_receiver: mpsc::Receiver<CancelFormatChange>,
-    ) {
+    ) -> Result<(), CancelFormatChange> {
         match self {
             // Perform any required upgrades, then mark the state as upgraded.
             Upgrade { .. } => self.apply_format_upgrade(
@@ -195,7 +196,7 @@ impl DbFormatChange {
                 initial_tip_height,
                 upgrade_db.clone(),
                 cancel_receiver,
-            ),
+            )?,
 
             NewlyCreated { .. } => {
                 Self::mark_as_newly_created(&config, network);
@@ -223,18 +224,21 @@ impl DbFormatChange {
         // - since this Zebra code knows how to de-duplicate trees, downgrades using this code
         //   still know how to make sure trees are unique
         Self::check_for_duplicate_trees(upgrade_db);
+
+        Ok(())
     }
 
     /// Apply any required format updates to the database.
     /// Format changes should be launched in an independent `std::thread`.
     ///
     /// If `cancel_receiver` gets a message, or its sender is dropped,
-    /// the format change stops running early.
+    /// the format change stops running early, and returns an error.
     ///
     /// See the format upgrade design docs for more details:
     /// <https://github.com/ZcashFoundation/zebra/blob/main/book/src/dev/state-db-upgrades.md#design>
     //
     // New format upgrades must be added to the *end* of this method.
+    #[allow(clippy::unwrap_in_result)]
     fn apply_format_upgrade(
         self,
         config: Config,
@@ -242,7 +246,7 @@ impl DbFormatChange {
         initial_tip_height: Option<Height>,
         db: ZebraDb,
         cancel_receiver: mpsc::Receiver<CancelFormatChange>,
-    ) {
+    ) -> Result<(), CancelFormatChange> {
         let Upgrade {
             newer_running_version,
             older_disk_version,
@@ -270,7 +274,7 @@ impl DbFormatChange {
                 "empty database is fully upgraded"
             );
 
-            return;
+            return Ok(());
         };
 
         // Start of a database upgrade task.
@@ -292,7 +296,7 @@ impl DbFormatChange {
             for (height, tree) in db.sapling_tree_by_height_range(Height(1)..=initial_tip_height) {
                 // Return early if there is a cancel signal.
                 if !matches!(cancel_receiver.try_recv(), Err(mpsc::TryRecvError::Empty)) {
-                    return;
+                    return Err(CancelFormatChange);
                 }
 
                 // Delete any duplicate trees.
@@ -319,7 +323,7 @@ impl DbFormatChange {
             for (height, tree) in db.orchard_tree_by_height_range(Height(1)..=initial_tip_height) {
                 // Return early if there is a cancel signal.
                 if !matches!(cancel_receiver.try_recv(), Err(mpsc::TryRecvError::Empty)) {
-                    return;
+                    return Err(CancelFormatChange);
                 }
 
                 // Delete any duplicate trees.
@@ -339,15 +343,9 @@ impl DbFormatChange {
 
             // Mark the database as upgraded. Zebra won't repeat the upgrade anymore once the
             // database is marked, so the upgrade MUST be complete at this point.
-            info!(
-                ?newer_running_version,
-                "Zebra automatically upgraded the database format to:"
-            );
             Self::mark_as_upgraded_to(&version_for_pruning_trees, &config, network);
         }
 
-        // End of a database upgrade task.
-
         // # New Upgrades Usually Go Here
         //
         // New code goes above this comment!
@@ -355,6 +353,12 @@ impl DbFormatChange {
         // Run the latest format upgrade code after the other upgrades are complete,
         // then mark the format as upgraded. The code should check `cancel_receiver`
         // every time it runs its inner update loop.
+        info!(
+            ?newer_running_version,
+            "Zebra automatically upgraded the database format to:"
+        );
+
+        Ok(())
     }
 
     /// Check that note commitment trees were correctly de-duplicated.