Ruby implement memsize functions for native types (#10291)

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.

cc @zhangskz @haberman

Closes #10291

COPYBARA_INTEGRATE_REVIEW=#10291 from casperisfine:ruby-sizes 9150795
PiperOrigin-RevId: 606718632

ruby/ext/google/protobuf_c/map.c

@@ -38,9 +38,11 @@ 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,
 };
 

ruby/ext/google/protobuf_c/message.c

@@ -44,9 +44,11 @@ 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,
 };
 

ruby/ext/google/protobuf_c/protobuf.c

@@ -164,11 +164,23 @@ static void Arena_free(void *data) {
   xfree(arena);
 }
 
+static size_t Arena_memsize(const void *data) {
+  const Arena *arena = data;
+  size_t fused_count;
+  size_t memsize = upb_Arena_SpaceAllocated(arena->arena, &fused_count);
+  if (fused_count > 1) {
+    // If other arena were fused we attribute an equal
+    // share of memory usage to each one.
+    memsize /= fused_count;
+  }
+  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,
 };
 

ruby/tests/memory_test.rb

@@ -0,0 +1,33 @@
+#!/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'
+
+$is_64bit = Google::Protobuf::Internal::SIZEOF_LONG == 8
+
+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
+    if $is_64bit
+      assert_operator 40, :<, ObjectSpace.memsize_of(Google::Protobuf::Internal::Arena.new)
+    end
+  end
+
+  def test_objspace_memsize_of_message
+    if $is_64bit
+      assert_operator 40, :<, ObjectSpace.memsize_of(FooBar::TestImportedMessage.new)
+    end
+  end
+
+  def test_objspace_memsize_of_map
+    if $is_64bit
+      assert_operator 40, :<, ObjectSpace.memsize_of(Google::Protobuf::Map.new(:string, :int32))
+    end
+  end
+end

upb/mem/arena.c

@@ -155,9 +155,10 @@ static upb_ArenaRoot _upb_Arena_FindRoot(upb_Arena* a) {
   return (upb_ArenaRoot){.root = ai, .tagged_count = poc};
 }
 
-size_t upb_Arena_SpaceAllocated(upb_Arena* arena) {
+size_t upb_Arena_SpaceAllocated(upb_Arena* arena, size_t* fused_count) {
   upb_ArenaInternal* ai = _upb_Arena_FindRoot(arena).root;
   size_t memsize = 0;
+  size_t local_fused_count = 0;
 
   while (ai != NULL) {
     upb_MemBlock* block = upb_Atomic_Load(&ai->blocks, memory_order_relaxed);
@@ -166,8 +167,10 @@ size_t upb_Arena_SpaceAllocated(upb_Arena* arena) {
       block = upb_Atomic_Load(&block->next, memory_order_relaxed);
     }
     ai = upb_Atomic_Load(&ai->next, memory_order_relaxed);
+    local_fused_count++;
   }
 
+  if (fused_count) *fused_count = local_fused_count;
   return memsize;
 }
 

upb/mem/arena.h

@@ -46,7 +46,7 @@ UPB_API bool upb_Arena_Fuse(upb_Arena* a, upb_Arena* b);
 bool upb_Arena_IncRefFor(upb_Arena* a, const void* owner);
 void upb_Arena_DecRefFor(upb_Arena* a, const void* owner);
 
-size_t upb_Arena_SpaceAllocated(upb_Arena* a);
+size_t upb_Arena_SpaceAllocated(upb_Arena* a, size_t* fused_count);
 uint32_t upb_Arena_DebugRefCount(upb_Arena* a);
 
 UPB_API_INLINE upb_Arena* upb_Arena_New(void) {