Implemented upb_arena_fuse() (#278)

* WIP.

* WIP.

* Tests are passing.

* Recover some perf: LIKELY doesn't propagate through functions. :(

* Added some more benchmarks.

* Simplify & optimize upb_arena_realloc().

* Only add owned blocks to the freelist.

* More optimization/simplification.

* Re-fixed the bug.

* Revert unintentional changes to parser.rl.

* Revert Lua changes for now.

* Revert the arena fuse changes for now.

* Added last_size to the arena representation.

* Re-applied Lua changes.

* Implemented upb_arena_fuse().

* Fix the compile by re-ordering statements.

* Improve comments.

tests/bindings/lua/test_upb.lua

@@ -467,6 +467,29 @@ function test_foo()
   assert_equal(set.file[1].name, "google/protobuf/descriptor.proto")
 end
 
+function test_gc()
+  local m = test_messages_proto3.TestAllTypesProto3()
+  local n = 1000
+
+  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
+  end
+
+  collectgarbage()
+
+  for i=1,n do
+    -- Verify we can touch all the messages again and without accessing freed
+    -- memory.
+    m = m.recursive_message
+    assert_not_nil(m)
+  end
+end
+
 local stats = lunit.main()
 
 if stats.failed > 0 or stats.errors > 0 then

upb/bindings/lua/msg.c

@@ -42,34 +42,29 @@
  * the wrappers can be collected if they are no longer needed.  A new wrapper
  * object can always be recreated later.
  *
- *                    arena
- *                 +->group
- *                 |
- *                 V        +-----+
+ *                          +-----+
  *            lupb_arena    |cache|-weak-+
  *                 |  ^     +-----+      |
  *                 |  |                  V
  * Lua level       |  +------------lupb_msg
- * ----------------|-----------------|-------------------------------------------
+ * ----------------|-----------------|------------------------------------------
  * upb level       |                 |
  *                 |            +----V------------------------------+
  *                 +->upb_arena | upb_msg  ...(empty arena storage) |
  *                              +-----------------------------------+
  *
  * If the user creates a reference between two objects that have different
- * arenas, we need to merge the arenas into a single, bigger arena group.  The
- * arena group will reference both arenas, and will inherit the longest lifetime
- * of anything in the arena.
+ * arenas, we need to fuse the two arenas together, so that the blocks will
+ * outlive both arenas.
  *
- *                                              arena
- *                 +--------------------------->group<-----------------+
+ *                 +-------------------------->(fused)<----------------+
  *                 |                                                   |
  *                 V                           +-----+                 V
  *            lupb_arena                +-weak-|cache|-weak-+     lupb_arena
  *                 |  ^                 |      +-----+      |        ^   |
  *                 |  |                 V                   V        |   |
  * Lua level       |  +------------lupb_msg              lupb_msg----+   |
- * ----------------|-----------------|-------------------------|---------|-------
+ * ----------------|-----------------|-------------------------|---------|------
  * upb level       |                 |                         |         |
  *                 |            +----V----+               +----V----+    V
  *                 +->upb_arena | upb_msg |               | upb_msg | upb_arena
@@ -77,7 +72,7 @@
  *                                     +---------------------+
  * Key invariants:
  *   1. every wrapper references the arena that contains it.
- *   2. every arena group references all arenas that own upb objects reachable
+ *   2. every fused arena includes all arenas that own upb objects reachable
  *      from that arena.  In other words, when a wrapper references an arena,
  *      this is sufficient to ensure that any upb object reachable from that
  *      wrapper will stay alive.
@@ -85,12 +80,6 @@
  * Additionally, every message object contains a strong reference to the
  * corresponding Descriptor object.  Likewise, array/map objects reference a
  * Descriptor object if they are typed to store message values.
- *
- * (The object cache could be per-arena-group.  This would keep individual cache
- * tables smaller, and when an arena group is freed the entire cache table(s) could
- * be collected in one fell swoop.  However this makes merging another arena
- * into the group an O(n) operation, since all entries would need to be copied
- * from the existing cache table.)
  */
 
 #define LUPB_ARENA "lupb.arena"
@@ -170,14 +159,8 @@ static void lupb_cacheset(lua_State *L, const void *key) {
 /* lupb_arena only exists to wrap a upb_arena.  It is never exposed to users; it
  * is an internal memory management detail.  Other wrapper objects refer to this
  * object from their userdata to keep the arena-owned data alive.
- *
- * The arena userval is a table representing the arena group.  Every arena in
- * the group points to the same table, and the table references all arenas in
- * the group.
  */
 
-#define LUPB_ARENAGROUP_INDEX 1
-
 typedef struct {
   upb_arena *arena;
 } lupb_arena;
@@ -190,63 +173,18 @@ static upb_arena *lupb_arena_check(lua_State *L, int narg) {
 upb_arena *lupb_arena_pushnew(lua_State *L) {
   lupb_arena *a = lupb_newuserdata(L, sizeof(lupb_arena), 1, LUPB_ARENA);
   a->arena = upb_arena_new();
-
-  /* Create arena group table and add this arena to it. */
-  lua_createtable(L, 0, 1);
-  lua_pushvalue(L, -2);
-  lua_rawseti(L, -2, 1);
-
-  /* Set arena group as this object's userval. */
-  lua_setiuservalue(L, -2, LUPB_ARENAGROUP_INDEX);
-
   return a->arena;
 }
 
 /**
- * lupb_arena_merge()
+ * lupb_arena_fuse()
  *
  * Merges |from| into |to| so that there is a single arena group that contains
  * both, and both arenas will point at this new table. */
-static void lupb_arena_merge(lua_State *L, int to, int from) {
-  int i, from_count, to_count;
-  lua_getiuservalue(L, to, LUPB_ARENAGROUP_INDEX);
-  lua_getiuservalue(L, from, LUPB_ARENAGROUP_INDEX);
-
-  if (lua_rawequal(L, -1, -2)) {
-    /* These arenas are already in the same group. */
-    lua_pop(L, 2);
-    return;
-  }
-
-  to_count = lua_rawlen(L, -2);
-  from_count = lua_rawlen(L, -1);
-
-  /* Add everything in |from|'s arena group. */
-  for (i = 1; i <= from_count; i++) {
-    lua_rawgeti(L, -1, i);
-    lua_rawseti(L, -3, i + to_count);
-  }
-
-  /* Make |from| point to |to|'s table. */
-  lua_pop(L, 1);
-  lua_setiuservalue(L, from, LUPB_ARENAGROUP_INDEX);
-}
-
-/**
- * lupb_arena_addobj()
- *
- * Creates a reference from the arena in |narg| to the object at the top of the
- * stack, and pops it.  This will guarantee that the object lives as long as
- * the arena.
- *
- * This is mainly useful for pinning strings we have parsed protobuf data from.
- * It will allow us to point directly to string data in the original string. */
-static void lupb_arena_addobj(lua_State *L, int narg) {
-  lua_getiuservalue(L, narg, LUPB_ARENAGROUP_INDEX);
-  int n = lua_rawlen(L, -1);
-  lua_pushvalue(L, -2);
-  lua_rawseti(L, -2, n + 1);
-  lua_pop(L, 2);  /* obj, arena group. */
+static void lupb_arena_fuse(lua_State *L, int to, int from) {
+  upb_arena *to_arena = lupb_arena_check(L, to);
+  upb_arena *from_arena = lupb_arena_check(L, from);
+  upb_arena_fuse(to_arena, from_arena);
 }
 
 static int lupb_arena_gc(lua_State *L) {
@@ -877,7 +815,7 @@ static int lupb_msg_newindex(lua_State *L) {
   if (merge_arenas) {
     lua_getiuservalue(L, 1, LUPB_ARENA_INDEX);
     lua_getiuservalue(L, 3, LUPB_ARENA_INDEX);
-    lupb_arena_merge(L, lua_absindex(L, -2), lua_absindex(L, -1));
+    lupb_arena_fuse(L, lua_absindex(L, -2), lua_absindex(L, -1));
     lua_pop(L, 2);
   }
 
@@ -940,6 +878,7 @@ static int lupb_decode(lua_State *L) {
   const upb_msgdef *m = lupb_msgdef_check(L, 1);
   const char *pb = lua_tolstring(L, 2, &len);
   const upb_msglayout *layout = upb_msgdef_layout(m);
+  char *buf;
   upb_msg *msg;
   upb_arena *arena;
   bool ok;
@@ -952,13 +891,13 @@ static int lupb_decode(lua_State *L) {
 
   lua_getiuservalue(L, -1, LUPB_ARENA_INDEX);
   arena = lupb_arena_check(L, -1);
-
-  /* Pin string data so we can reference it. */
-  lua_pushvalue(L, 2);
-  lupb_arena_addobj(L, -2);
   lua_pop(L, 1);
 
-  ok = upb_decode(pb, len, msg, layout, arena);
+  /* Copy input data to arena, message will reference it. */
+  buf = upb_arena_malloc(arena, len);
+  memcpy(buf, pb, len);
+
+  ok = upb_decode(buf, len, msg, layout, arena);
 
   if (!ok) {
     lua_pushstring(L, "Error decoding protobuf.");

upb/upb.c

@@ -83,10 +83,16 @@ struct upb_arena {
   /* Allocator to allocate arena blocks.  We are responsible for freeing these
    * when we are destroyed. */
   upb_alloc *block_alloc;
-  size_t last_size;
+  uint32_t last_size;
+
+  /* When multiple arenas are fused together, each arena points to a parent
+   * arena (root points to itself). The root tracks how many live arenas
+   * reference it. */
+  uint32_t refcount;  /* Only used when a->parent == a */
+  struct upb_arena *parent;
 
   /* Linked list of blocks to free/cleanup. */
-  mem_block *freelist;
+  mem_block *freelist, *freelist_tail;
 };
 
 static const size_t memblock_reserve = UPB_ALIGN_UP(sizeof(mem_block), 16);
@@ -99,6 +105,7 @@ static void upb_arena_addblock(upb_arena *a, void *ptr, size_t size) {
   block->cleanups = 0;
   a->freelist = block;
   a->last_size = size;
+  if (!a->freelist_tail) a->freelist_tail = block;
 
   a->head.ptr = UPB_PTR_AT(block, memblock_reserve, char);
   a->head.end = UPB_PTR_AT(block, size, char);
@@ -133,6 +140,17 @@ static void *upb_arena_doalloc(upb_alloc *alloc, void *ptr, size_t oldsize,
   return upb_arena_realloc(a, ptr, oldsize, size);
 }
 
+static upb_arena *arena_findroot(upb_arena *a) {
+  /* Path splitting keeps time complexity down, see:
+   *   https://en.wikipedia.org/wiki/Disjoint-set_data_structure */
+  while (a->parent != a) {
+    upb_arena *next = a->parent;
+    a->parent = next->parent;
+    a = next;
+  }
+  return a;
+}
+
 /* Public Arena API ***********************************************************/
 
 upb_arena *arena_initslow(void *mem, size_t n, upb_alloc *alloc) {
@@ -150,7 +168,10 @@ upb_arena *arena_initslow(void *mem, size_t n, upb_alloc *alloc) {
 
   a->head.alloc.func = &upb_arena_doalloc;
   a->block_alloc = alloc;
+  a->parent = a;
+  a->refcount = 1;
   a->freelist = NULL;
+  a->freelist_tail = NULL;
 
   upb_arena_addblock(a, mem, n);
 
@@ -173,6 +194,8 @@ upb_arena *upb_arena_init(void *mem, size_t n, upb_alloc *alloc) {
 
   a->head.alloc.func = &upb_arena_doalloc;
   a->block_alloc = alloc;
+  a->parent = a;
+  a->refcount = 1;
   a->last_size = 128;
   a->head.ptr = mem;
   a->head.end = UPB_PTR_AT(mem, n, char);
@@ -182,10 +205,10 @@ upb_arena *upb_arena_init(void *mem, size_t n, upb_alloc *alloc) {
   return a;
 }
 
-#undef upb_alignof
-
-void upb_arena_free(upb_arena *a) {
+static void arena_dofree(upb_arena *a) {
   mem_block *block = a->freelist;
+  UPB_ASSERT(a->parent == a);
+  UPB_ASSERT(a->refcount == 0);
 
   while (block) {
     /* Load first since we are deleting block. */
@@ -205,6 +228,11 @@ void upb_arena_free(upb_arena *a) {
   }
 }
 
+void upb_arena_free(upb_arena *a) {
+  a = arena_findroot(a);
+  if (--a->refcount == 0) arena_dofree(a);
+}
+
 bool upb_arena_addcleanup(upb_arena *a, void *ud, upb_cleanup_func *func) {
   cleanup_ent *ent;
 
@@ -222,3 +250,27 @@ bool upb_arena_addcleanup(upb_arena *a, void *ud, upb_cleanup_func *func) {
 
   return true;
 }
+
+void upb_arena_fuse(upb_arena *a1, upb_arena *a2) {
+  upb_arena *r1 = arena_findroot(a1);
+  upb_arena *r2 = arena_findroot(a2);
+
+  if (r1 == r2) return;  /* Already fused. */
+
+  /* We want to join the smaller tree to the larger tree.
+   * So swap first if they are backwards. */
+  if (r1->refcount < r2->refcount) {
+    upb_arena *tmp = r1;
+    r1 = r2;
+    r2 = tmp;
+  }
+
+  /* r1 takes over r2's freelist and refcount. */
+  r1->refcount += r2->refcount;
+  if (r2->freelist_tail) {
+    UPB_ASSERT(r2->freelist_tail->next == NULL);
+    r2->freelist_tail->next = r1->freelist;
+    r1->freelist = r2->freelist;
+  }
+  r2->parent = r1;
+}

upb/upb.h

@@ -155,6 +155,7 @@ typedef struct {
 upb_arena *upb_arena_init(void *mem, size_t n, upb_alloc *alloc);
 void upb_arena_free(upb_arena *a);
 bool upb_arena_addcleanup(upb_arena *a, void *ud, upb_cleanup_func *func);
+void upb_arena_fuse(upb_arena *a, upb_arena *b);
 void *_upb_arena_slowmalloc(upb_arena *a, size_t size);
 
 UPB_INLINE upb_alloc *upb_arena_alloc(upb_arena *a) { return (upb_alloc*)a; }