Make inline cache reads / writes atomic with object shapes

Prior to this commit, we were reading and writing ivar index and
shape ID in inline caches in two separate instructions when
getting and setting ivars. This meant there was a race condition
with ractors and these caches where one ractor could change
a value in the cache while another was still reading from it.

This commit instead reads and writes shape ID and ivar index to
inline caches atomically so there is no longer a race condition.

Co-Authored-By: Aaron Patterson <tenderlove@ruby-lang.org>
Co-Authored-By: John Hawthorn <john@hawthorn.email>

1 files changed, 21 insertions, 9 deletions

diff --git a/lib/mjit/compiler.rb b/lib/mjit/compiler.rb
index 06f018c934..0d2910bf4e 100644
--- a/lib/mjit/compiler.rb
+++ b/lib/mjit/compiler.rb
@@ -351,20 +351,27 @@ module RubyVM::MJIT
     # _mjit_compile_ivar.erb
     def compile_ivar(insn_name, stack_size, pos, status, operands, body)
       ic_copy = (status.is_entries + (C.iseq_inline_storage_entry.new(operands[1]) - body.is_entries)).iv_cache
+      dest_shape_id = ic_copy.value >> C.SHAPE_FLAG_SHIFT
+      attr_index = ic_copy.value & ((1 << C.SHAPE_FLAG_SHIFT) - 1)
+      source_shape_id = if dest_shape_id == C.INVALID_SHAPE_ID
+                          dest_shape_id
+                        else
+                          RubyVM::Shape.find_by_id(dest_shape_id).parent_id
+                        end
 
       src = +''
-      if !status.compile_info.disable_ivar_cache && ic_copy.source_shape_id != C.INVALID_SHAPE_ID
+      if !status.compile_info.disable_ivar_cache && source_shape_id != C.INVALID_SHAPE_ID
         # JIT: optimize away motion of sp and pc. This path does not call rb_warning() and so it's always leaf and not `handles_sp`.
         # compile_pc_and_sp(src, insn, stack_size, sp_inc, local_stack_p, next_pos)
 
         # JIT: prepare vm_getivar/vm_setivar arguments and variables
         src << "{\n"
         src << "    VALUE obj = GET_SELF();\n"
-        src << "    const shape_id_t source_shape_id = (rb_serial_t)#{ic_copy.source_shape_id};\n"
         # JIT: cache hit path of vm_getivar/vm_setivar, or cancel JIT (recompile it with exivar)
         if insn_name == :setinstancevariable
-          src << "    const uint32_t index = #{ic_copy.attr_index - 1};\n"
-          src << "    const shape_id_t dest_shape_id = (rb_serial_t)#{ic_copy.dest_shape_id};\n"
+          src << "    const shape_id_t source_shape_id = (shape_id_t)#{source_shape_id};\n"
+          src << "    const uint32_t index = #{attr_index - 1};\n"
+          src << "    const shape_id_t dest_shape_id = (shape_id_t)#{dest_shape_id};\n"
           src << "    if (source_shape_id == ROBJECT_SHAPE_ID(obj) && \n"
           src << "        dest_shape_id != ROBJECT_SHAPE_ID(obj)) {\n"
           src << "        if (UNLIKELY(index >= ROBJECT_NUMIV(obj))) {\n"
@@ -374,14 +381,19 @@ module RubyVM::MJIT
           src << "        VALUE *ptr = ROBJECT_IVPTR(obj);\n"
           src << "        RB_OBJ_WRITE(obj, &ptr[index], stack[#{stack_size - 1}]);\n"
           src << "    }\n"
+          src << "    else if (dest_shape_id == ROBJECT_SHAPE_ID(obj)) {\n"
+          src << "        VALUE *ptr = ROBJECT_IVPTR(obj);\n"
+          src << "        RB_OBJ_WRITE(obj, &ptr[index], stack[#{stack_size - 1}]);\n"
+          src << "    }\n"
         else
-          if ic_copy.attr_index == 0 # cache hit, but uninitialized iv
+          src << "    const shape_id_t source_shape_id = (shape_id_t)#{dest_shape_id};\n"
+          if attr_index == 0 # cache hit, but uninitialized iv
             src << "    /* Uninitialized instance variable */\n"
             src << "    if (source_shape_id == ROBJECT_SHAPE_ID(obj)) {\n"
             src << "        stack[#{stack_size}] = Qnil;\n"
             src << "    }\n"
           else
-            src << "    const uint32_t index = #{ic_copy.attr_index - 1};\n"
+            src << "    const uint32_t index = #{attr_index - 1};\n"
             src << "    if (source_shape_id == ROBJECT_SHAPE_ID(obj)) {\n"
             src << "        stack[#{stack_size}] = ROBJECT_IVPTR(obj)[index];\n"
             src << "    }\n"
@@ -394,15 +406,15 @@ module RubyVM::MJIT
         src << "    }\n"
         src << "}\n"
         return src
-      elsif insn_name == :getinstancevariable && !status.compile_info.disable_exivar_cache && ic_copy.source_shape_id != C.INVALID_SHAPE_ID
+      elsif insn_name == :getinstancevariable && !status.compile_info.disable_exivar_cache && source_shape_id != C.INVALID_SHAPE_ID
         # JIT: optimize away motion of sp and pc. This path does not call rb_warning() and so it's always leaf and not `handles_sp`.
         # compile_pc_and_sp(src, insn, stack_size, sp_inc, local_stack_p, next_pos)
 
         # JIT: prepare vm_getivar's arguments and variables
         src << "{\n"
         src << "    VALUE obj = GET_SELF();\n"
-        src << "    const shape_id_t source_shape_id = (rb_serial_t)#{ic_copy.source_shape_id};\n"
-        src << "    const uint32_t index = #{ic_copy.attr_index - 1};\n"
+        src << "    const shape_id_t source_shape_id = (shape_id_t)#{dest_shape_id};\n"
+        src << "    const uint32_t index = #{attr_index - 1};\n"
         # JIT: cache hit path of vm_getivar, or cancel JIT (recompile it without any ivar optimization)
         src << "    struct gen_ivtbl *ivtbl;\n"
         src << "    if (LIKELY(FL_TEST_RAW(obj, FL_EXIVAR) && source_shape_id == rb_shape_get_shape_id(obj) && rb_ivar_generic_ivtbl_lookup(obj, &ivtbl))) {\n"