[Scudo] Fix SizeClassAllocatorLocalCache::drain

It leaved few blocks in PerClassArray[0].

Reviewed By: cryptoad

Differential Revision: https://reviews.llvm.org/D99763

diff --git a/compiler-rt/lib/scudo/standalone/local_cache.h b/compiler-rt/lib/scudo/standalone/local_cache.h
index 8ee2b1c..c8667d2 100644 (file)
--- a/compiler-rt/lib/scudo/standalone/local_cache.h
+++ b/compiler-rt/lib/scudo/standalone/local_cache.h
@@ -101,12 +101,22 @@ template <class SizeClassAllocator> struct SizeClassAllocatorLocalCache {
     Stats.add(StatFree, ClassSize);
   }
 
+  bool isEmpty() const {
+    for (uptr I = 0; I < NumClasses; ++I)
+      if (PerClassArray[I].Count)
+        return false;
+    return true;
+  }
+
   void drain() {
-    for (uptr I = 0; I < NumClasses; I++) {
+    // Drain BatchClassId (0) the last as createBatch can refill it.
+    for (uptr I = NumClasses; I;) {
+      --I;
       PerClass *C = &PerClassArray[I];
       while (C->Count > 0)
         drain(C, I);
     }
+    DCHECK(isEmpty());
   }
 
   TransferBatch *createBatch(uptr ClassId, void *B) {

diff --git a/compiler-rt/lib/scudo/standalone/tests/combined_test.cpp b/compiler-rt/lib/scudo/standalone/tests/combined_test.cpp
index 703d9e1..3afc324 100644 (file)
--- a/compiler-rt/lib/scudo/standalone/tests/combined_test.cpp
+++ b/compiler-rt/lib/scudo/standalone/tests/combined_test.cpp
@@ -311,6 +311,14 @@ template <class Config> static void testAllocator() {
   EXPECT_NE(Stats.find("Stats: SizeClassAllocator"), std::string::npos);
   EXPECT_NE(Stats.find("Stats: MapAllocator"), std::string::npos);
   EXPECT_NE(Stats.find("Stats: Quarantine"), std::string::npos);
+
+  bool UnlockRequired;
+  auto *TSD = Allocator->getTSDRegistry()->getTSDAndLock(&UnlockRequired);
+  EXPECT_TRUE(!TSD->Cache.isEmpty());
+  TSD->Cache.drain();
+  EXPECT_TRUE(TSD->Cache.isEmpty());
+  if (UnlockRequired)
+    TSD->unlock();
 }
 
 // Test that multiple instantiations of the allocator have not messed up the