desks_restore: Restore windows to their parent desks.

- save a workspace in window's property.
- trigger saving to browser session when property is updated, but
need force saving in `WindowOpened()` for newly created window.
- restore each window to the right desk by assigning to the right
parent desk container in `GetDefaultParentForWindow()`.
- resolve desk activation conflict between browser session and
an active desk restore by checking `IsSessionRestoreInProgress()`
in ash to avoid every single window activation during restore.
- resolve the parent desk conflict between the first restored window
and a blank window by stop reusing the blank window for restore
and always create a new window during restoration.
- resolve clobbering tab restore by passing null browser when using
desks restore feature.
- add browser tests checking that each window restore to the parent
desk if DesksRestore enabled, otherwise to the first desk, which is a
default active desk after a crash.

`browser_tests --gtest_filter=*SessionRestoreTestChromeOS*`

Bug: 996999
Test: Manual test (a video in crbug comment 27) and browser tests
Change-Id: I0f9202e10a7b861a676725e7390b15fd71a00933
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/2572316
Commit-Queue: Cattalyya Nuengsigkapian <cattalyya@chromium.org>
Reviewed-by: Ahmed Fakhry <afakhry@chromium.org>
Reviewed-by: Scott Violet <sky@chromium.org>
Cr-Commit-Position: refs/heads/master@{#840778}

ash/public/cpp/autotest_desks_api.h

@@ -7,6 +7,7 @@
 
 #include "ash/ash_export.h"
 #include "base/callback_forward.h"
+#include "ui/aura/window.h"
 
 namespace ash {
 
@@ -42,6 +43,9 @@ class ASH_EXPORT AutotestDesksApi {
   // true otherwise.
   bool ActivateAdjacentDesksToTargetIndex(int index,
                                           base::OnceClosure on_complete);
+
+  // Check whether a window belongs to a desk at |desk_index| or not.
+  bool IsWindowInDesk(aura::Window* window, int desk_index);
 };
 
 }  // namespace ash

ash/shell_delegate.h

@@ -111,6 +111,9 @@ class ASH_EXPORT ShellDelegate {
   virtual media_session::MediaSessionService* GetMediaSessionService();
 
   virtual void OpenKeyboardShortcutHelpPage() const {}
+
+  // Returns if window browser sessions are restoring.
+  virtual bool IsSessionRestoreInProgress() const = 0;
 };
 
 }  // namespace ash

ash/test_shell_delegate.cc

@@ -78,4 +78,12 @@ TestShellDelegate::CreateNearbyShareDelegate(
   return std::make_unique<TestNearbyShareDelegate>();
 }
 
+bool TestShellDelegate::IsSessionRestoreInProgress() const {
+  return session_restore_in_progress_;
+}
+
+void TestShellDelegate::SetSessionRestoreInProgress(bool in_progress) {
+  session_restore_in_progress_ = in_progress;
+}
+
 }  // namespace ash

ash/test_shell_delegate.h

@@ -47,9 +47,11 @@ class TestShellDelegate : public ShellDelegate {
       override;
   std::unique_ptr<NearbyShareDelegate> CreateNearbyShareDelegate(
       NearbyShareController* controller) const override;
+  bool IsSessionRestoreInProgress() const override;
 
   void SetCanGoBack(bool can_go_back);
   void SetShouldWaitForTouchAck(bool should_wait_for_touch_ack);
+  void SetSessionRestoreInProgress(bool in_progress);
 
  private:
   // True if the current top window can go back.
@@ -64,6 +66,9 @@ class TestShellDelegate : public ShellDelegate {
   // description.
   bool should_wait_for_touch_ack_ = false;
 
+  // True if window browser sessions are restoring.
+  bool session_restore_in_progress_ = false;
+
   MultiDeviceSetupBinder multidevice_setup_binder_;
 
   DISALLOW_COPY_AND_ASSIGN(TestShellDelegate);

ash/wm/always_on_top_controller.cc

@@ -4,8 +4,10 @@
 
 #include "ash/wm/always_on_top_controller.h"
 
+#include "ash/public/cpp/ash_features.h"
 #include "ash/public/cpp/shell_window_ids.h"
 #include "ash/public/cpp/window_properties.h"
+#include "ash/wm/desks/desks_controller.h"
 #include "ash/wm/desks/desks_util.h"
 #include "ash/wm/window_state.h"
 #include "ash/wm/workspace/workspace_layout_manager.h"
@@ -54,6 +56,15 @@ aura::Window* AlwaysOnTopController::GetContainer(aura::Window* window) const {
     // TODO(afakhry): Do we need to worry about the context of |window| here? Or
     // is it safe to assume that |window| should always be parented to the
     // active desks' container.
+    int window_workspace =
+        window->GetProperty(aura::client::kWindowWorkspaceKey);
+    if (features::IsDesksRestoreEnabled() &&
+        window_workspace != aura::client::kUnassignedWorkspace) {
+      auto* desk_container =
+          DesksController::Get()->GetDeskContainer(root, window_workspace);
+      if (desk_container)
+        return desk_container;
+    }
     return desks_util::GetActiveDeskContainerForRoot(root);
   }
   if (window->parent() && WindowState::Get(window)->IsPip())

ash/wm/desks/autotest_desks_api.cc

@@ -201,4 +201,10 @@ bool AutotestDesksApi::ActivateAdjacentDesksToTargetIndex(
   return true;
 }
 
+bool AutotestDesksApi::IsWindowInDesk(aura::Window* window, int desk_index) {
+  aura::Window* desk_container = DesksController::Get()->GetDeskContainer(
+      window->GetRootWindow(), desk_index);
+  return desk_container->Contains(window);
+}
+
 }  // namespace ash

ash/wm/desks/desk.cc

@@ -8,9 +8,11 @@
 #include <utility>
 
 #include "ash/public/cpp/app_types.h"
+#include "ash/public/cpp/ash_features.h"
 #include "ash/public/cpp/shell_window_ids.h"
 #include "ash/public/cpp/window_properties.h"
 #include "ash/shell.h"
+#include "ash/wm/desks/desks_controller.h"
 #include "ash/wm/mru_window_tracker.h"
 #include "ash/wm/overview/overview_controller.h"
 #include "ash/wm/window_positioner.h"
@@ -23,6 +25,7 @@
 #include "base/containers/adapters.h"
 #include "base/containers/contains.h"
 #include "base/stl_util.h"
+#include "chromeos/ui/base/window_properties.h"
 #include "ui/aura/client/aura_constants.h"
 #include "ui/aura/window_tracker.h"
 #include "ui/display/screen.h"
@@ -197,6 +200,13 @@ void Desk::AddWindowToDesk(aura::Window* window) {
   // there in the first place.
   if (!window->GetProperty(kHideInDeskMiniViewKey))
     NotifyContentChanged();
+
+  // Update the window's workspace to this parent desk.
+  if (features::IsDesksRestoreEnabled() && !is_desk_being_removed_) {
+    auto* desks_controller = DesksController::Get();
+    window->SetProperty(aura::client::kWindowWorkspaceKey,
+                        desks_controller->GetDeskIndex(this));
+  }
 }
 
 void Desk::RemoveWindowFromDesk(aura::Window* window) {
@@ -415,6 +425,10 @@ void Desk::UpdateDeskBackdrops() {
     UpdateBackdropController(GetDeskContainerForRoot(root));
 }
 
+void Desk::SetDeskBeingRemoved() {
+  is_desk_being_removed_ = true;
+}
+
 void Desk::MoveWindowToDeskInternal(aura::Window* window,
                                     Desk* target_desk,
                                     aura::Window* target_root) {

ash/wm/desks/desk.h

@@ -124,6 +124,10 @@ class ASH_EXPORT Desk {
   // visibility on the containers (on all roots) associated with this desk.
   void UpdateDeskBackdrops();
 
+  // Set desk being removed to avoid unwanted action such as `GetDeskIndex()`
+  // when desk is already removed from |desks_| in DesksController.
+  void SetDeskBeingRemoved();
+
  private:
   void MoveWindowToDeskInternal(aura::Window* window,
                                 Desk* target_desk,
@@ -172,6 +176,9 @@ class ASH_EXPORT Desk {
   // are set by the user don't change.
   bool is_name_set_by_user_ = false;
 
+  // True if the desk is being removed.
+  bool is_desk_being_removed_ = false;
+
   DISALLOW_COPY_AND_ASSIGN(Desk);
 };
 

ash/wm/desks/desks_controller.cc

@@ -567,6 +567,13 @@ int DesksController::GetDeskIndex(const Desk* desk) const {
   return -1;
 }
 
+aura::Window* DesksController::GetDeskContainer(aura::Window* target_root,
+                                                int desk_index) {
+  if (desk_index < 0 || desk_index >= int{desks_.size()})
+    return nullptr;
+  return desks_[desk_index]->GetDeskContainerForRoot(target_root);
+}
+
 bool DesksController::BelongsToActiveDesk(aura::Window* window) {
   return desks_util::BelongsToActiveDesk(window);
 }
@@ -607,6 +614,14 @@ void DesksController::OnWindowActivating(ActivationReason reason,
   if (AreDesksBeingModified())
     return;
 
+  // Browser session restore opens all restored windows, so it activates
+  // every single window and activates the parent desk. Therefore, this check
+  // prevents repetitive desk activation. Moreover, when DesksRestore is
+  // enabled, it avoid switching desk back and forth when windows are restored
+  // to different desks.
+  if (Shell::Get()->shell_delegate()->IsSessionRestoreInProgress())
+    return;
+
   if (!gaining_active)
     return;
 
@@ -728,6 +743,7 @@ void DesksController::RemoveDeskInternal(const Desk* desk,
 
   // Keep the removed desk alive until the end of this function.
   std::unique_ptr<Desk> removed_desk = std::move(*iter);
+  removed_desk->SetDeskBeingRemoved();
   DCHECK_EQ(removed_desk.get(), desk);
   auto iter_after = desks_.erase(iter);
 

ash/wm/desks/desks_controller.h

@@ -171,6 +171,10 @@ class ASH_EXPORT DesksController : public DesksHelper,
 
   int GetDeskIndex(const Desk* desk) const;
 
+  // Gets the container of the desk at |desk_index| in a specific screen with a
+  // |target_root|. If desk_index is invalid, it returns nullptr.
+  aura::Window* GetDeskContainer(aura::Window* target_root, int desk_index);
+
   // DesksHelper:
   bool BelongsToActiveDesk(aura::Window* window) override;
   int GetActiveDeskIndex() const override;

chrome/browser/sessions/session_restore.cc

@@ -72,6 +72,7 @@
 #include "extensions/common/extension_set.h"
 
 #if BUILDFLAG(IS_CHROMEOS_ASH)
+#include "ash/public/cpp/ash_features.h"
 #include "chrome/browser/chromeos/boot_times_recorder.h"
 #endif
 
@@ -442,9 +443,17 @@ class SessionRestoreImpl : public BrowserListObserver {
 
       // 5. Restore tabs in |browser|. This will also call Show() on |browser|
       //    if its initial show state is not mimimized.
+      // However, with desks restore enabled, a window is restored to its parent
+      // desk, which can be non-active desk, and left invisible but unminimized.
       RestoreTabsToBrowser(*(*i), browser, initial_tab_count, created_contents);
+#if BUILDFLAG(IS_CHROMEOS_ASH)
+      DCHECK(browser->window()->IsVisible() ||
+             browser->window()->IsMinimized() ||
+             ash::features::IsDesksRestoreEnabled());
+#else
       DCHECK(browser->window()->IsVisible() ||
              browser->window()->IsMinimized());
+#endif
 
       // 6. Tabs will be grouped appropriately in RestoreTabsToBrowser. Now
       //    restore the groups' visual data.
@@ -784,9 +793,26 @@ Browser* SessionRestore::RestoreSession(
 
 // static
 void SessionRestore::RestoreSessionAfterCrash(Browser* browser) {
+  auto* profile = browser->profile();
+
+#if BUILDFLAG(IS_CHROMEOS_ASH)
+  // DesksRestore restores a window to the right desk, so we should not
+  // reuse any browser window. Otherwise, the conflict of the parent desk
+  // arises because tabs created in this |browser| should remain in the
+  // current active desk, but the first restored window should be restored
+  // to its saved parent desk before a crash. This also avoids users'
+  // confusion of the current window disappearing from the current desk
+  // after pressing a restore button.
+  if (ash::features::IsDesksRestoreEnabled())
+    browser = nullptr;
+#endif
+
   SessionRestore::BehaviorBitmask behavior =
-      HasSingleNewTabPage(browser) ? SessionRestore::CLOBBER_CURRENT_TAB : 0;
-  SessionRestore::RestoreSession(browser->profile(), browser, behavior,
+      browser && HasSingleNewTabPage(browser)
+          ? SessionRestore::CLOBBER_CURRENT_TAB
+          : 0;
+
+  SessionRestore::RestoreSession(profile, browser, behavior,
                                  std::vector<GURL>());
 }
 

chrome/browser/sessions/session_service.cc

@@ -273,6 +273,11 @@ void SessionService::WindowOpened(Browser* browser) {
   RestoreIfNecessary(std::vector<GURL>(), browser);
   SetWindowType(browser->session_id(), browser->type());
   SetWindowAppName(browser->session_id(), browser->app_name());
+
+  // Save a browser workspace after window is created in `Browser()`.
+  // DesksRestore feature in ash requires this line to restore correctly after
+  // creating a new browser window in a particular desk.
+  SetWindowWorkspace(browser->session_id(), browser->window()->GetWorkspace());
 }
 
 void SessionService::WindowClosing(const SessionID& window_id) {

chrome/browser/sessions/session_service.h

@@ -226,6 +226,7 @@ class SessionService : public sessions::CommandStorageManagerDelegate,
   FRIEND_TEST_ALL_PREFIXES(SessionServiceTest, RestoreActivation2);
   FRIEND_TEST_ALL_PREFIXES(SessionServiceTest, RemoveUnusedRestoreWindowsTest);
   FRIEND_TEST_ALL_PREFIXES(SessionServiceTest, Workspace);
+  FRIEND_TEST_ALL_PREFIXES(SessionServiceTest, WorkspaceSavedOnOpened);
   FRIEND_TEST_ALL_PREFIXES(NoStartupWindowTest, DontInitSessionServiceForApps);
 
   typedef std::map<SessionID, std::pair<int, int>> IdToRange;

chrome/browser/sessions/session_service_unittest.cc

@@ -1257,6 +1257,35 @@ TEST_F(SessionServiceTest, Workspace) {
   EXPECT_TRUE(found_workspace_command);
 }
 
+// Tests that the workspace is saved in the browser session during
+// `SessionService::WindowOpened(),` called in `Browser()` constructor to
+// save the current workspace to newly created browser.
+TEST_F(SessionServiceTest, WorkspaceSavedOnOpened) {
+  const std::string workspace = "xyz";
+  auto* test_browser_window =
+      static_cast<TestBrowserWindow*>(browser()->window());
+  test_browser_window->set_workspace(workspace);
+  service()->WindowOpened(browser());
+
+  sessions::CommandStorageManager* command_storage_manager =
+      service()->GetCommandStorageManagerForTest();
+  const std::vector<std::unique_ptr<sessions::SessionCommand>>&
+      pending_commands = command_storage_manager->pending_commands();
+  bool found_workspace_command = false;
+  std::unique_ptr<sessions::SessionCommand> workspace_command =
+      sessions::CreateSetWindowWorkspaceCommand(browser()->session_id(),
+                                                workspace);
+  for (const auto& command : pending_commands) {
+    if (command->id() == workspace_command->id() &&
+        command->contents_as_string_piece() ==
+            workspace_command->contents_as_string_piece()) {
+      found_workspace_command = true;
+      break;
+    }
+  }
+  EXPECT_TRUE(found_workspace_command);
+}
+
 // Functions used by GetSessionsAndDestroy.
 namespace {
 

chrome/browser/ui/ash/chrome_shell_delegate.cc

@@ -16,6 +16,7 @@
 #include "chrome/browser/nearby_sharing/nearby_share_delegate_impl.h"
 #include "chrome/browser/profiles/profile.h"
 #include "chrome/browser/profiles/profile_manager.h"
+#include "chrome/browser/sessions/session_restore.h"
 #include "chrome/browser/ui/ash/back_gesture_contextual_nudge_delegate.h"
 #include "chrome/browser/ui/ash/chrome_accessibility_delegate.h"
 #include "chrome/browser/ui/ash/chrome_capture_mode_delegate.h"
@@ -212,3 +213,8 @@ ChromeShellDelegate::CreateNearbyShareDelegate(
     ash::NearbyShareController* controller) const {
   return std::make_unique<NearbyShareDelegateImpl>(controller);
 }
+
+bool ChromeShellDelegate::IsSessionRestoreInProgress() const {
+  Profile* profile = ProfileManager::GetActiveUserProfile();
+  return SessionRestore::IsRestoring(profile);
+}

chrome/browser/ui/ash/chrome_shell_delegate.h

@@ -46,6 +46,7 @@ class ChromeShellDelegate : public ash::ShellDelegate {
   media_session::MediaSessionService* GetMediaSessionService() override;
   std::unique_ptr<ash::NearbyShareDelegate> CreateNearbyShareDelegate(
       ash::NearbyShareController* controller) const override;
+  bool IsSessionRestoreInProgress() const override;
 
  private:
   DISALLOW_COPY_AND_ASSIGN(ChromeShellDelegate);

ui/aura/client/aura_constants.cc

@@ -77,6 +77,7 @@ DEFINE_OWNED_UI_CLASS_PROPERTY_KEY(base::string16, kTitleKey, nullptr)
 DEFINE_UI_CLASS_PROPERTY_KEY(int, kTopViewInset, 0)
 DEFINE_OWNED_UI_CLASS_PROPERTY_KEY(gfx::ImageSkia, kWindowIconKey, nullptr)
 DEFINE_UI_CLASS_PROPERTY_KEY(int, kWindowCornerRadiusKey, -1)
+DEFINE_UI_CLASS_PROPERTY_KEY(int, kWindowWorkspaceKey, kUnassignedWorkspace)
 DEFINE_UI_CLASS_PROPERTY_KEY(ui::ZOrderLevel,
                              kZOrderingKey,
                              ui::ZOrderLevel::kNormal)

ui/aura/client/aura_constants.h

@@ -30,6 +30,9 @@ constexpr int kResizeBehaviorCanResize = 1 << 0;
 constexpr int kResizeBehaviorCanMaximize = 1 << 1;
 constexpr int kResizeBehaviorCanMinimize = 1 << 2;
 
+// A value used to represent an unassigned workspace for kWindowWorkspaceKey.
+constexpr int kUnassignedWorkspace = -1;
+
 // Alphabetical sort.
 
 // A property key to store whether accessibility focus falls back to widget or
@@ -158,6 +161,10 @@ AURA_EXPORT extern const WindowProperty<gfx::ImageSkia*>* const kWindowIconKey;
 // Default is -1, meaning "unspecified". 0 Ensures corners are square.
 AURA_EXPORT extern const WindowProperty<int>* const kWindowCornerRadiusKey;
 
+// A property key to indicate a desk index of a workspace this window belongs
+// to. The default value is kUnassignedWorkspace.
+AURA_EXPORT extern const WindowProperty<int>* const kWindowWorkspaceKey;
+
 // A property key to store the z-ordering.
 AURA_EXPORT extern const WindowProperty<ui::ZOrderLevel>* const kZOrderingKey;
 

ui/views/widget/native_widget_aura.cc

@@ -10,6 +10,7 @@
 #include "base/bind.h"
 #include "base/location.h"
 #include "base/single_thread_task_runner.h"
+#include "base/strings/string_number_conversions.h"
 #include "base/strings/string_util.h"
 #include "base/threading/thread_task_runner_handle.h"
 #include "build/build_config.h"
@@ -171,6 +172,14 @@ void NativeWidgetAura::InitNativeWidget(Widget::InitParams params) {
                          *params.corner_radius);
   }
   window_->SetProperty(aura::client::kShowStateKey, params.show_state);
+
+  int desk_index;
+  // Set workspace property of this window created with a specified workspace
+  // in InitParams. The desk index can be kActiveWorkspace=-1, representing
+  // an active desk.
+  if (base::StringToInt(params.workspace, &desk_index))
+    window_->SetProperty(aura::client::kWindowWorkspaceKey, desk_index);
+
   if (params.type == Widget::InitParams::TYPE_BUBBLE)
     wm::SetHideOnDeactivate(window_, true);
   window_->SetTransparent(params.opacity ==
@@ -491,7 +500,10 @@ gfx::Rect NativeWidgetAura::GetRestoredBounds() const {
 }
 
 std::string NativeWidgetAura::GetWorkspace() const {
-  return std::string();
+  int desk_index = window_->GetProperty(aura::client::kWindowWorkspaceKey);
+  return desk_index == aura::client::kUnassignedWorkspace
+             ? std::string()
+             : base::NumberToString(desk_index);
 }
 
 void NativeWidgetAura::SetBounds(const gfx::Rect& bounds) {
@@ -955,6 +967,9 @@ void NativeWidgetAura::OnWindowPropertyChanged(aura::Window* window,
                                                intptr_t old) {
   if (key == aura::client::kShowStateKey)
     delegate_->OnNativeWidgetWindowShowStateChanged();
+
+  if (key == aura::client::kWindowWorkspaceKey)
+    delegate_->OnNativeWidgetWorkspaceChanged();
 }
 
 void NativeWidgetAura::OnResizeLoopStarted(aura::Window* window) {