Strongly type alignment inside of TCMalloc.

PiperOrigin-RevId: 855477238
Change-Id: I3565a35159f1820b4815bbc3f956b2d52bd79e9b

Author: ckennelly

tcmalloc/allocation_sampling.h

@@ -116,9 +116,7 @@ ABSL_ATTRIBUTE_NOINLINE sized_ptr_t SampleifyAllocation(
   stack_trace.depth = absl::GetStackTrace(stack_trace.stack, kMaxStackDepth, 0);
 
   if (policy.has_explicit_alignment()) {
-    // TODO(b/404341539): Make policy.align() a std::align_val_t.
-    stack_trace.requested_alignment =
-        static_cast<std::align_val_t>(policy.align());
+    stack_trace.requested_alignment = policy.align();
   } else {
     stack_trace.requested_alignment = std::nullopt;
   }
@@ -146,8 +144,7 @@ ABSL_ATTRIBUTE_NOINLINE sized_ptr_t SampleifyAllocation(
     stack_trace.cold_allocated = IsExpandedSizeClass(size_class);
 
     Length num_pages = BytesToLengthCeil(stack_trace.allocated_size);
-    size_t sample_alignment = policy.align();
-    if (sample_alignment == 1) sample_alignment = 0;
+    std::align_val_t sample_alignment = policy.align();
     alloc_with_status = state.guardedpage_allocator().TrySample(
         requested_size, sample_alignment, num_pages, stack_trace);
     if (alloc_with_status.status == Profile::Sample::GuardedStatus::Guarded) {
@@ -382,8 +379,7 @@ void MaybeUnsampleAllocation(Static& state, Policy policy,
           std::align_val_t{1}));
   const std::optional<std::align_val_t> deallocated_alignment =
       policy.has_explicit_alignment()
-          ? std::make_optional<std::align_val_t>(
-                static_cast<std::align_val_t>(policy.align()))
+          ? std::make_optional<std::align_val_t>(policy.align())
           : std::nullopt;
 
   if ((size.has_value() ||

tcmalloc/guarded_page_allocator.cc

@@ -96,7 +96,7 @@ void GuardedPageAllocator::Reset() {
 }
 
 GuardedAllocWithStatus GuardedPageAllocator::TrySample(
-    size_t size, size_t alignment, Length num_pages,
+    size_t size, std::align_val_t alignment, Length num_pages,
     const StackTrace& stack_trace) {
   if (num_pages > Length(1)) {
     skipped_allocations_toolarge_.Add(1);
@@ -165,16 +165,17 @@ GuardedAllocWithStatus GuardedPageAllocator::TrySample(
 }
 
 GuardedAllocWithStatus GuardedPageAllocator::Allocate(
-    size_t size, size_t alignment, const StackTrace& stack_trace) {
+    size_t size, std::align_val_t alignment, const StackTrace& stack_trace) {
   const ssize_t free_slot = ReserveFreeSlot();
   if (free_slot == -1) {
     // All slots are reserved.
     return {nullptr, Profile::Sample::GuardedStatus::NoAvailableSlots};
   }
 
   TC_ASSERT_LE(size, page_size_);
-  TC_ASSERT_LE(alignment, page_size_);
-  TC_ASSERT(alignment == 0 || absl::has_single_bit(alignment));
+  TC_ASSERT_LE(static_cast<size_t>(alignment), page_size_);
+  TC_ASSERT(static_cast<size_t>(alignment) == 0 ||
+            absl::has_single_bit(static_cast<size_t>(alignment)));
   void* result = reinterpret_cast<void*>(SlotToAddr(free_slot));
 
   // For size == 0, the page remains protected.
@@ -206,13 +207,13 @@ GuardedAllocWithStatus GuardedPageAllocator::Allocate(
   d.dealloc_trace.depth = 0;
   d.requested_size = size;
   d.requested_alignment = static_cast<std::align_val_t>(std::max(
-      alignment,
-      std::max(static_cast<size_t>(kAlignment),
-               static_cast<size_t>(__STDCPP_DEFAULT_NEW_ALIGNMENT__))));
+      alignment, std::max(kAlignment, static_cast<std::align_val_t>(
+                                          __STDCPP_DEFAULT_NEW_ALIGNMENT__))));
   d.allocation_start = reinterpret_cast<uintptr_t>(result);
   d.dealloc_count.store(0, std::memory_order_relaxed);
   TC_ASSERT(!d.write_overflow_detected);
-  TC_ASSERT(!alignment || d.allocation_start % alignment == 0);
+  TC_ASSERT(static_cast<size_t>(alignment) == 0 ||
+            d.allocation_start % static_cast<size_t>(alignment) == 0);
 
   stacktrace_filter_.Add({stack_trace.stack, stack_trace.depth}, 1);
   return {result, Profile::Sample::GuardedStatus::Guarded};
@@ -535,7 +536,8 @@ size_t GuardedPageAllocator::AddrToSlot(uintptr_t addr) const {
 }
 
 void GuardedPageAllocator::MaybeRightAlign(size_t slot, size_t size,
-                                           size_t alignment, void** ptr) {
+                                           std::align_val_t alignment,
+                                           void** ptr) {
   if (!ShouldRightAlign(slot)) return;
   uintptr_t adjusted_ptr =
       reinterpret_cast<uintptr_t>(*ptr) + page_size_ - size;
@@ -552,8 +554,9 @@ void GuardedPageAllocator::MaybeRightAlign(size_t slot, size_t size,
                         static_cast<size_t>(__STDCPP_DEFAULT_NEW_ALIGNMENT__)));
 
   // Ensure valid alignment.
-  alignment = std::max(alignment, default_alignment);
-  uintptr_t alignment_padding = adjusted_ptr & (alignment - 1);
+  const uintptr_t mask =
+      std::max(static_cast<size_t>(alignment), default_alignment) - 1;
+  uintptr_t alignment_padding = adjusted_ptr & mask;
   adjusted_ptr -= alignment_padding;
 
   // Write magic bytes in alignment padding to detect small overflow writes.

tcmalloc/guarded_page_allocator.h

@@ -117,7 +117,7 @@ class GuardedPageAllocator {
   // returns an instance of GuardedAllocWithStatus, that includes guarded
   // allocation Span and guarded status. Otherwise, returns nullptr and the
   // status indicating why the allocation may not be guarded.
-  GuardedAllocWithStatus TrySample(size_t size, size_t alignment,
+  GuardedAllocWithStatus TrySample(size_t size, std::align_val_t alignment,
                                    Length num_pages,
                                    const StackTrace& stack_trace);
 
@@ -133,7 +133,7 @@ class GuardedPageAllocator {
   //
   // Precondition:  size and alignment <= page_size_
   // Precondition:  alignment is 0 or a power of 2
-  GuardedAllocWithStatus Allocate(size_t size, size_t alignment,
+  GuardedAllocWithStatus Allocate(size_t size, std::align_val_t alignment,
                                   const StackTrace& stack_trace)
       ABSL_LOCKS_EXCLUDED(guarded_page_lock_);
 
@@ -263,7 +263,7 @@ class GuardedPageAllocator {
   // If slot is marked for right alignment, moves the allocation in *ptr to the
   // right end of the slot, maintaining the specified size and alignment.  Magic
   // bytes are written in any alignment padding.
-  void MaybeRightAlign(size_t slot, size_t size, size_t alignment,
+  void MaybeRightAlign(size_t slot, size_t size, std::align_val_t alignment,
                        void** absl_nonnull ptr);
 
   uintptr_t SlotToAddr(size_t slot) const;

tcmalloc/guarded_page_allocator_benchmark.cc

@@ -17,6 +17,7 @@
 #include <cstddef>
 #include <functional>
 #include <memory>
+#include <new>
 #include <vector>
 
 #include "benchmark/benchmark.h"
@@ -79,7 +80,7 @@ void BM_AllocDealloc(benchmark::State& state) {
   auto gpa = GetGuardedPageAllocator();
   for (auto _ : state) {
     char* ptr = reinterpret_cast<char*>(
-        gpa->Allocate(alloc_size, 0, GetStackTrace(0)).alloc);
+        gpa->Allocate(alloc_size, std::align_val_t{0}, GetStackTrace(0)).alloc);
     TC_CHECK_NE(ptr, nullptr);
     ptr[0] = 'X';               // Page fault first page.
     ptr[alloc_size - 1] = 'X';  // Page fault last page.
@@ -103,7 +104,8 @@ void ReservePool(const benchmark::State&) {
   auto deleter = [gpa](void* p) { gpa->Deallocate(p); };
 
   for (size_t stack_idx = 0;;) {
-    auto alloc = gpa->TrySample(1, 0, Length(1), GetStackTrace(stack_idx));
+    auto alloc = gpa->TrySample(1, std::align_val_t{0}, Length(1),
+                                GetStackTrace(stack_idx));
     switch (alloc.status) {
       case GuardedStatus::NoAvailableSlots:
         TC_CHECK(!GetReserved().empty());
@@ -143,7 +145,8 @@ void BM_TrySample(benchmark::State& state) {
 
   for (auto _ : state) {
     StackTrace& stack_trace = GetStackTrace(stack_idx++ % GetReserved().size());
-    auto alloc = gpa->TrySample(alloc_size, 0, Length(1), stack_trace);
+    auto alloc =
+        gpa->TrySample(alloc_size, std::align_val_t{0}, Length(1), stack_trace);
 
     switch (alloc.status) {
       case GuardedStatus::RateLimited:

tcmalloc/guarded_page_allocator_test.cc

@@ -83,7 +83,8 @@ class GuardedPageAllocatorParamTest
 };
 
 TEST_F(GuardedPageAllocatorTest, SingleAllocDealloc) {
-  auto alloc_with_status = gpa_.Allocate(PageSize(), 0, GetStackTrace());
+  auto alloc_with_status =
+      gpa_.Allocate(PageSize(), std::align_val_t{0}, GetStackTrace());
   EXPECT_EQ(alloc_with_status.status, Profile::Sample::GuardedStatus::Guarded);
   EXPECT_EQ(gpa_.successful_allocations(), 1);
   char* buf = static_cast<char*>(alloc_with_status.alloc);
@@ -116,7 +117,8 @@ TEST_F(GuardedPageAllocatorTest, NoAlignmentProvided) {
       // Make several allocation attempts to encounter left/right-alignment in
       // the guarded region.
       for (int i = 0; i < kElements; i++) {
-        auto alloc_with_status = gpa_.Allocate(size, 0, GetStackTrace());
+        auto alloc_with_status =
+            gpa_.Allocate(size, std::align_val_t{0}, GetStackTrace());
         EXPECT_EQ(alloc_with_status.status,
                   Profile::Sample::GuardedStatus::Guarded);
         ptrs[i] = alloc_with_status.alloc;
@@ -140,7 +142,8 @@ TEST_F(GuardedPageAllocatorTest, NoAlignmentProvided) {
 TEST_F(GuardedPageAllocatorTest, AllocDeallocAligned) {
   for (size_t align = 1; align <= PageSize(); align <<= 1) {
     constexpr size_t alloc_size = 1;
-    auto alloc_with_status = gpa_.Allocate(alloc_size, align, GetStackTrace());
+    auto alloc_with_status = gpa_.Allocate(
+        alloc_size, static_cast<std::align_val_t>(align), GetStackTrace());
     EXPECT_EQ(alloc_with_status.status,
               Profile::Sample::GuardedStatus::Guarded);
     EXPECT_NE(alloc_with_status.alloc, nullptr);
@@ -157,8 +160,8 @@ TEST_F(GuardedPageAllocatorTest, MismatchedAlignment) {
   for (size_t align = 1; align <= PageSize(); align <<= 1) {
     for (size_t misalign = 1; misalign <= align; misalign <<= 1) {
       constexpr size_t alloc_size = 1;
-      auto alloc_with_status =
-          gpa_.Allocate(alloc_size, align, GetStackTrace());
+      auto alloc_with_status = gpa_.Allocate(
+          alloc_size, static_cast<std::align_val_t>(align), GetStackTrace());
       EXPECT_EQ(alloc_with_status.status,
                 Profile::Sample::GuardedStatus::Guarded);
       EXPECT_NE(alloc_with_status.alloc, nullptr);
@@ -182,20 +185,22 @@ TEST_P(GuardedPageAllocatorParamTest, AllocDeallocAllPages) {
   size_t num_pages = GetParam();
   char* bufs[kMaxGpaPages];
   for (size_t i = 0; i < num_pages; i++) {
-    auto alloc_with_status = gpa_.Allocate(1, 0, GetStackTrace());
+    auto alloc_with_status =
+        gpa_.Allocate(1, std::align_val_t{0}, GetStackTrace());
     EXPECT_EQ(alloc_with_status.status,
               Profile::Sample::GuardedStatus::Guarded);
     bufs[i] = reinterpret_cast<char*>(alloc_with_status.alloc);
     EXPECT_NE(bufs[i], nullptr);
     EXPECT_TRUE(gpa_.PointerIsMine(bufs[i]));
   }
   EXPECT_EQ(gpa_.successful_allocations(), num_pages);
-  auto alloc_with_status = gpa_.Allocate(1, 0, GetStackTrace());
+  auto alloc_with_status =
+      gpa_.Allocate(1, std::align_val_t{0}, GetStackTrace());
   EXPECT_EQ(alloc_with_status.status,
             Profile::Sample::GuardedStatus::NoAvailableSlots);
   EXPECT_EQ(alloc_with_status.alloc, nullptr);
   gpa_.Deallocate(bufs[0]);
-  alloc_with_status = gpa_.Allocate(1, 0, GetStackTrace());
+  alloc_with_status = gpa_.Allocate(1, std::align_val_t{0}, GetStackTrace());
   EXPECT_EQ(alloc_with_status.status, Profile::Sample::GuardedStatus::Guarded);
   bufs[0] = reinterpret_cast<char*>(alloc_with_status.alloc);
   EXPECT_NE(bufs[0], nullptr);
@@ -210,7 +215,8 @@ INSTANTIATE_TEST_SUITE_P(VaryNumPages, GuardedPageAllocatorParamTest,
                          testing::Values(1, kMaxGpaPages / 2, kMaxGpaPages));
 
 TEST_F(GuardedPageAllocatorTest, PointerIsMine) {
-  auto alloc_with_status = gpa_.Allocate(1, 0, GetStackTrace());
+  auto alloc_with_status =
+      gpa_.Allocate(1, std::align_val_t{0}, GetStackTrace());
   EXPECT_EQ(alloc_with_status.status, Profile::Sample::GuardedStatus::Guarded);
   EXPECT_EQ(gpa_.successful_allocations(), 1);
   void* buf = alloc_with_status.alloc;
@@ -229,7 +235,8 @@ TEST_F(GuardedPageAllocatorTest, Print) {
 }
 
 TEST_F(GuardedPageAllocatorTest, ZeroByteAllocationAndDeallocation) {
-  auto alloc_with_status = gpa_.Allocate(0, 0, GetStackTrace());
+  auto alloc_with_status =
+      gpa_.Allocate(0, std::align_val_t{0}, GetStackTrace());
   EXPECT_EQ(alloc_with_status.status, Profile::Sample::GuardedStatus::Guarded);
   EXPECT_NE(alloc_with_status.alloc, nullptr);
   void* ptr = alloc_with_status.alloc;
@@ -247,7 +254,8 @@ TEST_F(GuardedPageAllocatorTest, ZeroByteAllocationAndDeallocation) {
 }
 
 TEST_F(GuardedPageAllocatorTest, ZeroByteUseAfterFree) {
-  auto alloc_with_status = gpa_.Allocate(0, 0, GetStackTrace());
+  auto alloc_with_status =
+      gpa_.Allocate(0, std::align_val_t{0}, GetStackTrace());
   EXPECT_EQ(alloc_with_status.status, Profile::Sample::GuardedStatus::Guarded);
   ASSERT_NE(alloc_with_status.alloc, nullptr);
   void* ptr = alloc_with_status.alloc;
@@ -260,7 +268,8 @@ TEST_F(GuardedPageAllocatorTest, ZeroByteUseAfterFree) {
 }
 
 TEST_F(GuardedPageAllocatorTest, ZeroByteDoubleFree) {
-  auto alloc_with_status = gpa_.Allocate(0, 0, GetStackTrace());
+  auto alloc_with_status =
+      gpa_.Allocate(0, std::align_val_t{0}, GetStackTrace());
   EXPECT_EQ(alloc_with_status.status, Profile::Sample::GuardedStatus::Guarded);
   ASSERT_NE(alloc_with_status.alloc, nullptr);
   void* ptr = alloc_with_status.alloc;
@@ -280,7 +289,8 @@ TEST_F(GuardedPageAllocatorTest, ThreadedAllocCount) {
     for (size_t i = 0; i < kNumThreads; i++) {
       threads.push_back(std::thread([this, &allocations, i]() {
         for (size_t j = 0; j < kMaxGpaPages; j++) {
-          allocations[i][j] = gpa_.Allocate(1, 0, GetStackTrace()).alloc;
+          allocations[i][j] =
+              gpa_.Allocate(1, std::align_val_t{0}, GetStackTrace()).alloc;
         }
       }));
     }
@@ -311,7 +321,8 @@ TEST_F(GuardedPageAllocatorTest, ThreadedHighContention) {
       threads.push_back(std::thread([this]() {
         char* buf;
         while (true) {
-          auto alloc_with_status = gpa_.Allocate(1, 0, GetStackTrace());
+          auto alloc_with_status =
+              gpa_.Allocate(1, std::align_val_t{0}, GetStackTrace());
           if (alloc_with_status.status ==
               Profile::Sample::GuardedStatus::Guarded) {
             buf = reinterpret_cast<char*>(alloc_with_status.alloc);
@@ -341,7 +352,8 @@ TEST_F(GuardedPageAllocatorTest, ThreadedHighContention) {
   }
   // Verify all pages have been deallocated now that all threads are done.
   for (size_t i = 0; i < kMaxGpaPages; i++) {
-    auto alloc_with_status = gpa_.Allocate(1, 0, GetStackTrace());
+    auto alloc_with_status =
+        gpa_.Allocate(1, std::align_val_t{0}, GetStackTrace());
     EXPECT_EQ(alloc_with_status.status,
               Profile::Sample::GuardedStatus::Guarded);
     EXPECT_NE(alloc_with_status.alloc, nullptr);

tcmalloc/pages.h

@@ -19,6 +19,7 @@
 #include <cstddef>
 #include <cstdint>
 #include <limits>
+#include <new>
 #include <string>
 
 #include "absl/strings/numbers.h"
@@ -193,6 +194,11 @@ inline constexpr Length BytesToLengthCeil(size_t bytes) {
                 ((bytes & (kPageSize - 1)) > 0 ? 1 : 0));
 }
 
+TCMALLOC_ATTRIBUTE_CONST
+inline constexpr Length BytesToLengthCeil(std::align_val_t bytes) {
+  return BytesToLengthCeil(static_cast<size_t>(bytes));
+}
+
 TCMALLOC_ATTRIBUTE_CONST
 inline constexpr Length BytesToLengthFloor(size_t bytes) {
   return Length(bytes >> kPageShift);

tcmalloc/segv_handler_test.cc

@@ -18,6 +18,7 @@
 #include <sys/syscall.h>
 
 #include <cstddef>
+#include <new>
 
 #include "gtest/gtest.h"
 #include "tcmalloc/internal/logging.h"
@@ -40,7 +41,8 @@ TEST(SegvHandlerTest, SignalHandlerStackConsumption) {
   stack_trace.stack[0] = reinterpret_cast<void*>(&&self);
   stack_trace.depth = 1;
 
-  auto ptr = tc_globals.guardedpage_allocator().Allocate(1, 0, stack_trace);
+  auto ptr = tc_globals.guardedpage_allocator().Allocate(1, std::align_val_t{0},
+                                                         stack_trace);
   if (ptr.status != Profile::Sample::GuardedStatus::Guarded) {
     GTEST_SKIP() << "did not get a guarded allocation";
   }

tcmalloc/sizemap.h

@@ -223,7 +223,7 @@ class SizeMap {
   [[nodiscard]] ABSL_ATTRIBUTE_NO_SANITIZE_UNDEFINED
       ABSL_ATTRIBUTE_ALWAYS_INLINE inline SizeMapResult GetSizeClass(
           Policy policy, size_t size) const {
-    const size_t align = policy.align();
+    const size_t align = static_cast<size_t>(policy.align());
     TC_ASSERT(absl::has_single_bit(align));
 
     if (ABSL_PREDICT_FALSE(align > kPageSize)) {

tcmalloc/tcmalloc.cc

@@ -816,7 +816,8 @@ ABSL_ATTRIBUTE_NOINLINE static void free_non_normal(void* ptr, size_t size,
       policy.InSamePartitionAs(ptr).AccessAsCold(), size);
   if (ABSL_PREDICT_FALSE(!is_small)) {
     // We couldn't calculate the size class, which means size > kMaxSize.
-    TC_ASSERT(size > kMaxSize || policy.align() > alignof(std::max_align_t));
+    TC_ASSERT(size > kMaxSize ||
+              policy.align() > std::align_val_t{alignof(std::max_align_t)});
     static_assert(kMaxSize >= kPageSize, "kMaxSize must be at least kPageSize");
     return InvokeHooksAndFreePages(ptr, size, policy);
   }
@@ -876,7 +877,8 @@ inline ABSL_ATTRIBUTE_ALWAYS_INLINE void do_free_with_size(void* ptr,
       tc_globals.sizemap().GetSizeClass(policy.InSamePartitionAs(ptr), size);
   if (ABSL_PREDICT_FALSE(!is_small)) {
     // We couldn't calculate the size class, which means size > kMaxSize.
-    TC_ASSERT(size > kMaxSize || policy.align() > alignof(std::max_align_t));
+    TC_ASSERT(size > kMaxSize ||
+              policy.align() > std::align_val_t{alignof(std::max_align_t)});
     static_assert(kMaxSize >= kPageSize, "kMaxSize must be at least kPageSize");
     SLOW_PATH_BARRIER();
     return InvokeHooksAndFreePages(ptr, size, policy);
@@ -944,7 +946,7 @@ bool CorrectSize(void* ptr, const size_t provided_size, Policy policy) {
   }
 
   if (size_class > 0) {
-    if (policy.align() > static_cast<size_t>(kAlignment)) {
+    if (policy.align() > kAlignment) {
       // Nontrivial alignment.  We might have used a larger size to satisify it.
       minimum_size = 0;
     } else {