[Scudo] Fix SizeClassAllocatorLocalCache::drain

[vitalybuka]

Summary

It leaved few blocks in PerClassArray[0].

[cryptoad]

Thank you, good catch, how did this show up?
1 comment nit.

compiler-rt/lib/scudo/standalone/local_cache.h

112–114

s/cat/can

[cryptoad]

compiler-rt/lib/scudo/standalone/local_cache.h

113

Actually, could I please ask you to write this and not consider that the BatchClassId is 0? In other portions of the code I used a if I == BatchClassId continue and then a special case at the end. That would allow to change BatchClassId more easily if it ever came to this.

[vitalybuka]

typo

[vitalybuka]

In D99763#2664987, @cryptoad wrote:

Thank you, good catch, how did this show up?

I tried to run tests with --gtest_repeat=2. It still needs D99766

compiler-rt/lib/scudo/standalone/local_cache.h

void deallocate(uptr ClassId, void *P) {
     // See comment in allocate() about memory accesses.
     const uptr ClassSize = C->ClassSize;
     C->Chunks[C->Count++] =
         Allocator->compactPtr(ClassId, reinterpret_cast<uptr>(P));
     Stats.sub(StatAllocated, ClassSize);
     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) {
     if (ClassId != SizeClassMap::BatchClassId)
       B = allocate(SizeClassMap::BatchClassId);
     return reinterpret_cast<TransferBatch *>(B);
   }
 

compiler-rt/lib/scudo/standalone/tests/combined_test.cpp

while (ActualSize > BufferSize) {
     ActualSize = Allocator->getStats(Buffer.data(), BufferSize);
   }
   std::string Stats(Buffer.begin(), Buffer.end());
   // Basic checks on the contents of the statistics output, which also allows us
   // to verify that we got it all.
   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
 // process's signal handlers (GWP-ASan used to do this).
 void testSEGV() {
   const scudo::uptr Size = 4 * scudo::getPageSizeCached();
   scudo::MapPlatformData Data = {};
   void *P = scudo::map(nullptr, Size, "testSEGV", MAP_NOACCESS, &Data);