From 276add0b5c160fd8076cd67b47b5898d7d955400 Mon Sep 17 00:00:00 2001
From: zhangskz <89936743+zhangskz@users.noreply.github.com>
Date: Tue, 8 Mar 2022 12:05:34 -0500
Subject: [PATCH] Add ruby_upb_alloc using xrealloc/xfree so Ruby GC is aware
 of allocated memory for Arenas. (#9586)

* Add ruby-specific upb_alloc using xrealloc/xfree for use in Arena_alloc so Ruby GC is aware of allocated memory.
* Add RB_GC_GUARD to DescriptorPool_add_serialized_file  to ensure ruby does not aggressively garbage collect arena_rb due to lack of references.
---
 ruby/ext/google/protobuf_c/defs.c     |  1 +
 ruby/ext/google/protobuf_c/protobuf.c | 13 ++++++++++++-
 ruby/tests/gc_test.rb                 |  1 -
 3 files changed, 13 insertions(+), 2 deletions(-)

diff --git a/ruby/ext/google/protobuf_c/defs.c b/ruby/ext/google/protobuf_c/defs.c
index aaa8b4dc7a..3bd18e8400 100644
--- a/ruby/ext/google/protobuf_c/defs.c
+++ b/ruby/ext/google/protobuf_c/defs.c
@@ -159,6 +159,7 @@ VALUE DescriptorPool_add_serialized_file(VALUE _self,
     rb_raise(cTypeError, "Unable to build file to DescriptorPool: %s",
              upb_Status_ErrorMessage(&status));
   }
+  RB_GC_GUARD(arena_rb);
   return get_filedef_obj(_self, filedef);
 }
 
diff --git a/ruby/ext/google/protobuf_c/protobuf.c b/ruby/ext/google/protobuf_c/protobuf.c
index 4d3e1a5144..2135cca462 100644
--- a/ruby/ext/google/protobuf_c/protobuf.c
+++ b/ruby/ext/google/protobuf_c/protobuf.c
@@ -193,9 +193,20 @@ const rb_data_type_t Arena_type = {
     .flags = RUBY_TYPED_FREE_IMMEDIATELY,
 };
 
+static void* ruby_upb_allocfunc(upb_alloc* alloc, void* ptr, size_t oldsize, size_t size) {
+  if (size == 0) {
+    xfree(ptr);
+    return NULL;
+  } else {
+    return xrealloc(ptr, size);
+  }
+}
+
+upb_alloc ruby_upb_alloc = {&ruby_upb_allocfunc};
+
 static VALUE Arena_alloc(VALUE klass) {
   Arena *arena = ALLOC(Arena);
-  arena->arena = upb_Arena_New();
+  arena->arena = upb_Arena_Init(NULL, 0, &ruby_upb_alloc);
   arena->pinned_objs = Qnil;
   return TypedData_Wrap_Struct(klass, &Arena_type, arena);
 }
diff --git a/ruby/tests/gc_test.rb b/ruby/tests/gc_test.rb
index d7fecaeea1..5d48f464ae 100755
--- a/ruby/tests/gc_test.rb
+++ b/ruby/tests/gc_test.rb
@@ -94,7 +94,6 @@ class GCTest < Test::Unit::TestCase
     from = get_msg_proto3
     data = A::B::C::TestMessage.encode(from)
     to = A::B::C::TestMessage.decode(data)
-
     # This doesn't work for proto2 on JRuby because there is a nested required message.
     # A::B::Proto2::TestMessage has :required_msg which is of type:
     # A::B::Proto2::TestMessage so there is no way to generate a valid