From 7a812ab7826b17770ebda96941aeaa10ad543ddf Mon Sep 17 00:00:00 2001 From: Benjamin Peterson Date: Mon, 4 Nov 2019 21:39:50 -0800 Subject: [PATCH] asyncio: Add a pidfd child process watcher. --- Lib/asyncio/unix_events.py | 57 ++++++++++++++++++++++++ Lib/test/test_asyncio/test_subprocess.py | 5 +++ Modules/clinic/posixmodule.c.h | 39 +++++++++++++++- Modules/posixmodule.c | 22 +++++++++ 4 files changed, 122 insertions(+), 1 deletion(-) diff --git a/Lib/asyncio/unix_events.py b/Lib/asyncio/unix_events.py index d8f653045a..2551d69adc 100644 --- a/Lib/asyncio/unix_events.py +++ b/Lib/asyncio/unix_events.py @@ -878,6 +878,63 @@ class AbstractChildWatcher: raise NotImplementedError() +class PidfdChildWatcher(AbstractChildWatcher): + + def __init__(self): + self._loop = None + self._callbacks = {} + + def __enter__(self): + return self + + def __exit__(self, a, b, c): + pass + + def is_active(self): + return self._loop is not None and self._loop.is_running() + + def close(self): + self.attach_loop(None) + + def attach_loop(self, loop): + if self._loop is not None and loop is None and self._callbacks: + warnings.warn( + 'A loop is being detached ' + 'from a child watcher with pending handlers', + RuntimeWarning) + for pidfd, _, _ in self._callbacks.values(): + self._loop._remove_reader(pidfd) + os.close(pidfd) + self._callbacks.clear() + self._loop = loop + + def add_child_handler(self, pid, callback, *args): + existing = self._callbacks.get(pid) + if existing is not None: + self._callbacks[pid] = existing[0], callback, args + else: + pidfd = os.pidfd_open(pid) + self._loop._add_reader(pidfd, self._do_wait, pid) + self._callbacks[pid] = pidfd, callback, args + + def _do_wait(self, pid): + pidfd, callback, args = self._callbacks.pop(pid) + self._loop._remove_reader(pidfd) + _, status = os.waitpid(pid, 0) + os.close(pidfd) + returncode = _compute_returncode(status) + callback(pid, returncode, *args) + + def remove_child_handler(self, pid): + try: + pidfd, _, _ = self._callbacks.pop(pid) + except KeyError: + return False + self._loop._remove_reader(pidfd) + os.close(pidfd) + return True + + def _compute_returncode(status): if os.WIFSIGNALED(status): # The child process died because of a signal. diff --git a/Lib/test/test_asyncio/test_subprocess.py b/Lib/test/test_asyncio/test_subprocess.py index 17552d03f5..f92c64ccc8 100644 --- a/Lib/test/test_asyncio/test_subprocess.py +++ b/Lib/test/test_asyncio/test_subprocess.py @@ -691,6 +691,11 @@ if sys.platform != 'win32': Watcher = unix_events.FastChildWatcher + class SubprocessPidfdWatcherTests(SubprocessWatcherMixin, + test_utils.TestCase): + + Watcher = unix_events.PidfdChildWatcher + else: # Windows class SubprocessProactorTests(SubprocessMixin, test_utils.TestCase): diff --git a/Modules/clinic/posixmodule.c.h b/Modules/clinic/posixmodule.c.h index 561cbb0ca8..57a76f4e4c 100644 --- a/Modules/clinic/posixmodule.c.h +++ b/Modules/clinic/posixmodule.c.h @@ -3963,6 +3963,39 @@ os_wait(PyObject *module, PyObject *Py_UNUSED(ignored)) #endif /* defined(HAVE_WAIT) */ +#if defined(__linux__) + +PyDoc_STRVAR(os_pidfd_open__doc__, +"pidfd_open($module, /, pid)\n" +"--\n" +"\n"); + +#define OS_PIDFD_OPEN_METHODDEF \ + {"pidfd_open", (PyCFunction)(void(*)(void))os_pidfd_open, METH_FASTCALL|METH_KEYWORDS, os_pidfd_open__doc__}, + +static PyObject * +os_pidfd_open_impl(PyObject *module, pid_t pid); + +static PyObject * +os_pidfd_open(PyObject *module, PyObject *const *args, Py_ssize_t nargs, PyObject *kwnames) +{ + PyObject *return_value = NULL; + static const char * const _keywords[] = {"pid", NULL}; + static _PyArg_Parser _parser = {"" _Py_PARSE_PID ":pidfd_open", _keywords, 0}; + pid_t pid; + + if (!_PyArg_ParseStackAndKeywords(args, nargs, kwnames, &_parser, + &pid)) { + goto exit; + } + return_value = os_pidfd_open_impl(module, pid); + +exit: + return return_value; +} + +#endif /* defined(__linux__) */ + #if (defined(HAVE_READLINK) || defined(MS_WINDOWS)) PyDoc_STRVAR(os_readlink__doc__, @@ -8480,6 +8513,10 @@ exit: #define OS_WAIT_METHODDEF #endif /* !defined(OS_WAIT_METHODDEF) */ +#ifndef OS_PIDFD_OPEN_METHODDEF + #define OS_PIDFD_OPEN_METHODDEF +#endif /* !defined(OS_PIDFD_OPEN_METHODDEF) */ + #ifndef OS_READLINK_METHODDEF #define OS_READLINK_METHODDEF #endif /* !defined(OS_READLINK_METHODDEF) */ @@ -8731,4 +8768,4 @@ exit: #ifndef OS__REMOVE_DLL_DIRECTORY_METHODDEF #define OS__REMOVE_DLL_DIRECTORY_METHODDEF #endif /* !defined(OS__REMOVE_DLL_DIRECTORY_METHODDEF) */ -/*[clinic end generated code: output=fe7897441fed5402 input=a9049054013a1b77]*/ +/*[clinic end generated code: output=7d36df33e7ff55d2 input=a9049054013a1b77]*/ diff --git a/Modules/posixmodule.c b/Modules/posixmodule.c index dcd90d3a51..4bfcb5d710 100644 --- a/Modules/posixmodule.c +++ b/Modules/posixmodule.c @@ -7792,6 +7792,27 @@ os_wait_impl(PyObject *module) } #endif /* HAVE_WAIT */ +#ifdef __linux__ +/*[clinic input] +os.pidfd_open + pid: pid_t +[clinic start generated code]*/ + +static PyObject * +os_pidfd_open_impl(PyObject *module, pid_t pid) +/*[clinic end generated code: output=6f10e682fb292cc4 input=da1af61b3a170e95]*/ +{ +#ifndef __NR_pidfd_open +#define __NR_pidfd_open 434 +#endif + int fd = syscall(__NR_pidfd_open, pid, 0); + if (fd < 0) { + return posix_error(); + } + return PyLong_FromLong(fd); +} +#endif + #if defined(HAVE_READLINK) || defined(MS_WINDOWS) /*[clinic input] @@ -13641,6 +13662,7 @@ static PyMethodDef posix_methods[] = { OS_WAIT4_METHODDEF OS_WAITID_METHODDEF OS_WAITPID_METHODDEF + OS_PIDFD_OPEN_METHODDEF OS_GETSID_METHODDEF OS_SETSID_METHODDEF OS_SETPGID_METHODDEF -- 2.20.1