Move Extension messaging port release functionality from JS to C++.

extensions/renderer/messaging_bindings.cc

 // Copyright 2014 The Chromium Authors. All rights reserved.
 // Use of this source code is governed by a BSD-style license that can be
 // found in the LICENSE file.
 
 #include "extensions/renderer/messaging_bindings.h"
 
 #include <map>
 #include <string>
 
 #include "base/basictypes.h"
@@ skipping 101 matching lines @@
   // Function to run when |object_| bound to this GCCallback is GC'd.
   v8::Global<v8::Function> callback_;
 
   // Function to run if context is invalidated before we have a chance
   // to execute |callback_|.
   base::Closure fallback_;
 
   DISALLOW_COPY_AND_ASSIGN(GCCallback);
 };
 
-struct ExtensionData {
-  struct PortData {
-    int ref_count;  // how many contexts have a handle to this port
-    PortData() : ref_count(0) {}
-  };
-  std::map<int, PortData> ports;  // port ID -> data
+// Maintains the set of ScriptContexts that reference each Port, by ID.
+class PortTracker {
+ public:
+  PortTracker() {}
+
+  // Returns true if |context| references |port_id|.
+  bool HasReference(ScriptContext* context, int port_id) const {
+    auto ports = contexts_to_ports_.find(context);
+    if (ports == contexts_to_ports_.end())
+      return false;
+    return ports->second.count(port_id) > 0;
+  }
+
+  // Marks |context| and |port_id| as referencing each other.
+  void AddReference(ScriptContext* context, int port_id) {
+    if (contexts_to_ports_[context].insert(port_id).second)
+      ++port_refcount_[port_id];
+  }
+
+  // Removes the references between |context| and |port_id|.
+  void RemoveReference(ScriptContext* context, int port_id) {
+    auto ports = contexts_to_ports_.find(context);
+    if (ports == contexts_to_ports_.end()) {
+      return;
+    }
+    if (ports->second.erase(port_id) == 0)
+      return;
+    if (ports->second.size() == 0)
+      contexts_to_ports_.erase(context);
+    if (--port_refcount_[port_id] == 0)
+      port_refcount_.erase(port_id);
+  }
+
+  // Returns true if this tracker has any reference to |port_id|.
+  bool HasPort(int port_id) const {
+    return port_refcount_.find(port_id) != port_refcount_.end();
+  }
+
+  // Deletes all references to |port_id|.
+  void DeletePort(int port_id) {
+    port_refcount_.erase(port_id);
+    for (auto it = contexts_to_ports_.begin();
+         it != contexts_to_ports_.end();) {
+      if (it->second.erase(port_id) > 0 && it->second.empty())
+        contexts_to_ports_.erase(it++);
+      else
+        ++it;
+    }
+  }
+
+  // Gets every port ID that has a reference to |context|.
+  std::set<int> GetPorts(ScriptContext* context) const {
+    auto ports = contexts_to_ports_.find(context);
+    return ports == contexts_to_ports_.end() ? std::set<int>() : ports->second;
+  }
+
+ private:
+  // Maps ScriptContexts to the port IDs that have a reference to it.
+  std::map<ScriptContext*, std::set<int>> contexts_to_ports_;
+
+  // Maps port IDs to the number of ScriptContexts that have a reference to it.
+  std::map<int, int> port_refcount_;
+
+  DISALLOW_COPY_AND_ASSIGN(PortTracker);
 };
 
-base::LazyInstance<ExtensionData> g_extension_data = LAZY_INSTANCE_INITIALIZER;
-
-bool HasPortData(int port_id) {
-  return g_extension_data.Get().ports.find(port_id) !=
-         g_extension_data.Get().ports.end();
-}
-
-ExtensionData::PortData& GetPortData(int port_id) {
-  return g_extension_data.Get().ports[port_id];
-}
-
-void ClearPortData(int port_id) {
-  g_extension_data.Get().ports.erase(port_id);
-}
+base::LazyInstance<PortTracker> g_port_tracker = LAZY_INSTANCE_INITIALIZER;
 
 const char kPortClosedError[] = "Attempting to use a disconnected port object";
 const char kReceivingEndDoesntExistError[] =
     "Could not establish connection. Receiving end does not exist.";
 
 class ExtensionImpl : public ObjectBackedNativeHandler {
  public:
   ExtensionImpl(Dispatcher* dispatcher, ScriptContext* context)
       : ObjectBackedNativeHandler(context),
         dispatcher_(dispatcher),
         weak_ptr_factory_(this) {
     RouteFunction(
         "CloseChannel",
         base::Bind(&ExtensionImpl::CloseChannel, base::Unretained(this)));
     RouteFunction(
         "PortAddRef",
         base::Bind(&ExtensionImpl::PortAddRef, base::Unretained(this)));
     RouteFunction(
         "PortRelease",
         base::Bind(&ExtensionImpl::PortRelease, base::Unretained(this)));
     RouteFunction(
         "PostMessage",
         base::Bind(&ExtensionImpl::PostMessage, base::Unretained(this)));
     // TODO(fsamuel, kalman): Move BindToGC out of messaging natives.
     RouteFunction("BindToGC",
                   base::Bind(&ExtensionImpl::BindToGC, base::Unretained(this)));
+
+    // Observe |context| so that port references to it can be cleared.
+    context->AddInvalidationObserver(base::Bind(
+        &ExtensionImpl::OnContextInvalidated, weak_ptr_factory_.GetWeakPtr()));
   }
 
   ~ExtensionImpl() override {}
 
  private:
+  void OnContextInvalidated() {
+    for (int port_id : g_port_tracker.Get().GetPorts(context()))
+      ReleasePort(port_id);
+  }
+
   void ClearPortDataAndNotifyDispatcher(int port_id) {
-    ClearPortData(port_id);
+    g_port_tracker.Get().DeletePort(port_id);
     dispatcher_->ClearPortData(port_id);
   }
 
   // Sends a message along the given channel.
   void PostMessage(const v8::FunctionCallbackInfo<v8::Value>& args) {
     content::RenderFrame* renderframe = context()->GetRenderFrame();
     if (!renderframe)
       return;
 
     // Arguments are (int32 port_id, string message).
     CHECK(args.Length() == 2 && args[0]->IsInt32() && args[1]->IsString());
 
     int port_id = args[0]->Int32Value();
-    if (!HasPortData(port_id)) {
+    if (!g_port_tracker.Get().HasPort(port_id)) {
       args.GetIsolate()->ThrowException(v8::Exception::Error(
           v8::String::NewFromUtf8(args.GetIsolate(), kPortClosedError)));
       return;
     }
 
     renderframe->Send(new ExtensionHostMsg_PostMessage(
         renderframe->GetRoutingID(), port_id,
         Message(*v8::String::Utf8Value(args[1]),
                 blink::WebUserGestureIndicator::isProcessingUserGesture())));
   }
 
   // Forcefully disconnects a port.
   void CloseChannel(const v8::FunctionCallbackInfo<v8::Value>& args) {
     // Arguments are (int32 port_id, boolean notify_browser).
     CHECK_EQ(2, args.Length());
     CHECK(args[0]->IsInt32());
     CHECK(args[1]->IsBoolean());
 
     int port_id = args[0]->Int32Value();
-    if (!HasPortData(port_id))
+    if (!g_port_tracker.Get().HasPort(port_id))
       return;
 
     // Send via the RenderThread because the RenderFrame might be closing.
     bool notify_browser = args[1]->BooleanValue();
     if (notify_browser) {
       content::RenderThread::Get()->Send(
           new ExtensionHostMsg_CloseChannel(port_id, std::string()));
     }
 
     ClearPortDataAndNotifyDispatcher(port_id);
   }
 
   // A new port has been created for a context.  This occurs both when script
   // opens a connection, and when a connection is opened to this script.
   void PortAddRef(const v8::FunctionCallbackInfo<v8::Value>& args) {
     // Arguments are (int32 port_id).
     CHECK_EQ(1, args.Length());
     CHECK(args[0]->IsInt32());
 
     int port_id = args[0]->Int32Value();
-    ++GetPortData(port_id).ref_count;
+    g_port_tracker.Get().AddReference(context(), port_id);
   }
 
   // The frame a port lived in has been destroyed.  When there are no more
   // frames with a reference to a given port, we will disconnect it and notify
   // the other end of the channel.
   void PortRelease(const v8::FunctionCallbackInfo<v8::Value>& args) {
     // Arguments are (int32 port_id).
     CHECK(args.Length() == 1 && args[0]->IsInt32());
     ReleasePort(args[0]->Int32Value());
   }
 
-  // Implementation of both the PortRelease native handler call, and callback
-  // when contexts are invalidated to release their ports.
+  // Releases the reference to |port_id| for this context, and clears all port
+  // data if there are no more references.
   void ReleasePort(int port_id) {
-    if (HasPortData(port_id) && --GetPortData(port_id).ref_count == 0) {
+    g_port_tracker.Get().RemoveReference(context(), port_id);
+    if (!g_port_tracker.Get().HasPort(port_id)) {
       // Send via the RenderThread because the RenderFrame might be closing.
       content::RenderThread::Get()->Send(
           new ExtensionHostMsg_CloseChannel(port_id, std::string()));
       ClearPortDataAndNotifyDispatcher(port_id);
     }
   }
 
   // void BindToGC(object, callback, port_id)
   //
   // Binds |callback| to be invoked *sometime after* |object| is garbage
@@ skipping 227 matching lines @@
   // TODO(robwu): ScriptContextSet.ForEach should accept RenderFrame*.
   content::RenderView* restrict_to_render_view =
       restrict_to_render_frame ? restrict_to_render_frame->GetRenderView()
                                : NULL;
   context_set.ForEach(
       restrict_to_render_view,
       base::Bind(&DispatchOnDisconnectToScriptContext, port_id, error_message));
 }
 
 }  // namespace extensions