[scudo] Verify the size of free blocks in primary allocator

[Chia-hungDuan]

Summary

When all the blocks (local caches are included) are freed, the size of
free blocks should be equal to AllocatedUser.

[Chia-hungDuan]

This is useful when we support concurrent alloc/free and releaseToOSMaybe() in primary allocator.

[cferris]

LGTM

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

for (uptr I = 0; I < NumClasses; I++) {
       if (Sci->MinRegionIndex < MinRegionIndex)
         MinRegionIndex = Sci->MinRegionIndex;
       if (Sci->MaxRegionIndex > MaxRegionIndex)
         MaxRegionIndex = Sci->MaxRegionIndex;
       *Sci = {};
     }
 
     ScopedLock L(ByteMapMutex);
-    for (uptr I = MinRegionIndex; I < MaxRegionIndex; I++)
+    for (uptr I = MinRegionIndex; I <= MaxRegionIndex; I++)
       if (PossibleRegions[I])
         unmap(reinterpret_cast<void *>(I * RegionSize), RegionSize);
     PossibleRegions.unmapTestOnly();
   }
 
+  // When all blocks are freed, it has to be the same size as `AllocatedUser`.
+  void verifyAllBlocksAreReleasedTestOnly() {
+    // `BatchGroup` and `TransferBatch` also use the blocks from BatchClass.
+    uptr BatchClassUsedInFreeLists = 0;
+    for (uptr I = 0; I < NumClasses; I++) {
+      // We have to count BatchClassUsedInFreeLists in other regions first.
+      if (I == SizeClassMap::BatchClassId)
+        continue;
+      SizeClassInfo *Sci = getSizeClassInfo(I);
+      ScopedLock L1(Sci->Mutex);
+      uptr TotalBlocks = 0;
+      for (BatchGroup &BG : Sci->FreeListInfo.BlockList) {
+        // `BG::Batches` are `TransferBatches`. +1 for `BatchGroup`.
+        BatchClassUsedInFreeLists += BG.Batches.size() + 1;
+        for (const auto &It : BG.Batches)
+          TotalBlocks += It.getCount();
+      }
+
+      const uptr BlockSize = getSizeByClassId(I);
+      DCHECK_EQ(TotalBlocks, Sci->AllocatedUser / BlockSize);
+    }
+
+    SizeClassInfo *Sci = getSizeClassInfo(SizeClassMap::BatchClassId);
+    ScopedLock L1(Sci->Mutex);
+    uptr TotalBlocks = 0;
+    for (BatchGroup &BG : Sci->FreeListInfo.BlockList)
+      for (const auto &It : BG.Batches)
+        TotalBlocks += It.getCount();
+
+    const uptr BlockSize = getSizeByClassId(SizeClassMap::BatchClassId);
+    DCHECK_EQ(TotalBlocks + BatchClassUsedInFreeLists,
+              Sci->AllocatedUser / BlockSize);
+  }
+
   CompactPtrT compactPtr(UNUSED uptr ClassId, uptr Ptr) const {
     return static_cast<CompactPtrT>(Ptr);
   }
 
   void *decompactPtr(UNUSED uptr ClassId, CompactPtrT CompactPtr) const {
     return reinterpret_cast<void *>(static_cast<uptr>(CompactPtr));
   }
 

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

for (uptr I = 0; I < NumClasses; I++) {
       RegionInfo *Region = getRegionInfo(I);
       *Region = {};
     }
     if (PrimaryBase)
       ReservedMemory.release();
     PrimaryBase = 0U;
   }
 
+  // When all blocks are freed, it has to be the same size as `AllocatedUser`.
+  void verifyAllBlocksAreReleasedTestOnly() {
+    // `BatchGroup` and `TransferBatch` also use the blocks from BatchClass.
+    uptr BatchClassUsedInFreeLists = 0;
+    for (uptr I = 0; I < NumClasses; I++) {
+      // We have to count BatchClassUsedInFreeLists in other regions first.
+      if (I == SizeClassMap::BatchClassId)
+        continue;
+      RegionInfo *Region = getRegionInfo(I);
+      ScopedLock ML(Region->MMLock);
+      ScopedLock FL(Region->FLLock);
+      const uptr BlockSize = getSizeByClassId(I);
+      uptr TotalBlocks = 0;
+      for (BatchGroup &BG : Region->FreeListInfo.BlockList) {
+        // `BG::Batches` are `TransferBatches`. +1 for `BatchGroup`.
+        BatchClassUsedInFreeLists += BG.Batches.size() + 1;
+        for (const auto &It : BG.Batches)
+          TotalBlocks += It.getCount();
+      }
+
+      DCHECK_EQ(TotalBlocks, Region->MemMapInfo.AllocatedUser / BlockSize);
+    }
+
+    RegionInfo *Region = getRegionInfo(SizeClassMap::BatchClassId);
+    ScopedLock ML(Region->MMLock);
+    ScopedLock FL(Region->FLLock);
+    const uptr BlockSize = getSizeByClassId(SizeClassMap::BatchClassId);
+    uptr TotalBlocks = 0;
+    for (BatchGroup &BG : Region->FreeListInfo.BlockList)
+      for (const auto &It : BG.Batches)
+        TotalBlocks += It.getCount();
+    DCHECK_EQ(TotalBlocks + BatchClassUsedInFreeLists,
+              Region->MemMapInfo.AllocatedUser / BlockSize);
+  }
+
   TransferBatch *popBatch(CacheT *C, uptr ClassId) {
     DCHECK_LT(ClassId, NumClasses);
     RegionInfo *Region = getRegionInfo(ClassId);
 
     {
       ScopedLock L(Region->FLLock);
       TransferBatch *B = popBatchImpl(C, ClassId, Region);
       if (LIKELY(B))

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

 struct SizeClassAllocator
     : public scudo::SizeClassAllocator64<Config<BaseConfig, SizeClassMapT>> {};
 template <typename SizeClassMapT>
 struct SizeClassAllocator<TestConfig1, SizeClassMapT>
     : public scudo::SizeClassAllocator32<Config<TestConfig1, SizeClassMapT>> {};
 
 template <template <typename> class BaseConfig, typename SizeClassMapT>
 struct TestAllocator : public SizeClassAllocator<BaseConfig, SizeClassMapT> {
-  ~TestAllocator() { this->unmapTestOnly(); }
+  ~TestAllocator() {
+    this->verifyAllBlocksAreReleasedTestOnly();
+    this->unmapTestOnly();
+  }
 
   void *operator new(size_t size) {
     void *p = nullptr;
     EXPECT_EQ(0, posix_memalign(&p, alignof(TestAllocator), size));
     return p;
   }
 
   void operator delete(void *ptr) { free(ptr); }

 SCUDO_TYPED_TEST(ScudoPrimaryTest, PrimaryThreaded) {
   using Primary = TestAllocator<TypeParam, scudo::SvelteSizeClassMap>;
   std::unique_ptr<Primary> Allocator(new Primary);
   Allocator->init(/*ReleaseToOsInterval=*/-1);
   std::mutex Mutex;
   std::condition_variable Cv;
   bool Ready = false;
   std::thread Threads[32];
-  for (scudo::uptr I = 0; I < ARRAY_SIZE(Threads); I++)
+  for (scudo::uptr I = 0; I < ARRAY_SIZE(Threads); I++) {
     Threads[I] = std::thread([&]() {
       static thread_local typename Primary::CacheT Cache;
       Cache.init(nullptr, Allocator.get());
       std::vector<std::pair<scudo::uptr, void *>> V;
       {
         std::unique_lock<std::mutex> Lock(Mutex);
         while (!Ready)
           Cv.wait(Lock);
Threads[I] = std::thread([&]() {
       }
       while (!V.empty()) {
         auto Pair = V.back();
         Cache.deallocate(Pair.first, Pair.second);
         V.pop_back();
       }
       Cache.destroy(nullptr);
     });
+  }
   {
     std::unique_lock<std::mutex> Lock(Mutex);
     Ready = true;
     Cv.notify_all();
   }
   for (auto &T : Threads)
     T.join();
   Allocator->releaseToOS(scudo::ReleaseToOS::Force);
SCUDO_TYPED_TEST(ScudoPrimaryTest, MemoryGroup) {
   Cache.drain();
 
   for (scudo::uptr I = 0; I < FinalNumberOfAllocations; ++I)
     Blocks.push_back(reinterpret_cast<scudo::uptr>(Cache.allocate(ClassId)));
 
   EXPECT_LE(*std::max_element(Blocks.begin(), Blocks.end()) -
                 *std::min_element(Blocks.begin(), Blocks.end()),
             GroupSizeMem * 2);
+
+  while (!Blocks.empty()) {
+    Cache.deallocate(ClassId, reinterpret_cast<void *>(Blocks.back()));
+    Blocks.pop_back();
+  }
+  Cache.drain();
 }