From b529367b595dd4da7e70569cbea36d30ace4fa39 Mon Sep 17 00:00:00 2001
From: Shubhanshu Agrawal <shubhanshu@fb.com>
Date: Thu, 9 Nov 2017 17:35:30 -0800
Subject: [PATCH] adding a fibers compatible once flag

Summary:
The current folly::once_flag is not compatible with folly fibers
and when using it with fibers is inefficient and also cause deadlocks.

This diff makes the once flag's mutex be a templatable paramter
and overrides it in fibers library with a fiber compatible mtuex.

Reviewed By: yfeldblum

Differential Revision: D6288508

fbshipit-source-id: 6f82e1794d1f417f8d267061f1702a26a7b4ff12
---
 folly/fibers/CallOnce.h          | 25 +++++++++++
 folly/fibers/TimedMutex.h        |  2 +-
 folly/synchronization/CallOnce.h | 73 ++++++++++++++++++++------------
 3 files changed, 71 insertions(+), 29 deletions(-)
 create mode 100644 folly/fibers/CallOnce.h

diff --git a/folly/fibers/CallOnce.h b/folly/fibers/CallOnce.h
new file mode 100644
index 00000000..7b720208
--- /dev/null
+++ b/folly/fibers/CallOnce.h
@@ -0,0 +1,25 @@
+/*
+ * Copyright 2017 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 "folly/fibers/TimedMutex.h"
+#include "folly/synchronization/CallOnce.h"
+
+namespace folly {
+namespace fibers {
+
+using once_flag = ::folly::detail::once_flag<folly::fibers::TimedMutex>;
+}
+} // namespace folly
diff --git a/folly/fibers/TimedMutex.h b/folly/fibers/TimedMutex.h
index f965abe0..bec18b27 100644
--- a/folly/fibers/TimedMutex.h
+++ b/folly/fibers/TimedMutex.h
@@ -29,7 +29,7 @@ namespace fibers {
  **/
 class TimedMutex {
  public:
-  TimedMutex() {}
+  TimedMutex() noexcept {}
 
   ~TimedMutex() {
     DCHECK(threadWaiters_.empty());
diff --git a/folly/synchronization/CallOnce.h b/folly/synchronization/CallOnce.h
index da938beb..4d4bf0f4 100644
--- a/folly/synchronization/CallOnce.h
+++ b/folly/synchronization/CallOnce.h
@@ -39,28 +39,32 @@
 #include <folly/SharedMutex.h>
 
 namespace folly {
+namespace detail {
+template <typename Mutex>
+class once_flag;
 
-class once_flag {
- public:
-  constexpr once_flag() noexcept = default;
-  once_flag(const once_flag&) = delete;
-  once_flag& operator=(const once_flag&) = delete;
+// Implementation detail: out-of-line slow path
+template <class Mutex, class Callable, class... Args>
+void FOLLY_NOINLINE call_once_impl_no_inline(
+    detail::once_flag<Mutex>& flag,
+    Callable&& f,
+    Args&&... args) {
+  std::lock_guard<Mutex> lg(flag.mutex_);
+  if (flag.called_) {
+    return;
+  }
 
-  template <typename Callable, class... Args>
-  friend void call_once(once_flag& flag, Callable&& f, Args&&... args);
-  template <typename Callable, class... Args>
-  friend void call_once_impl_no_inline(once_flag& flag,
-                                       Callable&& f,
-                                       Args&&... args);
+  std::forward<Callable>(f)(std::forward<Args>(args)...);
 
- private:
-  std::atomic<bool> called_{false};
-  folly::SharedMutex mutex_;
-};
+  flag.called_.store(true, std::memory_order_release);
+}
+} // namespace detail
 
-template <class Callable, class... Args>
+using once_flag = detail::once_flag<folly::SharedMutex>;
+
+template <class Mutex, class Callable, class... Args>
 void FOLLY_ALWAYS_INLINE
-call_once(once_flag& flag, Callable&& f, Args&&... args) {
+call_once(detail::once_flag<Mutex>& flag, Callable&& f, Args&&... args) {
   if (LIKELY(flag.called_.load(std::memory_order_acquire))) {
     return;
   }
@@ -68,17 +72,30 @@ call_once(once_flag& flag, Callable&& f, Args&&... args) {
       flag, std::forward<Callable>(f), std::forward<Args>(args)...);
 }
 
-// Implementation detail: out-of-line slow path
-template <class Callable, class... Args>
-void FOLLY_NOINLINE
-call_once_impl_no_inline(once_flag& flag, Callable&& f, Args&&... args) {
-  std::lock_guard<folly::SharedMutex> lg(flag.mutex_);
-  if (flag.called_) {
-    return;
-  }
+namespace detail {
 
-  std::forward<Callable>(f)(std::forward<Args>(args)...);
+template <typename Mutex>
+class once_flag {
+ public:
+  constexpr once_flag() noexcept = default;
+  once_flag(const once_flag&) = delete;
+  once_flag& operator=(const once_flag&) = delete;
+
+  template <typename Mutex_, typename Callable, class... Args>
+  friend void ::folly::call_once(
+      once_flag<Mutex_>& flag,
+      Callable&& f,
+      Args&&... args);
+  template <typename Mutex_, typename Callable, class... Args>
+  friend void call_once_impl_no_inline(
+      once_flag<Mutex_>& flag,
+      Callable&& f,
+      Args&&... args);
+
+ private:
+  std::atomic<bool> called_{false};
+  Mutex mutex_;
+};
+} // namespace detail
 
-  flag.called_.store(true, std::memory_order_release);
-}
 } // namespace folly
-- 
2.34.1