Fix restoring windows to their previous workspace/virtual desktop.

When session restore rebuilds the session restore commands, it was
losing the workspace because it wasn't adding the workspace command to
the list of rebuilt commands. This is most noticeable if you have a lot
of windows open on different workspaces/virtual desktops.

This CL also makes BrowserDesktopWindowTreeHostWin update the workspace
on ::Show, if the workspace isn't set. This makes opening a new browser
window immediately update the workspace, so it will be restored
on session restore without the user having to focus the window.

Bug: 1150225
Change-Id: I38ab861147041f4557e5ae8b62ec0c04763f1f4e
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/2585990
Commit-Queue: David Bienvenu <davidbienvenu@chromium.org>
Reviewed-by: Nico Weber <thakis@chromium.org>
Cr-Commit-Position: refs/heads/master@{#837137}

chrome/browser/sessions/session_service.cc

@@ -776,8 +776,9 @@ void SessionService::BuildCommandsForBrowser(
                                                   browser->user_title()));
   }
 
+  command_storage_manager_->AppendRebuildCommand(
       sessions::CreateSetWindowWorkspaceCommand(
-      browser->session_id(), browser->window()->GetWorkspace());
+          browser->session_id(), browser->window()->GetWorkspace()));
 
   windows_to_track->insert(browser->session_id());
   TabStripModel* tab_strip = browser->tab_strip_model();

chrome/browser/sessions/session_service.h

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

chrome/browser/sessions/session_service_unittest.cc

@@ -7,6 +7,7 @@
 #include <stddef.h>
 
 #include <memory>
+#include <vector>
 
 #include "base/bind.h"
 #include "base/callback_helpers.h"
@@ -63,8 +64,7 @@ class SessionServiceTest : public BrowserWithTestWindowTest {
   void SetUp() override {
     BrowserWithTestWindowTest::SetUp();
 
-    std::string b = base::NumberToString(base::Time::Now().ToInternalValue());
-    TestingProfile* profile = profile_manager()->CreateTestingProfile(b);
+    Profile* profile = browser()->profile();
     SessionService* session_service = new SessionService(profile);
     path_ = profile->GetPath();
 
@@ -1229,6 +1229,34 @@ TEST_F(SessionServiceTest, TabGroupMetadataSaved) {
   }
 }
 
+TEST_F(SessionServiceTest, Workspace) {
+  auto* test_browser_window =
+      static_cast<TestBrowserWindow*>(browser()->window());
+  test_browser_window->set_workspace(window_workspace);
+  AddTab(browser(), GURL("http://foo/1"));
+  // Force a reset, to verify that SessionService::BuildCommandsForBrowser
+  // handles workspaces correctly.
+  service()->ResetFromCurrentBrowsers();
+
+  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(),
+                                                window_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/views/frame/browser_desktop_window_tree_host_win.cc

@@ -277,8 +277,12 @@ void BrowserDesktopWindowTreeHostWin::Show(ui::WindowShowState show_state,
   // and the session service is tracking the window.
   if (virtual_desktop_helper_ &&
       !virtual_desktop_helper_->GetInitialWorkspaceRemembered()) {
+    // If |virtual_desktop_helper_| has an empty workspace, kick off an update,
+    // which will eventually call OnHostWorkspaceChanged.
+    if (virtual_desktop_helper_->GetWorkspace().empty())
+      UpdateWorkspace();
+    else
       OnHostWorkspaceChanged();
-    virtual_desktop_helper_->SetInitialWorkspaceRemembered(true);
   }
   DesktopWindowTreeHostWin::Show(show_state, restore_bounds);
 }
@@ -419,13 +423,7 @@ void BrowserDesktopWindowTreeHostWin::PostHandleMSG(UINT message,
                                                     LPARAM l_param) {
   switch (message) {
     case WM_SETFOCUS: {
-      if (virtual_desktop_helper_) {
-        virtual_desktop_helper_->UpdateWindowDesktopId(
-            GetHWND(),
-            base::BindOnce(
-                &BrowserDesktopWindowTreeHostWin::OnHostWorkspaceChanged,
-                weak_factory_.GetWeakPtr()));
-      }
+      UpdateWorkspace();
       break;
     }
     case WM_CREATE:
@@ -550,6 +548,16 @@ void BrowserDesktopWindowTreeHostWin::OnProfileWasRemoved(
 
 ////////////////////////////////////////////////////////////////////////////////
 // BrowserDesktopWindowTreeHostWin, private:
+
+void BrowserDesktopWindowTreeHostWin::UpdateWorkspace() {
+  if (!virtual_desktop_helper_)
+    return;
+  virtual_desktop_helper_->UpdateWindowDesktopId(
+      GetHWND(),
+      base::BindOnce(&BrowserDesktopWindowTreeHostWin::OnHostWorkspaceChanged,
+                     weak_factory_.GetWeakPtr()));
+}
+
 bool BrowserDesktopWindowTreeHostWin::IsOpaqueHostedAppFrame() const {
   // TODO(https://crbug.com/868239): Support Windows 7 Aero glass for web-app
   // window titlebar controls.

chrome/browser/ui/views/frame/browser_desktop_window_tree_host_win.h

@@ -83,6 +83,10 @@ class BrowserDesktopWindowTreeHostWin
   void OnProfileWasRemoved(const base::FilePath& profile_path,
                            const base::string16& profile_name) override;
 
+  // Kicks off an asynchronous update of |workspace_|, and notifies
+  // WindowTreeHost of its value.
+  void UpdateWorkspace();
+
   bool IsOpaqueHostedAppFrame() const;
 
   void SetWindowIcon(bool badged);

chrome/test/base/test_browser_window.cc

@@ -243,7 +243,7 @@ ExclusiveAccessContext* TestBrowserWindow::GetExclusiveAccessContext() {
 }
 
 std::string TestBrowserWindow::GetWorkspace() const {
-  return std::string();
+  return workspace_;
 }
 
 bool TestBrowserWindow::IsVisibleOnAllWorkspaces() const {

chrome/test/base/test_browser_window.h

@@ -207,6 +207,8 @@ class TestBrowserWindow : public BrowserWindow {
   FeaturePromoController* SetFeaturePromoController(
       std::unique_ptr<FeaturePromoController> feature_promo_controller);
 
+  void set_workspace(std::string workspace) { workspace_ = workspace; }
+
  protected:
   void DestroyBrowser() override {}
 
@@ -240,6 +242,8 @@ class TestBrowserWindow : public BrowserWindow {
   TestLocationBar location_bar_;
   gfx::NativeWindow native_window_ = nullptr;
 
+  std::string workspace_;
+
   std::unique_ptr<FeaturePromoController> feature_promo_controller_;
 
   base::OnceClosure close_callback_;

components/sessions/core/session_command.h

@@ -12,6 +12,7 @@
 #include <string>
 
 #include "base/macros.h"
+#include "base/strings/string_piece.h"
 #include "components/sessions/core/sessions_export.h"
 
 namespace base {
@@ -50,7 +51,9 @@ class SESSIONS_EXPORT SessionCommand {
   // The contents of the command.
   char* contents() { return const_cast<char*>(contents_.c_str()); }
   const char* contents() const { return contents_.c_str(); }
-
+  base::StringPiece contents_as_string_piece() const {
+    return base::StringPiece(contents_);
+  }
   // Identifier for the command.
   id_type id() const { return id_; }