Optimize HugeRegion::Release using RangeTracker::PrevFreeRange.

tcmalloc/huge_region.h

@@ -331,18 +331,17 @@ inline void HugeRegion::Put(Range r, bool release) {
 // sophisticated mechanism similar to Filler/Cache, that accounts for a recent
 // peak while releasing pages.
 inline HugeLength HugeRegion::Release(Length desired, bool adaptive_release) {
-  if (desired == Length(0)) return NHugePages(0);
-
   const Length free_yet_backed = free_backed_count_.in_pages();
   const Length to_release = std::min(desired, free_yet_backed);
+  if (to_release == Length(0)) return NHugePages(0);
 
-  HugeLength release_target = NHugePages(0);
   bool should_unback[kNumHugePages] = {};
+#ifdef TCMALLOC_INTERNAL_LEGACY_LOCKING
+  HugeLength release_target = NHugePages(0);
   const int start = adaptive_release ? kNumHugePages - 1 : 0;
   const int end = adaptive_release ? -1 : kNumHugePages;
   const int step = adaptive_release ? -1 : 1;
 
-  // TODO(b/73749855): Consider optimizing this search by consulting tracker_.
   for (int i = start; i != end; i += step) {
     if (backed_[i] && pages_used_[i] == Length(0)) {
       should_unback[i] = true;
@@ -351,6 +350,47 @@ inline HugeLength HugeRegion::Release(Length desired, bool adaptive_release) {
 
     if (release_target.in_pages() >= to_release) break;
   }
+#else
+  const size_t needed = HLFromPages(to_release).raw_num();
+  size_t released = 0;
+
+  if (adaptive_release) {
+    size_t index = tracker_.size(), n;
+    while (tracker_.PrevFreeRange(index, &index, &n)) {
+      if (n < kPagesPerHugePage.raw_num()) continue;
+      const size_t first_hp = (index + kPagesPerHugePage.raw_num() - 1) /
+                              kPagesPerHugePage.raw_num();
+      const size_t last_hp = (index + n) / kPagesPerHugePage.raw_num();
+      for (size_t i = last_hp; i > first_hp; --i) {
+        const size_t hp = i - 1;
+        if (!backed_[hp]) continue;
+        TC_ASSERT_EQ(pages_used_[hp], Length(0));
+        should_unback[hp] = true;
+        if (++released == needed) break;
+      }
+      if (released == needed) break;
+    }
+  } else {
+    size_t index = 0, n;
+    while (tracker_.NextFreeRange(index, &index, &n)) {
+      if (n < kPagesPerHugePage.raw_num()) {
+        index += n;
+        continue;
+      }
+      const size_t first_hp = (index + kPagesPerHugePage.raw_num() - 1) /
+                              kPagesPerHugePage.raw_num();
+      const size_t last_hp = (index + n) / kPagesPerHugePage.raw_num();
+      for (size_t hp = first_hp; hp < last_hp; ++hp) {
+        if (!backed_[hp]) continue;
+        TC_ASSERT_EQ(pages_used_[hp], Length(0));
+        should_unback[hp] = true;
+        if (++released == needed) break;
+      }
+      if (released == needed) break;
+      index += n;
+    }
+  }
+#endif  // TCMALLOC_INTERNAL_LEGACY_LOCKING
   return UnbackHugepages(should_unback);
 }
 

tcmalloc/huge_region_test.cc

@@ -315,6 +315,66 @@ TEST_F(HugeRegionTest, ReleaseAdaptive) {
   CheckMock();
 }
 
+TEST_F(HugeRegionTest, ReleaseAdaptiveFragmented) {
+  const Length n = kPagesPerHugePage;
+  bool from_released;
+  std::optional<Alloc> allocs[8];
+  for (int i = 0; i < 8; ++i) {
+    allocs[i] = Allocate(n, &from_released);
+    EXPECT_TRUE(from_released);
+  }
+
+  // Delete hugepages: 0, 1, 3, 4, 6, 7 while keeping 2 and 5 allocated.
+  Delete(*allocs[0]);
+  Delete(*allocs[1]);
+  Delete(*allocs[3]);
+  Delete(*allocs[4]);
+  Delete(*allocs[6]);
+  Delete(*allocs[7]);
+
+  // Free backed hugepages are: [0, 1], [3, 4], [6, 7].
+  // Reverse order should release: 7, 6 from the last range, then 4 from the
+  // middle range.
+  ExpectUnback({p_ + NHugePages(4), NHugePages(1)});
+  ExpectUnback({p_ + NHugePages(6), NHugePages(2)});
+  EXPECT_EQ(NHugePages(3), region_.Release(NHugePages(3).in_pages(),
+                                           /*adaptive_release=*/true));
+  CheckMock();
+
+  Delete(*allocs[2]);
+  Delete(*allocs[5]);
+}
+
+TEST_F(HugeRegionTest, ReleaseNonAdaptiveFragmented) {
+  const Length n = kPagesPerHugePage;
+  bool from_released;
+  std::optional<Alloc> allocs[8];
+  for (int i = 0; i < 8; ++i) {
+    allocs[i] = Allocate(n, &from_released);
+    EXPECT_TRUE(from_released);
+  }
+
+  // Delete hugepages: 0, 1, 3, 4, 6, 7 while keeping 2 and 5 allocated.
+  Delete(*allocs[0]);
+  Delete(*allocs[1]);
+  Delete(*allocs[3]);
+  Delete(*allocs[4]);
+  Delete(*allocs[6]);
+  Delete(*allocs[7]);
+
+  // Free backed hugepages are: [0, 1], [3, 4], [6, 7].
+  // Forward order should release: 0, 1 from the first range, then 3 from the
+  // middle range.
+  ExpectUnback({p_, NHugePages(2)});
+  ExpectUnback({p_ + NHugePages(3), NHugePages(1)});
+  EXPECT_EQ(NHugePages(3), region_.Release(NHugePages(3).in_pages(),
+                                           /*adaptive_release=*/false));
+  CheckMock();
+
+  Delete(*allocs[2]);
+  Delete(*allocs[5]);
+}
+
 TEST_F(HugeRegionTest, ReleaseFailure) {
   const Length n = kPagesPerHugePage;
   bool from_released;

tcmalloc/internal/range_tracker.h

@@ -81,6 +81,12 @@ class Bitmap {
   // If there is at least one free range at or after <start>,
   // put it in *index, *length and return true; else return false.
   bool NextFreeRange(size_t start, size_t* index, size_t* length) const;
+  // If there is at least one free range before <end>,
+  // put it in *index, *length and return true; else return false.
+  // To iterate backwards through all free ranges:
+  //   size_t index = bitmap.size(), n;
+  //   while (bitmap.PrevFreeRange(index, &index, &n)) { ... }
+  bool PrevFreeRange(size_t end, size_t* index, size_t* length) const;
 
   // Returns index of the first {true, false} bit >= index, or N if none.
   size_t FindSet(size_t index) const;
@@ -166,6 +172,12 @@ class RangeTracker {
   // If there is at least one free range at or after <start>,
   // put it in *index, *length and return true; else return false.
   bool NextFreeRange(size_t start, size_t* index, size_t* length) const;
+  // If there is at least one free range before <end>,
+  // put it in *index, *length and return true; else return false.
+  // To iterate backwards through all free ranges:
+  //   size_t index = tracker.size(), n;
+  //   while (tracker.PrevFreeRange(index, &index, &n)) { ... }
+  bool PrevFreeRange(size_t end, size_t* index, size_t* length) const;
 
   void Clear();
 
@@ -331,6 +343,14 @@ inline bool RangeTracker<N>::NextFreeRange(size_t start, size_t* index,
   return bits_.NextFreeRange(start, index, length);
 }
 
+// If there is at least one free range before <end>,
+// put it in *index, *length and return true; else return false.
+template <size_t N>
+inline bool RangeTracker<N>::PrevFreeRange(size_t end, size_t* index,
+                                           size_t* length) const {
+  return bits_.PrevFreeRange(end, index, length);
+}
+
 template <size_t N>
 inline void RangeTracker<N>::Clear() {
   bits_.Clear();
@@ -495,6 +515,20 @@ inline bool Bitmap<N>::NextFreeRange(size_t start, size_t* index,
   return true;
 }
 
+template <size_t N>
+inline bool Bitmap<N>::PrevFreeRange(size_t end, size_t* index,
+                                     size_t* length) const {
+  if (end == 0) return false;
+  if (end > N) end = N;
+  ssize_t j = FindClearBackwards(end - 1);
+  if (j < 0) return false;
+  ssize_t i = FindSetBackwards(j);
+  size_t range_start = (i < 0) ? 0 : static_cast<size_t>(i + 1);
+  *index = range_start;
+  *length = static_cast<size_t>(j + 1) - range_start;
+  return true;
+}
+
 template <size_t N>
 inline size_t Bitmap<N>::FindSet(size_t index) const {
   return FindValue<true>(index);

tcmalloc/internal/range_tracker_test.cc

@@ -314,6 +314,71 @@ TEST_F(BitmapTest, PopBatch) {
   EXPECT_TRUE(map.IsZero());
 }
 
+TEST_F(BitmapTest, PrevFreeRange) {
+  // Empty bitmap (all free).
+  {
+    Bitmap<253> map;
+    size_t index = map.size(), len = 0;
+    EXPECT_TRUE(map.PrevFreeRange(index, &index, &len));
+    EXPECT_EQ(index, 0);
+    EXPECT_EQ(len, 253);
+    EXPECT_FALSE(map.PrevFreeRange(index, &index, &len));
+  }
+
+  // Full bitmap (no free ranges).
+  {
+    Bitmap<253> map;
+    map.SetRange(0, 253);
+    size_t index = map.size(), len = 0;
+    EXPECT_FALSE(map.PrevFreeRange(index, &index, &len));
+  }
+
+  // Bounds and edge conditions.
+  {
+    Bitmap<64> map;
+    size_t index = 0, len = 0;
+    EXPECT_FALSE(map.PrevFreeRange(0, &index, &len));
+
+    // end > N clamps to N.
+    EXPECT_TRUE(map.PrevFreeRange(100, &index, &len));
+    EXPECT_EQ(index, 0);
+    EXPECT_EQ(len, 64);
+
+    // Partial query: end inside a free range [0, 64).
+    EXPECT_TRUE(map.PrevFreeRange(30, &index, &len));
+    EXPECT_EQ(index, 0);
+    EXPECT_EQ(len, 30);
+  }
+
+  // Multiple ranges matching forward traversal in reverse.
+  {
+    Bitmap<253> map;
+    map.SetRange(0, 253);
+    // Open free ranges: [10, 20), [50, 80), [120, 150), [200, 253)
+    map.ClearRange(10, 10);
+    map.ClearRange(50, 30);
+    map.ClearRange(120, 30);
+    map.ClearRange(200, 53);
+
+    std::vector<std::pair<size_t, size_t>> forward_ranges;
+    size_t f_index = 0, f_len;
+    while (map.NextFreeRange(f_index, &f_index, &f_len)) {
+      forward_ranges.push_back({f_index, f_len});
+      f_index += f_len;
+    }
+    EXPECT_THAT(forward_ranges, ElementsAre(Pair(10, 10), Pair(50, 30),
+                                            Pair(120, 30), Pair(200, 53)));
+
+    std::vector<std::pair<size_t, size_t>> backward_ranges;
+    size_t b_index = map.size(), b_len;
+    while (map.PrevFreeRange(b_index, &b_index, &b_len)) {
+      backward_ranges.push_back({b_index, b_len});
+    }
+    EXPECT_THAT(backward_ranges, ElementsAre(Pair(200, 53), Pair(120, 30),
+                                             Pair(50, 30), Pair(10, 10)));
+  }
+}
+
 class RangeTrackerTest : public ::testing::Test {
  protected:
   std::vector<std::pair<size_t, size_t>> FreeRanges() {
@@ -325,6 +390,14 @@ class RangeTrackerTest : public ::testing::Test {
     }
     return ret;
   }
+  std::vector<std::pair<size_t, size_t>> FreeRangesBackwards() {
+    std::vector<std::pair<size_t, size_t>> ret;
+    size_t index = range_.size(), len;
+    while (range_.PrevFreeRange(index, &index, &len)) {
+      ret.push_back({index, len});
+    }
+    return ret;
+  }
   static constexpr size_t kBits = 1017;
   RangeTracker<kBits> range_;
 };
@@ -334,23 +407,28 @@ TEST_F(RangeTrackerTest, Trivial) {
   EXPECT_EQ(0, range_.used());
   EXPECT_EQ(kBits, range_.longest_free());
   EXPECT_THAT(FreeRanges(), ElementsAre(Pair(0, kBits)));
+  EXPECT_THAT(FreeRangesBackwards(), ElementsAre(Pair(0, kBits)));
   ASSERT_EQ(0, range_.FindAndMark(kBits));
   EXPECT_EQ(0, range_.longest_free());
   EXPECT_EQ(kBits, range_.used());
   EXPECT_THAT(FreeRanges(), ElementsAre());
+  EXPECT_THAT(FreeRangesBackwards(), ElementsAre());
   range_.Unmark(0, 100);
   EXPECT_EQ(100, range_.longest_free());
   EXPECT_EQ(kBits - 100, range_.used());
   EXPECT_THAT(FreeRanges(), ElementsAre(Pair(0, 100)));
+  EXPECT_THAT(FreeRangesBackwards(), ElementsAre(Pair(0, 100)));
   // non-contiguous - shouldn't increase longest
   range_.Unmark(200, 100);
   EXPECT_EQ(100, range_.longest_free());
   EXPECT_EQ(kBits - 200, range_.used());
   EXPECT_THAT(FreeRanges(), ElementsAre(Pair(0, 100), Pair(200, 100)));
+  EXPECT_THAT(FreeRangesBackwards(), ElementsAre(Pair(200, 100), Pair(0, 100)));
   range_.Unmark(100, 100);
   EXPECT_EQ(300, range_.longest_free());
   EXPECT_EQ(kBits - 300, range_.used());
   EXPECT_THAT(FreeRanges(), ElementsAre(Pair(0, 300)));
+  EXPECT_THAT(FreeRangesBackwards(), ElementsAre(Pair(0, 300)));
 }
 
 TEST_F(RangeTrackerTest, Mark) {
@@ -359,6 +437,8 @@ TEST_F(RangeTrackerTest, Mark) {
   EXPECT_EQ(range_.used(), 100);
   EXPECT_EQ(range_.longest_free(), kBits - 200);
   EXPECT_THAT(FreeRanges(), ElementsAre(Pair(0, 100), Pair(200, kBits - 200)));
+  EXPECT_THAT(FreeRangesBackwards(),
+              ElementsAre(Pair(200, kBits - 200), Pair(0, 100)));
   range_.Unmark(100, 100);
   EXPECT_EQ(range_.used(), 0);
   EXPECT_EQ(range_.longest_free(), kBits);