revise MicroLock not to access bytes outside of itself

Summary:
The `MicroLock` tests fail under libc++ with asan-abort `stack-buffer-overflow`.

Example:
```
SCARINESS: 32 (4-byte-read-stack-buffer-overflow)
    #0 unsigned int std::__2::__cxx_atomic_load[abi:ue170004]<unsigned int>(std::__2::__cxx_atomic_base_impl<unsigned int> const*, std::__2::memory_order) libcxx/include/c++/v1/__atomic/cxx_atomic_impl.h:375
    #1 std::__2::__atomic_base<unsigned int, false>::load[abi:ue170004](std::__2::memory_order) const libcxx/include/c++/v1/__atomic/atomic_base.h:60
    facebook#2 folly::MicroLockBase<1000u, 0u>::try_lock() folly/MicroLock.h:319
    facebook#3 SmallLocks_MicroLockTryLock_Test::TestBody() folly/synchronization/test/SmallLocksTest.cpp:314
```

The trouble is accessing the lock state as a 4-byte word, where the asan stack has allocated only 1 byte for the lock object - asan tracks this and aborts.

It is UB for `MicroLock` to perform atomic stores or atomic read-modify-write operations on the entire 4-byte aligned word containing the `MicroLock` which may be concurrent with other accesses to other bytes in the word. The 4-byte aligned word accesses are necessary in order to support using the `MicroLock` itself as a futex, on Linux platforms. The trouble is that Linux does not provide futex operations on all sizes of integer, only on 32-bit (4-byte) integers. If Linux were to provide futex operations on all sizes of integer, we would immediately switch to 8-bit (1-byte) integers here.

This diff fixes the UB and gets the tests passing under libc++ by switching to the `folly::ParkingLot`, as wrapped by `folly::atomic_notify_one` and `folly::atomic_wait` over `std::atomic<uint8_t>`, which indirects the 32-bit (4-byte) futex behind a sharded hashtable.

Alternatives are possible for getting the tests passing, although these alternatives would not fix the UB.
* We can disable sanitizers on even more functions, as well as switch to atomic builtins since the members of `std::atomic` are not marked as disabling the sanitizers.
* We can workaround just in the unit-tests by padding the on-stack `MicroLock` instances.

Reviewed By: praihan, ot, aary

Differential Revision: D76612658

fbshipit-source-id: 75e9f03e63c85c30c6201b9bce5cd303c9a2b530

Author: yfeldblum

third-party/folly/src/folly/MicroLock.cpp

@@ -23,16 +23,12 @@
 namespace folly {
 
 uint8_t MicroLockCore::lockSlowPath(
-    uint32_t oldWord,
-    detail::Futex<>* wordPtr,
-    unsigned baseShift,
-    unsigned maxSpins,
-    unsigned maxYields) noexcept {
-  uint32_t newWord;
+    uint8_t oldWord, unsigned maxSpins, unsigned maxYields) noexcept {
+  uint8_t newWord;
   unsigned spins = 0;
-  uint32_t heldBit = 1 << baseShift;
-  uint32_t waitBit = heldBit << 1;
-  uint32_t needWaitBit = 0;
+  uint8_t heldBit = 1;
+  uint8_t waitBit = heldBit << 1;
+  uint8_t needWaitBit = 0;
 
 retry:
   if ((oldWord & heldBit) != 0) {
@@ -43,34 +39,34 @@ uint8_t MicroLockCore::lockSlowPath(
       // lock holder knows to FUTEX_WAKE us.
       newWord = oldWord | waitBit;
       if (newWord != oldWord) {
-        if (!wordPtr->compare_exchange_weak(
+        if (!atomic_ref(lock_).compare_exchange_weak(
                 oldWord,
                 newWord,
                 std::memory_order_relaxed,
                 std::memory_order_relaxed)) {
           goto retry;
         }
       }
-      detail::futexWait(wordPtr, newWord, heldBit);
+      atomic_wait(&atomic_ref(lock_).atomic(), newWord);
       needWaitBit = waitBit;
     } else if (spins > maxSpins) {
       // sched_yield(), but more portable
       std::this_thread::yield();
     } else {
       folly::asm_volatile_pause();
     }
-    oldWord = wordPtr->load(std::memory_order_relaxed);
+    oldWord = atomic_ref(lock_).load(std::memory_order_relaxed);
     goto retry;
   }
 
   newWord = oldWord | heldBit | needWaitBit;
-  if (!wordPtr->compare_exchange_weak(
+  if (!atomic_ref(lock_).compare_exchange_weak(
           oldWord,
           newWord,
           std::memory_order_acquire,
           std::memory_order_relaxed)) {
     goto retry;
   }
-  return decodeDataFromWord(newWord, baseShift);
+  return decodeDataFromWord(newWord);
 }
 } // namespace folly

third-party/folly/src/folly/MicroLock.h

@@ -21,10 +21,10 @@
 #include <cstdint>
 #include <utility>
 
-#include <folly/CPortability.h>
 #include <folly/Optional.h>
 #include <folly/Portability.h>
-#include <folly/detail/Futex.h>
+#include <folly/synchronization/AtomicNotification.h>
+#include <folly/synchronization/AtomicRef.h>
 
 namespace folly {
 
@@ -51,12 +51,6 @@ namespace folly {
  * Unused bits in the lock can be used to store user data via
  * lockAndLoad() and unlockAndStore(), or LockGuardWithData.
  *
- * MicroLock uses a dirty trick: it actually operates on the full
- * 32-bit, four-byte-aligned bit of memory into which it is embedded.
- * It never modifies bits outside the ones it's defined to modify, but
- * it _accesses_ all the bits in the 32-bit memory location for
- * purposes of futex management.
- *
  * The MaxSpins template parameter controls the number of times we
  * spin trying to acquire the lock.  MaxYields controls the number of
  * times we call sched_yield; once we've tried to acquire the lock
@@ -96,10 +90,6 @@ namespace folly {
 class MicroLockCore {
  protected:
   uint8_t lock_{};
-  /**
-   * Arithmetic shift required to get to the byte from the word.
-   */
-  unsigned baseShift() const noexcept;
   /**
    * Mask for bit indicating that the flag is held.
    */
@@ -109,19 +99,8 @@ class MicroLockCore {
    */
   unsigned waitBit() const noexcept;
 
-  FOLLY_DISABLE_SANITIZERS static uint8_t lockSlowPath(
-      uint32_t oldWord,
-      detail::Futex<>* wordPtr,
-      unsigned baseShift,
-      unsigned maxSpins,
-      unsigned maxYields) noexcept;
-
-  /**
-   * The word (halfword on 64-bit systems) that this lock atomically operates
-   * on. Although the atomic operations access 4 bytes, only the byte used by
-   * the lock will be modified.
-   */
-  detail::Futex<>* word() const noexcept;
+  uint8_t lockSlowPath(
+      uint8_t oldWord, unsigned maxSpins, unsigned maxYields) noexcept;
 
   static constexpr unsigned kNumLockBits = 2;
   static constexpr uint8_t kLockBits =
@@ -140,40 +119,34 @@ class MicroLockCore {
     return static_cast<uint8_t>(data << kNumLockBits);
   }
 
-  static constexpr uint8_t decodeDataFromWord(
-      uint32_t word, unsigned baseShift) noexcept {
-    return static_cast<uint8_t>(
-        static_cast<uint8_t>(word >> baseShift) >> kNumLockBits);
-  }
-  uint8_t decodeDataFromWord(uint32_t word) const noexcept {
-    return decodeDataFromWord(word, baseShift());
+  static constexpr uint8_t decodeDataFromWord(uint8_t word) noexcept {
+    return static_cast<uint8_t>(word >> kNumLockBits);
   }
-  static constexpr uint32_t encodeDataToWord(
-      uint32_t word, unsigned shiftToByte, uint8_t value) noexcept {
-    const uint32_t preservedBits = word & ~(kDataBits << shiftToByte);
-    const uint32_t newBits = encodeDataToByte(value) << shiftToByte;
+  static constexpr uint8_t encodeDataToWord(
+      uint8_t word, uint8_t value) noexcept {
+    const uint8_t preservedBits = word & ~(kDataBits);
+    const uint8_t newBits = encodeDataToByte(value);
     return preservedBits | newBits;
   }
 
   template <typename Func>
-  FOLLY_DISABLE_SANITIZERS void unlockAndStoreWithModifier(
-      Func modifier) noexcept;
+  void unlockAndStoreWithModifier(Func modifier) noexcept;
 
  public:
   /**
    * Loads the data stored in the unused bits of the lock atomically.
    */
-  FOLLY_DISABLE_SANITIZERS uint8_t
-  load(std::memory_order order = std::memory_order_seq_cst) const noexcept {
-    return decodeDataFromWord(word()->load(order));
+  uint8_t load(
+      std::memory_order order = std::memory_order_seq_cst) const noexcept {
+    return decodeDataFromWord(atomic_ref(lock_).load(order));
   }
 
   /**
    * Stores the data in the unused bits of the lock atomically. Since 2 bits are
    * used by the lock, the most significant 2 bits of the provided value will be
    * ignored.
    */
-  FOLLY_DISABLE_SANITIZERS void store(
+  void store(
       uint8_t value,
       std::memory_order order = std::memory_order_seq_cst) noexcept;
 
@@ -186,73 +159,51 @@ class MicroLockCore {
   void unlock() noexcept;
 };
 
-inline detail::Futex<>* MicroLockCore::word() const noexcept {
-  uintptr_t lockptr = (uintptr_t)&lock_;
-  lockptr &= ~(sizeof(uint32_t) - 1);
-  return (detail::Futex<>*)lockptr;
-}
-
-inline unsigned MicroLockCore::baseShift() const noexcept {
-  unsigned offset_bytes = (unsigned)((uintptr_t)&lock_ - (uintptr_t)word());
-
-  return static_cast<unsigned>(
-      kIsLittleEndian ? CHAR_BIT * offset_bytes
-                      : CHAR_BIT * (sizeof(uint32_t) - offset_bytes - 1));
-}
-
 inline unsigned MicroLockCore::heldBit() const noexcept {
-  return 1U << (baseShift() + 0);
+  return 1U << 0;
 }
 
 inline unsigned MicroLockCore::waitBit() const noexcept {
-  return 1U << (baseShift() + 1);
+  return 1U << 1;
 }
 
-FOLLY_PUSH_WARNING
-FOLLY_GCC_DISABLE_WARNING("-Wattributes") // inline + [[gnu::noinline]]
 inline void MicroLockCore::store(
     uint8_t value, std::memory_order order) noexcept {
-  detail::Futex<>* wordPtr = word();
-
-  const auto shiftToByte = baseShift();
-  auto oldWord = wordPtr->load(std::memory_order_relaxed);
+  auto oldWord = atomic_ref(lock_).load(std::memory_order_relaxed);
   while (true) {
-    auto newWord = encodeDataToWord(oldWord, shiftToByte, value);
-    if (wordPtr->compare_exchange_weak(
+    auto newWord = encodeDataToWord(oldWord, value);
+    if (atomic_ref(lock_).compare_exchange_weak(
             oldWord, newWord, order, std::memory_order_relaxed)) {
       break;
     }
   }
 }
-FOLLY_POP_WARNING
 
 template <typename Func>
 void MicroLockCore::unlockAndStoreWithModifier(Func modifier) noexcept {
-  detail::Futex<>* wordPtr = word();
-  uint32_t oldWord;
-  uint32_t newWord;
+  uint8_t oldWord;
+  uint8_t newWord;
 
-  oldWord = wordPtr->load(std::memory_order_relaxed);
+  oldWord = atomic_ref(lock_).load(std::memory_order_relaxed);
   do {
     assert(oldWord & heldBit());
     newWord = modifier(oldWord) & ~(heldBit() | waitBit());
-  } while (!wordPtr->compare_exchange_weak(
+  } while (!atomic_ref(lock_).compare_exchange_weak(
       oldWord, newWord, std::memory_order_release, std::memory_order_relaxed));
 
   if (oldWord & waitBit()) {
-    detail::futexWake(wordPtr, 1, heldBit());
+    atomic_notify_one(&atomic_ref(lock_).atomic());
   }
 }
 
 inline void MicroLockCore::unlockAndStore(uint8_t value) noexcept {
-  unlockAndStoreWithModifier(
-      [value, shiftToByte = baseShift()](uint32_t oldWord) {
-        return encodeDataToWord(oldWord, shiftToByte, value);
-      });
+  unlockAndStoreWithModifier([value](uint8_t oldWord) {
+    return encodeDataToWord(oldWord, value);
+  });
 }
 
 inline void MicroLockCore::unlock() noexcept {
-  unlockAndStoreWithModifier([](uint32_t oldWord) { return oldWord; });
+  unlockAndStoreWithModifier(identity);
 }
 
 template <unsigned MaxSpins = 1000, unsigned MaxYields = 0>
@@ -264,9 +215,9 @@ class MicroLockBase : public MicroLockCore {
    * data, in which case reading and locking should be done in one atomic
    * operation.
    */
-  FOLLY_DISABLE_SANITIZERS uint8_t lockAndLoad() noexcept;
+  uint8_t lockAndLoad() noexcept;
   void lock() noexcept { lockAndLoad(); }
-  FOLLY_DISABLE_SANITIZERS bool try_lock() noexcept;
+  bool try_lock() noexcept;
 
   /**
    * A lock guard which allows reading and writing to the unused bits of the
@@ -315,13 +266,12 @@ bool MicroLockBase<MaxSpins, MaxYields>::try_lock() noexcept {
   // else has the lock (by looking at heldBit) or see our CAS succeed.
   // A failed CAS by itself does not indicate lock-acquire failure.
 
-  detail::Futex<>* wordPtr = word();
-  uint32_t oldWord = wordPtr->load(std::memory_order_relaxed);
+  uint8_t oldWord = atomic_ref(lock_).load(std::memory_order_relaxed);
   do {
     if (oldWord & heldBit()) {
       return false;
     }
-  } while (!wordPtr->compare_exchange_weak(
+  } while (!atomic_ref(lock_).compare_exchange_weak(
       oldWord,
       oldWord | heldBit(),
       std::memory_order_acquire,
@@ -334,11 +284,10 @@ template <unsigned MaxSpins, unsigned MaxYields>
 uint8_t MicroLockBase<MaxSpins, MaxYields>::lockAndLoad() noexcept {
   static_assert(MaxSpins + MaxYields < (unsigned)-1, "overflow");
 
-  detail::Futex<>* wordPtr = word();
-  uint32_t oldWord;
-  oldWord = wordPtr->load(std::memory_order_relaxed);
+  uint8_t oldWord;
+  oldWord = atomic_ref(lock_).load(std::memory_order_relaxed);
   if ((oldWord & heldBit()) == 0 &&
-      wordPtr->compare_exchange_weak(
+      atomic_ref(lock_).compare_exchange_weak(
           oldWord,
           oldWord | heldBit(),
           std::memory_order_acquire,
@@ -350,7 +299,7 @@ uint8_t MicroLockBase<MaxSpins, MaxYields>::lockAndLoad() noexcept {
     // sure its shifting produces the same result a call to waitBit would.
     assert(heldBit() << 1 == waitBit());
     // lockSlowPath emits its own memory barrier
-    return lockSlowPath(oldWord, wordPtr, baseShift(), MaxSpins, MaxYields);
+    return lockSlowPath(oldWord, MaxSpins, MaxYields);
   }
 }