[scudo] Pass the max number of blocks to popBlocks

[ChiaHungDuan]

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)

[llvmbot]

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

Author: None (ChiaHungDuan)

Changes

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

[cferris1000]

It doesn't appear this value is used any place, is that expected?

[ChiaHungDuan]

Yes, it'll be used in the case that TransferBatch stores more blocks than MaxBlockCount. Now the TransferBatch almost always stores the same amount as MaxBlocksCount so we only use it in the DCHECK at this moment

[cferris1000]

Do you need to add a [[maybe_unused]] in front of the variable? It seems like this could result in a build error when DCHECK is compiled out.

[ChiaHungDuan]

You're right. Given that the same check has been done in refill (local_cache.h), I just mark it as UNUSED and put a comment.

[cferris1000]

LGTM