Add ref count to service workers for extension API.

extensions/browser/extension_function_dispatcher.cc

Index: extensions/browser/extension_function_dispatcher.cc
diff --git a/extensions/browser/extension_function_dispatcher.cc b/extensions/browser/extension_function_dispatcher.cc
index b944303ad7315f31102346a3060cbc4d2a5847ba..5cfd6ab745b48e07ad39dd0010342333e3d893dd 100644
--- a/extensions/browser/extension_function_dispatcher.cc
+++ b/extensions/browser/extension_function_dispatcher.cc
@@ -24,6 +24,7 @@
 #include "content/public/browser/render_process_host.h"
 #include "content/public/browser/render_process_host_observer.h"
 #include "content/public/browser/render_view_host.h"
+#include "content/public/browser/service_worker_context.h"
 #include "content/public/browser/user_metrics.h"
 #include "content/public/browser/web_contents.h"
 #include "content/public/browser/web_contents_observer.h"
@@ -32,6 +33,7 @@
 #include "extensions/browser/extension_function_registry.h"
 #include "extensions/browser/extension_registry.h"
 #include "extensions/browser/extension_system.h"
+#include "extensions/browser/extension_util.h"
 #include "extensions/browser/extensions_browser_client.h"
 #include "extensions/browser/io_thread_extension_message_filter.h"
 #include "extensions/browser/process_manager.h"
@@ -474,6 +476,7 @@ void ExtensionFunctionDispatcher::DispatchWithCallbackInternal(
   }
   if (params.embedded_worker_id != -1) {
     function_ui->set_is_from_service_worker(true);
+    function_ui->set_embedded_worker_id(params.embedded_worker_id);
   } else {
     function_ui->SetRenderFrameHost(render_frame_host);
   }
@@ -542,7 +545,17 @@ void ExtensionFunctionDispatcher::DispatchWithCallbackInternal(
   // compatible with lazy background pages.
   // TODO(lazyboy): API functions from extension Service Worker will incorrectly
   // change keepalive count below.
-  process_manager->IncrementLazyKeepaliveCount(extension);
+  if (params.embedded_worker_id == -1) {
+    // Increment ref count for non-service worker extension API.
+    process_manager->IncrementLazyKeepaliveCount(extension);
+  } else {
+    // Increment ref count for service worker extension API.
+    LOG(ERROR) << "Increment ref count";
+    content::ServiceWorkerContext::IncrementEmbeddedWorkerRefCount(
+        extensions::util::GetStoragePartitionForExtensionId(params.extension_id,
+                                                            browser_context()),
+        params.embedded_worker_id);
+  }
 }
 
 void ExtensionFunctionDispatcher::RemoveWorkerCallbacksForProcess(
@@ -560,12 +573,21 @@ void ExtensionFunctionDispatcher::RemoveWorkerCallbacksForProcess(
 }
 
 void ExtensionFunctionDispatcher::OnExtensionFunctionCompleted(
-    const Extension* extension) {
-  // TODO(lazyboy): API functions from extension Service Worker will incorrectly
-  // change keepalive count below.
+    const Extension* extension,
+    int embedded_worker_id) {
   if (extension) {
-    ProcessManager::Get(browser_context_)
-        ->DecrementLazyKeepaliveCount(extension);
+    if (embedded_worker_id == -1) {
+      // Decrement ref count for non-service worker extension API.
+      ProcessManager::Get(browser_context_)
+          ->DecrementLazyKeepaliveCount(extension);
+    } else {
+      // Decrement ref count for service worker extension API.
+      LOG(ERROR) << "Decrement ref count.";
+      content::ServiceWorkerContext::DecrementEmbeddedWorkerRefCount(
+          extensions::util::GetStoragePartitionForExtensionId(extension->id(),
+                                                              browser_context_),
+          embedded_worker_id);
+    }
   }
 }