Reland "[scudo] Support partial concurrent page release in SizeClassAllocator64"

[Chia-hungDuan]

Summary

This reverts commit 2f04b688aadef747172a8f4a7d1658dc049b1e24.

[Chia-hungDuan]

Merge non-full TransferBatches

[Chia-hungDuan]

@cferris , FYI, Base is the original CL and Diff 2 includes the fix

[Chia-hungDuan]

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

1464–1465

will be removed

1476–1477

May add a

DCHECK_EQ(BG->Batches.front()->getCount(), MaxCachedPerBatch)

[cferris]

LGTM

[Chia-hungDuan]

Also add releaseToOS() in SCUDO_TYPED_TEST(ScudoCombinedTest, ThreadedCombined) which increases the test coverage

I have run > 10K runs of multithread tests and ensured all the code path have been verified.

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

1476–1477

This will be checked in markFreeblocks(). Don't need to add one here.

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

void setFromArray(CompactPtrT *Array, u16 N) {
       memcpy(Batch, Array, sizeof(Batch[0]) * Count);
     }
     void appendFromArray(CompactPtrT *Array, u16 N) {
       DCHECK_LE(N, MaxNumCached - Count);
       memcpy(Batch + Count, Array, sizeof(Batch[0]) * N);
       // u16 will be promoted to int by arithmetic type conversion.
       Count = static_cast<u16>(Count + N);
     }
+    void appendFromTransferBatch(TransferBatch *B, u16 N) {
+      DCHECK_LE(N, MaxNumCached - Count);
+      DCHECK_GE(B->Count, N);
+      // Append from the back of `B`.
+      memcpy(Batch + Count, B->Batch + (B->Count - N), sizeof(Batch[0]) * N);
+      // u16 will be promoted to int by arithmetic type conversion.
+      Count = static_cast<u16>(Count + N);
+      B->Count = static_cast<u16>(B->Count - N);
+    }
     void clear() { Count = 0; }
     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; }
+    bool isEmpty() const { return Count == 0U; }
     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

Str->append(
         Region->ReleaseInfo.LastReleasedBytes >> 10,
         RegionPushedBytesDelta >> 10, Region->RegionBeg,
         getRegionBaseByClassId(ClassId));
   }
 
   NOINLINE uptr releaseToOSMaybe(RegionInfo *Region, uptr ClassId,
                                  ReleaseToOS ReleaseType = ReleaseToOS::Normal)
       REQUIRES(Region->MMLock) EXCLUDES(Region->FLLock) {
-    // TODO(chiahungduan): Release `FLLock` in step 3 & 4 described in the
-    // comment below.
     ScopedLock L(Region->FLLock);
 
     const uptr BlockSize = getSizeByClassId(ClassId);
     const uptr BytesInFreeList =
         Region->MemMapInfo.AllocatedUser - (Region->FreeListInfo.PoppedBlocks -
                                             Region->FreeListInfo.PushedBlocks) *
                                                BlockSize;
     if (UNLIKELY(BytesInFreeList == 0))
       return false;
 
     const uptr AllocatedUserEnd =
         Region->MemMapInfo.AllocatedUser + Region->RegionBeg;
     const uptr CompactPtrBase = getCompactPtrBaseByClassId(ClassId);
 
     // ====================================================================== //
     // 1. Check if we have enough free blocks and if it's worth doing a page
     //    release.
     // ====================================================================== //
     if (ReleaseType != ReleaseToOS::ForceAll &&
         !hasChanceToReleasePages(Region, BlockSize, BytesInFreeList,
                                  ReleaseType)) {
       return 0;
     }
 
-    // This is only used for debugging to ensure the consistency of the number
-    // of groups.
-    uptr NumberOfBatchGroups = Region->FreeListInfo.BlockList.size();
-
     // ====================================================================== //
     // 2. Determine which groups can release the pages. Use a heuristic to
     //    gather groups that are candidates for doing a release.
     // ====================================================================== //
     SinglyLinkedList<BatchGroup> GroupsToRelease;
     if (ReleaseType == ReleaseToOS::ForceAll) {
       GroupsToRelease = Region->FreeListInfo.BlockList;
       Region->FreeListInfo.BlockList.clear();
     } else {
       GroupsToRelease = collectGroupsToRelease(
           Region, BlockSize, AllocatedUserEnd, CompactPtrBase);
     }
     if (GroupsToRelease.empty())
       return 0;
 
-    // ====================================================================== //
-    // 3. Mark the free blocks in `GroupsToRelease` in the `PageReleaseContext`.
-    //    Then we can tell which pages are in-use by querying
-    //    `PageReleaseContext`.
-    // ====================================================================== //
+    // Ideally, we should use a class like `ScopedUnlock`. However, this form of
+    // unlocking is not supported by the thread-safety analysis. See
+    // https://clang.llvm.org/docs/ThreadSafetyAnalysis.html#no-alias-analysis
+    // for more details.
+    // Put it as local class so that we can mark the ctor/dtor with proper
+    // annotations associated to the target lock. Note that accessing the
+    // function variable in local class only works in thread-safety annotations.
+    // TODO: Implement general `ScopedUnlock` when it's supported.
+    class FLLockScopedUnlock {
+    public:
+      FLLockScopedUnlock(RegionInfo *Region) RELEASE(Region->FLLock)
+          : R(Region) {
+        R->FLLock.assertHeld();
+        R->FLLock.unlock();
+      }
+      ~FLLockScopedUnlock() ACQUIRE(Region->FLLock) { R->FLLock.lock(); }
+
+    private:
+      RegionInfo *R;
+    };
+
+    // Note that we have extracted the `GroupsToRelease` from region freelist.
+    // It's safe to let pushBlocks()/popBatches() access the remaining region
+    // freelist. In the steps 3 and 4, we will temporarily release the FLLock
+    // and lock it again before step 5.
+
+    uptr ReleasedBytes = 0;
+    {
+      FLLockScopedUnlock UL(Region);
+      // ==================================================================== //
+      // 3. Mark the free blocks in `GroupsToRelease` in the
+      //    `PageReleaseContext`. Then we can tell which pages are in-use by
+      //    querying `PageReleaseContext`.
+      // ==================================================================== //
       PageReleaseContext Context = markFreeBlocks(
           Region, BlockSize, AllocatedUserEnd, CompactPtrBase, GroupsToRelease);
       if (UNLIKELY(!Context.hasBlockMarked())) {
-      mergeGroupsToReleaseBack(Region, GroupsToRelease, NumberOfBatchGroups);
+        ScopedLock L(Region->FLLock);
+        mergeGroupsToReleaseBack(Region, GroupsToRelease);
         return 0;
       }
 
-    // ====================================================================== //
+      // ==================================================================== //
       // 4. Release the unused physical pages back to the OS.
-    // ====================================================================== //
+      // ==================================================================== //
       RegionReleaseRecorder<MemMapT> Recorder(&Region->MemMapInfo.MemMap,
                                               Region->RegionBeg,
                                               Context.getReleaseOffset());
       auto SkipRegion = [](UNUSED uptr RegionIndex) { return false; };
       releaseFreeMemoryToOS(Context, Recorder, SkipRegion);
       if (Recorder.getReleasedRangesCount() > 0) {
         Region->ReleaseInfo.BytesInFreeListAtLastCheckpoint = BytesInFreeList;
         Region->ReleaseInfo.RangesReleased += Recorder.getReleasedRangesCount();
         Region->ReleaseInfo.LastReleasedBytes = Recorder.getReleasedBytes();
       }
       Region->ReleaseInfo.LastReleaseAtNs = getMonotonicTimeFast();
+      ReleasedBytes = Recorder.getReleasedBytes();
+    }
 
     // ====================================================================== //
     // 5. Merge the `GroupsToRelease` back to the freelist.
     // ====================================================================== //
-    mergeGroupsToReleaseBack(Region, GroupsToRelease, NumberOfBatchGroups);
+    mergeGroupsToReleaseBack(Region, GroupsToRelease);
 
-    return Recorder.getReleasedBytes();
+    return ReleasedBytes;
   }
 
   bool hasChanceToReleasePages(RegionInfo *Region, uptr BlockSize,
                                uptr BytesInFreeList, ReleaseToOS ReleaseType)
       REQUIRES(Region->MMLock, Region->FLLock) {
     DCHECK_GE(Region->FreeListInfo.PoppedBlocks,
               Region->FreeListInfo.PushedBlocks);
     const uptr PageSize = getPageSizeCached();
collectGroupsToRelease(RegionInfo *Region, const uptr BlockSize,
 
     return GroupsToRelease;
   }
 
   PageReleaseContext
   markFreeBlocks(RegionInfo *Region, const uptr BlockSize,
                  const uptr AllocatedUserEnd, const uptr CompactPtrBase,
                  SinglyLinkedList<BatchGroup> &GroupsToRelease)
-      REQUIRES(Region->MMLock, Region->FLLock) {
+      REQUIRES(Region->MMLock) EXCLUDES(Region->FLLock) {
     const uptr GroupSize = (1U << GroupSizeLog);
     auto DecompactPtr = [CompactPtrBase](CompactPtrT CompactPtr) {
       return decompactPtrInternal(CompactPtrBase, CompactPtr);
     };
 
     const uptr ReleaseBase = decompactGroupBase(
         CompactPtrBase, GroupsToRelease.front()->CompactPtrGroupBase);
     const uptr LastGroupEnd =
for (BatchGroup &BG : GroupsToRelease) {
       uptr MaxContainedBlocks = AllocatedGroupSize / BlockSize;
       if (!BlockAlignedWithUsedEnd)
         ++MaxContainedBlocks;
 
       const uptr NumBlocks = (BG.Batches.size() - 1) * BG.MaxCachedPerBatch +
                              BG.Batches.front()->getCount();
 
       if (NumBlocks == MaxContainedBlocks) {
-        for (const auto &It : BG.Batches)
+        for (const auto &It : BG.Batches) {
+          if (&It != BG.Batches.front())
+            DCHECK_EQ(It.getCount(), BG.MaxCachedPerBatch);
           for (u16 I = 0; I < It.getCount(); ++I)
             DCHECK_EQ(compactPtrGroup(It.get(I)), BG.CompactPtrGroupBase);
+        }
 
         Context.markRangeAsAllCounted(BatchGroupBase, BatchGroupUsedEnd,
                                       Region->RegionBeg, /*RegionIndex=*/0,
                                       Region->MemMapInfo.AllocatedUser);
       } else {
         DCHECK_LT(NumBlocks, MaxContainedBlocks);
         // Note that we don't always visit blocks in each BatchGroup so that we
         // may miss the chance of releasing certain pages that cross
         // BatchGroups.
         Context.markFreeBlocksInRegion(
             BG.Batches, DecompactPtr, Region->RegionBeg, /*RegionIndex=*/0,
             Region->MemMapInfo.AllocatedUser, MayContainLastBlockInRegion);
       }
     }
 
     DCHECK(Context.hasBlockMarked());
 
     return Context;
   }
 
   void mergeGroupsToReleaseBack(RegionInfo *Region,
-                                SinglyLinkedList<BatchGroup> &GroupsToRelease,
-                                const uptr NumberOfBatchGroups)
+                                SinglyLinkedList<BatchGroup> &GroupsToRelease)
       REQUIRES(Region->MMLock, Region->FLLock) {
+    // After merging two freelists, we may have redundant `BatchGroup`s that
+    // need to be recycled. The number of unused `BatchGroup`s is expected to be
+    // small. Pick a constant which is inferred from real programs.
+    constexpr uptr MaxUnusedSize = 8;
+    CompactPtrT Blocks[MaxUnusedSize];
+    u32 Idx = 0;
+    RegionInfo *BatchClassRegion = getRegionInfo(SizeClassMap::BatchClassId);
+    // We can't call pushBatchClassBlocks() to recycle the unused `BatchGroup`s
+    // when we are manipulating the freelist of `BatchClassRegion`. Instead, we
+    // should just push it back to the freelist when we merge two `BatchGroup`s.
+    // This logic hasn't been implemented because we haven't supported releasing
+    // pages in `BatchClassRegion`.
+    DCHECK_NE(BatchClassRegion, Region);
+
     // Merge GroupsToRelease back to the Region::FreeListInfo.BlockList. Note
     // that both `Region->FreeListInfo.BlockList` and `GroupsToRelease` are
     // sorted.
     for (BatchGroup *BG = Region->FreeListInfo.BlockList.front(),
                     *Prev = nullptr;
          ;) {
       if (BG == nullptr || GroupsToRelease.empty()) {
         if (!GroupsToRelease.empty())
           Region->FreeListInfo.BlockList.append_back(&GroupsToRelease);
         break;
       }
 
-      DCHECK_NE(BG->CompactPtrGroupBase,
-                GroupsToRelease.front()->CompactPtrGroupBase);
+      DCHECK(!BG->Batches.empty());
 
       if (BG->CompactPtrGroupBase <
           GroupsToRelease.front()->CompactPtrGroupBase) {
         Prev = BG;
         BG = BG->Next;
         continue;
       }
 
+      BatchGroup *Cur = GroupsToRelease.front();
+      TransferBatch *UnusedTransferBatch = nullptr;
+      GroupsToRelease.pop_front();
+
+      if (BG->CompactPtrGroupBase == Cur->CompactPtrGroupBase) {
+        BG->PushedBlocks += Cur->PushedBlocks;
+        // We have updated `BatchGroup::BytesInBGAtLastCheckpoint` while
+        // collecting the `GroupsToRelease`.
+        BG->BytesInBGAtLastCheckpoint = Cur->BytesInBGAtLastCheckpoint;
+        const uptr MaxCachedPerBatch = BG->MaxCachedPerBatch;
+
+        // Note that the first TransferBatches in both `Batches` may not be
+        // full and only the first TransferBatch can have non-full blocks. Thus
+        // we have to merge them before appending one to another.
+        if (Cur->Batches.front()->getCount() == MaxCachedPerBatch) {
+          BG->Batches.append_back(&Cur->Batches);
+        } else {
+          TransferBatch *NonFullBatch = Cur->Batches.front();
+          Cur->Batches.pop_front();
+          const u16 NonFullBatchCount = NonFullBatch->getCount();
+          // The remaining Batches in `Cur` are full.
+          BG->Batches.append_back(&Cur->Batches);
+
+          if (BG->Batches.front()->getCount() == MaxCachedPerBatch) {
+            // Only 1 non-full TransferBatch, push it to the front.
+            BG->Batches.push_front(NonFullBatch);
+          } else {
+            const u16 NumBlocksToMove = static_cast<u16>(
+                Min(static_cast<u16>(MaxCachedPerBatch -
+                                     BG->Batches.front()->getCount()),
+                    NonFullBatchCount));
+            BG->Batches.front()->appendFromTransferBatch(NonFullBatch,
+                                                         NumBlocksToMove);
+            if (NonFullBatch->isEmpty())
+              UnusedTransferBatch = NonFullBatch;
+            else
+              BG->Batches.push_front(NonFullBatch);
+          }
+        }
+
+        const u32 NeededSlots = UnusedTransferBatch == nullptr ? 1U : 2U;
+        if (UNLIKELY(Idx + NeededSlots > MaxUnusedSize)) {
+          ScopedLock L(BatchClassRegion->FLLock);
+          pushBatchClassBlocks(BatchClassRegion, Blocks, Idx);
+          Idx = 0;
+        }
+        Blocks[Idx++] =
+            compactPtr(SizeClassMap::BatchClassId, reinterpret_cast<uptr>(Cur));
+        if (UnusedTransferBatch) {
+          Blocks[Idx++] =
+              compactPtr(SizeClassMap::BatchClassId,
+                         reinterpret_cast<uptr>(UnusedTransferBatch));
+        }
+        Prev = BG;
+        BG = BG->Next;
+        continue;
+      }
+
       // At here, the `BG` is the first BatchGroup with CompactPtrGroupBase
       // larger than the first element in `GroupsToRelease`. We need to insert
       // `GroupsToRelease::front()` (which is `Cur` below)  before `BG`.
       //
       //   1. If `Prev` is nullptr, we simply push `Cur` to the front of
       //      FreeListInfo.BlockList.
       //   2. Otherwise, use `insert()` which inserts an element next to `Prev`.
       //
       // Afterwards, we don't need to advance `BG` because the order between
       // `BG` and the new `GroupsToRelease::front()` hasn't been checked.
-      BatchGroup *Cur = GroupsToRelease.front();
-      GroupsToRelease.pop_front();
       if (Prev == nullptr)
         Region->FreeListInfo.BlockList.push_front(Cur);
       else
         Region->FreeListInfo.BlockList.insert(Prev, Cur);
       DCHECK_EQ(Cur->Next, BG);
       Prev = Cur;
     }
 
-    DCHECK_EQ(Region->FreeListInfo.BlockList.size(), NumberOfBatchGroups);
-    (void)NumberOfBatchGroups;
+    if (Idx != 0) {
+      ScopedLock L(BatchClassRegion->FLLock);
+      pushBatchClassBlocks(BatchClassRegion, Blocks, Idx);
+    }
 
     if (SCUDO_DEBUG) {
       BatchGroup *Prev = Region->FreeListInfo.BlockList.front();
       for (BatchGroup *Cur = Prev->Next; Cur != nullptr;
            Prev = Cur, Cur = Cur->Next) {
         CHECK_LT(Prev->CompactPtrGroupBase, Cur->CompactPtrGroupBase);
       }
     }

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

Threads[I] = std::thread([&]() {
       std::vector<std::pair<void *, scudo::uptr>> V;
       for (scudo::uptr I = 0; I < 256U; I++) {
         const scudo::uptr Size = std::rand() % 4096U;
         void *P = Allocator->allocate(Size, Origin);
         // A region could have ran out of memory, resulting in a null P.
         if (P)
           V.push_back(std::make_pair(P, Size));
       }
+
+      // Try to interleave pushBlocks(), popBatch() and releaseToOS().
+      Allocator->releaseToOS(scudo::ReleaseToOS::Force);
+
       while (!V.empty()) {
         auto Pair = V.back();
         Allocator->deallocate(Pair.first, Origin, Pair.second);
         V.pop_back();
       }
     });
   {
     std::unique_lock<std::mutex> Lock(Mutex);

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

Threads[I] = std::thread([&]() {
         const scudo::uptr Size =
             std::rand() % Primary::SizeClassMap::MaxSize / 4;
         const scudo::uptr ClassId =
             Primary::SizeClassMap::getClassIdBySize(Size);
         void *P = Cache.allocate(ClassId);
         if (P)
           V.push_back(std::make_pair(ClassId, P));
       }
+
+      // Try to interleave pushBlocks(), popBatch() and releaseToOS().
+      Allocator->releaseToOS(scudo::ReleaseToOS::Force);
+
       while (!V.empty()) {
         auto Pair = V.back();
         Cache.deallocate(Pair.first, Pair.second);
         V.pop_back();
+        // This increases the chance of having non-full TransferBatches and it
+        // will jump into the code path of merging TransferBatches.
+        if (std::rand() % 8 == 0)
+          Cache.drain();
       }
       Cache.destroy(nullptr);
     });
   }
   {
     std::unique_lock<std::mutex> Lock(Mutex);
     Ready = true;
     Cv.notify_all();