Objspace memsize

PullRequest: truffleruby/740

Author: aardvark179

lib/truffle/objspace.rb

@@ -152,7 +152,9 @@ def dump_all(output: :file)
   module_function :dump_all
 
   def memsize_of(object)
-    Truffle::ObjSpace.memsize_of(object)
+    size = Truffle::ObjSpace.memsize_of(object)
+
+    size + Truffle::ObjSpace.sizer(object).call unless Truffle::ObjSpace.sizer(object).nil?
   end
   module_function :memsize_of
 

lib/truffle/truffle/cext.rb

@@ -1415,27 +1415,29 @@ def rb_data_object_wrap(ruby_class, data, mark, free)
     object
   end
 
-  def rb_data_typed_object_wrap(ruby_class, data, data_type, mark, free)
+  def rb_data_typed_object_wrap(ruby_class, data, data_type, mark, free, size)
     ruby_class = Object if Truffle::Interop.null?(ruby_class)
     object = BASIC_OBJECT_ALLOCATE.bind(ruby_class).call
     data_holder = DataHolder.new(data)
     hidden_variable_set object, :data_type, data_type
     hidden_variable_set object, :data_holder, data_holder
     ObjectSpace.define_finalizer object, data_finalizer(free, data_holder) unless free.nil?
+    Truffle::ObjSpace.define_sizer object, data_sizer(size, data_holder) unless size.nil?
     define_marker object, data_marker(mark, data_holder) unless mark.nil?
     object
   end
 
+  # These data function are created in separate methods to ensure they
+  # will not accidentally capture the objects they operate on, which
+  # might prevent garbage collection.
   def data_finalizer(free, data_holder)
-    # In a separate method to avoid capturing the object
     raise unless free.respond_to?(:call)
     proc {
       Truffle.invoke_primitive(:interop_call_c_with_mutex, free, [data_holder.data]) unless data_holder.data.nil?
     }
   end
 
   def data_marker(mark, data_holder)
-    # In a separate method to avoid capturing the object
     raise unless mark.respond_to?(:call)
     proc { |obj|
       create_mark_list
@@ -1444,6 +1446,13 @@ def data_marker(mark, data_holder)
     }
   end
 
+  def data_sizer(sizer, data_holder)
+    raise unless sizer.respond_to?(:call)
+    proc {
+      Truffle.invoke_primitive(:interop_call_c_with_mutex, sizer, [data_holder.data])
+    }
+  end
+
   def rb_ruby_verbose_ptr
     $VERBOSE
   end

spec/ruby/optional/capi/ext/typed_data_spec.c

@@ -43,7 +43,11 @@ void sample_typed_wrapped_struct_mark(void* st) {
 }
 
 size_t sample_typed_wrapped_struct_memsize(const void* st) {
-  return sizeof(struct sample_typed_wrapped_struct);
+  if (st == NULL) {
+    return 0;
+  } else {
+    return ((struct sample_typed_wrapped_struct *)st)->foo;
+  }
 }
 
 static const rb_data_type_t sample_typed_wrapped_struct_data_type = {

spec/ruby/optional/capi/typed_data_spec.rb

@@ -1,4 +1,5 @@
 require_relative 'spec_helper'
+require 'objspace'
 
 load_extension("typed_data")
 
@@ -7,6 +8,14 @@
     @s = CApiAllocTypedSpecs.new
     @s.typed_wrapped_data.should == 42 # not defined in initialize
   end
+
+  it "uses the specified memsize function for ObjectSpace.memsize" do
+    @s = CApiAllocTypedSpecs.new
+    # The defined memsize function for the type should return 42 as
+    # the size, and this should be added to the size of the object as
+    # known by Ruby.
+    ObjectSpace.memsize_of(@s).should > 42
+  end
 end
 
 describe "CApiWrappedTypedStruct" do

src/main/c/cext/ruby.c

@@ -2838,7 +2838,7 @@ VALUE rb_data_object_alloc_managed(VALUE klass, size_t size, RUBY_DATA_FUNC dmar
 
 VALUE rb_data_typed_object_wrap(VALUE ruby_class, void *data, const rb_data_type_t *data_type) {
   return rb_tr_wrap(polyglot_invoke(RUBY_CEXT, "rb_data_typed_object_wrap",
-                                            rb_tr_unwrap(ruby_class), data, data_type, data_type->function.dmark, rb_tr_free_function(data_type->function.dfree)));
+                                    rb_tr_unwrap(ruby_class), data, data_type, data_type->function.dmark, rb_tr_free_function(data_type->function.dfree), data_type->function.dsize));
 }
 
 VALUE rb_data_typed_object_zalloc(VALUE ruby_class, size_t size, const rb_data_type_t *data_type) {

src/main/java/org/truffleruby/Layouts.java

@@ -104,6 +104,7 @@ public abstract class Layouts {
     public static final HiddenKey ASSOCIATED_IDENTIFIER = new HiddenKey("associated");
     public static final HiddenKey FINALIZER_REF_IDENTIFIER = new HiddenKey("finalizerRef");
     public static final HiddenKey MARKED_OBJECTS_IDENTIFIER = new HiddenKey("marked_objects");
+    public static final HiddenKey MEMSIZER_IDENTIFIER = new HiddenKey("memsizer");
     public static final HiddenKey VALUE_WRAPPER_IDENTIFIER = new HiddenKey("value_wrapper");
 
     // Generated layouts

src/main/java/org/truffleruby/stdlib/ObjSpaceNodes.java

@@ -19,6 +19,10 @@
 import org.truffleruby.core.regexp.MatchDataNodes.ValuesNode;
 import org.truffleruby.core.string.StringOperations;
 import org.truffleruby.language.objects.ObjectGraph;
+import org.truffleruby.language.objects.ReadObjectFieldNode;
+import org.truffleruby.language.objects.ReadObjectFieldNodeGen;
+import org.truffleruby.language.objects.WriteObjectFieldNode;
+import org.truffleruby.language.objects.WriteObjectFieldNodeGen;
 
 import java.util.Set;
 
@@ -65,7 +69,29 @@ public int memsizeOfObject(DynamicObject object) {
         public int memsize(Object object) {
             return 0;
         }
+    }
+
+    @CoreMethod(names = "sizer", isModuleFunction = true, required = 1)
+    public abstract static class GetMemSizer extends CoreMethodArrayArgumentsNode {
+
+        @Child private ReadObjectFieldNode readSizerNode = ReadObjectFieldNodeGen.create(Layouts.MEMSIZER_IDENTIFIER, nil());
+
+        @Specialization
+        public Object hasSizer(DynamicObject object) {
+            return readSizerNode.execute(object);
+        }
+    }
+
+    @CoreMethod(names = "define_sizer", isModuleFunction = true, required = 2)
+    public abstract static class DefineMemSizer extends CoreMethodArrayArgumentsNode {
+
+        @Child private WriteObjectFieldNode setSizerNode = WriteObjectFieldNodeGen.create(Layouts.MEMSIZER_IDENTIFIER);
 
+        @Specialization
+        public DynamicObject defineSizer(DynamicObject object, DynamicObject sizer) {
+            setSizerNode.write(object, sizer);
+            return nil();
+        }
     }
 
     @CoreMethod(names = "adjacent_objects", isModuleFunction = true, required = 1)