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

newsfragments/1241.feature.rst

@@ -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.

trio/_subprocess.py

@@ -1,5 +1,6 @@
 import os
 import subprocess
+import sys
 from typing import Optional
 
 from ._abc import AsyncResource, SendStream, ReceiveStream
@@ -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.
@@ -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)
 
@@ -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):
@@ -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.