[Sketch] Rtsan/intercept more apple locks sanitizer dense map

compiler-rt/lib/rtsan/CMakeLists.txt

@@ -28,10 +28,10 @@ set(RTSAN_HEADERS
 set(RTSAN_DEPS)
 
 set(RTSAN_CFLAGS
-  ${COMPILER_RT_COMMON_CFLAGS}
+  ${SANITIZER_COMMON_CFLAGS}
   ${COMPILER_RT_CXX_CFLAGS}
   -DSANITIZER_COMMON_NO_REDEFINE_BUILTINS)
-set(RTSAN_LINK_FLAGS ${COMPILER_RT_COMMON_LINK_FLAGS})
+set(RTSAN_LINK_FLAGS ${SANITIZER_COMMON_LINK_FLAGS})
 set(RTSAN_DYNAMIC_LIBS
   ${COMPILER_RT_UNWINDER_LINK_LIBS}
   ${SANITIZER_CXX_ABI_LIBRARIES}

compiler-rt/lib/rtsan/rtsan.cpp

@@ -114,19 +114,19 @@ SANITIZER_INTERFACE_ATTRIBUTE bool __rtsan_is_initialized() {
 }
 
 SANITIZER_INTERFACE_ATTRIBUTE void __rtsan_realtime_enter() {
-  GetContextForThisThread().RealtimePush();
+  GetContext().RealtimePush();
 }
 
 SANITIZER_INTERFACE_ATTRIBUTE void __rtsan_realtime_exit() {
-  GetContextForThisThread().RealtimePop();
+  GetContext().RealtimePop();
 }
 
 SANITIZER_INTERFACE_ATTRIBUTE void __rtsan_disable() {
-  GetContextForThisThread().BypassPush();
+  GetContext().BypassPush();
 }
 
 SANITIZER_INTERFACE_ATTRIBUTE void __rtsan_enable() {
-  GetContextForThisThread().BypassPop();
+  GetContext().BypassPop();
 }
 
 SANITIZER_INTERFACE_ATTRIBUTE void
@@ -137,7 +137,7 @@ __rtsan_notify_intercepted_call(const char *func_name) {
 
   __rtsan_ensure_initialized();
   GET_CALLER_PC_BP;
-  ExpectNotRealtime(GetContextForThisThread(),
+  ExpectNotRealtime(GetContext(),
                     {DiagnosticsInfoType::InterceptedCall, func_name, pc, bp},
                     OnViolation);
 }
@@ -146,7 +146,7 @@ SANITIZER_INTERFACE_ATTRIBUTE void
 __rtsan_notify_blocking_call(const char *func_name) {
   __rtsan_ensure_initialized();
   GET_CALLER_PC_BP;
-  ExpectNotRealtime(GetContextForThisThread(),
+  ExpectNotRealtime(GetContext(),
                     {DiagnosticsInfoType::BlockingCall, func_name, pc, bp},
                     OnViolation);
 }

compiler-rt/lib/rtsan/rtsan_context.cpp

@@ -9,55 +9,68 @@
 //===----------------------------------------------------------------------===//
 
 #include "rtsan/rtsan_context.h"
-#include "rtsan/rtsan.h"
 
 #include "sanitizer_common/sanitizer_allocator_internal.h"
+#include "sanitizer_common/sanitizer_dense_map.h"
 
-#include <new>
 #include <pthread.h>
+#include <stddef.h>
+#include <stdio.h>
 
 using namespace __sanitizer;
 using namespace __rtsan;
+inline void *operator new(size_t, void *ptr) noexcept { return ptr; }
 
-static pthread_key_t context_key;
-static pthread_once_t key_once = PTHREAD_ONCE_INIT;
+static Context *context = nullptr;
 
-// InternalFree cannot be passed directly to pthread_key_create
-// because it expects a signature with only one arg
-static void InternalFreeWrapper(void *ptr) { __sanitizer::InternalFree(ptr); }
-
-static __rtsan::Context &GetContextForThisThreadImpl() {
-  auto MakeThreadLocalContextKey = []() {
-    CHECK_EQ(pthread_key_create(&context_key, InternalFreeWrapper), 0);
-  };
-
-  pthread_once(&key_once, MakeThreadLocalContextKey);
-  Context *current_thread_context =
-      static_cast<Context *>(pthread_getspecific(context_key));
-  if (current_thread_context == nullptr) {
-    current_thread_context =
-        static_cast<Context *>(InternalAlloc(sizeof(Context)));
-    new (current_thread_context) Context();
-    pthread_setspecific(context_key, current_thread_context);
-  }
-
-  return *current_thread_context;
+static void InitializeContext() {
+  context = static_cast<Context *>(InternalAlloc(sizeof(Context)));
+  new (context) Context();
 }
 
-__rtsan::Context::Context() = default;
+static __rtsan::Context &GetContextImpl() {
+  if (context == nullptr)
+    InitializeContext();
+  return *context;
+}
 
-void __rtsan::Context::RealtimePush() { realtime_depth_++; }
+void __rtsan::Context::RealtimePush() {
+  __sanitizer::SpinMutexLock lock{&spin_mutex_};
+  depths_[pthread_self()].realtime++;
+}
 
-void __rtsan::Context::RealtimePop() { realtime_depth_--; }
+void __rtsan::Context::RealtimePop() {
+  __sanitizer::SpinMutexLock lock{&spin_mutex_};
+  pthread_t const thread_id = pthread_self();
+  depths_[thread_id].realtime--;
+  MaybeCleanup(thread_id);
+}
 
-void __rtsan::Context::BypassPush() { bypass_depth_++; }
+void __rtsan::Context::BypassPush() {
+  __sanitizer::SpinMutexLock lock{&spin_mutex_};
+  depths_[pthread_self()].bypass++;
+}
 
-void __rtsan::Context::BypassPop() { bypass_depth_--; }
+void __rtsan::Context::BypassPop() {
+  __sanitizer::SpinMutexLock lock{&spin_mutex_};
+  pthread_t const thread_id = pthread_self();
+  depths_[thread_id].bypass--;
+  MaybeCleanup(thread_id);
+}
 
-bool __rtsan::Context::InRealtimeContext() const { return realtime_depth_ > 0; }
+bool __rtsan::Context::InRealtimeContext() const {
+  __sanitizer::SpinMutexLock lock{&spin_mutex_};
+  return depths_.lookup(pthread_self()).realtime > 0;
+}
 
-bool __rtsan::Context::IsBypassed() const { return bypass_depth_ > 0; }
+bool __rtsan::Context::IsBypassed() const {
+  __sanitizer::SpinMutexLock lock{&spin_mutex_};
+  return depths_.lookup(pthread_self()).bypass > 0;
+}
 
-Context &__rtsan::GetContextForThisThread() {
-  return GetContextForThisThreadImpl();
+void __rtsan::Context::MaybeCleanup(pthread_t thread_id) {
+  if (depths_[thread_id] == Depth{})
+    depths_.erase(thread_id);
 }
+
+Context &__rtsan::GetContext() { return GetContextImpl(); }

compiler-rt/lib/rtsan/rtsan_context.h

@@ -10,12 +10,14 @@
 
 #pragma once
 
+#include "sanitizer_common/sanitizer_dense_map.h"
+#include "sanitizer_common/sanitizer_mutex.h"
+#include <pthread.h>
+
 namespace __rtsan {
 
 class Context {
 public:
-  Context();
-
   void RealtimePush();
   void RealtimePop();
 
@@ -25,14 +27,25 @@ class Context {
   bool InRealtimeContext() const;
   bool IsBypassed() const;
 
-  Context(const Context &) = delete;
-  Context(Context &&) = delete;
-  Context &operator=(const Context &) = delete;
-  Context &operator=(Context &&) = delete;
-
 private:
-  int realtime_depth_{0};
-  int bypass_depth_{0};
+  static constexpr int max_concurrent_threads_{4096};
+  struct Depth {
+    int realtime{0};
+    int bypass{0};
+    bool operator==(Depth const &other) const {
+      return realtime == other.realtime && bypass == other.bypass;
+    }
+  };
+
+  void MaybeCleanup(pthread_t thread_id);
+
+  // This map serves as thread-local storage implemented entirely in user space.
+  // If an OS's implementation of pthread tls initialisation calls one of the
+  // intercepted functions in rtsan, an infinite recursion can occur when trying
+  // to initialize TLS for a thread. Using this user-space TLS avoids the problem
+  // entirely.
+  __sanitizer::DenseMap<pthread_t, Depth> depths_{max_concurrent_threads_};
+  mutable __sanitizer::SpinMutex spin_mutex_;
 };
 
 class ScopedBypass {
@@ -52,5 +65,5 @@ class ScopedBypass {
   Context &context_;
 };
 
-Context &GetContextForThisThread();
+Context &GetContext();
 } // namespace __rtsan

compiler-rt/lib/rtsan/rtsan_interceptors_posix.cpp

@@ -720,6 +720,11 @@ INTERCEPTOR(void, _os_nospin_lock_lock, _os_nospin_lock_t lock) {
   __rtsan_notify_intercepted_call("_os_nospin_lock_lock");
   return REAL(_os_nospin_lock_lock)(lock);
 }
+
+INTERCEPTOR(void, _os_nospin_lock_unlock, _os_nospin_lock_t lock) {
+  __rtsan_notify_intercepted_call("_os_nospin_lock_unlock");
+  return REAL(_os_nospin_lock_unlock)(lock);
+}
 #pragma clang diagnostic pop // "-Wdeprecated-declarations"
 #endif                       // SANITIZER_APPLE
 
@@ -729,6 +734,10 @@ INTERCEPTOR(void, os_unfair_lock_lock, os_unfair_lock_t lock) {
   return REAL(os_unfair_lock_lock)(lock);
 }
 
+INTERCEPTOR(void, os_unfair_lock_unlock, os_unfair_lock_t lock) {
+  __rtsan_notify_intercepted_call("os_unfair_lock_unlock");
+  return REAL(os_unfair_lock_unlock)(lock);
+}
 #define RTSAN_MAYBE_INTERCEPT_OS_UNFAIR_LOCK_LOCK                              \
   INTERCEPT_FUNCTION(os_unfair_lock_lock)
 #else

compiler-rt/lib/rtsan/rtsan_suppressions.cpp

@@ -19,7 +19,7 @@
 #include "sanitizer_common/sanitizer_suppressions.h"
 #include "sanitizer_common/sanitizer_symbolizer.h"
 
-#include <new>
+#include <stddef.h>
 
 using namespace __sanitizer;
 using namespace __rtsan;
@@ -52,6 +52,8 @@ static const char *ConvertTypeToFlagName(ErrorType Type) {
   UNREACHABLE("unknown ErrorType!");
 }
 
+inline void *operator new(size_t, void *ptr) noexcept { return ptr; }
+
 void __rtsan::InitializeSuppressions() {
   CHECK_EQ(nullptr, suppression_ctx);
 

compiler-rt/lib/rtsan/tests/CMakeLists.txt

@@ -34,8 +34,7 @@ set_target_properties(RtsanUnitTests PROPERTIES FOLDER "Compiler-RT Tests")
 set(RTSAN_UNITTEST_LINK_FLAGS
   ${COMPILER_RT_UNITTEST_LINK_FLAGS}
   ${COMPILER_RT_UNWINDER_LINK_LIBS}
-  ${SANITIZER_TEST_CXX_LIBRARIES}
-  -no-pie)
+  ${SANITIZER_TEST_CXX_LIBRARIES})
 
 if (COMPILER_RT_USE_ATOMIC_LIBRARY)
   list(APPEND RTSAN_UNITTEST_LINK_FLAGS ${COMPILER_RT_ATOMIC_LIBRARY})

compiler-rt/lib/rtsan/tests/rtsan_test_context.cpp

@@ -13,7 +13,10 @@
 #include "rtsan/rtsan.h"
 #include "rtsan/rtsan_context.h"
 
+#include <atomic>
+#include <chrono>
 #include <gtest/gtest.h>
+#include <thread>
 
 using namespace __rtsan;
 using namespace ::testing;
@@ -96,3 +99,63 @@ TEST_F(TestRtsanContext, BypassedStateIsStatefullyTracked) {
   context.BypassPush(); // depth 1
   ExpectBypassed(true);
 }
+
+TEST_F(TestRtsanContext, IsProbablyThreadSafe) {
+  std::atomic<int> num_threads_started{0};
+  std::atomic<bool> all_threads_wait{true};
+  std::atomic<bool> all_threads_continue{true};
+  Context context{};
+
+  auto const expect_context_state =
+      [&context](bool expected_in_realtime_context, bool expected_is_bypassed) {
+        EXPECT_THAT(context.InRealtimeContext(),
+                    Eq(expected_in_realtime_context));
+        EXPECT_THAT(context.IsBypassed(), Eq(expected_is_bypassed));
+      };
+
+  auto const test_thread_work = [&]() {
+    num_threads_started.fetch_add(1);
+    while (all_threads_wait.load())
+      std::this_thread::yield();
+
+    while (all_threads_continue.load()) {
+      context.RealtimePush();
+      expect_context_state(true, false);
+      context.RealtimePush();
+      expect_context_state(true, false);
+
+      context.BypassPush();
+      expect_context_state(true, true);
+      context.BypassPop();
+      expect_context_state(true, false);
+
+      context.RealtimePop();
+      expect_context_state(true, false);
+      context.RealtimePop();
+      expect_context_state(false, false);
+    }
+  };
+
+  auto const time_now = []() { return std::chrono::steady_clock::now(); };
+
+  int const num_threads = 32;
+  std::vector<std::thread> test_threads{};
+  auto const start_time = std::chrono::steady_clock::now();
+  for (int n = 0; n < num_threads; ++n)
+    test_threads.push_back(std::thread(test_thread_work));
+
+  std::chrono::duration timeout = std::chrono::milliseconds(100);
+  while (num_threads_started.load() != num_threads) {
+    if ((time_now() - start_time) > timeout) {
+        FAIL();
+    }
+    std::this_thread::yield();
+  }
+
+  all_threads_wait.store(false);
+  std::this_thread::sleep_for(std::chrono::milliseconds(10));
+  all_threads_continue.store(false);
+
+  for (auto &test_thread : test_threads)
+    test_thread.join();
+}