Fix bug where transient pages would miss events dispatched while it was

chrome/browser/extensions/extension_message_service.cc

Index: chrome/browser/extensions/extension_message_service.cc
diff --git a/chrome/browser/extensions/extension_message_service.cc b/chrome/browser/extensions/extension_message_service.cc
index ccb555e958f3f27a2915397a4b5acd945f0aa62a..2fe1656e12c8edebe4f126631fccd2fd43bfc487 100644
--- a/chrome/browser/extensions/extension_message_service.cc
+++ b/chrome/browser/extensions/extension_message_service.cc
@@ -13,6 +13,7 @@
 #include "chrome/browser/extensions/extension_host.h"
 #include "chrome/browser/extensions/extension_process_manager.h"
 #include "chrome/browser/extensions/extension_service.h"
+#include "chrome/browser/extensions/extension_system.h"
 #include "chrome/browser/extensions/extension_tab_util.h"
 #include "chrome/browser/extensions/lazy_background_task_queue.h"
 #include "chrome/browser/extensions/process_map.h"
@@ -47,16 +48,25 @@ using content::WebContents;
 struct ExtensionMessageService::MessagePort {
   content::RenderProcessHost* process;
   int routing_id;
-  explicit MessagePort(content::RenderProcessHost* process = NULL,
-                       int routing_id = MSG_ROUTING_CONTROL)
-     : process(process), routing_id(routing_id) {}
+  std::string extension_id;
+  void* background_host_ptr;  // used in IncrementLazyKeepaliveCount
+
+  MessagePort()
+     : process(NULL),
+       routing_id(MSG_ROUTING_CONTROL),
+       background_host_ptr(NULL) {}
+  MessagePort(content::RenderProcessHost* process,
+              int routing_id,
+              const std::string& extension_id)
+     : process(process),
+       routing_id(routing_id),
+       extension_id(extension_id),
+       background_host_ptr(NULL) {}
 };
 
 struct ExtensionMessageService::MessageChannel {
   ExtensionMessageService::MessagePort opener;
   ExtensionMessageService::MessagePort receiver;
-  std::string source_extension_id;
-  std::string target_extension_id;
 };
 
 struct ExtensionMessageService::OpenChannelParams {
@@ -89,8 +99,7 @@ namespace {
 static base::StaticAtomicSequenceNumber g_next_channel_id;
 
 static void DispatchOnConnect(const ExtensionMessageService::MessagePort& port,
-                              int dest_port_id,
-                              const std::string& channel_name,
+                              int dest_port_id, const std::string& channel_name,
                               const std::string& tab_json,
                               const std::string& source_extension_id,
                               const std::string& target_extension_id) {
@@ -124,24 +133,30 @@ static content::RenderProcessHost* GetExtensionProcess(
   return site_instance->GetProcess();
 }
 
-static void IncrementLazyKeepaliveCount(content::RenderProcessHost* process,
-                                        const std::string& extension_id) {
-  Profile* profile = Profile::FromBrowserContext(process->GetBrowserContext());
-  const Extension* extension = profile->GetExtensionService()->extensions()->
-      GetByID(extension_id);
-  if (extension)
-    profile->GetExtensionProcessManager()->IncrementLazyKeepaliveCount(
-        extension);
+static void IncrementLazyKeepaliveCount(
+    ExtensionMessageService::MessagePort* port) {
+  Profile* profile =
+      Profile::FromBrowserContext(port->process->GetBrowserContext());
+  ExtensionProcessManager* pm =
+      ExtensionSystem::Get(profile)->process_manager();
+  ExtensionHost* host = pm->GetBackgroundHostForExtension(port->extension_id);
+  if (host && host->extension()->has_lazy_background_page())
+    pm->IncrementLazyKeepaliveCount(host->extension());
+
+  // Keep track of the background host, so when we decrement, we only do so if
+  // the host hasn't reloaded.
+  port->background_host_ptr = host;
 }
 
-static void DecrementLazyKeepaliveCount(content::RenderProcessHost* process,
-                                        const std::string& extension_id) {
-  Profile* profile = Profile::FromBrowserContext(process->GetBrowserContext());
-  const Extension* extension = profile->GetExtensionService()->extensions()->
-      GetByID(extension_id);
-  if (extension)
-    profile->GetExtensionProcessManager()->DecrementLazyKeepaliveCount(
-        extension);
+static void DecrementLazyKeepaliveCount(
+    ExtensionMessageService::MessagePort* port) {
+  Profile* profile =
+      Profile::FromBrowserContext(port->process->GetBrowserContext());
+  ExtensionProcessManager* pm =
+      ExtensionSystem::Get(profile)->process_manager();
+  ExtensionHost* host = pm->GetBackgroundHostForExtension(port->extension_id);
+  if (host && host == port->background_host_ptr)
+    pm->DecrementLazyKeepaliveCount(host->extension());
 }
 
 }  // namespace
@@ -195,7 +210,8 @@ void ExtensionMessageService::OpenChannelToExtension(
   // which depends on whether the extension uses spanning or split mode.
   MessagePort receiver(
       GetExtensionProcess(profile, target_extension_id),
-      MSG_ROUTING_CONTROL);
+      MSG_ROUTING_CONTROL,
+      target_extension_id);
   WebContents* source_contents = tab_util::GetWebContentsByID(
       source_process_id, source_routing_id);
 
@@ -237,12 +253,13 @@ void ExtensionMessageService::OpenChannelToTab(
     receiver.process = contents->web_contents()->GetRenderProcessHost();
     receiver.routing_id =
         contents->web_contents()->GetRenderViewHost()->GetRoutingID();
+    receiver.extension_id = extension_id;
   }
 
   if (contents && contents->web_contents()->GetController().NeedsReload()) {
     // The tab isn't loaded yet. Don't attempt to connect. Treat this as a
     // disconnect.
-    DispatchOnDisconnect(MessagePort(source, MSG_ROUTING_CONTROL),
+    DispatchOnDisconnect(MessagePort(source, MSG_ROUTING_CONTROL, extension_id),
                          GET_OPPOSITE_PORT_ID(receiver_port_id), true);
     return;
   }
@@ -269,7 +286,7 @@ bool ExtensionMessageService::OpenChannelImpl(const OpenChannelParams& params) {
 
   if (!params.receiver.process) {
     // Treat it as a disconnect.
-    DispatchOnDisconnect(MessagePort(params.source, MSG_ROUTING_CONTROL),
+    DispatchOnDisconnect(MessagePort(params.source, MSG_ROUTING_CONTROL, ""),
                          GET_OPPOSITE_PORT_ID(params.receiver_port_id), true);
     return false;
   }
@@ -279,10 +296,9 @@ bool ExtensionMessageService::OpenChannelImpl(const OpenChannelParams& params) {
   CHECK(params.receiver.process);
 
   MessageChannel* channel(new MessageChannel);
-  channel->opener = MessagePort(params.source, MSG_ROUTING_CONTROL);
+  channel->opener = MessagePort(params.source, MSG_ROUTING_CONTROL,
+                                params.source_extension_id);
   channel->receiver = params.receiver;
-  channel->source_extension_id = params.source_extension_id;
-  channel->target_extension_id = params.target_extension_id;
 
   CHECK(params.receiver.process);
 
@@ -300,10 +316,8 @@ bool ExtensionMessageService::OpenChannelImpl(const OpenChannelParams& params) {
                     params.source_extension_id, params.target_extension_id);
 
   // Keep both ends of the channel alive until the channel is closed.
-  IncrementLazyKeepaliveCount(channel->opener.process,
-                              channel->source_extension_id);
-  IncrementLazyKeepaliveCount(channel->receiver.process,
-                              channel->target_extension_id);
+  IncrementLazyKeepaliveCount(&channel->opener);
+  IncrementLazyKeepaliveCount(&channel->receiver);
   return true;
 }
 
@@ -338,10 +352,8 @@ void ExtensionMessageService::CloseChannelImpl(
   }
 
   // Balance the addrefs in OpenChannelImpl.
-  DecrementLazyKeepaliveCount(channel->opener.process,
-                              channel->source_extension_id);
-  DecrementLazyKeepaliveCount(channel->receiver.process,
-                              channel->target_extension_id);
+  DecrementLazyKeepaliveCount(&channel->opener);
+  DecrementLazyKeepaliveCount(&channel->receiver);
 
   delete channel_iter->second;
   channels_.erase(channel_iter);
@@ -452,6 +464,7 @@ void ExtensionMessageService::PendingOpenChannel(
     return;
 
   params.receiver = MessagePort(host->render_process_host(),
-                                MSG_ROUTING_CONTROL);
+                                MSG_ROUTING_CONTROL,
+                                params.target_extension_id);
   OpenChannelImpl(params);
 }