ZJIT: Remove JITed code after TracePoint is enabled

Author: st0012

jit.c

@@ -492,3 +492,54 @@ rb_jit_vm_unlock(unsigned int *recursive_lock_level, const char *file, int line)
 {
     rb_vm_lock_leave(recursive_lock_level, file, line);
 }
+
+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

@@ -413,18 +413,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)
 {
@@ -643,41 +631,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

@@ -333,7 +333,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() {
         .allowlist_function("rb_jit_multi_ractor_p")
         .allowlist_function("rb_jit_vm_lock_then_barrier")
         .allowlist_function("rb_jit_vm_unlock")
+        .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

@@ -349,7 +349,6 @@ fn main() {
         .allowlist_function("rb_full_cfunc_return")
         .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")
@@ -368,6 +367,8 @@ fn main() {
         .allowlist_function("rb_jit_multi_ractor_p")
         .allowlist_function("rb_jit_vm_lock_then_barrier")
         .allowlist_function("rb_jit_vm_unlock")
+        .allowlist_function("rb_jit_for_each_iseq")
+        .allowlist_function("rb_iseq_reset_jit_func")
         .allowlist_type("robject_offsets")
 
         // from vm_sync.h