[rtsan] Refactor initialization state to prevent hangs

[cjappl]

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

Why?

Later, when introducing exit statistics, I add an Atexit handler. This handler is called, as the name suggests, at exit time. Our Atexit handler will print the exit stats if the user has requested.

We will install this Atexit handler when we are initializing for the first time. However during this process, intercepted functions are called - causing a hang.

We now early exit out of __rtsan_notify_intercepted_call if we are initializing

I do not think we need the same check in __rtsan_notify_blocking_call because no calls marked blocking will be made during init, these are mostly low level bootstrap calls.

[cjappl]

CC for review @davidtrevelyan

[llvmbot]

@llvm/pr-subscribers-compiler-rt-sanitizer

Author: Chris Apple (cjappl)

Changes

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

Why?

Later, when introducing exit statistics, I add an Atexit handler. This handler is called, as the name suggests, at exit time. Our Atexit handler will print the exit stats if the user has requested.

We will install this Atexit handler when we are initializing for the first time. However during this process, intercepted functions are called - causing a hang.

We now early exit out of __rtsan_notify_intercepted_call if we are initializing

I do not think we need the same check in __rtsan_notify_blocking_call because no calls marked blocking will be made during init, these are mostly low level bootstrap calls.

---

Full diff: https://github.com/llvm/llvm-project/pull/109830.diff

1 Files Affected:

• (modified) compiler-rt/lib/rtsan/rtsan.cpp (+24-5)

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(