subprocess: Support 'preexec_fn' and 'restore_signals' arguments

1st1 · 1st1 · commit b5b3b12dbba9 · 2016-06-08T17:18:06.000-04:00
diff --git a/tests/test_process.py b/tests/test_process.py
@@ -1,9 +1,11 @@
 import asyncio
 import contextlib
+import os
 import signal
 import subprocess
 import sys
 import tempfile
+import time
 
 from asyncio import test_utils
 from uvloop import _testbase as tb
@@ -46,6 +48,59 @@ async def test():
 
         self.loop.run_until_complete(test())
 
+    def test_process_preexec_fn_1(self):
+        # Copied from CPython/test_suprocess.py
+
+        # DISCLAIMER: Setting environment variables is *not* a good use
+        # of a preexec_fn.  This is merely a test.
+
+        async def test():
+            cmd = sys.executable
+            proc = await asyncio.create_subprocess_exec(
+                cmd, '-c',
+                'import os,sys;sys.stdout.write(os.getenv("FRUIT"))',
+                stdout=subprocess.PIPE,
+                preexec_fn=lambda: os.putenv("FRUIT", "apple"),
+                loop=self.loop)
+
+            out, _ = await proc.communicate()
+            self.assertEqual(out, b'apple')
+            self.assertEqual(proc.returncode, 0)
+
+        self.loop.run_until_complete(test())
+
+    def test_process_preexec_fn_2(self):
+        # Copied from CPython/test_suprocess.py
+
+        def raise_it():
+            raise ValueError("spam")
+
+        async def test():
+            cmd = sys.executable
+            proc = await asyncio.create_subprocess_exec(
+                cmd, '-c', 'import time; time.sleep(10)',
+                preexec_fn=raise_it,
+                loop=self.loop)
+
+            await proc.communicate()
+
+        started = time.time()
+        try:
+            self.loop.run_until_complete(test())
+        except subprocess.SubprocessError as ex:
+            self.assertIn('preexec_fn', ex.args[0])
+            if ex.__cause__ is not None:
+                # uvloop will set __cause__
+                self.assertIs(type(ex.__cause__), ValueError)
+                self.assertEqual(ex.__cause__.args[0], 'spam')
+        else:
+            self.fail(
+                'exception in preexec_fn did not propagate to the parent')
+
+        if time.time() - started > 5:
+            self.fail(
+                'exception in preexec_fn did not kill the child process')
+
     def test_process_executable_1(self):
         async def test():
             proc = await asyncio.create_subprocess_exec(
diff --git a/uvloop/handles/process.pxd b/uvloop/handles/process.pxd
@@ -3,6 +3,11 @@ cdef class UVProcess(UVHandle):
         object _returncode
         object _pid
 
+        object _errpipe_read
+        object _errpipe_write
+        object _preexec_fn
+        bint _restore_signals
+
         set _fds_to_close
 
         # Attributes used to compose uv_process_options_t:
@@ -18,7 +23,9 @@ cdef class UVProcess(UVHandle):
     cdef _init(self, Loop loop, list args, dict env, cwd,
                start_new_session,
                _stdin, _stdout, _stderr, pass_fds,
-               debug_flags)
+               debug_flags, preexec_fn, restore_signals)
+
+    cdef _after_fork(self)
 
     cdef char** __to_cstring_array(self, list arr)
     cdef _init_args(self, list args)
@@ -67,4 +74,5 @@ cdef class UVProcessTransport(UVProcess):
                                 start_new_session,
                                 _stdin, _stdout, _stderr, pass_fds,
                                 waiter,
-                                debug_flags)
+                                debug_flags,
+                                preexec_fn, restore_signals)
diff --git a/uvloop/handles/process.pyx b/uvloop/handles/process.pyx
@@ -8,11 +8,16 @@ cdef class UVProcess(UVHandle):
         self._returncode = None
         self._pid = None
         self._fds_to_close = set()
+        self._preexec_fn = None
+        self._restore_signals = True
 
     cdef _init(self, Loop loop, list args, dict env,
                cwd, start_new_session,
                _stdin, _stdout, _stderr,  # std* can be defined as macros in C
-               pass_fds, debug_flags):
+               pass_fds, debug_flags, preexec_fn, restore_signals):
+
+        global __forking
+        global __forking_loop
 
         cdef int err
 
@@ -43,14 +48,62 @@ cdef class UVProcess(UVHandle):
             self._abort_init()
             raise
 
-        err = uv.uv_spawn(loop.uvloop,
-                          <uv.uv_process_t*>self._handle,
-                          &self.options)
-        if err < 0:
+        if __forking or loop.active_process_handler is not None:
+            # Our pthread_atfork handlers won't work correctly when
+            # another loop is forking in another thread (even though
+            # GIL should help us to avoid that.)
             self._abort_init()
-            raise convert_error(err)
+            raise RuntimeError(
+                'Racing with another loop to spawn a process.')
+
+        self._errpipe_read, self._errpipe_write = os_pipe()
+        try:
+            os_set_inheritable(self._errpipe_write, True)
+
+            self._preexec_fn = preexec_fn
+            self._restore_signals = restore_signals
+
+            loop.active_process_handler = self
+            __forking = 1
+            __forking_loop = loop
+
+            _PyImport_AcquireLock()
+
+            err = uv.uv_spawn(loop.uvloop,
+                              <uv.uv_process_t*>self._handle,
+                              &self.options)
 
-        self._finish_init()
+            __forking = 0
+            __forking_loop = None
+            loop.active_process_handler = None
+
+            if _PyImport_ReleaseLock() < 0:
+                # See CPython/posixmodule.c for details
+                self._abort_init()
+                raise RuntimeError('not holding the import lock')
+
+            if err < 0:
+                self._abort_init()
+                raise convert_error(err)
+
+            self._finish_init()
+
+            os_close(self._errpipe_write)
+
+            errpipe_data = bytearray()
+            while True:
+                part = os_read(self._errpipe_read, 50000)
+                errpipe_data += part
+                if not part or len(errpipe_data) > 50000:
+                    break
+
+        finally:
+            os_close(self._errpipe_read)
+            try:
+                os_close(self._errpipe_write)
+            except OSError:
+                # Might be already closed
+                pass
 
         # asyncio caches the PID in BaseSubprocessTransport,
         # so that the transport knows what the PID was even
@@ -68,6 +121,52 @@ cdef class UVProcess(UVHandle):
         if debug_flags & __PROCESS_DEBUG_SLEEP_AFTER_FORK:
             time_sleep(1)
 
+        if errpipe_data:
+            # preexec_fn has raised an exception.  The child
+            # process must be dead now.
+            try:
+                exc_name, exc_msg = errpipe_data.split(b':', 1)
+                exc_name = exc_name.decode()
+                exc_msg = exc_msg.decode()
+            except:
+                self._close()
+                raise subprocess_SubprocessError(
+                    'Bad exception data from child: {!r}'.format(
+                        errpipe_data))
+            exc_cls = getattr(__builtins__, exc_name,
+                              subprocess_SubprocessError)
+
+            exc = subprocess_SubprocessError(
+                'Exception occurred in preexec_fn.')
+            exc.__cause__ = exc_cls(exc_msg)
+            self._close()
+            raise exc
+
+    cdef _after_fork(self):
+        # See CPython/_posixsubprocess.c for details
+
+        if self._restore_signals:
+            _Py_RestoreSignals()
+
+        if self._preexec_fn is not None:
+            PyOS_AfterFork()
+
+            try:
+                gc_disable()
+                self._preexec_fn()
+            except BaseException as ex:
+                try:
+                    with open(self._errpipe_write, 'wb') as f:
+                        f.write(str(ex.__class__.__name__).encode())
+                        f.write(b':')
+                        f.write(str(ex.args[0]).encode())
+                finally:
+                    system._exit(255)
+            else:
+                os_close(self._errpipe_write)
+        else:
+            os_close(self._errpipe_write)
+
     cdef _close_after_spawn(self, int fd):
         if self._fds_to_close is None:
             raise RuntimeError(
@@ -438,7 +537,9 @@ cdef class UVProcessTransport(UVProcess):
                                 cwd, start_new_session,
                                 _stdin, _stdout, _stderr, pass_fds,
                                 waiter,
-                                debug_flags):
+                                debug_flags,
+                                preexec_fn,
+                                restore_signals):
 
         cdef UVProcessTransport handle
         handle = UVProcessTransport.__new__(UVProcessTransport)
@@ -448,7 +549,9 @@ cdef class UVProcessTransport(UVProcess):
                      __process_convert_fileno(_stdout),
                      __process_convert_fileno(_stderr),
                      pass_fds,
-                     debug_flags)
+                     debug_flags,
+                     preexec_fn,
+                     restore_signals)
 
         if handle._init_futs:
             handle._stdio_ready = 0
diff --git a/uvloop/includes/__init__.py b/uvloop/includes/__init__.py
@@ -3,6 +3,7 @@
 import collections
 import concurrent.futures
 import functools
+import gc
 import inspect
 import itertools
 import os
diff --git a/uvloop/includes/python.pxd b/uvloop/includes/python.pxd
@@ -6,3 +6,8 @@ cdef extern from "Python.h":
 
     object PyUnicode_EncodeFSDefault(object)
     void PyErr_SetInterrupt() nogil
+
+    void PyOS_AfterFork()
+    void _PyImport_AcquireLock()
+    int _PyImport_ReleaseLock()
+    void _Py_RestoreSignals()
diff --git a/uvloop/includes/stdlib.pxi b/uvloop/includes/stdlib.pxi
@@ -4,6 +4,7 @@ import asyncio, asyncio.log, asyncio.base_events, \
 import collections
 import concurrent.futures
 import functools
+import gc
 import inspect
 import itertools
 import os
@@ -44,6 +45,8 @@ cdef cc_Future = concurrent.futures.Future
 
 cdef ft_partial = functools.partial
 
+cdef gc_disable = gc.disable
+
 cdef iter_chain = itertools.chain
 cdef inspect_isgenerator = inspect.isgenerator
 
@@ -88,6 +91,8 @@ cdef os_close = os.close
 cdef os_open = os.open
 cdef os_devnull = os.devnull
 cdef os_O_RDWR = os.O_RDWR
+cdef os_pipe = os.pipe
+cdef os_read = os.read
 
 cdef sys_ignore_environment = sys.flags.ignore_environment
 cdef sys_exc_info = sys.exc_info
@@ -102,6 +107,7 @@ cdef long MAIN_THREAD_ID = <long>threading.main_thread().ident
 cdef int subprocess_PIPE = subprocess.PIPE
 cdef int subprocess_STDOUT = subprocess.STDOUT
 cdef int subprocess_DEVNULL = subprocess.DEVNULL
+cdef subprocess_SubprocessError = subprocess.SubprocessError
 
 cdef int signal_NSIG = std_signal.NSIG
 
diff --git a/uvloop/includes/system.pxd b/uvloop/includes/system.pxd
@@ -48,6 +48,15 @@ cdef extern from "sys/socket.h" nogil:
 cdef extern from "unistd.h" nogil:
 
     ssize_t write(int fd, const void *buf, size_t count)
+    void _exit(int status)
+
+
+cdef extern from "pthread.h" nogil:
+
+    int pthread_atfork(
+        void (*prepare)() nogil,
+        void (*parent)() nogil,
+        void (*child)() nogil)
 
 
 cdef extern from "includes/compat.h" nogil:
diff --git a/uvloop/loop.pxd b/uvloop/loop.pxd
@@ -63,6 +63,8 @@ cdef class Loop:
         dict _signal_handlers
         bint _custom_sigint
 
+        UVProcess active_process_handler
+
         UVAsync handler_async
         UVIdle handler_idle
         UVCheck handler_check__exec_writes
diff --git a/uvloop/loop.pyx b/uvloop/loop.pyx
