Add method to detect holes in 2MiB regions of memory pages when a hug…

…e page aligned address is passed. On success, the method will return an ok status with the # of holes in the region or a boolean on whether the region is a huge page.

PiperOrigin-RevId: 684530002
Change-Id: Icba9ee370ee580df2c5905687b06062b49d4535a

tcmalloc/internal/BUILD

@@ -401,6 +401,7 @@ cc_test(
         "@com_google_absl//absl/log:check",
         "@com_google_absl//absl/random",
         "@com_google_absl//absl/random:distributions",
+        "@com_google_absl//absl/status",
         "@com_google_absl//absl/strings",
         "@com_google_absl//absl/strings:str_format",
         "@com_google_absl//absl/time",

tcmalloc/internal/pageflags.cc

@@ -14,14 +14,12 @@
 
 #include "tcmalloc/internal/pageflags.h"
 
-#include <errno.h>
 #include <fcntl.h>
 #include <linux/kernel-page-flags.h>
 #include <stddef.h>
 #include <unistd.h>
 
 #include <algorithm>
-#include <cinttypes>
 #include <cstdint>
 #include <cstdio>
 #include <optional>
@@ -41,6 +39,7 @@ namespace {
 // From include/uapi/linux/kernel-page-flags.h
 #define KPF_COMPOUND_HEAD 15
 #define KPF_COMPOUND_TAIL 16
+#define KPF_NOPAGE 20
 #define KPF_THP 22
 #ifndef KPF_MLOCKED
 #define KPF_MLOCKED 33
@@ -72,6 +71,10 @@ constexpr bool PageLocked(uint64_t flags) {
   constexpr uint64_t kPageUnevictable = (1UL << KPF_UNEVICTABLE);
   return (flags & (kPageMlocked | kPageUnevictable)) != 0;
 }
+constexpr bool PageHole(uint64_t flags) {
+  constexpr uint64_t kPageNoPage = (1UL << KPF_NOPAGE);
+  return (flags & kPageNoPage) == kPageNoPage;
+}
 void MaybeAddToStats(PageStats& stats, const uint64_t flags,
                      const size_t delta) {
   if (PageStale(flags)) stats.bytes_stale += delta;
@@ -327,6 +330,54 @@ uint64_t PageFlags::MaybeReadStaleScanSeconds(const char* filename) {
   return cached_scan_seconds_;
 }
 
+PageFlags::HoleInfo PageFlags::CountHolesInSinglePage(const void* const addr) {
+  uintptr_t currPage = reinterpret_cast<uintptr_t>(addr);
+  if ((currPage & kHugePageMask) != currPage) {
+    TC_LOG("Address is not hugepage aligned");
+    return HoleInfo{absl::StatusCode::kFailedPrecondition};
+  }
+  auto res = Seek(currPage);
+  if (res != absl::StatusCode::kOk) {
+    return HoleInfo{absl::StatusCode::kUnavailable};
+  }
+  const int kBufferSizeBasedOnEntries = kPagemapEntrySize * kPagesInHugePage;
+
+  TC_ASSERT(sizeof(buf_) >= kBufferSizeBasedOnEntries);
+  auto status = signal_safe_read(fd_, reinterpret_cast<char*>(buf_),
+                                 kBufferSizeBasedOnEntries, nullptr);
+  if (status != kBufferSizeBasedOnEntries) {
+    TC_LOG(
+        "Could not read from pageflags file due to unexpected number of bytes "
+        "read");
+    TC_LOG("Expected %d bytes, got %d bytes", kBufferSizeBasedOnEntries,
+           status);
+    return HoleInfo{absl::StatusCode::kUnavailable};
+  }
+
+  int32_t holes_count = 0;
+  for (int native_page_idx = 0; native_page_idx < kPagesInHugePage;
+       ++native_page_idx) {
+    uint64_t page_flags = buf_[native_page_idx];
+    // Case for hugepage
+    if (PageThp(page_flags)) {
+      if (ABSL_PREDICT_TRUE(native_page_idx == 0)) {
+        return HoleInfo{absl::StatusCode::kOk, /*num_holes=*/0,
+                        /*already_hugepage=*/true};
+        // Hugepage is found in the middle of a hugepage region
+      } else {
+        TC_LOG("Hugepage in middle of region, shouldn't happen");
+        return HoleInfo{absl::StatusCode::kFailedPrecondition};
+      }
+    }
+    // Case for a hole
+    if (PageHole(page_flags)) {
+      ++holes_count;
+    }
+  }
+  return HoleInfo{absl::StatusCode::kOk, holes_count,
+                  /*already_hugepage=*/false};
+}
+
 }  // namespace tcmalloc_internal
 }  // namespace tcmalloc
 GOOGLE_MALLOC_SECTION_END

tcmalloc/internal/pageflags.h

@@ -83,6 +83,12 @@ class PageFlags final : public PageFlagsBase {
   // reclaim it.
   PageFlags();
   ~PageFlags() override;
+
+  struct HoleInfo {
+    absl::StatusCode status;
+    int32_t num_holes;
+    bool already_hugepage;
+  };
   // Query a span of memory starting from `addr` for `size` bytes. The memory
   // span must consist of only native-size pages and THP hugepages; the behavior
   // is undefined if we encounter other hugepages (such as hugetlbfs). We try to
@@ -95,6 +101,11 @@ class PageFlags final : public PageFlagsBase {
   // use the function in places where memory allocation is prohibited.
   std::optional<PageStats> Get(const void* addr, size_t size) override;
 
+  // Count the number of holes in a single hugepage region when a hugepage
+  // aligned address is passed.  A HoleInfo object is returned with the status
+  // code, number of holes, and whether the region is  huge.
+  HoleInfo CountHolesInSinglePage(const void* addr);
+
  private:
   // Returns the offset in the pageflags file for the given virtual address.
   size_t GetOffset(uintptr_t vaddr);

tcmalloc/internal/pageflags_test.cc

@@ -28,6 +28,7 @@
 #include <cstring>
 #include <initializer_list>
 #include <iterator>
+#include <map>
 #include <memory>
 #include <optional>
 #include <ostream>
@@ -43,6 +44,7 @@
 #include "absl/log/log.h"
 #include "absl/random/distributions.h"
 #include "absl/random/random.h"
+#include "absl/status/status.h"
 #include "absl/strings/str_cat.h"
 #include "absl/strings/str_format.h"
 #include "absl/strings/string_view.h"
@@ -73,6 +75,10 @@ class PageFlagsFriend {
 
   decltype(auto) CachedScanSeconds() { return r_.cached_scan_seconds_; }
 
+  decltype(auto) CountHolesInSinglePage(const void* const addr) {
+    return r_.CountHolesInSinglePage(addr);
+  }
+
   void SetCachedScanSeconds(
       decltype(PageFlags::cached_scan_seconds_) scan_seconds) {
     r_.cached_scan_seconds_ = scan_seconds;
@@ -94,6 +100,7 @@ using ::testing::Optional;
 
 constexpr uint64_t kPageHead = (1UL << 15);
 constexpr uint64_t kPageTail = (1UL << 16);
+constexpr uint64_t kPageHole = (1UL << 20);
 constexpr uint64_t kPageThp = (1UL << 22);
 constexpr uint64_t kPageStale = (1UL << 44);
 
@@ -382,6 +389,157 @@ TEST(PageFlagsTest, OnlyTails) {
             std::nullopt);
 }
 
+// Method that can write a region with a single hugepage
+// a region with a single missing page, a region with every other page missing,
+// a region with all missing pages, or a region with a hugepage in the middle.
+void GenerateHolesInSinglePage(absl::string_view filename, int case_num,
+                               bool large_buffer) {
+  int write_fd = signal_safe_open(filename.data(), O_CREAT | O_WRONLY, S_IRUSR);
+  CHECK_NE(write_fd, -1) << errno;
+  // Create a large enough buffer to hold 4 hugepages worth of flags
+  constexpr size_t kPageSize = 4096;
+  int num_pages = large_buffer ? 2 << 13 : kHugePageSize / kPageSize;
+  std::vector<uint64_t> buf(num_pages, 0);
+  // constexpr int num_pages = static_cast<int>(kHugePageSize / kPageSize);
+  // std::array<uint64_t, num_pages> buf = {};
+  for (int i = 0; i < num_pages; i++) {
+    switch (case_num) {
+      case 0:
+        // First region is a thp hugepage
+        buf[i] = kPageThp;
+        if (i == 0) {
+          buf[i] |= kPageHead;
+        } else {
+          buf[i] |= kPageTail;
+        }
+        break;
+      case 1:
+        // First page is a hole, rest are stale
+        if (i == 0) {
+          buf[i] = kPageHole;
+        } else {
+          buf[i] = kPageStale;
+        }
+        break;
+      case 2:
+        // Every other page is a hole, rest are stale
+        if (i % 2 == 0) {
+          buf[i] = kPageHole;
+        } else if (i % 2 == 1) {
+          buf[i] = kPageStale;
+        }
+        break;
+      case 3:
+        // All pages are holes
+        buf[i] = kPageHole;
+        break;
+      case 4:
+        // Hugepage in the middle of region
+        if (i == 256) {
+          buf[i] = kPageThp;
+        } else {
+          buf[i] = kPageHole;
+        }
+        break;
+    }
+  }
+  int size_of_write = large_buffer ? (2 << 13) : kPageSize;
+  CHECK_EQ(write(write_fd, buf.data(), size_of_write), size_of_write);
+  CHECK_EQ(close(write_fd), 0) << errno;
+}
+
+TEST(PageFlagsTest, CountHolesInSinglePage) {
+  constexpr int kNumCases = 4;
+  std::array<PageFlags::HoleInfo, kNumCases> expected = {
+      PageFlags::HoleInfo{.status = absl::StatusCode::kOk,
+                          .num_holes = 0,
+                          .already_hugepage = true},
+      PageFlags::HoleInfo{.status = absl::StatusCode::kOk,
+                          .num_holes = 1,
+                          .already_hugepage = false},
+      PageFlags::HoleInfo{.status = absl::StatusCode::kOk,
+                          .num_holes = 256,
+                          .already_hugepage = false},
+      PageFlags::HoleInfo{.status = absl::StatusCode::kOk,
+                          .num_holes = 512,
+                          .already_hugepage = false},
+  };
+  std::optional<AllocationGuard> g;
+
+  for (int i = 0; i < kNumCases; ++i) {
+    std::string file_path =
+        absl::StrCat(testing::TempDir(), "/holes_in_single_page_", i);
+    GenerateHolesInSinglePage(file_path, /*case_num=*/i,
+                              /*large_buffer=*/false);
+
+    g.emplace();
+    PageFlagsFriend s(file_path);
+    PageFlags::HoleInfo res =
+        s.CountHolesInSinglePage(reinterpret_cast<void*>(0));
+    g.reset();
+    EXPECT_THAT(res.status, expected[i].status);
+    EXPECT_THAT(res.num_holes, expected[i].num_holes);
+    EXPECT_THAT(res.already_hugepage, expected[i].already_hugepage);
+  }
+}
+
+TEST(PageFlagsTest, CountHolesWithAddressBeyondFirstPage) {
+  std::optional<AllocationGuard> g;
+  std::string file_path =
+      absl::StrCat(testing::TempDir(), "/holes_in_single_page");
+  GenerateHolesInSinglePage(file_path, /*case_num=*/3,
+                            /*large_buffer=*/true);
+  PageFlags::HoleInfo expected =
+      PageFlags::HoleInfo{.status = absl::StatusCode::kOk,
+                          .num_holes = 512,
+                          .already_hugepage = false};
+  g.emplace();
+  PageFlagsFriend s(file_path);
+  PageFlags::HoleInfo res =
+      s.CountHolesInSinglePage(reinterpret_cast<void*>(2 << 21));
+  g.reset();
+  EXPECT_THAT(res.status, expected.status);
+  EXPECT_THAT(res.num_holes, expected.num_holes);
+  EXPECT_THAT(res.already_hugepage, expected.already_hugepage);
+}
+
+TEST(PageFlagsTest, VerifyAddressAlignmentCheckPasses) {
+  std::optional<AllocationGuard> g;
+  std::string file_path = absl::StrCat(testing::TempDir(), "asd");
+  GenerateHolesInSinglePage(file_path, /*case_num=*/0,
+                            /*large_buffer=*/false);
+  g.emplace();
+  PageFlagsFriend s(file_path);
+  PageFlags::HoleInfo non_align_addr_res =
+      s.CountHolesInSinglePage(reinterpret_cast<void*>(0x00001));
+  g.reset();
+  EXPECT_EQ(non_align_addr_res.status, absl::StatusCode::kFailedPrecondition);
+}
+
+TEST(PageFlagsTest, VerifyHugePageInMiddleOfRegionFails) {
+  std::optional<AllocationGuard> g;
+  std::string file_path =
+      absl::StrCat(testing::TempDir(), "/page_head_in_middle_of_region");
+  GenerateHolesInSinglePage(file_path, /*case_num=*/4,
+                            /*large_buffer=*/false);
+  g.emplace();
+  PageFlagsFriend s(file_path);
+  PageFlags::HoleInfo res =
+      s.CountHolesInSinglePage(reinterpret_cast<void*>(0));
+  g.reset();
+  EXPECT_EQ(res.status, absl::StatusCode::kFailedPrecondition);
+}
+
+TEST(PageFlagsTest, VerifyAddressAlignmentBeyondFirstPageFails) {
+  std::optional<AllocationGuard> g;
+  g.emplace();
+  PageFlagsFriend s;
+  PageFlags::HoleInfo res =
+      s.CountHolesInSinglePage(reinterpret_cast<void*>((2 << 21) + 1));
+  g.reset();
+  EXPECT_EQ(res.status, absl::StatusCode::kFailedPrecondition);
+}
+
 TEST(PageFlagsTest, TooManyTails) {
   const size_t kPageSize = getpagesize();
   std::vector<uint64_t> data(7 * kHugePageSize / kPageSize);