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.

2 files changed, 17 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
diff --git a/mjit_worker.c b/mjit_worker.c
index 3a73f14679..50f1b0787e 100644
--- a/mjit_worker.c
+++ b/mjit_worker.c
@@ -1396,6 +1396,8 @@ unload_units(void)
     }
 }
 
+static void mjit_add_iseq_to_process(const rb_iseq_t *iseq, const struct rb_mjit_compile_info *compile_info);
+
 // The function implementing a worker. It is executed in a separate
 // thread by rb_thread_create_mjit_thread. It compiles precompiled header
 // and then compiles requested ISeqs.
@@ -1445,6 +1447,8 @@ mjit_worker(void)
                         unit->stale_p = false;
                         remove_from_list(unit, &active_units);
                         add_to_list(unit, &stale_units);
+                        // Lazily put it to unit_queue as well to avoid race conditions on jit_unit with mjit_compile.
+                        mjit_add_iseq_to_process(unit->iseq, &unit->iseq->body->jit_unit->compile_info);
                     }
                 }
             }