[asan] Optimize initialization order checking (llvm#101837)

The pair `__asan_before_dynamic_init` and `__asan_after_dynamic_init` is
executed for each TU. `__asan_after_dynamic_init` unpoisons all globals,
which
makes the time complexity O(N^2), where N is the maximum of the global
count and
the TU count. This is expensive for large binaries.

This patch decreases the time complexity to O(N), when lld and static
runtime is
used on SANITIZER_CAN_USE_PREINIT_ARRAY platforms. This requires:

Enabling incremental poisoning (`__asan_before_dynamic_init` since
llvm#101597). Making most
`__asan_after_dynamic_init` calls do nothing.

compiler-rt/lib/asan/asan_globals.cpp

@@ -520,6 +520,43 @@ void __asan_before_dynamic_init(const char *module_name) {
   current_dynamic_init_module_name = module_name;
 }
 
+// Maybe SANITIZER_CAN_USE_PREINIT_ARRAY is to conservative for `.init_array`,
+// however we should not make mistake here. If `UnpoisonBeforeMain` was not
+// executed at all we will have false reports on globals.
+#if SANITIZER_CAN_USE_PREINIT_ARRAY
+// This optimization aims to reduce the overhead of `__asan_after_dynamic_init`
+// calls by leveraging incremental unpoisoning/poisoning in
+// `__asan_before_dynamic_init`. We expect most `__asan_after_dynamic_init
+// calls` to be no-ops. However, to ensure all globals are unpoisoned before the
+// `main`, we force `UnpoisonBeforeMain` to fully execute
+// `__asan_after_dynamic_init`.
+
+// With lld, `UnpoisonBeforeMain` runs after standard `.init_array`, making it
+// the final `__asan_after_dynamic_init` call for the static runtime. In
+// contrast, GNU ld executes it earlier, causing subsequent
+// `__asan_after_dynamic_init` calls to perform full unpoisoning, losing the
+// optimization.
+bool allow_after_dynamic_init SANITIZER_GUARDED_BY(mu_for_globals) = false;
+
+static void UnpoisonBeforeMain(void) {
+  {
+    Lock lock(&mu_for_globals);
+    if (allow_after_dynamic_init)
+      return;
+    allow_after_dynamic_init = true;
+  }
+  if (flags()->report_globals >= 3)
+    Printf("UnpoisonBeforeMain\n");
+  __asan_after_dynamic_init();
+}
+
+__attribute__((section(".init_array.65537"), used)) static void (
+    *asan_after_init_array)(void) = UnpoisonBeforeMain;
+#else
+// Incremental poisoning is disabled, unpoison globals immediately.
+static constexpr bool allow_after_dynamic_init = true;
+#endif  // SANITIZER_CAN_USE_PREINIT_ARRAY
+
 // This method runs immediately after dynamic initialization in each TU, when
 // all dynamically initialized globals except for those defined in the current
 // TU are poisoned.  It simply unpoisons all dynamically initialized globals.
@@ -528,6 +565,8 @@ void __asan_after_dynamic_init() {
     return;
   CHECK(AsanInited());
   Lock lock(&mu_for_globals);
+  if (!allow_after_dynamic_init)
+    return;
   if (!current_dynamic_init_module_name)
     return;
 

compiler-rt/test/asan/TestCases/Linux/initialization-nobug-lld.cpp

@@ -0,0 +1,14 @@
+// RUN: %clangxx_asan -O3 %S/../initialization-nobug.cpp %S/../Helpers/initialization-nobug-extra.cpp -fuse-ld=lld -o %t && %env_asan_opts=check_initialization_order=true:report_globals=3 %run %t 2>&1 | FileCheck %s --implicit-check-not "DynInit"
+
+// Same as initialization-nobug.cpp, but with lld we expect just one
+// `DynInitUnpoison` executed after `AfterDynamicInit` at the end.
+// REQUIRES: lld-available
+
+// With dynamic runtimes `AfterDynamicInit` will called before `executable`
+// contructors, with constructors of dynamic runtime.
+// XFAIL: asan-dynamic-runtime
+
+// CHECK: DynInitPoison
+// CHECK: DynInitPoison
+// CHECK: UnpoisonBeforeMain
+// CHECK: DynInitUnpoison

compiler-rt/test/asan/TestCases/initialization-nobug.cpp

@@ -1,10 +1,10 @@
 // A collection of various initializers which shouldn't trip up initialization
 // order checking.  If successful, this will just return 0.
 
-// RUN: %clangxx_asan -O0 %s %p/Helpers/initialization-nobug-extra.cpp -o %t && %env_asan_opts=check_initialization_order=true:report_globals=3 %run %t 2>&1 | FileCheck %s --implicit-check-not "DynInit"
-// RUN: %clangxx_asan -O1 %s %p/Helpers/initialization-nobug-extra.cpp -o %t && %env_asan_opts=check_initialization_order=true:report_globals=3 %run %t 2>&1 | FileCheck %s --implicit-check-not "DynInit"
-// RUN: %clangxx_asan -O2 %s %p/Helpers/initialization-nobug-extra.cpp -o %t && %env_asan_opts=check_initialization_order=true:report_globals=3 %run %t 2>&1 | FileCheck %s --implicit-check-not "DynInit"
-// RUN: %clangxx_asan -O3 %s %p/Helpers/initialization-nobug-extra.cpp -o %t && %env_asan_opts=check_initialization_order=true:report_globals=3 %run %t 2>&1 | FileCheck %s --implicit-check-not "DynInit"
+// RUN: %clangxx_asan -O0 %s %p/Helpers/initialization-nobug-extra.cpp -o %t && %env_asan_opts=check_initialization_order=true:report_globals=3 %run %t 2>&1 | FileCheck %s --implicit-check-not "DynInitPoison"
+// RUN: %clangxx_asan -O1 %s %p/Helpers/initialization-nobug-extra.cpp -o %t && %env_asan_opts=check_initialization_order=true:report_globals=3 %run %t 2>&1 | FileCheck %s --implicit-check-not "DynInitPoison"
+// RUN: %clangxx_asan -O2 %s %p/Helpers/initialization-nobug-extra.cpp -o %t && %env_asan_opts=check_initialization_order=true:report_globals=3 %run %t 2>&1 | FileCheck %s --implicit-check-not "DynInitPoison"
+// RUN: %clangxx_asan -O3 %s %p/Helpers/initialization-nobug-extra.cpp -o %t && %env_asan_opts=check_initialization_order=true:report_globals=3 %run %t 2>&1 | FileCheck %s --implicit-check-not "DynInitPoison"
 
 // Simple access:
 // Make sure that accessing a global in the same TU is safe
@@ -44,6 +44,10 @@ int getStructWithDtorValue() { return struct_with_dtor.value; }
 int main() { return 0; }
 
 // CHECK: DynInitPoison
-// CHECK: DynInitUnpoison
 // CHECK: DynInitPoison
+
+// In general case entire set of DynInitPoison must be followed by at lest one
+// DynInitUnpoison. In some cases we can limit the number of DynInitUnpoison,
+// see initialization-nobug-lld.cpp.
+
 // CHECK: DynInitUnpoison