From cb661d7d82984cdb54485ea3f4af01ac21960882 Mon Sep 17 00:00:00 2001 From: Alan Wu Date: Wed, 23 Oct 2024 10:17:08 -0400 Subject: [PATCH] YJIT: Check when gen_branch() fails We got some core dumps in the wild where a PendingBranch had everything as None, leading to a panic unwrapping in PendingBranch::into_branch(). This happened while compiling a `branchif`. It seems that the only way this can happen is when core::gen_branch() fails, but not due to OOM. We wouldn't have reach into_branch() when OOM, and the only way to not leave markers that would've set the branch's start_addr to some value in gen_branch() is for set_target() to fail, causing an early return. Unfortunately, it's hard to tell the exact sequence of events that led to this situation, but regardless, the dumps show us that we should check for errors in gen_branch(). Because gen_branch() is used deep in the stack during compilation (e.g. guard_known_class() -> jit_chain_guard() -> gen_branch()), it'd be bad for compile speed to propagate the error everywhere, not to mention the massive patch required. Opt for a flag checked near the end of compilation. --- yjit/src/codegen.rs | 32 ++++++++++++++++++++++++++------ yjit/src/core.rs | 13 +++++++------ yjit/src/stats.rs | 1 + 3 files changed, 34 insertions(+), 12 deletions(-) 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,