diff --git a/yjit/src/codegen.rs b/yjit/src/codegen.rs index 4447522567..ded89457c6 100644 --- a/yjit/src/codegen.rs +++ b/yjit/src/codegen.rs @@ -118,6 +118,12 @@ pub struct JITState<'a> { /// When true, this block is the first block compiled by gen_block_series(). first_block: bool, + + /// A killswitch for bailing out of compilation. Used in rare situations where we need to fail + /// compilation deep in the stack (e.g. codegen failed for some jump target, but not due to + /// OOM). Because these situations are so rare it's not worth it to check and propogate at each + /// site. Instead, we check this once at the end. + block_abandoned: bool, } impl<'a> JITState<'a> { @@ -145,6 +151,7 @@ impl<'a> JITState<'a> { perf_map: Rc::default(), perf_stack: vec![], first_block, + block_abandoned: false, } } @@ -1350,6 +1357,12 @@ pub fn gen_single_block( // doesn't go to the next instruction in the same iseq. assert!(!jit.record_boundary_patch_point); + // Bail when requested to. + if jit.block_abandoned { + incr_counter!(abandoned_block_count); + return Err(()); + } + // Pad the block if it has the potential to be invalidated if jit.block_entry_exit.is_some() { asm.pad_inval_patch(); @@ -2705,7 +2718,10 @@ fn jit_chain_guard( idx: jit.insn_idx, }; - gen_branch(jit, asm, bid, &deeper, None, None, target0_gen_fn); + // Bail if we can't generate the branch + if gen_branch(jit, asm, bid, &deeper, None, None, target0_gen_fn).is_none() { + jit.block_abandoned = true; + } } else { target0_gen_fn.call(asm, Target::side_exit(counter), None); } @@ -4564,7 +4580,7 @@ fn gen_branchif( Some(next_block), Some(&ctx), BranchGenFn::BranchIf(Cell::new(BranchShape::Default)), - ); + )?; } Some(EndBlock) @@ -4618,7 +4634,7 @@ fn gen_branchunless( Some(next_block), Some(&ctx), BranchGenFn::BranchUnless(Cell::new(BranchShape::Default)), - ); + )?; } Some(EndBlock) @@ -4669,7 +4685,7 @@ fn gen_branchnil( Some(next_block), Some(&ctx), BranchGenFn::BranchNil(Cell::new(BranchShape::Default)), - ); + )?; } Some(EndBlock) @@ -7988,7 +8004,7 @@ fn gen_send_iseq( return_asm.ctx.set_as_return_landing(); // Write the JIT return address on the callee frame - gen_branch( + if gen_branch( jit, asm, return_block, @@ -7996,7 +8012,11 @@ fn gen_send_iseq( None, None, BranchGenFn::JITReturn, - ); + ).is_none() { + // Returning None here would have send_dynamic() code following incomplete + // send code. Abandon the block instead. + jit.block_abandoned = true; + } // ec->cfp is updated after cfp->jit_return for rb_profile_frames() safety asm_comment!(asm, "switch to new CFP"); diff --git a/yjit/src/core.rs b/yjit/src/core.rs index f1e1c1d9a3..43ad688cc9 100644 --- a/yjit/src/core.rs +++ b/yjit/src/core.rs @@ -3845,6 +3845,7 @@ impl Assembler } } +#[must_use] pub fn gen_branch( jit: &mut JITState, asm: &mut Assembler, @@ -3853,17 +3854,17 @@ pub fn gen_branch( target1: Option, ctx1: Option<&Context>, gen_fn: BranchGenFn, -) { +) -> Option<()> { let branch = new_pending_branch(jit, gen_fn); // Get the branch targets or stubs - let target0_addr = branch.set_target(0, target0, ctx0, jit); + let target0_addr = branch.set_target(0, target0, ctx0, jit)?; let target1_addr = if let Some(ctx) = ctx1 { let addr = branch.set_target(1, target1.unwrap(), ctx, jit); if addr.is_none() { // target1 requested but we're out of memory. // Avoid unwrap() in gen_fn() - return; + return None; } addr @@ -3871,10 +3872,10 @@ pub fn gen_branch( // Call the branch generation function asm.mark_branch_start(&branch); - if let Some(dst_addr) = target0_addr { - branch.gen_fn.call(asm, Target::CodePtr(dst_addr), target1_addr.map(|addr| Target::CodePtr(addr))); - } + branch.gen_fn.call(asm, Target::CodePtr(target0_addr), target1_addr.map(|addr| Target::CodePtr(addr))); asm.mark_branch_end(&branch); + + Some(()) } pub fn gen_direct_jump(jit: &mut JITState, ctx: &Context, target0: BlockId, asm: &mut Assembler) { diff --git a/yjit/src/stats.rs b/yjit/src/stats.rs index a6dffe9103..3dc37d4bac 100644 --- a/yjit/src/stats.rs +++ b/yjit/src/stats.rs @@ -561,6 +561,7 @@ make_counters! { compiled_branch_count, compile_time_ns, compilation_failure, + abandoned_block_count, block_next_count, defer_count, defer_empty_count,