Better spinlock implementation

src/_profiler.cc

@@ -3,12 +3,12 @@
 
 #include <cstdlib>
 #include <memory>
-#include "third_party/google/tcmalloc/sampler.h"
 
 #include "heap.h"
 #include "log.h"
 #include "malloc_patch.h"
 #include "scoped_object.h"
+#include "third_party/google/tcmalloc/sampler.h"
 
 namespace {
 

src/heap.h

@@ -5,16 +5,16 @@
 #define MPROFILE_SRC_HEAP_H_
 
 #include <stdlib.h>
+
 #include <mutex>
 #include <vector>
 
+#include "spinlock.h"
+#include "stacktraces.h"
 #include "third_party/google/tcmalloc/addressmap.h"
 #include "third_party/google/tcmalloc/sampler.h"
 #include "third_party/greg7mdp/parallel-hashmap/phmap.h"
 
-#include "spinlock.h"
-#include "stacktraces.h"
-
 class HeapProfiler {
  public:
   HeapProfiler() : HeapProfiler(kMaxFramesToCapture) {}

src/heap_bench.cc

@@ -1,7 +1,6 @@
 // Copyright 2019 Timothy Palpant
 
 #include "benchmark/benchmark.h"
-
 #include "heap.h"
 
 static void BM_HandleMalloc(benchmark::State &state) {

src/heap_test.cc

@@ -1,9 +1,10 @@
 // Copyright 2019 Timothy Palpant
 //
-#include "gtest/gtest.h"
+#include "heap.h"
 
 #include <thread>
-#include "heap.h"
+
+#include "gtest/gtest.h"
 
 TEST(HeapProfiler, HandleMalloc) {
   Sampler::SetSamplePeriod(0);

src/log.cc

@@ -15,6 +15,7 @@
 #include "log.h"
 
 #include <Python.h>
+
 #include <iostream>
 
 void Log(const char *level, const char *fmt, ...) {

src/main_bench.cc

@@ -1,6 +1,7 @@
 // Copyright 2019 Timothy Palpant
 
 #include <Python.h>
+
 #include "benchmark/benchmark.h"
 
 int main(int argc, char** argv) {

src/main_test.cc

@@ -2,6 +2,7 @@
 //
 
 #include <Python.h>
+
 #include "gtest/gtest.h"
 
 int main(int argc, char **argv) {

src/malloc_patch.h

@@ -4,6 +4,7 @@
 #define MPROFILE_SRC_MALLOC_PATCH_H_
 
 #include <Python.h>
+
 #include <memory>
 
 #include "heap.h"

src/scoped_object.h

@@ -17,6 +17,7 @@
 #define MPROFILE_SRC_SCOPED_OBJECT_H_
 
 #include <Python.h>
+
 #include <memory>
 
 struct PyObjectDecReffer {

src/spinlock.h

@@ -5,8 +5,10 @@
 
 #include <atomic>
 
-// SpinLock is a simple adapter to use an std::atomic_flag
+// SpinLock is a simple adapter to use an std::atomic<bool>
 // with a std::lock_guard.
+//
+// Adapted from: https://rigtorp.se/spinlock/
 class SpinLock {
  public:
   SpinLock() {}
@@ -15,14 +17,25 @@ class SpinLock {
   SpinLock &operator=(const SpinLock &) = delete;
 
   void lock() {
-    while (flag_.test_and_set(std::memory_order_acquire))
-      ;
+    while (true) {
+      // Optimistically assume the lock is free on the first try.
+      if (!flag_.exchange(true, std::memory_order_acquire)) {
+        break;
+      }
+
+      // Wait for lock to be released without generating cache misses.
+      while (flag_.load(std::memory_order_relaxed)) {
+        // Issue X86 PAUSE or ARM YIELD instruction to reduce contention
+        // between hyper-threads.
+        __builtin_ia32_pause();
+      }
+    }
   }
 
-  void unlock() { flag_.clear(std::memory_order_release); }
+  void unlock() { flag_.store(false, std::memory_order_release); }
 
  private:
-  std::atomic_flag flag_ = ATOMIC_FLAG_INIT;
+  std::atomic<bool> flag_ = {false};
 };
 
 #endif  // MPROFILE_SRC_SPINLOCK_H

src/stacktraces.h

@@ -17,6 +17,7 @@
 #define MPROFILE_SRC_STACKTRACES_H_
 
 #include <Python.h>
+
 #include <vector>
 
 #include "third_party/greg7mdp/parallel-hashmap/phmap.h"

src/stacktraces_test.cc

@@ -1,8 +1,8 @@
 // Copyright 2019 Timothy Palpant
-#include "gtest/gtest.h"
+#include "stacktraces.h"
 
+#include "gtest/gtest.h"
 #include "scoped_object.h"
-#include "stacktraces.h"
 
 #if PY_MAJOR_VERSION >= 3
 #define STRING_FROMSTRING PyUnicode_FromString