YJIT: Invalidate JIT code only for ISEQ_TRACE_EVENTS (#6695)

2 files changed, 54 insertions, 4 deletions

diff --git a/test/ruby/test_yjit.rb b/test/ruby/test_yjit.rb
index 1a564889af..fab8768a7e 100644
--- a/test/ruby/test_yjit.rb
+++ b/test/ruby/test_yjit.rb
@@ -918,6 +918,54 @@ class TestYJIT < Test::Unit::TestCase
     RUBY
   end
 
+  def test_trace_script_compiled # not ISEQ_TRACE_EVENTS
+    assert_compiles(<<~'RUBY', exits: :any, result: :ok)
+      @eval_counter = 0
+      def eval_script
+        eval('@eval_counter += 1')
+      end
+
+      @trace_counter = 0
+      trace = TracePoint.new(:script_compiled) do |t|
+        @trace_counter += 1
+      end
+
+      eval_script # JIT without TracePoint
+      trace.enable
+      eval_script # call with TracePoint
+      trace.disable
+
+      return :"eval_#{@eval_counter}" if @eval_counter != 2
+      return :"trace_#{@trace_counter}" if @trace_counter != 1
+
+      :ok
+    RUBY
+  end
+
+  def test_trace_b_call # ISEQ_TRACE_EVENTS
+    assert_compiles(<<~'RUBY', exits: :any, result: :ok)
+      @call_counter = 0
+      def block_call
+        1.times { @call_counter += 1 }
+      end
+
+      @trace_counter = 0
+      trace = TracePoint.new(:b_call) do |t|
+        @trace_counter += 1
+      end
+
+      block_call # JIT without TracePoint
+      trace.enable
+      block_call # call with TracePoint
+      trace.disable
+
+      return :"call_#{@call_counter}" if @call_counter != 2
+      return :"trace_#{@trace_counter}" if @trace_counter != 1
+
+      :ok
+    RUBY
+  end
+
   private
 
   def code_gc_helpers
diff --git a/vm_trace.c b/vm_trace.c
index 93a8c1a4ed..caed71e96f 100644
--- a/vm_trace.c
+++ b/vm_trace.c
@@ -87,8 +87,9 @@ update_global_event_hook(rb_event_flag_t prev_events, rb_event_flag_t new_events
 {
     rb_event_flag_t new_iseq_events = new_events & ISEQ_TRACE_EVENTS;
     rb_event_flag_t enabled_iseq_events = ruby_vm_event_enabled_global_flags & ISEQ_TRACE_EVENTS;
+    bool trace_iseq_p = new_iseq_events & ~enabled_iseq_events;
 
-    if (new_iseq_events & ~enabled_iseq_events) {
+    if (trace_iseq_p) {
         // :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 != RUBY_EVENT_CLASS) {
@@ -111,10 +112,11 @@ update_global_event_hook(rb_event_flag_t prev_events, rb_event_flag_t new_events
     ruby_vm_event_enabled_global_flags |= new_events;
     rb_objspace_set_event_hook(new_events);
 
-    if (new_events & RUBY_EVENT_TRACEPOINT_ALL) {
-        // Invalidate all code if listening for any TracePoint event.
+    if (trace_iseq_p) {
+        // Invalidate all code when ISEQs are modified to use trace_* insns above.
         // Internal events fire inside C routines so don't need special handling.
-        // Do this last so other ractors see updated vm events when they wake up.
+        // Do this after event flags updates so other ractors see updated vm events
+        // when they wake up.
         rb_yjit_tracing_invalidate_all();
     }
 }