add EventBase-local storage abstraction

Summary: This has come up a couple times, and the implementation is never nice.  Just like we have thread-local storage, it will be useful to also have evb-local storage.  Provides at pretty simple get/set/delete interface (see unittest).

Reviewed By: @djwatson

Differential Revision: D2203063

folly/Makefile.am

@@ -199,6 +199,7 @@ nobase_follyinclude_HEADERS = \
 	io/async/DelayedDestructionBase.h \
 	io/async/DelayedDestruction.h \
 	io/async/EventBase.h \
+	io/async/EventBaseLocal.h \
 	io/async/EventBaseManager.h \
 	io/async/EventFDWrapper.h \
 	io/async/EventHandler.h \
@@ -340,6 +341,7 @@ libfolly_la_SOURCES = \
 	io/async/AsyncSocket.cpp \
 	io/async/AsyncSSLSocket.cpp \
 	io/async/EventBase.cpp \
+	io/async/EventBaseLocal.cpp \
 	io/async/EventBaseManager.cpp \
 	io/async/EventHandler.cpp \
 	io/async/SSLContext.cpp \

folly/io/async/EventBase.cpp

@@ -21,6 +21,7 @@
 #include <folly/io/async/EventBase.h>
 
 #include <folly/ThreadName.h>
+#include <folly/io/async/EventBaseLocal.h>
 #include <folly/io/async/NotificationQueue.h>
 
 #include <boost/static_assert.hpp>
@@ -237,6 +238,13 @@ EventBase::~EventBase() {
     std::lock_guard<std::mutex> lock(libevent_mutex_);
     event_base_free(evb_);
   }
+
+  {
+    std::lock_guard<std::mutex> lock(localStorageMutex_);
+    for (auto storage : localStorageToDtor_) {
+      storage->onEventBaseDestruction(*this);
+    }
+  }
   VLOG(5) << "EventBase(): Destroyed.";
 }
 

folly/io/async/EventBase.h

@@ -29,6 +29,9 @@
 #include <queue>
 #include <cstdlib>
 #include <set>
+#include <unordered_set>
+#include <unordered_map>
+#include <mutex>
 #include <utility>
 #include <boost/intrusive/list.hpp>
 #include <boost/utility.hpp>
@@ -44,6 +47,12 @@ typedef std::function<void()> Cob;
 template <typename MessageT>
 class NotificationQueue;
 
+namespace detail {
+class EventBaseLocalBase;
+}
+template <typename T>
+class EventBaseLocal;
+
 class EventBaseObserver {
  public:
   virtual ~EventBaseObserver() = default;
@@ -402,11 +411,11 @@ class EventBase : private boost::noncopyable,
     return runImmediatelyOrRunInEventBaseThreadAndWait(std::bind(fn, arg));
   }
 
-    /*
+  /*
    * Like runInEventBaseThreadAndWait, except if the caller is already in the
    * event base thread, the functor is simply run inline.
    */
-bool runImmediatelyOrRunInEventBaseThreadAndWait(const Cob& fn);
+  bool runImmediatelyOrRunInEventBaseThreadAndWait(const Cob& fn);
 
   /**
    * Runs the given Cob at some time after the specified number of
@@ -728,6 +737,13 @@ bool runImmediatelyOrRunInEventBaseThreadAndWait(const Cob& fn);
 
   // allow runOnDestruction() to be called from any threads
   std::mutex onDestructionCallbacksMutex_;
+
+  // see EventBaseLocal
+  friend class detail::EventBaseLocalBase;
+  template <typename T> friend class EventBaseLocal;
+  std::mutex localStorageMutex_;
+  std::unordered_map<uint64_t, std::shared_ptr<void>> localStorage_;
+  std::unordered_set<detail::EventBaseLocalBase*> localStorageToDtor_;
 };
 
 } // folly

folly/io/async/EventBaseLocal.cpp

+/*
+ * Copyright 2015 Facebook, Inc.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include <folly/io/async/EventBaseLocal.h>
+#include <atomic>
+#include <thread>
+
+namespace folly { namespace detail {
+
+EventBaseLocalBase::~EventBaseLocalBase() {
+  // There's a race condition if an EventBase and an EventBaseLocal destruct
+  // at the same time (each will lock eventBases_ and localStorageMutex_
+  // in the opposite order), so we dance around it with a loop and try_lock.
+  while (true) {
+    SYNCHRONIZED(eventBases_) {
+      auto it = eventBases_.begin();
+      while (it != eventBases_.end()) {
+        auto evb = *it;
+        if (evb->localStorageMutex_.try_lock()) {
+          evb->localStorage_.erase(key_);
+          evb->localStorageToDtor_.erase(this);
+          it = eventBases_.erase(it);
+          evb->localStorageMutex_.unlock();
+        } else {
+          ++it;
+        }
+      }
+
+      if (eventBases_.empty()) {
+        return;
+      }
+    }
+    std::this_thread::yield(); // let the other thread take the eventBases_ lock
+  }
+}
+
+void* EventBaseLocalBase::getVoid(EventBase& evb) {
+  std::lock_guard<std::mutex> lg(evb.localStorageMutex_);
+  auto it2 = evb.localStorage_.find(key_);
+  if (UNLIKELY(it2 != evb.localStorage_.end())) {
+    return it2->second.get();
+  }
+
+  return nullptr;
+}
+
+void EventBaseLocalBase::erase(EventBase& evb) {
+  std::lock_guard<std::mutex> lg(evb.localStorageMutex_);
+  evb.localStorage_.erase(key_);
+  evb.localStorageToDtor_.erase(this);
+
+  SYNCHRONIZED(eventBases_) {
+    eventBases_.erase(&evb);
+  }
+}
+
+void EventBaseLocalBase::onEventBaseDestruction(EventBase& evb) {
+  SYNCHRONIZED(eventBases_) {
+    eventBases_.erase(&evb);
+  }
+}
+
+void EventBaseLocalBase::setVoid(EventBase& evb, std::shared_ptr<void>&& ptr) {
+  std::lock_guard<std::mutex> lg(evb.localStorageMutex_);
+  setVoidUnlocked(evb, std::move(ptr));
+}
+
+void EventBaseLocalBase::setVoidUnlocked(
+    EventBase& evb, std::shared_ptr<void>&& ptr) {
+
+  auto alreadyExists =
+    evb.localStorage_.find(key_) != evb.localStorage_.end();
+
+  evb.localStorage_.emplace(key_, std::move(ptr));
+
+  if (!alreadyExists) {
+    SYNCHRONIZED(eventBases_) {
+      eventBases_.insert(&evb);
+    }
+    evb.localStorageToDtor_.insert(this);
+  }
+}
+
+std::atomic<uint64_t> EventBaseLocalBase::keyCounter_{0};
+}}

folly/io/async/EventBaseLocal.h

+/*
+ * Copyright 2015 Facebook, Inc.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#pragma once
+
+#include <boost/noncopyable.hpp>
+#include <folly/Synchronized.h>
+#include <folly/io/async/EventBase.h>
+#include <memory>
+#include <mutex>
+#include <unordered_map>
+#include <unordered_set>
+
+namespace folly {
+
+namespace detail {
+
+class EventBaseLocalBase : boost::noncopyable {
+ public:
+  EventBaseLocalBase() {}
+  virtual ~EventBaseLocalBase();
+  void erase(EventBase& evb);
+  void onEventBaseDestruction(EventBase& evb);
+
+ protected:
+  void setVoid(EventBase& evb, std::shared_ptr<void>&& ptr);
+  void setVoidUnlocked(EventBase& evb, std::shared_ptr<void>&& ptr);
+  void* getVoid(EventBase& evb);
+
+  folly::Synchronized<std::unordered_set<EventBase*>> eventBases_;
+  static std::atomic<uint64_t> keyCounter_;
+  uint64_t key_{keyCounter_++};
+};
+
+}
+
+/**
+ * A storage abstraction for data that should be tied to an EventBase.
+ *
+ *   struct Foo { Foo(int a, int b); };
+ *   EventBaseLocal<Foo> myFoo;
+ *   ...
+ *   EventBase evb;
+ *   myFoo.set(evb, new Foo(1, 2));
+ *   myFoo.set(evb, 1, 2);
+ *   Foo* foo = myFoo.get(evb);
+ *   myFoo.erase(evb);
+ *   Foo& foo = myFoo.getOrCreate(evb, 1, 2); // ctor
+ *   Foo& foo = myFoo.getOrCreate(evb, 1, 2); // no ctor
+ *   myFoo.erase(evb);
+ *   Foo& foo = myFoo.getOrCreateFn(evb, [] () { return new Foo(3, 4); })
+ *
+ * The objects will be deleted when the EventBaseLocal or the EventBase is
+ * destructed (whichever comes first).  All methods are thread-safe.
+ *
+ * The user is responsible for throwing away invalid references/ptrs returned
+ * by the get() method after set/erase is called.  If shared ownership is
+ * needed, use a EventBaseLocal<shared_ptr<...>>.
+ */
+template<typename T>
+class EventBaseLocal : public detail::EventBaseLocalBase {
+ public:
+  EventBaseLocal(): EventBaseLocalBase() {}
+
+  T* get(EventBase& evb) {
+    return static_cast<T*>(getVoid(evb));
+  }
+
+  void emplace(EventBase& evb, T* ptr) {
+    std::shared_ptr<T> smartPtr(ptr);
+    setVoid(evb, std::move(smartPtr));
+  }
+
+  template<typename... Args>
+  void emplace(EventBase& evb, Args... args) {
+    auto smartPtr = std::make_shared<T>(args...);
+    setVoid(evb, smartPtr);
+  }
+
+  template<typename... Args>
+  T& getOrCreate(EventBase& evb, Args... args) {
+    std::lock_guard<std::mutex> lg(evb.localStorageMutex_);
+
+    auto it2 = evb.localStorage_.find(key_);
+    if (LIKELY(it2 != evb.localStorage_.end())) {
+      return *static_cast<T*>(it2->second.get());
+    } else {
+      auto smartPtr = std::make_shared<T>(args...);
+      auto ptr = smartPtr.get();
+      setVoidUnlocked(evb, std::move(smartPtr));
+      return *ptr;
+    }
+  }
+
+  template <typename Func>
+  T& getOrCreateFn(EventBase& evb, Func& fn) {
+    // If this looks like it's copy/pasted from above, that's because it is.
+    // gcc has a bug (fixed in 4.9) that doesn't allow capturing variadic
+    // params in a lambda.
+    std::lock_guard<std::mutex> lg(evb.localStorageMutex_);
+
+    auto it2 = evb.localStorage_.find(key_);
+    if (LIKELY(it2 != evb.localStorage_.end())) {
+      return *static_cast<T*>(it2->second.get());
+    } else {
+      std::shared_ptr<T> smartPtr(fn());
+      auto ptr = smartPtr.get();
+      setVoidUnlocked(evb, std::move(smartPtr));
+      return *ptr;
+    }
+  }
+};
+
+
+}

folly/io/async/test/EventBaseLocalTest.cpp

+/*
+ * Copyright 2015 Facebook, Inc.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+#include <folly/io/async/EventBaseLocal.h>
+#include <folly/io/async/test/Util.h>
+#include <gtest/gtest.h>
+
+struct Foo {
+  Foo(int n, std::function<void()> dtorFn):
+    n(n), dtorFn(std::move(dtorFn)) {}
+  ~Foo() { dtorFn(); }
+
+  int n;
+  std::function<void()> dtorFn;
+};
+
+TEST(EventBaseLocalTest, Basic) {
+  int dtorCnt = 0;
+  folly::EventBase evb1;
+
+  {
+    folly::EventBaseLocal<Foo> foo;
+
+    EXPECT_EQ(foo.get(evb1), nullptr);
+
+    foo.emplace(evb1, new Foo(5, [&] () { ++dtorCnt; }));
+
+    EXPECT_EQ(foo.get(evb1)->n, 5);
+
+    {
+      folly::EventBase evb2;
+      foo.emplace(evb2, new Foo(6, [&] () { ++dtorCnt; }));
+      EXPECT_EQ(foo.get(evb2)->n, 6);
+      foo.erase(evb2);
+      EXPECT_EQ(dtorCnt, 1); // should dtor a Foo when we erase
+      EXPECT_EQ(foo.get(evb2), nullptr);
+      foo.emplace(evb2, 7, [&] () { ++dtorCnt; });
+      EXPECT_EQ(foo.get(evb2)->n, 7);
+    }
+
+    EXPECT_EQ(dtorCnt, 2); // should dtor a Foo when evb2 destructs
+
+  }
+  EXPECT_EQ(dtorCnt, 3); // should dtor a Foo when foo destructs
+}
+
+TEST(EventBaseLocalTest, getOrCreate) {
+  folly::EventBase evb1;
+  folly::EventBaseLocal<int> ints;
+
+  EXPECT_EQ(ints.getOrCreate(evb1), 0);
+  EXPECT_EQ(ints.getOrCreate(evb1, 5), 0);
+
+  folly::EventBase evb2;
+  EXPECT_EQ(ints.getOrCreate(evb2, 5), 5);
+  ints.erase(evb2);
+  auto creator = []() { return new int(4); };
+  EXPECT_EQ(ints.getOrCreateFn(evb2, creator), 4);
+}