[scudo] Fix missing one block in range marking

[Chia-hungDuan]

Summary

When a range contains only one block, we may not mark the pages touched
by the block as can-be-released. This happens in the last group and if
it only contains single block.

Also enhance the existing tests and add a new test for testing the last
block.

[cferris]

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

552

Remove the at.

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

475

Should this be a DCHECK? It seems like a check of an error condition that could be doing unnecessary work, so best to only do on a test build.

Although, if you want to leave it CHECK and change it to DCHECK later once the code has time to run on real systems, that sounds like a good plan.

571

contains a single block that will be

[Chia-hungDuan]

Address review comment

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

475

Oops, this is a debug left over. Restored.

[cferris]

LGTM

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

if (FirstBlockInRange != FromInRegion) {
           (FromInRegion + PageSize - FirstBlockInRange + BlockSize - 1) /
           BlockSize;
       PageMap.incN(RegionIndex, getPageIndex(FromInRegion),
                    NumBlocksInFirstPage);
       FromInRegion = roundUp(FromInRegion + 1, PageSize);
     }
 
     uptr LastBlockInRange = roundDownSlow(ToInRegion - 1, BlockSize);
-    if (LastBlockInRange < FromInRegion)
-      return;
+
+    // Note that LastBlockInRange may be smaller than `FromInRegion` at this
+    // point because it may contain only one block in the range.
 
     // When the last block sits across `To`, we can't just mark the pages
     // occupied by the last block as all counted. Instead, we increment the
     // counters of those pages by 1. The exception is that if it's the last
     // block in the region, it's fine to mark those pages as all counted.
     if (LastBlockInRange + BlockSize != RegionSize) {
       DCHECK_EQ(ToInRegion % PageSize, 0U);
       // The case below is like,
if (MayContainLastBlockInRegion) {
       //  -----+-----------------------+
       //       |      Last Page        | <- Rounded Region Boundary
       //  -----+-----------------------+
       //   |-----||- trailing blocks  -|
       //      ^
       //   last block
       const uptr RoundedRegionSize = roundUp(RegionSize, PageSize);
       const uptr TrailingBlockBase = LastBlockInRegion + BlockSize;
-      // Only the last page touched by the last block needs to mark the trailing
-      // blocks. If the difference between `RoundedRegionSize` and
+      // If the difference between `RoundedRegionSize` and
       // `TrailingBlockBase` is larger than a page, that implies the reported
       // `RegionSize` may not be accurate.
       DCHECK_LT(RoundedRegionSize - TrailingBlockBase, PageSize);
+
+      // Only the last page touched by the last block needs to mark the trailing
+      // blocks. Note that if the last "pretend" block straddles the boundary,
+      // we still have to count it in so that the logic of counting the number
+      // of blocks on a page is consistent.
       uptr NumTrailingBlocks =
-          roundUpSlow(RoundedRegionSize - TrailingBlockBase, BlockSize) /
+          (roundUpSlow(RoundedRegionSize - TrailingBlockBase, BlockSize) +
+           BlockSize - 1) /
           BlockSize;
       if (NumTrailingBlocks > 0) {
         PageMap.incN(RegionIndex, getPageIndex(TrailingBlockBase),
                      NumTrailingBlocks);
       }
     }
 
     // Iterate over free chunks and count how many free chunks affect each

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

for (scudo::uptr I = 1; I <= SizeClassMap::LargestClassId; I++) {
 
     const scudo::uptr GroupNum = 2;
     const scudo::uptr GroupSize = scudo::roundUp(BlockSize, PageSize) * 2;
     const scudo::uptr RegionSize =
         scudo::roundUpSlow(GroupSize * GroupNum, BlockSize);
     const scudo::uptr RoundedRegionSize = scudo::roundUp(RegionSize, PageSize);
 
     std::vector<scudo::uptr> Pages(RoundedRegionSize / PageSize, 0);
-    for (scudo::uptr Block = 0; Block + BlockSize <= RoundedRegionSize;
-         Block += BlockSize) {
-      for (scudo::uptr page = Block / PageSize;
-           page <= (Block + BlockSize - 1) / PageSize; ++page) {
-        ASSERT_LT(page, Pages.size());
-        ++Pages[page];
+    for (scudo::uptr Block = 0; Block < RoundedRegionSize; Block += BlockSize) {
+      for (scudo::uptr Page = Block / PageSize;
+           Page <= (Block + BlockSize - 1) / PageSize &&
+           Page < RoundedRegionSize / PageSize;
+           ++Page) {
+        ASSERT_LT(Page, Pages.size());
+        ++Pages[Page];
       }
     }
 
     for (scudo::uptr GroupId = 0; GroupId < GroupNum; ++GroupId) {
       const scudo::uptr GroupBeg = GroupId * GroupSize;
       const scudo::uptr GroupEnd = GroupBeg + GroupSize;
 
       scudo::PageReleaseContext Context(BlockSize, /*NumberOfRegions=*/1U,
for (scudo::uptr I = 1; I <= SizeClassMap::LargestClassId; I++) {
     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 ReleaseBase = scudo::roundUp(BlockSize, PageSize);
+    const scudo::uptr BasePageOffset = ReleaseBase / PageSize;
     const scudo::uptr RegionSize =
-        scudo::roundUpSlow(scudo::roundUp(BlockSize, PageSize) * 2, BlockSize) +
+        scudo::roundUpSlow(scudo::roundUp(BlockSize, PageSize) + ReleaseBase,
+                           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 remaining.
     std::vector<scudo::uptr> Pages(RoundedRegionSize / PageSize, 0);
-    for (scudo::uptr Block = scudo::roundUpSlow(PageSize, BlockSize);
+    for (scudo::uptr Block = scudo::roundUpSlow(ReleaseBase, 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 markFreeBlocksInRegion() handles the last block.
     if (RoundedRegionSize % BlockSize != 0)
       ++Pages.back();
 
     Batch *CurrentBatch = nullptr;
-    for (scudo::uptr Block = scudo::roundUpSlow(PageSize, BlockSize);
+    for (scudo::uptr Block = scudo::roundUpSlow(ReleaseBase, 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);
+      const scudo::uptr FirstBlock = scudo::roundUpSlow(ReleaseBase, 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, /*NumberOfRegions=*/1U,
                                         /*ReleaseSize=*/RegionSize - PageSize,

 
 TEST(ScudoReleaseTest, ReleasePartialRegion) {
   testReleasePartialRegion<scudo::DefaultSizeClassMap>();
   testReleasePartialRegion<scudo::AndroidSizeClassMap>();
   testReleasePartialRegion<scudo::FuchsiaSizeClassMap>();
   testReleasePartialRegion<scudo::SvelteSizeClassMap>();
 }
 
+template <class SizeClassMap> void testReleaseRangeWithSingleBlock() {
+  const scudo::uptr PageSize = scudo::getPageSizeCached();
+
+  // We want to test if a memory group only contains single block that will be
+  // handled properly. The case is like:
+  //
+  //   From                     To
+  //     +----------------------+
+  //  +------------+------------+
+  //  |            |            |
+  //  +------------+------------+
+  //                            ^
+  //                        RegionSize
+  //
+  // Note that `From` will be page aligned.
+  //
+  // If the second from the last block is aligned at `From`, then we expect all
+  // the pages after `From` will be marked as can-be-released. Otherwise, the
+  // pages only touched by the last blocks will be marked as can-be-released.
+  for (scudo::uptr I = 1; I <= SizeClassMap::LargestClassId; I++) {
+    const scudo::uptr BlockSize = SizeClassMap::getSizeByClassId(I);
+    const scudo::uptr From = scudo::roundUp(BlockSize, PageSize);
+    const scudo::uptr To =
+        From % BlockSize == 0
+            ? From + BlockSize
+            : scudo::roundDownSlow(From + BlockSize, BlockSize) + BlockSize;
+    const scudo::uptr RoundedRegionSize = scudo::roundUp(To, PageSize);
+
+    std::vector<scudo::uptr> Pages(RoundedRegionSize / PageSize, 0);
+    for (scudo::uptr Block = (To - BlockSize); Block < RoundedRegionSize;
+         Block += BlockSize) {
+      for (scudo::uptr Page = Block / PageSize;
+           Page <= (Block + BlockSize - 1) / PageSize &&
+           Page < RoundedRegionSize / PageSize;
+           ++Page) {
+        ASSERT_LT(Page, Pages.size());
+        ++Pages[Page];
+      }
+    }
+
+    scudo::PageReleaseContext Context(BlockSize, /*NumberOfRegions=*/1U,
+                                      /*ReleaseSize=*/To,
+                                      /*ReleaseBase=*/0U);
+    Context.markRangeAsAllCounted(From, To, /*Base=*/0U, /*RegionIndex=*/0,
+                                  /*RegionSize=*/To);
+
+    for (scudo::uptr Page = 0; Page < RoundedRegionSize; Page += PageSize) {
+      if (Context.PageMap.get(/*Region=*/0U, Page / PageSize) !=
+          Pages[Page / PageSize]) {
+        EXPECT_TRUE(
+            Context.PageMap.isAllCounted(/*Region=*/0U, Page / PageSize));
+      }
+    }
+  } // for each size class
+}
+
+TEST(ScudoReleaseTest, RangeReleaseRegionWithSingleBlock) {
+  testReleaseRangeWithSingleBlock<scudo::DefaultSizeClassMap>();
+  testReleaseRangeWithSingleBlock<scudo::AndroidSizeClassMap>();
+  testReleaseRangeWithSingleBlock<scudo::FuchsiaSizeClassMap>();
+  testReleaseRangeWithSingleBlock<scudo::SvelteSizeClassMap>();
+}
+
 TEST(ScudoReleaseTest, BufferPool) {
   constexpr scudo::uptr StaticBufferCount = SCUDO_WORDSIZE - 1;
   constexpr scudo::uptr StaticBufferSize = 512U;
 
   // Allocate the buffer pool on the heap because it is quite large (slightly
   // more than StaticBufferCount * StaticBufferSize * sizeof(uptr)) and it may
   // not fit in the stack on some platforms.
   using BufferPool = scudo::BufferPool<StaticBufferCount, StaticBufferSize>;