Stop using atomic_thread_fence in folly::Baton

Summary: Thread fences are not supported by TSAN

Reviewed By: yfeldblum

Differential Revision: D21989519

fbshipit-source-id: bb657c82a64abe025e8f4c030916c0cacc3c0b81

folly/synchronization/AtomicUtil-inl.h

@@ -62,28 +62,26 @@ constexpr std::memory_order atomic_compare_exchange_succ(
 
 } // namespace detail
 
-template <typename T>
+template <template <typename> class Atom, typename T>
 bool atomic_compare_exchange_weak_explicit(
-    std::atomic<T>* obj,
+    Atom<T>* obj,
     T* expected,
     T desired,
     std::memory_order succ,
     std::memory_order fail) {
   succ = detail::atomic_compare_exchange_succ(succ, fail);
-  return std::atomic_compare_exchange_weak_explicit(
-      obj, expected, desired, succ, fail);
+  return obj->compare_exchange_weak(*expected, desired, succ, fail);
 }
 
-template <typename T>
+template <template <typename> class Atom, typename T>
 bool atomic_compare_exchange_strong_explicit(
-    std::atomic<T>* obj,
+    Atom<T>* obj,
     T* expected,
     T desired,
     std::memory_order succ,
     std::memory_order fail) {
   succ = detail::atomic_compare_exchange_succ(succ, fail);
-  return std::atomic_compare_exchange_strong_explicit(
-      obj, expected, desired, succ, fail);
+  return obj->compare_exchange_strong(*expected, desired, succ, fail);
 }
 
 namespace detail {

folly/synchronization/AtomicUtil.h

@@ -27,9 +27,9 @@ namespace folly {
 //  Workaround for https://github.com/google/sanitizers/issues/970.
 //
 //  mimic: std::atomic_compare_exchange_weak
-template <typename T>
+template <template <typename> class Atom = std::atomic, typename T>
 bool atomic_compare_exchange_weak_explicit(
-    std::atomic<T>* obj,
+    Atom<T>* obj,
     T* expected,
     T desired,
     std::memory_order succ,
@@ -41,9 +41,9 @@ bool atomic_compare_exchange_weak_explicit(
 //  Workaround for https://github.com/google/sanitizers/issues/970.
 //
 //  mimic: std::atomic_compare_exchange_strong
-template <typename T>
+template <template <typename> class Atom = std::atomic, typename T>
 bool atomic_compare_exchange_strong_explicit(
-    std::atomic<T>* obj,
+    Atom<T>* obj,
     T* expected,
     T desired,
     std::memory_order succ,

folly/synchronization/Baton.h

@@ -27,6 +27,7 @@
 #include <folly/detail/Futex.h>
 #include <folly/detail/MemoryIdler.h>
 #include <folly/portability/Asm.h>
+#include <folly/synchronization/AtomicUtil.h>
 #include <folly/synchronization/WaitOptions.h>
 #include <folly/synchronization/detail/Spin.h>
 
@@ -292,15 +293,14 @@ class Baton {
 
     // guess we have to block :(
     uint32_t expected = INIT;
-    if (!state_.compare_exchange_strong(
-            expected,
-            WAITING,
+    if (!folly::atomic_compare_exchange_strong_explicit<Atom>(
+            &state_,
+            &expected,
+            static_cast<uint32_t>(WAITING),
             std::memory_order_relaxed,
-            std::memory_order_relaxed)) {
+            std::memory_order_acquire)) {
       // CAS failed, last minute reprieve
       assert(expected == EARLY_DELIVERY);
-      // TODO: move the acquire to the compare_exchange failure load after C++17
-      std::atomic_thread_fence(std::memory_order_acquire);
       return true;
     }