[scudo] Use MemMap in BufferPool and RegionPageMap (#66788)

Author: fabio-d

compiler-rt/lib/scudo/standalone/release.h

@@ -109,9 +109,22 @@ class BufferPool {
                           SCUDO_CACHE_LINE_SIZE),
                 "");
 
+  struct Buffer {
+    // Pointer to the buffer's memory, or nullptr if no buffer was allocated.
+    uptr *Data = nullptr;
+
+    // The index of the underlying static buffer, or StaticBufferCount if this
+    // buffer was dynamically allocated. This value is initially set to a poison
+    // value to aid debugging.
+    uptr BufferIndex = ~static_cast<uptr>(0);
+
+    // Only valid if BufferIndex == StaticBufferCount.
+    MemMapT MemMap = {};
+  };
+
   // Return a zero-initialized buffer which can contain at least the given
   // number of elements, or nullptr on failure.
-  uptr *getBuffer(const uptr NumElements) {
+  Buffer getBuffer(const uptr NumElements) {
     if (UNLIKELY(NumElements > StaticBufferNumElements))
       return getDynamicBuffer(NumElements);
 
@@ -130,53 +143,33 @@ class BufferPool {
     if (index >= StaticBufferCount)
       return getDynamicBuffer(NumElements);
 
-    const uptr Offset = index * StaticBufferNumElements;
-    memset(&RawBuffer[Offset], 0, StaticBufferNumElements * sizeof(uptr));
-    return &RawBuffer[Offset];
+    Buffer Buf;
+    Buf.Data = &RawBuffer[index * StaticBufferNumElements];
+    Buf.BufferIndex = index;
+    memset(Buf.Data, 0, StaticBufferNumElements * sizeof(uptr));
+    return Buf;
   }
 
-  void releaseBuffer(uptr *Buffer, const uptr NumElements) {
-    const uptr index = getStaticBufferIndex(Buffer, NumElements);
-    if (index < StaticBufferCount) {
+  void releaseBuffer(Buffer Buf) {
+    DCHECK_NE(Buf.Data, nullptr);
+    DCHECK_LE(Buf.BufferIndex, StaticBufferCount);
+    if (Buf.BufferIndex != StaticBufferCount) {
       ScopedLock L(Mutex);
-      DCHECK_EQ((Mask & (static_cast<uptr>(1) << index)), 0U);
-      Mask |= static_cast<uptr>(1) << index;
+      DCHECK_EQ((Mask & (static_cast<uptr>(1) << Buf.BufferIndex)), 0U);
+      Mask |= static_cast<uptr>(1) << Buf.BufferIndex;
     } else {
-      const uptr MappedSize =
-          roundUp(NumElements * sizeof(uptr), getPageSizeCached());
-      unmap(reinterpret_cast<void *>(Buffer), MappedSize);
+      Buf.MemMap.unmap(Buf.MemMap.getBase(), Buf.MemMap.getCapacity());
     }
   }
 
-  bool isStaticBufferTestOnly(uptr *Buffer, uptr NumElements) {
-    return getStaticBufferIndex(Buffer, NumElements) < StaticBufferCount;
+  bool isStaticBufferTestOnly(const Buffer &Buf) {
+    DCHECK_NE(Buf.Data, nullptr);
+    DCHECK_LE(Buf.BufferIndex, StaticBufferCount);
+    return Buf.BufferIndex != StaticBufferCount;
   }
 
 private:
-  uptr getStaticBufferIndex(uptr *Buffer, uptr NumElements) {
-    if (UNLIKELY(NumElements > StaticBufferNumElements))
-      return StaticBufferCount;
-
-    const uptr BufferBase = reinterpret_cast<uptr>(Buffer);
-    const uptr RawBufferBase = reinterpret_cast<uptr>(RawBuffer);
-
-    if (BufferBase < RawBufferBase ||
-        BufferBase >= RawBufferBase + sizeof(RawBuffer)) {
-      return StaticBufferCount;
-    }
-
-    DCHECK_LE(NumElements, StaticBufferNumElements);
-    DCHECK_LE(BufferBase + NumElements * sizeof(uptr),
-              RawBufferBase + sizeof(RawBuffer));
-
-    const uptr StaticBufferSize = StaticBufferNumElements * sizeof(uptr);
-    DCHECK_EQ((BufferBase - RawBufferBase) % StaticBufferSize, 0U);
-    const uptr index = (BufferBase - RawBufferBase) / StaticBufferSize;
-    DCHECK_LT(index, StaticBufferCount);
-    return index;
-  }
-
-  uptr *getDynamicBuffer(const uptr NumElements) {
+  Buffer getDynamicBuffer(const uptr NumElements) {
     // When using a heap-based buffer, precommit the pages backing the
     // Vmar by passing |MAP_PRECOMMIT| flag. This allows an optimization
     // where page fault exceptions are skipped as the allocated memory
@@ -185,15 +178,18 @@ class BufferPool {
     const uptr MmapFlags = MAP_ALLOWNOMEM | (SCUDO_FUCHSIA ? MAP_PRECOMMIT : 0);
     const uptr MappedSize =
         roundUp(NumElements * sizeof(uptr), getPageSizeCached());
-    return reinterpret_cast<uptr *>(
-        map(nullptr, MappedSize, "scudo:counters", MmapFlags, &MapData));
+    Buffer Buf;
+    if (Buf.MemMap.map(/*Addr=*/0, MappedSize, "scudo:counters", MmapFlags)) {
+      Buf.Data = reinterpret_cast<uptr *>(Buf.MemMap.getBase());
+      Buf.BufferIndex = StaticBufferCount;
+    }
+    return Buf;
   }
 
   HybridMutex Mutex;
   // '1' means that buffer index is not used. '0' means the buffer is in use.
   uptr Mask GUARDED_BY(Mutex) = ~static_cast<uptr>(0);
   uptr RawBuffer[StaticBufferCount * StaticBufferNumElements] GUARDED_BY(Mutex);
-  [[no_unique_address]] MapPlatformData MapData = {};
 };
 
 // A Region page map is used to record the usage of pages in the regions. It
@@ -210,15 +206,15 @@ class RegionPageMap {
   RegionPageMap()
       : Regions(0), NumCounters(0), CounterSizeBitsLog(0), CounterMask(0),
         PackingRatioLog(0), BitOffsetMask(0), SizePerRegion(0),
-        BufferNumElements(0), Buffer(nullptr) {}
+        BufferNumElements(0) {}
   RegionPageMap(uptr NumberOfRegions, uptr CountersPerRegion, uptr MaxValue) {
     reset(NumberOfRegions, CountersPerRegion, MaxValue);
   }
   ~RegionPageMap() {
     if (!isAllocated())
       return;
-    Buffers.releaseBuffer(Buffer, BufferNumElements);
-    Buffer = nullptr;
+    Buffers.releaseBuffer(Buffer);
+    Buffer = {};
   }
 
   // Lock of `StaticBuffer` is acquired conditionally and there's no easy way to
@@ -233,7 +229,7 @@ class RegionPageMap {
     Regions = NumberOfRegion;
     NumCounters = CountersPerRegion;
 
-    constexpr uptr MaxCounterBits = sizeof(*Buffer) * 8UL;
+    constexpr uptr MaxCounterBits = sizeof(*Buffer.Data) * 8UL;
     // Rounding counter storage size up to the power of two allows for using
     // bit shifts calculating particular counter's Index and offset.
     const uptr CounterSizeBits =
@@ -254,7 +250,7 @@ class RegionPageMap {
     Buffer = Buffers.getBuffer(BufferNumElements);
   }
 
-  bool isAllocated() const { return !!Buffer; }
+  bool isAllocated() const { return Buffer.Data != nullptr; }
 
   uptr getCount() const { return NumCounters; }
 
@@ -263,7 +259,8 @@ class RegionPageMap {
     DCHECK_LT(I, NumCounters);
     const uptr Index = I >> PackingRatioLog;
     const uptr BitOffset = (I & BitOffsetMask) << CounterSizeBitsLog;
-    return (Buffer[Region * SizePerRegion + Index] >> BitOffset) & CounterMask;
+    return (Buffer.Data[Region * SizePerRegion + Index] >> BitOffset) &
+           CounterMask;
   }
 
   void inc(uptr Region, uptr I) const {
@@ -272,8 +269,8 @@ class RegionPageMap {
     const uptr BitOffset = (I & BitOffsetMask) << CounterSizeBitsLog;
     DCHECK_LT(BitOffset, SCUDO_WORDSIZE);
     DCHECK_EQ(isAllCounted(Region, I), false);
-    Buffer[Region * SizePerRegion + Index] += static_cast<uptr>(1U)
-                                              << BitOffset;
+    Buffer.Data[Region * SizePerRegion + Index] += static_cast<uptr>(1U)
+                                                   << BitOffset;
   }
 
   void incN(uptr Region, uptr I, uptr N) const {
@@ -284,7 +281,7 @@ class RegionPageMap {
     const uptr BitOffset = (I & BitOffsetMask) << CounterSizeBitsLog;
     DCHECK_LT(BitOffset, SCUDO_WORDSIZE);
     DCHECK_EQ(isAllCounted(Region, I), false);
-    Buffer[Region * SizePerRegion + Index] += N << BitOffset;
+    Buffer.Data[Region * SizePerRegion + Index] += N << BitOffset;
   }
 
   void incRange(uptr Region, uptr From, uptr To) const {
@@ -303,7 +300,7 @@ class RegionPageMap {
     const uptr Index = I >> PackingRatioLog;
     const uptr BitOffset = (I & BitOffsetMask) << CounterSizeBitsLog;
     DCHECK_LT(BitOffset, SCUDO_WORDSIZE);
-    Buffer[Region * SizePerRegion + Index] |= CounterMask << BitOffset;
+    Buffer.Data[Region * SizePerRegion + Index] |= CounterMask << BitOffset;
   }
   void setAsAllCountedRange(uptr Region, uptr From, uptr To) const {
     DCHECK_LE(From, To);
@@ -329,6 +326,13 @@ class RegionPageMap {
   uptr getBufferNumElements() const { return BufferNumElements; }
 
 private:
+  // We may consider making this configurable if there are cases which may
+  // benefit from this.
+  static const uptr StaticBufferCount = 2U;
+  static const uptr StaticBufferNumElements = 512U;
+  using BufferPoolT = BufferPool<StaticBufferCount, StaticBufferNumElements>;
+  static BufferPoolT Buffers;
+
   uptr Regions;
   uptr NumCounters;
   uptr CounterSizeBitsLog;
@@ -338,13 +342,7 @@ class RegionPageMap {
 
   uptr SizePerRegion;
   uptr BufferNumElements;
-  uptr *Buffer;
-
-  // We may consider making this configurable if there are cases which may
-  // benefit from this.
-  static const uptr StaticBufferCount = 2U;
-  static const uptr StaticBufferNumElements = 512U;
-  static BufferPool<StaticBufferCount, StaticBufferNumElements> Buffers;
+  BufferPoolT::Buffer Buffer;
 };
 
 template <class ReleaseRecorderT> class FreePagesRangeTracker {

compiler-rt/lib/scudo/standalone/tests/release_test.cpp

@@ -637,18 +637,18 @@ TEST(ScudoReleaseTest, BufferPool) {
       scudo::BufferPool<StaticBufferCount, StaticBufferNumElements>;
   std::unique_ptr<BufferPool> Pool(new BufferPool());
 
-  std::vector<std::pair<scudo::uptr *, scudo::uptr>> Buffers;
+  std::vector<BufferPool::Buffer> Buffers;
   for (scudo::uptr I = 0; I < StaticBufferCount; ++I) {
-    scudo::uptr *P = Pool->getBuffer(StaticBufferNumElements);
-    EXPECT_TRUE(Pool->isStaticBufferTestOnly(P, StaticBufferNumElements));
-    Buffers.emplace_back(P, StaticBufferNumElements);
+    BufferPool::Buffer Buffer = Pool->getBuffer(StaticBufferNumElements);
+    EXPECT_TRUE(Pool->isStaticBufferTestOnly(Buffer));
+    Buffers.push_back(Buffer);
   }
 
   // The static buffer is supposed to be used up.
-  scudo::uptr *P = Pool->getBuffer(StaticBufferNumElements);
-  EXPECT_FALSE(Pool->isStaticBufferTestOnly(P, StaticBufferNumElements));
+  BufferPool::Buffer Buffer = Pool->getBuffer(StaticBufferNumElements);
+  EXPECT_FALSE(Pool->isStaticBufferTestOnly(Buffer));
 
-  Pool->releaseBuffer(P, StaticBufferNumElements);
+  Pool->releaseBuffer(Buffer);
   for (auto &Buffer : Buffers)
-    Pool->releaseBuffer(Buffer.first, Buffer.second);
+    Pool->releaseBuffer(Buffer);
 }