Fix MicroLock data storage when it is not word/half-word aligned

Summary:
The current data storage API code assumes that the MicroLock is 4-byte aligned and that the lock lives in the lowest byte of `word()`. This assumption is incorrect (see added test). We should instead rely on `baseShift()`.

This API is only used for `folly::compact_once_flag` which is new so low chance something broke.

Reviewed By: andriigrynenko, davidtgoldblatt

Differential Revision: D26091484

fbshipit-source-id: c6e6e9aa7a963cada16e0ef7783fcc6bf6d2aed6

folly/MicroLock.cpp

@@ -24,11 +24,12 @@ namespace folly {
 uint8_t MicroLockCore::lockSlowPath(
     uint32_t oldWord,
     detail::Futex<>* wordPtr,
-    uint32_t heldBit,
+    unsigned baseShift,
     unsigned maxSpins,
     unsigned maxYields) noexcept {
   uint32_t newWord;
   unsigned spins = 0;
+  uint32_t heldBit = 1 << baseShift;
   uint32_t waitBit = heldBit << 1;
   uint32_t needWaitBit = 0;
 
@@ -69,6 +70,6 @@ retry:
           std::memory_order_relaxed)) {
     goto retry;
   }
-  return decodeDataFromWord(newWord);
+  return decodeDataFromWord(newWord, baseShift);
 }
 } // namespace folly

folly/MicroLock.h

@@ -111,7 +111,7 @@ class MicroLockCore {
   static uint8_t lockSlowPath(
       uint32_t oldWord,
       detail::Futex<>* wordPtr,
-      uint32_t heldBit,
+      unsigned baseShift,
       unsigned maxSpins,
       unsigned maxYields) noexcept;
 
@@ -121,13 +121,6 @@ class MicroLockCore {
    * the lock will be modified.
    */
   detail::Futex<>* word() const noexcept;
-  /**
-   * Extract the lock byte from word() value.
-   */
-  static constexpr uint8_t byteFromWord(uint32_t word) noexcept {
-    return kIsLittleEndian ? static_cast<uint8_t>(word & 0xff)
-                           : static_cast<uint8_t>((word >> 24) & 0xff);
-  }
 
   static constexpr unsigned kNumLockBits = 2;
   static constexpr uint8_t kLockBits =
@@ -146,8 +139,14 @@ class MicroLockCore {
     return static_cast<uint8_t>(data << kNumLockBits);
   }
 
-  static constexpr uint8_t decodeDataFromWord(uint32_t word) noexcept {
-    return decodeDataFromByte(byteFromWord(word));
+  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 uint32_t encodeDataToWord(
       uint32_t word,
@@ -358,7 +357,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, heldBit(), MaxSpins, MaxYields);
+    return lockSlowPath(oldWord, wordPtr, baseShift(), MaxSpins, MaxYields);
   }
 }
 

folly/synchronization/test/SmallLocksTest.cpp

@@ -358,6 +358,20 @@ TEST(SmallLocks, MicroLockWithData) {
   EXPECT_EQ(lock.load(std::memory_order_relaxed), 0b00101110);
 }
 
+TEST(SmallLocks, MicroLockDataAlignment) {
+  struct alignas(uint32_t) Thing {
+    uint8_t padding1[2];
+    MicroLock lock;
+    uint8_t padding2;
+  } thing;
+  auto& lock = thing.lock;
+  lock.init();
+
+  EXPECT_EQ(lock.lockAndLoad(), 0);
+  lock.unlockAndStore(60);
+  EXPECT_EQ(lock.load(std::memory_order_relaxed), 60);
+}
+
 namespace {
 template <typename Mutex, typename Duration>
 void simpleStressTest(Duration duration, int numThreads) {