[compiler-rt] ce9e1a3 - [Scudo] Fix SizeClassAllocatorLocalCache::drain

[Vitaly Buka via llvm-commits]

Author: Vitaly Buka
Date: 2021-04-01T13:27:03-07:00
New Revision: ce9e1a3632779294077766d6646ef3d79b35d414

URL: https://github.com/llvm/llvm-project/commit/ce9e1a3632779294077766d6646ef3d79b35d414
DIFF: https://github.com/llvm/llvm-project/commit/ce9e1a3632779294077766d6646ef3d79b35d414.diff

LOG: [Scudo] Fix SizeClassAllocatorLocalCache::drain

It leaved few blocks in PerClassArray[0].

Reviewed By: cryptoad

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

Added: 
    

Modified: 
    compiler-rt/lib/scudo/standalone/local_cache.h
    compiler-rt/lib/scudo/standalone/tests/combined_test.cpp

Removed: 
    


################################################################################
diff  --git a/compiler-rt/lib/scudo/standalone/local_cache.h b/compiler-rt/lib/scudo/standalone/local_cache.h
index 8ee2b1cae84e..c8667d2e5622 100644
--- 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 703d9e12bdca..3afc32408e8e 100644
--- 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