gc.c: Fix a race condition in object_id for shareable objects

If an object is shareable and has no capacity left, it isn't
safe to store the object ID in fields as it requires an object
resize which can't be done unless all field reads are synchronized.

So in this case we have to store the ID externally like we used to.

Author: byroot

gc.c

@@ -1879,6 +1879,61 @@ static const rb_data_type_t id2ref_tbl_type = {
     .flags = RUBY_TYPED_WB_PROTECTED | RUBY_TYPED_FREE_IMMEDIATELY
 };
 
+static VALUE obj_to_id_value = 0;
+static st_table *obj_to_id_tbl = NULL;
+
+static void mark_hash_values(st_table *tbl);
+
+static void
+obj_to_id_tbl_mark(void *data)
+{
+    st_table *table = (st_table *)data;
+    if (UNLIKELY(!RB_POSFIXABLE(LAST_OBJECT_ID()))) {
+        // It's very unlikely, but if enough object ids were generated, keys may be T_BIGNUM
+        mark_hash_values(table);
+    }
+    // We purposedly don't mark keys, as they are weak references.
+    // rb_gc_obj_free_vm_weak_references takes care of cleaning them up.
+}
+
+static size_t
+obj_to_id_tbl_memsize(const void *data)
+{
+    return rb_st_memsize(data);
+}
+
+static void
+obj_to_id_tbl_compact(void *data)
+{
+    st_table *table = (st_table *)data;
+    if (LIKELY(RB_POSFIXABLE(LAST_OBJECT_ID()))) {
+        // We know values are all FIXNUM, so no need to update them.
+        gc_ref_update_table_keys_only(table);
+    }
+    else {
+        gc_update_table_refs(table);
+    }
+}
+
+static void
+obj_to_id_tbl_free(void *data)
+{
+    id2ref_tbl = NULL; // clear global ref
+    st_table *table = (st_table *)data;
+    st_free_table(table);
+}
+
+static const rb_data_type_t obj_to_id_tbl_type = {
+    .wrap_struct_name = "VM/obj_to_id_table",
+    .function = {
+        .dmark = obj_to_id_tbl_mark,
+        .dfree = obj_to_id_tbl_free,
+        .dsize = obj_to_id_tbl_memsize,
+        .dcompact = obj_to_id_tbl_compact,
+    },
+    .flags = RUBY_TYPED_WB_PROTECTED | RUBY_TYPED_FREE_IMMEDIATELY
+};
+
 #define RUBY_ATOMIC_VALUE_LOAD(x) (VALUE)(RUBY_ATOMIC_PTR_LOAD(x))
 
 static VALUE
@@ -1901,23 +1956,46 @@ class_object_id(VALUE klass)
 }
 
 static VALUE
-object_id0(VALUE obj)
+object_id0(VALUE obj, bool shareable)
 {
     VALUE id = Qfalse;
 
-    if (rb_shape_has_object_id(rb_obj_shape(obj))) {
-        shape_id_t object_id_shape_id = rb_shape_transition_object_id(obj);
-        id = rb_obj_field_get(obj, object_id_shape_id);
+    rb_shape_t *current_shape = rb_obj_shape(obj);
+
+    if (rb_shape_has_object_id(current_shape)) {
+        shape_id_t object_id_shape_id = rb_shape_object_id(obj);
+
+        if (RB_UNLIKELY(shareable && RSHAPE(object_id_shape_id)->type == SHAPE_EXTERNAL_OBJ_ID)) {
+            RUBY_ASSERT(obj_to_id_tbl);
+            st_lookup(obj_to_id_tbl, obj, &id);
+        }
+        else {
+            id = rb_obj_field_get(obj, object_id_shape_id);
+        }
         RUBY_ASSERT(id, "object_id missing");
         return id;
     }
 
-    // rb_shape_object_id_shape may lock if the current shape has
-    // multiple children.
-    shape_id_t object_id_shape_id = rb_shape_transition_object_id(obj);
-
     id = generate_next_object_id();
-    rb_obj_field_set(obj, object_id_shape_id, id);
+    // If the object is shareable and doesn't have free capacity we can't safely
+    // resize it. So we have to store the object_id externally.
+    if (shareable && current_shape->next_field_index == current_shape->capacity) {
+        shape_id_t object_id_shape_id = rb_shape_transition_external_object_id(obj);
+
+        if (RB_UNLIKELY(!obj_to_id_tbl)) {
+            obj_to_id_tbl = st_init_numtable();
+            obj_to_id_value = TypedData_Wrap_Struct(0, &obj_to_id_tbl_type, obj_to_id_tbl);
+        }
+        st_insert(obj_to_id_tbl, obj, id);
+        rb_shape_set_shape_id(obj, object_id_shape_id);
+    }
+    else {
+        // rb_shape_object_id_shape may lock if the current shape has
+        // multiple children.
+        shape_id_t object_id_shape_id = rb_shape_transition_object_id(obj);
+        rb_obj_field_set(obj, object_id_shape_id, id);
+    }
+
     if (RB_UNLIKELY(id2ref_tbl)) {
         st_insert(id2ref_tbl, (st_data_t)id, (st_data_t)obj);
     }
@@ -1940,12 +2018,12 @@ object_id(VALUE obj)
 
     if (UNLIKELY(rb_gc_multi_ractor_p() && rb_ractor_shareable_p(obj))) {
         unsigned int lock_lev = rb_gc_vm_lock();
-        VALUE id = object_id0(obj);
+        VALUE id = object_id0(obj, true);
         rb_gc_vm_unlock(lock_lev);
         return id;
     }
 
-    return object_id0(obj);
+    return object_id0(obj, false);
 }
 
 static void
@@ -2717,6 +2795,14 @@ mark_key(st_data_t key, st_data_t value, st_data_t data)
     return ST_CONTINUE;
 }
 
+static int
+mark_value(st_data_t key, st_data_t value, st_data_t data)
+{
+    gc_mark_internal((VALUE)value);
+
+    return ST_CONTINUE;
+}
+
 void
 rb_mark_set(st_table *tbl)
 {
@@ -2725,6 +2811,14 @@ rb_mark_set(st_table *tbl)
     st_foreach(tbl, mark_key, (st_data_t)rb_gc_get_objspace());
 }
 
+static void
+mark_hash_values(st_table *tbl)
+{
+    if (!tbl) return;
+
+    st_foreach(tbl, mark_value, 0);
+}
+
 static int
 mark_keyvalue(st_data_t key, st_data_t value, st_data_t data)
 {
@@ -5512,6 +5606,7 @@ Init_GC(void)
 {
 #undef rb_intern
     rb_gc_register_address(&id2ref_value);
+    rb_gc_register_address(&obj_to_id_value);
 
     malloc_offset = gc_compute_malloc_offset();
 

gc/gc.h

@@ -135,6 +135,34 @@ gc_ref_update_table_values_only(st_table *tbl)
     }
 }
 
+static int
+hash_foreach_replace_key(st_data_t key, st_data_t value, st_data_t argp, int error)
+{
+    if (rb_gc_location((VALUE)key) != (VALUE)key) {
+        return ST_REPLACE;
+    }
+    return ST_CONTINUE;
+}
+
+static int
+hash_replace_ref_key(st_data_t *key, st_data_t *value, st_data_t argp, int existing)
+{
+    *key = rb_gc_location((VALUE)*key);
+
+    return ST_CONTINUE;
+}
+
+static void
+gc_ref_update_table_keys_only(st_table *tbl)
+{
+    if (!tbl || tbl->num_entries == 0) return;
+
+    // FIXME: this certainly isn't correct. If a key moved, we need to re-hash.
+    if (st_foreach_with_replace(tbl, hash_foreach_replace_key, hash_replace_ref_key, 0)) {
+        rb_raise(rb_eRuntimeError, "hash modified during iteration");
+    }
+}
+
 static int
 gc_mark_tbl_no_pin_i(st_data_t key, st_data_t value, st_data_t data)
 {

shape.c

@@ -47,6 +47,7 @@
 
 static ID id_frozen;
 static ID id_t_object;
+static ID id_external_object_id;
 ID ruby_internal_object_id; // extern
 
 #define LEAF 0
@@ -469,6 +470,7 @@ rb_shape_alloc_new_child(ID id, rb_shape_t *shape, enum shape_type shape_type)
 
     switch (shape_type) {
       case SHAPE_OBJ_ID:
+      case SHAPE_EXTERNAL_OBJ_ID:
         new_shape->flags |= SHAPE_FL_HAS_OBJECT_ID;
         // fallthrough
       case SHAPE_IVAR:
@@ -745,22 +747,57 @@ rb_shape_has_object_id(rb_shape_t *shape)
     return shape->flags & SHAPE_FL_HAS_OBJECT_ID;
 }
 
+shape_id_t
+rb_shape_object_id(VALUE obj)
+{
+    rb_shape_t *shape = rb_obj_shape(obj);
+    RUBY_ASSERT(shape);
+    RUBY_ASSERT(shape->flags & SHAPE_FL_HAS_OBJECT_ID);
+
+    while (shape->type != SHAPE_OBJ_ID && shape->type != SHAPE_EXTERNAL_OBJ_ID) {
+        shape = RSHAPE(shape->parent_id);
+    }
+    return rb_shape_id(shape);
+}
+
 shape_id_t
 rb_shape_transition_object_id(VALUE obj)
 {
-    rb_shape_t* shape = rb_obj_shape(obj);
+    rb_shape_t *shape = rb_obj_shape(obj);
     RUBY_ASSERT(shape);
 
     if (shape->flags & SHAPE_FL_HAS_OBJECT_ID) {
         while (shape->type != SHAPE_OBJ_ID) {
+            RUBY_ASSERT(shape->type != SHAPE_EXTERNAL_OBJ_ID);
             shape = RSHAPE(shape->parent_id);
         }
     }
     else {
         bool dont_care;
         shape = get_next_shape_internal(shape, ruby_internal_object_id, SHAPE_OBJ_ID, &dont_care, true);
     }
+
+    RUBY_ASSERT(shape && shape->type == SHAPE_OBJ_ID);
+    return rb_shape_id(shape);
+}
+
+shape_id_t
+rb_shape_transition_external_object_id(VALUE obj)
+{
+    rb_shape_t* shape = rb_obj_shape(obj);
     RUBY_ASSERT(shape);
+
+    if (shape->flags & SHAPE_FL_HAS_OBJECT_ID) {
+        while (shape->type != SHAPE_EXTERNAL_OBJ_ID) {
+            RUBY_ASSERT(shape->type != SHAPE_OBJ_ID);
+            shape = RSHAPE(shape->parent_id);
+        }
+    }
+    else {
+        bool dont_care;
+        shape = get_next_shape_internal(shape, id_external_object_id, SHAPE_EXTERNAL_OBJ_ID, &dont_care, true);
+    }
+    RUBY_ASSERT(shape && shape->type == SHAPE_EXTERNAL_OBJ_ID);
     return rb_shape_id(shape);
 }
 
@@ -924,6 +961,7 @@ shape_get_iv_index(rb_shape_t *shape, ID id, attr_index_t *value)
                 return false;
               case SHAPE_OBJ_TOO_COMPLEX:
               case SHAPE_OBJ_ID:
+              case SHAPE_EXTERNAL_OBJ_ID:
               case SHAPE_FROZEN:
                 rb_bug("Ivar should not exist on transition");
             }
@@ -1009,6 +1047,7 @@ shape_traverse_from_new_root(rb_shape_t *initial_shape, rb_shape_t *dest_shape)
     switch ((enum shape_type)dest_shape->type) {
       case SHAPE_IVAR:
       case SHAPE_OBJ_ID:
+      case SHAPE_EXTERNAL_OBJ_ID:
       case SHAPE_FROZEN:
         if (!next_shape->edges) {
             return NULL;
@@ -1080,6 +1119,7 @@ rb_shape_rebuild_shape(rb_shape_t *initial_shape, rb_shape_t *dest_shape)
         midway_shape = shape_get_next_iv_shape(midway_shape, dest_shape->edge_name);
         break;
       case SHAPE_OBJ_ID:
+      case SHAPE_EXTERNAL_OBJ_ID:
       case SHAPE_ROOT:
       case SHAPE_FROZEN:
       case SHAPE_T_OBJECT:
@@ -1368,6 +1408,7 @@ Init_default_shapes(void)
 
     id_frozen = rb_make_internal_id();
     id_t_object = rb_make_internal_id();
+    id_external_object_id = rb_make_internal_id();
     ruby_internal_object_id = rb_make_internal_id();
 
 #ifdef HAVE_MMAP

shape.h

@@ -67,6 +67,7 @@ enum shape_type {
     SHAPE_ROOT,
     SHAPE_IVAR,
     SHAPE_OBJ_ID,
+    SHAPE_EXTERNAL_OBJ_ID,
     SHAPE_FROZEN,
     SHAPE_T_OBJECT,
     SHAPE_OBJ_TOO_COMPLEX,
@@ -169,6 +170,9 @@ bool rb_shape_transition_remove_ivar(VALUE obj, ID id, VALUE *removed);
 shape_id_t rb_shape_transition_add_ivar(VALUE obj, ID id);
 shape_id_t rb_shape_transition_add_ivar_no_warnings(VALUE obj, ID id);
 shape_id_t rb_shape_transition_object_id(VALUE obj);
+shape_id_t rb_shape_transition_external_object_id(VALUE obj);
+shape_id_t rb_shape_object_id(VALUE obj);
+
 
 bool rb_shape_has_object_id(rb_shape_t *shape);
 void rb_shape_free_all(void);

test/ruby/test_object_id.rb

@@ -195,3 +195,34 @@ def setup
     end
   end
 end
+
+class TestObjectIdRactor < Test::Unit::TestCase
+  def test_object_id_race_free
+    assert_separately([], "#{<<~"begin;"}\n#{<<~'end;'}")
+    begin;
+      Warning[:experimental] = false
+      class MyClass
+        attr_reader :a, :b, :c
+        def initialize
+          @a = @b = @c = nil
+        end
+      end
+
+      N = 10_000
+      objs = Ractor.make_shareable(N.times.map { MyClass.new })
+
+      results = 4.times.map{
+        Ractor.new(objs) { |objs|
+          vars = []
+          ids = []
+          objs.each do |obj|
+            vars << obj.a << obj.b << obj.c
+            ids << obj.object_id
+          end
+          [vars, ids]
+        }
+      }.map(&:take)
+      assert_equal 1, results.uniq.size
+    end;
+  end
+end