[scudo] Make BufferPool typed and always express sizes in terms of element count

[fabio-d]

This fixes the issue that resulted in getBuffer interpreting its
argument as a number of elements and getDynamicBuffer interpreting it
as a number of bytes.

[llvmbot]

@llvm/pr-subscribers-compiler-rt-sanitizer

Changes

This fixes the issue that resulted in getBuffer interpreting its
argument as a number of elements and getDynamicBuffer interpreting it
as a number of bytes.

---

Full diff: https://github.com/llvm/llvm-project/pull/66896.diff

3 Files Affected:

• (modified) compiler-rt/lib/scudo/standalone/release.cpp (+2-1)
• (modified) compiler-rt/lib/scudo/standalone/release.h (+45-37)
• (modified) compiler-rt/lib/scudo/standalone/tests/release_test.cpp (+16-15)

diff --git a/compiler-rt/lib/scudo/standalone/release.cpp b/compiler-rt/lib/scudo/standalone/release.cpp
index 938bb41faf69561..87dbc55e40df927 100644
--- a/compiler-rt/lib/scudo/standalone/release.cpp
+++ b/compiler-rt/lib/scudo/standalone/release.cpp
@@ -10,7 +10,8 @@
 
 namespace scudo {
 
-BufferPool<RegionPageMap::StaticBufferCount, RegionPageMap::StaticBufferSize>
+BufferPool<uptr, RegionPageMap::StaticBufferCount,
+           RegionPageMap::StaticBufferNumElements>
     RegionPageMap::Buffers;
 
 } // namespace scudo
diff --git a/compiler-rt/lib/scudo/standalone/release.h b/compiler-rt/lib/scudo/standalone/release.h
index e138570e7993336..98f59535e282783 100644
--- a/compiler-rt/lib/scudo/standalone/release.h
+++ b/compiler-rt/lib/scudo/standalone/release.h
@@ -95,20 +95,25 @@ class FragmentationRecorder {
   uptr ReleasedPagesCount = 0;
 };
 
-// A buffer pool which holds a fixed number of static buffers for fast buffer
-// allocation. If the request size is greater than `StaticBufferSize`, it'll
+// A buffer pool which holds a fixed number of static buffers of `T` elements
+// for fast buffer allocation. If the request size is greater than
+// `StaticBufferNumElements` or if all the static buffers are in use, it'll
 // delegate the allocation to map().
-template <uptr StaticBufferCount, uptr StaticBufferSize> class BufferPool {
+template <typename T, uptr StaticBufferCount, uptr StaticBufferNumElements>
+class BufferPool {
 public:
   // Preserve 1 bit in the `Mask` so that we don't need to do zero-check while
   // extracting the least significant bit from the `Mask`.
   static_assert(StaticBufferCount < SCUDO_WORDSIZE, "");
-  static_assert(isAligned(StaticBufferSize, SCUDO_CACHE_LINE_SIZE), "");
+  static_assert(isAligned(StaticBufferNumElements * sizeof(T),
+                          SCUDO_CACHE_LINE_SIZE),
+                "");
 
-  // Return a buffer which is at least `BufferSize`.
-  uptr *getBuffer(const uptr BufferSize) {
-    if (UNLIKELY(BufferSize > StaticBufferSize))
-      return getDynamicBuffer(BufferSize);
+  // Return a zero-initialized buffer which can contain at least the given
+  // number of elements, or nullptr on failure.
+  T *getBuffer(const uptr NumElements) {
+    if (UNLIKELY(NumElements > StaticBufferNumElements))
+      return getDynamicBuffer(NumElements);
 
     uptr index;
     {
@@ -123,32 +128,33 @@ template <uptr StaticBufferCount, uptr StaticBufferSize> class BufferPool {
     }
 
     if (index >= StaticBufferCount)
-      return getDynamicBuffer(BufferSize);
+      return getDynamicBuffer(NumElements);
 
-    const uptr Offset = index * StaticBufferSize;
-    memset(&RawBuffer[Offset], 0, StaticBufferSize);
+    const uptr Offset = index * StaticBufferNumElements;
+    memset(&RawBuffer[Offset], 0, StaticBufferNumElements * sizeof(T));
     return &RawBuffer[Offset];
   }
 
-  void releaseBuffer(uptr *Buffer, const uptr BufferSize) {
-    const uptr index = getStaticBufferIndex(Buffer, BufferSize);
+  void releaseBuffer(T *Buffer, const uptr NumElements) {
+    const uptr index = getStaticBufferIndex(Buffer, NumElements);
     if (index < StaticBufferCount) {
       ScopedLock L(Mutex);
       DCHECK_EQ((Mask & (static_cast<uptr>(1) << index)), 0U);
       Mask |= static_cast<uptr>(1) << index;
     } else {
-      unmap(reinterpret_cast<void *>(Buffer),
-            roundUp(BufferSize, getPageSizeCached()));
+      const uptr MappedSize =
+          roundUp(NumElements * sizeof(T), getPageSizeCached());
+      unmap(reinterpret_cast<void *>(Buffer), MappedSize);
     }
   }
 
-  bool isStaticBufferTestOnly(uptr *Buffer, uptr BufferSize) {
-    return getStaticBufferIndex(Buffer, BufferSize) < StaticBufferCount;
+  bool isStaticBufferTestOnly(T *Buffer, uptr NumElements) {
+    return getStaticBufferIndex(Buffer, NumElements) < StaticBufferCount;
   }
 
 private:
-  uptr getStaticBufferIndex(uptr *Buffer, uptr BufferSize) {
-    if (UNLIKELY(BufferSize > StaticBufferSize))
+  uptr getStaticBufferIndex(T *Buffer, uptr NumElements) {
+    if (UNLIKELY(NumElements > StaticBufferNumElements))
       return StaticBufferCount;
 
     const uptr BufferBase = reinterpret_cast<uptr>(Buffer);
@@ -159,32 +165,34 @@ template <uptr StaticBufferCount, uptr StaticBufferSize> class BufferPool {
       return StaticBufferCount;
     }
 
-    DCHECK_LE(BufferSize, StaticBufferSize);
-    DCHECK_LE(BufferBase + BufferSize, RawBufferBase + sizeof(RawBuffer));
-    DCHECK_EQ((BufferBase - RawBufferBase) % StaticBufferSize, 0U);
+    DCHECK_LE(NumElements, StaticBufferNumElements);
+    DCHECK_LE(BufferBase + NumElements * sizeof(T),
+              RawBufferBase + sizeof(RawBuffer));
 
-    const uptr index =
-        (BufferBase - RawBufferBase) / (StaticBufferSize * sizeof(uptr));
+    const uptr StaticBufferSize = StaticBufferNumElements * sizeof(T);
+    DCHECK_EQ((BufferBase - RawBufferBase) % StaticBufferSize, 0U);
+    const uptr index = (BufferBase - RawBufferBase) / StaticBufferSize;
     DCHECK_LT(index, StaticBufferCount);
     return index;
   }
 
-  uptr *getDynamicBuffer(const uptr BufferSize) {
+  T *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
     // is accessed. So far, this is only enabled on Fuchsia. It hasn't proven a
     // performance benefit on other platforms.
     const uptr MmapFlags = MAP_ALLOWNOMEM | (SCUDO_FUCHSIA ? MAP_PRECOMMIT : 0);
-    return reinterpret_cast<uptr *>(
-        map(nullptr, roundUp(BufferSize, getPageSizeCached()), "scudo:counters",
-            MmapFlags, &MapData));
+    const uptr MappedSize =
+        roundUp(NumElements * sizeof(T), getPageSizeCached());
+    return reinterpret_cast<T *>(
+        map(nullptr, MappedSize, "scudo:counters", MmapFlags, &MapData));
   }
 
   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 * StaticBufferSize] GUARDED_BY(Mutex);
+  T RawBuffer[StaticBufferCount * StaticBufferNumElements] GUARDED_BY(Mutex);
   [[no_unique_address]] MapPlatformData MapData = {};
 };
 
@@ -207,7 +215,7 @@ class RegionPageMap {
         PackingRatioLog(0),
         BitOffsetMask(0),
         SizePerRegion(0),
-        BufferSize(0),
+        BufferNumElements(0),
         Buffer(nullptr) {}
   RegionPageMap(uptr NumberOfRegions, uptr CountersPerRegion, uptr MaxValue) {
     reset(NumberOfRegions, CountersPerRegion, MaxValue);
@@ -215,7 +223,7 @@ class RegionPageMap {
   ~RegionPageMap() {
     if (!isAllocated())
       return;
-    Buffers.releaseBuffer(Buffer, BufferSize);
+    Buffers.releaseBuffer(Buffer, BufferNumElements);
     Buffer = nullptr;
   }
 
@@ -248,8 +256,8 @@ class RegionPageMap {
     SizePerRegion =
         roundUp(NumCounters, static_cast<uptr>(1U) << PackingRatioLog) >>
         PackingRatioLog;
-    BufferSize = SizePerRegion * sizeof(*Buffer) * Regions;
-    Buffer = Buffers.getBuffer(BufferSize);
+    BufferNumElements = SizePerRegion * Regions;
+    Buffer = Buffers.getBuffer(BufferNumElements);
   }
 
   bool isAllocated() const { return !!Buffer; }
@@ -324,7 +332,7 @@ class RegionPageMap {
     return get(Region, I) == CounterMask;
   }
 
-  uptr getBufferSize() const { return BufferSize; }
+  uptr getBufferNumElements() const { return BufferNumElements; }
 
 private:
   uptr Regions;
@@ -335,14 +343,14 @@ class RegionPageMap {
   uptr BitOffsetMask;
 
   uptr SizePerRegion;
-  uptr BufferSize;
+  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 StaticBufferSize = 512U;
-  static BufferPool<StaticBufferCount, StaticBufferSize> Buffers;
+  static const uptr StaticBufferNumElements = 512U;
+  static BufferPool<uptr, StaticBufferCount, StaticBufferNumElements> Buffers;
 };
 
 template <class ReleaseRecorderT> class FreePagesRangeTracker {
diff --git a/compiler-rt/lib/scudo/standalone/tests/release_test.cpp b/compiler-rt/lib/scudo/standalone/tests/release_test.cpp
index 0ab49460807f297..11f1bb1e37875ef 100644
--- a/compiler-rt/lib/scudo/standalone/tests/release_test.cpp
+++ b/compiler-rt/lib/scudo/standalone/tests/release_test.cpp
@@ -23,17 +23,16 @@ TEST(ScudoReleaseTest, RegionPageMap) {
     // Various valid counter's max values packed into one word.
     scudo::RegionPageMap PageMap2N(1U, 1U, 1UL << I);
     ASSERT_TRUE(PageMap2N.isAllocated());
-    EXPECT_EQ(sizeof(scudo::uptr), PageMap2N.getBufferSize());
+    EXPECT_EQ(1, PageMap2N.getBufferNumElements());
     // Check the "all bit set" values too.
     scudo::RegionPageMap PageMap2N1_1(1U, 1U, ~0UL >> I);
     ASSERT_TRUE(PageMap2N1_1.isAllocated());
-    EXPECT_EQ(sizeof(scudo::uptr), PageMap2N1_1.getBufferSize());
+    EXPECT_EQ(1, PageMap2N1_1.getBufferNumElements());
     // Verify the packing ratio, the counter is Expected to be packed into the
     // closest power of 2 bits.
     scudo::RegionPageMap PageMap(1U, SCUDO_WORDSIZE, 1UL << I);
     ASSERT_TRUE(PageMap.isAllocated());
-    EXPECT_EQ(sizeof(scudo::uptr) * scudo::roundUpPowerOfTwo(I + 1),
-              PageMap.getBufferSize());
+    EXPECT_EQ(scudo::roundUpPowerOfTwo(I + 1), PageMap.getBufferNumElements());
   }
 
   // Go through 1, 2, 4, 8, .. {32,64} bits per counter.
@@ -533,7 +532,8 @@ template <class SizeClassMap> void testReleasePartialRegion() {
                                         ReleaseBase);
       Partial.ensurePageMapAllocated();
 
-      EXPECT_GE(Full.PageMap.getBufferSize(), Partial.PageMap.getBufferSize());
+      EXPECT_GE(Full.PageMap.getBufferNumElements(),
+                Partial.PageMap.getBufferNumElements());
     }
 
     while (!FreeList.empty()) {
@@ -628,26 +628,27 @@ TEST(ScudoReleaseTest, RangeReleaseRegionWithSingleBlock) {
 
 TEST(ScudoReleaseTest, BufferPool) {
   constexpr scudo::uptr StaticBufferCount = SCUDO_WORDSIZE - 1;
-  constexpr scudo::uptr StaticBufferSize = 512U;
+  constexpr scudo::uptr StaticBufferNumElements = 512U;
 
   // Allocate the buffer pool on the heap because it is quite large (slightly
-  // more than StaticBufferCount * StaticBufferSize * sizeof(uptr)) and it may
-  // not fit in the stack on some platforms.
-  using BufferPool = scudo::BufferPool<StaticBufferCount, StaticBufferSize>;
+  // more than StaticBufferCount * StaticBufferNumElements * sizeof(uptr)) and
+  // it may not fit in the stack on some platforms.
+  using BufferPool = scudo::BufferPool<scudo::uptr, StaticBufferCount,
+                                       StaticBufferNumElements>;
   std::unique_ptr<BufferPool> Pool(new BufferPool());
 
   std::vector<std::pair<scudo::uptr *, scudo::uptr>> Buffers;
   for (scudo::uptr I = 0; I < StaticBufferCount; ++I) {
-    scudo::uptr *P = Pool->getBuffer(StaticBufferSize);
-    EXPECT_TRUE(Pool->isStaticBufferTestOnly(P, StaticBufferSize));
-    Buffers.emplace_back(P, StaticBufferSize);
+    scudo::uptr *P = Pool->getBuffer(StaticBufferNumElements);
+    EXPECT_TRUE(Pool->isStaticBufferTestOnly(P, StaticBufferNumElements));
+    Buffers.emplace_back(P, StaticBufferNumElements);
   }
 
   // The static buffer is supposed to be used up.
-  scudo::uptr *P = Pool->getBuffer(StaticBufferSize);
-  EXPECT_FALSE(Pool->isStaticBufferTestOnly(P, StaticBufferSize));
+  scudo::uptr *P = Pool->getBuffer(StaticBufferNumElements);
+  EXPECT_FALSE(Pool->isStaticBufferTestOnly(P, StaticBufferNumElements));
 
-  Pool->releaseBuffer(P, StaticBufferSize);
+  Pool->releaseBuffer(P, StaticBufferNumElements);
   for (auto &Buffer : Buffers)
     Pool->releaseBuffer(Buffer.first, Buffer.second);
 }

[hctim]

Looks reasonable to me, two things:

1. Seems unnecessary to template the uptr type of the buffer class, unless you can elaborate on why you need that?
2. Can you please add a small justification in your commit message about how you discovered the mismatch? It's nice to know the backstory here.

[fabio-d]

Hi @hctim, thanks for reviewing :) I'll reply in reverse order:

2 -> Unfortunately I don't have a nice backstory as I'm not aware of any nasty issues causes by this. 😄 I've just found the mismatch while migrating BufferPool to the new platform abstraction layer (i.e. from map/unmap to MemMap, see #66788)

1 -> Yes, the only user is RegionPageMap, in the same .h file, and it only instantiates it with uptr elements. The idea to use a template came up while discussing with @ChiaHungDuan on whether to fix the issue by uniformly interpreting arguments as a number of bytes or as a number of elements.
While we initially considered always interpreting them as a number of bytes and returning void* from getBuffer, we realized that that it would have made the class harder to use (because the caller would be responsible for always requesting a size that is multiple of the intended element size) and also to implement (because the underlying storage would have had to be a simple byte array, to avoid aliasing issues, while suitably aligned for the intended type). The conclusion was that requesting the type as a template and always interpreting sizes as number of elements make both things simpler.

[hctim]

1 -> Yes, the only user is RegionPageMap, in the same .h file, and it only instantiates it with uptr elements. The idea to use a template came up while discussing with @ChiaHungDuan on whether to fix the issue by uniformly interpreting arguments as a number of bytes or as a number of elements. While we initially considered always interpreting them as a number of bytes and returning void* from getBuffer, we realized that that it would have made the class harder to use (because the caller would be responsible for always requesting a size that is multiple of the intended element size) and also to implement (because the underlying storage would have had to be a simple byte array, to avoid aliasing issues, while suitably aligned for the intended type). The conclusion was that requesting the type as a template and always interpreting sizes as number of elements make both things simpler.

IMHO there's enough templates already floating around in scudo. If there's no plans to add additional callsites, and it seems fine that getBuffer can return uptr*, seems like a good idea to just remove the template, and make this patch a pure migration from BufferSize -> NumElements.

[ChiaHungDuan]

There's plan for different types in BufferPool but I agree that we can stick on uptr at this moment

[github-actions]

✅ With the latest revision this PR passed the C/C++ code formatter.

[ChiaHungDuan]

Could you also update the title?

[ChiaHungDuan]

Feel free to submit it after the title/commit message updated