Add MicroSpinLock TSAN annotations

Summary: Add MicroSpinLock TSAN annotations

Reviewed By: yfeldblum

Differential Revision: D14578441

fbshipit-source-id: e2c9bc1806550d003cc659a86d214c089464f126

folly/synchronization/MicroSpinLock.h

@@ -46,6 +46,7 @@
 
 #include <folly/Portability.h>
 #include <folly/lang/Align.h>
+#include <folly/synchronization/SanitizeThread.h>
 #include <folly/synchronization/detail/Sleeper.h>
 
 namespace folly {
@@ -74,21 +75,28 @@ struct MicroSpinLock {
   }
 
   bool try_lock() noexcept {
-    return cas(FREE, LOCKED);
+    bool ret = cas(FREE, LOCKED);
+    annotate_rwlock_try_acquired(
+        this, annotate_rwlock_level::wrlock, ret, __FILE__, __LINE__);
+    return ret;
   }
 
   void lock() noexcept {
     detail::Sleeper sleeper;
-    while (!try_lock()) {
+    while (!cas(FREE, LOCKED)) {
       do {
         sleeper.wait();
       } while (payload()->load(std::memory_order_relaxed) == LOCKED);
     }
     assert(payload()->load() == LOCKED);
+    annotate_rwlock_acquired(
+        this, annotate_rwlock_level::wrlock, __FILE__, __LINE__);
   }
 
   void unlock() noexcept {
     assert(payload()->load() == LOCKED);
+    annotate_rwlock_released(
+        this, annotate_rwlock_level::wrlock, __FILE__, __LINE__);
     payload()->store(FREE, std::memory_order_release);
   }
 

folly/synchronization/test/SmallLocksTest.cpp

@@ -32,6 +32,7 @@
 #include <folly/portability/GTest.h>
 #include <folly/portability/PThread.h>
 #include <folly/portability/Unistd.h>
+#include <folly/test/TestUtils.h>
 
 using folly::MicroLock;
 using folly::MicroSpinLock;
@@ -393,6 +394,57 @@ TEST(SmallLocks, MicroSpinLockStressTestTryLockHardwareConcurrency) {
   simpleStressTestTryLock<MicroSpinLock>(duration, threads);
 }
 
+TEST(SmallLocksk, MicroSpinLockThreadSanitizer) {
+  SKIP_IF(!folly::kIsSanitizeThread) << "Enabled in TSAN mode only";
+
+  uint8_t val = 0;
+  static_assert(sizeof(uint8_t) == sizeof(MicroSpinLock), "sanity check");
+  // make sure TSAN handles this case too:
+  // same lock but initialized via setting a value
+  for (int i = 0; i < 10; i++) {
+    val = 0;
+    std::lock_guard<MicroSpinLock> g(*reinterpret_cast<MicroSpinLock*>(&val));
+  }
+
+  {
+    MicroSpinLock a;
+    MicroSpinLock b;
+    a.init();
+    b.init();
+    {
+      std::lock_guard<MicroSpinLock> ga(a);
+      std::lock_guard<MicroSpinLock> gb(b);
+    }
+    {
+      std::lock_guard<MicroSpinLock> gb(b);
+      EXPECT_DEATH(
+          [&]() { std::lock_guard<MicroSpinLock> ga(a); }(),
+          "Cycle in lock order graph");
+    }
+  }
+
+  {
+    uint8_t a = 0;
+    uint8_t b = 0;
+    {
+      std::lock_guard<MicroSpinLock> ga(*reinterpret_cast<MicroSpinLock*>(&a));
+      std::lock_guard<MicroSpinLock> gb(*reinterpret_cast<MicroSpinLock*>(&b));
+    }
+
+    a = 0;
+    b = 0;
+    {
+      std::lock_guard<MicroSpinLock> gb(*reinterpret_cast<MicroSpinLock*>(&b));
+      EXPECT_DEATH(
+          [&]() {
+            std::lock_guard<MicroSpinLock> ga(
+                *reinterpret_cast<MicroSpinLock*>(&a));
+          }(),
+          "Cycle in lock order graph");
+    }
+  }
+}
+
 TEST(SmallLocks, PicoSpinLockStressTestTryLockTwoThreads) {
   auto duration = std::chrono::seconds{FLAGS_stress_test_seconds};
   simpleStressTestTryLock<PicoSpinLock<std::uint16_t>>(duration, 2);