Add support for dispatch contexts on BindingSet

mojo/public/cpp/bindings/binding_set.h

Index: mojo/public/cpp/bindings/binding_set.h
diff --git a/mojo/public/cpp/bindings/binding_set.h b/mojo/public/cpp/bindings/binding_set.h
index b1baca6a808dc3e6e637bbd408e7f424dfae828d..dc66860b4e0d78fba19a3de4727f59bb6e235a59 100644
--- a/mojo/public/cpp/bindings/binding_set.h
+++ b/mojo/public/cpp/bindings/binding_set.h
@@ -5,107 +5,172 @@
 #ifndef MOJO_PUBLIC_CPP_BINDINGS_BINDING_SET_H_
 #define MOJO_PUBLIC_CPP_BINDINGS_BINDING_SET_H_
 
-#include <algorithm>
 #include <utility>
-#include <vector>
 
 #include "base/bind.h"
 #include "base/callback.h"
 #include "base/macros.h"
-#include "base/memory/weak_ptr.h"
 #include "mojo/public/cpp/bindings/binding.h"
+#include "mojo/public/cpp/bindings/interface_ptr.h"
+#include "mojo/public/cpp/bindings/interface_request.h"
+#include "mojo/public/cpp/bindings/message.h"
 
 namespace mojo {
 
+template <typename BindingType>
+struct BindingSetTraits;
+
+template <typename Interface>
+struct BindingSetTraits<Binding<Interface>> {
+  using ProxyType = InterfacePtr<Interface>;
+  using RequestType = InterfaceRequest<Interface>;
+
+  static RequestType GetProxy(ProxyType* proxy) {
+    return mojo::GetProxy(proxy);
+  }
+};
+
+enum class BindingSetDispatchMode {
+  WITHOUT_CONTEXT,
+  WITH_CONTEXT,
+};
+
 // Use this class to manage a set of bindings, which are automatically destroyed
 // and removed from the set when the pipe they are bound to is disconnected.
-template <typename Interface>
+template <typename Interface, typename BindingType = Binding<Interface>>
 class BindingSet {
  public:
-  BindingSet() {}
-  ~BindingSet() { CloseAllBindings(); }
+  using Traits = BindingSetTraits<BindingType>;
+  using ProxyType = typename Traits::ProxyType;
+  using RequestType = typename Traits::RequestType;
+
+  BindingSet() : BindingSet(BindingSetDispatchMode::WITHOUT_CONTEXT) {}
+
+  // Constructs a new BindingSet operating in |dispatch_mode|. If |WITH_CONTEXT|
+  // is used, AddBinding() supports a |context| argument, and dispatch_context()
+  // may be called during message or error dispatch to identify which specific
+  // binding received the message or error.
+  explicit BindingSet(BindingSetDispatchMode dispatch_mode)
+      : dispatch_mode_(dispatch_mode) {}
 
   void set_connection_error_handler(const base::Closure& error_handler) {
     error_handler_ = error_handler;
   }
 
-  void AddBinding(Interface* impl, InterfaceRequest<Interface> request) {
-    auto binding = new Element(impl, std::move(request));
-    binding->set_connection_error_handler(
-        base::Bind(&BindingSet::OnConnectionError, base::Unretained(this)));
-    bindings_.push_back(binding->GetWeakPtr());
+  // Adds a new binding to the set which binds |request| to |impl|. If |context|
+  // is non-null, dispatch_context() will reflect this value during the extent
+  // of any message or error dispatch targeting this specific binding. Note that
+  // |context| may only be non-null if the BindingSet was constructed with
+  // |BindingSetDispatchMode::WITH_CONTEXT|.
+  void AddBinding(Interface* impl,
+                  RequestType request,
+                  void* context = nullptr) {
+    DCHECK(!context || SupportsContext());
+    std::unique_ptr<Entry> entry =
+        base::MakeUnique<Entry>(impl, std::move(request), this, context);
+    bindings_.insert(std::make_pair(entry.get(), std::move(entry)));
   }
 
-  // Returns an InterfacePtr bound to one end of a pipe whose other end is
-  // bound to |this|.
-  InterfacePtr<Interface> CreateInterfacePtrAndBind(Interface* impl) {
-    InterfacePtr<Interface> interface_ptr;
-    AddBinding(impl, GetProxy(&interface_ptr));
-    return interface_ptr;
+  // Returns a proxy bound to one end of a pipe whose other end is bound to
+  // |this|.
+  ProxyType CreateInterfacePtrAndBind(Interface* impl) {
+    ProxyType proxy;
+    AddBinding(impl, Traits::GetProxy(&proxy));
+    return proxy;
   }
 
-  void CloseAllBindings() {
-    for (const auto& it : bindings_) {
-      if (it) {
-        it->Close();
-        delete it.get();
-      }
-    }
-    bindings_.clear();
-  }
+  void CloseAllBindings() { bindings_.clear(); }
 
   bool empty() const { return bindings_.empty(); }
 
+  // Implementations may call this when processing a dispatched message or
+  // error. During the extent of message or error dispatch, this will return the
+  // context associated with the specific binding which received the message or
+  // error. Use AddBinding() to associated a context with a specific binding.
+  //
+  // Note that this may ONLY be called if the BindingSet was constructed with
+  // |BindingSetDispatchMode::WITH_CONTEXT|.
+  void* dispatch_context() const {
+    DCHECK(SupportsContext());
+    return dispatch_context_;
+  }
+
  private:
-  class Element {
+  friend class Entry;
+
+  class Entry {
    public:
-    Element(Interface* impl, InterfaceRequest<Interface> request)
-        : binding_(impl, std::move(request)), weak_ptr_factory_(this) {
-      binding_.set_connection_error_handler(
-          base::Bind(&Element::OnConnectionError, base::Unretained(this)));
+    Entry(Interface* impl,
+          RequestType request,
+          BindingSet* binding_set,
+          void* context)
+        : binding_(impl, std::move(request)),
+          binding_set_(binding_set),
+          context_(context) {
+      if (binding_set->SupportsContext())
+        binding_.AddFilter(base::MakeUnique<DispatchFilter>(this));
+      binding_.set_connection_error_handler(base::Bind(
+          &Entry::OnConnectionError, base::Unretained(this)));
     }
 
-    ~Element() {}
+   private:
+    class DispatchFilter : public MessageReceiver {
+     public:
+      explicit DispatchFilter(Entry* entry) : entry_(entry) {}
+      ~DispatchFilter() override {}
+
+     private:
+      // MessageReceiver:
+      bool Accept(Message* message) override {
+        entry_->WillDispatch();
+        return true;
+      }
 
-    void set_connection_error_handler(const base::Closure& error_handler) {
-      error_handler_ = error_handler;
-    }
+      Entry* entry_;
 
-    base::WeakPtr<Element> GetWeakPtr() {
-      return weak_ptr_factory_.GetWeakPtr();
-    }
+      DISALLOW_COPY_AND_ASSIGN(DispatchFilter);
+    };
 
-    void Close() { binding_.Close(); }
+    void WillDispatch() {
+      DCHECK(binding_set_->SupportsContext());
+      binding_set_->SetDispatchContext(context_);
+    }
 
     void OnConnectionError() {
-      base::Closure error_handler = error_handler_;
-      delete this;
-      if (!error_handler.is_null())
-        error_handler.Run();
+      if (binding_set_->SupportsContext())
+        WillDispatch();
+      binding_set_->OnConnectionError(this);
     }
 
-   private:
-    Binding<Interface> binding_;
-    base::Closure error_handler_;
-    base::WeakPtrFactory<Element> weak_ptr_factory_;
+    BindingType binding_;
+    BindingSet* const binding_set_;
+    void* const context_;
 
-    DISALLOW_COPY_AND_ASSIGN(Element);
+    DISALLOW_COPY_AND_ASSIGN(Entry);
   };
 
-  void OnConnectionError() {
-    // Clear any deleted bindings.
-    bindings_.erase(std::remove_if(bindings_.begin(), bindings_.end(),
-                                   [](const base::WeakPtr<Element>& p) {
-                                     return p.get() == nullptr;
-                                   }),
-                    bindings_.end());
+  void SetDispatchContext(void* context) {
+    DCHECK(SupportsContext());
+    dispatch_context_ = context;
+  }
+
+  bool SupportsContext() const {
+    return dispatch_mode_ == BindingSetDispatchMode::WITH_CONTEXT;
+  }
+
+  void OnConnectionError(Entry* entry) {
+    auto it = bindings_.find(entry);
+    DCHECK(it != bindings_.end());
+    bindings_.erase(it);
 
     if (!error_handler_.is_null())
       error_handler_.Run();
   }
 
+  BindingSetDispatchMode dispatch_mode_;
   base::Closure error_handler_;
-  std::vector<base::WeakPtr<Element>> bindings_;
+  std::map<Entry*, std::unique_ptr<Entry>> bindings_;
+  void* dispatch_context_ = nullptr;
 
   DISALLOW_COPY_AND_ASSIGN(BindingSet);
 };