[asan] Optimize initialization order checking

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