mjit_worker.c: support MJIT in forked Ruby process

by launching MJIT worker thread in child Ruby process.

See the comment before `mjit_child_after_fork` for details.

git-svn-id: svn+ssh://ci.ruby-lang.org/ruby/trunk@65785 b2dd03c8-39d4-4d8f-98ff-823fe69b080e

5 files changed, 115 insertions, 20 deletions

diff --git a/mjit.c b/mjit.c
index 047971f162..25b2bd9322 100644
--- a/mjit.c
+++ b/mjit.c
@@ -496,18 +496,6 @@ init_header_filename(void)
     return TRUE;
 }
 
-/* This is called after each fork in the child in to switch off MJIT
-   engine in the child as it does not inherit MJIT threads.  */
-void
-mjit_child_after_fork(void)
-{
-    if (mjit_enabled) {
-        verbose(3, "Switching off MJIT in a forked child");
-        mjit_enabled = FALSE;
-    }
-    /* TODO: Should we initiate MJIT in the forked Ruby.  */
-}
-
 static enum rb_id_table_iterator_result
 valid_class_serials_add_i(ID key, VALUE v, void *unused)
 {
@@ -661,6 +649,7 @@ mjit_init(struct mjit_options *opts)
         verbose(1, "Failure in MJIT header file name initialization\n");
         return;
     }
+    pch_owner_pid = getpid();
 
     init_list(&unit_queue);
     init_list(&active_units);
@@ -748,6 +737,54 @@ mjit_resume(void)
     return Qtrue;
 }
 
+/* Skip calling `clean_object_files` for units which currently exist in the list. */
+static void
+skip_cleaning_object_files(struct rb_mjit_unit_list *list)
+{
+    struct rb_mjit_unit *unit = NULL, *next;
+
+    /* No mutex for list, assuming MJIT worker does not exist yet since it's immediately after fork. */
+    list_for_each_safe(&list->head, unit, next, unode) {
+#ifndef _MSC_VER /* Actualy mswin does not reach here since it doesn't have fork */
+        if (unit->o_file) unit->o_file_inherited_p = TRUE;
+#endif
+
+#if defined(_WIN32) /* mswin doesn't reach here either. This is for MinGW. */
+        if (unit->so_file) unit->so_file = NULL;
+#endif
+    }
+}
+
+/* This is called after fork initiated by Ruby's method to launch MJIT worker thread
+   for child Ruby process.
+
+   In multi-process Ruby applications, child Ruby processes do most of the jobs.
+   Thus we want child Ruby processes to enqueue ISeqs to MJIT worker's queue and
+   call the JIT-ed code.
+
+   But unfortunately current MJIT-generated code is process-specific. After the fork,
+   JIT-ed code created by parent Ruby process cannnot be used in child Ruby process
+   because the code could rely on inline cache values (ivar's IC, send's CC) which
+   may vary between processes after fork or embed some process-specific addresses.
+
+   So child Ruby process can't request parent process to JIT an ISeq and use the code.
+   Instead of that, MJIT worker thread is created for all child Ruby processes, even
+   while child processes would end up with compiling the same ISeqs.
+ */
+void
+mjit_child_after_fork(void)
+{
+    if (!mjit_enabled)
+        return;
+
+    /* Let parent process delete the already-compiled object files.
+       This must be done before starting MJIT worker on child process. */
+    skip_cleaning_object_files(&active_units);
+
+    /* MJIT worker thread is not inherited on fork. Start it for this child process. */
+    start_worker();
+}
+
 /* Finish the threads processing units and creating PCH, finalize
    and free MJIT data.  It should be called last during MJIT
    life.  */
@@ -781,7 +818,7 @@ mjit_finish(void)
     rb_native_cond_destroy(&mjit_gc_wakeup);
 
 #ifndef _MSC_VER /* mswin has prebuilt precompiled header */
-    if (!mjit_opts.save_temps)
+    if (!mjit_opts.save_temps && getpid() == pch_owner_pid)
         remove_file(pch_file);
 
     xfree(header_file); header_file = NULL;
diff --git a/mjit_worker.c b/mjit_worker.c
index b0dddc9f22..d85afb5ebe 100644
--- a/mjit_worker.c
+++ b/mjit_worker.c
@@ -132,6 +132,10 @@ struct rb_mjit_unit {
 #ifndef _MSC_VER
     /* This value is always set for `compact_all_jit_code`. Also used for lazy deletion. */
     char *o_file;
+    /* TRUE if it's inherited from parent Ruby process and lazy deletion should be skipped.
+       `o_file = NULL` can't be used to skip lazy deletion because `o_file` could be used
+       by child for `compact_all_jit_code`. */
+    int o_file_inherited_p;
 #endif
 #if defined(_WIN32)
     /* DLL cannot be removed while loaded on Windows. If this is set, it'll be lazily deleted. */
@@ -213,6 +217,8 @@ static VALUE valid_class_serials;
 static const char *cc_path;
 /* Name of the precompiled header file.  */
 static char *pch_file;
+/* The process id which should delete the pch_file on mjit_finish. */
+static rb_pid_t pch_owner_pid;
 /* Status of the precompiled header creation.  The status is
    shared by the workers and the pch thread.  */
 static enum {PCH_NOT_READY, PCH_FAILED, PCH_SUCCESS} pch_status;
@@ -347,7 +353,7 @@ clean_object_files(struct rb_mjit_unit *unit)
         unit->o_file = NULL;
         /* For compaction, unit->o_file is always set when compilation succeeds.
            So save_temps needs to be checked here. */
-        if (!mjit_opts.save_temps)
+        if (!mjit_opts.save_temps && !unit->o_file_inherited_p)
             remove_file(o_file);
         free(o_file);
     }
diff --git a/process.c b/process.c
index fafcf11ad2..5f1d3891e8 100644
--- a/process.c
+++ b/process.c
@@ -1502,12 +1502,26 @@ after_exec(void)
 }
 
 #if defined HAVE_WORKING_FORK || defined HAVE_DAEMON
-#define before_fork_ruby() before_exec()
 static void
-after_fork_ruby(void)
+before_fork_ruby(void)
+{
+    if (mjit_enabled) {
+        /* avoid leaving locked mutex and units being modified for child process. */
+        mjit_pause(FALSE);
+    }
+
+    before_exec();
+}
+
+static void
+after_fork_ruby(int parent_p)
 {
     rb_threadptr_pending_interrupt_clear(GET_THREAD());
     after_exec();
+
+    if (mjit_enabled && parent_p) { /* child is cared by `rb_thread_atfork` */
+        mjit_resume();
+    }
 }
 #endif
 
@@ -3997,7 +4011,7 @@ rb_fork_ruby(int *status)
 	before_fork_ruby();
 	pid = fork();
 	err = errno;
-	after_fork_ruby();
+	after_fork_ruby(pid > 0);
 	disable_child_handler_fork_parent(&old); /* yes, bad name */
 	if (pid >= 0) /* fork succeed */
 	    return pid;
@@ -6422,7 +6436,7 @@ rb_daemon(int nochdir, int noclose)
 #ifdef HAVE_DAEMON
     before_fork_ruby();
     err = daemon(nochdir, noclose);
-    after_fork_ruby();
+    after_fork_ruby(TRUE);
     rb_thread_atfork();
 #else
     int n;
diff --git a/test/ruby/test_jit.rb b/test/ruby/test_jit.rb
index 451ed91a56..b25a9b15a2 100644
--- a/test/ruby/test_jit.rb
+++ b/test/ruby/test_jit.rb
@@ -567,7 +567,7 @@ class TestJIT < Test::Unit::TestCase
     assert_match(/^Successful MJIT finish$/, err)
   end
 
-  def test_unload_units
+  def test_unload_units_and_compaction
     Dir.mktmpdir("jit_test_unload_units_") do |dir|
       # MIN_CACHE_SIZE is 10
       out, err = eval_with_jit({"TMPDIR"=>dir}, "#{<<~"begin;"}\n#{<<~'end;'}", verbose: 1, min_calls: 1, max_cache: 10)
@@ -582,6 +582,12 @@ class TestJIT < Test::Unit::TestCase
           EOS
           i += 1
         end
+
+        if defined?(fork)
+          # test the child does not try to delete files which are deleted by parent,
+          # and test possible deadlock on fork during MJIT unload and JIT compaction on child
+          Process.waitpid(Process.fork {})
+        end
       end;
 
       debug_info = "stdout:\n```\n#{out}\n```\n\nstderr:\n```\n#{err}```\n"
@@ -598,7 +604,7 @@ class TestJIT < Test::Unit::TestCase
       # On --jit-wait, when the number of JIT-ed code reaches --jit-max-cache,
       # it should trigger compaction.
       unless RUBY_PLATFORM.match?(/mswin|mingw/) # compaction is not supported on Windows yet
-        assert_equal(2, compactions.size, debug_info)
+        assert_equal(3, compactions.size, debug_info)
       end
 
       if appveyor_mswin?
@@ -838,6 +844,36 @@ class TestJIT < Test::Unit::TestCase
     assert_equal("-e:8:in `a'\n", lines[1])
   end
 
+  def test_fork_with_mjit_worker_thread
+    Dir.mktmpdir("jit_test_fork_with_mjit_worker_thread_") do |dir|
+      # min_calls: 2 to skip fork block
+      out, err = eval_with_jit({ "TMPDIR" => dir }, "#{<<~"begin;"}\n#{<<~"end;"}", min_calls: 2, verbose: 1)
+      begin;
+        def before_fork; end
+        def after_fork; end
+
+        before_fork; before_fork # the child should not delete this .o file
+        pid = Process.fork do # this child should not delete shared .pch file
+          after_fork; after_fork # this child does not share JIT-ed after_fork with parent
+        end
+        after_fork; after_fork # this parent does not share JIT-ed after_fork with child
+
+        Process.waitpid(pid)
+      end;
+      success_count = err.scan(/^#{JIT_SUCCESS_PREFIX}:/).size
+      assert_equal(3, success_count)
+
+      # assert no remove error
+      lines = err.lines
+      assert_match(/^Successful MJIT finish$/, lines[3])
+      assert_match(/^Successful MJIT finish$/, lines[4])
+
+      # ensure objects are deleted
+      debug_info = "stdout:\n```\n#{out}\n```\n\nstderr:\n```\n#{err}```\n"
+      assert_send([Dir, :empty?, dir], debug_info)
+    end
+  end if defined?(fork)
+
   private
 
   def appveyor_mswin?
diff --git a/thread.c b/thread.c
index 37950f0139..8420808c40 100644
--- a/thread.c
+++ b/thread.c
@@ -4438,6 +4438,8 @@ rb_thread_atfork(void)
 
     /* We don't want reproduce CVE-2003-0900. */
     rb_reset_random_seed();
+
+    /* For child, starting MJIT worker thread in this place which is safer than `after_fork_ruby`. */
     mjit_child_after_fork();
 }