Use Baton (again) in EventBase::runInEventBaseThreadAndWait

[folly.git] / folly / io / async / EventBaseLocal.cpp

diff --git a/folly/io/async/EventBaseLocal.cpp b/folly/io/async/EventBaseLocal.cpp
index 8dd551856471906abd516e6ae0affcaa17697f6b..cf5cf6f86e123f0f23eb399a812a5f177e152a09 100644 (file)
--- a/folly/io/async/EventBaseLocal.cpp
+++ b/folly/io/async/EventBaseLocal.cpp
@@ -15,50 +15,30 @@
  */
 
 #include <folly/io/async/EventBaseLocal.h>
+#include <folly/MapUtil.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
+  for (auto* evb : *eventBases_.rlock()) {
+    evb->runInEventBaseThread([ this, evb, key = key_ ] {
+      evb->localStorage_.erase(key);
+      evb->localStorageToDtor_.erase(this);
+    });
   }
 }
 
 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();
-  }
+  evb.dcheckIsInEventBaseThread();
 
-  return nullptr;
+  return folly::get_default(evb.localStorage_, key_, {}).get();
 }
 
 void EventBaseLocalBase::erase(EventBase& evb) {
-  std::lock_guard<std::mutex> lg(evb.localStorageMutex_);
+  evb.dcheckIsInEventBaseThread();
+
   evb.localStorage_.erase(key_);
   evb.localStorageToDtor_.erase(this);
 
@@ -68,18 +48,15 @@ void EventBaseLocalBase::erase(EventBase& evb) {
 }
 
 void EventBaseLocalBase::onEventBaseDestruction(EventBase& evb) {
+  evb.dcheckIsInEventBaseThread();
+
   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) {
+  evb.dcheckIsInEventBaseThread();
 
   auto alreadyExists =
     evb.localStorage_.find(key_) != evb.localStorage_.end();
@@ -87,9 +64,7 @@ void EventBaseLocalBase::setVoidUnlocked(
   evb.localStorage_.emplace(key_, std::move(ptr));
 
   if (!alreadyExists) {
-    SYNCHRONIZED(eventBases_) {
-      eventBases_.insert(&evb);
-    }
+    eventBases_.wlock()->insert(&evb);
     evb.localStorageToDtor_.insert(this);
   }
 }