Add a mechanism to retrieve the PageNode associated with a WC

Change-Id: I811bb177eab0c3f274d7ee62c6967d3ec6d5c520
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/1879710Reviewed-by: Sigurður Ásgeirsson <siggi@chromium.org>
Commit-Queue: Sébastien Marchand <sebmarchand@chromium.org>
Cr-Commit-Position: refs/heads/master@{#713444}

components/performance_manager/BUILD.gn

@@ -103,6 +103,7 @@ source_set("unit_tests") {
     "performance_manager_tab_helper_unittest.cc",
     "performance_manager_test_harness.cc",
     "performance_manager_test_harness.h",
+    "performance_manager_unittest.cc",
     "shared_worker_watcher_unittest.cc",
     "web_contents_proxy_unittest.cc",
   ]

components/performance_manager/graph/page_node_impl.cc

@@ -271,6 +271,12 @@ void PageNodeImpl::set_has_nonempty_beforeunload(
 
 void PageNodeImpl::JoinGraph() {
   DCHECK_CALLED_ON_VALID_SEQUENCE(sequence_checker_);
+#if DCHECK_IS_ON()
+  // Dereferencing the WeakPtr associated with this node will bind it to the
+  // current sequence (all subsequent calls to |GetWeakPtr| will return the
+  // same WeakPtr).
+  GetWeakPtr()->GetImpl();
+#endif
 
   NodeBase::JoinGraph();
 }

components/performance_manager/graph/page_node_impl.h

@@ -11,6 +11,7 @@
 
 #include "base/containers/flat_set.h"
 #include "base/macros.h"
+#include "base/memory/weak_ptr.h"
 #include "base/time/time.h"
 #include "components/performance_manager/graph/node_attached_data.h"
 #include "components/performance_manager/graph/node_base.h"
@@ -106,6 +107,10 @@ class PageNodeImpl
     SetPageAlmostIdle(page_almost_idle);
   }
 
+  base::WeakPtr<PageNodeImpl> GetWeakPtr() {
+    return weak_factory_.GetWeakPtr();
+  }
+
  private:
   friend class FrameNodeImpl;
   friend class PageAggregatorAccess;
@@ -251,6 +256,8 @@ class PageNodeImpl
   // Inline storage for PageAggregatorAccess user data.
   InternalNodeAttachedDataStorage<sizeof(uintptr_t) + 24> page_aggregator_data_;
 
+  base::WeakPtrFactory<PageNodeImpl> weak_factory_{this};
+
   DISALLOW_COPY_AND_ASSIGN(PageNodeImpl);
 };
 

components/performance_manager/performance_manager.cc

@@ -4,7 +4,9 @@
 
 #include "components/performance_manager/public/performance_manager.h"
 
+#include "components/performance_manager/graph/page_node_impl.h"
 #include "components/performance_manager/performance_manager_impl.h"
+#include "components/performance_manager/performance_manager_tab_helper.h"
 
 namespace performance_manager {
 
@@ -45,4 +47,16 @@ void PerformanceManager::PassToGraph(const base::Location& from_here,
           std::move(graph_owned)));
 }
 
+// static
+base::WeakPtr<PageNode> PerformanceManager::GetPageNodeForWebContents(
+    content::WebContents* wc) {
+  DCHECK(wc);
+  PerformanceManagerTabHelper* helper =
+      PerformanceManagerTabHelper::FromWebContents(wc);
+  if (!helper)
+    return nullptr;
+
+  return helper->page_node()->GetWeakPtr();
+}
+
 }  // namespace performance_manager

components/performance_manager/performance_manager_impl_unittest.cc

@@ -19,11 +19,11 @@
 
 namespace performance_manager {
 
-class PerformanceManagerTest : public testing::Test {
+class PerformanceManagerImplTest : public testing::Test {
  public:
-  PerformanceManagerTest() {}
+  PerformanceManagerImplTest() {}
 
-  ~PerformanceManagerTest() override {}
+  ~PerformanceManagerImplTest() override {}
 
   void SetUp() override {
     EXPECT_EQ(nullptr, PerformanceManagerImpl::GetInstance());
@@ -52,10 +52,10 @@ class PerformanceManagerTest : public testing::Test {
   std::unique_ptr<PerformanceManagerImpl> performance_manager_;
   base::test::TaskEnvironment task_environment_;
 
-  DISALLOW_COPY_AND_ASSIGN(PerformanceManagerTest);
+  DISALLOW_COPY_AND_ASSIGN(PerformanceManagerImplTest);
 };
 
-TEST_F(PerformanceManagerTest, InstantiateNodes) {
+TEST_F(PerformanceManagerImplTest, InstantiateNodes) {
   int next_render_frame_id = 0;
 
   std::unique_ptr<ProcessNodeImpl> process_node =
@@ -78,7 +78,7 @@ TEST_F(PerformanceManagerTest, InstantiateNodes) {
   performance_manager()->DeleteNode(std::move(process_node));
 }
 
-TEST_F(PerformanceManagerTest, BatchDeleteNodes) {
+TEST_F(PerformanceManagerImplTest, BatchDeleteNodes) {
   int next_render_frame_id = 0;
   // Create a page node and a small hierarchy of frames.
   std::unique_ptr<ProcessNodeImpl> process_node =
@@ -125,7 +125,7 @@ TEST_F(PerformanceManagerTest, BatchDeleteNodes) {
   performance_manager()->BatchDeleteNodes(std::move(nodes));
 }
 
-TEST_F(PerformanceManagerTest, CallOnGraphImpl) {
+TEST_F(PerformanceManagerImplTest, CallOnGraphImpl) {
   // Create a page node for something to target.
   std::unique_ptr<PageNodeImpl> page_node =
       performance_manager()->CreatePageNode(WebContentsProxy(), std::string(),
@@ -147,7 +147,7 @@ TEST_F(PerformanceManagerTest, CallOnGraphImpl) {
   performance_manager()->DeleteNode(std::move(page_node));
 }
 
-TEST_F(PerformanceManagerTest, CallOnGraphAndReplyWithResult) {
+TEST_F(PerformanceManagerImplTest, CallOnGraphAndReplyWithResult) {
   // Create a page node for something to target.
   std::unique_ptr<PageNodeImpl> page_node =
       performance_manager()->CreatePageNode(WebContentsProxy(), std::string(),

components/performance_manager/performance_manager_unittest.cc

+// Copyright 2019 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 "components/performance_manager/public/performance_manager.h"
+
+#include "base/run_loop.h"
+#include "base/task/post_task.h"
+#include "base/test/bind_test_util.h"
+#include "components/performance_manager/performance_manager_test_harness.h"
+#include "components/performance_manager/public/graph/page_node.h"
+#include "components/performance_manager/public/web_contents_proxy.h"
+#include "content/public/browser/browser_task_traits.h"
+#include "testing/gtest/include/gtest/gtest.h"
+
+namespace performance_manager {
+
+class PerformanceManagerTest : public PerformanceManagerTestHarness {
+ public:
+  using Super = PerformanceManagerTestHarness;
+
+  PerformanceManagerTest() {}
+
+  void SetUp() override {
+    EXPECT_FALSE(PerformanceManager::IsAvailable());
+    Super::SetUp();
+    EXPECT_TRUE(PerformanceManager::IsAvailable());
+  }
+
+  void TearDown() override {
+    EXPECT_TRUE(PerformanceManager::IsAvailable());
+    Super::TearDown();
+    EXPECT_FALSE(PerformanceManager::IsAvailable());
+  }
+
+  ~PerformanceManagerTest() override {}
+
+  DISALLOW_COPY_AND_ASSIGN(PerformanceManagerTest);
+};
+
+TEST_F(PerformanceManagerTest, GetPageNodeForWebContents) {
+  auto contents = CreateTestWebContents();
+
+  base::WeakPtr<PageNode> page_node =
+      PerformanceManager::GetPageNodeForWebContents(contents.get());
+
+  // Post a task to the Graph and make it call a function on the UI thread that
+  // will ensure that |page_node| is really associated with |contents|.
+
+  base::RunLoop run_loop;
+  auto check_wc_on_main_thread =
+      base::BindLambdaForTesting([&](const WebContentsProxy& wc_proxy) {
+        EXPECT_EQ(contents.get(), wc_proxy.Get());
+        run_loop.Quit();
+      });
+
+  auto call_on_graph_cb = base::BindLambdaForTesting([&](Graph* unused) {
+    EXPECT_TRUE(page_node.get());
+    base::PostTask(FROM_HERE, {content::BrowserThread::UI},
+                   base::BindOnce(std::move(check_wc_on_main_thread),
+                                  page_node->GetContentsProxy()));
+  });
+
+  PerformanceManager::CallOnGraph(FROM_HERE, call_on_graph_cb);
+
+  // Wait for |check_wc_on_main_thread| to be called.
+  run_loop.Run();
+
+  contents.reset();
+
+  // After deleting |contents| the corresponding PageNode WeakPtr should be
+  // invalid.
+  base::RunLoop run_loop_after_contents_reset;
+  auto quit_closure = run_loop_after_contents_reset.QuitClosure();
+  auto call_on_graph_cb_2 = base::BindLambdaForTesting([&](Graph* unused) {
+    EXPECT_FALSE(page_node.get());
+    std::move(quit_closure).Run();
+  });
+
+  PerformanceManager::CallOnGraph(FROM_HERE, call_on_graph_cb_2);
+  run_loop_after_contents_reset.Run();
+}
+
+}  // namespace performance_manager

components/performance_manager/public/performance_manager.h

@@ -7,11 +7,17 @@
 
 #include "base/callback.h"
 #include "base/location.h"
+#include "base/memory/weak_ptr.h"
+
+namespace content {
+class WebContents;
+}
 
 namespace performance_manager {
 
 class Graph;
 class GraphOwned;
+class PageNode;
 
 // The performance manager is a rendezvous point for communicating with the
 // performance manager graph on its dedicated sequence.
@@ -33,6 +39,14 @@ class PerformanceManager {
   static void PassToGraph(const base::Location& from_here,
                           std::unique_ptr<GraphOwned> graph_owned);
 
+  // Returns a WeakPtr to the PageNode associated with a given WebContents,
+  // or a null WeakPtr if there's no PageNode for this WebContents.
+  // Valid to call from the main thread only, the returned WeakPtr should only
+  // be dereferenced on the PM sequence (e.g. it can be used in a
+  // CallOnGraph callback).
+  static base::WeakPtr<PageNode> GetPageNodeForWebContents(
+      content::WebContents* wc);
+
  protected:
   PerformanceManager();
   virtual ~PerformanceManager();