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

[via llvm-commits]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

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

Author: None (ChiaHungDuan)

<details>
<summary>Changes</summary>

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)

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


3 Files Affected:

- (modified) compiler-rt/lib/scudo/standalone/local_cache.h (+8-8) 
- (modified) compiler-rt/lib/scudo/standalone/primary32.h (+3-1) 
- (modified) compiler-rt/lib/scudo/standalone/primary64.h (+3-1) 


``````````diff
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)

``````````

</details>


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