Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Use pidfd to implement trio.Process.wait, when available (2nd try) #1482

Merged
merged 4 commits into from
Apr 29, 2020
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 4 additions & 0 deletions newsfragments/1241.feature.rst
Original file line number Diff line number Diff line change
@@ -0,0 +1,4 @@
On Linux kernels v5.3 or newer, `trio.Process.wait` now uses `the
pidfd API <https://lwn.net/Articles/794707/>`__ to track child
processes. This shouldn't have any user-visible change, but it makes
working with subprocesses faster and use less memory.
79 changes: 72 additions & 7 deletions trio/_subprocess.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
import os
import subprocess
import sys
from typing import Optional

from ._abc import AsyncResource, SendStream, ReceiveStream
Expand All @@ -11,6 +12,39 @@
)
import trio

# Linux-specific, but has complex lifetime management stuff so we hard-code it
# here instead of hiding it behind the _subprocess_platform abstraction
can_try_pidfd_open = True
try:
from os import pidfd_open
except ImportError:
if sys.platform == "linux":
import ctypes
_cdll_for_pidfd_open = ctypes.CDLL(None, use_errno=True)
_cdll_for_pidfd_open.syscall.restype = ctypes.c_long
# pid and flags are actually int-sized, but the syscall() function
# always takes longs. (Except on x32 where long is 32-bits and syscall
# takes 64-bit arguments. But in the unlikely case that anyone is
# using x32, this will still work, b/c we only need to pass in 32 bits
# of data, and the C ABI doesn't distinguish between passing 32-bit vs
# 64-bit integers; our 32-bit values will get loaded into 64-bit
# registers where syscall() will find them.)
_cdll_for_pidfd_open.syscall.argtypes = [
ctypes.c_long, # syscall number
ctypes.c_long, # pid
ctypes.c_long, # flags
]
__NR_pidfd_open = 434

def pidfd_open(fd, flags):
result = _cdll_for_pidfd_open.syscall(__NR_pidfd_open, fd, flags)
if result < 0:
err = ctypes.get_errno()
raise OSError(err, os.strerror(err))
return result
else:
can_try_pidfd_open = False


class Process(AsyncResource):
r"""A child process. Like :class:`subprocess.Popen`, but async.
Expand Down Expand Up @@ -151,6 +185,21 @@ def _init(
if self.stderr is not None:
os.close(stderr)

self._pidfd = None
if can_try_pidfd_open:
try:
fd = pidfd_open(self._proc.pid, 0)
except OSError:
# Well, we tried, but it didn't work (probably because we're
# running on an older kernel, or in an older sandbox, that
# hasn't been updated to support pidfd_open). We'll fall back
# on waitid instead.
pass
else:
# It worked! Wrap the raw fd up in a Python file object to
# make sure it'll get closed.
self._pidfd = open(fd)

if self.stdin is not None and self.stdout is not None:
self.stdio = StapledStream(self.stdin, self.stdout)

Expand Down Expand Up @@ -205,19 +254,32 @@ async def aclose(self):
with trio.CancelScope(shield=True):
await self.wait()

def _close_pidfd(self):
if self._pidfd is not None:
self._pidfd.close()
self._pidfd = None

async def wait(self):
"""Block until the process exits.

Returns:
The exit status of the process; see :attr:`returncode`.
"""
if self.poll() is None:
async with self._wait_lock:
if self.poll() is None:
async with self._wait_lock:
if self.poll() is None:
if self._pidfd is not None:
await trio.hazmat.wait_readable(self._pidfd)
else:
await wait_child_exiting(self)
self._proc.wait()
else:
await trio.hazmat.checkpoint()
# We have to use .wait() here, not .poll(), because on macOS
# (and maybe other systems, who knows), there's a race
# condition inside the kernel that creates a tiny window where
# kqueue reports that the process has exited, but
# waitpid(WNOHANG) can't yet reap it. So this .wait() may
# actually block for a tiny fraction of a second.
self._proc.wait()
self._close_pidfd()
assert self.returncode is not None
return self.returncode

def poll(self):
Expand All @@ -227,7 +289,10 @@ def poll(self):
The exit status of the process, or ``None`` if it is still
running; see :attr:`returncode`.
"""
return self._proc.poll()
result = self._proc.poll()
if result is not None:
self._close_pidfd()
return result

def send_signal(self, sig):
"""Send signal ``sig`` to the process.
Expand Down