Call appropriate hooks on MJIT's fork

This takes care of signal_self_pipe and other things.

2 files changed, 27 insertions, 18 deletions

diff --git a/mjit.c b/mjit.c
index 662ef99b32..0aa6d6e7cf 100644
--- a/mjit.c
+++ b/mjit.c
@@ -742,23 +742,17 @@ mjit_compact_unit(struct rb_mjit_unit *unit)
     return 1;
 }
 
+extern pid_t rb_mjit_fork();
+
 static pid_t
 start_mjit_compact(struct rb_mjit_unit *unit)
 {
-    rb_vm_t *vm = GET_VM();
-    rb_native_mutex_lock(&vm->waitpid_lock);
-
-    pid_t pid = rb_fork();
+    pid_t pid = rb_mjit_fork();
     if (pid == 0) {
-        rb_native_mutex_unlock(&vm->waitpid_lock);
-
         int exit_code = mjit_compact_unit(unit);
         exit(exit_code);
     }
     else {
-        mjit_add_waiting_pid(vm, pid);
-        rb_native_mutex_unlock(&vm->waitpid_lock);
-
         return pid;
     }
 }
@@ -908,20 +902,12 @@ mjit_compile_unit(struct rb_mjit_unit *unit)
 static pid_t
 start_mjit_compile(struct rb_mjit_unit *unit)
 {
-    rb_vm_t *vm = GET_VM();
-    rb_native_mutex_lock(&vm->waitpid_lock);
-
-    pid_t pid = rb_fork();
+    pid_t pid = rb_mjit_fork();
     if (pid == 0) {
-        rb_native_mutex_unlock(&vm->waitpid_lock);
-
         int exit_code = mjit_compile_unit(unit);
         exit(exit_code);
     }
     else {
-        mjit_add_waiting_pid(vm, pid);
-        rb_native_mutex_unlock(&vm->waitpid_lock);
-
         return pid;
     }
 }
diff --git a/process.c b/process.c
index 57df2dc06f..405c7edcb3 100644
--- a/process.c
+++ b/process.c
@@ -4204,6 +4204,29 @@ retry_fork_async_signal_safe(struct rb_process_status *status, int *ep,
     }
 }
 
+#if USE_MJIT
+// This is used to create MJIT's child Ruby process
+pid_t
+rb_mjit_fork(void)
+{
+    struct child_handler_disabler_state old;
+    rb_vm_t *vm = GET_VM();
+    prefork();
+    disable_child_handler_before_fork(&old);
+    before_fork_ruby();
+
+    rb_native_mutex_lock(&vm->waitpid_lock);
+    pid_t pid = rb_fork();
+    if (pid > 0) mjit_add_waiting_pid(vm, pid);
+    rb_native_mutex_unlock(&vm->waitpid_lock);
+
+    after_fork_ruby();
+    disable_child_handler_fork_parent(&old);
+
+    return pid;
+}
+#endif
+
 static rb_pid_t
 fork_check_err(struct rb_process_status *status, int (*chfunc)(void*, char *, size_t), void *charg,
         VALUE fds, char *errmsg, size_t errmsg_buflen,