From ab9b741f9ef914c69039f75fc00b75cc83c585fd Mon Sep 17 00:00:00 2001
From: Jean Boussier <jean.boussier@gmail.com>
Date: Thu, 21 Jul 2022 08:34:52 +0200
Subject: [PATCH] Ruby implement memsize functions for native types

Fix: https://github.com/protocolbuffers/protobuf/issues/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 |  7 ++++++-
 ruby/ext/google/protobuf_c/ruby-upb.c | 19 +++++++++++++++++++
 ruby/ext/google/protobuf_c/ruby-upb.h |  1 +
 ruby/tests/memory_test.rb             | 25 +++++++++++++++++++++++++
 6 files changed, 61 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 5d30319c198e..2415bd078b71 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 6b8bbaa3c5e5..10e71db4e30a 100644
--- a/ruby/ext/google/protobuf_c/message.c
+++ b/ruby/ext/google/protobuf_c/message.c
@@ -63,9 +63,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 = {
     "Message",
-    {Message_mark, RUBY_DEFAULT_FREE, NULL},
+    {Message_mark, RUBY_DEFAULT_FREE, Message_memsize},
     .flags = RUBY_TYPED_FREE_IMMEDIATELY,
 };
 
diff --git a/ruby/ext/google/protobuf_c/protobuf.c b/ruby/ext/google/protobuf_c/protobuf.c
index 3c765c564d23..d605d7e91de7 100644
--- a/ruby/ext/google/protobuf_c/protobuf.c
+++ b/ruby/ext/google/protobuf_c/protobuf.c
@@ -185,11 +185,16 @@ static void Arena_free(void *data) {
   xfree(arena);
 }
 
+static size_t Arena_memsize(const void *data) {
+  const Arena *arena = data;
+  return _upb_ArenaMemsize(arena->arena) + 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,
 };
 
diff --git a/ruby/ext/google/protobuf_c/ruby-upb.c b/ruby/ext/google/protobuf_c/ruby-upb.c
index d50bd99f9c42..6efe9f6e4f16 100755
--- a/ruby/ext/google/protobuf_c/ruby-upb.c
+++ b/ruby/ext/google/protobuf_c/ruby-upb.c
@@ -3349,6 +3349,25 @@ static upb_Arena* arena_findroot(upb_Arena* a) {
   return a;
 }
 
+size_t _upb_ArenaMemsize(upb_Arena *arena) {
+  size_t memsize = 0;
+
+  mem_block* block = arena->freelist;
+
+  while (block) {
+    memsize += sizeof(mem_block) + block->size;
+    block = block->next;
+  }
+
+  if (arena->refcount > 1) {
+    // If other arena were fused we attirbute an equal
+    // share of memory usage to each one.
+    memsize /= arena->refcount;
+  }
+
+  return memsize;
+}
+
 static void upb_Arena_addblock(upb_Arena* a, upb_Arena* root, void* ptr,
                                size_t size) {
   mem_block* block = ptr;
diff --git a/ruby/ext/google/protobuf_c/ruby-upb.h b/ruby/ext/google/protobuf_c/ruby-upb.h
index e57eb0edad00..4b7bf34c5f40 100755
--- a/ruby/ext/google/protobuf_c/ruby-upb.h
+++ b/ruby/ext/google/protobuf_c/ruby-upb.h
@@ -450,6 +450,7 @@ void upb_Arena_Free(upb_Arena* a);
 bool upb_Arena_AddCleanup(upb_Arena* a, void* ud, upb_CleanupFunc* func);
 bool upb_Arena_Fuse(upb_Arena* a, upb_Arena* b);
 void* _upb_Arena_SlowMalloc(upb_Arena* a, size_t size);
+size_t _upb_ArenaMemsize(upb_Arena *arena);
 
 UPB_INLINE upb_alloc* upb_Arena_Alloc(upb_Arena* a) { return (upb_alloc*)a; }
 
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