[compiler-rt] [scudo] Add two missing locks to enable/disable. (PR #79670)

[Evgenii Stepanov via llvm-commits]

https://github.com/eugenis created https://github.com/llvm/llvm-project/pull/79670

The StackDepot lock was actually observed deadlocking in the code that uses malloc after fork but before exec. The other one was found by code inspection.

>From cccdbb69ed588099aa1366c17514a1c4bfb97367 Mon Sep 17 00:00:00 2001
From: Evgenii Stepanov <eugenis at google.com>
Date: Fri, 26 Jan 2024 16:29:13 -0800
Subject: [PATCH] [scudo] Add two missing locks to enable/disable.

The StackDepot lock was actually observed deadlocking in the code that uses
malloc after fork but before exec. The other one was found by code
inspection.
---
 compiler-rt/lib/scudo/standalone/combined.h    |  2 ++
 compiler-rt/lib/scudo/standalone/primary32.h   |  2 ++
 compiler-rt/lib/scudo/standalone/primary64.h   |  2 ++
 compiler-rt/lib/scudo/standalone/release.h     | 16 ++++++++++++++++
 compiler-rt/lib/scudo/standalone/stack_depot.h |  8 ++++++++
 5 files changed, 30 insertions(+)

diff --git a/compiler-rt/lib/scudo/standalone/combined.h b/compiler-rt/lib/scudo/standalone/combined.h
index 4624f83d142a0d..80774073522a72 100644
--- a/compiler-rt/lib/scudo/standalone/combined.h
+++ b/compiler-rt/lib/scudo/standalone/combined.h
@@ -691,10 +691,12 @@ class Allocator {
     Quarantine.disable();
     Primary.disable();
     Secondary.disable();
+    Depot.disable();
   }
 
   void enable() NO_THREAD_SAFETY_ANALYSIS {
     initThreadMaybe();
+    Depot.enable();
     Secondary.enable();
     Primary.enable();
     Quarantine.enable();
diff --git a/compiler-rt/lib/scudo/standalone/primary32.h b/compiler-rt/lib/scudo/standalone/primary32.h
index 4d03b282d000de..89fffc50dcd009 100644
--- a/compiler-rt/lib/scudo/standalone/primary32.h
+++ b/compiler-rt/lib/scudo/standalone/primary32.h
@@ -261,6 +261,7 @@ template <typename Config> class SizeClassAllocator32 {
   }
 
   void disable() NO_THREAD_SAFETY_ANALYSIS {
+    RegionPageMap::disable();
     // The BatchClassId must be locked last since other classes can use it.
     for (sptr I = static_cast<sptr>(NumClasses) - 1; I >= 0; I--) {
       if (static_cast<uptr>(I) == SizeClassMap::BatchClassId)
@@ -281,6 +282,7 @@ template <typename Config> class SizeClassAllocator32 {
         continue;
       getSizeClassInfo(I)->Mutex.unlock();
     }
+    RegionPageMap::enable();
   }
 
   template <typename F> void iterateOverBlocks(F Callback) {
diff --git a/compiler-rt/lib/scudo/standalone/primary64.h b/compiler-rt/lib/scudo/standalone/primary64.h
index 9a642d23620e1e..4adb53b3748ffe 100644
--- a/compiler-rt/lib/scudo/standalone/primary64.h
+++ b/compiler-rt/lib/scudo/standalone/primary64.h
@@ -341,9 +341,11 @@ template <typename Config> class SizeClassAllocator64 {
     }
     getRegionInfo(SizeClassMap::BatchClassId)->MMLock.lock();
     getRegionInfo(SizeClassMap::BatchClassId)->FLLock.lock();
+    RegionPageMap::disable();
   }
 
   void enable() NO_THREAD_SAFETY_ANALYSIS {
+    RegionPageMap::enable();
     getRegionInfo(SizeClassMap::BatchClassId)->FLLock.unlock();
     getRegionInfo(SizeClassMap::BatchClassId)->MMLock.unlock();
     for (uptr I = 0; I < NumClasses; I++) {
diff --git a/compiler-rt/lib/scudo/standalone/release.h b/compiler-rt/lib/scudo/standalone/release.h
index b6f76a4d20585a..6e90ec97e0f3c3 100644
--- a/compiler-rt/lib/scudo/standalone/release.h
+++ b/compiler-rt/lib/scudo/standalone/release.h
@@ -168,6 +168,14 @@ class BufferPool {
     return Buf.BufferIndex != StaticBufferCount;
   }
 
+  void disable() NO_THREAD_SAFETY_ANALYSIS {
+    Mutex.lock();
+  }
+
+  void enable() NO_THREAD_SAFETY_ANALYSIS {
+    Mutex.unlock();
+  }
+
 private:
   Buffer getDynamicBuffer(const uptr NumElements) {
     // When using a heap-based buffer, precommit the pages backing the
@@ -325,6 +333,14 @@ class RegionPageMap {
 
   uptr getBufferNumElements() const { return BufferNumElements; }
 
+  static void disable() NO_THREAD_SAFETY_ANALYSIS {
+    Buffers.disable();
+  }
+
+  static void enable() NO_THREAD_SAFETY_ANALYSIS {
+    Buffers.enable();
+  }
+
 private:
   // We may consider making this configurable if there are cases which may
   // benefit from this.
diff --git a/compiler-rt/lib/scudo/standalone/stack_depot.h b/compiler-rt/lib/scudo/standalone/stack_depot.h
index 12c35eb2a4f338..ab5e96680cab58 100644
--- a/compiler-rt/lib/scudo/standalone/stack_depot.h
+++ b/compiler-rt/lib/scudo/standalone/stack_depot.h
@@ -136,6 +136,14 @@ class StackDepot {
   u64 operator[](uptr RingPos) const {
     return atomic_load_relaxed(&Ring[RingPos & RingMask]);
   }
+
+  void disable() NO_THREAD_SAFETY_ANALYSIS {
+    RingEndMu.lock();
+  }
+
+  void enable() NO_THREAD_SAFETY_ANALYSIS {
+    RingEndMu.unlock();
+  }
 };
 
 } // namespace scudo