Separate client surface reference tracking from FrameGenerator.

The logic to track surface references for surface embeddings was built
as part of FrameGenerator. Generalize the logic and remove from
FrameGenerator, creating the class EmbeddedSurfaceTracker.

EmbeddedSurfaceTracker keeps track of surface references from one client
surface to many embedded surfaces. It also handles generating new
references when the client surface changes. Includes unit tests.

This simplifies FrameGenerator and should make it easier to use surface
references in other clients.

BUG=659227
CQ_INCLUDE_TRYBOTS=master.tryserver.blink:linux_trusty_blink_rel

Review-Url: https://codereview.chromium.org/2610063007
Cr-Commit-Position: refs/heads/master@{#442997}

cc/BUILD.gn

@@ -936,6 +936,7 @@ cc_test("cc_unittests") {
     "surfaces/direct_compositor_frame_sink_unittest.cc",
     "surfaces/display_scheduler_unittest.cc",
     "surfaces/display_unittest.cc",
+    "surfaces/embedded_surface_tracker_unittest.cc",
     "surfaces/surface_aggregator_unittest.cc",
     "surfaces/surface_factory_unittest.cc",
     "surfaces/surface_hittest_unittest.cc",

cc/surfaces/BUILD.gn

@@ -44,6 +44,8 @@ cc_component("surfaces") {
     "display_client.h",
     "display_scheduler.cc",
     "display_scheduler.h",
+    "embedded_surface_tracker.cc",
+    "embedded_surface_tracker.h",
     "sequence_surface_reference_factory.cc",
     "sequence_surface_reference_factory.h",
     "surface.cc",

cc/surfaces/embedded_surface_tracker.cc

+// Copyright 2016 The Chromium Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style license that can be
+// found in the LICENSE file.
+
+#include "cc/surfaces/embedded_surface_tracker.h"
+
+#include <utility>
+
+#include "base/logging.h"
+
+namespace cc {
+
+EmbeddedSurfaceTracker::EmbeddedSurfaceTracker() {}
+
+EmbeddedSurfaceTracker::~EmbeddedSurfaceTracker() {}
+
+void EmbeddedSurfaceTracker::SetCurrentSurfaceId(const SurfaceId& surface_id) {
+  DCHECK(surface_id.is_valid());
+  DCHECK_NE(current_surface_id_, surface_id);
+
+  // The last submitted CompositorFrame for |current_surface_id_| still embeds
+  // the surfaces in |references_to_remove_| and hasn't embedded the surfaces in
+  // |references_to_add_| yet.
+  references_to_remove_.clear();
+  references_to_add_.clear();
+
+  current_surface_id_ = surface_id;
+
+  // Add references from updated |current_surface_id_| to all embedded surfaces.
+  for (auto& map_entry : embedded_surfaces_)
+    AddReference(current_surface_id_, map_entry.second);
+}
+
+bool EmbeddedSurfaceTracker::EmbedSurface(const SurfaceId& embedded_id) {
+  auto iter = embedded_surfaces_.find(embedded_id.frame_sink_id());
+
+  // This is the first surface for the FrameSinkId.
+  if (iter == embedded_surfaces_.end()) {
+    if (HasValidSurfaceId())
+      AddReference(current_surface_id_, embedded_id);
+
+    // Add record that surface is embedded.
+    embedded_surfaces_[embedded_id.frame_sink_id()] = embedded_id;
+    return true;
+  }
+
+  // This is not the first surface for the FrameSinkId, we are unembedding the
+  // old surface and embedding the new surface.
+  SurfaceId& stored_surface_id = iter->second;
+  DCHECK_NE(stored_surface_id, embedded_id);
+
+  if (HasValidSurfaceId()) {
+    RemoveReference(current_surface_id_, stored_surface_id);
+    AddReference(current_surface_id_, embedded_id);
+  }
+
+  stored_surface_id = embedded_id;
+  return false;
+}
+
+void EmbeddedSurfaceTracker::UnembedSurface(const FrameSinkId& frame_sink_id) {
+  auto iter = embedded_surfaces_.find(frame_sink_id);
+  if (iter == embedded_surfaces_.end())
+    return;
+
+  if (HasValidSurfaceId())
+    RemoveReference(current_surface_id_, iter->second);
+
+  // Remove record that surface is embedded.
+  embedded_surfaces_.erase(iter);
+}
+
+std::vector<SurfaceReference> EmbeddedSurfaceTracker::GetReferencesToAdd() {
+  std::vector<SurfaceReference> references;
+  std::swap(references, references_to_add_);
+  return references;
+}
+
+std::vector<SurfaceReference> EmbeddedSurfaceTracker::GetReferencesToRemove() {
+  std::vector<SurfaceReference> references;
+  std::swap(references, references_to_remove_);
+  return references;
+}
+
+void EmbeddedSurfaceTracker::AddReference(const SurfaceId& parent_id,
+                                          const SurfaceId& child_id) {
+  references_to_add_.push_back(SurfaceReference(parent_id, child_id));
+}
+
+void EmbeddedSurfaceTracker::RemoveReference(const SurfaceId& parent_id,
+                                             const SurfaceId& child_id) {
+  references_to_remove_.push_back(SurfaceReference(parent_id, child_id));
+}
+
+}  // namespace cc

cc/surfaces/embedded_surface_tracker.h

+// Copyright 2016 The Chromium Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style license that can be
+// found in the LICENSE file.
+
+#ifndef CC_SURFACES_EMBEDDED_SURFACE_TRACKER_H_
+#define CC_SURFACES_EMBEDDED_SURFACE_TRACKER_H_
+
+#include <unordered_map>
+#include <vector>
+
+#include "base/macros.h"
+#include "cc/surfaces/frame_sink_id.h"
+#include "cc/surfaces/surface_id.h"
+#include "cc/surfaces/surface_reference.h"
+#include "cc/surfaces/surfaces_export.h"
+
+namespace cc {
+
+namespace test {
+class EmbeddedSurfaceTrackerTest;
+}
+
+// Tracks surface references from a client surface to embedded surfaces. Handles
+// updating the references when the client surface resizes.
+class CC_SURFACES_EXPORT EmbeddedSurfaceTracker {
+ public:
+  EmbeddedSurfaceTracker();
+  ~EmbeddedSurfaceTracker();
+
+  const SurfaceId& current_surface_id() const { return current_surface_id_; }
+
+  // Checks if the client has a valid SurfaceId. This will become true after
+  // SetCurrentSurfaceId() is called for the first time. No references will be
+  // added or removed until this is true.
+  bool HasValidSurfaceId() const { return current_surface_id_.is_valid(); }
+
+  // Sets |current_surface_id_| when the client switches surfaces. Clears all
+  // references waiting to be added or removed, since no CompositorFrame will be
+  // submitted for the previous surface now. Generates new references from the
+  // new surface to all embedded surfaces. Does not send IPC.
+  void SetCurrentSurfaceId(const SurfaceId& surface_id);
+
+  // Adds a reference from |current_surface_id_| to |embedded_id| and records
+  // |embedded_id|. If an existing surface with the same FrameSinkId is embedded
+  // removes the reference to it. Does not send IPC. Returns true for the first
+  // surface from a FrameSinkId and false for all subsequent surfaces.
+  bool EmbedSurface(const SurfaceId& embedded_id);
+
+  // Removes reference from |current_surface_id_| to latest embedded surface for
+  // |frame_sink_id|. Does not send IPC.
+  void UnembedSurface(const FrameSinkId& frame_sink_id);
+
+  // Returns |references_to_add_| and clears contents. It's expected this will
+  // be called immediately before submitting a CompositorFrame and caller will
+  // add references via IPC.
+  std::vector<SurfaceReference> GetReferencesToAdd();
+  bool HasReferencesToAdd() const { return !references_to_add_.empty(); }
+
+  // Returns |references_to_remove_| and clears contents. It's expected this
+  // will be called immediately after submitting a CompositorFrame and caller
+  // will remove references via IPC.
+  std::vector<SurfaceReference> GetReferencesToRemove();
+  bool HasReferencesToRemove() const { return !references_to_remove_.empty(); }
+
+ private:
+  friend class test::EmbeddedSurfaceTrackerTest;
+
+  // Records adding a reference from |parent_id| to |child_id|. References are
+  // not actually added until SendAddSurfaceReferences() is called.
+  void AddReference(const SurfaceId& parent_id, const SurfaceId& child_id);
+
+  // Records removing a reference from |parent_id| to |child_id|. References
+  // are not actually removed until SendRemoveSurfaceReferences() is called.
+  void RemoveReference(const SurfaceId& parent_id, const SurfaceId& child_id);
+
+  // The id of the client surface that is embedding other surfaces.
+  SurfaceId current_surface_id_;
+
+  // Surfaces that are embedded by the client surface. A reference from
+  // |current_surface_id_| will be added if it's valid.
+  std::unordered_map<FrameSinkId, SurfaceId, FrameSinkIdHash>
+      embedded_surfaces_;
+
+  // References to surfaces that should be removed after the next
+  // CompositorFrame has been submitted and the surfaces are no longer embedded.
+  std::vector<SurfaceReference> references_to_remove_;
+
+  // References that should be added before the next CompositorFrame is
+  // submitted.
+  std::vector<SurfaceReference> references_to_add_;
+
+  DISALLOW_COPY_AND_ASSIGN(EmbeddedSurfaceTracker);
+};
+
+}  // namespace cc
+
+#endif  // CC_SURFACES_EMBEDDED_SURFACE_TRACKER_H_

cc/surfaces/embedded_surface_tracker_unittest.cc

+// Copyright 2017 The Chromium Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style license that can be
+// found in the LICENSE file.
+
+#include "cc/surfaces/embedded_surface_tracker.h"
+
+#include <memory>
+
+#include "base/logging.h"
+#include "base/macros.h"
+#include "base/memory/ptr_util.h"
+#include "cc/surfaces/frame_sink_id.h"
+#include "cc/surfaces/surface_id.h"
+#include "cc/surfaces/surface_reference.h"
+#include "testing/gmock/include/gmock/gmock.h"
+#include "testing/gtest/include/gtest/gtest.h"
+
+using testing::ElementsAre;
+using testing::IsEmpty;
+using testing::SizeIs;
+
+namespace cc {
+namespace test {
+namespace {
+
+constexpr FrameSinkId kParentFrameSink(2, 1);
+constexpr FrameSinkId kEmbeddedFrameSink(65563, 1);
+
+SurfaceId MakeSurfaceId(const FrameSinkId& frame_sink_id, uint32_t local_id) {
+  return SurfaceId(
+      frame_sink_id,
+      LocalFrameId(local_id, base::UnguessableToken::Deserialize(0, 1u)));
+}
+
+}  // namespace
+
+class EmbeddedSurfaceTrackerTest : public testing::Test {
+ public:
+  EmbeddedSurfaceTrackerTest() {}
+  ~EmbeddedSurfaceTrackerTest() override {}
+
+  EmbeddedSurfaceTracker& tracker() { return *tracker_; }
+
+  // Returns references to add without clearing them.
+  const std::vector<SurfaceReference>& peek_to_add() {
+    return tracker_->references_to_add_;
+  }
+
+  // Returns references to remove without clearing them.
+  const std::vector<SurfaceReference>& peek_to_remove() {
+    return tracker_->references_to_remove_;
+  }
+
+  // testing::Test:
+  void SetUp() override {
+    testing::Test::SetUp();
+    tracker_ = base::MakeUnique<EmbeddedSurfaceTracker>();
+  }
+
+  void TearDown() override { tracker_.reset(); }
+
+ private:
+  std::unique_ptr<EmbeddedSurfaceTracker> tracker_;
+
+  DISALLOW_COPY_AND_ASSIGN(EmbeddedSurfaceTrackerTest);
+};
+
+TEST_F(EmbeddedSurfaceTrackerTest, SetCurrentSurfaceId) {
+  const SurfaceId parent_id = MakeSurfaceId(kParentFrameSink, 1);
+
+  // Initially HasValidSurfaceId() should be false.
+  EXPECT_FALSE(tracker().HasValidSurfaceId());
+
+  // After setting current SurfaceId then HasValidSurfaceId() should be true.
+  tracker().SetCurrentSurfaceId(parent_id);
+  EXPECT_TRUE(tracker().HasValidSurfaceId());
+  EXPECT_EQ(parent_id, tracker().current_surface_id());
+}
+
+TEST_F(EmbeddedSurfaceTrackerTest, EmbedSurface) {
+  const SurfaceId parent_id = MakeSurfaceId(kParentFrameSink, 1);
+  const SurfaceId embedded_id = MakeSurfaceId(kEmbeddedFrameSink, 1);
+  const SurfaceReference reference(parent_id, embedded_id);
+
+  // Set parent surface id then embed another surface. The tracker should have a
+  // reference to add from parent surface to embedded surface.
+  tracker().SetCurrentSurfaceId(parent_id);
+  EXPECT_TRUE(tracker().EmbedSurface(embedded_id));
+  EXPECT_THAT(peek_to_add(), ElementsAre(reference));
+
+  // Getting the references to add should return the correct reference and empty
+  // the reference the add.
+  EXPECT_THAT(tracker().GetReferencesToAdd(), ElementsAre(reference));
+  EXPECT_THAT(peek_to_add(), IsEmpty());
+}
+
+TEST_F(EmbeddedSurfaceTrackerTest, EmbedNewSurfaceForFrameSink) {
+  const SurfaceId parent_id = MakeSurfaceId(kParentFrameSink, 1);
+  const SurfaceId embedded_id_first = MakeSurfaceId(kEmbeddedFrameSink, 1);
+  const SurfaceId embedded_id_second = MakeSurfaceId(kEmbeddedFrameSink, 2);
+  const SurfaceReference reference_first(parent_id, embedded_id_first);
+  const SurfaceReference reference_second(parent_id, embedded_id_second);
+
+  // Embed a surface and add references.
+  tracker().SetCurrentSurfaceId(parent_id);
+  EXPECT_TRUE(tracker().EmbedSurface(embedded_id_first));
+  EXPECT_THAT(tracker().GetReferencesToAdd(), ElementsAre(reference_first));
+
+  // Embed a newer surface with the same FrameSinkId as the first embedded
+  // surface. The first reference should be in the list to remove and the new
+  // reference in the last to add.
+  EXPECT_FALSE(tracker().EmbedSurface(embedded_id_second));
+  EXPECT_THAT(peek_to_remove(), ElementsAre(reference_first));
+  EXPECT_THAT(peek_to_add(), ElementsAre(reference_second));
+}
+
+TEST_F(EmbeddedSurfaceTrackerTest, UnembedSurface) {
+  const SurfaceId parent_id = MakeSurfaceId(kParentFrameSink, 1);
+  const SurfaceId embedded_id = MakeSurfaceId(kEmbeddedFrameSink, 1);
+  const SurfaceReference reference(parent_id, embedded_id);
+
+  tracker().SetCurrentSurfaceId(parent_id);
+  tracker().EmbedSurface(embedded_id);
+  EXPECT_THAT(tracker().GetReferencesToAdd(), ElementsAre(reference));
+  EXPECT_THAT(peek_to_remove(), IsEmpty());
+
+  // Unembed the surface. The reference should be in the list to remove.
+  tracker().UnembedSurface(embedded_id.frame_sink_id());
+  EXPECT_THAT(peek_to_remove(), ElementsAre(reference));
+
+  // Getting the references to remove should return the correct reference and
+  // empty the reference the remove.
+  EXPECT_THAT(tracker().GetReferencesToRemove(), ElementsAre(reference));
+  EXPECT_THAT(peek_to_remove(), IsEmpty());
+}
+
+TEST_F(EmbeddedSurfaceTrackerTest, EmbedSurfaceBeforeSetCurrentSurfaceId) {
+  const SurfaceId parent_id = MakeSurfaceId(kParentFrameSink, 1);
+  const SurfaceId embedded_id = MakeSurfaceId(kEmbeddedFrameSink, 1);
+  const SurfaceReference reference(parent_id, embedded_id);
+
+  // Embed a surface before the parent id is set. There should be no references
+  // to add because we don't know the parent id yet.
+  tracker().EmbedSurface(embedded_id);
+  EXPECT_THAT(peek_to_add(), IsEmpty());
+
+  // Set the parent id, this should add a reference to add to the embedded
+  // surface.
+  tracker().SetCurrentSurfaceId(parent_id);
+  EXPECT_THAT(peek_to_add(), ElementsAre(reference));
+}
+
+TEST_F(EmbeddedSurfaceTrackerTest, UpdateCurrentSurfaceId) {
+  const SurfaceId parent_id_first = MakeSurfaceId(kParentFrameSink, 1);
+  const SurfaceId parent_id_second = MakeSurfaceId(kParentFrameSink, 2);
+  const SurfaceId embedded_id = MakeSurfaceId(kEmbeddedFrameSink, 1);
+  const SurfaceReference reference(parent_id_second, embedded_id);
+
+  // Set parent id and embed a surface.
+  tracker().SetCurrentSurfaceId(parent_id_first);
+  tracker().EmbedSurface(embedded_id);
+  EXPECT_THAT(tracker().GetReferencesToAdd(), SizeIs(1));
+
+  // Update the current parent id. There should be a reference to add from the
+  // new parent id to the embedded surface.
+  tracker().SetCurrentSurfaceId(parent_id_second);
+  EXPECT_THAT(peek_to_add(), ElementsAre(reference));
+}
+
+TEST_F(EmbeddedSurfaceTrackerTest, UpdateCurrentSurfaceIdBeforeAdding) {
+  const SurfaceId parent_id_first = MakeSurfaceId(kParentFrameSink, 1);
+  const SurfaceId parent_id_second = MakeSurfaceId(kParentFrameSink, 2);
+  const SurfaceId embedded_id = MakeSurfaceId(kEmbeddedFrameSink, 1);
+  const SurfaceReference reference_first(parent_id_first, embedded_id);
+  const SurfaceReference reference_second(parent_id_second, embedded_id);
+
+  // Set parent id and embed a surface. There should be a reference to add from
+  // the parent to embedded surface.
+  tracker().SetCurrentSurfaceId(parent_id_first);
+  tracker().EmbedSurface(embedded_id);
+  EXPECT_THAT(peek_to_add(), ElementsAre(reference_first));
+
+  // Update the parent id before sending IPC to add references. There should be
+  // a reference to add from the new parent id to the embedded surface and the
+  // previous reference to add was deleted.
+  tracker().SetCurrentSurfaceId(parent_id_second);
+  EXPECT_THAT(peek_to_add(), ElementsAre(reference_second));
+  EXPECT_THAT(tracker().GetReferencesToAdd(), ElementsAre(reference_second));
+}
+
+}  // namespace test
+}  // namespace cc

services/ui/ws/frame_generator.cc

@@ -5,6 +5,7 @@
 #include "services/ui/ws/frame_generator.h"
 
 #include <utility>
+#include <vector>
 
 #include "base/containers/adapters.h"
 #include "cc/output/compositor_frame.h"
@@ -31,7 +32,16 @@ FrameGenerator::FrameGenerator(FrameGeneratorDelegate* delegate,
 }
 
 FrameGenerator::~FrameGenerator() {
-  RemoveAllSurfaceReferences();
+  // Remove reference from top level root to the display root surface, if one
+  // exists. This will make everything referenced by the display surface
+  // unreachable so it can be garbage collected.
+  if (surface_tracker_.HasValidSurfaceId()) {
+    compositor_frame_sink_->RemoveSurfaceReferences(
+        std::vector<cc::SurfaceReference>{
+            cc::SurfaceReference(root_window_->delegate()->GetRootSurfaceId(),
+                                 surface_tracker_.current_surface_id())});
+  }
+
   // Invalidate WeakPtrs now to avoid callbacks back into the
   // FrameGenerator during destruction of |compositor_frame_sink_|.
   weak_factory_.InvalidateWeakPtrs();
@@ -57,34 +67,16 @@ void FrameGenerator::OnSurfaceCreated(const cc::SurfaceId& surface_id,
                                       ServerWindow* window) {
   DCHECK(surface_id.is_valid());
 
-  auto iter = active_references_.find(surface_id.frame_sink_id());
-  if (iter == active_references_.end()) {
-    AddFirstReference(surface_id, window);
-    return;
-  }
-
-  cc::SurfaceReference& ref = iter->second;
-
-  // This shouldn't be called multiple times for the same SurfaceId.
-  DCHECK_EQ(surface_id.frame_sink_id(), ref.child_id().frame_sink_id());
-  DCHECK_NE(surface_id.local_frame_id(), ref.child_id().local_frame_id());
-
-  // The current reference will be removed after the next CompositorFrame is
-  // submitted or FrameGenerator is destroyed.
-  references_to_remove_.push_back(ref);
-  cc::SurfaceId old_surface_id = ref.child_id();
+  // TODO(samans): Clients are actually embedded in the WM and only the WM is
+  // embedded here. This needs to be fixed.
 
-  // New surface reference is recorded and will be added at end of this method.
-  ref = cc::SurfaceReference(ref.parent_id(), surface_id);
-  references_to_add_.push_back(ref);
-
-  // If the display root surface has changed, add references from the new
-  // SurfaceId to all embedded surfaces. For example, this would happen when the
-  // display resolution or device scale factor changes.
-  if (window == root_window_)
-    AddNewParentReferences(old_surface_id, surface_id);
-
-  PerformAddSurfaceReferences();
+  // Only handle embedded surfaces changing here. The display root surface
+  // changing is handled immediately after the CompositorFrame is submitted.
+  if (window != root_window_) {
+    // Add observer for window the first time it's seen.
+    if (surface_tracker_.EmbedSurface(surface_id))
+      Add(window);
+  }
 }
 
 void FrameGenerator::DidReceiveCompositorFrameAck() {}
@@ -100,16 +92,37 @@ void FrameGenerator::OnBeginFrame(const cc::BeginFrameArgs& begin_frame_arags) {
     gfx::Size frame_size = last_submitted_frame_size_;
     if (!frame.render_pass_list.empty())
       frame_size = frame.render_pass_list[0]->output_rect.size();
-    if (!local_frame_id_.is_valid() || frame_size != last_submitted_frame_size_)
+
+    bool display_surface_changed = false;
+    if (!local_frame_id_.is_valid() ||
+        frame_size != last_submitted_frame_size_) {
       local_frame_id_ = id_allocator_.GenerateId();
+      display_surface_changed = true;
+    } else {
+      // If the display surface is changing then we shouldn't add references
+      // from the old display surface. We want to add references from the new
+      // display surface, this happens after we submit the first CompositorFrame
+      // so the new display surface exists.
+      if (surface_tracker_.HasReferencesToAdd()) {
+        compositor_frame_sink_->AddSurfaceReferences(
+            surface_tracker_.GetReferencesToAdd());
+      }
+    }
+
     compositor_frame_sink_->SubmitCompositorFrame(local_frame_id_,
                                                   std::move(frame));
     last_submitted_frame_size_ = frame_size;
 
-    // Remove dead references after we submit a frame. This has to happen after
-    // the frame is submitted otherwise we could end up deleting a surface that
-    // is still embedded in the last frame.
-    PerformRemoveSurfaceReferences();
+    if (display_surface_changed)
+      UpdateDisplaySurfaceId();
+
+    // Remove references to surfaces that are no longer embedded. This has to
+    // happen after the frame is submitted otherwise we could end up deleting
+    // a surface that is still embedded in the last submitted frame.
+    if (surface_tracker_.HasReferencesToRemove()) {
+      compositor_frame_sink_->RemoveSurfaceReferences(
+          surface_tracker_.GetReferencesToRemove());
+    }
   }
 }
 
@@ -123,6 +136,43 @@ void FrameGenerator::WillDrawSurface() {
   // TODO(fsamuel, staraz): Implement this.
 }
 
+void FrameGenerator::UpdateDisplaySurfaceId() {
+  // FrameGenerator owns the display root surface and is a bit of a special
+  // case. There is no surface that embeds the display surface, so nothing
+  // would reference it. Instead, a reference from the top level root is added
+  // to mark the display surface as visible. As a result, FrameGenerator is
+  // responsible for maintaining a reference from the top level root to the
+  // display surface, in addition to references from the display surface to
+  // embedded surfaces.
+  const cc::SurfaceId old_surface_id = surface_tracker_.current_surface_id();
+  const cc::SurfaceId new_surface_id(
+      cc::FrameSinkId(WindowIdToTransportId(root_window_->id()), 0),
+      local_frame_id_);
+
+  DCHECK_NE(old_surface_id, new_surface_id);
+
+  // Set new SurfaceId for the display surface. This will add references from
+  // the new display surface to all embedded surfaces.
+  surface_tracker_.SetCurrentSurfaceId(new_surface_id);
+  std::vector<cc::SurfaceReference> references_to_add =
+      surface_tracker_.GetReferencesToAdd();
+
+  // Adds a reference from the top level root to the new display surface.
+  references_to_add.push_back(cc::SurfaceReference(
+      root_window_->delegate()->GetRootSurfaceId(), new_surface_id));
+
+  compositor_frame_sink_->AddSurfaceReferences(references_to_add);
+
+  // Remove the reference from the top level root to the old display surface
+  // after we have added references from the new display surface. Not applicable
+  // for the first display surface.
+  if (old_surface_id.is_valid()) {
+    compositor_frame_sink_->RemoveSurfaceReferences(
+        std::vector<cc::SurfaceReference>{cc::SurfaceReference(
+            root_window_->delegate()->GetRootSurfaceId(), old_surface_id)});
+  }
+}
+
 cc::CompositorFrame FrameGenerator::GenerateCompositorFrame(
     const gfx::Rect& output_rect) {
   const int render_pass_id = 1;
@@ -189,98 +239,6 @@ void FrameGenerator::DrawWindow(cc::RenderPass* pass, ServerWindow* window) {
                default_surface_id);
 }
 
-cc::SurfaceId FrameGenerator::FindParentSurfaceId(ServerWindow* window) {
-  if (window == root_window_)
-    return root_window_->delegate()->GetRootSurfaceId();
-
-  // The root window holds the parent SurfaceId. This SurfaceId will have an
-  // invalid LocalFrameId before FrameGenerator has submitted a CompositorFrame.
-  // After the first frame is submitted it will always be a valid SurfaceId.
-  return root_window_->compositor_frame_sink_manager()->GetLatestSurfaceId();
-}
-
-void FrameGenerator::AddSurfaceReference(const cc::SurfaceId& parent_id,
-                                         const cc::SurfaceId& child_id) {
-  DCHECK_NE(parent_id, child_id);
-
-  // Add new reference from parent to surface and record reference.
-  cc::SurfaceReference ref(parent_id, child_id);
-  active_references_[child_id.frame_sink_id()] = ref;
-  references_to_add_.push_back(ref);
-}
-
-void FrameGenerator::AddFirstReference(const cc::SurfaceId& surface_id,
-                                       ServerWindow* window) {
-  cc::SurfaceId parent_id = FindParentSurfaceId(window);
-
-  if (parent_id.local_frame_id().is_valid()) {
-    AddSurfaceReference(parent_id, surface_id);
-
-    // For the first display root surface, add references to any child surfaces
-    // that were created before it, since no reference has been added yet.
-    if (window == root_window_) {
-      for (auto& child_surface_id : waiting_for_references_)
-        AddSurfaceReference(surface_id, child_surface_id);
-      waiting_for_references_.clear();
-    }
-
-    PerformAddSurfaceReferences();
-  } else {
-    // This isn't the display root surface and display root surface hasn't
-    // submitted a CF yet. We can't add a reference to an unknown SurfaceId.
-    waiting_for_references_.push_back(surface_id);
-  }
-
-  // Observe |window| so that we can remove references when it's destroyed.
-  Add(window);
-}
-
-void FrameGenerator::AddNewParentReferences(
-    const cc::SurfaceId& old_surface_id,
-    const cc::SurfaceId& new_surface_id) {
-  DCHECK_EQ(old_surface_id.frame_sink_id(), new_surface_id.frame_sink_id());
-
-  for (auto& map_entry : active_references_) {
-    cc::SurfaceReference& ref = map_entry.second;
-    if (ref.parent_id() == old_surface_id) {
-      ref = cc::SurfaceReference(new_surface_id, ref.child_id());
-      references_to_add_.push_back(ref);
-    }
-  }
-}
-
-void FrameGenerator::PerformAddSurfaceReferences() {
-  if (references_to_add_.empty())
-    return;
-
-  compositor_frame_sink_->AddSurfaceReferences(references_to_add_);
-  references_to_add_.clear();
-}
-
-void FrameGenerator::PerformRemoveSurfaceReferences() {
-  if (references_to_remove_.empty())
-    return;
-
-  compositor_frame_sink_->RemoveSurfaceReferences(references_to_remove_);
-  references_to_remove_.clear();
-}
-
-void FrameGenerator::RemoveFrameSinkReference(
-    const cc::FrameSinkId& frame_sink_id) {
-  auto it = active_references_.find(frame_sink_id);
-  if (it == active_references_.end())
-    return;
-  references_to_remove_.push_back(it->second);
-  active_references_.erase(it);
-}
-
-void FrameGenerator::RemoveAllSurfaceReferences() {
-  for (auto& map_entry : active_references_)
-    references_to_remove_.push_back(map_entry.second);
-  active_references_.clear();
-  PerformRemoveSurfaceReferences();
-}
-
 void FrameGenerator::OnWindowDestroying(ServerWindow* window) {
   Remove(window);
   ServerWindowCompositorFrameSinkManager* compositor_frame_sink_manager =
@@ -293,7 +251,7 @@ void FrameGenerator::OnWindowDestroying(ServerWindow* window) {
   cc::SurfaceId surface_id =
       window->compositor_frame_sink_manager()->GetLatestSurfaceId();
   if (surface_id.is_valid())
-    RemoveFrameSinkReference(surface_id.frame_sink_id());
+    surface_tracker_.UnembedSurface(surface_id.frame_sink_id());
 }
 
 }  // namespace ws

services/ui/ws/frame_generator.h

@@ -6,12 +6,11 @@
 #define SERVICES_UI_WS_FRAME_GENERATOR_H_
 
 #include <memory>
-#include <unordered_map>
-#include <vector>
 
 #include "base/macros.h"
 #include "base/timer/timer.h"
 #include "cc/ipc/display_compositor.mojom.h"
+#include "cc/surfaces/embedded_surface_tracker.h"
 #include "cc/surfaces/frame_sink_id.h"
 #include "cc/surfaces/surface_id.h"
 #include "cc/surfaces/surface_id_allocator.h"
@@ -77,6 +76,9 @@ class FrameGenerator : public ServerWindowTracker,
   void ReclaimResources(const cc::ReturnedResourceArray& resources) override;
   void WillDrawSurface() override;
 
+  // Updates the display surface SurfaceId using new value in |local_frame_id_|.
+  void UpdateDisplaySurfaceId();
+
   // Generates the CompositorFrame.
   cc::CompositorFrame GenerateCompositorFrame(const gfx::Rect& output_rect);
 
@@ -84,36 +86,6 @@ class FrameGenerator : public ServerWindowTracker,
   // appends it to the provided cc::RenderPass.
   void DrawWindow(cc::RenderPass* pass, ServerWindow* window);
 
-  // Finds the parent surface id for |window|.
-  cc::SurfaceId FindParentSurfaceId(ServerWindow* window);
-
-  // Adds surface reference to local cache and surface manager.
-  void AddSurfaceReference(const cc::SurfaceId& parent_id,
-                           const cc::SurfaceId& child_id);
-
-  // Does work necessary for adding the first surface reference.
-  void AddFirstReference(const cc::SurfaceId& surface_id, ServerWindow* window);
-
-  // Finds all Surfaces with references from |old_surface_id| and adds a new
-  // reference from |new_surface_id|. The caller should remove any references
-  // to |old_surface_id| afterwards to finish cleanup.
-  void AddNewParentReferences(const cc::SurfaceId& old_surface_id,
-                              const cc::SurfaceId& new_surface_id);
-
-  // Sends IPC to add references in |references_to_add_|.
-  void PerformAddSurfaceReferences();
-
-  // Sends IPC to remove all references in |references_to_remove_|.
-  void PerformRemoveSurfaceReferences();
-
-  // Removes any retained references for |frame_sink_id_|.
-  void RemoveFrameSinkReference(const cc::FrameSinkId& frame_sink_id);
-
-  // Removes all retained references to surfaces.
-  void RemoveAllSurfaceReferences();
-
-  cc::mojom::DisplayCompositor* GetDisplayCompositor();
-
   // ServerWindowObserver implementation.
   void OnWindowDestroying(ServerWindow* window) override;
 
@@ -127,24 +99,8 @@ class FrameGenerator : public ServerWindowTracker,
   cc::mojom::MojoCompositorFrameSinkPtr compositor_frame_sink_;
   cc::mojom::DisplayPrivatePtr display_private_;
 
-  // Active references held by this client to surfaces that could be embedded in
-  // a CompositorFrame submitted from FrameGenerator.
-  std::unordered_map<cc::FrameSinkId, cc::SurfaceReference, cc::FrameSinkIdHash>
-      active_references_;
-
-  // References to surfaces that should be removed after a CompositorFrame has
-  // been submitted and the surfaces are not being used.
-  std::vector<cc::SurfaceReference> references_to_remove_;
-
-  // References that should be added before the next CompositorFrame is
-  // submitted.
-  std::vector<cc::SurfaceReference> references_to_add_;
-
-  // If a CompositorFrame for a child surface is submitted before the first
-  // display root CompositorFrame, we can't add a reference from the unknown
-  // display root SurfaceId. Track the child SurfaceId here and add a reference
-  // to it when the display root SurfaceId is available.
-  std::vector<cc::SurfaceId> waiting_for_references_;
+  // Tracks surface references for embedded surfaces.
+  cc::EmbeddedSurfaceTracker surface_tracker_;
 
   mojo::Binding<cc::mojom::MojoCompositorFrameSinkClient> binding_;