Only shutdown once for pending update

If an update is pending, restarting the browser will restart the whole
system. Instead of first restarting the browser and then restarting the
system (which will restart the browser again), simply restart the system
once.

Bug: b/216129138
Change-Id: I2543123014fdac0143c46e7ce252b14f5d0a7031
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/3500223
Reviewed-by: Xiyuan Xia <xiyuan@chromium.org>
Reviewed-by: Scott Violet <sky@chromium.org>
Reviewed-by: Mitsuru Oshima <oshima@chromium.org>
Commit-Queue: Miriam Polzer <mpolzer@google.com>
Cr-Commit-Position: refs/heads/main@{#981517}

chrome/browser/lifetime/application_lifetime.cc

@@ -173,6 +173,15 @@ void AttemptRestartInternal(IgnoreUnloadHandlers ignore_unload_handlers) {
   DCHECK(!g_send_stop_request_to_session_manager);
   // Make sure we don't send stop request to the session manager.
   g_send_stop_request_to_session_manager = false;
+
+  // If an update is pending NotifyAndTerminate() will trigger a system reboot,
+  // which in turn will send SIGTERM to Chrome, and that ends up processing
+  // unload handlers.
+  if (browser_shutdown::UpdatePending()) {
+    browser_shutdown::NotifyAndTerminate(true);
+    return;
+  }
+
   // Run exit process in clean stack.
   content::GetUIThreadTaskRunner({})->PostTask(
       FROM_HERE, base::BindOnce(&ExitIgnoreUnloadHandlers));

chrome/browser/lifetime/application_lifetime_chromeos_browsertest.cc

@@ -13,6 +13,7 @@
 #include "chromeos/dbus/dbus_thread_manager.h"
 #include "chromeos/dbus/power/fake_power_manager_client.h"
 #include "chromeos/dbus/session_manager/fake_session_manager_client.h"
+#include "chromeos/dbus/update_engine/fake_update_engine_client.h"
 #include "components/keep_alive_registry/keep_alive_registry.h"
 #include "components/prefs/pref_service.h"
 #include "content/public/test/browser_test.h"
@@ -27,6 +28,9 @@ class ApplicationLifetimeTest : public InProcessBrowserTest,
  public:
   void SetUpOnMainThread() override {
     InProcessBrowserTest::SetUpOnMainThread();
+    fake_update_engine_client_ = new chromeos::FakeUpdateEngineClient;
+    chromeos::DBusThreadManager::GetSetterForTesting()->SetUpdateEngineClient(
+        std::unique_ptr<ash::UpdateEngineClient>(fake_update_engine_client_));
     BrowserList::AddObserver(this);
   }
 
@@ -40,13 +44,33 @@ class ApplicationLifetimeTest : public InProcessBrowserTest,
     quits_on_browser_closing_->Run();
   }
 
+ protected:
+  void FakePendingUpdate() {
+    update_engine::StatusResult update_status;
+    update_status.set_current_operation(
+        update_engine::Operation::UPDATED_NEED_REBOOT);
+    fake_update_engine_client_->set_default_status(update_status);
+
+    // Shutdown code in termination_notification.cc is only run once.
+    // Exit Chrome after reboot after update would have been sent to
+    // update_engine. Otherwise the browsertest framework wont be able to stop
+    // Chrome gracefully.
+    fake_update_engine_client_->set_reboot_after_update_callback(
+        base::BindOnce(&ExitIgnoreUnloadHandlers));
+  }
+
+  bool RequestedRebootAfterUpdate() {
+    return fake_update_engine_client_->reboot_after_update_call_count() > 0;
+  }
+
  private:
   void OnBrowserClosing(Browser* browser) override {
     if (quits_on_browser_closing_)
       quits_on_browser_closing_->Quit();
   }
 
   absl::optional<base::RunLoop> quits_on_browser_closing_;
+  chromeos::FakeUpdateEngineClient* fake_update_engine_client_ = nullptr;
 };
 
 IN_PROC_BROWSER_TEST_F(ApplicationLifetimeTest,
@@ -70,6 +94,32 @@ IN_PROC_BROWSER_TEST_F(ApplicationLifetimeTest,
   WaitForBrowserToClose();
 }
 
+IN_PROC_BROWSER_TEST_F(ApplicationLifetimeTest,
+                       AttemptRestartWithPendingUpdateReboots) {
+  FakePendingUpdate();
+
+  AttemptRestart();
+
+  // Session Manager is not going to stop session.
+  EXPECT_FALSE(IsAttemptingShutdown());
+  auto* fake_session_manager_client = chromeos::FakeSessionManagerClient::Get();
+  EXPECT_FALSE(fake_session_manager_client->session_stopped());
+
+  // No reboot requested via power manager.
+  auto* fake_power_manager_client = chromeos::FakePowerManagerClient::Get();
+  EXPECT_EQ(fake_power_manager_client->num_request_restart_calls(), 0);
+
+  // Reboot requested via update engine client.
+  EXPECT_TRUE(RequestedRebootAfterUpdate());
+
+  // Restart flags are set.
+  PrefService* pref_service = g_browser_process->local_state();
+  EXPECT_TRUE(pref_service->GetBoolean(prefs::kWasRestarted));
+  EXPECT_TRUE(KeepAliveRegistry::GetInstance()->IsRestarting());
+
+  WaitForBrowserToClose();
+}
+
 IN_PROC_BROWSER_TEST_F(ApplicationLifetimeTest, AttemptRelaunchRelaunchesOs) {
   AttemptRelaunch();
 

chrome/browser/lifetime/termination_notification.cc

@@ -21,6 +21,19 @@
 #endif
 
 namespace browser_shutdown {
+namespace {
+
+#if BUILDFLAG(IS_CHROMEOS_ASH)
+chromeos::UpdateEngineClient* GetUpdateEngineClient() {
+  DCHECK(chromeos::DBusThreadManager::IsInitialized());
+  auto* update_engine_client =
+      chromeos::DBusThreadManager::Get()->GetUpdateEngineClient();
+  DCHECK(update_engine_client);
+  return update_engine_client;
+}
+#endif
+
+}  // namespace
 
 void NotifyAppTerminating() {
   static bool notified = false;
@@ -53,23 +66,29 @@ void NotifyAndTerminate(bool fast_path, RebootPolicy reboot_policy) {
   if (chromeos::PowerPolicyController::IsInitialized())
     chromeos::PowerPolicyController::Get()->NotifyChromeIsExiting();
 
-  if (chromeos::DBusThreadManager::IsInitialized()) {
-    // If we're on a ChromeOS device, reboot if an update has been applied,
-    // or else signal the session manager to log out.
-    chromeos::UpdateEngineClient* update_engine_client =
-        chromeos::DBusThreadManager::Get()->GetUpdateEngineClient();
-    if (update_engine_client->GetLastStatus().current_operation() ==
-            update_engine::Operation::UPDATED_NEED_REBOOT ||
-        reboot_policy == RebootPolicy::kForceReboot) {
-      update_engine_client->RebootAfterUpdate();
-    } else if (chrome::IsAttemptingShutdown()) {
-      // Don't ask SessionManager to stop session if the shutdown request comes
-      // from session manager.
-      ash::SessionTerminationManager::Get()->StopSession(
-          login_manager::SessionStopReason::REQUEST_FROM_SESSION_MANAGER);
-    }
+  // Reboot if an update has been applied.
+  if (UpdatePending() || reboot_policy == RebootPolicy::kForceReboot) {
+    GetUpdateEngineClient()->RebootAfterUpdate();
+    return;
+  }
+
+  // Signal session manager to stop the session if Chrome has initiated an
+  // attempt to do so.
+  if (chrome::IsAttemptingShutdown() && ash::SessionTerminationManager::Get()) {
+    ash::SessionTerminationManager::Get()->StopSession(
+        login_manager::SessionStopReason::REQUEST_FROM_SESSION_MANAGER);
   }
 #endif
 }
 
+#if BUILDFLAG(IS_CHROMEOS_ASH)
+bool UpdatePending() {
+  if (!chromeos::DBusThreadManager::IsInitialized())
+    return false;
+
+  return GetUpdateEngineClient()->GetLastStatus().current_operation() ==
+         update_engine::Operation::UPDATED_NEED_REBOOT;
+}
+#endif
+
 }  // namespace browser_shutdown

chrome/browser/lifetime/termination_notification.h

@@ -5,6 +5,8 @@
 #ifndef CHROME_BROWSER_LIFETIME_TERMINATION_NOTIFICATION_H_
 #define CHROME_BROWSER_LIFETIME_TERMINATION_NOTIFICATION_H_
 
+#include "build/chromeos_buildflags.h"
+
 namespace browser_shutdown {
 
 // Emits APP_TERMINATING notification. It is guaranteed that the
@@ -21,6 +23,10 @@ enum class RebootPolicy { kForceReboot, kOptionalReboot };
 void NotifyAndTerminate(bool fast_path);
 void NotifyAndTerminate(bool fast_path, RebootPolicy reboot_policy);
 
+#if BUILDFLAG(IS_CHROMEOS_ASH)
+bool UpdatePending();
+#endif
+
 }  // namespace browser_shutdown
 
 #endif  // CHROME_BROWSER_LIFETIME_TERMINATION_NOTIFICATION_H_

chromeos/dbus/update_engine/fake_update_engine_client.cc

@@ -5,6 +5,7 @@
 #include "chromeos/dbus/update_engine/fake_update_engine_client.h"
 
 #include "base/bind.h"
+#include "base/callback.h"
 #include "base/threading/thread_task_runner_handle.h"
 
 namespace chromeos {
@@ -48,6 +49,9 @@ void FakeUpdateEngineClient::CanRollbackCheck(RollbackCheckCallback callback) {
 }
 
 void FakeUpdateEngineClient::RebootAfterUpdate() {
+  if (reboot_after_update_callback_) {
+    std::move(reboot_after_update_callback_).Run();
+  }
   reboot_after_update_call_count_++;
 }
 

chromeos/dbus/update_engine/fake_update_engine_client.h

@@ -8,6 +8,7 @@
 #include <map>
 #include <string>
 
+#include "base/callback_forward.h"
 #include "base/component_export.h"
 #include "base/containers/queue.h"
 #include "base/observer_list.h"
@@ -79,6 +80,10 @@ class COMPONENT_EXPORT(CHROMEOS_DBUS_UPDATE_ENGINE) FakeUpdateEngineClient
   void set_update_check_result(
       const UpdateEngineClient::UpdateCheckResult& result);
 
+  void set_reboot_after_update_callback(base::OnceClosure callback) {
+    reboot_after_update_callback_ = std::move(callback);
+  }
+
   void set_can_rollback_check_result(bool result) {
     can_rollback_stub_result_ = result;
   }
@@ -123,6 +128,7 @@ class COMPONENT_EXPORT(CHROMEOS_DBUS_UPDATE_ENGINE) FakeUpdateEngineClient
   base::queue<update_engine::StatusResult> status_queue_;
   update_engine::StatusResult default_status_;
   UpdateCheckResult update_check_result_ = UPDATE_RESULT_SUCCESS;
+  base::OnceClosure reboot_after_update_callback_;
   bool can_rollback_stub_result_ = false;
   int reboot_after_update_call_count_ = 0;
   int request_update_check_call_count_ = 0;