desks_restore: Restore an active desk after a crash and re-login

- Add a feature flag for desks-restore.
- Store a primary user active desk in user prefs and restore it
after a crash and a logging out.

# TODOs
- crbug.com/1143997 : secondary users should restore to a
default active desk when the crash happened in their sessions.
- crbug.com/1145404 : consider adding an UMA metric for restoring
desks.

Bug: 996999
Test: Manual test, ash_unittests --gtest_filter=DesksRestoreMultiUserTest.DesksRestoredFromPrimaryUserPrefsOnly
Change-Id: I301b394b147468474233af507a263a2d9d38fe53
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/2507399Reviewed-by: Ahmed Fakhry <afakhry@chromium.org>
Commit-Queue: Cattalyya Nuengsigkapian <cattalyya@chromium.org>
Cr-Commit-Position: refs/heads/master@{#826895}

ash/public/cpp/ash_features.cc

@@ -27,6 +27,9 @@ const base::Feature kContextualNudges{"ContextualNudges",
 const base::Feature kDarkLightMode{"DarkLightMode",
                                    base::FEATURE_DISABLED_BY_DEFAULT};
 
+const base::Feature kDesksRestore{"DesksRestore",
+                                  base::FEATURE_DISABLED_BY_DEFAULT};
+
 const base::Feature kDisplayAlignAssist{"DisplayAlignAssist",
                                         base::FEATURE_DISABLED_BY_DEFAULT};
 
@@ -218,6 +221,10 @@ bool IsPipRoundedCornersEnabled() {
   return base::FeatureList::IsEnabled(kPipRoundedCorners);
 }
 
+bool IsDesksRestoreEnabled() {
+  return base::FeatureList::IsEnabled(kDesksRestore);
+}
+
 bool IsSeparateNetworkIconsEnabled() {
   return base::FeatureList::IsEnabled(kSeparateNetworkIcons);
 }

ash/public/cpp/ash_features.h

@@ -128,6 +128,10 @@ ASH_PUBLIC_EXPORT extern const base::Feature kPipRoundedCorners;
 // Enables suppression of Displays notifications other than resolution change.
 ASH_PUBLIC_EXPORT extern const base::Feature kReduceDisplayNotifications;
 
+// Enables desks restoring, including an active desk and windows belonging to
+// them.
+ASH_PUBLIC_EXPORT extern const base::Feature kDesksRestore;
+
 // Enables displaying separate network icons for different networks types.
 // https://crbug.com/902409
 ASH_PUBLIC_EXPORT extern const base::Feature kSeparateNetworkIcons;
@@ -203,6 +207,8 @@ ASH_PUBLIC_EXPORT bool IsCaptureModeEnabled();
 
 ASH_PUBLIC_EXPORT bool IsDarkLightModeEnabled();
 
+ASH_PUBLIC_EXPORT bool IsDesksRestoreEnabled();
+
 ASH_PUBLIC_EXPORT bool IsEnhancedDeskAnimations();
 
 ASH_PUBLIC_EXPORT bool IsFullRestoreEnabled();

ash/public/cpp/ash_pref_names.cc

@@ -151,6 +151,8 @@ const char kContextualTooltips[] = "settings.contextual_tooltip.shown_info";
 // name will appear in this list as an empty string. The desk names are stored
 // as UTF8 strings.
 const char kDesksNamesList[] = "ash.desks.desks_names_list";
+// An integer index of a user's active desk.
+const char kDesksActiveDesk[] = "ash.desks.active_desk";
 
 // A boolean pref storing the enabled status of the Docked Magnifier feature.
 const char kDockedMagnifierEnabled[] = "ash.docked_magnifier.enabled";

ash/public/cpp/ash_pref_names.h

@@ -57,6 +57,7 @@ ASH_PUBLIC_EXPORT extern const char kShouldAlwaysShowAccessibilityMenu[];
 ASH_PUBLIC_EXPORT extern const char kContextualTooltips[];
 
 ASH_PUBLIC_EXPORT extern const char kDesksNamesList[];
+ASH_PUBLIC_EXPORT extern const char kDesksActiveDesk[];
 
 ASH_PUBLIC_EXPORT extern const char kDockedMagnifierEnabled[];
 ASH_PUBLIC_EXPORT extern const char kDockedMagnifierScale[];

ash/wm/desks/desks_controller.cc

@@ -248,6 +248,20 @@ const Desk* DesksController::GetTargetActiveDesk() const {
   return active_desk();
 }
 
+void DesksController::RestorePrimaryUserActiveDeskIndex(int active_desk_index) {
+  DCHECK_GE(active_desk_index, 0);
+  DCHECK_LT(active_desk_index, int{desks_.size()});
+  user_to_active_desk_index_[Shell::Get()
+                                 ->session_controller()
+                                 ->GetPrimaryUserSession()
+                                 ->user_info.account_id] = active_desk_index;
+  // Following |OnActiveUserSessionChanged| approach, restoring uses
+  // DesksSwitchSource::kUserSwitch as a desk switch source.
+  // TODO(crbug.com/1145404): consider adding an UMA metric for desks
+  // restoring to change the source to kDeskRestored.
+  ActivateDesk(desks_[active_desk_index].get(), DesksSwitchSource::kUserSwitch);
+}
+
 void DesksController::Shutdown() {
   animation_.reset();
 }
@@ -561,16 +575,15 @@ void DesksController::OnWindowActivated(ActivationReason reason,
 
 void DesksController::OnActiveUserSessionChanged(const AccountId& account_id) {
   // TODO(afakhry): Remove this when multi-profile support goes away.
-  if (!current_account_id_.is_valid()) {
-    // This is the login of the first primary user. No need to switch any desks.
-    current_account_id_ = account_id;
+  DCHECK(current_account_id_.is_valid());
+  if (current_account_id_ == account_id) {
     return;
   }
 
   user_to_active_desk_index_[current_account_id_] = GetDeskIndex(active_desk_);
   current_account_id_ = account_id;
 
-  // Note the following constraints:
+  // Note the following constraints for secondary users:
   // - Simultaneously logged-in users share the same number of desks.
   // - We don't sync and restore the number of desks nor the active desk
   //   position from previous login sessions.
@@ -593,6 +606,8 @@ void DesksController::OnActiveUserSessionChanged(const AccountId& account_id) {
 }
 
 void DesksController::OnFirstSessionStarted() {
+  current_account_id_ =
+      Shell::Get()->session_controller()->GetActiveAccountId();
   desks_restore_util::RestorePrimaryUserDesks();
 }
 
@@ -633,13 +648,21 @@ void DesksController::ActivateDeskInternal(const Desk* desk,
 
   // If in the middle of a window cycle gesture, reset the window cycle list
   // contents so it contains the new active desk's windows.
+  auto* shell = Shell::Get();
   if (features::IsAltTabLimitedToActiveDesk()) {
-    auto* window_cycle_controller = Shell::Get()->window_cycle_controller();
+    auto* window_cycle_controller = shell->window_cycle_controller();
     window_cycle_controller->MaybeResetCycleList();
   }
 
   for (auto& observer : observers_)
     observer.OnDeskActivationChanged(active_desk_, old_active);
+
+  // Only update active desk prefs when a primary user switches a desk.
+  if (features::IsDesksRestoreEnabled() &&
+      shell->session_controller()->IsUserPrimary()) {
+    desks_restore_util::UpdatePrimaryUserActiveDeskPrefs(
+        GetDeskIndex(active_desk_));
+  }
 }
 
 void DesksController::RemoveDeskInternal(const Desk* desk,

ash/wm/desks/desks_controller.h

@@ -79,6 +79,9 @@ class ASH_EXPORT DesksController : public DesksHelper,
   // switch animation is in progress.
   const Desk* GetTargetActiveDesk() const;
 
+  // Restores the primary user's activate desk at active_desk_index.
+  void RestorePrimaryUserActiveDeskIndex(int active_desk_index);
+
   // Destroys any pending animations in preparation for shutdown.
   void Shutdown();
 

ash/wm/desks/desks_restore_util.cc

@@ -4,6 +4,7 @@
 
 #include "ash/wm/desks/desks_restore_util.h"
 
+#include "ash/public/cpp/ash_features.h"
 #include "ash/public/cpp/ash_pref_names.h"
 #include "ash/session/session_controller_impl.h"
 #include "ash/shell.h"
@@ -32,10 +33,23 @@ PrefService* GetPrimaryUserPrefService() {
   return Shell::Get()->session_controller()->GetPrimaryUserPrefService();
 }
 
+// Check if the desk index is valid against a list of existing desks in
+// DesksController.
+bool IsValidDeskIndex(int desk_index) {
+  return desk_index >= 0 &&
+         desk_index < int{DesksController::Get()->desks().size()} &&
+         desk_index < int{desks_util::kMaxNumberOfDesks};
+}
+
 }  // namespace
 
 void RegisterProfilePrefs(PrefRegistrySimple* registry) {
+  constexpr int kDefaultActiveDeskIndex = 0;
   registry->RegisterListPref(prefs::kDesksNamesList);
+  if (features::IsDesksRestoreEnabled()) {
+    registry->RegisterIntegerPref(prefs::kDesksActiveDesk,
+                                  kDefaultActiveDeskIndex);
+  }
 }
 
 void RestorePrimaryUserDesks() {
@@ -75,6 +89,19 @@ void RestorePrimaryUserDesks() {
     }
     ++index;
   }
+
+  // Restore an active desk for the primary user.
+  if (features::IsDesksRestoreEnabled()) {
+    const int active_desk_index =
+        primary_user_prefs->GetInteger(prefs::kDesksActiveDesk);
+
+    // A crash in between prefs::kDesksNamesList and prefs::kDesksActiveDesk
+    // can cause an invalid active desk index.
+    if (!IsValidDeskIndex(active_desk_index))
+      return;
+
+    desks_controller->RestorePrimaryUserActiveDeskIndex(active_desk_index);
+  }
 }
 
 void UpdatePrimaryUserDesksPrefs() {
@@ -103,6 +130,22 @@ void UpdatePrimaryUserDesksPrefs() {
   DCHECK_EQ(pref_data->GetSize(), desks.size());
 }
 
+void UpdatePrimaryUserActiveDeskPrefs(int active_desk_index) {
+  DCHECK(features::IsDesksRestoreEnabled());
+  DCHECK(Shell::Get()->session_controller()->IsUserPrimary());
+  DCHECK(IsValidDeskIndex(active_desk_index));
+  if (g_pause_desks_prefs_updates)
+    return;
+
+  PrefService* primary_user_prefs = GetPrimaryUserPrefService();
+  if (!primary_user_prefs) {
+    // Can be null in tests.
+    return;
+  }
+
+  primary_user_prefs->SetInteger(prefs::kDesksActiveDesk, active_desk_index);
+}
+
 }  // namespace desks_restore_util
 
 }  // namespace ash

ash/wm/desks/desks_restore_util.h

@@ -22,6 +22,10 @@ void RestorePrimaryUserDesks();
 // change to the desks count or their names is triggered.
 void UpdatePrimaryUserDesksPrefs();
 
+// Called to update the active desk restore prefs for the primary user, whenever
+// the primary user switches an active desk.
+void UpdatePrimaryUserActiveDeskPrefs(int active_desk_index);
+
 }  // namespace desks_restore_util
 
 }  // namespace ash

ash/wm/desks/desks_unittests.cc

@@ -2879,7 +2879,8 @@ constexpr char kUser2Email[] = "user2@desks";
 }  // namespace
 
 class DesksMultiUserTest : public NoSessionAshTestBase,
-                           public MultiUserWindowManagerDelegate {
+                           public MultiUserWindowManagerDelegate,
+                           public ::testing::WithParamInterface<bool> {
  public:
   DesksMultiUserTest() = default;
   ~DesksMultiUserTest() override = default;
@@ -2892,7 +2893,10 @@ class DesksMultiUserTest : public NoSessionAshTestBase,
 
   // AshTestBase:
   void SetUp() override {
+    if (GetParam())
+      scoped_feature_list_.InitAndEnableFeature(features::kDesksRestore);
     NoSessionAshTestBase::SetUp();
+
     TestSessionControllerClient* session_controller =
         GetSessionControllerClient();
     session_controller->Reset();
@@ -2929,6 +2933,8 @@ class DesksMultiUserTest : public NoSessionAshTestBase,
                                  bool teleported) override {}
   void OnTransitionUserShelfToNewAccount() override {}
 
+  bool IsDesksRestoreEnabled() const { return GetParam(); }
+
   AccountId GetUser1AccountId() const {
     return AccountId::FromUserEmail(kUser1Email);
   }
@@ -2964,6 +2970,8 @@ class DesksMultiUserTest : public NoSessionAshTestBase,
   }
 
  private:
+  base::test::ScopedFeatureList scoped_feature_list_;
+
   std::unique_ptr<MultiUserWindowManager> multi_user_window_manager_;
 
   TestingPrefServiceSimple* user_1_prefs_ = nullptr;
@@ -2972,7 +2980,7 @@ class DesksMultiUserTest : public NoSessionAshTestBase,
   DISALLOW_COPY_AND_ASSIGN(DesksMultiUserTest);
 };
 
-TEST_F(DesksMultiUserTest, SwitchUsersBackAndForth) {
+TEST_P(DesksMultiUserTest, SwitchUsersBackAndForth) {
   SimulateUserLogin(GetUser1AccountId());
   auto* controller = DesksController::Get();
   NewDesk();
@@ -3029,7 +3037,7 @@ TEST_F(DesksMultiUserTest, SwitchUsersBackAndForth) {
   EXPECT_TRUE(win3->IsVisible());
 }
 
-TEST_F(DesksMultiUserTest, RemoveDesks) {
+TEST_P(DesksMultiUserTest, RemoveDesks) {
   SimulateUserLogin(GetUser1AccountId());
   // Create two desks with several windows with different app types that
   // belong to different users.
@@ -3126,7 +3134,7 @@ TEST_F(DesksMultiUserTest, RemoveDesks) {
   EXPECT_TRUE(win6->IsVisible());
 }
 
-TEST_F(DesksMultiUserTest, SwitchingUsersEndsOverview) {
+TEST_P(DesksMultiUserTest, SwitchingUsersEndsOverview) {
   SimulateUserLogin(GetUser1AccountId());
   OverviewController* overview_controller = Shell::Get()->overview_controller();
   EXPECT_TRUE(overview_controller->StartOverview());
@@ -3137,8 +3145,15 @@ TEST_F(DesksMultiUserTest, SwitchingUsersEndsOverview) {
 
 using DesksRestoreMultiUserTest = DesksMultiUserTest;
 
-TEST_F(DesksRestoreMultiUserTest, DesksRestoredFromPrimaryUserPrefsOnly) {
+TEST_P(DesksRestoreMultiUserTest, DesksRestoredFromPrimaryUserPrefsOnly) {
+  constexpr int kDefaultActiveDesk = 0;
+  constexpr int kUser1StoredActiveDesk = 2;
   InitPrefsWithDesksRestoreData(user_1_prefs());
+
+  // Set the primary user1's active desk prefs to kUser1StoredActiveDesk.
+  if (IsDesksRestoreEnabled())
+    user_1_prefs()->SetInteger(prefs::kDesksActiveDesk, kUser1StoredActiveDesk);
+
   SimulateUserLogin(GetUser1AccountId());
   // User 1 is the first to login, hence the primary user.
   auto* controller = DesksController::Get();
@@ -3157,14 +3172,35 @@ TEST_F(DesksRestoreMultiUserTest, DesksRestoredFromPrimaryUserPrefsOnly) {
   };
 
   verify_desks("Before switching users");
+  if (IsDesksRestoreEnabled()) {
+    // The primary user1 should restore the saved active desk from its pref.
+    EXPECT_EQ(desks[kUser1StoredActiveDesk]->container_id(),
+              desks_util::GetActiveDeskContainerId());
+  }
 
   // Switching users should not change anything as restoring happens only at
   // the time when the first user signs in.
   SwitchActiveUser(GetUser2AccountId());
   verify_desks("After switching users");
+  if (IsDesksRestoreEnabled()) {
+    // The secondary user2 should start with a default active desk.
+    EXPECT_EQ(desks[kDefaultActiveDesk]->container_id(),
+              desks_util::GetActiveDeskContainerId());
+
+    // Activating the second desk in the secondary user session should not
+    // affect the primary user1's active desk pref. Moreover, switching back to
+    // user1 session should activate the user1's previously active desk
+    // correctly.
+    ActivateDesk(desks[1].get());
+    EXPECT_EQ(user_1_prefs()->GetInteger(prefs::kDesksActiveDesk),
+              kUser1StoredActiveDesk);
+    SwitchActiveUser(GetUser1AccountId());
+    EXPECT_EQ(desks[kUser1StoredActiveDesk]->container_id(),
+              desks_util::GetActiveDeskContainerId());
+  }
 }
 
-TEST_F(DesksRestoreMultiUserTest,
+TEST_P(DesksRestoreMultiUserTest,
        ChangesMadeBySecondaryUserAffectsOnlyPrimaryUserPrefs) {
   InitPrefsWithDesksRestoreData(user_1_prefs());
   SimulateUserLogin(GetUser1AccountId());
@@ -3733,6 +3769,8 @@ TEST_F(DesksMockTimeTest, DeskTraversalNonTouchpadMetrics) {
 // Instantiate the parametrized tests.
 INSTANTIATE_TEST_SUITE_P(All, DesksTest, ::testing::Bool());
 INSTANTIATE_TEST_SUITE_P(All, PerDeskShelfTest, ::testing::Bool());
+INSTANTIATE_TEST_SUITE_P(All, DesksMultiUserTest, ::testing::Bool());
+INSTANTIATE_TEST_SUITE_P(All, DesksRestoreMultiUserTest, ::testing::Bool());
 
 }  // namespace