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..708ed2d078491d0f653271c9d3ecdf10283d6c8c 100644
--- a/extensions/browser/extension_function_dispatcher.cc
+++ b/extensions/browser/extension_function_dispatcher.cc
@@ -19,11 +19,13 @@
 #include "base/scoped_observer.h"
 #include "base/values.h"
 #include "build/build_config.h"
+#include "content/common/service_worker/service_worker_types.h"
 #include "content/public/browser/browser_thread.h"
 #include "content/public/browser/render_frame_host.h"
 #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 +34,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"
@@ -60,6 +63,12 @@ void NotifyApiFunctionCalled(const std::string& extension_id,
                                         args);
 }
 
+bool IsRequestFromServiceWorker(
+    const ExtensionHostMsg_Request_Params& request_params) {
+  return content::kInvalidServiceWorkerVersionId !=

[Devlin, 2016/10/05 17:02:25]

nit: prefer inverting these

[lazyboy, 2016/10/06 01:02:06]

Done.

+         request_params.service_worker_version_id;
+}
+
 // Separate copy of ExtensionAPI used for IO thread extension functions. We need
 // this because ExtensionAPI has mutable data. It should be possible to remove
 // this once all the extension APIs are updated to the feature system.
@@ -277,17 +286,18 @@ class ExtensionFunctionDispatcher::UIThreadWorkerResponseCallbackWrapper
 };
 
 struct ExtensionFunctionDispatcher::WorkerResponseCallbackMapKey {
-  WorkerResponseCallbackMapKey(int render_process_id, int embedded_worker_id)
+  WorkerResponseCallbackMapKey(int render_process_id,
+                               int64_t service_worker_version_id)
       : render_process_id(render_process_id),
-        embedded_worker_id(embedded_worker_id) {}
+        service_worker_version_id(service_worker_version_id) {}
 
   bool operator<(const WorkerResponseCallbackMapKey& other) const {
-    return std::tie(render_process_id, embedded_worker_id) <
-           std::tie(other.render_process_id, other.embedded_worker_id);
+    return std::tie(render_process_id, service_worker_version_id) <
+           std::tie(other.render_process_id, other.service_worker_version_id);
   }
 
   int render_process_id;
-  int embedded_worker_id;
+  int64_t service_worker_version_id;
 };
 
 WindowController*
@@ -418,9 +428,10 @@ void ExtensionFunctionDispatcher::Dispatch(
         callback_wrapper->CreateCallback(params.request_id));
   } else {
     // Extension API from Service Worker.
-    DCHECK_GE(params.embedded_worker_id, 0);
+    DCHECK_NE(params.service_worker_version_id,
+              content::kInvalidServiceWorkerVersionId);

[Devlin, 2016/10/05 17:02:25]

nit: invert

[lazyboy, 2016/10/06 01:02:06]

Done.

     WorkerResponseCallbackMapKey key(render_process_id,
-                                     params.embedded_worker_id);
+                                     params.service_worker_version_id);
     UIThreadWorkerResponseCallbackWrapperMap::const_iterator iter =
         ui_thread_response_callback_wrappers_for_worker_.find(key);
     UIThreadWorkerResponseCallbackWrapper* callback_wrapper = nullptr;
@@ -472,8 +483,9 @@ void ExtensionFunctionDispatcher::DispatchWithCallbackInternal(
     NOTREACHED();
     return;
   }
-  if (params.embedded_worker_id != -1) {
-    function_ui->set_is_from_service_worker(true);
+  if (IsRequestFromServiceWorker(params)) {
+    function_ui->set_service_worker_version_id(
+        params.service_worker_version_id);
   } else {
     function_ui->SetRenderFrameHost(render_frame_host);
   }
@@ -536,13 +548,18 @@ void ExtensionFunctionDispatcher::DispatchWithCallbackInternal(
   if (!registry->enabled_extensions().GetByID(params.extension_id))
     return;
 
-  // We only adjust the keepalive count for UIThreadExtensionFunction for
-  // now, largely for simplicity's sake. This is OK because currently, only
-  // the webRequest API uses IOThreadExtensionFunction, and that API is not
-  // compatible with lazy background pages.
-  // TODO(lazyboy): API functions from extension Service Worker will incorrectly
-  // change keepalive count below.
-  process_manager->IncrementLazyKeepaliveCount(extension);
+  if (IsRequestFromServiceWorker(params)) {

[Devlin, 2016/10/05 17:02:25]

Similar comment to line 583:
if (!IsRequestFromSW()) {
  // Increment ref count for non-service worker extension API. Ref counts for
service...
}

[lazyboy, 2016/10/06 01:02:06]

Done.

+    // Incrementing ref count of service worker extension API is done separately
+    // on IO thread via IPC.
+  } else {
+    // We only adjust the keepalive count for UIThreadExtensionFunction for
+    // now, largely for simplicity's sake. This is OK because currently, only
+    // the webRequest API uses IOThreadExtensionFunction, and that API is not
+    // compatible with lazy background pages.
+
+    // Increment ref count for non-service worker extension API.
+    process_manager->IncrementLazyKeepaliveCount(function->extension());
+  }
 }
 
 void ExtensionFunctionDispatcher::RemoveWorkerCallbacksForProcess(
@@ -560,12 +577,17 @@ 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,
+    int64_t service_worker_version_id) {

[Devlin, 2016/10/05 17:02:25]

Perhaps just pass in bool is_from_service_worker, since we don't need the id?

[lazyboy, 2016/10/06 01:02:06]

Done.

   if (extension) {
-    ProcessManager::Get(browser_context_)
-        ->DecrementLazyKeepaliveCount(extension);
+    if (service_worker_version_id == content::kInvalidServiceWorkerVersionId) {

[Devlin, 2016/10/05 17:02:25]

I think this is cleaner as:
if (extension && !is_from_service_worker) {
  // Decrement ref count for non-service worker extension API. Service worker
ref counts
  // are handled separately on the IO thread directly via IPC.
}

[lazyboy, 2016/10/06 01:02:06]

Done.

+      // Decrement ref count for non-service worker extension API.
+      ProcessManager::Get(browser_context_)
+          ->DecrementLazyKeepaliveCount(extension);
+    } else {
+      // Decrement ref count for service worker extension API is done separately
+      // on IO thread directly via IPC.
+    }
   }
 }