From a1c2caeb25b21644700b9423da573b1ccddc35a7 Mon Sep 17 00:00:00 2001
From: Joshua Haberman <jhaberman@gmail.com>
Date: Fri, 8 May 2020 18:50:31 -0700
Subject: [PATCH] More arena tests. (#279)

---
 tests/bindings/lua/test_upb.lua | 23 +++++++++++++++--------
 tests/test_cpp.cc               | 26 ++++++++++++++++++++++++++
 upb/upb.hpp                     |  7 +++++--
 3 files changed, 46 insertions(+), 10 deletions(-)

diff --git a/tests/bindings/lua/test_upb.lua b/tests/bindings/lua/test_upb.lua
index 1180e6b3c3..d1c32b1674 100644
--- a/tests/bindings/lua/test_upb.lua
+++ b/tests/bindings/lua/test_upb.lua
@@ -468,16 +468,23 @@ function test_foo()
 end
 
 function test_gc()
-  local m = test_messages_proto3.TestAllTypesProto3()
-  local n = 1000
+  local top = test_messages_proto3.TestAllTypesProto3()
+  local n = 100
+  local m
 
   for i=1,n do
-    local tmp = m
-    m = test_messages_proto3.TestAllTypesProto3()
-    -- This will cause the arenas to fuse. But we stop referring to the child,
-    -- so the Lua object is eligible for collection (and therefore its original
-    -- arena can be collected too). Only the fusing will keep the C mem alivd.
-    m.recursive_message = tmp
+    local inner = test_messages_proto3.TestAllTypesProto3()
+    m = inner
+    for j=1,n do
+      local tmp = m
+      m = test_messages_proto3.TestAllTypesProto3()
+      -- This will cause the arenas to fuse. But we stop referring to the child,
+      -- so the Lua object is eligible for collection (and therefore its original
+      -- arena can be collected too). Only the fusing will keep the C mem alivd.
+      m.recursive_message = tmp
+
+    end
+    top.recursive_message = m
   end
 
   collectgarbage()
diff --git a/tests/test_cpp.cc b/tests/test_cpp.cc
index abbafdabc8..55cc077c88 100644
--- a/tests/test_cpp.cc
+++ b/tests/test_cpp.cc
@@ -914,6 +914,31 @@ void TestIteration() {
   ASSERT(oneof_count == md.oneof_count());
 }
 
+void TestArena() {
+  int n = 100000;
+
+  struct Decrementer {
+    Decrementer(int* _p) : p(_p) {}
+    ~Decrementer() { (*p)--; }
+    int* p;
+  };
+
+  {
+    upb::Arena arena;
+    for (int i = 0; i < n; i++) {
+      arena.Own(new Decrementer(&n));
+
+      // Intersperse allocation and ensure we can write to it.
+      int* val = static_cast<int*>(upb_arena_malloc(arena.ptr(), sizeof(int)));
+      *val = i;
+    }
+
+    // Test a large allocation.
+    upb_arena_malloc(arena.ptr(), 1000000);
+  }
+  ASSERT(n == 0);
+}
+
 extern "C" {
 
 int run_tests(int argc, char *argv[]) {
@@ -950,6 +975,7 @@ int run_tests(int argc, char *argv[]) {
 
   TestHandlerDataDestruction();
   TestIteration();
+  TestArena();
 
   return 0;
 }
diff --git a/upb/upb.hpp b/upb/upb.hpp
index d8f9d4a6f8..60ef29b296 100644
--- a/upb/upb.hpp
+++ b/upb/upb.hpp
@@ -53,8 +53,11 @@ class Arena {
 
   // Add a cleanup function to run when the arena is destroyed.
   // Returns false on out-of-memory.
-  bool AddCleanup(void *ud, upb_cleanup_func* func) {
-    return upb_arena_addcleanup(ptr_.get(), ud, func);
+  template <class T>
+  bool Own(T *obj) {
+    return upb_arena_addcleanup(ptr_.get(), obj, [](void* obj) {
+      delete static_cast<T*>(obj);
+    });
   }
 
  private: