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

@@ -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");

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<BlockId>,
     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) {

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,