From cb2fe8d01f9b3eb746a852cc96bc930287d95f68 Mon Sep 17 00:00:00 2001
From: Yedidya Feldblum <yfeldblum@fb.com>
Date: Mon, 30 Mar 2015 18:17:31 -0700
Subject: [PATCH] EventBase::runImediatelyOrRunInEventBaseThreadAndWait.

Summary:
[Folly] EventBase::runImediatelyOrRunInEventBaseThreadAndWait.

When you have code that you need to run in the event loop and wait for it to finish, and you yourself might be in the event loop but might be in another thread instead.

Test Plan:
Unit tests:
* `folly/io/async/test/EventBaseTest.cpp`.

Reviewed By: davejwatson@fb.com

Subscribers: folly-diffs@, yfeldblum, chalfant, agallagher, njormrod, ldbrandy, brettp, dougw

FB internal diff: D1953110

Signature: t1:1953110:1427751076:ad2c9f4d5b65ef952cba94b06be887672ee44583
---
 folly/io/async/EventBase.cpp          |  9 ++++++++
 folly/io/async/EventBase.h            | 25 +++++++++++++++++++++
 folly/io/async/test/EventBaseTest.cpp | 32 ++++++++++++++++++++++++++-
 3 files changed, 65 insertions(+), 1 deletion(-)

diff --git a/folly/io/async/EventBase.cpp b/folly/io/async/EventBase.cpp
index 8bc09523..5ff171eb 100644
--- a/folly/io/async/EventBase.cpp
+++ b/folly/io/async/EventBase.cpp
@@ -610,6 +610,15 @@ bool EventBase::runInEventBaseThreadAndWait(const Cob& fn) {
   return true;
 }
 
+bool EventBase::runImmediatelyOrRunInEventBaseThreadAndWait(const Cob& fn) {
+  if (inRunningEventBaseThread()) {
+    fn();
+    return true;
+  } else {
+    return runInEventBaseThreadAndWait(fn);
+  }
+}
+
 void EventBase::runAfterDelay(const Cob& cob,
                               int milliseconds,
                               TimeoutManager::InternalEnum in) {
diff --git a/folly/io/async/EventBase.h b/folly/io/async/EventBase.h
index ca895c3a..0bdc428a 100644
--- a/folly/io/async/EventBase.h
+++ b/folly/io/async/EventBase.h
@@ -382,6 +382,31 @@ class EventBase : private boost::noncopyable,
    */
   bool runInEventBaseThreadAndWait(const Cob& fn);
 
+  /*
+   * Like runInEventBaseThreadAndWait, except if the caller is already in the
+   * event base thread, the functor is simply run inline.
+   */
+  template<typename T>
+  bool runImmediatelyOrRunInEventBaseThreadAndWait(void (*fn)(T*), T* arg) {
+    return runImmediatelyOrRunInEventBaseThreadAndWait(
+        reinterpret_cast<void (*)(void*)>(fn), reinterpret_cast<void*>(arg));
+  }
+
+  /*
+   * Like runInEventBaseThreadAndWait, except if the caller is already in the
+   * event base thread, the functor is simply run inline.
+   */
+  bool runImmediatelyOrRunInEventBaseThreadAndWait(
+      void (*fn)(void*), void* arg) {
+    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);
+
   /**
    * Runs the given Cob at some time after the specified number of
    * milliseconds.  (No guarantees exactly when.)
diff --git a/folly/io/async/test/EventBaseTest.cpp b/folly/io/async/test/EventBaseTest.cpp
index a5fbcad5..b9716a7c 100644
--- a/folly/io/async/test/EventBaseTest.cpp
+++ b/folly/io/async/test/EventBaseTest.cpp
@@ -1180,7 +1180,7 @@ TEST(EventBaseTest, RunInThread) {
 //  This test simulates some calls, and verifies that the waiting happens by
 //  triggering what otherwise would be race conditions, and trying to detect
 //  whether any of the race conditions happened.
-TEST(EventBaseTest, RunInEventLoopThreadAndWait) {
+TEST(EventBaseTest, RunInEventBaseThreadAndWait) {
   const size_t c = 256;
   vector<unique_ptr<atomic<size_t>>> atoms(c);
   for (size_t i = 0; i < c; ++i) {
@@ -1216,6 +1216,36 @@ TEST(EventBaseTest, RunInEventLoopThreadAndWait) {
   EXPECT_EQ(c, sum);
 }
 
+TEST(EventBaseTest, RunImmediatelyOrRunInEventBaseThreadAndWaitCross) {
+  EventBase eb;
+  thread th(&EventBase::loopForever, &eb);
+  SCOPE_EXIT {
+    eb.terminateLoopSoon();
+    th.join();
+  };
+  auto mutated = false;
+  eb.runImmediatelyOrRunInEventBaseThreadAndWait([&] {
+      mutated = true;
+  });
+  EXPECT_TRUE(mutated);
+}
+
+TEST(EventBaseTest, RunImmediatelyOrRunInEventBaseThreadAndWaitWithin) {
+  EventBase eb;
+  thread th(&EventBase::loopForever, &eb);
+  SCOPE_EXIT {
+    eb.terminateLoopSoon();
+    th.join();
+  };
+  eb.runInEventBaseThreadAndWait([&] {
+      auto mutated = false;
+      eb.runImmediatelyOrRunInEventBaseThreadAndWait([&] {
+          mutated = true;
+      });
+      EXPECT_TRUE(mutated);
+  });
+}
+
 ///////////////////////////////////////////////////////////////////////////
 // Tests for runInLoop()
 ///////////////////////////////////////////////////////////////////////////
-- 
2.34.1