ZJIT: Bail gen_ccall_variadic when c_call and c_return are traced

Author: st0012

jit.c

@@ -493,6 +493,24 @@ rb_jit_vm_unlock(unsigned int *recursive_lock_level, const char *file, int line)
     rb_vm_lock_leave(recursive_lock_level, file, line);
 }
 
+// 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)
 {

vm.c

@@ -445,7 +445,7 @@ jit_compile(rb_execution_context_t *ec)
 
         // At call-threshold, compile the ISEQ with ZJIT.
         if (body->jit_entry_calls == rb_zjit_call_threshold) {
-            rb_zjit_compile_iseq(iseq, false);
+            rb_zjit_compile_iseq(iseq, ec, false);
         }
     }
 #endif
@@ -505,7 +505,7 @@ jit_compile_exception(rb_execution_context_t *ec)
 
         // At call-threshold, compile the ISEQ with ZJIT.
         if (body->jit_exception_calls == rb_zjit_call_threshold) {
-            rb_zjit_compile_iseq(iseq, true);
+            rb_zjit_compile_iseq(iseq, ec, true);
         }
     }
 #endif

yjit.c

@@ -169,24 +169,6 @@ rb_yjit_exit_locations_dict(VALUE *yjit_raw_samples, int *yjit_line_samples, int
     return result;
 }
 
-// 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.

zjit.c

@@ -34,14 +34,14 @@
 void rb_zjit_profile_disable(const rb_iseq_t *iseq);
 
 void
-rb_zjit_compile_iseq(const rb_iseq_t *iseq, bool jit_exception)
+rb_zjit_compile_iseq(const rb_iseq_t *iseq, rb_execution_context_t *ec, bool jit_exception)
 {
     RB_VM_LOCKING() {
         rb_vm_barrier();
 
         // Compile a block version starting at the current instruction
-        uint8_t *rb_zjit_iseq_gen_entry_point(const rb_iseq_t *iseq, bool jit_exception); // defined in Rust
-        uintptr_t code_ptr = (uintptr_t)rb_zjit_iseq_gen_entry_point(iseq, jit_exception);
+        uint8_t *rb_zjit_iseq_gen_entry_point(const rb_iseq_t *iseq, rb_execution_context_t *ec, bool jit_exception); // defined in Rust
+        uintptr_t code_ptr = (uintptr_t)rb_zjit_iseq_gen_entry_point(iseq, ec, jit_exception);
 
         if (jit_exception) {
             iseq->body->jit_exception = (rb_jit_func_t)code_ptr;

zjit.h

@@ -13,7 +13,7 @@
 extern bool rb_zjit_enabled_p;
 extern uint64_t rb_zjit_call_threshold;
 extern uint64_t rb_zjit_profile_threshold;
-void rb_zjit_compile_iseq(const rb_iseq_t *iseq, bool jit_exception);
+void rb_zjit_compile_iseq(const rb_iseq_t *iseq, rb_execution_context_t *ec, bool jit_exception);
 void rb_zjit_profile_insn(uint32_t insn, rb_execution_context_t *ec);
 void rb_zjit_profile_enable(const rb_iseq_t *iseq);
 void rb_zjit_bop_redefined(int redefined_flag, enum ruby_basic_operators bop);
@@ -26,7 +26,7 @@ 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, bool jit_exception) {}
+static inline void rb_zjit_compile_iseq(const rb_iseq_t *iseq, rb_execution_context_t *ec, bool jit_exception) {}
 static inline void rb_zjit_profile_insn(uint32_t insn, rb_execution_context_t *ec) {}
 static inline void rb_zjit_profile_enable(const rb_iseq_t *iseq) {}
 static inline void rb_zjit_bop_redefined(int redefined_flag, enum ruby_basic_operators bop) {}

zjit/src/codegen.rs

@@ -38,17 +38,21 @@ struct JITState {
 
     /// The number of bytes allocated for basic block arguments spilled onto the C stack
     c_stack_slots: usize,
+
+    /// Execution context
+    ec: EcPtr,
 }
 
 impl JITState {
     /// Create a new JITState instance
-    fn new(iseq: IseqPtr, num_insns: usize, num_blocks: usize, c_stack_slots: usize) -> Self {
+    fn new(iseq: IseqPtr, num_insns: usize, num_blocks: usize, c_stack_slots: usize, ec: EcPtr) -> Self {
         JITState {
             iseq,
             opnds: vec![None; num_insns],
             labels: vec![None; num_blocks],
             iseq_calls: Vec::default(),
             c_stack_slots,
+            ec,
         }
     }
 
@@ -74,7 +78,7 @@ impl JITState {
 /// If jit_exception is true, compile JIT code for handling exceptions.
 /// See jit_compile_exception() for details.
 #[unsafe(no_mangle)]
-pub extern "C" fn rb_zjit_iseq_gen_entry_point(iseq: IseqPtr, jit_exception: bool) -> *const u8 {
+pub extern "C" fn rb_zjit_iseq_gen_entry_point(iseq: IseqPtr, ec: EcPtr, jit_exception: bool) -> *const u8 {
     // Do not test the JIT code in HIR tests
     if cfg!(test) {
         return std::ptr::null();
@@ -84,7 +88,7 @@ pub extern "C" fn rb_zjit_iseq_gen_entry_point(iseq: IseqPtr, jit_exception: boo
     // with_vm_lock() does nothing if the program doesn't use Ractors.
     with_vm_lock(src_loc!(), || {
         let cb = ZJITState::get_code_block();
-        let mut code_ptr = with_time_stat(compile_time_ns, || gen_iseq_entry_point(cb, iseq, jit_exception));
+        let mut code_ptr = with_time_stat(compile_time_ns, || gen_iseq_entry_point(cb, iseq, ec, jit_exception));
 
         if let Err(err) = &code_ptr {
             // Assert that the ISEQ compiles if RubyVM::ZJIT.assert_compiles is enabled.
@@ -110,7 +114,7 @@ pub extern "C" fn rb_zjit_iseq_gen_entry_point(iseq: IseqPtr, jit_exception: boo
 }
 
 /// Compile an entry point for a given ISEQ
-fn gen_iseq_entry_point(cb: &mut CodeBlock, iseq: IseqPtr, jit_exception: bool) -> Result<CodePtr, CompileError> {
+fn gen_iseq_entry_point(cb: &mut CodeBlock, iseq: IseqPtr, ec: EcPtr, jit_exception: bool) -> Result<CodePtr, CompileError> {
     // We don't support exception handlers yet
     if jit_exception {
         return Err(CompileError::ExceptionHandler);
@@ -122,7 +126,7 @@ fn gen_iseq_entry_point(cb: &mut CodeBlock, iseq: IseqPtr, jit_exception: bool)
     })?;
 
     // Compile the High-level IR
-    let start_ptr = gen_iseq(cb, iseq, Some(&function)).inspect_err(|err| {
+    let start_ptr = gen_iseq(cb, iseq, ec, Some(&function)).inspect_err(|err| {
         debug!("{err:?}: gen_iseq failed: {}", iseq_get_location(iseq, 0));
     })?;
 
@@ -196,7 +200,7 @@ fn gen_entry(cb: &mut CodeBlock, iseq: IseqPtr, function: &Function, function_pt
 }
 
 /// Compile an ISEQ into machine code if not compiled yet
-fn gen_iseq(cb: &mut CodeBlock, iseq: IseqPtr, function: Option<&Function>) -> Result<CodePtr, CompileError> {
+fn gen_iseq(cb: &mut CodeBlock, iseq: IseqPtr, ec: EcPtr, function: Option<&Function>) -> Result<CodePtr, CompileError> {
     // Return an existing pointer if it's already compiled
     let payload = get_or_create_iseq_payload(iseq);
     match &payload.status {
@@ -206,7 +210,7 @@ fn gen_iseq(cb: &mut CodeBlock, iseq: IseqPtr, function: Option<&Function>) -> R
     }
 
     // Compile the ISEQ
-    let code_ptr = gen_iseq_body(cb, iseq, function, payload);
+    let code_ptr = gen_iseq_body(cb, iseq, ec, function, payload);
     match &code_ptr {
         Ok(start_ptr) => {
             payload.status = IseqStatus::Compiled(*start_ptr);
@@ -221,15 +225,15 @@ fn gen_iseq(cb: &mut CodeBlock, iseq: IseqPtr, function: Option<&Function>) -> R
 }
 
 /// Compile an ISEQ into machine code
-fn gen_iseq_body(cb: &mut CodeBlock, iseq: IseqPtr, function: Option<&Function>, payload: &mut IseqPayload) -> Result<CodePtr, CompileError> {
+fn gen_iseq_body(cb: &mut CodeBlock, iseq: IseqPtr, ec: EcPtr, function: Option<&Function>, payload: &mut IseqPayload) -> Result<CodePtr, CompileError> {
     // Convert ISEQ into optimized High-level IR if not given
     let function = match function {
         Some(function) => function,
         None => &compile_iseq(iseq)?,
     };
 
     // Compile the High-level IR
-    let (start_ptr, gc_offsets, iseq_calls) = gen_function(cb, iseq, function)?;
+    let (start_ptr, gc_offsets, iseq_calls) = gen_function(cb, iseq, ec, function)?;
 
     // Stub callee ISEQs for JIT-to-JIT calls
     for iseq_call in iseq_calls.iter() {
@@ -243,9 +247,9 @@ fn gen_iseq_body(cb: &mut CodeBlock, iseq: IseqPtr, function: Option<&Function>,
 }
 
 /// Compile a function
-fn gen_function(cb: &mut CodeBlock, iseq: IseqPtr, function: &Function) -> Result<(CodePtr, Vec<CodePtr>, Vec<IseqCallRef>), CompileError> {
+fn gen_function(cb: &mut CodeBlock, iseq: IseqPtr, ec: EcPtr, function: &Function) -> Result<(CodePtr, Vec<CodePtr>, Vec<IseqCallRef>), CompileError> {
     let c_stack_slots = max_num_params(function).saturating_sub(ALLOC_REGS.len());
-    let mut jit = JITState::new(iseq, function.num_insns(), function.num_blocks(), c_stack_slots);
+    let mut jit = JITState::new(iseq, function.num_insns(), function.num_blocks(), c_stack_slots, ec);
     let mut asm = Assembler::new();
 
     // Compile each basic block
@@ -342,6 +346,10 @@ fn gen_insn(cb: &mut CodeBlock, jit: &mut JITState, asm: &mut Assembler, functio
         asm_comment!(asm, "Insn: {insn_id} {insn}");
     }
 
+    unsafe extern "C" {
+        fn rb_c_method_tracing_currently_enabled(ec: EcPtr) -> bool;
+    }
+
     let out_opnd = match insn {
         Insn::Const { val: Const::Value(val) } => gen_const(*val),
         Insn::Const { .. } => panic!("Unexpected Const in gen_insn: {insn}"),
@@ -397,6 +405,7 @@ fn gen_insn(cb: &mut CodeBlock, jit: &mut JITState, asm: &mut Assembler, functio
         &Insn::GuardBlockParamProxy { level, state } => no_output!(gen_guard_block_param_proxy(jit, asm, level, &function.frame_state(state))),
         Insn::PatchPoint { invariant, state } => no_output!(gen_patch_point(jit, asm, invariant, &function.frame_state(*state))),
         Insn::CCall { cfun, args, name: _, return_type: _, elidable: _ } => gen_ccall(asm, *cfun, opnds!(args)),
+        Insn::CCallVariadic { state, .. } if unsafe { rb_c_method_tracing_currently_enabled(jit.ec) } => return Err(*state),
         Insn::CCallVariadic { cfun, recv, args, name: _, cme, state } => {
             gen_ccall_variadic(jit, asm, *cfun, opnd!(recv), opnds!(args), *cme, &function.frame_state(*state))
         }
@@ -1816,7 +1825,7 @@ c_callable! {
     /// This function is expected to be called repeatedly when ZJIT fails to compile the stub.
     /// We should be able to compile most (if not all) function stubs by side-exiting at unsupported
     /// instructions, so this should be used primarily for cb.has_dropped_bytes() situations.
-    fn function_stub_hit(iseq_call_ptr: *const c_void, cfp: CfpPtr, sp: *mut VALUE) -> *const u8 {
+    fn function_stub_hit(iseq_call_ptr: *const c_void, cfp: CfpPtr, sp: *mut VALUE, ec: EcPtr) -> *const u8 {
         with_vm_lock(src_loc!(), || {
             // gen_push_frame() doesn't set PC, so we need to set them before exit.
             // function_stub_hit_body() may allocate and call gc_validate_pc(), so we always set PC.
@@ -1864,7 +1873,7 @@ c_callable! {
             }
 
             // Otherwise, attempt to compile the ISEQ. We have to mark_all_executable() beyond this point.
-            let code_ptr = with_time_stat(compile_time_ns, || function_stub_hit_body(cb, &iseq_call));
+            let code_ptr = with_time_stat(compile_time_ns, || function_stub_hit_body(cb, &iseq_call, ec));
             let code_ptr = code_ptr.unwrap_or_else(|compile_error| {
                 prepare_for_exit(iseq, cfp, sp, &compile_error);
                 ZJITState::get_exit_trampoline_with_counter()
@@ -1876,9 +1885,9 @@ c_callable! {
 }
 
 /// Compile an ISEQ for a function stub
-fn function_stub_hit_body(cb: &mut CodeBlock, iseq_call: &Rc<RefCell<IseqCall>>) -> Result<CodePtr, CompileError> {
+fn function_stub_hit_body(cb: &mut CodeBlock, iseq_call: &Rc<RefCell<IseqCall>>, ec: EcPtr) -> Result<CodePtr, CompileError> {
     // Compile the stubbed ISEQ
-    let code_ptr = gen_iseq(cb, iseq_call.borrow().iseq, None).inspect_err(|err| {
+    let code_ptr = gen_iseq(cb, iseq_call.borrow().iseq, ec, None).inspect_err(|err| {
         debug!("{err:?}: gen_iseq failed: {}", iseq_get_location(iseq_call.borrow().iseq, 0));
     })?;
 
@@ -1924,7 +1933,7 @@ pub fn gen_function_stub_hit_trampoline(cb: &mut CodeBlock) -> Result<CodePtr, C
     const { assert!(ALLOC_REGS.len() % 2 == 0, "x86_64 would need to push one more if we push an odd number of regs"); }
 
     // Compile the stubbed ISEQ
-    let jump_addr = asm_ccall!(asm, function_stub_hit, SCRATCH_OPND, CFP, SP);
+    let jump_addr = asm_ccall!(asm, function_stub_hit, SCRATCH_OPND, CFP, SP, EC);
     asm.mov(SCRATCH_OPND, jump_addr);
 
     asm_comment!(asm, "restore argument registers");