[NFC][sanitizer] Rename Lock{Before,After}Fork suffixes locking StackDepotBase

[vitalybuka]

This is preporation for performance optimization.

We need to highlight that this is very specific lock, and should not be
used for other purposes.
Add fork_child parameter to distinguish processes after fork.

[llvmbot]

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

Author: Vitaly Buka (vitalybuka)

Changes

This is preporation for performance optimization.

We need to highlight that this is very specific lock, and should not be
used for other purposes.
Add fork_child parameter to distinguish processes after fork.

---

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

11 Files Affected:

• (modified) compiler-rt/lib/asan/asan_posix.cpp (+26-22)
• (modified) compiler-rt/lib/dfsan/dfsan_chained_origin_depot.cpp (+6)
• (modified) compiler-rt/lib/dfsan/dfsan_chained_origin_depot.h (+3)
• (modified) compiler-rt/lib/dfsan/dfsan_custom.cpp (+6-6)
• (modified) compiler-rt/lib/hwasan/hwasan_linux.cpp (+25-21)
• (modified) compiler-rt/lib/lsan/lsan_posix.cpp (+17-13)
• (modified) compiler-rt/lib/msan/msan_chained_origin_depot.cpp (+2-4)
• (modified) compiler-rt/lib/msan/msan_chained_origin_depot.h (+2-2)
• (modified) compiler-rt/lib/msan/msan_linux.cpp (+18-14)
• (modified) compiler-rt/lib/sanitizer_common/sanitizer_stackdepot.cpp (+2-2)
• (modified) compiler-rt/lib/sanitizer_common/sanitizer_stackdepot.h (+2-2)

diff --git a/compiler-rt/lib/asan/asan_posix.cpp b/compiler-rt/lib/asan/asan_posix.cpp
index a5b87b7fbf1b5a..76564538bd5d77 100644
--- a/compiler-rt/lib/asan/asan_posix.cpp
+++ b/compiler-rt/lib/asan/asan_posix.cpp
@@ -146,33 +146,37 @@ void PlatformTSDDtor(void *tsd) {
 #    endif
   AsanThread::TSDDtor(tsd);
 }
-#endif
+#  endif
+
+static void BeforeFork() {
+  if (CAN_SANITIZE_LEAKS) {
+    __lsan::LockGlobal();
+  }
+  // `_lsan` functions defined regardless of `CAN_SANITIZE_LEAKS` and lock the
+  // stuff we need.
+  __lsan::LockThreads();
+  __lsan::LockAllocator();
+  StackDepotLockBeforeFork();
+}
+
+static void AfterFork(bool fork_child) {
+  StackDepotUnlockAfterFork(fork_child);
+  // `_lsan` functions defined regardless of `CAN_SANITIZE_LEAKS` and unlock
+  // the stuff we need.
+  __lsan::UnlockAllocator();
+  __lsan::UnlockThreads();
+  if (CAN_SANITIZE_LEAKS) {
+    __lsan::UnlockGlobal();
+  }
+}
 
 void InstallAtForkHandler() {
 #  if SANITIZER_SOLARIS || SANITIZER_NETBSD || SANITIZER_APPLE
   return;  // FIXME: Implement FutexWait.
 #  endif
-  auto before = []() {
-    if (CAN_SANITIZE_LEAKS) {
-      __lsan::LockGlobal();
-    }
-    // `_lsan` functions defined regardless of `CAN_SANITIZE_LEAKS` and lock the
-    // stuff we need.
-    __lsan::LockThreads();
-    __lsan::LockAllocator();
-    StackDepotLockAll();
-  };
-  auto after = []() {
-    StackDepotUnlockAll();
-    // `_lsan` functions defined regardless of `CAN_SANITIZE_LEAKS` and unlock
-    // the stuff we need.
-    __lsan::UnlockAllocator();
-    __lsan::UnlockThreads();
-    if (CAN_SANITIZE_LEAKS) {
-      __lsan::UnlockGlobal();
-    }
-  };
-  pthread_atfork(before, after, after);
+  pthread_atfork(
+      &BeforeFork, []() { AfterFork(/* fork_child= */ false); },
+      []() { AfterFork(/* fork_child= */ true); });
 }
 
 void InstallAtExitCheckLeaks() {
diff --git a/compiler-rt/lib/dfsan/dfsan_chained_origin_depot.cpp b/compiler-rt/lib/dfsan/dfsan_chained_origin_depot.cpp
index 9ec598bf2ce9e4..6644bd6a7c6c0c 100644
--- a/compiler-rt/lib/dfsan/dfsan_chained_origin_depot.cpp
+++ b/compiler-rt/lib/dfsan/dfsan_chained_origin_depot.cpp
@@ -19,4 +19,10 @@ static ChainedOriginDepot chainedOriginDepot;
 
 ChainedOriginDepot* GetChainedOriginDepot() { return &chainedOriginDepot; }
 
+void ChainedOriginDepotLockBeforeFork() { chainedOriginDepot.LockAll(); }
+
+void ChainedOriginDepotUnlockAfterFork(bool fork_child) {
+  chainedOriginDepot.UnlockAll();
+}
+
 }  // namespace __dfsan
diff --git a/compiler-rt/lib/dfsan/dfsan_chained_origin_depot.h b/compiler-rt/lib/dfsan/dfsan_chained_origin_depot.h
index d715ef707f41b7..83b9e29e1b710d 100644
--- a/compiler-rt/lib/dfsan/dfsan_chained_origin_depot.h
+++ b/compiler-rt/lib/dfsan/dfsan_chained_origin_depot.h
@@ -21,6 +21,9 @@ namespace __dfsan {
 
 ChainedOriginDepot* GetChainedOriginDepot();
 
+void ChainedOriginDepotLockBeforeFork();
+void ChainedOriginDepotUnlockAfterFork(bool fork_child);
+
 }  // namespace __dfsan
 
 #endif  // DFSAN_CHAINED_ORIGIN_DEPOT_H
diff --git a/compiler-rt/lib/dfsan/dfsan_custom.cpp b/compiler-rt/lib/dfsan/dfsan_custom.cpp
index 38371d35336818..05b48fd0525e33 100644
--- a/compiler-rt/lib/dfsan/dfsan_custom.cpp
+++ b/compiler-rt/lib/dfsan/dfsan_custom.cpp
@@ -2893,13 +2893,13 @@ int __dfso___isoc99_sscanf(char *str, const char *format, dfsan_label str_label,
 }
 
 static void BeforeFork() {
-  StackDepotLockAll();
-  GetChainedOriginDepot()->LockAll();
+  StackDepotLockBeforeFork();
+  ChainedOriginDepotLockBeforeFork();
 }
 
-static void AfterFork() {
-  GetChainedOriginDepot()->UnlockAll();
-  StackDepotUnlockAll();
+static void AfterFork(bool fork_child) {
+  ChainedOriginDepotUnlockAfterFork(fork_child);
+  StackDepotUnlockAfterFork(fork_child);
 }
 
 SANITIZER_INTERFACE_ATTRIBUTE
@@ -2913,7 +2913,7 @@ SANITIZER_INTERFACE_ATTRIBUTE
 pid_t __dfso_fork(dfsan_label *ret_label, dfsan_origin *ret_origin) {
   BeforeFork();
   pid_t pid = __dfsw_fork(ret_label);
-  AfterFork();
+  AfterFork(/* fork_child= */ pid == 0);
   return pid;
 }
 
diff --git a/compiler-rt/lib/hwasan/hwasan_linux.cpp b/compiler-rt/lib/hwasan/hwasan_linux.cpp
index 3271a955e7ed10..e6aa60b324fa7f 100644
--- a/compiler-rt/lib/hwasan/hwasan_linux.cpp
+++ b/compiler-rt/lib/hwasan/hwasan_linux.cpp
@@ -521,28 +521,32 @@ uptr TagMemoryAligned(uptr p, uptr size, tag_t tag) {
   return AddTagToPointer(p, tag);
 }
 
+static void BeforeFork() {
+  if (CAN_SANITIZE_LEAKS) {
+    __lsan::LockGlobal();
+  }
+  // `_lsan` functions defined regardless of `CAN_SANITIZE_LEAKS` and lock the
+  // stuff we need.
+  __lsan::LockThreads();
+  __lsan::LockAllocator();
+  StackDepotLockBeforeFork();
+}
+
+static void AfterFork(bool fork_child) {
+  StackDepotUnlockAfterFork(fork_child);
+  // `_lsan` functions defined regardless of `CAN_SANITIZE_LEAKS` and unlock
+  // the stuff we need.
+  __lsan::UnlockAllocator();
+  __lsan::UnlockThreads();
+  if (CAN_SANITIZE_LEAKS) {
+    __lsan::UnlockGlobal();
+  }
+}
+
 void HwasanInstallAtForkHandler() {
-  auto before = []() {
-    if (CAN_SANITIZE_LEAKS) {
-      __lsan::LockGlobal();
-    }
-    // `_lsan` functions defined regardless of `CAN_SANITIZE_LEAKS` and lock the
-    // stuff we need.
-    __lsan::LockThreads();
-    __lsan::LockAllocator();
-    StackDepotLockAll();
-  };
-  auto after = []() {
-    StackDepotUnlockAll();
-    // `_lsan` functions defined regardless of `CAN_SANITIZE_LEAKS` and unlock
-    // the stuff we need.
-    __lsan::UnlockAllocator();
-    __lsan::UnlockThreads();
-    if (CAN_SANITIZE_LEAKS) {
-      __lsan::UnlockGlobal();
-    }
-  };
-  pthread_atfork(before, after, after);
+  pthread_atfork(
+      &BeforeFork, []() { AfterFork(/* fork_child= */ false); },
+      []() { AfterFork(/* fork_child= */ true); });
 }
 
 void InstallAtExitCheckLeaks() {
diff --git a/compiler-rt/lib/lsan/lsan_posix.cpp b/compiler-rt/lib/lsan/lsan_posix.cpp
index 4bfadf1ef809c3..422c29acca69f3 100644
--- a/compiler-rt/lib/lsan/lsan_posix.cpp
+++ b/compiler-rt/lib/lsan/lsan_posix.cpp
@@ -100,23 +100,27 @@ void InstallAtExitCheckLeaks() {
     Atexit(DoLeakCheck);
 }
 
+static void BeforeFork() {
+  LockGlobal();
+  LockThreads();
+  LockAllocator();
+  StackDepotLockBeforeFork();
+}
+
+static void AfterFork(bool fork_child) {
+  StackDepotUnlockAfterFork(fork_child);
+  UnlockAllocator();
+  UnlockThreads();
+  UnlockGlobal();
+}
+
 void InstallAtForkHandler() {
 #  if SANITIZER_SOLARIS || SANITIZER_NETBSD || SANITIZER_APPLE
   return;  // FIXME: Implement FutexWait.
 #  endif
-  auto before = []() {
-    LockGlobal();
-    LockThreads();
-    LockAllocator();
-    StackDepotLockAll();
-  };
-  auto after = []() {
-    StackDepotUnlockAll();
-    UnlockAllocator();
-    UnlockThreads();
-    UnlockGlobal();
-  };
-  pthread_atfork(before, after, after);
+  pthread_atfork(
+      &BeforeFork, []() { AfterFork(/* fork_child= */ false); },
+      []() { AfterFork(/* fork_child= */ true); });
 }
 
 }  // namespace __lsan
diff --git a/compiler-rt/lib/msan/msan_chained_origin_depot.cpp b/compiler-rt/lib/msan/msan_chained_origin_depot.cpp
index 49b14131a89be0..c3bd54141e6c38 100644
--- a/compiler-rt/lib/msan/msan_chained_origin_depot.cpp
+++ b/compiler-rt/lib/msan/msan_chained_origin_depot.cpp
@@ -31,11 +31,9 @@ u32 ChainedOriginDepotGet(u32 id, u32 *other) {
   return chainedOriginDepot.Get(id, other);
 }
 
-void ChainedOriginDepotLockAll() {
-  chainedOriginDepot.LockAll();
-}
+void ChainedOriginDepotBeforeFork() { chainedOriginDepot.LockAll(); }
 
-void ChainedOriginDepotUnlockAll() {
+void ChainedOriginDepotAfterFork(bool fork_child) {
   chainedOriginDepot.UnlockAll();
 }
 
diff --git a/compiler-rt/lib/msan/msan_chained_origin_depot.h b/compiler-rt/lib/msan/msan_chained_origin_depot.h
index ea51c77a905b5d..7518745dc8520a 100644
--- a/compiler-rt/lib/msan/msan_chained_origin_depot.h
+++ b/compiler-rt/lib/msan/msan_chained_origin_depot.h
@@ -30,8 +30,8 @@ bool ChainedOriginDepotPut(u32 here_id, u32 prev_id, u32 *new_id);
 // Retrieves the stored StackDepot ID for the given origin ID.
 u32 ChainedOriginDepotGet(u32 id, u32 *other);
 
-void ChainedOriginDepotLockAll();
-void ChainedOriginDepotUnlockAll();
+void ChainedOriginDepotBeforeFork();
+void ChainedOriginDepotAfterFork(bool fork_child);
 
 }  // namespace __msan
 
diff --git a/compiler-rt/lib/msan/msan_linux.cpp b/compiler-rt/lib/msan/msan_linux.cpp
index 04af6f4b27ac89..c7ecb7cad56661 100644
--- a/compiler-rt/lib/msan/msan_linux.cpp
+++ b/compiler-rt/lib/msan/msan_linux.cpp
@@ -256,22 +256,26 @@ void MsanTSDDtor(void *tsd) {
   atomic_signal_fence(memory_order_seq_cst);
   MsanThread::TSDDtor(tsd);
 }
-#endif
+#  endif
+
+static void BeforeFork() {
+  // Usually we lock ThreadRegistry, but msan does not have one.
+  LockAllocator();
+  StackDepotLockBeforeFork();
+  ChainedOriginDepotBeforeFork();
+}
+
+static void AfterFork(bool fork_child) {
+  ChainedOriginDepotAfterFork(fork_child);
+  StackDepotUnlockAfterFork(fork_child);
+  UnlockAllocator();
+  // Usually we unlock ThreadRegistry, but msan does not have one.
+}
 
 void InstallAtForkHandler() {
-  auto before = []() {
-    // Usually we lock ThreadRegistry, but msan does not have one.
-    LockAllocator();
-    StackDepotLockAll();
-    ChainedOriginDepotLockAll();
-  };
-  auto after = []() {
-    ChainedOriginDepotUnlockAll();
-    StackDepotUnlockAll();
-    UnlockAllocator();
-    // Usually we unlock ThreadRegistry, but msan does not have one.
-  };
-  pthread_atfork(before, after, after);
+  pthread_atfork(
+      &BeforeFork, []() { AfterFork(/* fork_child= */ false); },
+      []() { AfterFork(/* fork_child= */ true); });
 }
 
 } // namespace __msan
diff --git a/compiler-rt/lib/sanitizer_common/sanitizer_stackdepot.cpp b/compiler-rt/lib/sanitizer_common/sanitizer_stackdepot.cpp
index a746d4621936c6..ce21f3c178bce0 100644
--- a/compiler-rt/lib/sanitizer_common/sanitizer_stackdepot.cpp
+++ b/compiler-rt/lib/sanitizer_common/sanitizer_stackdepot.cpp
@@ -215,13 +215,13 @@ StackTrace StackDepotGet(u32 id) {
   return theDepot.Get(id);
 }
 
-void StackDepotLockAll() {
+void StackDepotLockBeforeFork() {
   theDepot.LockAll();
   compress_thread.LockAndStop();
   stackStore.LockAll();
 }
 
-void StackDepotUnlockAll() {
+void StackDepotUnlockAfterFork(bool fork_child) {
   stackStore.UnlockAll();
   compress_thread.Unlock();
   theDepot.UnlockAll();
diff --git a/compiler-rt/lib/sanitizer_common/sanitizer_stackdepot.h b/compiler-rt/lib/sanitizer_common/sanitizer_stackdepot.h
index cca6fd53468839..82cf7578d0fb9b 100644
--- a/compiler-rt/lib/sanitizer_common/sanitizer_stackdepot.h
+++ b/compiler-rt/lib/sanitizer_common/sanitizer_stackdepot.h
@@ -39,8 +39,8 @@ StackDepotHandle StackDepotPut_WithHandle(StackTrace stack);
 // Retrieves a stored stack trace by the id.
 StackTrace StackDepotGet(u32 id);
 
-void StackDepotLockAll();
-void StackDepotUnlockAll();
+void StackDepotLockBeforeFork();
+void StackDepotUnlockAfterFork(bool fork_child);
 void StackDepotPrintAll();
 void StackDepotStopBackgroundThread();