4 files changed, 17 insertions, 12 deletions
diff --git a/lib/lsan/lsan_common.cpp b/lib/lsan/lsan_common.cpp
index 86bd120e7..9ff9f4c5d 100644
--- a/lib/lsan/lsan_common.cpp
+++ b/lib/lsan/lsan_common.cpp
@@ -570,11 +570,7 @@ static bool CheckForLeaks() {
   EnsureMainThreadIDIsCorrect();
   CheckForLeaksParam param;
   param.success = false;
-  LockThreadRegistry();
-  LockAllocator();
-  DoStopTheWorld(CheckForLeaksCallback, &param);
-  UnlockAllocator();
-  UnlockThreadRegistry();
+  LockStuffAndStopTheWorld(CheckForLeaksCallback, &param);
 
   if (!param.success) {
     Report("LeakSanitizer has encountered a fatal error.\n");
diff --git a/lib/lsan/lsan_common.h b/lib/lsan/lsan_common.h
index 58dc00faa..d24abe31b 100644
--- a/lib/lsan/lsan_common.h
+++ b/lib/lsan/lsan_common.h
@@ -129,8 +129,9 @@ struct RootRegion {
 InternalMmapVector<RootRegion> const *GetRootRegions();
 void ScanRootRegion(Frontier *frontier, RootRegion const &region,
                     uptr region_begin, uptr region_end, bool is_readable);
-// Run stoptheworld while holding any platform-specific locks.
-void DoStopTheWorld(StopTheWorldCallback callback, void* argument);
+// Run stoptheworld while holding any platform-specific locks, as well as the
+// allocator and thread registry locks.
+void LockStuffAndStopTheWorld(StopTheWorldCallback callback, void* argument);
 
 void ScanRangeForPointers(uptr begin, uptr end,
                           Frontier *frontier,
diff --git a/lib/lsan/lsan_common_linux.cpp b/lib/lsan/lsan_common_linux.cpp
index 9ce27a983..ea1a4a2f5 100644
--- a/lib/lsan/lsan_common_linux.cpp
+++ b/lib/lsan/lsan_common_linux.cpp
@@ -115,10 +115,14 @@ void HandleLeaks() {
   if (common_flags()->exitcode) Die();
 }
 
-static int DoStopTheWorldCallback(struct dl_phdr_info *info, size_t size,
-                                  void *data) {
+static int LockStuffAndStopTheWorldCallback(struct dl_phdr_info *info,
+                                            size_t size, void *data) {
+  LockThreadRegistry();
+  LockAllocator();
   DoStopTheWorldParam *param = reinterpret_cast<DoStopTheWorldParam *>(data);
   StopTheWorld(param->callback, param->argument);
+  UnlockAllocator();
+  UnlockThreadRegistry();
   return 1;
 }
 
@@ -130,9 +134,9 @@ static int DoStopTheWorldCallback(struct dl_phdr_info *info, size_t size,
 // while holding the libdl lock in the parent thread, we can safely reenter it
 // in the tracer. The solution is to run stoptheworld from a dl_iterate_phdr()
 // callback in the parent thread.
-void DoStopTheWorld(StopTheWorldCallback callback, void *argument) {
+void LockStuffAndStopTheWorld(StopTheWorldCallback callback, void *argument) {
   DoStopTheWorldParam param = {callback, argument};
-  dl_iterate_phdr(DoStopTheWorldCallback, &param);
+  dl_iterate_phdr(LockStuffAndStopTheWorldCallback, &param);
 }
 
 } // namespace __lsan
diff --git a/lib/lsan/lsan_common_mac.cpp b/lib/lsan/lsan_common_mac.cpp
index 5204a6624..c1804e93c 100644
--- a/lib/lsan/lsan_common_mac.cpp
+++ b/lib/lsan/lsan_common_mac.cpp
@@ -193,8 +193,12 @@ void ProcessPlatformSpecificAllocations(Frontier *frontier) {
 // causes rare race conditions.
 void HandleLeaks() {}
 
-void DoStopTheWorld(StopTheWorldCallback callback, void *argument) {
+void LockStuffAndStopTheWorld(StopTheWorldCallback callback, void *argument) {
+  LockThreadRegistry();
+  LockAllocator();
   StopTheWorld(callback, argument);
+  UnlockAllocator();
+  UnlockThreadRegistry();
 }
 
 } // namespace __lsan