[scudo] PopBatch after populateFreeList()

[Chia-hungDuan]

Summary

Ensure the thread that refills freelist will get the Batch without
contending the lock in SizeClassAllocator64.

[Chia-hungDuan]

Rebase only

[cferris]

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

872

Wouldn't it be better to pull this out of the if and else statements like:

TransferBatch *B = popBatchImpl(C, ClassId, Region);

[Chia-hungDuan]

Address review comment

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

872

Oops, definitely it'll be better. The FLLock was held in the if-clause. I moved it out later and forgot to move the popBatchImp(). Done.

[cferris]

LGTM.

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

struct TransferBatch {
     void add(CompactPtrT P) {
       DCHECK_LT(Count, MaxNumCached);
       Batch[Count++] = P;
     }
     void copyToArray(CompactPtrT *Array) const {
       memcpy(Array, Batch, sizeof(Batch[0]) * Count);
     }
     u16 getCount() const { return Count; }
+    CompactPtrT *getRawArray() { return Batch; }
     CompactPtrT get(u16 I) const {
       DCHECK_LE(I, Count);
       return Batch[I];
     }
     static u16 getMaxCached(uptr Size) {
       return Min(MaxNumCached, SizeClassMap::getMaxCachedHint(Size));
     }
     TransferBatch *Next;

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

TransferBatch *popBatch(CacheT *C, uptr ClassId) {
     {
       ScopedLock L(Region->FLLock);
       TransferBatch *B = popBatchImpl(C, ClassId, Region);
       if (LIKELY(B))
         return B;
     }
 
     bool PrintStats = false;
+    TransferBatch *B = nullptr;
 
     while (true) {
       // When two threads compete for `Region->MMLock`, we only want one of them
-      // does the populateFreeList(). To avoid both of them doing that, always
-      // check the freelist before mapping new pages.
+      // does the populateFreeListAndPopBatch(). To avoid both of them doing
+      // that, always check the freelist before mapping new pages.
       //
       // TODO(chiahungduan): Use condition variable so that we don't need to
       // hold `Region->MMLock` here.
       ScopedLock ML(Region->MMLock);
       {
         ScopedLock FL(Region->FLLock);
-        TransferBatch *B = popBatchImpl(C, ClassId, Region);
+        B = popBatchImpl(C, ClassId, Region);
         if (LIKELY(B))
           return B;
       }
 
       const bool RegionIsExhausted = Region->Exhausted;
-      if (!RegionIsExhausted) {
-        // TODO: Make sure the one who does populateFreeList() get one Batch
-        // here.
-        populateFreeList(C, ClassId, Region);
-      }
+      if (!RegionIsExhausted)
+        B = populateFreeListAndPopBatch(C, ClassId, Region);
       PrintStats = !RegionIsExhausted && Region->Exhausted;
-      if (Region->Exhausted)
-        break;
     }
 
     if (UNLIKELY(PrintStats)) {
       ScopedString Str;
       getStats(&Str);
       Str.append(
           "Scudo OOM: The process has exhausted %zuM for size class %zu.\n",
           RegionSize >> 20, getSizeByClassId(ClassId));
       Str.output();
     }
 
-    return nullptr;
+    return B;
   }
 
   // Push the array of free blocks to the designated batch group.
   void pushBlocks(CacheT *C, uptr ClassId, CompactPtrT *Array, u32 Size) {
     DCHECK_LT(ClassId, NumClasses);
     DCHECK_GT(Size, 0);
 
     RegionInfo *Region = getRegionInfo(ClassId);
     if (ClassId == SizeClassMap::BatchClassId) {
       // Constructing a batch group in the free list will use two blocks in
       // BatchClassId. If we are pushing BatchClassId blocks, we will use the
       // blocks in the array directly (can't delegate local cache which will
       // cause a recursive allocation). However, The number of free blocks may
       // be less than two. Therefore, populate the freelist before inserting the
       // blocks.
+      TransferBatch *B = nullptr;
       while (true) {
         // TODO(chiahungduan): Move the lock right before the call of
-        // populateFreeList() by using condition variable. See more details in
-        // the comment of popBatch().
+        // populateFreeListAndPopBatch() by using condition variable. See more
+        // details in the comment of popBatch().
         ScopedLock L(Region->MMLock);
         {
           ScopedLock L(Region->FLLock);
-          const bool NeedToRefill =
-              Size == 1U && Region->FreeListInfo.BlockList.empty();
-          if (!NeedToRefill) {
+          const bool NeedToRefill = Size == 1U &&
+                                    Region->FreeListInfo.BlockList.empty() &&
+                                    B == nullptr;
+          if (UNLIKELY(!NeedToRefill)) {
+            if (UNLIKELY(B)) {
+              pushBlocksImpl(C, SizeClassMap::BatchClassId, Region,
+                             B->getRawArray(), B->getCount());
+              CHECK(!Region->FreeListInfo.BlockList.empty());
+            }
             pushBlocksImpl(C, SizeClassMap::BatchClassId, Region, Array, Size);
             return;
           }
         }
 
-        if (UNLIKELY(!populateFreeList(C, SizeClassMap::BatchClassId, Region)))
+        // Note that this batch will be pushed again and it's only used to
+        // ensure we have enough blocks to construct batch group.
+        B = populateFreeListAndPopBatch(C, SizeClassMap::BatchClassId, Region);
+        if (!B)
           break;
       }
 
       // At here, it hints that the BatchClass region has been used up. Dump the
       // region stats before the program shutdown.
       ScopedString Str;
       getStats(&Str);
       Str.append(
if (Batches.empty()) {
         C->deallocate(SizeClassMap::BatchClassId, BG);
     }
 
     Region->FreeListInfo.PoppedBlocks += B->getCount();
 
     return B;
   }
 
-  NOINLINE bool populateFreeList(CacheT *C, uptr ClassId, RegionInfo *Region)
+  // Refill the freelist and return one batch.
+  NOINLINE TransferBatch *populateFreeListAndPopBatch(CacheT *C, uptr ClassId,
+                                                      RegionInfo *Region)
       REQUIRES(Region->MMLock) EXCLUDES(Region->FLLock) {
     const uptr Size = getSizeByClassId(ClassId);
     const u16 MaxCount = TransferBatch::getMaxCached(Size);
 
     const uptr RegionBeg = Region->RegionBeg;
     const uptr MappedUser = Region->MemMapInfo.MappedUser;
     const uptr TotalUserBytes =
         Region->MemMapInfo.AllocatedUser + MaxCount * Size;
     // Map more space for blocks, if necessary.
     if (TotalUserBytes > MappedUser) {
       // Do the mmap for the user memory.
       const uptr MapSize =
           roundUp(TotalUserBytes - MappedUser, MapSizeIncrement);
       const uptr RegionBase = RegionBeg - getRegionBaseByClassId(ClassId);
       if (UNLIKELY(RegionBase + MappedUser + MapSize > RegionSize)) {
         Region->Exhausted = true;
-        return false;
+        return nullptr;
       }
       // TODO: Consider allocating MemMapInfo.MemMap in init().
       if (!Region->MemMapInfo.MemMap.isAllocated()) {
         // TODO: Ideally, a region should reserve RegionSize because the memory
         // between `RegionBeg` and region base is still belong to a region and
         // it's just not used. In order to make it work on every platform (some
         // of them don't support `remap()` across the unused range), dispatch
         // from `RegionBeg` for now.
         const uptr ReserveSize =
             RegionSize - (RegionBeg - getRegionBaseByClassId(ClassId));
         Region->MemMapInfo.MemMap =
             ReservedMemory.dispatch(RegionBeg, ReserveSize);
       }
       DCHECK(Region->MemMapInfo.MemMap.isAllocated());
 
       if (UNLIKELY(!Region->MemMapInfo.MemMap.remap(
               RegionBeg + MappedUser, MapSize, "scudo:primary",
               MAP_ALLOWNOMEM | MAP_RESIZABLE |
                   (useMemoryTagging<Config>(Options.load()) ? MAP_MEMTAG
                                                             : 0)))) {
-        return false;
+        return nullptr;
       }
       Region->MemMapInfo.MappedUser += MapSize;
       C->getStats().add(StatMapped, MapSize);
     }
 
     const u32 NumberOfBlocks =
         Min(MaxNumBatches * MaxCount,
             static_cast<u32>((Region->MemMapInfo.MappedUser -
                               Region->MemMapInfo.AllocatedUser) /
                              Size));
     DCHECK_GT(NumberOfBlocks, 0);
 
     constexpr u32 ShuffleArraySize =
         MaxNumBatches * TransferBatch::MaxNumCached;
     CompactPtrT ShuffleArray[ShuffleArraySize];
     DCHECK_LE(NumberOfBlocks, ShuffleArraySize);
 
     const uptr CompactPtrBase = getCompactPtrBaseByClassId(ClassId);
     uptr P = RegionBeg + Region->MemMapInfo.AllocatedUser;
     for (u32 I = 0; I < NumberOfBlocks; I++, P += Size)
       ShuffleArray[I] = compactPtrInternal(CompactPtrBase, P);
 
     ScopedLock L(Region->FLLock);
 
+    TransferBatch *B = nullptr;
+
     if (ClassId != SizeClassMap::BatchClassId) {
       u32 N = 1;
       uptr CurGroup = compactPtrGroup(ShuffleArray[0]);
       for (u32 I = 1; I < NumberOfBlocks; I++) {
         if (UNLIKELY(compactPtrGroup(ShuffleArray[I]) != CurGroup)) {
           shuffle(ShuffleArray + I - N, N, &Region->RandState);
           pushBlocksImpl(C, ClassId, Region, ShuffleArray + I - N, N,
                          /*SameGroup=*/true);
           N = 1;
           CurGroup = compactPtrGroup(ShuffleArray[I]);
         } else {
           ++N;
         }
       }
 
       shuffle(ShuffleArray + NumberOfBlocks - N, N, &Region->RandState);
       pushBlocksImpl(C, ClassId, Region, &ShuffleArray[NumberOfBlocks - N], N,
                      /*SameGroup=*/true);
+      B = popBatchImpl(C, ClassId, Region);
     } else {
       pushBlocksImpl(C, ClassId, Region, ShuffleArray, NumberOfBlocks,
                      /*SameGroup=*/true);
+      B = popBatchImpl(C, ClassId, Region);
     }
+    DCHECK_NE(B, nullptr);
 
     // Note that `PushedBlocks` and `PoppedBlocks` are supposed to only record
     // the requests from `PushBlocks` and `PopBatch` which are external
     // interfaces. `populateFreeListAndPopBatch` is internal interface so we
     // should set the values back to avoid messing the stats.
     Region->FreeListInfo.PushedBlocks -= NumberOfBlocks;
 
     const uptr AllocatedUser = Size * NumberOfBlocks;
     C->getStats().add(StatFree, AllocatedUser);
     Region->MemMapInfo.AllocatedUser += AllocatedUser;
 
-    return true;
+    return B;
   }
 
   void getStats(ScopedString *Str, uptr ClassId, RegionInfo *Region, uptr Rss)
       REQUIRES(Region->MMLock, Region->FLLock) {
     if (Region->MemMapInfo.MappedUser == 0)
       return;
     const uptr InUse =
         Region->FreeListInfo.PoppedBlocks - Region->FreeListInfo.PushedBlocks;