From 1dcb3874369f9640fe9ccb11a0cdf861636a4d77 Mon Sep 17 00:00:00 2001
From: ClaytonKnittel <35512940+ClaytonKnittel@users.noreply.github.com>
Date: Sun, 17 Nov 2024 14:34:01 -0800
Subject: [PATCH] Guard TestSysAlloc map with mutex. (#76)

* Guard TestSysAlloc map with mutex.

* Fix mutex cycle.
---
 src/ckmalloc/BUILD      |  1 +
 src/ckmalloc/testlib.cc | 28 ++++++++++++++++++++--------
 src/ckmalloc/testlib.h  | 23 ++++++++++++++---------
 3 files changed, 35 insertions(+), 17 deletions(-)

diff --git a/src/ckmalloc/BUILD b/src/ckmalloc/BUILD
index 5f112e26..126193c3 100644
--- a/src/ckmalloc/BUILD
+++ b/src/ckmalloc/BUILD
@@ -624,6 +624,7 @@ cc_library(
         "@abseil-cpp//absl/status:statusor",
         "@abseil-cpp//absl/strings",
         "@abseil-cpp//absl/strings:str_format",
+        "@abseil-cpp//absl/synchronization",
     ],
 )
 
diff --git a/src/ckmalloc/testlib.cc b/src/ckmalloc/testlib.cc
index 1389bb7f..7f87df3e 100644
--- a/src/ckmalloc/testlib.cc
+++ b/src/ckmalloc/testlib.cc
@@ -6,6 +6,7 @@
 #include <vector>
 
 #include "absl/status/statusor.h"
+#include "absl/synchronization/mutex.h"
 
 #include "src/ckmalloc/common.h"
 #include "src/ckmalloc/slab.h"
@@ -122,13 +123,16 @@ TestHeap* RandomHeapFromFactory(bench::HeapFactory& heap_factory) {
 
 TestSysAlloc::TestSysAlloc(bench::HeapFactory* heap_factory)
     : heap_factory_(heap_factory) {
-  heap_factory_->WithInstances<void>([this](const auto& instances) {
-    for (const auto& heap : instances) {
-      // Assume all already-created heaps are metadata heaps.
-      heap_map_.emplace(heap->Start(),
-                        std::make_pair(HeapType::kMetadataHeap, heap.get()));
-    }
-  });
+  absl::MutexLock lock(&mutex_);
+  heap_factory_->WithInstances<void>(
+      [this](const auto& instances) CK_EXCLUSIVE_LOCKS_REQUIRED(mutex_) {
+        for (const auto& heap : instances) {
+          // Assume all already-created heaps are metadata heaps.
+          heap_map_.emplace(
+              heap->Start(),
+              std::make_pair(HeapType::kMetadataHeap, heap.get()));
+        }
+      });
 }
 
 /* static */
@@ -160,12 +164,16 @@ void* TestSysAlloc::Mmap(void* start_hint, size_t size, HeapType type) {
 
   bench::Heap* heap = result.value();
   void* heap_start = heap->Start();
-  heap_map_.emplace(heap_start, std::make_pair(type, heap));
+  {
+    absl::MutexLock lock(&mutex_);
+    heap_map_.emplace(heap_start, std::make_pair(type, heap));
+  }
 
   return heap_start;
 }
 
 void TestSysAlloc::Munmap(void* ptr, size_t size) {
+  absl::MutexLock lock(&mutex_);
   auto it = heap_map_.find(ptr);
   CK_ASSERT_TRUE(it != heap_map_.end());
   bench::Heap* heap = it->second.second;
@@ -188,20 +196,24 @@ void TestSysAlloc::Sbrk(void* heap_start, size_t increment, void* current_end) {
 }
 
 bench::Heap* TestSysAlloc::HeapFromStart(void* heap_start) {
+  absl::MutexLock lock(&mutex_);
   auto it = heap_map_.find(heap_start);
   CK_ASSERT_TRUE(it != heap_map_.end());
   return it->second.second;
 }
 
 size_t TestSysAlloc::Size() const {
+  absl::MutexLock lock(&mutex_);
   return heap_map_.size();
 }
 
 TestSysAlloc::const_iterator TestSysAlloc::begin() const {
+  absl::MutexLock lock(&mutex_);
   return heap_map_.begin();
 }
 
 TestSysAlloc::const_iterator TestSysAlloc::end() const {
+  absl::MutexLock lock(&mutex_);
   return heap_map_.end();
 }
 
diff --git a/src/ckmalloc/testlib.h b/src/ckmalloc/testlib.h
index 8d56034a..17ebbbaa 100644
--- a/src/ckmalloc/testlib.h
+++ b/src/ckmalloc/testlib.h
@@ -9,6 +9,7 @@
 #include "absl/status/statusor.h"
 #include "absl/strings/str_cat.h"
 #include "absl/strings/str_format.h"
+#include "absl/synchronization/mutex.h"
 
 #include "src/ckmalloc/block.h"
 #include "src/ckmalloc/common.h"
@@ -248,27 +249,31 @@ class TestSysAlloc : public SysAlloc {
 
   static void Reset();
 
-  void* Mmap(void* start_hint, size_t size, HeapType type) override;
+  void* Mmap(void* start_hint, size_t size, HeapType type) override
+      CK_LOCKS_EXCLUDED(mutex_);
 
-  void Munmap(void* ptr, size_t size) override;
+  void Munmap(void* ptr, size_t size) override CK_LOCKS_EXCLUDED(mutex_);
 
-  void Sbrk(void* heap_start, size_t increment, void* current_end) override;
+  void Sbrk(void* heap_start, size_t increment, void* current_end) override
+      CK_LOCKS_EXCLUDED(mutex_);
 
-  bench::Heap* HeapFromStart(void* heap_start);
+  bench::Heap* HeapFromStart(void* heap_start) CK_LOCKS_EXCLUDED(mutex_);
 
-  size_t Size() const;
+  size_t Size() const CK_LOCKS_EXCLUDED(mutex_);
 
-  const_iterator begin() const;
-  const_iterator end() const;
+  const_iterator begin() const CK_LOCKS_EXCLUDED(mutex_);
+  const_iterator end() const CK_LOCKS_EXCLUDED(mutex_);
 
-  auto Find(void* heap_start) const {
+  auto Find(void* heap_start) const CK_LOCKS_EXCLUDED(mutex_) {
+    absl::MutexLock lock(&mutex_);
     return heap_map_.find(heap_start);
   }
 
  private:
   bench::HeapFactory* heap_factory_;
 
-  MapT heap_map_;
+  mutable absl::Mutex mutex_;
+  MapT heap_map_ CK_GUARDED_BY(mutex_);
 };
 
 class CkMallocTest {