[NFC][sanitizer] Rename Lock{Before,After}Fork suffixes locking StackDepotBase (#76279)

This is preparation 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.

Author: vitalybuka

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() {

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

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

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;
 }
 

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() {

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

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();
 }
 

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
 

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

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();

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();