process.c: Add Process._fork (#5017)

* process.c: Add Process._fork

This API is supposed for application monitoring libraries to hook fork
event.

[Feature #17795]

Co-authored-by: Nobuyoshi Nakada <nobu@ruby-lang.org>

5 files changed, 136 insertions, 11 deletions

diff --git a/common.mk b/common.mk
index debef394b8..c3c599d24d 100644
--- a/common.mk
+++ b/common.mk
@@ -10360,15 +10360,18 @@ process.$(OBJEXT): $(hdrdir)/ruby.h
 process.$(OBJEXT): $(hdrdir)/ruby/ruby.h
 process.$(OBJEXT): $(top_srcdir)/include/ruby/fiber/scheduler.h
 process.$(OBJEXT): $(top_srcdir)/internal/array.h
+process.$(OBJEXT): $(top_srcdir)/internal/bignum.h
 process.$(OBJEXT): $(top_srcdir)/internal/bits.h
 process.$(OBJEXT): $(top_srcdir)/internal/class.h
 process.$(OBJEXT): $(top_srcdir)/internal/compilers.h
 process.$(OBJEXT): $(top_srcdir)/internal/dir.h
 process.$(OBJEXT): $(top_srcdir)/internal/error.h
 process.$(OBJEXT): $(top_srcdir)/internal/eval.h
+process.$(OBJEXT): $(top_srcdir)/internal/fixnum.h
 process.$(OBJEXT): $(top_srcdir)/internal/gc.h
 process.$(OBJEXT): $(top_srcdir)/internal/hash.h
 process.$(OBJEXT): $(top_srcdir)/internal/imemo.h
+process.$(OBJEXT): $(top_srcdir)/internal/numeric.h
 process.$(OBJEXT): $(top_srcdir)/internal/object.h
 process.$(OBJEXT): $(top_srcdir)/internal/process.h
 process.$(OBJEXT): $(top_srcdir)/internal/serial.h
diff --git a/internal/process.h b/internal/process.h
index 75e7db5b42..ceadfdcbbb 100644
--- a/internal/process.h
+++ b/internal/process.h
@@ -75,7 +75,7 @@ struct rb_execarg {
 };
 
 /* process.c */
-rb_pid_t rb_fork_ruby(int *status);
+rb_pid_t rb_call_proc__fork(void);
 void rb_last_status_clear(void);
 static inline char **ARGVSTR2ARGV(VALUE argv_str);
 static inline size_t ARGVSTR2ARGC(VALUE argv_str);
diff --git a/io.c b/io.c
index 4cd8a4fa5d..fceb7019f1 100644
--- a/io.c
+++ b/io.c
@@ -6897,7 +6897,7 @@ pipe_open(VALUE execarg_obj, const char *modestr, int fmode,
     }
     else {
 # if defined(HAVE_WORKING_FORK)
-	pid = rb_fork_ruby(&status);
+	pid = rb_call_proc__fork();
 	if (pid == 0) {		/* child */
 	    popen_redirect(&arg);
 	    rb_io_synchronized(RFILE(orig_stdout)->fptr);
diff --git a/process.c b/process.c
index 5e37a29460..6b83767738 100644
--- a/process.c
+++ b/process.c
@@ -103,6 +103,7 @@ int initgroups(const char *, rb_gid_t);
 #include "internal/error.h"
 #include "internal/eval.h"
 #include "internal/hash.h"
+#include "internal/numeric.h"
 #include "internal/object.h"
 #include "internal/process.h"
 #include "internal/thread.h"
@@ -4345,11 +4346,42 @@ rb_fork_ruby(int *status)
     return pid;
 }
 
+rb_pid_t
+rb_call_proc__fork(void)
+{
+    VALUE pid = rb_funcall(rb_mProcess, rb_intern("_fork"), 0);
+
+    return NUM2PIDT(pid);
+}
 #endif
 
 #if defined(HAVE_WORKING_FORK) && !defined(CANNOT_FORK_WITH_PTHREAD)
 /*
  *  call-seq:
+ *     Process._fork   -> integer
+ *
+ *  An internal API for fork. Do not call this method directly.
+ *  Currently, this is called via +Kernel.#fork+, +Process.fork+, and
+ *  +popen+ with +"-"+.
+ *
+ *  This method is not for casual code but for application monitoring
+ *  libraries. You can add custom code before and after fork events
+ *  by overriding this method.
+ */
+VALUE
+rb_proc__fork(VALUE _obj)
+{
+    rb_pid_t pid = rb_fork_ruby(NULL);
+
+    if (pid == -1) {
+	rb_sys_fail("fork(2)");
+    }
+
+    return PIDT2NUM(pid);
+}
+
+/*
+ *  call-seq:
  *     Kernel.fork  [{ block }]   -> integer or nil
  *     Process.fork [{ block }]   -> integer or nil
  *
@@ -4378,24 +4410,21 @@ rb_f_fork(VALUE obj)
 {
     rb_pid_t pid;
 
-    switch (pid = rb_fork_ruby(NULL)) {
-      case 0:
+    pid = rb_call_proc__fork();
+
+    if (pid == 0) {
 	if (rb_block_given_p()) {
 	    int status;
 	    rb_protect(rb_yield, Qundef, &status);
 	    ruby_stop(status);
 	}
 	return Qnil;
-
-      case -1:
-	rb_sys_fail("fork(2)");
-	return Qnil;
-
-      default:
-	return PIDT2NUM(pid);
     }
+
+    return PIDT2NUM(pid);
 }
 #else
+#define rb_proc__fork rb_f_notimplement
 #define rb_f_fork rb_f_notimplement
 #endif
 
@@ -8700,6 +8729,7 @@ InitVM_process(void)
     rb_define_singleton_method(rb_mProcess, "exit", f_exit, -1);
     rb_define_singleton_method(rb_mProcess, "abort", f_abort, -1);
     rb_define_singleton_method(rb_mProcess, "last_status", proc_s_last_status, 0);
+    rb_define_singleton_method(rb_mProcess, "_fork", rb_proc__fork, 0);
 
     rb_define_module_function(rb_mProcess, "kill", proc_rb_f_kill, -1);
     rb_define_module_function(rb_mProcess, "wait", proc_m_wait, -1);
diff --git a/test/ruby/test_process.rb b/test/ruby/test_process.rb
index 98c934945c..636871f822 100644
--- a/test/ruby/test_process.rb
+++ b/test/ruby/test_process.rb
@@ -2545,4 +2545,96 @@ EOS
     end
     assert_empty(Process.waitall)
   end
+
+  def test__fork
+    r, w = IO.pipe
+    pid = Process._fork
+    if pid == 0
+      begin
+        r.close
+        w << "ok: #$$"
+        w.close
+      ensure
+        exit!
+      end
+    else
+      w.close
+      assert_equal("ok: #{pid}", r.read)
+      r.close
+      Process.waitpid(pid)
+    end
+  end if Process.respond_to?(:_fork)
+
+  def test__fork_hook
+    %w(fork Process.fork).each do |method|
+      feature17795 = '[ruby-core:103400] [Feature #17795]'
+      assert_in_out_err([], <<-"end;", [], [], feature17795, timeout: 60) do |r, e|
+        module ForkHook
+          def _fork
+            p :before
+            ret = super
+            p :after
+            ret
+          end
+        end
+
+        Process.singleton_class.prepend(ForkHook)
+
+        pid = #{ method }
+        p pid
+        Process.waitpid(pid) if pid
+      end;
+        assert_equal([], e)
+        assert_equal(":before", r.shift)
+        assert_equal(":after", r.shift)
+        s = r.map {|s| s.chomp }.sort #=> [pid, ":after", "nil"]
+        assert_match(/^\d+$/, s[0]) # pid
+        assert_equal(":after", s[1])
+        assert_equal("nil", s[2])
+      end
+    end
+  end if Process.respond_to?(:_fork)
+
+  def test__fork_hook_popen
+    feature17795 = '[ruby-core:103400] [Feature #17795]'
+    assert_in_out_err([], <<-"end;", %w(:before :after :after foo bar), [], feature17795, timeout: 60)
+      module ForkHook
+        def _fork
+          p :before
+          ret = super
+          p :after
+          ret
+        end
+      end
+
+      Process.singleton_class.prepend(ForkHook)
+
+      IO.popen("-") {|io|
+        if !io
+          puts "foo"
+        else
+          puts io.read + "bar"
+        end
+      }
+    end;
+  end if Process.respond_to?(:_fork)
+
+  def test__fork_wrong_type_hook
+    feature17795 = '[ruby-core:103400] [Feature #17795]'
+    assert_in_out_err([], <<-"end;", ["OK"], [], feature17795, timeout: 60)
+      module ForkHook
+        def _fork
+          "BOO"
+        end
+      end
+
+      Process.singleton_class.prepend(ForkHook)
+
+      begin
+        fork
+      rescue TypeError
+        puts "OK"
+      end
+    end;
+  end if Process.respond_to?(:_fork)
 end