Ruby implement memsize functions for native types

byroot · byroot · commit 915079532826 · 2023-02-28T08:58:09.000+01:00
Fix: #10280 This allows Ruby to report a more correct estimation of the memory used by these objects. It's useful when running memory profilers against applications.
diff --git a/ruby/ext/google/protobuf_c/map.c b/ruby/ext/google/protobuf_c/map.c
@@ -61,9 +61,13 @@ static void Map_mark(void* _self) {
   rb_gc_mark(self->arena);
 }
 
+static size_t Map_memsize(const void* _self) {
+  return sizeof(Map);
+}
+
 const rb_data_type_t Map_type = {
     "Google::Protobuf::Map",
-    {Map_mark, RUBY_DEFAULT_FREE, NULL},
+    {Map_mark, RUBY_DEFAULT_FREE, Map_memsize},
     .flags = RUBY_TYPED_FREE_IMMEDIATELY,
 };
 
diff --git a/ruby/ext/google/protobuf_c/message.c b/ruby/ext/google/protobuf_c/message.c
@@ -66,9 +66,13 @@ static void Message_mark(void* _self) {
   rb_gc_mark(self->arena);
 }
 
+static size_t Message_memsize(const void* _self) {
+  return sizeof(Message);
+}
+
 static rb_data_type_t Message_type = {
     "Google::Protobuf::Message",
-    {Message_mark, RUBY_DEFAULT_FREE, NULL},
+    {Message_mark, RUBY_DEFAULT_FREE, Message_memsize},
     .flags = RUBY_TYPED_FREE_IMMEDIATELY | RUBY_TYPED_WB_PROTECTED,
 };
 
diff --git a/ruby/ext/google/protobuf_c/protobuf.c b/ruby/ext/google/protobuf_c/protobuf.c
@@ -187,11 +187,22 @@ static void Arena_free(void *data) {
   xfree(arena);
 }
 
+static size_t Arena_memsize(const void *data) {
+  const Arena *arena = data;
+  size_t memsize = upb_Arena_SpaceAllocated(arena->arena);
+  if (arena->arena->refcount > 1) {
+      // If other arena were fused we attribute an equal
+      // share of memory usage to each one.
+      memsize /= arena->arena->refcount;
+  }
+  return memsize + sizeof(Arena);
+}
+
 static VALUE cArena;
 
 const rb_data_type_t Arena_type = {
     "Google::Protobuf::Internal::Arena",
-    {Arena_mark, Arena_free, NULL},
+    {Arena_mark, Arena_free, Arena_memsize},
     .flags = RUBY_TYPED_FREE_IMMEDIATELY | RUBY_TYPED_WB_PROTECTED,
 };
 
diff --git a/ruby/tests/memory_test.rb b/ruby/tests/memory_test.rb
@@ -0,0 +1,25 @@
+#!/usr/bin/ruby
+#
+# generated_code.rb is in the same directory as this test.
+$LOAD_PATH.unshift(File.expand_path(File.dirname(__FILE__)))
+
+require 'test/unit'
+require 'objspace'
+require 'test_import_pb'
+
+class MemoryTest < Test::Unit::TestCase
+  # 40 byte is the default object size. But the real size is dependent on many things
+  # such as arch etc, so there's no point trying to assert the exact return value here.
+  # We merely assert that we return something other than the default.
+  def test_objspace_memsize_of_arena
+    assert_operator 40, :<, ObjectSpace.memsize_of(Google::Protobuf::Internal::Arena.new)
+  end
+
+  def test_objspace_memsize_of_message
+    assert_operator 40, :<, ObjectSpace.memsize_of(FooBar::TestImportedMessage.new)
+  end
+
+  def test_objspace_memsize_of_map
+    assert_operator 40, :<, ObjectSpace.memsize_of(Google::Protobuf::Map.new(:string, :int32))
+  end
+end
