From 25f6eabb762ae435bb65c6475f63cacfa07cc181 Mon Sep 17 00:00:00 2001
From: Protobuf Team Bot <protobuf-github-bot@google.com>
Date: Tue, 29 Oct 2024 16:12:37 -0700
Subject: [PATCH] Expose function to determine whether arenas are fused

PiperOrigin-RevId: 691196963
---
 upb/mem/arena.c       | 11 ++++++++
 upb/mem/arena.h       |  1 +
 upb/mem/arena_test.cc | 62 +++++++++++++++++++++++++++++++++----------
 3 files changed, 60 insertions(+), 14 deletions(-)

diff --git a/upb/mem/arena.c b/upb/mem/arena.c
index 784794335b..975520b7c7 100644
--- a/upb/mem/arena.c
+++ b/upb/mem/arena.c
@@ -531,6 +531,17 @@ bool upb_Arena_Fuse(const upb_Arena* a1, const upb_Arena* a2) {
   }
 }
 
+bool upb_Arena_IsFused(const upb_Arena* a, const upb_Arena* b) {
+  if (a == b) return true;  // trivial fuse
+  while (true) {
+    upb_ArenaRoot ra = _upb_Arena_FindRoot(a);
+    if (ra.root == _upb_Arena_FindRoot(b).root) return true;
+    if (ra.root == _upb_Arena_FindRoot(a).root) return false;
+
+    // a's root changed since we last checked.  Retry.
+  }
+}
+
 bool upb_Arena_IncRefFor(const upb_Arena* a, const void* owner) {
   upb_ArenaInternal* ai = upb_Arena_Internal(a);
   if (_upb_ArenaInternal_HasInitialBlock(ai)) return false;
diff --git a/upb/mem/arena.h b/upb/mem/arena.h
index 846cb37bbd..63fad6d3ff 100644
--- a/upb/mem/arena.h
+++ b/upb/mem/arena.h
@@ -42,6 +42,7 @@ UPB_API upb_Arena* upb_Arena_Init(void* mem, size_t n, upb_alloc* alloc);
 
 UPB_API void upb_Arena_Free(upb_Arena* a);
 UPB_API bool upb_Arena_Fuse(const upb_Arena* a, const upb_Arena* b);
+UPB_API bool upb_Arena_IsFused(const upb_Arena* a, const upb_Arena* b);
 
 bool upb_Arena_IncRefFor(const upb_Arena* a, const void* owner);
 void upb_Arena_DecRefFor(const upb_Arena* a, const void* owner);
diff --git a/upb/mem/arena_test.cc b/upb/mem/arena_test.cc
index 43570b7a5c..1dce05f28b 100644
--- a/upb/mem/arena_test.cc
+++ b/upb/mem/arena_test.cc
@@ -69,9 +69,9 @@ TEST(ArenaTest, FuseWithInitialBlock) {
 
 class Environment {
  public:
-  void RandomNewFree(absl::BitGen& gen) {
+  void RandomNewFree(absl::BitGen& gen, size_t min_index = 0) {
     auto a = std::make_shared<const upb::Arena>();
-    SwapRandomArena(gen, a);
+    SwapRandomArena(gen, a, min_index);
   }
 
   void RandomIncRefCount(absl::BitGen& gen) {
@@ -86,10 +86,10 @@ class Environment {
     EXPECT_TRUE(upb_Arena_Fuse(a->ptr(), b->ptr()));
   }
 
-  void RandomPoke(absl::BitGen& gen) {
+  void RandomPoke(absl::BitGen& gen, size_t min_index = 0) {
     switch (absl::Uniform(gen, 0, 2)) {
       case 0:
-        RandomNewFree(gen);
+        RandomNewFree(gen, min_index);
         break;
       case 1:
         RandomFuse(gen);
@@ -99,15 +99,23 @@ class Environment {
     }
   }
 
+  std::shared_ptr<const upb::Arena> IndexedNonNullArena(size_t index) {
+    absl::MutexLock lock(&mutex_);
+    std::shared_ptr<const upb::Arena>& ret = arenas_[index];
+    if (!ret) ret = std::make_shared<const upb::Arena>();
+    return ret;
+  }
+
  private:
-  size_t RandomIndex(absl::BitGen& gen) {
-    return absl::Uniform<size_t>(gen, 0, std::tuple_size<ArenaArray>::value);
+  size_t RandomIndex(absl::BitGen& gen, size_t min_index = 0) {
+    return absl::Uniform<size_t>(gen, min_index,
+                                 std::tuple_size<ArenaArray>::value);
   }
 
   // Swaps a random arena from the set with the given arena.
-  void SwapRandomArena(absl::BitGen& gen,
-                       std::shared_ptr<const upb::Arena>& a) {
-    size_t i = RandomIndex(gen);
+  void SwapRandomArena(absl::BitGen& gen, std::shared_ptr<const upb::Arena>& a,
+                       size_t min_index) {
+    size_t i = RandomIndex(gen, min_index);
     absl::MutexLock lock(&mutex_);
     arenas_[i].swap(a);
   }
@@ -118,11 +126,7 @@ class Environment {
   // Note that the returned arena is shared and may be accessed concurrently
   // by other threads.
   std::shared_ptr<const upb::Arena> RandomNonNullArena(absl::BitGen& gen) {
-    size_t i = RandomIndex(gen);
-    absl::MutexLock lock(&mutex_);
-    std::shared_ptr<const upb::Arena>& ret = arenas_[i];
-    if (!ret) ret = std::make_shared<const upb::Arena>();
-    return ret;
+    return IndexedNonNullArena(RandomIndex(gen));
   }
 
   using ArenaArray = std::array<std::shared_ptr<const upb::Arena>, 100>;
@@ -292,6 +296,36 @@ TEST(ArenaTest, IncRefCountShouldFailForInitialBlock) {
   EXPECT_FALSE(upb_Arena_IncRefFor(arena, nullptr));
 }
 
+TEST(ArenaTest, FuzzFuseIsFusedRace) {
+  Environment env;
+
+  // Create two arenas and fuse them.
+  std::shared_ptr<const upb::Arena> a = env.IndexedNonNullArena(0);
+  std::shared_ptr<const upb::Arena> b = env.IndexedNonNullArena(1);
+  upb_Arena_Fuse(a->ptr(), b->ptr());
+  EXPECT_TRUE(upb_Arena_IsFused(a->ptr(), b->ptr()));
+
+  absl::Notification done;
+  std::vector<std::thread> threads;
+  for (int i = 0; i < 10; ++i) {
+    threads.emplace_back([&]() {
+      absl::BitGen gen;
+      while (!done.HasBeenNotified()) {
+        env.RandomPoke(gen, 2);
+      }
+    });
+  }
+
+  absl::BitGen gen;
+  auto end = absl::Now() + absl::Seconds(2);
+  while (absl::Now() < end) {
+    // Verify that the two arenas are still fused.
+    EXPECT_TRUE(upb_Arena_IsFused(a->ptr(), b->ptr()));
+  }
+  done.Notify();
+  for (auto& t : threads) t.join();
+}
+
 #endif
 
 }  // namespace