[compiler-rt] [scudo] Pass the max number of blocks to popBlocks (PR #70243)

[via llvm-commits]

https://github.com/ChiaHungDuan updated https://github.com/llvm/llvm-project/pull/70243

>From ae8b0ae6b40117bb988ce5b6b2d6bfef7bd18b5f Mon Sep 17 00:00:00 2001
From: Chia-hung Duan <chiahungduan at google.com>
Date: Wed, 25 Oct 2023 18:52:45 +0000
Subject: [PATCH 1/2] [scudo] Pass the max number of blocks to popBlocks

Make the cache have the fully control on how many blocks to be popped
(At before, it depended the number of blocks stored in the
TransferBatch)
---
 compiler-rt/lib/scudo/standalone/local_cache.h | 16 ++++++++--------
 compiler-rt/lib/scudo/standalone/primary32.h   |  4 +++-
 compiler-rt/lib/scudo/standalone/primary64.h   |  4 +++-
 3 files changed, 14 insertions(+), 10 deletions(-)

diff --git a/compiler-rt/lib/scudo/standalone/local_cache.h b/compiler-rt/lib/scudo/standalone/local_cache.h
index 1814272277ff436..6898840eb2bcba4 100644
--- a/compiler-rt/lib/scudo/standalone/local_cache.h
+++ b/compiler-rt/lib/scudo/standalone/local_cache.h
@@ -40,7 +40,11 @@ template <class SizeClassAllocator> struct SizeClassAllocatorLocalCache {
     DCHECK_LT(ClassId, NumClasses);
     PerClass *C = &PerClassArray[ClassId];
     if (C->Count == 0) {
-      if (UNLIKELY(!refill(C, ClassId)))
+      initCacheMaybe(C);
+
+      // Refill half of the number of max cached.
+      DCHECK_GT(C->MaxCount / 2, 0U);
+      if (UNLIKELY(!refill(C, ClassId, C->MaxCount / 2)))
         return nullptr;
       DCHECK_GT(C->Count, 0);
     }
@@ -173,14 +177,10 @@ template <class SizeClassAllocator> struct SizeClassAllocatorLocalCache {
       deallocate(BatchClassId, B);
   }
 
-  NOINLINE bool refill(PerClass *C, uptr ClassId) {
-    initCacheMaybe(C);
-
-    // TODO(chiahungduan): Pass the max number cached for each size class.
+  NOINLINE bool refill(PerClass *C, uptr ClassId, u16 MaxRefill) {
     const u16 NumBlocksRefilled =
-        Allocator->popBlocks(this, ClassId, C->Chunks);
-    DCHECK_LE(NumBlocksRefilled,
-              getMaxCached(SizeClassAllocator::getSizeByClassId(ClassId)));
+        Allocator->popBlocks(this, ClassId, C->Chunks, MaxRefill);
+    DCHECK_LE(NumBlocksRefilled, MaxRefill);
     C->Count = static_cast<u16>(C->Count + NumBlocksRefilled);
     return NumBlocksRefilled != 0;
   }
diff --git a/compiler-rt/lib/scudo/standalone/primary32.h b/compiler-rt/lib/scudo/standalone/primary32.h
index c900550ac675e47..137ac2f02ebd95d 100644
--- a/compiler-rt/lib/scudo/standalone/primary32.h
+++ b/compiler-rt/lib/scudo/standalone/primary32.h
@@ -191,13 +191,15 @@ template <typename Config> class SizeClassAllocator32 {
     return BlockSize > PageSize;
   }
 
-  u16 popBlocks(CacheT *C, uptr ClassId, CompactPtrT *ToArray) {
+  u16 popBlocks(CacheT *C, uptr ClassId, CompactPtrT *ToArray,
+                const u16 MaxBlockCount) {
     TransferBatchT *B = popBatch(C, ClassId);
     if (!B)
       return 0;
 
     const u16 Count = B->getCount();
     DCHECK_GT(Count, 0U);
+    DCHECK_LE(Count, MaxBlockCount);
     B->moveToArray(ToArray);
 
     if (ClassId != SizeClassMap::BatchClassId)
diff --git a/compiler-rt/lib/scudo/standalone/primary64.h b/compiler-rt/lib/scudo/standalone/primary64.h
index 0f4ba88ee1f4b9d..db404e40b9b0007 100644
--- a/compiler-rt/lib/scudo/standalone/primary64.h
+++ b/compiler-rt/lib/scudo/standalone/primary64.h
@@ -221,13 +221,15 @@ template <typename Config> class SizeClassAllocator64 {
     DCHECK_EQ(BlocksInUse, BatchClassUsedInFreeLists);
   }
 
-  u16 popBlocks(CacheT *C, uptr ClassId, CompactPtrT *ToArray) {
+  u16 popBlocks(CacheT *C, uptr ClassId, CompactPtrT *ToArray,
+                const u16 MaxBlockCount) {
     TransferBatchT *B = popBatch(C, ClassId);
     if (!B)
       return 0;
 
     const u16 Count = B->getCount();
     DCHECK_GT(Count, 0U);
+    DCHECK_LE(Count, MaxBlockCount);
     B->moveToArray(ToArray);
 
     if (ClassId != SizeClassMap::BatchClassId)

>From 2f16b5898cd8a7a63a2eca5f4c32ebe1a9422007 Mon Sep 17 00:00:00 2001
From: Chia-hung Duan <chiahungduan at google.com>
Date: Thu, 26 Oct 2023 21:13:39 +0000
Subject: [PATCH 2/2] fixup! [scudo] Pass the max number of blocks to popBlocks

---
 compiler-rt/lib/scudo/standalone/primary32.h | 6 ++++--
 compiler-rt/lib/scudo/standalone/primary64.h | 6 ++++--
 2 files changed, 8 insertions(+), 4 deletions(-)

diff --git a/compiler-rt/lib/scudo/standalone/primary32.h b/compiler-rt/lib/scudo/standalone/primary32.h
index 137ac2f02ebd95d..8281e02ba164c13 100644
--- a/compiler-rt/lib/scudo/standalone/primary32.h
+++ b/compiler-rt/lib/scudo/standalone/primary32.h
@@ -191,15 +191,17 @@ template <typename Config> class SizeClassAllocator32 {
     return BlockSize > PageSize;
   }
 
+  // Note that the `MaxBlockCount` will be used when we support arbitrary blocks
+  // count. Now it's the same as the number of blocks stored in the
+  // `TransferBatch`.
   u16 popBlocks(CacheT *C, uptr ClassId, CompactPtrT *ToArray,
-                const u16 MaxBlockCount) {
+                UNUSED const u16 MaxBlockCount) {
     TransferBatchT *B = popBatch(C, ClassId);
     if (!B)
       return 0;
 
     const u16 Count = B->getCount();
     DCHECK_GT(Count, 0U);
-    DCHECK_LE(Count, MaxBlockCount);
     B->moveToArray(ToArray);
 
     if (ClassId != SizeClassMap::BatchClassId)
diff --git a/compiler-rt/lib/scudo/standalone/primary64.h b/compiler-rt/lib/scudo/standalone/primary64.h
index db404e40b9b0007..d1929ff7212f478 100644
--- a/compiler-rt/lib/scudo/standalone/primary64.h
+++ b/compiler-rt/lib/scudo/standalone/primary64.h
@@ -221,15 +221,17 @@ template <typename Config> class SizeClassAllocator64 {
     DCHECK_EQ(BlocksInUse, BatchClassUsedInFreeLists);
   }
 
+  // Note that the `MaxBlockCount` will be used when we support arbitrary blocks
+  // count. Now it's the same as the number of blocks stored in the
+  // `TransferBatch`.
   u16 popBlocks(CacheT *C, uptr ClassId, CompactPtrT *ToArray,
-                const u16 MaxBlockCount) {
+                UNUSED const u16 MaxBlockCount) {
     TransferBatchT *B = popBatch(C, ClassId);
     if (!B)
       return 0;
 
     const u16 Count = B->getCount();
     DCHECK_GT(Count, 0U);
-    DCHECK_LE(Count, MaxBlockCount);
     B->moveToArray(ToArray);
 
     if (ClassId != SizeClassMap::BatchClassId)