MJIT: Cancel all on disastrous situations (#7019)

I noticed this while running test_yjit with --mjit-call-threshold=1, 
which redefines `Integer#<`. When Ruby is monkey-patched, 
MJIT itself could be broken.

Similarly, Ruby scripts could break MJIT in many different ways. I
prepared the same set of hooks as YJIT so that we could possibly
override it and disable it on those moments. Every constant under
RubyVM::MJIT is private and thus it's an unsupported behavior though.

.github/workflows/mjit.yml

@@ -69,10 +69,9 @@ jobs:
       - run: make incs
       - run: make
       - run: sudo make -s install
-      - run: sudo apt-get install gdb # used by test / test-all failure
       - name: Run test
         run: |
-          ulimit -c unlimited
+          unset GNUMAKEFLAGS
           make -s test RUN_OPTS="$RUN_OPTS"
         timeout-minutes: 60
       # - name: Run test-all
@@ -82,7 +81,7 @@ jobs:
       #   timeout-minutes: 60
       - name: Run test-spec
         run: |
-          ulimit -c unlimited
+          unset GNUMAKEFLAGS
           make -s test-spec RUN_OPTS="$RUN_OPTS"
         timeout-minutes: 60
       - uses: ruby/action-slack@b6882ea6ef8f556f9f9af9ec1220d3f1ced74acf # v3.0.0

bootstraptest/test_yjit.rb

@@ -2173,7 +2173,7 @@ assert_equal '[[:c_return, :String, :string_alias, "events_to_str"]]', %q{
   events.compiled(events)
 
   events
-}
+} unless defined?(RubyVM::MJIT) && RubyVM::MJIT.enabled? # MJIT calls extra Ruby methods
 
 # test enabling a TracePoint that targets a particular line in a C method call
 assert_equal '[true]', %q{
@@ -2255,7 +2255,7 @@ assert_equal '[[:c_call, :itself]]', %q{
   tp.enable { shouldnt_compile }
 
   events
-}
+} unless defined?(RubyVM::MJIT) && RubyVM::MJIT.enabled? # MJIT calls extra Ruby methods
 
 # test enabling c_return tracing before compiling
 assert_equal '[[:c_return, :itself, main]]', %q{
@@ -2270,7 +2270,7 @@ assert_equal '[[:c_return, :itself, main]]', %q{
   tp.enable { shouldnt_compile }
 
   events
-}
+} unless defined?(RubyVM::MJIT) && RubyVM::MJIT.enabled? # MJIT calls extra Ruby methods
 
 # test c_call invalidation
 assert_equal '[[:c_call, :itself]]', %q{

common.mk

@@ -11515,6 +11515,7 @@ ractor.$(OBJEXT): $(CCAN_DIR)/check_type/check_type.h
 ractor.$(OBJEXT): $(CCAN_DIR)/container_of/container_of.h
 ractor.$(OBJEXT): $(CCAN_DIR)/list/list.h
 ractor.$(OBJEXT): $(CCAN_DIR)/str/str.h
+ractor.$(OBJEXT): $(hdrdir)/ruby.h
 ractor.$(OBJEXT): $(hdrdir)/ruby/ruby.h
 ractor.$(OBJEXT): $(top_srcdir)/internal/array.h
 ractor.$(OBJEXT): $(top_srcdir)/internal/basic_operators.h
@@ -11710,6 +11711,7 @@ ractor.$(OBJEXT): {$(VPATH)}internal/warning_push.h
 ractor.$(OBJEXT): {$(VPATH)}internal/xmalloc.h
 ractor.$(OBJEXT): {$(VPATH)}method.h
 ractor.$(OBJEXT): {$(VPATH)}missing.h
+ractor.$(OBJEXT): {$(VPATH)}mjit.h
 ractor.$(OBJEXT): {$(VPATH)}node.h
 ractor.$(OBJEXT): {$(VPATH)}onigmo.h
 ractor.$(OBJEXT): {$(VPATH)}oniguruma.h

lib/ruby_vm/mjit/hooks.rb

@@ -0,0 +1,32 @@
+module RubyVM::MJIT::Hooks # :nodoc: all
+  C = RubyVM::MJIT.const_get(:C, false)
+
+  def self.on_bop_redefined(_redefined_flag, _bop)
+    C.mjit_cancel_all("BOP is redefined")
+  end
+
+  def self.on_cme_invalidate(_cme)
+    # to be used later
+  end
+
+  def self.on_ractor_spawn
+    C.mjit_cancel_all("Ractor is spawned")
+  end
+
+  def self.on_constant_state_changed(_id)
+    # to be used later
+  end
+
+  def self.on_constant_ic_update(_iseq, _ic, _insn_idx)
+    # to be used later
+  end
+
+  def self.on_tracing_invalidate_all(new_iseq_events)
+    # Stop calling all JIT-ed code. We can't rewrite existing JIT-ed code to trace_ insns for now.
+    # :class events are triggered only in ISEQ_TYPE_CLASS, but mjit_target_iseq_p ignores such iseqs.
+    # Thus we don't need to cancel JIT-ed code for :class events.
+    if new_iseq_events != C.RUBY_EVENT_CLASS
+      C.mjit_cancel_all("TracePoint is enabled")
+    end
+  end
+end

mjit.c

@@ -142,6 +142,8 @@ bool mjit_enabled = false;
 // true if JIT-ed code should be called. When `ruby_vm_event_enabled_global_flags & ISEQ_TRACE_EVENTS`
 // and `mjit_call_p == false`, any JIT-ed code execution is cancelled as soon as possible.
 bool mjit_call_p = false;
+// A flag to communicate that mjit_call_p should be disabled while it's temporarily false.
+bool mjit_cancel_p = false;
 // There's an ISEQ in unit_queue whose total_calls reached 2 * call_threshold.
 // If this is true, check_unit_queue will start compiling ISEQs in unit_queue.
 static bool mjit_compile_p = false;
@@ -963,6 +965,7 @@ mjit_cancel_all(const char *reason)
         return;
 
     mjit_call_p = false;
+    mjit_cancel_p = true;
     if (mjit_opts.warnings || mjit_opts.verbose) {
         fprintf(stderr, "JIT cancel: Disabled JIT-ed code because %s\n", reason);
     }
@@ -1220,31 +1223,115 @@ static VALUE rb_mMJITC = 0;
 static VALUE rb_cMJITCompiler = 0;
 // RubyVM::MJIT::CPointer::Struct_rb_iseq_t
 static VALUE rb_cMJITIseqPtr = 0;
+// RubyVM::MJIT::CPointer::Struct_IC
+static VALUE rb_cMJITICPtr = 0;
+// RubyVM::MJIT::Compiler
+static VALUE rb_mMJITHooks = 0;
+
+#define WITH_MJIT_DISABLED(stmt) do { \
+    bool original_call_p = mjit_call_p; \
+    mjit_call_p = false; \
+    stmt; \
+    mjit_call_p = original_call_p; \
+    if (mjit_cancel_p) mjit_call_p = false; \
+} while (0);
+
+// Hook MJIT when BOP is redefined.
+MJIT_FUNC_EXPORTED void
+rb_mjit_bop_redefined(int redefined_flag, enum ruby_basic_operators bop)
+{
+    if (!mjit_enabled || !mjit_call_p || !rb_mMJITHooks) return;
+    WITH_MJIT_DISABLED({
+        rb_funcall(rb_mMJITHooks, rb_intern("on_bop_redefined"), 2, INT2NUM(redefined_flag), INT2NUM((int)bop));
+    });
+}
+
+// Hook MJIT when CME is invalidated.
+MJIT_FUNC_EXPORTED void
+rb_mjit_cme_invalidate(rb_callable_method_entry_t *cme)
+{
+    if (!mjit_enabled || !mjit_call_p || !rb_mMJITHooks) return;
+    WITH_MJIT_DISABLED({
+        VALUE cme_klass = rb_funcall(rb_mMJITC, rb_intern("rb_callable_method_entry_struct"), 0);
+        VALUE cme_ptr = rb_funcall(cme_klass, rb_intern("new"), 1, SIZET2NUM((size_t)cme));
+        rb_funcall(rb_mMJITHooks, rb_intern("on_cme_invalidate"), 1, cme_ptr);
+    });
+}
+
+// Hook MJIT when Ractor is spawned.
+void
+rb_mjit_before_ractor_spawn(void)
+{
+    if (!mjit_enabled || !mjit_call_p || !rb_mMJITHooks) return;
+    WITH_MJIT_DISABLED({
+        rb_funcall(rb_mMJITHooks, rb_intern("on_ractor_spawn"), 0);
+    });
+}
+
+static void
+mjit_constant_state_changed(void *data)
+{
+    if (!mjit_enabled || !mjit_call_p || !rb_mMJITHooks) return;
+    ID id = (ID)data;
+    WITH_MJIT_DISABLED({
+        rb_funcall(rb_mMJITHooks, rb_intern("on_constant_state_changed"), 1, ID2SYM(id));
+    });
+}
+
+// Hook MJIT when constant state is changed.
+MJIT_FUNC_EXPORTED void
+rb_mjit_constant_state_changed(ID id)
+{
+    if (!mjit_enabled || !mjit_call_p || !rb_mMJITHooks) return;
+    // Asynchronously hook the Ruby code since this is hooked during a "Ruby critical section".
+    extern int rb_workqueue_register(unsigned flags, rb_postponed_job_func_t func, void *data);
+    rb_workqueue_register(0, mjit_constant_state_changed, (void *)id);
+}
+
+// Hook MJIT when constant IC is updated.
+MJIT_FUNC_EXPORTED void
+rb_mjit_constant_ic_update(const rb_iseq_t *const iseq, IC ic, unsigned insn_idx)
+{
+    if (!mjit_enabled || !mjit_call_p || !rb_mMJITHooks) return;
+    WITH_MJIT_DISABLED({
+        VALUE iseq_ptr = rb_funcall(rb_cMJITIseqPtr, rb_intern("new"), 1, SIZET2NUM((size_t)iseq));
+        VALUE ic_ptr = rb_funcall(rb_cMJITICPtr, rb_intern("new"), 1, SIZET2NUM((size_t)ic));
+        rb_funcall(rb_mMJITHooks, rb_intern("on_constant_ic_update"), 3, iseq_ptr, ic_ptr, UINT2NUM(insn_idx));
+    });
+}
+
+// Hook MJIT when TracePoint is enabled.
+MJIT_FUNC_EXPORTED void
+rb_mjit_tracing_invalidate_all(rb_event_flag_t new_iseq_events)
+{
+    if (!mjit_enabled || !mjit_call_p || !rb_mMJITHooks) return;
+    WITH_MJIT_DISABLED({
+        rb_funcall(rb_mMJITHooks, rb_intern("on_tracing_invalidate_all"), 1, UINT2NUM(new_iseq_events));
+    });
+}
 
 // [experimental] Call custom RubyVM::MJIT.compile if defined
 static void
 mjit_hook_custom_compile(const rb_iseq_t *iseq)
 {
-    bool original_call_p = mjit_call_p;
-    mjit_call_p = false; // Avoid impacting JIT metrics by itself
-
-    VALUE iseq_class = rb_funcall(rb_mMJITC, rb_intern("rb_iseq_t"), 0);
-    VALUE iseq_ptr = rb_funcall(iseq_class, rb_intern("new"), 1, ULONG2NUM((size_t)iseq));
-    VALUE jit_func = rb_funcall(rb_mMJIT, rb_intern("compile"), 1, iseq_ptr);
-    ISEQ_BODY(iseq)->jit_func = (jit_func_t)NUM2ULONG(jit_func);
-
-    mjit_call_p = original_call_p;
+    WITH_MJIT_DISABLED({
+        VALUE iseq_class = rb_funcall(rb_mMJITC, rb_intern("rb_iseq_t"), 0);
+        VALUE iseq_ptr = rb_funcall(iseq_class, rb_intern("new"), 1, ULONG2NUM((size_t)iseq));
+        VALUE jit_func = rb_funcall(rb_mMJIT, rb_intern("compile"), 1, iseq_ptr);
+        ISEQ_BODY(iseq)->jit_func = (jit_func_t)NUM2ULONG(jit_func);
+    });
 }
 
 static void
 mjit_add_iseq_to_process(const rb_iseq_t *iseq, const struct rb_mjit_compile_info *compile_info)
 {
-    if (!mjit_enabled || pch_status != PCH_SUCCESS || !rb_ractor_main_p()) // TODO: Support non-main Ractors
-        return;
+    if (!mjit_enabled) return;
     if (mjit_opts.custom) { // Hook custom RubyVM::MJIT.compile if defined
         mjit_hook_custom_compile(iseq);
         return;
     }
+    if (pch_status != PCH_SUCCESS || !rb_ractor_main_p()) // TODO: Support non-main Ractors
+        return;
     if (!mjit_target_iseq_p(iseq)) {
         ISEQ_BODY(iseq)->jit_func = (jit_func_t)MJIT_FUNC_FAILED; // skip mjit_wait
         return;
@@ -1646,6 +1733,9 @@ mjit_init(const struct mjit_options *opts)
     rb_mMJITC = rb_const_get(rb_mMJIT, rb_intern("C"));
     rb_cMJITCompiler = rb_funcall(rb_const_get(rb_mMJIT, rb_intern("Compiler")), rb_intern("new"), 0);
     rb_cMJITIseqPtr = rb_funcall(rb_mMJITC, rb_intern("rb_iseq_t"), 0);
+    rb_cMJITICPtr = rb_funcall(rb_mMJITC, rb_intern("IC"), 0);
+    rb_funcall(rb_cMJITICPtr, rb_intern("new"), 1, SIZET2NUM(0)); // Trigger no-op constant events before enabling hooks
+    rb_mMJITHooks = rb_const_get(rb_mMJIT, rb_intern("Hooks"));
 
     mjit_call_p = true;
     mjit_pid = getpid();
@@ -1845,6 +1935,8 @@ mjit_mark(void)
     // Mark objects used by the MJIT compiler
     rb_gc_mark(rb_cMJITCompiler);
     rb_gc_mark(rb_cMJITIseqPtr);
+    rb_gc_mark(rb_cMJITICPtr);
+    rb_gc_mark(rb_mMJITHooks);
 
     // Mark JIT-compiled ISEQs
     struct rb_mjit_unit *unit = NULL;

mjit.h

@@ -103,6 +103,13 @@ extern void mjit_mark(void);
 extern void mjit_mark_cc_entries(const struct rb_iseq_constant_body *const body);
 extern void mjit_notify_waitpid(int exit_code);
 
+extern void rb_mjit_bop_redefined(int redefined_flag, enum ruby_basic_operators bop);
+extern void rb_mjit_cme_invalidate(rb_callable_method_entry_t *cme);
+extern void rb_mjit_before_ractor_spawn(void);
+extern void rb_mjit_constant_state_changed(ID id);
+extern void rb_mjit_constant_ic_update(const rb_iseq_t *const iseq, IC ic, unsigned insn_idx);
+extern void rb_mjit_tracing_invalidate_all(rb_event_flag_t new_iseq_events);
+
 void mjit_child_after_fork(void);
 
 #  ifdef MJIT_HEADER
@@ -122,6 +129,13 @@ static inline void mjit_mark(void){}
 static inline VALUE jit_exec(rb_execution_context_t *ec) { return Qundef; /* unreachable */ }
 static inline void mjit_child_after_fork(void){}
 
+static inline void rb_mjit_bop_redefined(int redefined_flag, enum ruby_basic_operators bop) {}
+static inline void rb_mjit_cme_invalidate(rb_callable_method_entry_t *cme) {}
+static inline void rb_mjit_before_ractor_spawn(void) {}
+static inline void rb_mjit_constant_state_changed(ID id) {}
+static inline void rb_mjit_constant_ic_update(const rb_iseq_t *const iseq, IC ic, unsigned insn_idx) {}
+static inline void rb_mjit_tracing_invalidate_all(rb_event_flag_t new_iseq_events) {}
+
 #define mjit_enabled false
 static inline VALUE mjit_pause(bool wait_p){ return Qnil; } // unreachable
 static inline VALUE mjit_resume(void){ return Qnil; } // unreachable

mjit.rb

@@ -29,6 +29,7 @@ if RubyVM::MJIT.enabled?
   require 'ruby_vm/mjit/c_type'
   require 'ruby_vm/mjit/instruction'
   require 'ruby_vm/mjit/compiler'
+  require 'ruby_vm/mjit/hooks'
 
   module RubyVM::MJIT
     private_constant(*constants)

mjit_c.rb

@@ -109,6 +109,13 @@ module RubyVM::MJIT # :nodoc: all
       }
     end
 
+    def mjit_cancel_all(reason)
+      Primitive.cstmt! %{
+        mjit_cancel_all(RSTRING_PTR(reason));
+        return Qnil;
+      }
+    end
+
     # Convert encoded VM pointers to insn BINs.
     def rb_vm_insn_decode(encoded)
       Primitive.cexpr! 'INT2NUM(rb_vm_insn_decode(NUM2PTR(encoded)))'
@@ -170,6 +177,10 @@ module RubyVM::MJIT # :nodoc: all
     Primitive.cexpr! %q{ INT2NUM(VM_METHOD_TYPE_ISEQ) }
   end
 
+  def C.RUBY_EVENT_CLASS
+    Primitive.cexpr! %q{ UINT2NUM(RUBY_EVENT_CLASS) }
+  end
+
   def C.SHAPE_CAPACITY_CHANGE
     Primitive.cexpr! %q{ UINT2NUM(SHAPE_CAPACITY_CHANGE) }
   end

ractor.c

@@ -17,6 +17,7 @@
 #include "gc.h"
 #include "transient_heap.h"
 #include "yjit.h"
+#include "mjit.h"
 
 VALUE rb_cRactor;
 
@@ -1606,6 +1607,7 @@ ractor_create(rb_execution_context_t *ec, VALUE self, VALUE loc, VALUE name, VAL
     r->debug = cr->debug;
 
     rb_yjit_before_ractor_spawn();
+    rb_mjit_before_ractor_spawn();
     rb_thread_create_ractor(r, args, block);
 
     RB_GC_GUARD(rv);

test/ruby/test_mjit.rb

@@ -1177,6 +1177,18 @@ class TestMJIT < Test::Unit::TestCase
     end;
   end
 
+  def test_cancel_by_bop_redefinition
+    assert_eval_with_jit("#{<<~"begin;"}\n#{<<~"end;"}", success_count: 0, call_threshold: 2)
+    begin;
+      class Integer
+        def <(x)
+          true
+        end
+      end
+      2.times {}
+    end;
+  end
+
   def test_caller_locations_without_catch_table
     out, _ = eval_with_jit("#{<<~"begin;"}\n#{<<~"end;"}", call_threshold: 2)
     begin;

tool/mjit/bindgen.rb

@@ -353,10 +353,11 @@ generator = BindingGenerator.new(
       VM_METHOD_TYPE_ISEQ
     ],
     UINT: %w[
-      SHAPE_ID_NUM_BITS
+      RUBY_EVENT_CLASS
       SHAPE_CAPACITY_CHANGE
       SHAPE_FLAG_SHIFT
       SHAPE_FROZEN
+      SHAPE_ID_NUM_BITS
       SHAPE_INITIAL_CAPACITY
       SHAPE_IVAR
       SHAPE_ROOT

vm.c

@@ -1975,6 +1975,7 @@ rb_vm_check_redefinition_opt_method(const rb_method_entry_t *me, VALUE klass)
             int flag = vm_redefinition_check_flag(klass);
             if (flag != 0) {
                 rb_yjit_bop_redefined(flag, (enum ruby_basic_operators)bop);
+                rb_mjit_bop_redefined(flag, (enum ruby_basic_operators)bop);
                 ruby_vm_redefined_flag[bop] |= flag;
             }
         }

vm_insnhelper.c

@@ -5324,13 +5324,11 @@ vm_ic_update(const rb_iseq_t *iseq, IC ic, VALUE val, const VALUE *reg_ep, const
     ice->ic_cref = vm_get_const_key_cref(reg_ep);
     if (rb_ractor_shareable_p(val)) ice->flags |= IMEMO_CONST_CACHE_SHAREABLE;
     RB_OBJ_WRITE(iseq, &ic->entry, ice);
-#ifndef MJIT_HEADER
-    // MJIT and YJIT can't be on at the same time, so there is no need to
-    // notify YJIT about changes to the IC when running inside MJIT code.
+
     RUBY_ASSERT(pc >= ISEQ_BODY(iseq)->iseq_encoded);
     unsigned pos = (unsigned)(pc - ISEQ_BODY(iseq)->iseq_encoded);
     rb_yjit_constant_ic_update(iseq, ic, pos);
-#endif
+    rb_mjit_constant_ic_update(iseq, ic, pos);
 }
 
 static VALUE

vm_method.c

@@ -4,6 +4,7 @@
 
 #include "id_table.h"
 #include "yjit.h"
+#include "mjit.h"
 
 #define METHOD_DEBUG 0
 
@@ -124,6 +125,7 @@ vm_cme_invalidate(rb_callable_method_entry_t *cme)
     RB_DEBUG_COUNTER_INC(cc_cme_invalidate);
 
     rb_yjit_cme_invalidate(cme);
+    rb_mjit_cme_invalidate(cme);
 }
 
 static int
@@ -149,6 +151,7 @@ rb_clear_constant_cache_for_id(ID id)
     }
 
     rb_yjit_constant_state_changed(id);
+    rb_mjit_constant_state_changed(id);
 }
 
 static void
@@ -188,6 +191,7 @@ clear_method_cache_by_id_in_class(VALUE klass, ID mid)
         if (cc_tbl && rb_id_table_lookup(cc_tbl, mid, &ccs_data)) {
             struct rb_class_cc_entries *ccs = (struct rb_class_cc_entries *)ccs_data;
             rb_yjit_cme_invalidate((rb_callable_method_entry_t *)ccs->cme);
+            rb_mjit_cme_invalidate((rb_callable_method_entry_t *)ccs->cme);
             if (NIL_P(ccs->cme->owner)) invalidate_negative_cache(mid);
             rb_vm_ccs_free(ccs);
             rb_id_table_delete(cc_tbl, mid);
@@ -200,6 +204,7 @@ clear_method_cache_by_id_in_class(VALUE klass, ID mid)
             VALUE cme;
             if (rb_yjit_enabled_p() && rb_id_table_lookup(cm_tbl, mid, &cme)) {
                 rb_yjit_cme_invalidate((rb_callable_method_entry_t *)cme);
+                rb_mjit_cme_invalidate((rb_callable_method_entry_t *)cme);
             }
             rb_id_table_delete(cm_tbl, mid);
             RB_DEBUG_COUNTER_INC(cc_invalidate_leaf_callable);

vm_trace.c

@@ -106,12 +106,6 @@ update_global_event_hook(rb_event_flag_t prev_events, rb_event_flag_t new_events
     rb_objspace_set_event_hook(new_events);
 
     // Invalidate JIT code as needed
-    if (first_time_iseq_events_p && new_iseq_events != RUBY_EVENT_CLASS) {
-        // Stop calling all JIT-ed code. We can't rewrite existing JIT-ed code to trace_ insns for now.
-        // :class events are triggered only in ISEQ_TYPE_CLASS, but mjit_target_iseq_p ignores such iseqs.
-        // Thus we don't need to cancel JIT-ed code for :class events.
-        mjit_cancel_all("TracePoint is enabled");
-    }
     if (first_time_iseq_events_p || enable_c_call || enable_c_return) {
         // Invalidate all code when ISEQs are modified to use trace_* insns above.
         // Also invalidate when enabling c_call or c_return because generated code
@@ -120,6 +114,7 @@ update_global_event_hook(rb_event_flag_t prev_events, rb_event_flag_t new_events
         // Do this after event flags updates so other ractors see updated vm events
         // when they wake up.
         rb_yjit_tracing_invalidate_all();
+        rb_mjit_tracing_invalidate_all(new_iseq_events);
     }
 }
 
@@ -1257,6 +1252,7 @@ rb_tracepoint_enable_for_target(VALUE tpval, VALUE target, VALUE target_line)
     }
 
     rb_yjit_tracing_invalidate_all();
+    rb_mjit_tracing_invalidate_all(tp->events);
 
     ruby_vm_event_local_num++;