4 files changed, 26 insertions, 24 deletions

diff --git a/mjit.c b/mjit.c
index 4895e42d7d..d25fa79d24 100644
--- a/mjit.c
+++ b/mjit.c
@@ -1259,18 +1259,18 @@ check_unit_queue(void)
 
     current_cc_ms = real_ms_time();
     current_cc_unit = unit;
-    current_cc_pid = start_mjit_compile(unit);
-
-    // JIT failure
-    if (current_cc_pid == -1) {
-        current_cc_pid = 0;
-        current_cc_unit->iseq->body->jit_func = (mjit_func_t)NOT_COMPILED_JIT_ISEQ_FUNC; // TODO: consider unit->compact_p
-        current_cc_unit = NULL;
-        return;
-    }
-
     if (mjit_opts.wait) {
-        mjit_wait(unit->iseq->body);
+        int exit_code = mjit_compile_unit(unit);
+        mjit_notify_waitpid(exit_code);
+    }
+    else {
+        current_cc_pid = start_mjit_compile(unit);
+        if (current_cc_pid == -1) { // JIT failure
+            current_cc_pid = 0;
+            current_cc_unit->iseq->body->jit_func = (mjit_func_t)NOT_COMPILED_JIT_ISEQ_FUNC; // TODO: consider unit->compact_p
+            current_cc_unit = NULL;
+            return;
+        }
     }
 }
 
@@ -1315,7 +1315,13 @@ check_compaction(void)
             // TODO: assert unit is null
             current_cc_ms = real_ms_time();
             current_cc_unit = unit;
-            current_cc_pid = start_mjit_compact(unit);
+            if (mjit_opts.wait) {
+                int exit_code = mjit_compact_unit(unit);
+                mjit_notify_waitpid(exit_code);
+            }
+            else {
+                current_cc_pid = start_mjit_compact(unit);
+            }
             // TODO: check -1
         }
     }
@@ -1323,7 +1329,7 @@ check_compaction(void)
 
 // Check the current CC process if any, and start a next C compiler process as needed.
 void
-mjit_notify_waitpid(int status)
+mjit_notify_waitpid(int exit_code)
 {
     // TODO: check current_cc_pid?
     current_cc_pid = 0;
@@ -1333,11 +1339,7 @@ mjit_notify_waitpid(int status)
     sprint_uniq_filename(c_file, (int)sizeof(c_file), current_cc_unit->id, MJIT_TMP_PREFIX, ".c");
 
     // Check the result
-    bool success = false;
-    if (WIFEXITED(status)) {
-        success = (WEXITSTATUS(status) == 0);
-    }
-    if (!success) {
+    if (exit_code != 0) {
         verbose(2, "Failed to generate so");
         if (!current_cc_unit->compact_p) {
             current_cc_unit->iseq->body->jit_func = (mjit_func_t)NOT_COMPILED_JIT_ISEQ_FUNC;
@@ -1438,8 +1440,8 @@ rb_mjit_add_iseq_to_process(const rb_iseq_t *iseq)
     check_unit_queue();
 }
 
-// For this timeout seconds, --jit-wait will wait for JIT compilation finish.
-#define MJIT_WAIT_TIMEOUT_SECONDS 600
+// For this timeout seconds, mjit_finish will wait for JIT compilation finish.
+#define MJIT_WAIT_TIMEOUT_SECONDS 5
 
 static void
 mjit_wait(struct rb_iseq_constant_body *body)
diff --git a/mjit.h b/mjit.h
index 66db417daf..55b9fbcfcd 100644
--- a/mjit.h
+++ b/mjit.h
@@ -100,7 +100,7 @@ extern void mjit_mark(void);
 extern struct mjit_cont *mjit_cont_new(rb_execution_context_t *ec);
 extern void mjit_cont_free(struct mjit_cont *cont);
 extern void mjit_mark_cc_entries(const struct rb_iseq_constant_body *const body);
-extern void mjit_notify_waitpid(int status);
+extern void mjit_notify_waitpid(int exit_code);
 
 void mjit_child_after_fork(void);
 
diff --git a/thread.c b/thread.c
index b271128193..c42c79914d 100644
--- a/thread.c
+++ b/thread.c
@@ -2323,7 +2323,7 @@ rb_threadptr_execute_interrupts(rb_thread_t *th, int blocking_timing)
             // outside ruby_sigchld_handler to avoid recursively relying on the SIGCHLD handler.
             if (mjit_waitpid_finished) {
                 mjit_waitpid_finished = false;
-                mjit_notify_waitpid(mjit_waitpid_status);
+                mjit_notify_waitpid(WIFEXITED(mjit_waitpid_status) ? WEXITSTATUS(mjit_waitpid_status) : -1);
             }
 #endif
         }
diff --git a/vm.c b/vm.c
index 15b6fa3a57..0de461392f 100644
--- a/vm.c
+++ b/vm.c
@@ -393,8 +393,8 @@ mjit_check_iseq(rb_execution_context_t *ec, const rb_iseq_t *iseq, struct rb_ise
         RB_DEBUG_COUNTER_INC(mjit_exec_not_added);
         if (body->total_calls == mjit_opts.min_calls) {
             rb_mjit_add_iseq_to_process(iseq);
-            if (UNLIKELY(mjit_opts.wait)) {
-                return rb_mjit_wait_call(ec, body);
+            if (UNLIKELY(mjit_opts.wait && (uintptr_t)body->jit_func > LAST_JIT_ISEQ_FUNC)) {
+                return body->jit_func(ec, ec->cfp);
             }
         }
         break;