Fixed documentation in folly/Random.h

[folly.git] / folly / PicoSpinLock.h

diff --git a/folly/PicoSpinLock.h b/folly/PicoSpinLock.h
index 0a0f28c213bb1dc4b438b6a4174900153c62162e..485050f136c9ee895af8e28a3d2111583a49b9da 100644 (file)
--- a/folly/PicoSpinLock.h
+++ b/folly/PicoSpinLock.h
@@ -1,5 +1,5 @@
 /*
- * Copyright 2016 Facebook, Inc.
+ * Copyright 2017 Facebook, Inc.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -16,7 +16,7 @@
 
 /*
  * N.B. You most likely do _not_ want to use PicoSpinLock or any other
- * kind of spinlock.
+ * kind of spinlock.  Consider MicroLock instead.
  *
  * In short, spinlocks in preemptive multi-tasking operating systems
  * have serious problems and fast mutexes like std::mutex are almost
@@ -31,6 +31,7 @@
  */
 
 #pragma once
+#define FOLLY_PICO_SPIN_LOCK_H_
 
 /*
  * @author Keith Adams <kma@fb.com>
@@ -38,16 +39,15 @@
  */
 
 #include <array>
+#include <atomic>
 #include <cinttypes>
-#include <type_traits>
 #include <cstdlib>
-#include <pthread.h>
+#include <folly/Portability.h>
 #include <mutex>
-#include <atomic>
+#include <type_traits>
 
-#include <glog/logging.h>
 #include <folly/detail/Sleeper.h>
-#include <folly/Portability.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_ = 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_ = (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,28 @@ struct PicoSpinLock {
   bool try_lock() const {
     bool ret = false;
 
-#if FOLLY_X64
+#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 :(
+        ret = _InterlockedOr16(
+            (volatile short*)&lock_, (short)kLockBitMask_) & kLockBitMask_;
+        break;
+      case 4:
+        ret = _interlockedbittestandset((volatile long*)&lock_, Bit);
+        break;
+      case 8:
+        ret = _interlockedbittestandset64((volatile long long*)&lock_, Bit);
+        break;
+    }
+#elif FOLLY_X64
 #define FB_DOBTS(size)                                  \
   asm volatile("lock; bts" #size " %1, (%2); setnc %0"  \
                : "=r" (ret)                             \
@@ -143,7 +171,9 @@ struct PicoSpinLock {
 
 #undef FB_DOBTS
 #elif FOLLY_A64
-    ret = __atomic_fetch_or(&lock_, 1 << Bit, __ATOMIC_SEQ_CST);
+    ret =
+        !(__atomic_fetch_or(&lock_, kLockBitMask_, __ATOMIC_SEQ_CST) &
+          kLockBitMask_);
 #elif FOLLY_PPC64
 #define FB_DOBTS(size)                                 \
     asm volatile("\teieio\n"                           \
@@ -192,7 +222,20 @@ struct PicoSpinLock {
    * integer.
    */
   void unlock() const {
-#if FOLLY_X64
+#ifdef _MSC_VER
+    switch (sizeof(IntType)) {
+      case 2:
+        // There is no _interlockedbittestandreset16 for some reason :(
+        _InterlockedAnd16((volatile short*)&lock_, (short)~kLockBitMask_);
+        break;
+      case 4:
+        _interlockedbittestandreset((volatile long*)&lock_, Bit);
+        break;
+      case 8:
+        _interlockedbittestandreset64((volatile long long*)&lock_, Bit);
+        break;
+    }
+#elif FOLLY_X64
 #define FB_DOBTR(size)                          \
   asm volatile("lock; btr" #size " %0, (%1)"    \
                :                                \
@@ -211,7 +254,7 @@ struct PicoSpinLock {
 
 #undef FB_DOBTR
 #elif FOLLY_A64
-    __atomic_fetch_and(&lock_, ~(1 << Bit), __ATOMIC_SEQ_CST);
+    __atomic_fetch_and(&lock_, ~kLockBitMask_, __ATOMIC_SEQ_CST);
 #elif FOLLY_PPC64
 #define FB_DOBTR(size)                                 \
     asm volatile("\teieio\n"                           \