Separate client surface reference tracking from FrameGenerator.

services/ui/ws/frame_generator.cc

Index: services/ui/ws/frame_generator.cc
diff --git a/services/ui/ws/frame_generator.cc b/services/ui/ws/frame_generator.cc
index 7229a299d751fdbbffde7e42542d11099c94012c..e8531346aa7867d90e12df22e49cab2a4446faf4 100644
--- a/services/ui/ws/frame_generator.cc
+++ b/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,10 @@ 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();
-
-  // 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_)
+    surface_tracker_.EmbedSurface(surface_id);
 }
 
 void FrameGenerator::DidReceiveCompositorFrameAck() {}
@@ -100,16 +86,33 @@ 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.
+      surface_tracker_.SendAddSurfaceReferences(compositor_frame_sink_.get());
+    }
+
     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();
+    } else {
+      // 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.
+      surface_tracker_.SendRemoveSurfaceReferences(
+          compositor_frame_sink_.get());
+    }
   }
 }
 
@@ -123,6 +126,40 @@ 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()), 1u),
+      local_frame_id_);
+
+  DCHECK_NE(old_surface_id, new_surface_id);
+
+  // Adds a reference from the top level root to the new display surface.
+  compositor_frame_sink_->AddSurfaceReferences(
+      std::vector<cc::SurfaceReference>{cc::SurfaceReference(
+          root_window_->delegate()->GetRootSurfaceId(), 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);
+  surface_tracker_.SendAddSurfaceReferences(compositor_frame_sink_.get());
+
+  // Remove the reference from the top level root to the old display surface
+  // last. 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;
@@ -231,99 +268,6 @@ void FrameGenerator::DrawWindowTree(
   }
 }
 
-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(
-      mojom::CompositorFrameSinkType::DEFAULT);
-}
-
-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 =
@@ -337,13 +281,13 @@ void FrameGenerator::OnWindowDestroying(ServerWindow* window) {
       window->compositor_frame_sink_manager()->GetLatestSurfaceId(
           mojom::CompositorFrameSinkType::DEFAULT);
   if (default_surface_id.is_valid())
-    RemoveFrameSinkReference(default_surface_id.frame_sink_id());
+    surface_tracker_.UnembedSurface(default_surface_id.frame_sink_id());
 
   cc::SurfaceId underlay_surface_id =
       window->compositor_frame_sink_manager()->GetLatestSurfaceId(
           mojom::CompositorFrameSinkType::UNDERLAY);
   if (underlay_surface_id.is_valid())
-    RemoveFrameSinkReference(underlay_surface_id.frame_sink_id());
+    surface_tracker_.UnembedSurface(underlay_surface_id.frame_sink_id());
 }
 
 }  // namespace ws