Pin object ivars while they are being copied in a hash

When transitioning an object to TOO_COMPLEX we copy all its
ivar in a table, but if GC (and compaction) trigger in the middle
of the transition, the old `iv_ptr` might still be the canonical
ivar list and will be updated by the GC, but the reference we
copied in the table will be outdated.

So we must pin these reference until they are all copied and
the object `iv_ptr` is pointing to the table.

To do this we allocate a temporary TOO_COMPLEX object
which we use as the host for the copy. TOO_COMPLEX objects
are marked with `mark_tbl` which does pin references.

Co-Authored-By: Alan Wu <[email protected]>

Author: byroot

internal/variable.h

@@ -47,7 +47,8 @@ VALUE rb_mod_set_temporary_name(VALUE, VALUE);
 
 struct gen_ivtbl;
 int rb_gen_ivtbl_get(VALUE obj, ID id, struct gen_ivtbl **ivtbl);
-void rb_obj_copy_ivs_to_hash_table(VALUE obj, st_table *table);
+VALUE rb_obj_copy_ivs_to_hash_table(VALUE obj, st_table *table);
+void rb_obj_copy_ivs_to_hash_table_complete(VALUE ivar_pinner);
 void rb_obj_convert_to_too_complex(VALUE obj, st_table *table);
 void rb_evict_ivars_to_hash(VALUE obj);
 

object.c

@@ -326,8 +326,9 @@ rb_obj_copy_ivar(VALUE dest, VALUE obj)
         shape_to_set_on_dest = rb_shape_rebuild_shape(initial_shape, src_shape);
         if (UNLIKELY(rb_shape_id(shape_to_set_on_dest) == OBJ_TOO_COMPLEX_SHAPE_ID)) {
             st_table * table = rb_st_init_numtable_with_size(src_num_ivs);
-            rb_obj_copy_ivs_to_hash_table(obj, table);
+            VALUE ivar_pinner = rb_obj_copy_ivs_to_hash_table(obj, table);
             rb_obj_convert_to_too_complex(dest, table);
+            rb_obj_copy_ivs_to_hash_table_complete(ivar_pinner);
 
             return;
         }

test/ruby/test_shapes.rb

@@ -256,9 +256,9 @@ def test_run_out_of_shape_for_class_ivar
     assert_separately([], "#{<<~"begin;"}\n#{<<~'end;'}")
     begin;
       i = 0
+      o = Object.new
       while RubyVM::Shape.shapes_available > 0
-        c = Class.new
-        c.instance_variable_set(:"@i#{i}", 1)
+        o.instance_variable_set(:"@i#{i}", 1)
         i += 1
       end
 
@@ -275,6 +275,103 @@ def test_run_out_of_shape_for_class_ivar
     end;
   end
 
+  def test_evacuate_class_ivar_and_compaction
+    assert_separately([], "#{<<~"begin;"}\n#{<<~'end;'}")
+
+    begin;
+      count = 20
+
+      c = Class.new
+      count.times do |ivar|
+        c.instance_variable_set("@i#{ivar}", "ivar-#{ivar}")
+      end
+
+      i = 0
+      o = Object.new
+      while RubyVM::Shape.shapes_available > 0
+        o.instance_variable_set("@i#{i}", 1)
+        i += 1
+      end
+
+      GC.auto_compact = true
+      GC.stress = true
+      # Cause evacuation
+      c.instance_variable_set(:@a, o = Object.new)
+      assert_equal(o, c.instance_variable_get(:@a))
+      GC.stress = false
+
+      count.times do |ivar|
+        assert_equal "ivar-#{ivar}", c.instance_variable_get("@i#{ivar}")
+      end
+    end;
+  end
+
+  def test_evacuate_generic_ivar_and_compaction
+    assert_separately([], "#{<<~"begin;"}\n#{<<~'end;'}")
+
+    begin;
+      count = 20
+
+      c = Hash.new
+      count.times do |ivar|
+        c.instance_variable_set("@i#{ivar}", "ivar-#{ivar}")
+      end
+
+      i = 0
+      o = Object.new
+      while RubyVM::Shape.shapes_available > 0
+        o.instance_variable_set("@i#{i}", 1)
+        i += 1
+      end
+
+      GC.auto_compact = true
+      GC.stress = true
+
+      # Cause evacuation
+      c.instance_variable_set(:@a, o = Object.new)
+      assert_equal(o, c.instance_variable_get(:@a))
+
+      GC.stress = false
+
+      count.times do |ivar|
+        assert_equal "ivar-#{ivar}", c.instance_variable_get("@i#{ivar}")
+      end
+    end;
+  end
+
+  def test_evacuate_object_ivar_and_compaction
+    assert_separately([], "#{<<~"begin;"}\n#{<<~'end;'}")
+
+    begin;
+      count = 20
+
+      c = Object.new
+      count.times do |ivar|
+        c.instance_variable_set("@i#{ivar}", "ivar-#{ivar}")
+      end
+
+      i = 0
+      o = Object.new
+      while RubyVM::Shape.shapes_available > 0
+        o.instance_variable_set("@i#{i}", 1)
+        i += 1
+      end
+
+      GC.auto_compact = true
+      GC.stress = true
+
+      # Cause evacuation
+      c.instance_variable_set(:@a, o = Object.new)
+      assert_equal(o, c.instance_variable_get(:@a))
+
+      GC.stress = false
+
+      count.times do |ivar|
+        assert_equal "ivar-#{ivar}", c.instance_variable_get("@i#{ivar}")
+      end
+    end;
+  end
+
   def test_run_out_of_shape_for_module_ivar
     assert_separately([], "#{<<~"begin;"}\n#{<<~'end;'}")
     begin;

variable.c

@@ -1370,13 +1370,22 @@ rb_attr_delete(VALUE obj, ID id)
     return rb_ivar_delete(obj, id, Qnil);
 }
 
+static int
+rb_obj_write_barrier_ivars_i(ID key, VALUE val, st_data_t arg)
+{
+    RB_OBJ_WRITTEN((VALUE)arg, Qundef, val);
+    return ST_CONTINUE;
+}
+
 void
 rb_obj_convert_to_too_complex(VALUE obj, st_table *table)
 {
     RUBY_ASSERT(!rb_shape_obj_too_complex(obj));
 
     VALUE *old_ivptr = NULL;
 
+    rb_ivar_foreach(obj, rb_obj_write_barrier_ivars_i, (st_data_t)obj);
+
     switch (BUILTIN_TYPE(obj)) {
         case T_OBJECT:
           if (!(RBASIC(obj)->flags & ROBJECT_EMBED)) {
@@ -1422,8 +1431,9 @@ rb_evict_ivars_to_hash(VALUE obj)
     st_table *table = st_init_numtable_with_size(rb_ivar_count(obj));
 
     // Evacuate all previous values from shape into id_table
-    rb_obj_copy_ivs_to_hash_table(obj, table);
+    VALUE ivar_pinner = rb_obj_copy_ivs_to_hash_table(obj, table);
     rb_obj_convert_to_too_complex(obj, table);
+    rb_obj_copy_ivs_to_hash_table_complete(ivar_pinner);
 
     RUBY_ASSERT(rb_shape_obj_too_complex(obj));
 }
@@ -1640,17 +1650,44 @@ rb_ensure_iv_list_size(VALUE obj, uint32_t current_capacity, uint32_t new_capaci
     }
 }
 
-int
-rb_obj_copy_ivs_to_hash_table_i(ID key, VALUE val, st_data_t arg)
+struct rb_evacuate_arg {
+    st_table *table;
+    VALUE host;
+};
+
+static int
+copy_ivs_to_hash_table_i(ID key, VALUE val, st_data_t arg)
 {
-    st_insert((st_table *)arg, (st_data_t)key, (st_data_t)val);
+    struct rb_evacuate_arg *evac_arg = (struct rb_evacuate_arg *)arg;
+    st_insert(evac_arg->table, (st_data_t)key, (st_data_t)val);
+    RB_OBJ_WRITTEN(evac_arg->host, Qundef, val);
     return ST_CONTINUE;
 }
 
-void
+VALUE
 rb_obj_copy_ivs_to_hash_table(VALUE obj, st_table *table)
 {
-    rb_ivar_foreach(obj, rb_obj_copy_ivs_to_hash_table_i, (st_data_t)table);
+    // There can be compaction runs between each ivar we copy out of obj, so we
+    // need an object to mark each ivar to make sure every reference is valid
+    // by the time we're done. Use a special T_OBJECT for marking.
+    VALUE ivar_pinner = rb_wb_protected_newobj_of(GET_EC(), rb_cBasicObject, T_OBJECT | ROBJECT_EMBED, RVALUE_SIZE);
+    rb_shape_set_shape_id(ivar_pinner, OBJ_TOO_COMPLEX_SHAPE_ID);
+    ROBJECT_SET_IV_HASH(ivar_pinner, table);
+
+    // Evacuate all previous values from shape into id_table
+    struct rb_evacuate_arg evac_arg = { .table = table, .host = ivar_pinner };
+    rb_ivar_foreach(obj, copy_ivs_to_hash_table_i, (st_data_t)&evac_arg);
+
+    return ivar_pinner;
+}
+
+void
+rb_obj_copy_ivs_to_hash_table_complete(VALUE ivar_pinner)
+{
+    // done with pinning, now set pinner to 0 ivars
+    ROBJECT_SET_IV_HASH(ivar_pinner, NULL);
+    rb_shape_set_shape(ivar_pinner, rb_shape_get_root_shape());
+    RB_GC_GUARD(ivar_pinner);
 }
 
 static VALUE *