Desks: Fix querying the active desk when Alt-tab ends

When Alt-tab ends, leading to a desk switch, we should query
the traget active desk, not the currently active desk, since
the desk switch animation will be in progress.

BUG=1042486
TEST=Added new tests

Change-Id: I17481169bf9cd1b9a66bec702a9fafafeac6cc85
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/2004059
Commit-Queue: Ahmed Fakhry <afakhry@chromium.org>
Reviewed-by: James Cook <jamescook@chromium.org>
Cr-Commit-Position: refs/heads/master@{#732704}

ash/wm/desks/desks_controller.cc

@@ -118,6 +118,8 @@ class DesksController::DeskAnimationBase
  public:
   ~DeskAnimationBase() override = default;
 
+  const Desk* ending_desk() const { return ending_desk_; }
+
   // Launches the animation. This should be done once all animators
   // are created and added to `desk_switch_animators_`. This is to avoid any
   // potential race conditions that might happen if one animator finished phase
@@ -200,8 +202,11 @@ class DesksController::DeskAnimationBase
   }
 
  protected:
-  explicit DeskAnimationBase(DesksController* controller)
-      : controller_(controller) {}
+  DeskAnimationBase(DesksController* controller, const Desk* ending_desk)
+      : controller_(controller), ending_desk_(ending_desk) {
+    DCHECK(controller_);
+    DCHECK(ending_desk_);
+  }
 
   // Abstract functions that can be overridden by child classes to do different
   // things when phase (1), and phase (3) completes. Note that
@@ -224,6 +229,9 @@ class DesksController::DeskAnimationBase
   std::vector<std::unique_ptr<RootWindowDeskSwitchAnimator>>
       desk_switch_animators_;
 
+  // The desk that will be active after this animation ends.
+  const Desk* const ending_desk_;
+
  private:
   // Computes the animation smoothness and reports an UMA stat for it.
   void ComputeAnimationSmoothnessAndReport() {
@@ -249,7 +257,7 @@ class DesksController::DeskActivationAnimation
   DeskActivationAnimation(DesksController* controller,
                           const Desk* ending_desk,
                           bool move_left)
-      : DeskAnimationBase(controller) {
+      : DeskAnimationBase(controller, ending_desk) {
     for (auto* root : Shell::GetAllRootWindows()) {
       desk_switch_animators_.emplace_back(
           std::make_unique<RootWindowDeskSwitchAnimator>(root, ending_desk,
@@ -262,6 +270,7 @@ class DesksController::DeskActivationAnimation
 
   // DesksController::AbstractDeskSwitchAnimation:
   void OnStartingDeskScreenshotTakenInternal(const Desk* ending_desk) override {
+    DCHECK_EQ(ending_desk_, ending_desk);
     // The order here matters. Overview must end before ending tablet split view
     // before switching desks. (If clamshell split view is active on one or more
     // displays, then it simply will end when we end overview.) That's because
@@ -312,7 +321,7 @@ class DesksController::DeskRemovalAnimation
                        const Desk* desk_to_activate,
                        bool move_left,
                        DesksCreationRemovalSource source)
-      : DeskAnimationBase(controller),
+      : DeskAnimationBase(controller, desk_to_activate),
         desk_to_remove_(desk_to_remove),
         request_source_(source) {
     DCHECK(!Shell::Get()->overview_controller()->InOverviewSession());
@@ -330,6 +339,7 @@ class DesksController::DeskRemovalAnimation
 
   // DesksController::AbstractDeskSwitchAnimation:
   void OnStartingDeskScreenshotTakenInternal(const Desk* ending_desk) override {
+    DCHECK_EQ(ending_desk_, ending_desk);
     DCHECK_EQ(controller_->active_desk(), desk_to_remove_);
     // We are removing the active desk, which may have tablet split view active.
     // We will restore the split view state of the newly activated desk at the
@@ -393,6 +403,12 @@ DesksController* DesksController::Get() {
   return Shell::Get()->desks_controller();
 }
 
+const Desk* DesksController::GetTargetActiveDesk() const {
+  if (!animations_.empty())
+    return animations_.back()->ending_desk();
+  return active_desk();
+}
+
 void DesksController::Shutdown() {
   animations_.clear();
 }

ash/wm/desks/desks_controller.h

@@ -71,6 +71,10 @@ class ASH_EXPORT DesksController : public DesksHelper,
 
   const Desk* active_desk() const { return active_desk_; }
 
+  // Returns the current |active_desk()| or the soon-to-be active desk if a desk
+  // switch animation is in progress.
+  const Desk* GetTargetActiveDesk() const;
+
   // Destroys any pending animations in preparation for shutdown.
   void Shutdown();
 

ash/wm/desks/desks_test_util.cc

@@ -43,8 +43,9 @@ void DeskSwitchAnimationWaiter::OnDeskSwitchAnimationFinished() {
 void ActivateDesk(const Desk* desk) {
   ASSERT_FALSE(desk->is_active());
   DeskSwitchAnimationWaiter waiter;
-  DesksController::Get()->ActivateDesk(desk,
-                                       DesksSwitchSource::kMiniViewButton);
+  auto* desks_controller = DesksController::Get();
+  desks_controller->ActivateDesk(desk, DesksSwitchSource::kMiniViewButton);
+  EXPECT_EQ(desk, desks_controller->GetTargetActiveDesk());
   waiter.Wait();
   ASSERT_TRUE(desk->is_active());
 }

ash/wm/desks/desks_util.h

@@ -29,6 +29,11 @@ ASH_EXPORT bool IsDeskContainer(const aura::Window* container);
 
 ASH_EXPORT bool IsDeskContainerId(int id);
 
+// NOTE: The below *ActiveDesk* functions work with the currently active desk.
+// If they can be called during a desk-switch animation, you might be interested
+// in the soon-to-be active desk when the animation ends.
+// See `DesksController::GetTargetActiveDesk()`.
+
 ASH_EXPORT int GetActiveDeskContainerId();
 
 ASH_EXPORT bool IsActiveDeskContainer(const aura::Window* container);

ash/wm/mru_window_tracker.h

@@ -21,6 +21,11 @@ enum DesksMruType {
   kAllDesks,
 
   // The MRU window list will exclude windows from the inactive desks.
+  //
+  // NOTE: During an on-going desk-switch animation, getting the MRU window list
+  // for the active desk can be inconsistent, depending on at which stage of the
+  // animation it is done. If you want the MRU windows in the soon-to-be active
+  // desk, then wait for the animation to finish.
   kActiveDesk,
 };
 

ash/wm/window_cycle_controller.cc

@@ -10,6 +10,7 @@
 #include "ash/public/cpp/shell_window_ids.h"
 #include "ash/session/session_controller_impl.h"
 #include "ash/shell.h"
+#include "ash/wm/desks/desk.h"
 #include "ash/wm/desks/desks_controller.h"
 #include "ash/wm/desks/desks_util.h"
 #include "ash/wm/mru_window_tracker.h"
@@ -33,11 +34,18 @@ aura::Window* GetActiveWindow(const WindowCycleList::WindowList& window_list) {
 void ReportPossibleDesksSwitchStats(int active_desk_container_id_before_cycle) {
   // Report only for users who have 2 or more desks, since we're only interested
   // in seeing how users of Virtual Desks use window cycling.
-  if (DesksController::Get()->desks().size() < 2)
+  auto* desks_controller = DesksController::Get();
+  if (!desks_controller)
     return;
 
+  if (desks_controller->desks().size() < 2)
+    return;
+
+  // Note that this functions is called while a potential desk switch animation
+  // is starting, in this case we want the target active desk (i.e. the soon-to-
+  // be active desk after the animation finishes).
   const int active_desk_container_id_after_cycle =
-      desks_util::GetActiveDeskContainerId();
+      desks_controller->GetTargetActiveDesk()->container_id();
   DCHECK_NE(active_desk_container_id_before_cycle, kShellWindowId_Invalid);
   DCHECK_NE(active_desk_container_id_after_cycle, kShellWindowId_Invalid);
 
@@ -119,8 +127,12 @@ void WindowCycleController::Step(Direction direction) {
 void WindowCycleController::StopCycling() {
   window_cycle_list_.reset();
 
-  aura::Window* active_window_after_window_cycle = GetActiveWindow(
-      Shell::Get()->mru_window_tracker()->BuildMruWindowList(kActiveDesk));
+  // We can't use the MRU window list here to get the active window, since
+  // cycling can activate a window on a different desk, leading to a desk-switch
+  // animation launching. Getting the MRU window list for the active desk now
+  // will always be for the current active desk, not the target active desk.
+  aura::Window* active_window_after_window_cycle =
+      window_util::GetActiveWindow();
 
   // Remove our key event filter.
   event_filter_.reset();

ash/wm/window_cycle_controller_unittest.cc

@@ -30,6 +30,7 @@
 #include "ash/wm/window_state.h"
 #include "ash/wm/window_util.h"
 #include "ash/wm/wm_event.h"
+#include "base/test/metrics/histogram_tester.h"
 #include "base/test/scoped_feature_list.h"
 #include "ui/aura/client/aura_constants.h"
 #include "ui/aura/client/screen_position_client.h"
@@ -719,11 +720,15 @@ TEST_F(DesksWindowCyclingTest, CycleShowsAllDesksWindows) {
   auto win1 = CreateAppWindow(gfx::Rect(50, 50, 200, 200));
   auto* desks_controller = DesksController::Get();
   desks_controller->NewDesk(DesksCreationRemovalSource::kButton);
-  ASSERT_EQ(2u, desks_controller->desks().size());
+  desks_controller->NewDesk(DesksCreationRemovalSource::kButton);
+  ASSERT_EQ(3u, desks_controller->desks().size());
   const Desk* desk_2 = desks_controller->desks()[1].get();
   ActivateDesk(desk_2);
   EXPECT_EQ(desk_2, desks_controller->active_desk());
   auto win2 = CreateAppWindow(gfx::Rect(0, 0, 300, 200));
+  const Desk* desk_3 = desks_controller->desks()[2].get();
+  ActivateDesk(desk_3);
+  EXPECT_EQ(desk_3, desks_controller->active_desk());
   auto win3 = CreateAppWindow(gfx::Rect(10, 30, 400, 200));
 
   WindowCycleController* cycle_controller =
@@ -740,13 +745,36 @@ TEST_F(DesksWindowCyclingTest, CycleShowsAllDesksWindows) {
   // The MRU order is {win3, win2, win1, win0}. We're now at win2. Cycling one
   // more time and completing the cycle, will activate win1 which exists on a
   // desk_1. This should activate desk_1.
-  DeskSwitchAnimationWaiter waiter;
-  cycle_controller->HandleCycleWindow(WindowCycleController::FORWARD);
-  cycle_controller->CompleteCycling();
-  waiter.Wait();
-  Desk* desk_1 = desks_controller->desks()[0].get();
-  EXPECT_EQ(desk_1, desks_controller->active_desk());
-  EXPECT_EQ(win1.get(), window_util::GetActiveWindow());
+  {
+    base::HistogramTester histogram_tester;
+    DeskSwitchAnimationWaiter waiter;
+    cycle_controller->HandleCycleWindow(WindowCycleController::FORWARD);
+    cycle_controller->CompleteCycling();
+    waiter.Wait();
+    Desk* desk_1 = desks_controller->desks()[0].get();
+    EXPECT_EQ(desk_1, desks_controller->active_desk());
+    EXPECT_EQ(win1.get(), window_util::GetActiveWindow());
+    histogram_tester.ExpectUniqueSample(
+        "Ash.WindowCycleController.DesksSwitchDistance",
+        /*desk distance of 3 - 1 = */ 2, /*expected_count=*/1);
+  }
+
+  // Cycle again and activate win2, which exist on desk_2. Expect that desk to
+  // be activated, and a histogram sample of distance of 1 is recorded.
+  // MRU is {win1, win3, win2, win0}.
+  {
+    base::HistogramTester histogram_tester;
+    DeskSwitchAnimationWaiter waiter;
+    cycle_controller->HandleCycleWindow(WindowCycleController::FORWARD);
+    cycle_controller->HandleCycleWindow(WindowCycleController::FORWARD);
+    cycle_controller->CompleteCycling();
+    waiter.Wait();
+    EXPECT_EQ(desk_2, desks_controller->active_desk());
+    EXPECT_EQ(win2.get(), window_util::GetActiveWindow());
+    histogram_tester.ExpectUniqueSample(
+        "Ash.WindowCycleController.DesksSwitchDistance",
+        /*desk distance of 2 - 1 = */ 1, /*expected_count=*/1);
+  }
 }
 
 }  // namespace ash