ZJIT: Clear ZJIT compiled code when TracePoint is enabled

Author: st0012

jit.c

@@ -466,3 +466,72 @@ rb_jit_shape_too_complex_p(shape_id_t shape_id)
 {
     return rb_shape_too_complex_p(shape_id);
 }
+
+// Is anyone listening for :c_call and :c_return event currently?
+bool
+rb_c_method_tracing_currently_enabled(const rb_execution_context_t *ec)
+{
+    rb_event_flag_t tracing_events;
+    if (rb_multi_ractor_p()) {
+        tracing_events = ruby_vm_event_enabled_global_flags;
+    }
+    else {
+        // At the time of writing, events are never removed from
+        // ruby_vm_event_enabled_global_flags so always checking using it would
+        // mean we don't compile even after tracing is disabled.
+        tracing_events = rb_ec_ractor_hooks(ec)->events;
+    }
+
+    return tracing_events & (RUBY_EVENT_C_CALL | RUBY_EVENT_C_RETURN);
+}
+
+void
+rb_iseq_reset_jit_func(const rb_iseq_t *iseq)
+{
+    RUBY_ASSERT_ALWAYS(IMEMO_TYPE_P(iseq, imemo_iseq));
+    iseq->body->jit_entry = NULL;
+    iseq->body->jit_exception = NULL;
+    // Enable re-compiling this ISEQ. Event when it's invalidated for TracePoint,
+    // we'd like to re-compile ISEQs that haven't been converted to trace_* insns.
+    iseq->body->jit_entry_calls = 0;
+    iseq->body->jit_exception_calls = 0;
+}
+
+// Callback type for rb_jit_for_each_iseq
+typedef void (*rb_iseq_callback)(const rb_iseq_t *iseq, void *data);
+
+// Callback data for rb_jit_for_each_iseq
+struct iseq_callback_data {
+    rb_iseq_callback callback;
+    void *data;
+};
+
+// Heap-walking callback for rb_jit_for_each_iseq
+static int
+for_each_iseq_i(void *vstart, void *vend, size_t stride, void *data)
+{
+    const struct iseq_callback_data *callback_data = (struct iseq_callback_data *)data;
+    VALUE v = (VALUE)vstart;
+    for (; v != (VALUE)vend; v += stride) {
+        void *ptr = rb_asan_poisoned_object_p(v);
+        rb_asan_unpoison_object(v, false);
+
+        if (rb_obj_is_iseq(v)) {
+            rb_iseq_t *iseq = (rb_iseq_t *)v;
+            callback_data->callback(iseq, callback_data->data);
+        }
+
+        if (ptr) {
+            rb_asan_poison_object(v);
+        }
+    }
+    return 0;
+}
+
+// Walk all ISEQs in the heap and invoke the callback - shared between YJIT and ZJIT
+void
+rb_jit_for_each_iseq(rb_iseq_callback callback, void *data)
+{
+    struct iseq_callback_data callback_data = { .callback = callback, .data = data };
+    rb_objspace_each_objects(for_each_iseq_i, (void *)&callback_data);
+}

test/ruby/test_zjit.rb

@@ -2148,6 +2148,44 @@ def test
     }
   end
 
+  def test_global_tracepoint
+    assert_compiles 'true', %q{
+      def foo = 1
+
+      foo
+      foo
+
+      called = false
+
+      tp = TracePoint.new(:return) { |event|
+        if event.method_id == :foo
+          called = true
+        end
+      }
+      tp.enable do
+        foo
+      end
+      called
+    }
+  end
+
+  def test_local_tracepoint
+    assert_compiles 'true', %q{
+      def foo = 1
+
+      foo
+      foo
+
+      called = false
+
+      tp = TracePoint.new(:return) { |_| called = true }
+      tp.enable(target: method(:foo)) do
+        foo
+      end
+      called
+    }
+  end
+
   private
 
   # Assert that every method call in `test_script` can be compiled by ZJIT

vm_trace.c

@@ -35,6 +35,7 @@
 #include "vm_core.h"
 #include "ruby/ractor.h"
 #include "yjit.h"
+#include "zjit.h"
 
 #include "builtin.h"
 
@@ -135,6 +136,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_zjit_tracing_invalidate_all();
     }
 }
 
@@ -1285,6 +1287,7 @@ rb_tracepoint_enable_for_target(VALUE tpval, VALUE target, VALUE target_line)
     }
 
     rb_yjit_tracing_invalidate_all();
+    rb_zjit_tracing_invalidate_all();
 
     ruby_vm_event_local_num++;
 

yjit.c

@@ -343,24 +343,6 @@ rb_yjit_reserve_addr_space(uint32_t mem_size)
 #endif
 }
 
-// Is anyone listening for :c_call and :c_return event currently?
-bool
-rb_c_method_tracing_currently_enabled(const rb_execution_context_t *ec)
-{
-    rb_event_flag_t tracing_events;
-    if (rb_multi_ractor_p()) {
-        tracing_events = ruby_vm_event_enabled_global_flags;
-    }
-    else {
-        // At the time of writing, events are never removed from
-        // ruby_vm_event_enabled_global_flags so always checking using it would
-        // mean we don't compile even after tracing is disabled.
-        tracing_events = rb_ec_ractor_hooks(ec)->events;
-    }
-
-    return tracing_events & (RUBY_EVENT_C_CALL | RUBY_EVENT_C_RETURN);
-}
-
 // The code we generate in gen_send_cfunc() doesn't fire the c_return TracePoint event
 // like the interpreter. When tracing for c_return is enabled, we patch the code after
 // the C method return to call into this to fire the event.
@@ -414,18 +396,6 @@ rb_iseq_set_yjit_payload(const rb_iseq_t *iseq, void *payload)
     iseq->body->yjit_payload = payload;
 }
 
-void
-rb_iseq_reset_jit_func(const rb_iseq_t *iseq)
-{
-    RUBY_ASSERT_ALWAYS(IMEMO_TYPE_P(iseq, imemo_iseq));
-    iseq->body->jit_entry = NULL;
-    iseq->body->jit_exception = NULL;
-    // Enable re-compiling this ISEQ. Event when it's invalidated for TracePoint,
-    // we'd like to re-compile ISEQs that haven't been converted to trace_* insns.
-    iseq->body->jit_entry_calls = 0;
-    iseq->body->jit_exception_calls = 0;
-}
-
 rb_proc_t *
 rb_yjit_get_proc_ptr(VALUE procv)
 {
@@ -650,41 +620,6 @@ rb_yjit_constcache_shareable(const struct iseq_inline_constant_cache_entry *ice)
     return (ice->flags & IMEMO_CONST_CACHE_SHAREABLE) != 0;
 }
 
-// Used for passing a callback and other data over rb_objspace_each_objects
-struct iseq_callback_data {
-    rb_iseq_callback callback;
-    void *data;
-};
-
-// Heap-walking callback for rb_yjit_for_each_iseq().
-static int
-for_each_iseq_i(void *vstart, void *vend, size_t stride, void *data)
-{
-    const struct iseq_callback_data *callback_data = (struct iseq_callback_data *)data;
-    VALUE v = (VALUE)vstart;
-    for (; v != (VALUE)vend; v += stride) {
-        void *ptr = rb_asan_poisoned_object_p(v);
-        rb_asan_unpoison_object(v, false);
-
-        if (rb_obj_is_iseq(v)) {
-            rb_iseq_t *iseq = (rb_iseq_t *)v;
-            callback_data->callback(iseq, callback_data->data);
-        }
-
-        asan_poison_object_if(ptr, v);
-    }
-    return 0;
-}
-
-// Iterate through the whole GC heap and invoke a callback for each iseq.
-// Used for global code invalidation.
-void
-rb_yjit_for_each_iseq(rb_iseq_callback callback, void *data)
-{
-    struct iseq_callback_data callback_data = { .callback = callback, .data = data };
-    rb_objspace_each_objects(for_each_iseq_i, (void *)&callback_data);
-}
-
 // For running write barriers from Rust. Required when we add a new edge in the
 // object graph from `old` to `young`.
 void

yjit/bindgen/src/main.rs

@@ -336,7 +336,6 @@ fn main() {
         .allowlist_function("rb_yjit_constcache_shareable")
         .allowlist_function("rb_iseq_reset_jit_func")
         .allowlist_function("rb_yjit_dump_iseq_loc")
-        .allowlist_function("rb_yjit_for_each_iseq")
         .allowlist_function("rb_yjit_obj_written")
         .allowlist_function("rb_yjit_str_simple_append")
         .allowlist_function("rb_RSTRING_PTR")
@@ -355,6 +354,7 @@ fn main() {
         // From jit.c
         .allowlist_function("rb_assert_holding_vm_lock")
         .allowlist_function("rb_jit_shape_too_complex_p")
+        .allowlist_function("rb_jit_for_each_iseq")
         .allowlist_type("robject_offsets")
 
         // from vm_sync.h

yjit/src/core.rs

@@ -1818,7 +1818,7 @@ pub fn for_each_iseq<F: FnMut(IseqPtr)>(mut callback: F) {
         callback(iseq);
     }
     let mut data: &mut dyn FnMut(IseqPtr) = &mut callback;
-    unsafe { rb_yjit_for_each_iseq(Some(callback_wrapper), (&mut data) as *mut _ as *mut c_void) };
+    unsafe { rb_jit_for_each_iseq(Some(callback_wrapper), (&mut data) as *mut _ as *mut c_void) };
 }
 
 /// Iterate over all on-stack ISEQs

yjit/src/cruby_bindings.inc.rs

@@ -23,6 +23,7 @@ void rb_zjit_constant_state_changed(ID id);
 void rb_zjit_iseq_mark(void *payload);
 void rb_zjit_iseq_update_references(void *payload);
 void rb_zjit_before_ractor_spawn(void);
+void rb_zjit_tracing_invalidate_all(void);
 #else
 #define rb_zjit_enabled_p false
 static inline void rb_zjit_compile_iseq(const rb_iseq_t *iseq, rb_execution_context_t *ec, bool jit_exception) {}
@@ -33,6 +34,7 @@ static inline void rb_zjit_cme_invalidate(const rb_callable_method_entry_t *cme)
 static inline void rb_zjit_invalidate_ep_is_bp(const rb_iseq_t *iseq) {}
 static inline void rb_zjit_constant_state_changed(ID id) {}
 static inline void rb_zjit_before_ractor_spawn(void) {}
+static inline void rb_zjit_tracing_invalidate_all(void) {}
 #endif // #if USE_ZJIT
 
 #endif // #ifndef ZJIT_H

zjit.h

@@ -351,7 +351,6 @@ fn main() {
         .allowlist_function("rb_zjit_vm_unlock")
         .allowlist_function("rb_assert_(iseq|cme)_handle")
         .allowlist_function("rb_IMEMO_TYPE_P")
-        .allowlist_function("rb_iseq_reset_jit_func")
         .allowlist_function("rb_RSTRING_PTR")
         .allowlist_function("rb_RSTRING_LEN")
         .allowlist_function("rb_ENCODING_GET")
@@ -367,6 +366,8 @@ fn main() {
         // From jit.c
         .allowlist_function("rb_assert_holding_vm_lock")
         .allowlist_function("rb_jit_shape_too_complex_p")
+        .allowlist_function("rb_jit_for_each_iseq")
+        .allowlist_function("rb_iseq_reset_jit_func")
         .allowlist_type("robject_offsets")
 
         // from vm_sync.h

zjit/bindgen/src/main.rs

@@ -67,12 +67,17 @@ impl JITState {
 
 /// CRuby API to compile a given ISEQ
 #[unsafe(no_mangle)]
-pub extern "C" fn rb_zjit_iseq_gen_entry_point(iseq: IseqPtr, _ec: EcPtr) -> *const u8 {
+pub extern "C" fn rb_zjit_iseq_gen_entry_point(iseq: IseqPtr, ec: EcPtr) -> *const u8 {
     // Do not test the JIT code in HIR tests
     if cfg!(test) {
         return std::ptr::null();
     }
 
+    // Check if c_call or c_return tracing is enabled - if so, don't compile
+    if unsafe { rb_c_method_tracing_currently_enabled(ec) } {
+        return std::ptr::null();
+    }
+
     // Take a lock to avoid writing to ISEQ in parallel with Ractors.
     // with_vm_lock() does nothing if the program doesn't use Ractors.
     with_vm_lock(src_loc!(), || {