[scudo][standalone] Simplify populateFreelist

[cryptoad]

Summary

populateFreelist was more complicated that it needed to be. We used
to call to populateBatches that would do some internal shuffling and
add pointers one by one to the batches, but ultimately this was not
needed. We can get rid of populateBatches, and do processing in
bulk. This doesn't necessarily make things faster as this is not on the
hot path, but it makes the function cleaner.

Additionally clean up a couple of items, like UNLIKELYs and setting
Exhausted to false which can't happen.

[cryptoad]

Re-add a comment that got removed.

[cryptoad]

Changing some lines around.

[cryptoad]

Restoring a couple of UNLIKELY in the loops.

[cryptoad]

And another, sorry for the back & forth, getting numbers as they come.

[pcc]

LGTM

With this change we will end up with blocks distributed randomly among all of the transfer batches that we create at one time instead of having all of a transfer batch's blocks be consecutive. So we improve randomness at the cost of some locality which could impact performance. I'd be fine with letting this land though and we can see if it significantly impacts performance in practice once it's picked up downstream.

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

uptr allocateRegion(uptr ClassId) {
     return Region;
   }
 
   SizeClassInfo *getSizeClassInfo(uptr ClassId) {
     DCHECK_LT(ClassId, NumClasses);
     return &SizeClassInfoArray[ClassId];
   }
 
-  bool populateBatches(CacheT *C, SizeClassInfo *Sci, uptr ClassId,
-                       TransferBatch **CurrentBatch, u32 MaxCount,
-                       void **PointersArray, u32 Count) {
-    if (ClassId != SizeClassMap::BatchClassId)
-      shuffle(PointersArray, Count, &Sci->RandState);
-    TransferBatch *B = *CurrentBatch;
-    for (uptr I = 0; I < Count; I++) {
-      if (B && B->getCount() == MaxCount) {
-        Sci->FreeList.push_back(B);
-        B = nullptr;
-      }
-      if (!B) {
-        B = C->createBatch(ClassId, PointersArray[I]);
-        if (UNLIKELY(!B))
-          return false;
-        B->clear();
-      }
-      B->add(PointersArray[I]);
-    }
-    *CurrentBatch = B;
-    return true;
-  }
-
   NOINLINE TransferBatch *populateFreeList(CacheT *C, uptr ClassId,
                                            SizeClassInfo *Sci) {
     uptr Region;
     uptr Offset;
     // If the size-class currently has a region associated to it, use it. The
     // newly created blocks will be located after the currently allocated memory
     // for that region (up to RegionSize). Otherwise, create a new region, where
     // the new blocks will be carved from the beginning.
     if (Sci->CurrentRegion) {
       Region = Sci->CurrentRegion;
       DCHECK_GT(Sci->CurrentRegionAllocated, 0U);
       Offset = Sci->CurrentRegionAllocated;
     } else {
       DCHECK_EQ(Sci->CurrentRegionAllocated, 0U);
       Region = allocateRegion(ClassId);
-      if (UNLIKELY(!Region))
+      if (!Region)
         return nullptr;
       C->getStats().add(StatMapped, RegionSize);
       Sci->CurrentRegion = Region;
       Offset = 0;
     }
 
     const uptr Size = getSizeByClassId(ClassId);
     const u32 MaxCount = TransferBatch::getMaxCached(Size);
     DCHECK_GT(MaxCount, 0U);
     // The maximum number of blocks we should carve in the region is dictated
     // by the maximum number of batches we want to fill, and the amount of
     // memory left in the current region (we use the lowest of the two). This
     // will not be 0 as we ensure that a region can at least hold one block (via
     // static_assert and at the end of this function).
     const u32 NumberOfBlocks =
         Min(MaxNumBatches * MaxCount,
             static_cast<u32>((RegionSize - Offset) / Size));
     DCHECK_GT(NumberOfBlocks, 0U);
 
-    TransferBatch *B = nullptr;
     constexpr u32 ShuffleArraySize =
         MaxNumBatches * TransferBatch::MaxNumCached;
-    // Fill the transfer batches and put them in the size-class freelist. We
-    // need to randomize the blocks for security purposes, so we first fill a
-    // local array that we then shuffle before populating the batches.
     void *ShuffleArray[ShuffleArraySize];
-    u32 Count = 0;
+    DCHECK_LE(NumberOfBlocks, ShuffleArraySize);
+
     const uptr AllocatedUser = Size * NumberOfBlocks;
-    for (uptr I = Region + Offset; I < Region + Offset + AllocatedUser;
-         I += Size) {
-      ShuffleArray[Count++] = reinterpret_cast<void *>(I);
-      if (Count == ShuffleArraySize) {
-        if (UNLIKELY(!populateBatches(C, Sci, ClassId, &B, MaxCount,
-                                      ShuffleArray, Count)))
-          return nullptr;
-        Count = 0;
-      }
-    }
-    if (Count) {
-      if (UNLIKELY(!populateBatches(C, Sci, ClassId, &B, MaxCount, ShuffleArray,
-                                    Count)))
+    uptr P = Region + Offset;
+    for (u32 I = 0; I < NumberOfBlocks; I++, P += Size)
+      ShuffleArray[I] = reinterpret_cast<void *>(P);
+    // No need to shuffle the batches size class.
+    if (ClassId != SizeClassMap::BatchClassId)
+      shuffle(ShuffleArray, NumberOfBlocks, &Sci->RandState);
+    for (u32 I = 0; I < NumberOfBlocks;) {
+      TransferBatch *B = C->createBatch(ClassId, ShuffleArray[I]);
+      if (!B)
         return nullptr;
-    }
-    DCHECK(B);
-    if (!Sci->FreeList.empty()) {
+      const u32 N = Min(MaxCount, NumberOfBlocks - I);
+      B->setFromArray(&ShuffleArray[I], N);
       Sci->FreeList.push_back(B);
-      B = Sci->FreeList.front();
-      Sci->FreeList.pop_front();
+      I += N;
     }
+    TransferBatch *B = Sci->FreeList.front();
+    Sci->FreeList.pop_front();
+    DCHECK(B);
     DCHECK_GT(B->getCount(), 0);
 
     C->getStats().add(StatFree, AllocatedUser);
     DCHECK_LE(Sci->CurrentRegionAllocated + AllocatedUser, RegionSize);
     // If there is not enough room in the region currently associated to fit
     // more blocks, we deassociate the region by resetting CurrentRegion and
     // CurrentRegionAllocated. Otherwise, update the allocated amount.
     if (RegionSize - (Sci->CurrentRegionAllocated + AllocatedUser) < Size) {

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

RegionInfo *getRegionInfo(uptr ClassId) {
     DCHECK_LT(ClassId, NumClasses);
     return &RegionInfoArray[ClassId];
   }
 
   uptr getRegionBaseByClassId(uptr ClassId) const {
     return PrimaryBase + (ClassId << RegionSizeLog);
   }
 
-  bool populateBatches(CacheT *C, RegionInfo *Region, uptr ClassId,
-                       TransferBatch **CurrentBatch, u32 MaxCount,
-                       void **PointersArray, u32 Count) {
-    // No need to shuffle the batches size class.
-    if (ClassId != SizeClassMap::BatchClassId)
-      shuffle(PointersArray, Count, &Region->RandState);
-    TransferBatch *B = *CurrentBatch;
-    for (uptr I = 0; I < Count; I++) {
-      if (B && B->getCount() == MaxCount) {
-        Region->FreeList.push_back(B);
-        B = nullptr;
-      }
-      if (!B) {
-        B = C->createBatch(ClassId, PointersArray[I]);
-        if (UNLIKELY(!B))
-          return false;
-        B->clear();
-      }
-      B->add(PointersArray[I]);
-    }
-    *CurrentBatch = B;
-    return true;
-  }
-
   NOINLINE TransferBatch *populateFreeList(CacheT *C, uptr ClassId,
                                            RegionInfo *Region) {
     const uptr Size = getSizeByClassId(ClassId);
     const u32 MaxCount = TransferBatch::getMaxCached(Size);
 
     const uptr RegionBeg = Region->RegionBeg;
     const uptr MappedUser = Region->MappedUser;
     const uptr TotalUserBytes = Region->AllocatedUser + MaxCount * Size;
     // Map more space for blocks, if necessary.
-    if (TotalUserBytes > MappedUser) {
+    if (UNLIKELY(TotalUserBytes > MappedUser)) {
       // Do the mmap for the user memory.
       const uptr UserMapSize =
           roundUpTo(TotalUserBytes - MappedUser, MapSizeIncrement);
       const uptr RegionBase = RegionBeg - getRegionBaseByClassId(ClassId);
-      if (UNLIKELY(RegionBase + MappedUser + UserMapSize > RegionSize)) {
+      if (RegionBase + MappedUser + UserMapSize > RegionSize) {
         if (!Region->Exhausted) {
           Region->Exhausted = true;
           ScopedString Str(1024);
           getStats(&Str);
           Str.append(
-              "Scudo OOM: The process has Exhausted %zuM for size class %zu.\n",
+              "Scudo OOM: The process has exhausted %zuM for size class %zu.\n",
               RegionSize >> 20, Size);
           Str.output();
         }
         return nullptr;
       }
-      if (UNLIKELY(MappedUser == 0))
+      if (MappedUser == 0)
         Region->Data = Data;
-      if (UNLIKELY(!map(reinterpret_cast<void *>(RegionBeg + MappedUser),
-                        UserMapSize, "scudo:primary",
-                        MAP_ALLOWNOMEM | MAP_RESIZABLE |
-                            (useMemoryTagging(Options.load()) ? MAP_MEMTAG : 0),
-                        &Region->Data)))
+      if (!map(reinterpret_cast<void *>(RegionBeg + MappedUser), UserMapSize,
+               "scudo:primary",
+               &Region->Data))
         return nullptr;
       Region->MappedUser += UserMapSize;
       C->getStats().add(StatMapped, UserMapSize);
     }
 
     const u32 NumberOfBlocks = Min(
         MaxNumBatches * MaxCount,
         static_cast<u32>((Region->MappedUser - Region->AllocatedUser) / Size));
     DCHECK_GT(NumberOfBlocks, 0);
 
-    TransferBatch *B = nullptr;
     constexpr u32 ShuffleArraySize =
         MaxNumBatches * TransferBatch::MaxNumCached;
     void *ShuffleArray[ShuffleArraySize];
-    u32 Count = 0;
-    const uptr P = RegionBeg + Region->AllocatedUser;
+    DCHECK_LE(NumberOfBlocks, ShuffleArraySize);
+
     const uptr AllocatedUser = Size * NumberOfBlocks;
-    for (uptr I = P; I < P + AllocatedUser; I += Size) {
-      ShuffleArray[Count++] = reinterpret_cast<void *>(I);
-      if (Count == ShuffleArraySize) {
-        if (UNLIKELY(!populateBatches(C, Region, ClassId, &B, MaxCount,
-                                      ShuffleArray, Count)))
-          return nullptr;
-        Count = 0;
-      }
-    }
-    if (Count) {
-      if (UNLIKELY(!populateBatches(C, Region, ClassId, &B, MaxCount,
-                                    ShuffleArray, Count)))
+    uptr P = RegionBeg + Region->AllocatedUser;
+    for (u32 I = 0; I < NumberOfBlocks; I++, P += Size)
+      ShuffleArray[I] = reinterpret_cast<void *>(P);
+    // No need to shuffle the batches size class.
+    if (ClassId != SizeClassMap::BatchClassId)
+      shuffle(ShuffleArray, NumberOfBlocks, &Region->RandState);
+    for (u32 I = 0; I < NumberOfBlocks;) {
+      TransferBatch *B = C->createBatch(ClassId, ShuffleArray[I]);
+      if (!B)
         return nullptr;
-    }
-    DCHECK(B);
-    if (!Region->FreeList.empty()) {
+      const u32 N = Min(MaxCount, NumberOfBlocks - I);
+      B->setFromArray(&ShuffleArray[I], N);
       Region->FreeList.push_back(B);
-      B = Region->FreeList.front();
-      Region->FreeList.pop_front();
+      I += N;
     }
+    TransferBatch *B = Region->FreeList.front();
+    Region->FreeList.pop_front();
+    DCHECK(B);
     DCHECK_GT(B->getCount(), 0);
 
     C->getStats().add(StatFree, AllocatedUser);
     Region->AllocatedUser += AllocatedUser;
-    Region->Exhausted = false;
 
     return B;
   }
 
   void getStats(ScopedString *Str, uptr ClassId, uptr Rss) {
     RegionInfo *Region = getRegionInfo(ClassId);
     if (Region->MappedUser == 0)
       return;