Explicitly handle Visual Studio needing a restart

Author: ChrisDenton

src/cli/self_update.rs

@@ -380,14 +380,21 @@ pub(crate) fn install(
                 "\nAutomatically download and install Visual Studio 2022 Community edition? (Y/n)",
                 true,
             )? {
-                if let Err(e) = try_install_msvc(&opts) {
-                    // Make sure the console doesn't exit before the user can
-                    // see the error and give the option to continue anyway.
-                    report_error(&e);
-                    if !common::confirm("\nContinue installing rustup? (y/N)", false)? {
-                        info!("aborting installation");
+                match try_install_msvc(&opts) {
+                    Err(e) => {
+                        // Make sure the console doesn't exit before the user can
+                        // see the error and give the option to continue anyway.
+                        report_error(&e);
+                        if !common::confirm("\nContinue installing rustup? (y/N)", false)? {
+                            info!("aborting installation");
+                            return Ok(utils::ExitCode(0));
+                        }
+                    }
+                    Ok(ContinueInstall::No) => {
+                        ensure_prompt()?;
                         return Ok(utils::ExitCode(0));
                     }
+                    _ => {}
                 }
             } else {
                 md(&mut term, MSVC_MANUAL_INSTALL_MESSAGE);

src/cli/self_update/windows.rs

@@ -66,7 +66,7 @@ pub(crate) fn do_msvc_check(opts: &InstallOpts<'_>) -> Option<VsInstallPlan> {
     }
 }
 
-#[derive(Debug)]
+#[derive(Debug, Eq, PartialEq)]
 struct VsInstallError(i32);
 impl std::error::Error for VsInstallError {}
 impl fmt::Display for VsInstallError {
@@ -97,8 +97,21 @@ impl fmt::Display for VsInstallError {
         write!(f, "{} (exit code {})", message, self.0)
     }
 }
+impl VsInstallError {
+    const REBOOTING_NOW: Self = Self(1641);
+    const REBOOT_REQUIRED: Self = Self(3010);
+}
+
+pub(crate) enum ContinueInstall {
+    Yes,
+    No,
+}
 
-pub(crate) fn try_install_msvc(opts: &InstallOpts<'_>) -> Result<()> {
+/// Tries to install the needed Visual Studio components.
+///
+/// Returns `Ok(ContinueInstall::No)` if installing Visual Studio was successful
+/// but the rustup install should not be continued at this time.
+pub(crate) fn try_install_msvc(opts: &InstallOpts<'_>) -> Result<ContinueInstall> {
     // download the installer
     let visual_studio_url = utils::parse_url("https://aka.ms/vs/17/release/vs_community.exe")?;
 
@@ -145,20 +158,35 @@ pub(crate) fn try_install_msvc(opts: &InstallOpts<'_>) -> Result<()> {
         .context("error running Visual Studio installer")?;
 
     if exit_status.success() {
-        Ok(())
+        Ok(ContinueInstall::Yes)
     } else {
-        let err = VsInstallError(exit_status.code().unwrap());
-        // It's possible that the installer returned a non-zero exit code
-        // even though the required components were successfully installed.
-        // In that case we warn about the error but continue on.
-        let have_msvc = do_msvc_check(opts).is_none();
-        let has_libs = has_windows_sdk_libs();
-        if have_msvc && has_libs {
-            warn!("Visual Studio is installed but a problem ocurred during installation");
-            warn!("{}", err);
-            Ok(())
-        } else {
-            Err(err).context("failed to install Visual Studio")
+        match VsInstallError(exit_status.code().unwrap()) {
+            err @ VsInstallError::REBOOT_REQUIRED => {
+                // A reboot is required but the user opted to delay it.
+                warn!("{}", err);
+                Ok(ContinueInstall::Yes)
+            }
+            err @ VsInstallError::REBOOTING_NOW => {
+                // The user is wanting to reboot right now, so we should
+                // not continue the install.
+                warn!("{}", err);
+                info!("\nRun rustup-init after restart to continue install");
+                Ok(ContinueInstall::No)
+            }
+            err => {
+                // It's possible that the installer returned a non-zero exit code
+                // even though the required components were successfully installed.
+                // In that case we warn about the error but continue on.
+                let have_msvc = do_msvc_check(opts).is_none();
+                let has_libs = has_windows_sdk_libs();
+                if have_msvc && has_libs {
+                    warn!("Visual Studio is installed but a problem ocurred during installation");
+                    warn!("{}", err);
+                    Ok(ContinueInstall::Yes)
+                } else {
+                    Err(err).context("failed to install Visual Studio")
+                }
+            }
         }
     }
 }