[asan] Optimize initialization order checking

[vitalybuka]

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
#101597). Making most
__asan_after_dynamic_init calls do nothing.

[llvmbot]

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

Author: Vitaly Buka (vitalybuka)

Changes

Before the patch each TU was executing pair
__asan_before_dynamic_init and
__asan_after_dynamic_init.

__asan_before_dynamic_init supports incremental
poisoning, but __asan_after_dynamic_init
unpoisons all globals, so
__asan_before_dynamic_init must repoison
everything again.

So it's O(N^2) where N is the number of globals
(or TUs). It can be expensive for large binaries.

The patch will make it O(N).

The patch requires LLD and platform with
SANITIZER_CAN_USE_PREINIT_ARRAY. For the rest it
continue to be O(N^2).

---

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

3 Files Affected:

• (modified) compiler-rt/lib/asan/asan_globals.cpp (+40)
• (added) compiler-rt/test/asan/TestCases/initialization-nobug-lld.cpp (+14)
• (modified) compiler-rt/test/asan/TestCases/initialization-nobug.cpp (+8-5)

diff --git a/compiler-rt/lib/asan/asan_globals.cpp b/compiler-rt/lib/asan/asan_globals.cpp
index 293e771e1dc06..373b7eaa3c537 100644
--- a/compiler-rt/lib/asan/asan_globals.cpp
+++ b/compiler-rt/lib/asan/asan_globals.cpp
@@ -520,6 +520,44 @@ 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 `AfterDynamicInit` was not
+// executed at all we will have false reports on globals.
+#if SANITIZER_CAN_USE_PREINIT_ARRAY
+// This is optimization. We will ignore all `__asan_after_dynamic_init`, but the
+// last `__asan_after_dynamic_init`. We don't need number of
+// `__asan_{before,after}_dynamic_init` matches, but we need that the last call
+// was to `__asan_after_dynamic_init`, as it will unpoison all global preparing
+// program for `main` execution. To run `__asan_after_dynamic_init` later, we
+// will register in `.init_array`.
+static bool allow_after_dynamic_init SANITIZER_GUARDED_BY(mu_for_globals) =
+    false;
+
+static void __attribute__((used)) AfterDynamicInit(void) {
+  {
+    Lock lock(&mu_for_globals);
+    if (allow_after_dynamic_init)
+      return;
+    allow_after_dynamic_init = true;
+  }
+  if (flags()->report_globals >= 3)
+    Printf("AfterDynamicInit\n");
+  __asan_after_dynamic_init();
+}
+
+// 65537 will make it run after constructors with default priority, but it
+// requires ld.lld. With ld.bfd it can be called to early, and fail the
+// optimization. However, correctness should not be affected, as after the first
+// call all subsequent `__asan_after_dynamic_init` will be allowed.
+__attribute__((section(".init_array.65537"), used)) static void (
+    *asan_after_init_array)(void) = AfterDynamicInit;
+#else
+// Allow all `__asan_after_dynamic_init` if `AfterDynamicInit` is not set.
+// Compiler still generates `__asan_{before,after}_dynamic_init`in pairs, and
+// it's guaranteed that `__asan_after_dynamic_init` will be the last.
+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 +566,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;
 
diff --git a/compiler-rt/test/asan/TestCases/initialization-nobug-lld.cpp b/compiler-rt/test/asan/TestCases/initialization-nobug-lld.cpp
new file mode 100644
index 0000000000000..ece3118de4ad9
--- /dev/null
+++ b/compiler-rt/test/asan/TestCases/initialization-nobug-lld.cpp
@@ -0,0 +1,14 @@
+// RUN: %clangxx_asan -O3 %p/initialization-nobug.cpp %p/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 module: {{.*}}initialization-nobug.cpp
+// CHECK: DynInitPoison module: {{.*}}initialization-nobug-extra.cpp
+// CHECK: AfterDynamicInit
+// CHECK: DynInitUnpoison
diff --git a/compiler-rt/test/asan/TestCases/initialization-nobug.cpp b/compiler-rt/test/asan/TestCases/initialization-nobug.cpp
index 5659db088096b..f58f38203d0b2 100644
--- a/compiler-rt/test/asan/TestCases/initialization-nobug.cpp
+++ b/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,9 @@ int getStructWithDtorValue() { return struct_with_dtor.value; }
 int main() { return 0; }
 
 // CHECK: DynInitPoison module: {{.*}}initialization-nobug.cpp
-// CHECK: DynInitUnpoison
 // CHECK: DynInitPoison module: {{.*}}initialization-nobug-extra.cpp
+
+// In general case we only need the last of DynInit{Poison,Unpoison} is always
+// Unpoison.
+
 // CHECK: DynInitUnpoison

[vitalybuka]

FYI @artempyanykh