From 9150795328267e69a0b8e241dda675402e5633e2 Mon Sep 17 00:00:00 2001
From: Jean Boussier <jean.boussier@gmail.com>
Date: Fri, 17 Feb 2023 08:19:05 +0100
Subject: [PATCH] Ruby implement memsize functions for native types

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.
---
 ruby/ext/google/protobuf_c/map.c      |  6 +++++-
 ruby/ext/google/protobuf_c/message.c  |  6 +++++-
 ruby/ext/google/protobuf_c/protobuf.c | 13 ++++++++++++-
 ruby/tests/memory_test.rb             | 25 +++++++++++++++++++++++++
 4 files changed, 47 insertions(+), 3 deletions(-)
 create mode 100755 ruby/tests/memory_test.rb

diff --git a/ruby/ext/google/protobuf_c/map.c b/ruby/ext/google/protobuf_c/map.c
index 18617ed1c2d6..a4f0979d52c2 100644
--- 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
index 067a44661576..79886e56305e 100644
--- 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
index bee8ac288f71..96c4e538fcd5 100644
--- 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
new file mode 100755
index 000000000000..6d3f79ac12f1
--- /dev/null
+++ 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