[MBI] Track IPC listeners in AgentSchedulingGroup(Host).

For now the listener map is only used for obtaining a RouteProvider
(with a fallback to the process/thread). In the future it will also be
used for routing incoming IPCs for the ASG's listeners.

Bug: 1111231
Change-Id: I0a82f5681767b48751407ecbe8be7e7b6ecf506d
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/2513958
Commit-Queue: Tal Pressman <talp@chromium.org>
Reviewed-by: Kentaro Hara <haraken@chromium.org>
Reviewed-by: Kouhei Ueno <kouhei@chromium.org>
Reviewed-by: Dominic Farolino <dom@chromium.org>
Cr-Commit-Position: refs/heads/master@{#826548}

content/browser/renderer_host/agent_scheduling_group_host.cc

@@ -49,6 +49,10 @@ class AgentGroupHostUserData : public base::SupportsUserData::Data {
   std::unique_ptr<AgentSchedulingGroupHost> agent_group_;
 };
 
+RenderProcessHostImpl& ToImpl(RenderProcessHost& process) {
+  return static_cast<RenderProcessHostImpl&>(process);
+}
+
 }  // namespace
 
 // MaybeAssociatedReceiver:
@@ -232,10 +236,14 @@ bool AgentSchedulingGroupHost::Send(IPC::Message* message) {
 
 void AgentSchedulingGroupHost::AddRoute(int32_t routing_id,
                                         Listener* listener) {
+  DCHECK(!listener_map_.Lookup(routing_id));
+  listener_map_.AddWithID(listener, routing_id);
   process_.AddRoute(routing_id, listener);
 }
 
 void AgentSchedulingGroupHost::RemoveRoute(int32_t routing_id) {
+  DCHECK(listener_map_.Lookup(routing_id));
+  listener_map_.Remove(routing_id);
   process_.RemoveRoute(routing_id);
 }
 
@@ -292,10 +300,8 @@ void AgentSchedulingGroupHost::GetAssociatedInterface(
         receiver) {
   int32_t routing_id =
       associated_interface_provider_receivers_.current_context();
-  IPC::Listener* listener =
-      static_cast<RenderProcessHostImpl&>(process_).GetListener(PassKey(),
-                                                                routing_id);
-  if (listener)
+
+  if (auto* listener = GetListener(routing_id))
     listener->OnAssociatedInterfaceRequest(name, receiver.PassHandle());
 }
 
@@ -342,4 +348,16 @@ void AgentSchedulingGroupHost::SetUpMojoIfNeeded() {
       remote_route_provider_.BindNewEndpointAndPassReceiver());
 }
 
+Listener* AgentSchedulingGroupHost::GetListener(int32_t routing_id) {
+  if (routing_id == MSG_ROUTING_CONTROL)
+    return &process_;
+
+  if (auto* listener = listener_map_.Lookup(routing_id))
+    return listener;
+
+  // TODO(crbug.com/1111231): Can/should we log it when we find the listener on
+  // the process but not here?
+  return ToImpl(process_).GetListener(PassKey(), routing_id);
+}
+
 }  // namespace content

content/browser/renderer_host/agent_scheduling_group_host.h

@@ -7,6 +7,7 @@
 
 #include <stdint.h>
 
+#include "base/containers/id_map.h"
 #include "content/common/agent_scheduling_group.mojom.h"
 #include "content/common/associated_interfaces.mojom.h"
 #include "content/common/content_export.h"
@@ -175,11 +176,16 @@ class CONTENT_EXPORT AgentSchedulingGroupHost
   void ResetMojo();
   void SetUpMojoIfNeeded();
 
+  IPC::Listener* GetListener(int32_t routing_id);
+
   // The RenderProcessHost this AgentSchedulingGroup is assigned to.
   RenderProcessHost& process_;
 
   const bool should_associate_;
 
+  // Map of registered IPC listeners.
+  base::IDMap<IPC::Listener*> listener_map_;
+
   // Implementation of `mojom::AgentSchedulingGroupHost`, used for responding to
   // calls from the (renderer-side) `AgentSchedulingGroup`.
   MaybeAssociatedReceiver receiver_;

content/renderer/agent_scheduling_group.cc

@@ -128,20 +128,16 @@ bool AgentSchedulingGroup::Send(IPC::Message* message) {
   return RenderThread::Get()->Send(message);
 }
 
-// IPC messages to be forwarded to the `RenderThread`, for now. In the future
-// they will be handled directly by the `AgentSchedulingGroup`.
 void AgentSchedulingGroup::AddRoute(int32_t routing_id, Listener* listener) {
-  // TODO(crbug.com/1111231): For some reason, changing this to use
-  // render_thread_ causes trybots to time out (not specific tests).
-  RenderThread::Get()->AddRoute(routing_id, listener);
+  DCHECK(!listener_map_.Lookup(routing_id));
+  listener_map_.AddWithID(listener, routing_id);
+  render_thread_.AddRoute(routing_id, listener);
 }
 
-// IPC messages to be forwarded to the `RenderThread`, for now. In the future
-// they will be handled directly by the `AgentSchedulingGroup`.
 void AgentSchedulingGroup::RemoveRoute(int32_t routing_id) {
-  // TODO(crbug.com/1111231): For some reason, changing this to use
-  // render_thread_ causes trybots to time out (not specific tests).
-  RenderThread::Get()->RemoveRoute(routing_id);
+  DCHECK(listener_map_.Lookup(routing_id));
+  listener_map_.Remove(routing_id);
+  render_thread_.RemoveRoute(routing_id);
 }
 
 mojom::RouteProvider* AgentSchedulingGroup::GetRemoteRouteProvider() {
@@ -232,9 +228,21 @@ void AgentSchedulingGroup::GetAssociatedInterface(
         receiver) {
   int32_t routing_id =
       associated_interface_provider_receivers_.current_context();
-  IPC::Listener* listener = ToImpl(render_thread_).GetListener(routing_id);
-  if (listener)
+
+  if (auto* listener = GetListener(routing_id))
     listener->OnAssociatedInterfaceRequest(name, receiver.PassHandle());
 }
 
+Listener* AgentSchedulingGroup::GetListener(int32_t routing_id) {
+  if (routing_id == MSG_ROUTING_CONTROL)
+    return &ToImpl(render_thread_);
+
+  if (auto* listener = listener_map_.Lookup(routing_id))
+    return listener;
+
+  // TODO(crbug.com/1111231): Can/should we log it when we find the listener on
+  // the render thread but not here?
+  return ToImpl(render_thread_).GetListener(routing_id);
+}
+
 }  // namespace content

content/renderer/agent_scheduling_group.h

@@ -5,6 +5,7 @@
 #ifndef CONTENT_RENDERER_AGENT_SCHEDULING_GROUP_H_
 #define CONTENT_RENDERER_AGENT_SCHEDULING_GROUP_H_
 
+#include "base/containers/id_map.h"
 #include "content/common/agent_scheduling_group.mojom.h"
 #include "content/common/associated_interfaces.mojom.h"
 #include "content/common/content_export.h"
@@ -137,6 +138,11 @@ class CONTENT_EXPORT AgentSchedulingGroup
       mojo::PendingAssociatedReceiver<blink::mojom::AssociatedInterface>
           receiver) override;
 
+  IPC::Listener* GetListener(int32_t routing_id);
+
+  // Map of registered IPC listeners.
+  base::IDMap<IPC::Listener*> listener_map_;
+
   // A dedicated scheduler for this AgentSchedulingGroup.
   std::unique_ptr<blink::scheduler::WebAgentGroupScheduler>
       agent_group_scheduler_;