SignalFD.areaditer: support iterative reading

1 files changed, 36 insertions, 2 deletions

diff --git a/linuxnamespaces/syscalls.py b/linuxnamespaces/syscalls.py
index 9528b07..c4edad2 100644
--- a/linuxnamespaces/syscalls.py
+++ b/linuxnamespaces/syscalls.py
@@ -762,6 +762,8 @@ def setns(fd: int, nstype: CloneFlags = CloneFlags.NONE) -> None:
 class SignalFD:
     """Represent a file descriptor returned from signalfd(2)."""
 
+    _ReadIterFut = asyncio.Future[tuple[list[SignalFDSigInfo], "_ReadIterFut"]]
+
     def __init__(
         self,
         sigmask: typing.Iterable[signal.Signals],
@@ -801,7 +803,7 @@ class SignalFD:
         res = self.readv(1)
         return res[0]
 
-    def __handle_readable(
+    def __handle_read(
         self, fd: int, fut: asyncio.Future[SignalFDSigInfo]
     ) -> None:
         try:
@@ -826,9 +828,41 @@ class SignalFD:
             raise ValueError("attempt to read from closed signalfd")
         loop = asyncio.get_running_loop()
         fut: asyncio.Future[SignalFDSigInfo] = loop.create_future()
-        loop.add_reader(self.fd, self.__handle_readable, self.fd, fut)
+        loop.add_reader(self.fd, self.__handle_read, self.fd, fut)
         return fut
 
+    def __handle_readiter(self, fd: int, fut: _ReadIterFut) -> None:
+        loop = fut.get_loop()
+        try:
+            if fd != self.fd:
+                raise RuntimeError("SignalFD file descriptor changed")
+            try:
+                # Attempt to read a full page worth of queued signals.
+                results = self.readv(32)
+            except OSError as err:
+                if err.errno == errno.EAGAIN:
+                    return
+                raise
+        except Exception as exc:
+            loop.remove_reader(fd)
+            fut.set_exception(exc)
+        else:
+            nextfut: SignalFD._ReadIterFut = loop.create_future()
+            loop.add_reader(fd, self.__handle_readiter, self.fd, nextfut)
+            fut.set_result((results, nextfut))
+
+    async def areaditer(self) -> typing.AsyncIterator[SignalFDSigInfo]:
+        """Asynchronously read signals from the signalfd forever."""
+        if self.fd < 0:
+            raise ValueError("attempt to read from closed signalfd")
+        loop = asyncio.get_running_loop()
+        fut: SignalFD._ReadIterFut = loop.create_future()
+        loop.add_reader(self.fd, self.__handle_readiter, self.fd, fut)
+        while True:
+            results, fut = await fut
+            for result in results:
+                yield result
+
     def fileno(self) -> FileDescriptor:
         """Return the underlying file descriptor."""
         return self.fd