Instrument Multiprocessing forkserver

This depends on gevent/gevent#2126 being fixed.

Update documentation

Author: arcivanov

.idea/geventmp.iml

@@ -121,11 +121,6 @@ Supported Platforms
 * CPython 3.9, 3.10, 3.11, 3.12, 3.13
 * PyPy 3.9, 3.10
 
-Known Issues
-============
-
-* Multiprocessing `forkserver` works in GeventMP_, but the spawned child isn't green.
-
 TODO
 ====
 1. Monkey patch Windows to the extent possible.

.idea/misc.xml

@@ -73,6 +73,9 @@ def test_manager_fork(self):
     def test_manager_spawn(self):
         self.run_manager_test("spawn")
 
+    def test_manager_forkserver(self):
+        self.run_manager_test("forkserver")
+
     def run_manager_test(self, context, do_trace=False, remote_trace=False):
         ctx = mp.get_context(context)
         if do_trace:

README.rst

@@ -59,12 +59,18 @@ def test_mp_queues_fork(self):
     def test_mp_queues_spawn(self):
         self.run_test_mp_queues("spawn", _mp_test_gevent.test_queues)
 
+    def test_mp_queues_forkserver(self):
+        self.run_test_mp_queues("forkserver", _mp_test_gevent.test_queues)
+
     def test_mp_jqueues_fork(self):
         self.run_test_mp_jqueues("fork", _mp_test_gevent.test_queues)
 
     def test_mp_jqueues_spawn(self):
         self.run_test_mp_jqueues("spawn", _mp_test_gevent.test_queues)
 
+    def test_mp_jqueues_forkserver(self):
+        self.run_test_mp_jqueues("forkserver", _mp_test_gevent.test_queues)
+
     def test_mp_no_args_fork(self):
         self.run_test_mp_no_args("fork", _mp_test_gevent.test_no_args)
 

src/integrationtest/python/monkey_manager_tests.py

@@ -13,12 +13,14 @@
 #   See the License for the specific language governing permissions and
 #   limitations under the License.
 
+import os
+import signal
 import sys
 from multiprocessing.forkserver import ForkServer as _ForkServer
 
 from gevent.os import make_nonblocking, nb_read, nb_write
 
-__implements__ = ["ForkServer", "_forkserver", "ensure_running",
+__implements__ = ["ForkServer", "_forkserver", "ensure_running", "_serve_one",
                   "get_inherited_fds", "connect_to_new_process",
                   "set_forkserver_preload"]
 __target__ = "multiprocessing.forkserver"
@@ -85,6 +87,32 @@ def write_unsigned(fd, n):
                 raise RuntimeError('should not get here')
             msg = msg[nbytes:]
 
+
+    def _serve_one(child_r, fds, unused_fds, handlers):
+        from multiprocessing import resource_tracker, spawn
+        # close unnecessary stuff and reset signal handlers
+        signal.set_wakeup_fd(-1)
+        for sig, val in handlers.items():
+            signal.signal(sig, val)
+        for fd in unused_fds:
+            os.close(fd)
+
+        (_forkserver._forkserver_alive_fd,
+         resource_tracker._resource_tracker._fd,
+         *_forkserver._inherited_fds) = fds
+        make_nonblocking(_forkserver._forkserver_alive_fd)
+        make_nonblocking(resource_tracker._resource_tracker._fd)
+        for ifd in _forkserver._inherited_fds:
+            make_nonblocking(ifd)
+
+        # Run process object received over pipe
+        parent_sentinel = os.dup(child_r)
+        make_nonblocking(child_r)
+        make_nonblocking(parent_sentinel)
+        code = spawn._main(child_r, parent_sentinel)
+
+        return code
+
 _forkserver = ForkServer()
 ensure_running = _forkserver.ensure_running
 get_inherited_fds = _forkserver.get_inherited_fds

src/integrationtest/python/monkey_mp_tests.py

@@ -0,0 +1,31 @@
+#   -*- coding: utf-8 -*-
+#   Copyright 2022 Karellen, Inc. and contributors
+#
+#   Licensed under the Apache License, Version 2.0 (the "License");
+#   you may not use this file except in compliance with the License.
+#   You may obtain a copy of the License at
+#
+#       http://www.apache.org/licenses/LICENSE-2.0
+#
+#   Unless required by applicable law or agreed to in writing, software
+#   distributed under the License is distributed on an "AS IS" BASIS,
+#   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#   See the License for the specific language governing permissions and
+#   limitations under the License.
+
+from multiprocessing.popen_forkserver import Popen as _Popen
+
+from gevent.os import make_nonblocking
+
+__implements__ = ["Popen"]
+__target__ = "multiprocessing.popen_forkserver"
+
+
+class Popen(_Popen):
+    def _launch(self, process_obj):
+        self.sentinel = None
+        try:
+            super()._launch(process_obj)
+        finally:
+            if self.sentinel is not None:
+                make_nonblocking(self.sentinel)

src/main/python/geventmp/_mp/3/_mp_forkserver.py

@@ -28,20 +28,14 @@ def get_command_line(**kwds):
         return ([sys.executable, '--multiprocessing-fork'] +
                 ['%s=%r' % item for item in kwds.items()])
     else:
-        from gevent.monkey import saved
-        from gevent import config
-        from gevent._config import ImportableSetting
-        from geventmp.monkey import GEVENT_SAVED_MODULE_SETTINGS
-
-        prog = 'from gevent import monkey; monkey.patch_all(**%r); ' + \
-               'from gevent import config; [setattr(config, k, v) for k, v in %r.items()]; ' + \
-               'from multiprocessing.spawn import spawn_main; ' + \
-               'spawn_main(%s);'
-
-        prog %= (saved[GEVENT_SAVED_MODULE_SETTINGS],
-                 {k: getattr(config, k) for k in dir(config)
-                  if not isinstance(config.settings[k], ImportableSetting)},
-                 ', '.join('%s=%r' % item for item in kwds.items()))
+        prog, args = util.get_command_line_gevent_preamble()
+
+        prog += 'from multiprocessing.spawn import spawn_main; ' + \
+                'spawn_main(%s);'
+
+        args += (', '.join('%s=%r' % item for item in kwds.items()),)
+
+        prog %= args
 
         opts = util._args_from_interpreter_flags()
         return [spawn._python_exe] + opts + ['-c', prog, '--multiprocessing-fork']

src/main/python/geventmp/_mp/3/_mp_popen_forkserver.py

@@ -13,15 +13,39 @@
 #   See the License for the specific language governing permissions and
 #   limitations under the License.
 
+from multiprocessing.util import spawnv_passfds as _spawnv_passfd, register_after_fork
+
 from gevent.os import _watch_child
 from gevent.threading import local
-from multiprocessing.util import spawnv_passfds as _spawnv_passfd, register_after_fork
 
-__implements__ = ["spawnv_passfds", "ForkAwareLocal"]
+__implements__ = ["spawnv_passfds", "ForkAwareLocal", "get_command_line_gevent_preamble"]
 __target__ = "multiprocessing.util"
 
 
+def get_command_line_gevent_preamble():
+    from gevent.monkey import saved
+    from gevent import config
+    from gevent._config import ImportableSetting
+    from geventmp.monkey import GEVENT_SAVED_MODULE_SETTINGS
+
+    prog = 'from gevent import monkey; monkey.patch_all(**%r); ' + \
+           'from gevent import config; [setattr(config, k, v) for k, v in %r.items()]; '
+
+    args = (saved[GEVENT_SAVED_MODULE_SETTINGS],
+            {k: getattr(config, k) for k in dir(config)
+             if not isinstance(config.settings[k], ImportableSetting)})
+
+    return prog, args
+
+
 def spawnv_passfds(path, args, passfds):
+    launch_args = args[2] if len(args) >= 3 else None
+    if launch_args and launch_args.startswith("from multiprocessing.forkserver import main;"):
+        prog, prog_args = get_command_line_gevent_preamble()
+        prog += "%s"
+        launch_args = prog % (prog_args + (launch_args,))
+        args[2] = launch_args
+
     cpid = _spawnv_passfd(path, args, passfds)
     _watch_child(cpid)
     return cpid

src/main/python/geventmp/_mp/3/_mp_spawn.py

@@ -72,6 +72,7 @@ def _patch_mp(will_patch_all):
         _patch_module("_mp.3._mp_connection", _package_prefix='geventmp.')
         _patch_module("_mp.3._mp_synchronize", _package_prefix='geventmp.')
         _patch_module("_mp.3._mp_popen_fork", _package_prefix='geventmp.')
+        _patch_module("_mp.3._mp_popen_forkserver", _package_prefix='geventmp.')
         _patch_module("_mp.3._mp_popen_spawn_posix", _package_prefix='geventmp.')
         _patch_module("_mp.3._mp_forkserver", _package_prefix='geventmp.')
         _patch_module("_mp.3._mp_resource_tracker", _package_prefix='geventmp.')