Add threading tests for arenas

PiperOrigin-RevId: 518009578

BUILD

@@ -825,6 +825,18 @@ cc_library(
     ],
 )
 
+cc_test(
+    name = "arena_test",
+    srcs = ["upb/mem/arena_test.cc"],
+    deps = [
+        ":upb",
+        "@com_google_absl//absl/random",
+        "@com_google_absl//absl/random:distributions",
+        "@com_google_absl//absl/synchronization",
+        "@com_google_googletest//:gtest_main",
+    ],
+)
+
 cc_library(
     name = "mem_internal",
     srcs = [

upb/mem/arena_test.cc

@@ -0,0 +1,187 @@
+#include "upb/mem/arena.h"
+
+#include <array>
+#include <atomic>
+#include <thread>
+#include <vector>
+
+#include "gmock/gmock.h"
+#include "gtest/gtest.h"
+#include "absl/random/distributions.h"
+#include "absl/random/random.h"
+#include "absl/synchronization/notification.h"
+#include "upb/upb.hpp"
+
+namespace {
+
+static void decrement_int(void* ptr) {
+  int* iptr = static_cast<int*>(ptr);
+  (*iptr)--;
+}
+
+TEST(ArenaTest, ArenaFuse) {
+  int i1 = 5;
+  int i2 = 5;
+  int i3 = 5;
+  int i4 = 5;
+
+  upb_Arena* arena1 = upb_Arena_New();
+  upb_Arena* arena2 = upb_Arena_New();
+
+  upb_Arena_AddCleanup(arena1, &i1, decrement_int);
+  upb_Arena_AddCleanup(arena2, &i2, decrement_int);
+
+  EXPECT_TRUE(upb_Arena_Fuse(arena1, arena2));
+
+  upb_Arena_AddCleanup(arena1, &i3, decrement_int);
+  upb_Arena_AddCleanup(arena2, &i4, decrement_int);
+
+  upb_Arena_Free(arena1);
+  EXPECT_EQ(5, i1);
+  EXPECT_EQ(5, i2);
+  EXPECT_EQ(5, i3);
+  EXPECT_EQ(5, i4);
+  upb_Arena_Free(arena2);
+  EXPECT_EQ(4, i1);
+  EXPECT_EQ(4, i2);
+  EXPECT_EQ(4, i3);
+  EXPECT_EQ(4, i4);
+}
+
+/* Do nothing allocator for testing */
+extern "C" void* TestAllocFunc(upb_alloc* alloc, void* ptr, size_t oldsize,
+                               size_t size) {
+  return upb_alloc_global.func(alloc, ptr, oldsize, size);
+}
+ABSL_CONST_INIT upb_alloc test_alloc = {&TestAllocFunc};
+
+TEST(ArenaTest, FuseWithInitialBlock) {
+  char buf1[1024];
+  char buf2[1024];
+  upb_Arena* arenas[] = {upb_Arena_Init(buf1, 1024, &upb_alloc_global),
+                         upb_Arena_Init(buf2, 1024, &upb_alloc_global),
+                         upb_Arena_Init(NULL, 0, &test_alloc),
+                         upb_Arena_Init(NULL, 0, &upb_alloc_global)};
+  int size = sizeof(arenas) / sizeof(arenas[0]);
+  for (int i = 0; i < size; ++i) {
+    for (int j = 0; j < size; ++j) {
+      if (i == j) {
+        EXPECT_TRUE(upb_Arena_Fuse(arenas[i], arenas[j]));
+      } else {
+        EXPECT_FALSE(upb_Arena_Fuse(arenas[i], arenas[j]));
+      }
+    }
+  }
+
+  for (int i = 0; i < size; ++i) upb_Arena_Free(arenas[i]);
+}
+
+class Environment {
+ public:
+  ~Environment() {
+    for (auto& atom : arenas_) {
+      auto* a = atom.load(std::memory_order_relaxed);
+      if (a != nullptr) upb_Arena_Free(a);
+    }
+  }
+
+  void RandomNewFree(absl::BitGen& gen) {
+    auto* old = SwapRandomly(gen, upb_Arena_New());
+    if (old != nullptr) upb_Arena_Free(old);
+  }
+
+  void RandomFuse(absl::BitGen& gen) {
+    std::array<upb_Arena*, 2> old;
+    for (auto& o : old) {
+      o = SwapRandomly(gen, nullptr);
+      if (o == nullptr) o = upb_Arena_New();
+    }
+
+    upb_Arena_Fuse(old[0], old[1]);
+    for (auto& o : old) {
+      o = SwapRandomly(gen, o);
+      if (o != nullptr) upb_Arena_Free(o);
+    }
+  }
+
+  void RandomPoke(absl::BitGen& gen) {
+    switch (absl::Uniform(gen, 0, 2)) {
+      case 0:
+        RandomNewFree(gen);
+        break;
+      case 1:
+        RandomFuse(gen);
+        break;
+      default:
+        break;
+    }
+  }
+
+ private:
+  upb_Arena* SwapRandomly(absl::BitGen& gen, upb_Arena* a) {
+    return arenas_[absl::Uniform<size_t>(gen, 0, arenas_.size())].exchange(
+        a, std::memory_order_acq_rel);
+  }
+
+  std::array<std::atomic<upb_Arena*>, 100> arenas_ = {};
+};
+
+TEST(ArenaTest, FuzzSingleThreaded) {
+  Environment env;
+
+  absl::BitGen gen;
+  auto end = absl::Now() + absl::Seconds(0.5);
+  while (absl::Now() < end) {
+    env.RandomPoke(gen);
+  }
+}
+
+// Disabled because this test currently fails.
+TEST(ArenaTest, DISABLED_FuzzFuseFreeRace) {
+  Environment env;
+
+  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.RandomNewFree(gen);
+      }
+    });
+  }
+
+  absl::BitGen gen;
+  auto end = absl::Now() + absl::Seconds(2);
+  while (absl::Now() < end) {
+    env.RandomFuse(gen);
+  }
+  done.Notify();
+  for (auto& t : threads) t.join();
+}
+
+// Disabled because this operation is currently unsupported.
+TEST(ArenaTest, DISABLED_FuzzFuseFuseRace) {
+  Environment env;
+
+  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);
+      }
+    });
+  }
+
+  absl::BitGen gen;
+  auto end = absl::Now() + absl::Seconds(2);
+  while (absl::Now() < end) {
+    env.RandomPoke(gen);
+  }
+  done.Notify();
+  for (auto& t : threads) t.join();
+}
+
+}  // namespace