[scudo] Only prepare PageMap entry for partial region

[Chia-hungDuan]

Summary

This reduces the size of PageMap and we are more likely to use the
static local buffer. Note that now this is only supported for single
region case, i.e. on SizeClassAllocator64. For SizeClassAllocator32,
it needs a different way to save the PageMap.

[cferris]

Is there any performance penalty of going from a single pass to two pass approach? Is the number of blocks usually small so there isn't much of an issue?

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

804

Might be better worded as:

visiting BlockGroups twice.

813

of the number

821

Is there a reason you are using a comma, rather than making this two statements?

It's a bit odd to see this construct, but it might just be me.

885

I think you mean that the Next pointer can be modified. Either that or maybe use destructive to indicate the Next pointer could be changed.

888

After this loop, Prev is still pointing to the current previous block instead of Cur. Is this intended?

908

Should this be roundUpSlow?

I assume from the name that his is a slower way of doing the round up, but I might be incorrect.

910

As with the other change, making all of these Beg suffixes be Base instead make the code easier to read.

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

443

logic of how

[Chia-hungDuan]

Rebase

[Chia-hungDuan]

Address review comment

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

821

TBH, I simply wanted to make the induction variables updating compact (here and below). Changed them two separate lines instead.

888

Yes, we extract the current one and the Cur->Next will be the next of Prev

908

The BlockSize may not always be power-of-2 so we need to use the slower version here.

910

Done. BTW, I'm considering to also change the RegionBeg to RegionBase in another CL

[Chia-hungDuan]

In D142659#4101366, @cferris wrote:

Is there any performance penalty of going from a single pass to two pass approach? Is the number of blocks usually small so there isn't much of an issue?

The preliminary measurement showed that the cost of two pass was neglectable. I will bring some numbers later.

[cferris]

One last question about a bit of code.

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

888

But BG has already been set to BG->Next. So if you hit this part of the loop twice in a row, the second time:

Prev is the same as in the previous loop
BG = BG->Next (no relation to Cur)

Did you mean BG->Next will be the next batch of Prev?

At the very least, I think a comment would help here since this is kind of confusing and look wrong.

[Chia-hungDuan]

1. Add more comments on BG extraction
2. Fix a bug in releaseToOs and add more tests

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

888

I guess the confusion comes from how we update Prev. Prev will be pointed to BG->Next by FreeList::extract() so the relation between Prev and BG->Next is updated.

Added more comments.

[Chia-hungDuan]

Remove debugging header

[Chia-hungDuan]

Remove more debugging stuff...

[Chia-hungDuan]

Add a performance number: The time spent in extracting/merging BatchGroup

02-21 21:04:56.996   515   515 I scudo   : -- Average Operation Time -- -- Name (# of Calls) --
02-21 21:04:56.996   515   515 I scudo   :            68.3(ns)             Extract/MergeBlock (7700)

[cferris]

Just a couple of comment nits.

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

916

a

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

430

remaining

[Chia-hungDuan]

Address review comment

[Chia-hungDuan]

A failed test on sanitizer-x86_64-linux-qemu. Looking into it now

[MaskRay]

For relands, we usually use Reland $original_subject instead of Revert Revert $original_subject :)
Likewise, one Revert is sufficient, no need for Revert Revert Revert $original_subject (3)

[MaskRay]

I usually use Revert Dxxxx $original_subject instead of Revert Dxxxx $original_subject so that the revert commit is associated with this page. You can see the commit in "Commits" in the "Details" section.

[Chia-hungDuan]

In D142659#4165452, @MaskRay wrote:

I usually use Revert Dxxxx $original_subject instead of Revert Dxxxx $original_subject so that the revert commit is associated with this page. You can see the commit in "Commits" in the "Details" section.

Got it! Sorry I didn't notice that. Thanks

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

NOINLINE uptr releaseToOSMaybe(SizeClassInfo *Sci, uptr ClassId,
     uptr TotalReleasedBytes = 0;
     const uptr Base = First * RegionSize;
     const uptr NumberOfRegions = Last - First + 1U;
     const uptr GroupSize = (1U << GroupSizeLog);
     const uptr CurRegionGroupId =
         compactPtrGroup(compactPtr(ClassId, Sci->CurrentRegion));
 
     ReleaseRecorder Recorder(Base);
-    PageReleaseContext Context(BlockSize, RegionSize, NumberOfRegions);
+    PageReleaseContext Context(BlockSize, RegionSize, NumberOfRegions,
+                               /*ReleaseSize=*/RegionSize);
 
     auto DecompactPtr = [](CompactPtrT CompactPtr) {
       return reinterpret_cast<uptr>(CompactPtr);
     };
     for (BatchGroup &BG : Sci->FreeList) {
       const uptr PushedBytesDelta =
           BG.PushedBlocks - BG.PushedBlocksAtLastCheckpoint;
       if (PushedBytesDelta * BlockSize < PageSize)

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

if (!Force) {
               static_cast<u64>(IntervalMs) * 1000000 >
           getMonotonicTime()) {
         return 0; // Memory was returned recently.
       }
     }
 
     const uptr GroupSize = (1U << GroupSizeLog);
     const uptr AllocatedUserEnd = Region->AllocatedUser + Region->RegionBeg;
-    ReleaseRecorder Recorder(Region->RegionBeg, &Region->Data);
-    PageReleaseContext Context(BlockSize, Region->AllocatedUser,
-                               /*NumberOfRegions=*/1U);
-
     const uptr CompactPtrBase = getCompactPtrBaseByClassId(ClassId);
     auto DecompactPtr = [CompactPtrBase](CompactPtrT CompactPtr) {
       return decompactPtrInternal(CompactPtrBase, CompactPtr);
     };
-    for (BatchGroup &BG : Region->FreeList) {
+
+    // Instead of always preparing PageMap for the entire region, we only do it
+    // for the range of releasing groups. To do that, the free-block marking
+    // process includes visiting BlockGroups twice.
+
+    // The first visit is to determine the range of BatchGroups we are going to
+    // release. And we will extract those BatchGroups out and push into
+    // `GroupToRelease`.
+    SinglyLinkedList<BatchGroup> GroupToRelease;
+    GroupToRelease.clear();
+    BatchGroup *Prev = nullptr;
+
+    // This is only used for debugging to ensure the consistency of the number
+    // of groups.
+    uptr NumberOfBatchGroups = Region->FreeList.size();
+
+    for (BatchGroup *BG = Region->FreeList.front(); BG != nullptr;) {
       const uptr PushedBytesDelta =
-          BG.PushedBlocks - BG.PushedBlocksAtLastCheckpoint;
-      if (PushedBytesDelta * BlockSize < PageSize)
+          BG->PushedBlocks - BG->PushedBlocksAtLastCheckpoint;
+      if (PushedBytesDelta * BlockSize < PageSize) {
+        Prev = BG;
+        BG = BG->Next;
         continue;
+      }
 
       // Group boundary does not necessarily have the same alignment as Region.
       // It may sit across a Region boundary. Which means that we may have the
       // following two cases,
       //
       // 1. Group boundary sits before RegionBeg.
       //
-      //                (BatchGroupBeg)
+      //                (BatchGroupBase)
       // batchGroupBase  RegionBeg       BatchGroupEnd
       //        |            |                |
       //        v            v                v
       //        +------------+----------------+
       //         \                           /
       //          ------   GroupSize   ------
       //
       // 2. Group boundary sits after RegionBeg.
       //
-      //               (BatchGroupBeg)
+      //               (BatchGroupBase)
       //    RegionBeg  batchGroupBase               BatchGroupEnd
       //        |           |                             |
       //        v           v                             v
       //        +-----------+-----------------------------+
       //                     \                           /
       //                      ------   GroupSize   ------
       //
       // Note that in the first case, the group range before RegionBeg is never
       // used. Therefore, while calculating the used group size, we should
       // exclude that part to get the correct size.
-      const uptr BatchGroupBeg =
-          Max(batchGroupBase(CompactPtrBase, BG.GroupId), Region->RegionBeg);
-      DCHECK_GE(AllocatedUserEnd, BatchGroupBeg);
+      const uptr BatchGroupBase =
+          Max(batchGroupBase(CompactPtrBase, BG->GroupId), Region->RegionBeg);
+      DCHECK_GE(AllocatedUserEnd, BatchGroupBase);
       const uptr BatchGroupEnd =
-          batchGroupBase(CompactPtrBase, BG.GroupId) + GroupSize;
+          batchGroupBase(CompactPtrBase, BG->GroupId) + GroupSize;
       const uptr AllocatedGroupSize = AllocatedUserEnd >= BatchGroupEnd
-                                          ? BatchGroupEnd - BatchGroupBeg
-                                          : AllocatedUserEnd - BatchGroupBeg;
-      if (AllocatedGroupSize == 0)
+                                          ? BatchGroupEnd - BatchGroupBase
+                                          : AllocatedUserEnd - BatchGroupBase;
+      if (AllocatedGroupSize == 0) {
+        Prev = BG;
+        BG = BG->Next;
         continue;
+      }
 
       // TransferBatches are pushed in front of BG.Batches. The first one may
       // not have all caches used.
-      const uptr NumBlocks = (BG.Batches.size() - 1) * BG.MaxCachedPerBatch +
-                             BG.Batches.front()->getCount();
+      const uptr NumBlocks = (BG->Batches.size() - 1) * BG->MaxCachedPerBatch +
+                             BG->Batches.front()->getCount();
       const uptr BytesInBG = NumBlocks * BlockSize;
       // Given the randomness property, we try to release the pages only if the
       // bytes used by free blocks exceed certain proportion of group size. Note
       // that this heuristic only applies when all the spaces in a BatchGroup
       // are allocated.
       if (CheckDensity && (BytesInBG * 100U) / AllocatedGroupSize <
                               (100U - 1U - BlockSize / 16U)) {
+        Prev = BG;
+        BG = BG->Next;
         continue;
       }
 
+      // If `BG` is the first BatchGroup in the list, we only need to advance
+      // `BG` and call FreeList::pop_front(). No update is needed for `Prev`.
+      //
+      //         (BG)   (BG->Next)
+      // Prev     Cur      BG
+      //   |       |       |
+      //   v       v       v
+      //  nil     +--+    +--+
+      //          |X | -> |  | -> ...
+      //          +--+    +--+
+      //
+      // Otherwise, `Prev` will be used to extract the `Cur` from the
+      // `FreeList`.
+      //
+      //         (BG)   (BG->Next)
+      // Prev     Cur      BG
+      //   |       |       |
+      //   v       v       v
+      //  +--+    +--+    +--+
+      //  |  | -> |X | -> |  | -> ...
+      //  +--+    +--+    +--+
+      //
+      // After FreeList::extract(),
+      //
+      // Prev     Cur       BG
+      //   |       |        |
+      //   v       v        v
+      //  +--+    +--+     +--+
+      //  |  |-+  |X |  +->|  | -> ...
+      //  +--+ |  +--+  |  +--+
+      //       +--------+
+      //
+      // Note that we need to advance before pushing this BatchGroup to
+      // GroupToRelease because it's an destructive operation.
+
+      BatchGroup *Cur = BG;
+      BG = BG->Next;
+
+      if (Prev != nullptr)
+        Region->FreeList.extract(Prev, Cur);
+      else
+        Region->FreeList.pop_front();
+      GroupToRelease.push_back(Cur);
+    }
+
+    if (GroupToRelease.empty())
+      return 0;
+
+    const uptr ReleaseBase =
+        Max(batchGroupBase(CompactPtrBase, GroupToRelease.front()->GroupId),
+            Region->RegionBeg);
+    const uptr LastGroupEnd =
+        Min(batchGroupBase(CompactPtrBase, GroupToRelease.back()->GroupId) +
+                GroupSize,
+            AllocatedUserEnd);
+    // The last block may straddle the group boundary. Rounding up to BlockSize
+    // to get the exact range.
+    const uptr ReleaseEnd =
+        roundUpSlow(LastGroupEnd - Region->RegionBeg, BlockSize) +
+        Region->RegionBeg;
+    const uptr ReleaseRangeSize = ReleaseEnd - ReleaseBase;
+
+    ReleaseRecorder Recorder(ReleaseBase, &Region->Data);
+    PageReleaseContext Context(
+        BlockSize, Region->AllocatedUser, /*NumberOfRegions=*/1U,
+        ReleaseRangeSize, /*ReleaseOffset=*/ReleaseBase - Region->RegionBeg);
+
+    for (BatchGroup &BG : GroupToRelease) {
       BG.PushedBlocksAtLastCheckpoint = BG.PushedBlocks;
 
-      const uptr BatchGroupUsedEnd = BatchGroupBeg + AllocatedGroupSize;
+      // TODO(chiahungduan): Replace GroupId with BatchGroupBase to simplify the
+      // calculation here and above (where we determine the set of groups to
+      // release).
+      const uptr BatchGroupBase =
+          Max(batchGroupBase(CompactPtrBase, BG.GroupId), Region->RegionBeg);
+      const uptr BatchGroupEnd =
+          batchGroupBase(CompactPtrBase, BG.GroupId) + GroupSize;
+      const uptr AllocatedGroupSize = AllocatedUserEnd >= BatchGroupEnd
+                                          ? BatchGroupEnd - BatchGroupBase
+                                          : AllocatedUserEnd - BatchGroupBase;
+      const uptr BatchGroupUsedEnd = BatchGroupBase + AllocatedGroupSize;
       const bool BlockAlignedWithUsedEnd =
           (BatchGroupUsedEnd - Region->RegionBeg) % BlockSize == 0;
 
       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 (u16 I = 0; I < It.getCount(); ++I)
             DCHECK_EQ(compactPtrGroup(It.get(I)), BG.GroupId);
 
-        Context.markRangeAsAllCounted(BatchGroupBeg, BatchGroupUsedEnd,
+        Context.markRangeAsAllCounted(BatchGroupBase, BatchGroupUsedEnd,
                                       Region->RegionBeg);
       } 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.markFreeBlocks(BG.Batches, DecompactPtr, Region->RegionBeg);
       }
     }
 
-    if (!Context.hasBlockMarked())
-      return 0;
+    DCHECK(Context.hasBlockMarked());
 
     auto SkipRegion = [](UNUSED uptr RegionIndex) { return false; };
     releaseFreeMemoryToOS(Context, Recorder, SkipRegion);
 
     if (Recorder.getReleasedRangesCount() > 0) {
       Region->ReleaseInfo.PushedBlocksAtLastRelease =
           Region->Stats.PushedBlocks;
       Region->ReleaseInfo.RangesReleased += Recorder.getReleasedRangesCount();
       Region->ReleaseInfo.LastReleasedBytes = Recorder.getReleasedBytes();
     }
     Region->ReleaseInfo.LastReleaseAtNs = getMonotonicTime();
+
+    // Merge GroupToRelease back to the Region::FreeList. Note that both
+    // `Region->FreeList` and `GroupToRelease` are sorted.
+    for (BatchGroup *BG = Region->FreeList.front();; BG = BG->Next) {
+      if (BG == nullptr || GroupToRelease.empty()) {
+        if (!GroupToRelease.empty())
+          Region->FreeList.append_back(&GroupToRelease);
+        break;
+      }
+
+      DCHECK_NE(BG->GroupId, GroupToRelease.front()->GroupId);
+
+      if (BG->GroupId < GroupToRelease.front()->GroupId)
+        continue;
+
+      BatchGroup *Cur = GroupToRelease.front();
+      GroupToRelease.pop_front();
+      Region->FreeList.insert(BG, Cur);
+    }
+
+    DCHECK_EQ(Region->FreeList.size(), NumberOfBatchGroups);
+    (void)NumberOfBatchGroups;
+
     return Recorder.getReleasedBytes();
   }
 };
 
 } // namespace scudo
 
 #endif // SCUDO_PRIMARY64_H_

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

private:
   ReleaseRecorderT &Recorder;
   const uptr PageSizeLog;
   bool InRange = false;
   uptr CurrentPage = 0;
   uptr CurrentRangeStatePage = 0;
 };
 
 struct PageReleaseContext {
-  PageReleaseContext(uptr BlockSize, uptr RegionSize, uptr NumberOfRegions) :
-      BlockSize(BlockSize),
-      RegionSize(RegionSize),
+  PageReleaseContext(uptr BlockSize, uptr RegionSize, uptr NumberOfRegions,
+                     uptr ReleaseSize, uptr ReleaseOffset = 0)
+      : BlockSize(BlockSize), RegionSize(RegionSize),
         NumberOfRegions(NumberOfRegions) {
     PageSize = getPageSizeCached();
     if (BlockSize <= PageSize) {
       if (PageSize % BlockSize == 0) {
         // Same number of chunks per page, no cross overs.
         FullPagesBlockCountMax = PageSize / BlockSize;
         SameBlockCountPerPage = true;
       } else if (BlockSize % (PageSize % BlockSize) == 0) {
         // Some chunks are crossing page boundaries, which means that the page
if (BlockSize <= PageSize) {
       } else {
         // One chunk covers multiple pages, Some chunks are crossing page
         // boundaries. Some pages contain one chunk, some contain two.
         FullPagesBlockCountMax = 2;
         SameBlockCountPerPage = false;
       }
     }
 
-    PagesCount = roundUp(RegionSize, PageSize) / PageSize;
+    // TODO: For multiple regions, it's more complicated to support partial
+    // region marking (which includes the complexity of how to handle the last
+    // block in a region). We may consider this after markFreeBlocks() accepts
+    // only free blocks from the same region.
+    if (NumberOfRegions != 1) {
+      DCHECK_EQ(ReleaseSize, RegionSize);
+      DCHECK_EQ(ReleaseOffset, 0U);
+    }
+
+    PagesCount = roundUp(ReleaseSize, PageSize) / PageSize;
     PageSizeLog = getLog2(PageSize);
-    RoundedRegionSize = PagesCount << PageSizeLog;
+    RoundedRegionSize = roundUp(RegionSize, PageSize);
     RoundedSize = NumberOfRegions * RoundedRegionSize;
+    ReleasePageOffset = ReleaseOffset >> PageSizeLog;
   }
 
   // PageMap is lazily allocated when markFreeBlocks() is invoked.
   bool hasBlockMarked() const {
     return PageMap.isAllocated();
   }
 
   void ensurePageMapAllocated() {
void markRangeAsAllCounted(uptr From, uptr To, uptr Base) {
     // Therefore, we can't just mark the first page as all counted. Instead, we
     // increment the number of blocks in the first page in the page map and
     // then round up the `From` to the next page.
     if (FirstBlockInRange != FromInRegion) {
       DCHECK_GT(FromInRegion + PageSize, FirstBlockInRange);
       uptr NumBlocksInFirstPage =
           (FromInRegion + PageSize - FirstBlockInRange + BlockSize - 1) /
           BlockSize;
-      PageMap.incN(RegionIndex, FromInRegion >> PageSizeLog,
+      PageMap.incN(RegionIndex, getPageIndex(FromInRegion),
                    NumBlocksInFirstPage);
       FromInRegion = roundUp(FromInRegion + 1, PageSize);
     }
 
     uptr LastBlockInRange = roundDownSlow(ToInRegion - 1, BlockSize);
     if (LastBlockInRange < FromInRegion)
       return;
 
if (LastBlockInRange + BlockSize != RegionSize) {
       //                          +-----+-----+-----+-----+
       //                          |     |     |     |     | ...
       //                          +-----+-----+-----+-----+
       //                    ... -||-    last page    -||-    next page    -|
       //
       // The last block is not aligned to `To`, we need to increment the
       // counter of `next page` by 1.
       if (LastBlockInRange + BlockSize != ToInRegion) {
-        PageMap.incRange(RegionIndex, ToInRegion >> PageSizeLog,
-                         (LastBlockInRange + BlockSize - 1) >> PageSizeLog);
+        PageMap.incRange(RegionIndex, getPageIndex(ToInRegion),
+                         getPageIndex(LastBlockInRange + BlockSize - 1));
       }
     } else {
       ToInRegion = RegionSize;
     }
 
     // After handling the first page and the last block, it's safe to mark any
     // page in between the range [From, To).
     if (FromInRegion < ToInRegion) {
-      PageMap.setAsAllCountedRange(RegionIndex, FromInRegion >> PageSizeLog,
-                                   (ToInRegion - 1) >> PageSizeLog);
+      PageMap.setAsAllCountedRange(RegionIndex, getPageIndex(FromInRegion),
+                                   getPageIndex(ToInRegion - 1));
     }
   }
 
   template<class TransferBatchT, typename DecompactPtrT>
   void markFreeBlocks(const IntrusiveList<TransferBatchT> &FreeList,
                       DecompactPtrT DecompactPtr, uptr Base) {
     ensurePageMapAllocated();
 
+    const uptr LastBlockInRegion = ((RegionSize / BlockSize) - 1U) * BlockSize;
+
+    // The last block in a region may not use the entire page, so if it's free,
+    // we mark the following "pretend" memory block(s) as free.
+    auto markLastBlock = [this, LastBlockInRegion](const uptr RegionIndex) {
+      uptr PInRegion = LastBlockInRegion + BlockSize;
+      while (PInRegion < RoundedRegionSize) {
+        PageMap.incRange(RegionIndex, getPageIndex(PInRegion),
+                         getPageIndex(PInRegion + BlockSize - 1));
+        PInRegion += BlockSize;
+      }
+    };
+
     // Iterate over free chunks and count how many free chunks affect each
     // allocated page.
     if (BlockSize <= PageSize && PageSize % BlockSize == 0) {
       // Each chunk affects one page only.
       for (const auto &It : FreeList) {
         for (u16 I = 0; I < It.getCount(); I++) {
           const uptr P = DecompactPtr(It.get(I)) - Base;
           if (P >= RoundedSize)
             continue;
           const uptr RegionIndex = NumberOfRegions == 1U ? 0 : P / RegionSize;
           const uptr PInRegion = P - RegionIndex * RegionSize;
-          PageMap.inc(RegionIndex, PInRegion >> PageSizeLog);
+          PageMap.inc(RegionIndex, getPageIndex(PInRegion));
+          if (PInRegion == LastBlockInRegion)
+            markLastBlock(RegionIndex);
         }
       }
     } else {
       // In all other cases chunks might affect more than one page.
       DCHECK_GE(RegionSize, BlockSize);
-      const uptr LastBlockInRegion =
-          ((RegionSize / BlockSize) - 1U) * BlockSize;
       for (const auto &It : FreeList) {
         for (u16 I = 0; I < It.getCount(); I++) {
           const uptr P = DecompactPtr(It.get(I)) - Base;
           if (P >= RoundedSize)
             continue;
           const uptr RegionIndex = NumberOfRegions == 1U ? 0 : P / RegionSize;
           uptr PInRegion = P - RegionIndex * RegionSize;
-          PageMap.incRange(RegionIndex, PInRegion >> PageSizeLog,
-                            (PInRegion + BlockSize - 1) >> PageSizeLog);
-          // The last block in a region might straddle a page, so if it's
-          // free, we mark the following "pretend" memory block(s) as free.
-          if (PInRegion == LastBlockInRegion) {
-            PInRegion += BlockSize;
-            while (PInRegion < RoundedRegionSize) {
-              PageMap.incRange(RegionIndex, PInRegion >> PageSizeLog,
-                                (PInRegion + BlockSize - 1) >> PageSizeLog);
-              PInRegion += BlockSize;
-            }
-          }
+          PageMap.incRange(RegionIndex, getPageIndex(PInRegion),
+                           getPageIndex(PInRegion + BlockSize - 1));
+          if (PInRegion == LastBlockInRegion)
+            markLastBlock(RegionIndex);
         }
       }
     }
   }
 
+  uptr getPageIndex(uptr P) { return (P >> PageSizeLog) - ReleasePageOffset; }
+
   uptr BlockSize;
   uptr RegionSize;
   uptr NumberOfRegions;
+  // For partial region marking, some pages in front are not needed to be
+  // counted.
+  uptr ReleasePageOffset;
   uptr PageSize;
   uptr PagesCount;
   uptr PageSizeLog;
   uptr RoundedRegionSize;
   uptr RoundedSize;
   uptr FullPagesBlockCountMax;
   bool SameBlockCountPerPage;
   RegionPageMap PageMap;
 };
 
 // Try to release the page which doesn't have any in-used block, i.e., they are
 // all free blocks. The `PageMap` will record the number of free blocks in each
 // page.
 template <class ReleaseRecorderT, typename SkipRegionT>
 NOINLINE void
 releaseFreeMemoryToOS(PageReleaseContext &Context,
                       ReleaseRecorderT &Recorder, SkipRegionT SkipRegion) {
   const uptr PageSize = Context.PageSize;
   const uptr BlockSize = Context.BlockSize;
   const uptr PagesCount = Context.PagesCount;
   const uptr NumberOfRegions = Context.NumberOfRegions;
+  const uptr ReleasePageOffset = Context.ReleasePageOffset;
   const uptr FullPagesBlockCountMax = Context.FullPagesBlockCountMax;
   const bool SameBlockCountPerPage = Context.SameBlockCountPerPage;
   RegionPageMap &PageMap = Context.PageMap;
 
   // Iterate over pages detecting ranges of pages with chunk Counters equal
   // to the expected number of chunks for the particular page.
   FreePagesRangeTracker<ReleaseRecorderT> RangeTracker(Recorder);
   if (SameBlockCountPerPage) {
if (SameBlockCountPerPage) {
     // whether the page boundary was crossed.
     for (uptr I = 0; I < NumberOfRegions; I++) {
       if (SkipRegion(I)) {
         RangeTracker.skipPages(PagesCount);
         continue;
       }
       uptr PrevPageBoundary = 0;
       uptr CurrentBoundary = 0;
+      if (ReleasePageOffset > 0) {
+        PrevPageBoundary = ReleasePageOffset * PageSize;
+        CurrentBoundary = roundUpSlow(PrevPageBoundary, BlockSize);
+      }
       for (uptr J = 0; J < PagesCount; J++) {
         const uptr PageBoundary = PrevPageBoundary + PageSize;
         uptr BlocksPerPage = Pn;
         if (CurrentBoundary < PageBoundary) {
           if (CurrentBoundary > PrevPageBoundary)
             BlocksPerPage++;
           CurrentBoundary += Pnc;
           if (CurrentBoundary < PageBoundary) {

 // information after releasing.
 template <class TransferBatchT, class ReleaseRecorderT, typename DecompactPtrT,
           typename SkipRegionT>
 NOINLINE void
 releaseFreeMemoryToOS(const IntrusiveList<TransferBatchT> &FreeList,
                       uptr RegionSize, uptr NumberOfRegions, uptr BlockSize,
                       ReleaseRecorderT &Recorder, DecompactPtrT DecompactPtr,
                       SkipRegionT SkipRegion) {
-  PageReleaseContext Context(BlockSize, RegionSize, NumberOfRegions);
+  PageReleaseContext Context(BlockSize, /*ReleaseSize=*/RegionSize, RegionSize,
+                             NumberOfRegions);
   Context.markFreeBlocks(FreeList, DecompactPtr, Recorder.getBase());
   releaseFreeMemoryToOS(Context, Recorder, SkipRegion);
 }
 
 } // namespace scudo
 
 #endif // SCUDO_RELEASE_H_

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

for (auto TestCase : TestCases) {
     std::string Expected(TestCase,
                          LastX == nullptr ? 0 : (LastX - TestCase + 1));
     EXPECT_STREQ(Expected.c_str(), Recorder.ReportedPages.c_str());
   }
 }
 
 class ReleasedPagesRecorder {
 public:
+  ReleasedPagesRecorder() = default;
+  explicit ReleasedPagesRecorder(scudo::uptr Base) : Base(Base) {}
   std::set<scudo::uptr> ReportedPages;
 
   void releasePageRangeToOS(scudo::uptr From, scudo::uptr To) {
     const scudo::uptr PageSize = scudo::getPageSizeCached();
     for (scudo::uptr I = From; I < To; I += PageSize)
-      ReportedPages.insert(I);
+      ReportedPages.insert(I + getBase());
   }
 
-  scudo::uptr getBase() const { return 0; }
+  scudo::uptr getBase() const { return Base; }
+  scudo::uptr Base = 0;
 };
 
 // Simplified version of a TransferBatch.
 template <class SizeClassMap> struct FreeBatch {
   static const scudo::u16 MaxCount = SizeClassMap::MaxNumCachedHint;
   void clear() { Count = 0; }
   void add(scudo::uptr P) {
     DCHECK_LT(Count, MaxCount);
for (scudo::uptr I = 1; I <= SizeClassMap::LargestClassId; I++) {
     }
 
     // Release the memory.
     auto SkipRegion = [](UNUSED scudo::uptr RegionIndex) { return false; };
     auto DecompactPtr = [](scudo::uptr P) { return P; };
     ReleasedPagesRecorder Recorder;
     scudo::PageReleaseContext Context(BlockSize,
                                       /*RegionSize=*/MaxBlocks * BlockSize,
-                                      /*NumberOfRegions=*/1U);
+                                      /*NumberOfRegions=*/1U,
+                                      /*ReleaseSize=*/MaxBlocks * BlockSize);
     ASSERT_FALSE(Context.hasBlockMarked());
     Context.markFreeBlocks(FreeList, DecompactPtr, Recorder.getBase());
     ASSERT_TRUE(Context.hasBlockMarked());
     releaseFreeMemoryToOS(Context, Recorder, SkipRegion);
     scudo::RegionPageMap &PageMap = Context.PageMap;
 
     // Verify that there are no released pages touched by used chunks and all
     // ranges of free chunks big enough to contain the entire memory pages had
for (scudo::uptr Block = 0; Block + BlockSize <= RoundedRegionSize;
       }
     }
 
     for (scudo::uptr GroupId = 0; GroupId < GroupNum; ++GroupId) {
       const scudo::uptr GroupBeg = GroupId * GroupSize;
       const scudo::uptr GroupEnd = GroupBeg + GroupSize;
 
       scudo::PageReleaseContext Context(BlockSize, RegionSize,
-                                        /*NumberOfRegions=*/1U);
-      Context.markRangeAsAllCounted(GroupBeg, GroupEnd, /*Base=*/0);
+                                        /*NumberOfRegions=*/1U,
+                                        /*ReleaseSize=*/RegionSize);
+      Context.markRangeAsAllCounted(GroupBeg, GroupEnd, /*Base=*/0U);
 
       scudo::uptr FirstBlock =
           ((GroupBeg + BlockSize - 1) / BlockSize) * BlockSize;
 
       // All the pages before first block page are not supposed to be marked.
       if (FirstBlock / PageSize > 0) {
         for (scudo::uptr Page = 0; Page <= FirstBlock / PageSize - 1; ++Page)
           EXPECT_EQ(Context.PageMap.get(/*Region=*/0, Page), 0U);
for (scudo::uptr GroupId = 0; GroupId < GroupNum; ++GroupId) {
       for (scudo::uptr Page = FirstUncountedPage;
            Page <= RoundedRegionSize / PageSize; ++Page) {
         EXPECT_EQ(Context.PageMap.get(/*Region=*/0U, Page), 0U);
       }
     } // Iterate each Group
 
     // Release the entire region. This is to ensure the last page is counted.
     scudo::PageReleaseContext Context(BlockSize, RegionSize,
-                                      /*NumberOfRegions=*/1U);
+                                      /*NumberOfRegions=*/1U,
+                                      /*ReleaseSize=*/RegionSize);
     Context.markRangeAsAllCounted(/*From=*/0U, /*To=*/RegionSize, /*Base=*/0);
     for (scudo::uptr Page = 0; Page < RoundedRegionSize / PageSize; ++Page)
       EXPECT_TRUE(Context.PageMap.isAllCounted(/*Region=*/0, Page));
   } // Iterate each size class
 }
 
+template <class SizeClassMap> void testReleasePartialRegion() {
+  typedef FreeBatch<SizeClassMap> Batch;
+  const scudo::uptr PageSize = scudo::getPageSizeCached();
+  const scudo::uptr ReleaseBase = PageSize;
+  const scudo::uptr BasePageOffset = ReleaseBase / PageSize;
+
+  for (scudo::uptr I = 1; I <= SizeClassMap::LargestClassId; I++) {
+    // In the following, we want to ensure the region includes at least 2 pages
+    // and we will release all the pages except the first one. The handling of
+    // the last block is tricky, so we always test the case that includes the
+    // last block.
+    const scudo::uptr BlockSize = SizeClassMap::getSizeByClassId(I);
+    const scudo::uptr RegionSize =
+        scudo::roundUpSlow(scudo::roundUp(BlockSize, PageSize) * 2, BlockSize) +
+        BlockSize;
+    const scudo::uptr RoundedRegionSize = scudo::roundUp(RegionSize, PageSize);
+
+    scudo::SinglyLinkedList<Batch> FreeList;
+    FreeList.clear();
+
+    // Skip the blocks in the first page and add the remainings.
+    std::vector<scudo::uptr> Pages(RoundedRegionSize / PageSize, 0);
+    for (scudo::uptr Block = scudo::roundUpSlow(PageSize, BlockSize);
+         Block + BlockSize <= RoundedRegionSize; Block += BlockSize) {
+      for (scudo::uptr Page = Block / PageSize;
+           Page <= (Block + BlockSize - 1) / PageSize; ++Page) {
+        ASSERT_LT(Page, Pages.size());
+        ++Pages[Page];
+      }
+    }
+
+    // This follows the logic how we count the last page. It should be
+    // consistent with how markFreeBlocks() handles the last block.
+    if (RoundedRegionSize % BlockSize != 0)
+      ++Pages.back();
+
+    Batch *CurrentBatch = nullptr;
+    for (scudo::uptr Block = scudo::roundUpSlow(PageSize, BlockSize);
+         Block < RegionSize; Block += BlockSize) {
+      if (CurrentBatch == nullptr ||
+          CurrentBatch->getCount() == Batch::MaxCount) {
+        CurrentBatch = new Batch;
+        CurrentBatch->clear();
+        FreeList.push_back(CurrentBatch);
+      }
+      CurrentBatch->add(Block);
+    }
+
+    auto VerifyReleaseToOs = [&](scudo::PageReleaseContext &Context) {
+      auto SkipRegion = [](UNUSED scudo::uptr RegionIndex) { return false; };
+      ReleasedPagesRecorder Recorder(ReleaseBase);
+      releaseFreeMemoryToOS(Context, Recorder, SkipRegion);
+      const scudo::uptr FirstBlock = scudo::roundUpSlow(PageSize, BlockSize);
+
+      for (scudo::uptr P = 0; P < RoundedRegionSize; P += PageSize) {
+        if (P < FirstBlock) {
+          // If FirstBlock is not aligned with page boundary, the first touched
+          // page will not be released either.
+          EXPECT_TRUE(Recorder.ReportedPages.find(P) ==
+                      Recorder.ReportedPages.end());
+        } else {
+          EXPECT_TRUE(Recorder.ReportedPages.find(P) !=
+                      Recorder.ReportedPages.end());
+        }
+      }
+    };
+
+    // Test marking by visiting each block.
+    {
+      auto DecompactPtr = [](scudo::uptr P) { return P; };
+      scudo::PageReleaseContext Context(
+          BlockSize, RegionSize, /*NumberOfRegions=*/1U,
+          /*ReleaseSize=*/RegionSize - PageSize, ReleaseBase);
+      Context.markFreeBlocks(FreeList, DecompactPtr, /*Base=*/0U);
+      for (const Batch &It : FreeList) {
+        for (scudo::u16 I = 0; I < It.getCount(); I++) {
+          scudo::uptr Block = It.get(I);
+          for (scudo::uptr Page = Block / PageSize;
+               Page <= (Block + BlockSize - 1) / PageSize; ++Page) {
+            EXPECT_EQ(Pages[Page], Context.PageMap.get(/*Region=*/0U,
+                                                       Page - BasePageOffset));
+          }
+        }
+      }
+
+      VerifyReleaseToOs(Context);
+    }
+
+    // Test range marking.
+    {
+      scudo::PageReleaseContext Context(
+          BlockSize, RegionSize, /*NumberOfRegions=*/1U,
+          /*ReleaseSize=*/RegionSize - PageSize, ReleaseBase);
+      Context.markRangeAsAllCounted(ReleaseBase, RegionSize, /*Base=*/0U);
+      for (scudo::uptr Page = ReleaseBase / PageSize;
+           Page < RoundedRegionSize / PageSize; ++Page) {
+        if (Context.PageMap.get(/*Region=*/0, Page - BasePageOffset) !=
+            Pages[Page]) {
+          EXPECT_TRUE(Context.PageMap.isAllCounted(/*Region=*/0,
+                                                   Page - BasePageOffset));
+        }
+      }
+
+      VerifyReleaseToOs(Context);
+    }
+
+    // Check the buffer size of PageMap.
+    {
+      scudo::PageReleaseContext Full(BlockSize, RegionSize,
+                                     /*NumberOfRegions=*/1U,
+                                     /*ReleaseSize=*/RegionSize);
+      Full.ensurePageMapAllocated();
+      scudo::PageReleaseContext Partial(
+          BlockSize, RegionSize, /*NumberOfRegions=*/1U,
+          /*ReleaseSize=*/RegionSize - PageSize, ReleaseBase);
+      Partial.ensurePageMapAllocated();
+
+      EXPECT_GE(Full.PageMap.getBufferSize(), Partial.PageMap.getBufferSize());
+    }
+
+    while (!FreeList.empty()) {
+      CurrentBatch = FreeList.front();
+      FreeList.pop_front();
+      delete CurrentBatch;
+    }
+  } // Iterate each size class
+}
+
 TEST(ScudoReleaseTest, ReleaseFreeMemoryToOSDefault) {
   testReleaseFreeMemoryToOS<scudo::DefaultSizeClassMap>();
 }
 
 TEST(ScudoReleaseTest, ReleaseFreeMemoryToOSAndroid) {
   testReleaseFreeMemoryToOS<scudo::AndroidSizeClassMap>();
 }
 
 TEST(ScudoReleaseTest, ReleaseFreeMemoryToOSSvelte) {
   testReleaseFreeMemoryToOS<scudo::SvelteSizeClassMap>();
 }
 
 TEST(ScudoReleaseTest, PageMapMarkRange) {
   testPageMapMarkRange<scudo::DefaultSizeClassMap>();
   testPageMapMarkRange<scudo::AndroidSizeClassMap>();
   testPageMapMarkRange<scudo::FuchsiaSizeClassMap>();
   testPageMapMarkRange<scudo::SvelteSizeClassMap>();
 }
+
+TEST(ScudoReleaseTest, ReleasePartialRegion) {
+  testReleasePartialRegion<scudo::DefaultSizeClassMap>();
+  testReleasePartialRegion<scudo::AndroidSizeClassMap>();
+  testReleasePartialRegion<scudo::FuchsiaSizeClassMap>();
+  testReleasePartialRegion<scudo::SvelteSizeClassMap>();
+}