From a0c640e8e2ff45bc43b9f8d96cd108bb66635bee Mon Sep 17 00:00:00 2001
From: Dave Watson <davejwatson@fb.com>
Date: Wed, 26 Apr 2017 10:04:04 -0700
Subject: [PATCH] Make folly pass TSAN checks

Summary:
Currently, contbuild has a blanket TSAN suppression for folly.
Fix PicoSpinLock instead

This should fix TSAN errors as an alternative to D4781776

Some of the tests even had TSAN errors, fixed those.

Reviewed By: davidtgoldblatt

Differential Revision: D4795284

fbshipit-source-id: 9f0fc6868399da2f86be355ce3c081990260a649
---
 folly/PicoSpinLock.h          | 26 ++++++++++++++++++++------
 folly/test/SmallLocksTest.cpp |  6 +++---
 folly/test/SpinLockTest.cpp   |  7 +++++--
 3 files changed, 28 insertions(+), 11 deletions(-)

diff --git a/folly/PicoSpinLock.h b/folly/PicoSpinLock.h
index 68385880..8ee98b84 100644
--- a/folly/PicoSpinLock.h
+++ b/folly/PicoSpinLock.h
@@ -46,8 +46,8 @@
 #include <mutex>
 #include <type_traits>
 
-#include <glog/logging.h>
 #include <folly/detail/Sleeper.h>
+#include <glog/logging.h>
 
 #if !FOLLY_X64 && !FOLLY_A64 && !FOLLY_PPC64
 # error "PicoSpinLock.h is currently x64, aarch64 and ppc64 only."
@@ -81,7 +81,7 @@ struct PicoSpinLock {
 
  public:
   static const UIntType kLockBitMask_ = UIntType(1) << Bit;
-  UIntType lock_;
+  mutable UIntType lock_;
 
   /*
    * You must call this function before using this class, if you
@@ -93,7 +93,8 @@ struct PicoSpinLock {
    */
   void init(IntType initialValue = 0) {
     CHECK(!(initialValue & kLockBitMask_));
-    lock_ = UIntType(initialValue);
+    reinterpret_cast<std::atomic<UIntType>*>(&lock_)->store(
+        UIntType(initialValue), std::memory_order_release);
   }
 
   /*
@@ -106,7 +107,10 @@ struct PicoSpinLock {
    * as you normally get.)
    */
   IntType getData() const {
-    return static_cast<IntType>(lock_ & ~kLockBitMask_);
+    auto res = reinterpret_cast<std::atomic<UIntType>*>(&lock_)->load(
+                   std::memory_order_relaxed) &
+        ~kLockBitMask_;
+    return res;
   }
 
   /*
@@ -117,7 +121,10 @@ struct PicoSpinLock {
    */
   void setData(IntType w) {
     CHECK(!(w & kLockBitMask_));
-    lock_ = UIntType((lock_ & kLockBitMask_) | w);
+    auto l = reinterpret_cast<std::atomic<UIntType>*>(&lock_);
+    l->store(
+        (l->load(std::memory_order_relaxed) & kLockBitMask_) | w,
+        std::memory_order_relaxed);
   }
 
   /*
@@ -127,7 +134,14 @@ struct PicoSpinLock {
   bool try_lock() const {
     bool ret = false;
 
-#ifdef _MSC_VER
+#if defined(FOLLY_SANITIZE_THREAD)
+    // TODO: Might be able to fully move to std::atomic when gcc emits lock btr:
+    // https://gcc.gnu.org/bugzilla/show_bug.cgi?id=49244
+    ret =
+        !(reinterpret_cast<std::atomic<UIntType>*>(&lock_)->fetch_or(
+              kLockBitMask_, std::memory_order_acquire) &
+          kLockBitMask_);
+#elif _MSC_VER
     switch (sizeof(IntType)) {
       case 2:
         // There is no _interlockedbittestandset16 for some reason :(
diff --git a/folly/test/SmallLocksTest.cpp b/folly/test/SmallLocksTest.cpp
index 1618b9f9..9da8b2bd 100644
--- a/folly/test/SmallLocksTest.cpp
+++ b/folly/test/SmallLocksTest.cpp
@@ -93,7 +93,7 @@ template<class T> struct PslTest {
   void doTest() {
     using UT = typename std::make_unsigned<T>::type;
     T ourVal = rand() % T(UT(1) << (sizeof(UT) * 8 - 1));
-    for (int i = 0; i < 10000; ++i) {
+    for (int i = 0; i < 100; ++i) {
       std::lock_guard<PicoSpinLock<T>> guard(lock);
       lock.setData(ourVal);
       for (int n = 0; n < 10; ++n) {
@@ -231,9 +231,9 @@ TEST(SmallLocks, MicroLock) {
   // affect bits outside the ones MicroLock is defined to affect.
   struct {
     uint8_t a;
-    volatile uint8_t b;
+    std::atomic<uint8_t> b;
     MicroLock alock;
-    volatile uint8_t d;
+    std::atomic<uint8_t> d;
   } x;
 
   uint8_t origB = 'b';
diff --git a/folly/test/SpinLockTest.cpp b/folly/test/SpinLockTest.cpp
index b2387ea7..67279a70 100644
--- a/folly/test/SpinLockTest.cpp
+++ b/folly/test/SpinLockTest.cpp
@@ -68,16 +68,19 @@ template <typename LOCK>
 void trylockTestThread(TryLockState<LOCK>* state, size_t count) {
   while (true) {
     folly::asm_pause();
+    bool ret = state->lock2.try_lock();
     SpinLockGuardImpl<LOCK> g(state->lock1);
     if (state->obtained >= count) {
+      if (ret) {
+        state->lock2.unlock();
+      }
       break;
     }
 
-    bool ret = state->lock2.try_lock();
-    EXPECT_NE(state->locked, ret);
 
     if (ret) {
       // We got lock2.
+      EXPECT_NE(state->locked, ret);
       ++state->obtained;
       state->locked = true;
 
-- 
2.34.1