[compiler-rt] 0907552 - [rtsan] Refactor initialization state to prevent hangs (#109830)

[via llvm-commits]

Author: Chris Apple
Date: 2024-09-24T10:01:12-07:00
New Revision: 090755234e5033de67be994e99e31f4d13dcdcc5

URL: https://github.com/llvm/llvm-project/commit/090755234e5033de67be994e99e31f4d13dcdcc5
DIFF: https://github.com/llvm/llvm-project/commit/090755234e5033de67be994e99e31f4d13dcdcc5.diff

LOG: [rtsan] Refactor initialization state to prevent hangs (#109830)

Move to tri state (Uninitialized, Initialized, Initializing) for our
init state. We currently just have 2 (Uninitialized and Initialized).

Added: 
    

Modified: 
    compiler-rt/lib/rtsan/rtsan.cpp

Removed: 
    


################################################################################
diff  --git a/compiler-rt/lib/rtsan/rtsan.cpp b/compiler-rt/lib/rtsan/rtsan.cpp
index b288da64ffbe25..e6d2481b2c2a3d 100644
--- a/compiler-rt/lib/rtsan/rtsan.cpp
+++ b/compiler-rt/lib/rtsan/rtsan.cpp
@@ -22,11 +22,25 @@
 using namespace __rtsan;
 using namespace __sanitizer;
 
+namespace {
+enum class InitializationState : u8 {
+  Uninitialized,
+  Initializing,
+  Initialized,
+};
+} // namespace
+
 static StaticSpinMutex rtsan_inited_mutex;
 static atomic_uint8_t rtsan_initialized = {0};
 
-static void SetInitialized() {
-  atomic_store(&rtsan_initialized, 1, memory_order_release);
+static void SetInitializationState(InitializationState state) {
+  atomic_store(&rtsan_initialized, static_cast<u8>(state),
+               memory_order_release);
+}
+
+static InitializationState GetInitializationState() {
+  return static_cast<InitializationState>(
+      atomic_load(&rtsan_initialized, memory_order_acquire));
 }
 
 static auto PrintDiagnosticsAndDieAction(DiagnosticsInfo info) {
@@ -39,13 +53,14 @@ static auto PrintDiagnosticsAndDieAction(DiagnosticsInfo info) {
 extern "C" {
 
 SANITIZER_INTERFACE_ATTRIBUTE void __rtsan_init() {
-  CHECK(!__rtsan_is_initialized());
+  CHECK(GetInitializationState() == InitializationState::Uninitialized);
+  SetInitializationState(InitializationState::Initializing);
 
   SanitizerToolName = "RealtimeSanitizer";
   InitializeFlags();
   InitializeInterceptors();
 
-  SetInitialized();
+  SetInitializationState(InitializationState::Initialized);
 }
 
 SANITIZER_INTERFACE_ATTRIBUTE void __rtsan_ensure_initialized() {
@@ -62,7 +77,7 @@ SANITIZER_INTERFACE_ATTRIBUTE void __rtsan_ensure_initialized() {
 }
 
 SANITIZER_INTERFACE_ATTRIBUTE bool __rtsan_is_initialized() {
-  return atomic_load(&rtsan_initialized, memory_order_acquire) == 1;
+  return GetInitializationState() == InitializationState::Initialized;
 }
 
 SANITIZER_INTERFACE_ATTRIBUTE void __rtsan_realtime_enter() {
@@ -83,6 +98,10 @@ SANITIZER_INTERFACE_ATTRIBUTE void __rtsan_enable() {
 
 SANITIZER_INTERFACE_ATTRIBUTE void
 __rtsan_notify_intercepted_call(const char *func_name) {
+  // While initializing, we need all intercepted functions to behave normally
+  if (GetInitializationState() == InitializationState::Initializing)
+    return;
+
   __rtsan_ensure_initialized();
   GET_CALLER_PC_BP;
   ExpectNotRealtime(