YJIT: Fail gracefully while OOM for new entry points

Previously, YJIT crashes with rb_bug() when asked to compile new methods
while out of executable memory.

To handle this situation gracefully, this change keeps track of all the
blocks compiled each invocation in case YJIT runs out of memory in the
middle of a compliation sequence. The list is used to free all blocks in
case compilation fails.

yjit_gen_block() is renamed to gen_single_block() to make it distinct from
gen_block_version(). Call to limit_block_version() and block_t
allocation is moved into the function to help tidy error checking in the
outer loop.

limit_block_version() now returns by value. I feel that an out parameter
with conditional mutation is unnecessarily hard to read in code that
does not need to go for last drop performance. There is a good chance
that the optimizer is able to output identical code anyways.

1 files changed, 69 insertions, 44 deletions

diff --git a/yjit_core.c b/yjit_core.c
index 4460d325fc..f19b83c5ff 100644
--- a/yjit_core.c
+++ b/yjit_core.c
@@ -542,9 +542,10 @@ static size_t get_num_versions(blockid_t blockid)
 
 // Keep track of a block version. Block should be fully constructed.
 static void
-add_block_version(blockid_t blockid, block_t *block)
+add_block_version(block_t *block)
 {
-    const rb_iseq_t *iseq = block->blockid.iseq;
+    const blockid_t blockid = block->blockid;
+    const rb_iseq_t *iseq = blockid.iseq;
     struct rb_iseq_constant_body *body = iseq->body;
 
     // Function entry blocks must have stack size 0
@@ -704,57 +705,66 @@ find_block_version(blockid_t blockid, const ctx_t *ctx)
 
 // Produce a generic context when the block version limit is hit for a blockid
 // Note that this will mutate the ctx argument
-static void
-limit_block_versions(blockid_t blockid, ctx_t *ctx)
+static ctx_t
+limit_block_versions(blockid_t blockid, const ctx_t *ctx)
 {
     // Guard chains implement limits separately, do nothing
     if (ctx->chain_depth > 0)
-        return;
+        return *ctx;
 
     // If this block version we're about to add will hit the version limit
-    if (get_num_versions(blockid) + 1 >= rb_yjit_opts.max_versions)
-    {
+    if (get_num_versions(blockid) + 1 >= rb_yjit_opts.max_versions) {
         // Produce a generic context that stores no type information,
-        // but still respects the stack_size and sp_offset constraints
+        // but still respects the stack_size and sp_offset constraints.
         // This new context will then match all future requests.
         ctx_t generic_ctx = DEFAULT_CTX;
         generic_ctx.stack_size = ctx->stack_size;
         generic_ctx.sp_offset = ctx->sp_offset;
 
         // Mutate the incoming context
-        *ctx = generic_ctx;
+        return generic_ctx;
     }
+
+    return *ctx;
 }
 
-// Compile a new block version immediately
+static void yjit_free_block(block_t *block);
+
+// Immediately compile a series of block versions at a starting point and
+// return the starting block.
 static block_t *
 gen_block_version(blockid_t blockid, const ctx_t *start_ctx, rb_execution_context_t *ec)
 {
-    // Allocate a new block version object
-    block_t *block = calloc(1, sizeof(block_t));
-    block->blockid = blockid;
-    memcpy(&block->ctx, start_ctx, sizeof(ctx_t));
-
-    // Store a pointer to the first block (returned by this function)
-    block_t *first_block = block;
-
-    // Limit the number of specialized versions for this block
-    limit_block_versions(block->blockid, &block->ctx);
+    // Small array to keep track of all the blocks compiled per invocation. We
+    // tend to have small batches since we often break up compilation with lazy
+    // stubs. Compilation is successful only if the whole batch is successful.
+    enum { MAX_PER_BATCH = 64 };
+    block_t *batch[MAX_PER_BATCH];
+    int compiled_count = 0;
+    bool batch_success = true;
+    block_t *block;
 
     // Generate code for the first block
-    yjit_gen_block(block, ec);
+    block = gen_single_block(blockid, start_ctx, ec);
+    batch_success = block && compiled_count < MAX_PER_BATCH;
+
+    if (batch_success) {
+        // Track the block
+        add_block_version(block);
 
-    // Keep track of the new block version
-    add_block_version(block->blockid, block);
+        batch[compiled_count] = block;
+        compiled_count++;
+    }
 
     // For each successor block to compile
-    for (;;) {
+    while (batch_success) {
         // If the previous block compiled doesn't have outgoing branches, stop
         if (rb_darray_size(block->outgoing) == 0) {
             break;
         }
 
-        // Get the last outgoing branch from the previous block
+        // Get the last outgoing branch from the previous block. Blocks can use
+        // gen_direct_jump() to request a block to be placed immediately after.
         branch_t *last_branch = rb_darray_back(block->outgoing);
 
         // If there is no next block to compile, stop
@@ -766,32 +776,48 @@ gen_block_version(blockid_t blockid, const ctx_t *start_ctx, rb_execution_contex
             rb_bug("invalid target for last branch");
         }
 
-        // Allocate a new block version object
-        // Use the context from the branch
-        block = calloc(1, sizeof(block_t));
-        block->blockid = last_branch->targets[0];
-        block->ctx = last_branch->target_ctxs[0];
-        //memcpy(&block->ctx, ctx, sizeof(ctx_t));
+        // Generate code for the current block using context from the last branch.
+        blockid_t requested_id = last_branch->targets[0];
+        const ctx_t *requested_ctx = &last_branch->target_ctxs[0];
+        block = gen_single_block(requested_id, requested_ctx, ec);
+        batch_success = block && compiled_count < MAX_PER_BATCH;
 
-        // Limit the number of specialized versions for this block
-        limit_block_versions(block->blockid, &block->ctx);
-
-        // Generate code for the current block
-        yjit_gen_block(block, ec);
-
-        // Keep track of the new block version
-        add_block_version(block->blockid, block);
+        // If the batch failed, stop
+        if (!batch_success) {
+            break;
+        }
 
-        // Patch the last branch address
+        // Connect the last branch and the new block
         last_branch->dst_addrs[0] = block->start_addr;
         rb_darray_append(&block->incoming, last_branch);
         last_branch->blocks[0] = block;
 
         // This block should immediately follow the last branch
         RUBY_ASSERT(block->start_addr == last_branch->end_addr);
+
+        // Track the block
+        add_block_version(block);
+
+        batch[compiled_count] = block;
+        compiled_count++;
     }
 
-    return first_block;
+    if (batch_success) {
+        // Success. Return first block in the batch.
+        RUBY_ASSERT(compiled_count > 0);
+        return batch[0];
+    }
+    else {
+        // The batch failed. Free everything in the batch
+        for (int block_idx = 0; block_idx < compiled_count; block_idx++) {
+            yjit_free_block(batch[block_idx]);
+        }
+
+#if YJIT_STATS
+        yjit_runtime_counters.compilation_failure++;
+#endif
+        return NULL;
+    }
 }
 
 // Generate a block version that is an entry point inserted into an iseq
@@ -807,15 +833,14 @@ gen_entry_point(const rb_iseq_t *iseq, uint32_t insn_idx, rb_execution_context_t
     // The entry context makes no assumptions about types
     blockid_t blockid = { iseq, insn_idx };
 
-    // Write the interpreter entry prologue
+    // Write the interpreter entry prologue. Might be NULL when out of memory.
     uint8_t *code_ptr = yjit_entry_prologue(cb, iseq);
 
     // Try to generate code for the entry block
     block_t *block = gen_block_version(blockid, &DEFAULT_CTX, ec);
 
     // If we couldn't generate any code
-    if (block->end_idx == insn_idx)
-    {
+    if (!block || block->end_idx == insn_idx) {
         return NULL;
     }