From 9567b33fa15e5349cb522a3b3e39abc300c7644c Mon Sep 17 00:00:00 2001
From: Vitaly Buka <vitalybuka@google.com>
Date: Tue, 12 Dec 2023 13:57:24 -0800
Subject: [PATCH] [asan] Switch initialization to "double-checked locking"

This allows to remove `asan_init_is_running` which likely had a data
race.

Simplifies https://github.com/llvm/llvm-project/pull/74086 and reduces a difference between platforms.

Reviewers: zacklj89, eugenis, dvyukov

Reviewed By: zacklj89, dvyukov

Pull Request: https://github.com/llvm/llvm-project/pull/74387
---
 compiler-rt/lib/asan/asan_rtl.cpp | 49 ++++++++++++++++---------------
 1 file changed, 25 insertions(+), 24 deletions(-)

diff --git a/compiler-rt/lib/asan/asan_rtl.cpp b/compiler-rt/lib/asan/asan_rtl.cpp
index 04ecd20821fa6d..b28f9f181239b3 100644
--- a/compiler-rt/lib/asan/asan_rtl.cpp
+++ b/compiler-rt/lib/asan/asan_rtl.cpp
@@ -71,16 +71,16 @@ static void CheckUnwind() {
 }
 
 // -------------------------- Globals --------------------- {{{1
-static int asan_inited = 0;
-static int asan_init_is_running = 0;
+static StaticSpinMutex asan_inited_mutex;
+static atomic_uint8_t asan_inited = {0};
 
-static void SetAsanInited() { asan_inited = 1; }
-
-static void SetAsanInitIsRunning(u32 val) { asan_init_is_running = val; }
-
-bool AsanInited() { return asan_inited == 1; }
+static void SetAsanInited() {
+  atomic_store(&asan_inited, 1, memory_order_release);
+}
 
-static bool AsanInitIsRunning() { return asan_init_is_running == 1; }
+bool AsanInited() {
+  return atomic_load(&asan_inited, memory_order_acquire) == 1;
+}
 
 bool replace_intrin_cached;
 
@@ -390,12 +390,10 @@ void PrintAddressSpaceLayout() {
           kHighShadowBeg > kMidMemEnd);
 }
 
-static void AsanInitInternal() {
+static bool AsanInitInternal() {
   if (LIKELY(AsanInited()))
-    return;
+    return true;
   SanitizerToolName = "AddressSanitizer";
-  CHECK(!AsanInitIsRunning() && "ASan init calls itself!");
-  SetAsanInitIsRunning(1);
 
   CacheBinaryName();
 
@@ -408,9 +406,8 @@ static void AsanInitInternal() {
   // Stop performing init at this point if we are being loaded via
   // dlopen() and the platform supports it.
   if (SANITIZER_SUPPORTS_INIT_FOR_DLOPEN && UNLIKELY(HandleDlopenInit())) {
-    SetAsanInitIsRunning(0);
     VReport(1, "AddressSanitizer init is being performed for dlopen().\n");
-    return;
+    return false;
   }
 
   AsanCheckIncompatibleRT();
@@ -471,7 +468,6 @@ static void AsanInitInternal() {
   // should be set to 1 prior to initializing the threads.
   replace_intrin_cached = flags()->replace_intrin;
   SetAsanInited();
-  SetAsanInitIsRunning(0);
 
   if (flags()->atexit)
     Atexit(asan_atexit);
@@ -515,22 +511,27 @@ static void AsanInitInternal() {
   VReport(1, "AddressSanitizer Init done\n");
 
   WaitForDebugger(flags()->sleep_after_init, "after init");
+
+  return true;
 }
 
 // Initialize as requested from some part of ASan runtime library (interceptors,
 // allocator, etc).
 void AsanInitFromRtl() {
-  CHECK(!AsanInitIsRunning());
-  if (UNLIKELY(!AsanInited()))
-    AsanInitInternal();
+  if (LIKELY(AsanInited()))
+    return;
+  SpinMutexLock lock(&asan_inited_mutex);
+  AsanInitInternal();
 }
 
 bool TryAsanInitFromRtl() {
-  if (UNLIKELY(AsanInitIsRunning()))
+  if (LIKELY(AsanInited()))
+    return true;
+  if (!asan_inited_mutex.TryLock())
     return false;
-  if (UNLIKELY(!AsanInited()))
-    AsanInitInternal();
-  return true;
+  bool result = AsanInitInternal();
+  asan_inited_mutex.Unlock();
+  return result;
 }
 
 #if ASAN_DYNAMIC
@@ -603,7 +604,7 @@ static void UnpoisonFakeStack() {
 using namespace __asan;
 
 void NOINLINE __asan_handle_no_return() {
-  if (AsanInitIsRunning())
+  if (UNLIKELY(!AsanInited()))
     return;
 
   if (!PlatformUnpoisonStacks())
@@ -633,7 +634,7 @@ void NOINLINE __asan_set_death_callback(void (*callback)(void)) {
 // We use this call as a trigger to wake up ASan from deactivated state.
 void __asan_init() {
   AsanActivate();
-  AsanInitInternal();
+  AsanInitFromRtl();
 }
 
 void __asan_version_mismatch_check() {