slight rearrangement

rs/sns/governance/src/governance.rs

@@ -4682,10 +4682,6 @@ impl Governance {
 
         self.process_proposals();
 
-        if self.should_check_upgrade_status() {
-            self.check_upgrade_status().await;
-        }
-
         let should_distribute_rewards = self.should_distribute_rewards();
 
         // Getting the total governance token supply from the ledger is expensive enough
@@ -4716,71 +4712,110 @@ impl Governance {
 
         self.maybe_gc();
 
+        if self.should_check_upgrade_status() {
+            self.check_upgrade_status().await;
+        }
+
         if self.should_refresh_cached_upgrade_steps() {
             self.temporarily_lock_refresh_cached_upgrade_steps();
             self.refresh_cached_upgrade_steps().await;
         }
 
-        if self.proto.pending_version.is_none() {
-            if let Some(deployed_version) = &self.proto.deployed_version {
-                if let Some(target_version) = &self.proto.target_version {
-                    if let Some(CachedUpgradeSteps {
-                        upgrade_steps:
-                            Some(Versions {
-                                versions: upgrade_steps,
-                            }),
-                        ..
-                    }) = &self.proto.cached_upgrade_steps
-                    {
-                        if let Some(current_position) =
-                            upgrade_steps.iter().position(|v| v == deployed_version)
-                        {
-                            if let Some(target_position) =
-                                upgrade_steps.iter().position(|v| v == target_version)
-                            {
-                                if current_position + 1 < upgrade_steps.len()
-                                    && target_position > current_position
-                                {
-                                    let next_version = upgrade_steps[current_position + 1].clone();
-                                    let (next_canister_type, next_hash) =
-                                        match canister_type_and_wasm_hash_for_upgrade(
-                                            deployed_version,
-                                            &next_version,
-                                        ) {
-                                            Ok((canister_type, wasm_hash)) => {
-                                                (canister_type, wasm_hash)
-                                            }
-                                            Err(err) => {
-                                                log!(ERROR, "Upgrade attempt failed: {}", err);
-                                                self.proto.target_version = None;
-                                                return;
-                                            }
-                                        };
-
-                                    self.proto.pending_version = Some(UpgradeInProgress {
-                                        target_version: Some(next_version),
-                                        mark_failed_at_seconds: self.env.now() + 5 * 60,
-                                        checking_upgrade_lock: 0,
-                                        proposal_id: None,
-                                    });
-                                    println!("Initiating upgrade!!!");
-                                    let upgrade_attempt = self
-                                        .upgrade_test_hooray(next_hash, next_canister_type)
-                                        .await;
-                                    if let Err(err) = upgrade_attempt {
-                                        log!(ERROR, "Upgrade attempt failed: {}", err);
-                                        self.proto.pending_version = None;
-                                        self.proto.target_version = None;
-                                    }
-                                }
-                            }
-                        }
-                    }
+        self.initiate_upgrade_if_needed().await;
+    }
+
+    /// Checks if an automatic upgrade is needed and initiates it.
+    /// An automatic upgrade may be needed if `target_version` is set to a future version
+    async fn initiate_upgrade_if_needed(&mut self) {
+        if self.proto.pending_version.is_some() {
+            // An upgrade is already in progress
+            return;
+        }
+
+        if let Some(deployed_version) = self.proto.deployed_version.clone() {
+            if let Some(target_version) = self.proto.target_version.clone() {
+                self.validate_and_prepare_upgrade(deployed_version, target_version)
+                    .await;
+            }
+        }
+    }
+
+    /// Validates that the target version is after the deployed version
+    /// and prepares the next upgrade step.
+    async fn validate_and_prepare_upgrade(
+        &mut self,
+        deployed_version: Version,
+        target_version: Version,
+    ) {
+        let Some(CachedUpgradeSteps {
+            upgrade_steps: Some(Versions {
+                versions: upgrade_steps,
+            }),
+            ..
+        }) = &self.proto.cached_upgrade_steps
+        else {
+            log!(ERROR, "Cached upgrade steps set to None");
+            return;
+        };
+
+        // Find the current position of the deployed version
+        let Some(deployed_position) = upgrade_steps.iter().position(|v| v == &deployed_version)
+        else {
+            log!(ERROR, "Deployed version is not on the upgrade path");
+            self.invalidate_cached_upgrade_steps();
+            return;
+        };
+
+        // Find the target position of the target version
+        let Some(target_position) = upgrade_steps.iter().position(|v| v == &target_version) else {
+            log!(ERROR, "Target version is not on the upgrade path");
+            return;
+        };
+
+        if target_position > deployed_position {
+            // since `target_position > deployed_position`, `deployed_position + 1 < upgrade_steps.len()`
+            let next_version = upgrade_steps[deployed_position + 1].clone();
+
+            match canister_type_and_wasm_hash_for_upgrade(&deployed_version, &next_version) {
+                Ok((canister_type, wasm_hash)) => {
+                    self.prepare_upgrade(next_version, canister_type, wasm_hash)
+                        .await;
+                }
+                Err(err) => {
+                    log!(ERROR, "Upgrade attempt failed: {}", err);
+                    self.proto.target_version = None;
                 }
             }
         }
     }
 
+    async fn prepare_upgrade(
+        &mut self,
+        next_version: Version,
+        canister_type: SnsCanisterType,
+        wasm_hash: Vec<u8>,
+    ) {
+        self.proto.pending_version = Some(UpgradeInProgress {
+            target_version: Some(next_version.clone()),
+            mark_failed_at_seconds: self.env.now() + 5 * 60,
+            checking_upgrade_lock: 0,
+            proposal_id: None,
+        });
+
+        println!("Initiating upgrade to version: {:?}", next_version);
+        let upgrade_attempt = self.upgrade_test_hooray(wasm_hash, canister_type).await;
+        if let Err(err) = upgrade_attempt {
+            log!(ERROR, "Upgrade attempt failed: {}", err);
+            self.proto.pending_version = None;
+            self.proto.target_version = None;
+        }
+    }
+
+    /// Invalidates the cached upgrade steps.
+    fn invalidate_cached_upgrade_steps(&mut self) {
+        self.proto.cached_upgrade_steps = None;
+    }
+
     pub fn temporarily_lock_refresh_cached_upgrade_steps(&mut self) {
         if let Some(ref mut cached_upgrade_steps) = self.proto.cached_upgrade_steps {
             cached_upgrade_steps.requested_timestamp_seconds = Some(self.env.now());