Extract compare-exchange hack under TSAN

Summary:
[Folly] Extract compare-exchange hack under TSAN to a common location and make its interface mimic `std::atomic_compare_exchange_strong_explicit`. Limit it to Clang TSAN (no knowledge of whether it applies to other implementations of TSAN), but extend it to handle all pairs of success and failure orders.

The TSAN bug is described in https://github.com/google/sanitizers/issues/970. High level, Clang TSAN ignores the explicit failure order and infers a failure order from the success order - which is broken if the explicit failure order is in the relevant sense stronger than the success order.

Reviewed By: aary, nbronson

Differential Revision: D16854217

fbshipit-source-id: 18f6458520bbd5f482e41c10d7229e6cfae1db2a

folly/futures/detail/Core.h

@@ -33,6 +33,7 @@
 #include <folly/futures/detail/Types.h>
 #include <folly/lang/Assume.h>
 #include <folly/lang/Exception.h>
+#include <folly/synchronization/AtomicUtil.h>
 #include <folly/synchronization/MicroSpinLock.h>
 #include <glog/logging.h>
 
@@ -74,20 +75,6 @@ struct SpinLock : private MicroSpinLock {
 };
 static_assert(sizeof(SpinLock) == 1, "missized");
 
-template <typename T>
-bool compare_exchange_strong_release_acquire(
-    std::atomic<T>& self,
-    T& expected,
-    T desired) {
-  if (kIsSanitizeThread) {
-    // Workaround for https://github.com/google/sanitizers/issues/970
-    return self.compare_exchange_strong(
-        expected, desired, std::memory_order_acq_rel);
-  }
-  return self.compare_exchange_strong(
-      expected, desired, std::memory_order_release, std::memory_order_acquire);
-}
-
 class DeferredExecutor;
 
 class UniqueDeleter {
@@ -608,8 +595,12 @@ class Core final {
         : State::OnlyCallback;
 
     if (state == State::Start) {
-      if (detail::compare_exchange_strong_release_acquire(
-              state_, state, nextState)) {
+      if (atomic_compare_exchange_strong_explicit(
+              &state_,
+              &state,
+              nextState,
+              std::memory_order_release,
+              std::memory_order_acquire)) {
         return;
       }
       assume(state == State::OnlyResult || state == State::Proxy);
@@ -642,8 +633,12 @@ class Core final {
     auto state = state_.load(std::memory_order_acquire);
     switch (state) {
       case State::Start:
-        if (detail::compare_exchange_strong_release_acquire(
-                state_, state, State::Proxy)) {
+        if (atomic_compare_exchange_strong_explicit(
+                &state_,
+                &state,
+                State::Proxy,
+                std::memory_order_release,
+                std::memory_order_acquire)) {
           break;
         }
         assume(
@@ -692,8 +687,12 @@ class Core final {
     auto state = state_.load(std::memory_order_acquire);
     switch (state) {
       case State::Start:
-        if (detail::compare_exchange_strong_release_acquire(
-                state_, state, State::OnlyResult)) {
+        if (atomic_compare_exchange_strong_explicit(
+                &state_,
+                &state,
+                State::OnlyResult,
+                std::memory_order_release,
+                std::memory_order_acquire)) {
           return;
         }
         assume(

folly/synchronization/AtomicUtil-inl.h

@@ -30,6 +30,62 @@
 #endif
 
 namespace folly {
+
+namespace detail {
+
+constexpr std::memory_order atomic_compare_exchange_succ(
+    bool cond,
+    std::memory_order succ,
+    std::memory_order fail) {
+  constexpr auto const relaxed = std::memory_order_relaxed;
+  constexpr auto const release = std::memory_order_release;
+  constexpr auto const acq_rel = std::memory_order_acq_rel;
+
+  assert(fail != release);
+  assert(fail != acq_rel);
+
+  //  Clang TSAN ignores the passed failure order and infers failure order from
+  //  success order in atomic compare-exchange operations, which is broken for
+  //  cases like success-release/failure-acquire, so return a success order with
+  //  the failure order mixed in.
+  auto const bump = succ == release ? acq_rel : succ;
+  auto const high = fail < bump ? bump : fail;
+  return !cond || fail == relaxed ? succ : high;
+}
+
+constexpr std::memory_order atomic_compare_exchange_succ(
+    std::memory_order succ,
+    std::memory_order fail) {
+  constexpr auto const cond = kIsSanitizeThread && kIsClang;
+  return atomic_compare_exchange_succ(cond, succ, fail);
+}
+
+} // namespace detail
+
+template <typename T>
+bool atomic_compare_exchange_weak_explicit(
+    std::atomic<T>* obj,
+    typename std::atomic<T>::value_type* expected,
+    typename std::atomic<T>::value_type 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);
+}
+
+template <typename T>
+bool atomic_compare_exchange_strong_explicit(
+    std::atomic<T>* obj,
+    typename std::atomic<T>::value_type* expected,
+    typename std::atomic<T>::value_type 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);
+}
+
 namespace detail {
 
 // TODO: Remove the non-default implementations when both gcc and clang

folly/synchronization/AtomicUtil.h

@@ -21,6 +21,34 @@
 
 namespace folly {
 
+//  atomic_compare_exchange_weak_explicit
+//
+//  Fix TSAN bug in std::atomic_compare_exchange_weak_explicit.
+//  Workaround for https://github.com/google/sanitizers/issues/970.
+//
+//  mimic: std::atomic_compare_exchange_weak
+template <typename T>
+bool atomic_compare_exchange_weak_explicit(
+    std::atomic<T>* obj,
+    typename std::atomic<T>::value_type* expected,
+    typename std::atomic<T>::value_type desired,
+    std::memory_order succ,
+    std::memory_order fail);
+
+//  atomic_compare_exchange_strong_explicit
+//
+//  Fix TSAN bug in std::atomic_compare_exchange_strong_explicit.
+//  Workaround for https://github.com/google/sanitizers/issues/970.
+//
+//  mimic: std::atomic_compare_exchange_strong
+template <typename T>
+bool atomic_compare_exchange_weak_explicit(
+    std::atomic<T>* obj,
+    typename std::atomic<T>::value_type* expected,
+    typename std::atomic<T>::value_type desired,
+    std::memory_order succ,
+    std::memory_order fail);
+
 //  atomic_fetch_set
 //
 //  Sets the bit at the given index in the binary representation of the integer

folly/synchronization/test/AtomicUtilTest.cpp

@@ -25,6 +25,85 @@
 
 namespace folly {
 
+class AtomicCompareExchangeSuccTest : public testing::Test {};
+
+TEST_F(AtomicCompareExchangeSuccTest, examples) {
+  using detail::atomic_compare_exchange_succ;
+
+  auto const relaxed = std::memory_order_relaxed;
+  auto const consume = std::memory_order_consume;
+  auto const acquire = std::memory_order_acquire;
+  auto const release = std::memory_order_release;
+  auto const acq_rel = std::memory_order_acq_rel;
+  auto const seq_cst = std::memory_order_seq_cst;
+
+  // noop table
+  EXPECT_EQ(relaxed, atomic_compare_exchange_succ(false, relaxed, relaxed));
+  EXPECT_EQ(consume, atomic_compare_exchange_succ(false, consume, relaxed));
+  EXPECT_EQ(acquire, atomic_compare_exchange_succ(false, acquire, relaxed));
+  EXPECT_EQ(release, atomic_compare_exchange_succ(false, release, relaxed));
+  EXPECT_EQ(acq_rel, atomic_compare_exchange_succ(false, acq_rel, relaxed));
+  EXPECT_EQ(seq_cst, atomic_compare_exchange_succ(false, seq_cst, relaxed));
+  EXPECT_EQ(relaxed, atomic_compare_exchange_succ(false, relaxed, consume));
+  EXPECT_EQ(consume, atomic_compare_exchange_succ(false, consume, consume));
+  EXPECT_EQ(acquire, atomic_compare_exchange_succ(false, acquire, consume));
+  EXPECT_EQ(release, atomic_compare_exchange_succ(false, release, consume));
+  EXPECT_EQ(acq_rel, atomic_compare_exchange_succ(false, acq_rel, consume));
+  EXPECT_EQ(seq_cst, atomic_compare_exchange_succ(false, seq_cst, consume));
+  EXPECT_EQ(relaxed, atomic_compare_exchange_succ(false, relaxed, acquire));
+  EXPECT_EQ(consume, atomic_compare_exchange_succ(false, consume, acquire));
+  EXPECT_EQ(acquire, atomic_compare_exchange_succ(false, acquire, acquire));
+  EXPECT_EQ(release, atomic_compare_exchange_succ(false, release, acquire));
+  EXPECT_EQ(acq_rel, atomic_compare_exchange_succ(false, acq_rel, acquire));
+  EXPECT_EQ(seq_cst, atomic_compare_exchange_succ(false, seq_cst, acquire));
+  EXPECT_EQ(relaxed, atomic_compare_exchange_succ(false, relaxed, seq_cst));
+  EXPECT_EQ(consume, atomic_compare_exchange_succ(false, consume, seq_cst));
+  EXPECT_EQ(acquire, atomic_compare_exchange_succ(false, acquire, seq_cst));
+  EXPECT_EQ(release, atomic_compare_exchange_succ(false, release, seq_cst));
+  EXPECT_EQ(acq_rel, atomic_compare_exchange_succ(false, acq_rel, seq_cst));
+  EXPECT_EQ(seq_cst, atomic_compare_exchange_succ(false, seq_cst, seq_cst));
+
+  // xform table
+  EXPECT_EQ(relaxed, atomic_compare_exchange_succ(true, relaxed, relaxed));
+  EXPECT_EQ(consume, atomic_compare_exchange_succ(true, consume, relaxed));
+  EXPECT_EQ(acquire, atomic_compare_exchange_succ(true, acquire, relaxed));
+  EXPECT_EQ(release, atomic_compare_exchange_succ(true, release, relaxed));
+  EXPECT_EQ(acq_rel, atomic_compare_exchange_succ(true, acq_rel, relaxed));
+  EXPECT_EQ(seq_cst, atomic_compare_exchange_succ(true, seq_cst, relaxed));
+  EXPECT_EQ(consume, atomic_compare_exchange_succ(true, relaxed, consume));
+  EXPECT_EQ(consume, atomic_compare_exchange_succ(true, consume, consume));
+  EXPECT_EQ(acquire, atomic_compare_exchange_succ(true, acquire, consume));
+  EXPECT_EQ(acq_rel, atomic_compare_exchange_succ(true, release, consume));
+  EXPECT_EQ(acq_rel, atomic_compare_exchange_succ(true, acq_rel, consume));
+  EXPECT_EQ(seq_cst, atomic_compare_exchange_succ(true, seq_cst, consume));
+  EXPECT_EQ(acquire, atomic_compare_exchange_succ(true, relaxed, acquire));
+  EXPECT_EQ(acquire, atomic_compare_exchange_succ(true, consume, acquire));
+  EXPECT_EQ(acquire, atomic_compare_exchange_succ(true, acquire, acquire));
+  EXPECT_EQ(acq_rel, atomic_compare_exchange_succ(true, release, acquire));
+  EXPECT_EQ(acq_rel, atomic_compare_exchange_succ(true, acq_rel, acquire));
+  EXPECT_EQ(seq_cst, atomic_compare_exchange_succ(true, seq_cst, acquire));
+  EXPECT_EQ(seq_cst, atomic_compare_exchange_succ(true, relaxed, seq_cst));
+  EXPECT_EQ(seq_cst, atomic_compare_exchange_succ(true, consume, seq_cst));
+  EXPECT_EQ(seq_cst, atomic_compare_exchange_succ(true, acquire, seq_cst));
+  EXPECT_EQ(seq_cst, atomic_compare_exchange_succ(true, release, seq_cst));
+  EXPECT_EQ(seq_cst, atomic_compare_exchange_succ(true, acq_rel, seq_cst));
+  EXPECT_EQ(seq_cst, atomic_compare_exchange_succ(true, seq_cst, seq_cst));
+
+  // properties
+  for (auto succ : {relaxed, consume, acquire, release, acq_rel, seq_cst}) {
+    SCOPED_TRACE(succ);
+    for (auto fail : {relaxed, consume, acquire, seq_cst}) {
+      SCOPED_TRACE(fail);
+      EXPECT_EQ(succ, atomic_compare_exchange_succ(false, succ, fail));
+      auto const sfix = atomic_compare_exchange_succ(true, succ, fail);
+      EXPECT_GE(sfix, succ);
+      EXPECT_GE(sfix, fail);
+      EXPECT_TRUE(fail != relaxed || sfix == succ);
+      EXPECT_TRUE(fail == relaxed || succ != release || sfix != release);
+    }
+  }
+}
+
 namespace {
 auto default_fetch_set = [](auto&&... args) {
   return atomic_fetch_set(args...);