Fix a race condition around mjit_recompile

This fixes SEGVs like https://github.com/ruby/ruby/runs/2715166621?check_suite_focus=true.

When mjit_recompile is called when mjit_compile is compiling the exact
same iseq (and after it called mjit_capture_cc_entries), iseq->body->jit_unit
is re-created and its cc_entries becomes NULL. Then, when it tries to
lookup cc_entries through iseq->body->jit_unit, it fails.

1 files changed, 13 insertions, 8 deletions

diff --git a/mjit.c b/mjit.c
index 035353ddae..d9c83c4829 100644
--- a/mjit.c
+++ b/mjit.c
@@ -350,17 +350,22 @@ mjit_recompile(const rb_iseq_t *iseq)
             RSTRING_PTR(rb_iseq_path(iseq)), FIX2INT(iseq->body->location.first_lineno));
     assert(iseq->body->jit_unit != NULL);
 
-    // Lazily move active_units to stale_units to avoid race conditions around active_units with compaction
-    CRITICAL_SECTION_START(3, "in rb_mjit_recompile_iseq");
-    iseq->body->jit_unit->stale_p = true;
-    pending_stale_p = true;
-    CRITICAL_SECTION_FINISH(3, "in rb_mjit_recompile_iseq");
-
-    iseq->body->jit_func = (mjit_func_t)NOT_ADDED_JIT_ISEQ_FUNC;
-    mjit_add_iseq_to_process(iseq, &iseq->body->jit_unit->compile_info);
     if (UNLIKELY(mjit_opts.wait)) {
+        remove_from_list(iseq->body->jit_unit, &active_units);
+        add_to_list(iseq->body->jit_unit, &stale_units);
+        mjit_add_iseq_to_process(iseq, &iseq->body->jit_unit->compile_info);
         mjit_wait(iseq->body);
     }
+    else {
+        // Lazily move active_units to stale_units to avoid race conditions around active_units with compaction.
+        // Also, it's lazily moved to unit_queue as well because otherwise it won't be added to stale_units properly.
+        // It's good to avoid a race condition between mjit_add_iseq_to_process and mjit_compile around jit_unit as well.
+        CRITICAL_SECTION_START(3, "in rb_mjit_recompile_iseq");
+        iseq->body->jit_unit->stale_p = true;
+        iseq->body->jit_func = (mjit_func_t)NOT_ADDED_JIT_ISEQ_FUNC;
+        pending_stale_p = true;
+        CRITICAL_SECTION_FINISH(3, "in rb_mjit_recompile_iseq");
+    }
 }
 
 // Recompile iseq, disabling send optimization