[headless] Improved Simple CDP Client callback handling.

Previous implementation of the asynchronous callbacks posted a task
for each result and event which is suboptimal. This CL changes logic
so that all the result and event callbacks are called from a single
task posted on the browser UI thread.

Drive by: fixed occasional use-after-move in event handler.

Bug: 1382993
Change-Id: I031044a5bfdaab2b88455d1649f53e4adf053140
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/4026625
Reviewed-by: Dmitry Gozman <dgozman@chromium.org>
Commit-Queue: Peter Kvitek <kvitekp@chromium.org>
Cr-Commit-Position: refs/heads/main@{#1071300}

components/devtools/simple_devtools_protocol_client/BUILD.gn

@@ -26,6 +26,7 @@ source_set("unit_tests") {
     ":simple_devtools_protocol_client",
     "//base",
     "//content/public/browser",
+    "//content/test:test_support",
     "//testing/gmock",
     "//testing/gtest",
   ]

components/devtools/simple_devtools_protocol_client/DEPS

@@ -1,3 +1,4 @@
 include_rules = [
   "+content/public/browser",
+  "+content/public/test",
 ]

components/devtools/simple_devtools_protocol_client/simple_devtools_protocol_client.cc

@@ -35,6 +35,7 @@ int g_next_message_id = 0;
 }  // namespace
 
 SimpleDevToolsProtocolClient::SimpleDevToolsProtocolClient() = default;
+
 SimpleDevToolsProtocolClient::SimpleDevToolsProtocolClient(
     const std::string& session_id)
     : session_id_(session_id) {}
@@ -44,10 +45,6 @@ SimpleDevToolsProtocolClient::~SimpleDevToolsProtocolClient() {
     parent_client_->sessions_.erase(session_id_);
 }
 
-void SimpleDevToolsProtocolClient::InitBrowserMainThread() {
-  browser_main_thread_ = content::GetUIThreadTaskRunner({});
-}
-
 void SimpleDevToolsProtocolClient::AttachClient(
     scoped_refptr<content::DevToolsAgentHost> agent_host) {
   DCHECK(!agent_host_);
@@ -94,12 +91,19 @@ void SimpleDevToolsProtocolClient::DispatchProtocolMessage(
   if (const std::string* session_id = message.FindString("sessionId")) {
     auto it = sessions_.find(*session_id);
     if (it != sessions_.cend()) {
-      it->second->DispatchProtocolMessage(std::move(message));
+      content::GetUIThreadTaskRunner({})->PostTask(
+          FROM_HERE,
+          base::BindOnce(
+              &SimpleDevToolsProtocolClient::DispatchProtocolMessageTask,
+              base::Unretained(it->second), std::move(message)));
       return;
     }
   }
 
-  DispatchProtocolMessage(std::move(message));
+  content::GetUIThreadTaskRunner({})->PostTask(
+      FROM_HERE,
+      base::BindOnce(&SimpleDevToolsProtocolClient::DispatchProtocolMessageTask,
+                     base::Unretained(this), std::move(message)));
 }
 
 void SimpleDevToolsProtocolClient::AgentHostClosed(
@@ -110,7 +114,7 @@ void SimpleDevToolsProtocolClient::AgentHostClosed(
   }
 }
 
-void SimpleDevToolsProtocolClient::DispatchProtocolMessage(
+void SimpleDevToolsProtocolClient::DispatchProtocolMessageTask(
     base::Value::Dict message) {
   // Handle response message shutting down the host if it's unexpected.
   if (absl::optional<int> id = message.FindInt(kId)) {
@@ -127,12 +131,7 @@ void SimpleDevToolsProtocolClient::DispatchProtocolMessage(
     ResponseCallback callback(std::move(it->second));
     pending_response_map_.erase(it);
 
-    if (browser_main_thread_) {
-      browser_main_thread_->PostTask(
-          FROM_HERE, base::BindOnce(std::move(callback), std::move(message)));
-    } else {
-      std::move(callback).Run(std::move(message));
-    }
+    std::move(callback).Run(std::move(message));
     return;
   }
 
@@ -154,12 +153,7 @@ void SimpleDevToolsProtocolClient::DispatchProtocolMessage(
   for (auto& callback : handlers) {
     if (first_callback || HasEventHandler(*event_name, callback)) {
       first_callback = false;
-      if (browser_main_thread_) {
-        browser_main_thread_->PostTask(
-            FROM_HERE, base::BindOnce(std::move(callback), std::move(message)));
-      } else {
-        callback.Run(std::move(message));
-      }
+      callback.Run(message);
     }
   }
 }

components/devtools/simple_devtools_protocol_client/simple_devtools_protocol_client.h

@@ -13,7 +13,6 @@
 #include "base/containers/span.h"
 #include "base/memory/raw_ptr.h"
 #include "base/memory/scoped_refptr.h"
-#include "base/task/sequenced_task_runner.h"
 #include "base/values.h"
 #include "content/public/browser/devtools_agent_host.h"
 
@@ -33,8 +32,6 @@ class SimpleDevToolsProtocolClient : public content::DevToolsAgentHostClient {
 
   ~SimpleDevToolsProtocolClient() override;
 
-  void InitBrowserMainThread();
-
   void AttachClient(scoped_refptr<content::DevToolsAgentHost> agent_host);
   void DetachClient();
 
@@ -66,7 +63,7 @@ class SimpleDevToolsProtocolClient : public content::DevToolsAgentHostClient {
                                base::span<const uint8_t> json_message) override;
   void AgentHostClosed(content::DevToolsAgentHost* agent_host) override;
 
-  void DispatchProtocolMessage(base::Value::Dict message);
+  void DispatchProtocolMessageTask(base::Value::Dict message);
 
   void SendProtocolMessage(base::Value::Dict message);
 
@@ -74,7 +71,6 @@ class SimpleDevToolsProtocolClient : public content::DevToolsAgentHostClient {
                        const EventCallback& event_callback);
 
   const std::string session_id_;
-  scoped_refptr<base::SequencedTaskRunner> browser_main_thread_;
   base::raw_ptr<SimpleDevToolsProtocolClient> parent_client_ = nullptr;
   base::flat_map<std::string, SimpleDevToolsProtocolClient*> sessions_;
 

components/devtools/simple_devtools_protocol_client/simple_devtools_protocol_client_unittest.cc

@@ -13,6 +13,7 @@
 #include "base/memory/scoped_refptr.h"
 #include "base/values.h"
 #include "content/public/browser/devtools_agent_host.h"
+#include "content/public/test/browser_task_environment.h"
 #include "testing/gmock/include/gmock/gmock.h"
 #include "testing/gtest/include/gtest/gtest.h"
 
@@ -86,6 +87,11 @@ class SimpleDevToolsProtocolClientTest : public SimpleDevToolsProtocolClient,
   SimpleDevToolsProtocolClientTest() {
     AttachClient(new MockDevToolsAgentHost);
   }
+
+  void RunUntilIdle() { task_environment_.RunUntilIdle(); }
+
+ private:
+  content::BrowserTaskEnvironment task_environment_;
 };
 
 class SimpleDevToolsProtocolClientSendCommandTest
@@ -99,6 +105,7 @@ class SimpleDevToolsProtocolClientSendCommandTest
                 base::BindOnce(&SimpleDevToolsProtocolClientSendCommandTest::
                                    OnSendCommand1Response,
                                base::Unretained(this)));
+    RunUntilIdle();
   }
 
   void OnSendCommand1Response(base::Value::Dict params) {
@@ -108,6 +115,7 @@ class SimpleDevToolsProtocolClientSendCommandTest
                 base::BindOnce(&SimpleDevToolsProtocolClientSendCommandTest::
                                    OnSendCommand2Response,
                                base::Unretained(this)));
+    RunUntilIdle();
   }
 
   void OnSendCommand2Response(base::Value::Dict params) {
@@ -120,6 +128,7 @@ class SimpleDevToolsProtocolClientSendCommandTest
                       self->OnSendCommand3Response(std::move(params));
                     },
                     base::Unretained(this)));
+    RunUntilIdle();
   }
 
   void OnSendCommand3Response(base::Value::Dict params) {
@@ -147,6 +156,7 @@ class SimpleDevToolsProtocolClientEventHandlerTest
     base::JSONWriter::Write(base::Value(std::move(params)), &json);
     DispatchProtocolMessage(agent_host_.get(),
                             base::as_bytes(base::make_span(json)));
+    RunUntilIdle();
   }
 };
 

headless/app/headless_shell.cc

@@ -242,8 +242,6 @@ void HeadlessShell::OnBrowserStart(HeadlessBrowser* browser) {
   }
 #endif
 
-  devtools_client_.InitBrowserMainThread();
-
   file_task_runner_ = base::ThreadPool::CreateSequencedTaskRunner(
       {base::MayBlock(), base::TaskPriority::BEST_EFFORT});