add an asyncio waitid(P_PIDFD, ...) helper

1 files changed, 34 insertions, 0 deletions

diff --git a/linuxnamespaces/__init__.py b/linuxnamespaces/__init__.py
index 1a8b0ee..3302867 100644
--- a/linuxnamespaces/__init__.py
+++ b/linuxnamespaces/__init__.py
@@ -673,3 +673,37 @@ def async_copyfd(
     ):
         return _AsyncSplicer(from_fd, to_fd, count).fut
     return _AsyncCopier(from_fd, to_fd, count).fut
+
+
+class _AsyncPidfdWaiter:
+    def __init__(self, pidfd: int, flags: int):
+        self.pidfd = pidfd
+        self.flags = flags
+        self.loop = asyncio.get_running_loop()
+        self.fut: asyncio.Future[
+            os.waitid_result | None
+        ] = self.loop.create_future()
+        self.loop.add_reader(pidfd, self.handle_readable)
+
+    def handle_readable(self) -> None:
+        try:
+            result = os.waitid(os.P_PIDFD, self.pidfd, self.flags)
+        except OSError as err:
+            if err.errno != errno.EAGAIN:
+                self.loop.remove_reader(self.pidfd)
+                self.fut.set_exception(err)
+        except Exception as err:
+            self.loop.remove_reader(self.pidfd)
+            self.fut.set_exception(err)
+        else:
+            self.loop.remove_reader(self.pidfd)
+            self.fut.set_result(result)
+
+
+def async_waitpidfd(
+    pidfd: int, flags: int
+) -> asyncio.Future[os.waitid_result | None]:
+    """Asynchronously wait for a process represented as a pidfd. This is an
+    async variant of waitid(P_PIDFD, pidfd, flags).
+    """
+    return _AsyncPidfdWaiter(pidfd, flags).fut