Synchronize PageMap radix-tree expansion with atomic operations.

Store and load radix-tree interior pointers and hugepage pointers in `PageMap3` using atomic operations with release/acquire semantics.

Radix-tree expansion in `Ensure` allocates and initializes new interior and leaf nodes under the page heap lock before publishing them with release stores. Fast-path allocation and deallocation traverse the radix tree locklessly; atomic loads prevent observing partially initialized pointer slots or stale nulls. Relaxed loads are used in `MaybeIndex` to avoid load-acquire overhead on architectures where hardware address dependencies preserve ordering.

`PageMeta` descriptor words remain non-atomic to preserve the 32-cycle deletion fast path, with ThreadSanitizer clean conversion tracked under b/193887621.

Co-authored-by: Chris Kennelly <ckennelly@google.com>
PiperOrigin-RevId: 1007121156

tcmalloc/pagemap.h

@@ -31,6 +31,8 @@
 #include <string.h>
 
 #include <algorithm>
+#include <atomic>
+#include <new>
 #include <optional>
 #include <tuple>
 #include <utility>
@@ -186,17 +188,17 @@ class PageMap {
     // avoid two separate memory loads when fetching both the descriptor and the
     // sizeclass.
     PageMeta page[kLeafLength];
-    void* hugepage[kLeafHugepages];
+    std::atomic<void*> hugepage[kLeafHugepages];
   };
 
   struct Node {
     // Mid-level structure that holds pointers to leafs
-    Leaf* absl_nullable leafs[kMidLength];
+    std::atomic<Leaf* absl_nullable> leafs[kMidLength];
   };
 
   typedef uintptr_t Number;
 
-  Node* absl_nullable root_[kRootLength];  // Top-level node
+  std::atomic<Node* absl_nullable> root_[kRootLength];  // Top-level node
 
   [[nodiscard]] ABSL_ATTRIBUTE_ALWAYS_INLINE std::tuple<Number, Number, Number>
   Index(PageId p) const {
@@ -212,12 +214,18 @@ class PageMap {
       MustIndex(PageId p) const {
     TC_ASSERT_EQ(p.index() >> BITS, 0);
     auto [i1, i2, i3] = Index(p);
-    TC_ASSERT_NE(root_[i1], nullptr);
-    Leaf* leaf = root_[i1]->leafs[i2];
+    Node* node = root_[i1].load(std::memory_order_acquire);
+    TC_ASSERT_NE(node, nullptr);
+    Leaf* leaf = node->leafs[i2].load(std::memory_order_acquire);
     TC_ASSERT_NE(leaf, nullptr);
     return {leaf, i3};
   }
 
+  // Relaxed loads are used here to avoid emitting load-acquire instructions
+  // (such as LDAR on ARM64) on the deallocation fast path. On weakly ordered
+  // architectures, the hardware address dependency from the pointer load to
+  // its dereference ensures that newly published Node and Leaf structures
+  // (published via release stores in Ensure) are observed initialized.
   [[nodiscard]] ABSL_ATTRIBUTE_ALWAYS_INLINE
       std::pair<Leaf* absl_nullable, Number>
       MaybeIndex(PageId p) const {
@@ -226,11 +234,11 @@ class PageMap {
       return {nullptr, 0};
     }
     auto [i1, i2, i3] = Index(p);
-    const Node* node = root_[i1];
+    const Node* node = root_[i1].load(std::memory_order_relaxed);
     if (ABSL_PREDICT_FALSE(node == nullptr)) {
       return {nullptr, 0};
     }
-    return {node->leafs[i2], i3};
+    return {node->leafs[i2].load(std::memory_order_relaxed), i3};
   }
 
  public:
@@ -285,36 +293,46 @@ class PageMap {
   void SetSmall(PageId p, Span* absl_nonnull span, CompactSizeClass sc,
                 bool donated) {
     auto [leaf, i3] = MustIndex(p);
-    leaf->page[i3].set_small(span, sc, donated);
+    PageMeta meta;
+    meta.set_small(span, sc, donated);
+    leaf->page[i3] = meta;
   }
 
   void SetLarge(PageId p, Length size, bool donated) {
     auto [leaf, i3] = MustIndex(p);
     TC_ASSERT_EQ(leaf->page[i3].sizeclass(), 0);
-    leaf->page[i3].set_large(size, donated);
+    PageMeta meta;
+    meta.set_large(size, donated);
+    leaf->page[i3] = meta;
   }
 
   void SetSampled(PageId p, SampledAllocation* absl_nonnull sampled,
                   bool donated) {
     auto [leaf, i3] = MustIndex(p);
     TC_ASSERT_EQ(leaf->page[i3].sizeclass(), 0);
-    leaf->page[i3].set_sampled(sampled, donated);
+    PageMeta meta;
+    meta.set_sampled(sampled, donated);
+    leaf->page[i3] = meta;
   }
 
   void SetFreed(PageId p) {
     auto [leaf, i3] = MustIndex(p);
     TC_ASSERT_EQ(leaf->page[i3].sizeclass(), 0);
-    leaf->page[i3].set_freed();
+    PageMeta meta;
+    meta.set_freed();
+    leaf->page[i3] = meta;
   }
 
   [[nodiscard]] void* GetHugepage(PageId p) const {
     auto [leaf, i3] = MustIndex(p);
-    return leaf->hugepage[i3 >> (kLeafBits - kLeafHugeBits)];
+    return leaf->hugepage[i3 >> (kLeafBits - kLeafHugeBits)].load(
+        std::memory_order_acquire);
   }
 
   void SetHugepage(PageId p, void* v) {
     auto [leaf, i3] = MustIndex(p);
-    leaf->hugepage[i3 >> (kLeafBits - kLeafHugeBits)] = v;
+    leaf->hugepage[i3 >> (kLeafBits - kLeafHugeBits)].store(
+        v, std::memory_order_release);
   }
 
   [[nodiscard]] bool HasLeaf(PageId p) const {
@@ -326,11 +344,13 @@ class PageMap {
   [[nodiscard]] std::optional<PageId> get_next_set_page(PageId p) const {
     auto [i1, i2, i3] = Index(p + Length(1));
     for (; i1 < kRootLength; ++i1, i2 = 0, i3 = 0) {
-      if (root_[i1] == nullptr) continue;
+      const Node* node = root_[i1].load(std::memory_order_acquire);
+      if (node == nullptr) continue;
       for (; i2 < kMidLength; ++i2, i3 = 0) {
-        if (root_[i1]->leafs[i2] == nullptr) continue;
+        const Leaf* leaf = node->leafs[i2].load(std::memory_order_acquire);
+        if (leaf == nullptr) continue;
         for (; i3 < kLeafLength; ++i3) {
-          if (root_[i1]->leafs[i2]->page[i3].valid())
+          if (leaf->page[i3].valid())
             return PageId((i1 << (kLeafBits + kMidBits)) | (i2 << kLeafBits) |
                           i3);
         }
@@ -353,19 +373,20 @@ class PageMap {
       if (i1 >= kRootLength) return false;
 
       // Allocate Node if necessary
-      if (root_[i1] == nullptr) {
-        Node* node = reinterpret_cast<Node*>(Allocator(sizeof(Node)));
-        if (node == nullptr) return false;
-        memset(node, 0, sizeof(*node));
-        root_[i1] = node;
+      Node* node = root_[i1].load(std::memory_order_acquire);
+      if (node == nullptr) {
+        void* mem = Allocator(sizeof(Node));
+        if (mem == nullptr) return false;
+        node = new (mem) Node();
+        root_[i1].store(node, std::memory_order_release);
       }
 
       // Allocate Leaf if necessary
-      if (root_[i1]->leafs[i2] == nullptr) {
-        Leaf* leaf = reinterpret_cast<Leaf*>(Allocator(sizeof(Leaf)));
-        if (leaf == nullptr) return false;
-        memset(leaf, 0, sizeof(*leaf));
-        root_[i1]->leafs[i2] = leaf;
+      if (node->leafs[i2].load(std::memory_order_acquire) == nullptr) {
+        void* mem = Allocator(sizeof(Leaf));
+        if (mem == nullptr) return false;
+        Leaf* leaf = new (mem) Leaf();
+        node->leafs[i2].store(leaf, std::memory_order_release);
       }
 
       // Advance p past whatever is covered by this leaf node
@@ -405,9 +426,11 @@ class PageMap {
     const PageMeta first_meta = GetExistingDescriptor(first);
     TC_ASSERT_EQ(first_meta.span(), span);
     const bool donated = first_meta.donated();
+    PageMeta freed_meta;
+    freed_meta.set_freed();
     for (PageId p = first; p <= last; ++p) {
       auto [leaf, i3] = MustIndex(p);
-      leaf->page[i3].set_freed();
+      leaf->page[i3] = freed_meta;
     }
     return donated;
   }

tcmalloc/pagemap_test.cc

@@ -21,9 +21,12 @@
 #include <unistd.h>
 
 #include <algorithm>
+#include <atomic>
 #include <limits>
+#include <memory>
 #include <new>
 #include <string>
+#include <thread>  // NOLINT(build/c++11)
 #include <utility>
 #include <vector>
 
@@ -50,8 +53,8 @@ class PageMapTest : public ::testing::TestWithParam<int> {
  public:
   PageMapTest() {
     // Arrange to pass zero-filled memory as the backing store for map.
-    memset(storage, 0, sizeof(Map));
-    map = new (storage) Map();
+    memset(storage_, 0, sizeof(Map));
+    map_ = new (storage_) Map();
   }
 
   ~PageMapTest() override {
@@ -76,66 +79,66 @@ class PageMapTest : public ::testing::TestWithParam<int> {
  public:
   static constexpr int kTestBits = kAddressBits - kPageShift;
   using Map = PageMap<kTestBits, alloc>;
-  Map* map;
+  Map* map_;
 
  private:
-  alignas(Map) char storage[sizeof(Map)];
+  alignas(Map) char storage_[sizeof(Map)];
 };
 
 TEST_P(PageMapTest, Sequential) {
   const intptr_t limit = GetParam();
 
   for (intptr_t i = 0; i < limit; i++) {
-    ASSERT_TRUE(map->Ensure(Range(PageId(i), Length(1))));
-    map->SetLarge(PageId(i), Length(i + 1), (i & 1) != 0);
-    ASSERT_TRUE(map->GetDescriptor(PageId(i)).valid());
-    ASSERT_FALSE(map->GetDescriptor(PageId(i)).sampled());
-    ASSERT_FALSE(map->GetDescriptor(PageId(i)).freed());
-    ASSERT_EQ(map->GetDescriptor(PageId(i)).size(), Length(i + 1));
-    ASSERT_EQ(map->GetDescriptor(PageId(i)).donated(), (i & 1) != 0);
-    ASSERT_EQ(map->GetExistingDescriptor(PageId(i)).size(), Length(i + 1));
+    ASSERT_TRUE(map_->Ensure(Range(PageId(i), Length(1))));
+    map_->SetLarge(PageId(i), Length(i + 1), (i & 1) != 0);
+    ASSERT_TRUE(map_->GetDescriptor(PageId(i)).valid());
+    ASSERT_FALSE(map_->GetDescriptor(PageId(i)).sampled());
+    ASSERT_FALSE(map_->GetDescriptor(PageId(i)).freed());
+    ASSERT_EQ(map_->GetDescriptor(PageId(i)).size(), Length(i + 1));
+    ASSERT_EQ(map_->GetDescriptor(PageId(i)).donated(), (i & 1) != 0);
+    ASSERT_EQ(map_->GetExistingDescriptor(PageId(i)).size(), Length(i + 1));
 
     // Test size class handling
-    ASSERT_EQ(0, map->sizeclass(PageId(i)));
-    map->SetSmall(PageId(i), span(i), sc(i), false);
-    ASSERT_EQ(sc(i), map->sizeclass(PageId(i)));
+    ASSERT_EQ(0, map_->sizeclass(PageId(i)));
+    map_->SetSmall(PageId(i), span(i), sc(i), false);
+    ASSERT_EQ(sc(i), map_->sizeclass(PageId(i)));
   }
   for (intptr_t i = 0; i < limit; i++) {
-    ASSERT_EQ(map->GetDescriptor(PageId(i)).span(), span(i));
-    ASSERT_EQ(map->GetExistingDescriptor(PageId(i)).span(), span(i));
-    ASSERT_EQ(map->sizeclass(PageId(i)), sc(i));
-    ASSERT_EQ(map->GetDescriptor(PageId(i)).sizeclass(), sc(i));
+    ASSERT_EQ(map_->GetDescriptor(PageId(i)).span(), span(i));
+    ASSERT_EQ(map_->GetExistingDescriptor(PageId(i)).span(), span(i));
+    ASSERT_EQ(map_->sizeclass(PageId(i)), sc(i));
+    ASSERT_EQ(map_->GetDescriptor(PageId(i)).sizeclass(), sc(i));
   }
 }
 
 TEST_P(PageMapTest, Bulk) {
   const intptr_t limit = GetParam();
 
-  ASSERT_TRUE(map->Ensure(Range(PageId(0), Length(limit))));
+  ASSERT_TRUE(map_->Ensure(Range(PageId(0), Length(limit))));
   for (intptr_t i = 0; i < limit; i++) {
-    map->SetSmall(PageId(i), span(i), sc(i), false);
-    ASSERT_EQ(map->GetDescriptor(PageId(i)).span(), span(i));
+    map_->SetSmall(PageId(i), span(i), sc(i), false);
+    ASSERT_EQ(map_->GetDescriptor(PageId(i)).span(), span(i));
   }
   for (intptr_t i = 0; i < limit; i++) {
-    ASSERT_EQ(map->GetDescriptor(PageId(i)).span(), span(i));
+    ASSERT_EQ(map_->GetDescriptor(PageId(i)).span(), span(i));
   }
 }
 
 TEST_P(PageMapTest, Overflow) {
   const uintptr_t kLimit = uintptr_t{1} << kTestBits;
-  ASSERT_FALSE(map->Ensure(Range(PageId(kLimit), Length(kLimit + 1))));
+  ASSERT_FALSE(map_->Ensure(Range(PageId(kLimit), Length(kLimit + 1))));
 }
 
 TEST_P(PageMapTest, EnsureOverflowAndPartialBounds) {
   const uintptr_t kLimit = uintptr_t{1} << kTestBits;
   const Length kMaxLen(std::numeric_limits<uintptr_t>::max());
-  EXPECT_FALSE(map->Ensure(Range(PageId(2), kMaxLen)));
-  EXPECT_FALSE(map->HasLeaf(PageId(2)));
-  EXPECT_FALSE(map->Ensure(Range(PageId(kLimit), kMaxLen)));
-  EXPECT_FALSE(map->Ensure(Range(PageId(kLimit - 1), Length(2))));
-  EXPECT_FALSE(map->HasLeaf(PageId(kLimit - 1)));
-  EXPECT_FALSE(map->Ensure(Range(PageId(1), kMaxLen)));
-  EXPECT_FALSE(map->HasLeaf(PageId(1)));
+  EXPECT_FALSE(map_->Ensure(Range(PageId(2), kMaxLen)));
+  EXPECT_FALSE(map_->HasLeaf(PageId(2)));
+  EXPECT_FALSE(map_->Ensure(Range(PageId(kLimit), kMaxLen)));
+  EXPECT_FALSE(map_->Ensure(Range(PageId(kLimit - 1), Length(2))));
+  EXPECT_FALSE(map_->HasLeaf(PageId(kLimit - 1)));
+  EXPECT_FALSE(map_->Ensure(Range(PageId(1), kMaxLen)));
+  EXPECT_FALSE(map_->HasLeaf(PageId(1)));
 }
 
 TEST_P(PageMapTest, RandomAccess) {
@@ -148,17 +151,66 @@ TEST_P(PageMapTest, RandomAccess) {
   std::shuffle(elements.begin(), elements.end(), absl::BitGen());
 
   for (intptr_t i = 0; i < limit; i++) {
-    ASSERT_TRUE(map->Ensure(Range(PageId(elements[i]), Length(1))));
-    map->SetSmall(PageId(elements[i]), span(elements[i]), sc(elements[i]),
-                  false);
-    ASSERT_EQ(map->GetDescriptor(PageId(elements[i])).span(),
+    ASSERT_TRUE(map_->Ensure(Range(PageId(elements[i]), Length(1))));
+    map_->SetSmall(PageId(elements[i]), span(elements[i]), sc(elements[i]),
+                   false);
+    ASSERT_EQ(map_->GetDescriptor(PageId(elements[i])).span(),
               span(elements[i]));
   }
   for (intptr_t i = 0; i < limit; i++) {
-    ASSERT_EQ(map->GetDescriptor(PageId(i)).span(), span(i));
+    ASSERT_EQ(map_->GetDescriptor(PageId(i)).span(), span(i));
   }
 }
 
+TEST_P(PageMapTest, ConcurrentUpdateAndReadDescriptor) {
+#if defined(ABSL_HAVE_THREAD_SANITIZER)
+  GTEST_SKIP() << "PageMap does not run cleanly under TSan (b/193887621).";
+#else
+  const intptr_t limit = std::min<intptr_t>(GetParam(), 256);
+  std::vector<std::unique_ptr<Span> > spans(limit);
+  for (intptr_t i = 0; i < limit; ++i) {
+    spans[i] = std::make_unique<Span>(PageId(i + 1));
+  }
+
+  std::atomic<bool> stop{false};
+  auto reader_fn = [&]() {
+    while (!stop.load(std::memory_order_relaxed)) {
+      for (intptr_t i = 0; i < limit; ++i) {
+        const PageId p(i + 1);
+        const CompactSizeClass cls = map_->sizeclass(p);
+        EXPECT_TRUE(cls == 0 || cls == sc(i));
+        const PageMeta meta = map_->GetDescriptor(p);
+        if (meta.valid()) {
+          if (meta.sizeclass() != 0) {
+            EXPECT_EQ(meta.sizeclass(), sc(i));
+            EXPECT_EQ(meta.span(), spans[i].get());
+          } else {
+            EXPECT_EQ(meta.size(), Length(i + 1));
+          }
+        }
+      }
+      (void)map_->get_next_set_page(PageId(0));
+    }
+  };
+  std::thread reader(reader_fn);
+
+  for (int iter = 0; iter < 50; ++iter) {
+    for (intptr_t i = 0; i < limit; ++i) {
+      const PageId p(i + 1);
+      ASSERT_TRUE(map_->Ensure(Range(p, Length(1))));
+      map_->RegisterSizeClass(spans[i].get(), Length(1), sc(i),
+                              /*donated=*/false);
+      EXPECT_FALSE(map_->UnregisterSizeClass(spans[i].get(), Length(1)));
+      map_->SetLarge(p, Length(i + 1), /*donated=*/(i & 1) != 0);
+      map_->SetFreed(p);
+    }
+  }
+
+  stop.store(true, std::memory_order_relaxed);
+  reader.join();
+#endif
+}
+
 INSTANTIATE_TEST_SUITE_P(Limits, PageMapTest, ::testing::Values(100, 1 << 16));
 
 // Surround pagemap with unused memory. This isolates it so that it does not