From fd75ecc960796a5f49207ae69cfefc299d3953e8 Mon Sep 17 00:00:00 2001
From: Jean Boussier <jean.boussier@gmail.com>
Date: Mon, 6 Mar 2023 11:31:04 +0100
Subject: Implement Write Barrier and dsize for FFI::Function

Ref: https://github.com/ffi/ffi/pull/991

Write barrier protected objects are allowed to be promoted to the old generation,
which means they only get marked on major GC.

The downside is that the RB_BJ_WRITE macro MUST be used to set references,
otherwise the referenced object may be garbaged collected.

This commit also implement a `dsize` function so that these instance report
a more relevant size in various memory profilers. It's not counting everything
because some types are opaque right now, so a larger refactoring would be needed.
---
 ext/ffi_c/Function.c      | 33 +++++++++++++++++++++++++--------
 spec/ffi/function_spec.rb |  8 ++++++++
 2 files changed, 33 insertions(+), 8 deletions(-)

diff --git a/ext/ffi_c/Function.c b/ext/ffi_c/Function.c
index 20822da..b73e8b7 100644
--- a/ext/ffi_c/Function.c
+++ b/ext/ffi_c/Function.c
@@ -77,6 +77,7 @@ typedef struct Function_ {
 
 static void function_mark(void *data);
 static void function_free(void *data);
+static size_t function_memsize(const void *data);
 static VALUE function_init(VALUE self, VALUE rbFunctionInfo, VALUE rbProc);
 static void callback_invoke(ffi_cif* cif, void* retval, void** parameters, void* user_data);
 static bool callback_prep(void* ctx, void* code, Closure* closure, char* errmsg, size_t errmsgsize);
@@ -100,10 +101,12 @@ static const rb_data_type_t function_data_type = {
     .function = {
         .dmark = function_mark,
         .dfree = function_free,
-        .dsize = NULL,
+        .dsize = function_memsize,
     },
     .parent = &rbffi_pointer_data_type,
-    .flags = RUBY_TYPED_FREE_IMMEDIATELY
+    // IMPORTANT: WB_PROTECTED objects must only use the RB_OBJ_WRITE()
+    // macro to update VALUE references, as to trigger write barriers.
+    .flags = RUBY_TYPED_FREE_IMMEDIATELY | RUBY_TYPED_WB_PROTECTED
 };
 
 VALUE rbffi_FunctionClass = Qnil;
@@ -153,9 +156,9 @@ function_allocate(VALUE klass)
     obj = TypedData_Make_Struct(klass, Function, &function_data_type, fn);
 
     fn->base.memory.flags = MEM_RD;
-    fn->base.rbParent = Qnil;
-    fn->rbProc = Qnil;
-    fn->rbFunctionInfo = Qnil;
+    RB_OBJ_WRITE(obj, &fn->base.rbParent, Qnil);
+    RB_OBJ_WRITE(obj, &fn->rbProc, Qnil);
+    RB_OBJ_WRITE(obj, &fn->rbFunctionInfo, Qnil);
     fn->autorelease = true;
 
     return obj;
@@ -185,6 +188,20 @@ function_free(void *data)
     xfree(fn);
 }
 
+static size_t
+function_memsize(const void *data)
+{
+    const Function *fn = (const Function *)data;
+    size_t memsize = sizeof(Function);
+
+    // Would be nice to better account for MethodHandle and Closure too.
+    if (fn->closure) {
+        memsize += sizeof(Closure);
+    }
+
+    return memsize;
+}
+
 /*
  * @param [Type, Symbol] return_type return type for the function
  * @param [Array<Type, Symbol>] param_types array of parameters types
@@ -313,7 +330,7 @@ function_init(VALUE self, VALUE rbFunctionInfo, VALUE rbProc)
 
     TypedData_Get_Struct(self, Function, &function_data_type, fn);
 
-    fn->rbFunctionInfo = rbFunctionInfo;
+    RB_OBJ_WRITE(self, &fn->rbFunctionInfo, rbFunctionInfo);
 
     TypedData_Get_Struct(fn->rbFunctionInfo, FunctionType, &rbffi_fntype_data_type, fn->info);
 
@@ -321,7 +338,7 @@ function_init(VALUE self, VALUE rbFunctionInfo, VALUE rbProc)
         Pointer* orig;
         TypedData_Get_Struct(rbProc, Pointer, &rbffi_pointer_data_type, orig);
         fn->base.memory = orig->memory;
-        fn->base.rbParent = rbProc;
+        RB_OBJ_WRITE(self, &fn->base.rbParent, rbProc);
 
     } else if (rb_obj_is_kind_of(rbProc, rb_cProc) || rb_respond_to(rbProc, id_call)) {
         if (fn->info->closurePool == NULL) {
@@ -357,7 +374,7 @@ function_init(VALUE self, VALUE rbFunctionInfo, VALUE rbProc)
                 rb_obj_classname(rbProc));
     }
 
-    fn->rbProc = rbProc;
+    RB_OBJ_WRITE(self, &fn->rbProc, rbProc);
 
     return self;
 }
diff --git a/spec/ffi/function_spec.rb b/spec/ffi/function_spec.rb
index 5bc2421..a1a50f3 100644
--- a/spec/ffi/function_spec.rb
+++ b/spec/ffi/function_spec.rb
@@ -103,4 +103,12 @@ describe FFI::Function do
     fp = FFI::Function.new(:int, [:int, :int], @libtest.find_function('testAdd'))
     expect { fp.free }.to raise_error RuntimeError
   end
+
+  it 'has a memsize function', skip: RUBY_ENGINE != "ruby" do
+    base_size = ObjectSpace.memsize_of(Object.new)
+
+    function = FFI::Function.new(:int, [:int, :int], @libtest.find_function('testAdd'))
+    size = ObjectSpace.memsize_of(function)
+    expect(size).to be > base_size
+  end
 end
-- 
cgit v1.2.1