SocketPool: keep a descriptor's registration when a closed socket's is reused

FlyingSocks/Sources/SocketPool+ePoll.swift

@@ -107,7 +107,10 @@ public struct ePoll: EventQueue {
     public mutating func addEvents(_ events: Socket.Events, for socket: Socket.FileDescriptor) throws {
         var socketEvents = existing[socket] ?? []
         socketEvents.formUnion(events)
-        try setEvents(socketEvents, for: socket)
+        // Tell the kernel even when the events are recorded as registered: a descriptor
+        // closed while registered has left the epoll set, and a socket since given the
+        // descriptor would otherwise never be registered.
+        try setEvents(socketEvents, for: socket, evenIfRecorded: true)
     }
 
     public mutating func removeEvents(_ events: Socket.Events, for socket: Socket.FileDescriptor) throws {
@@ -118,32 +121,39 @@ public struct ePoll: EventQueue {
         try setEvents(socketEvents, for: socket)
     }
 
-    mutating func setEvents(_ events: Socket.Events, for socket: Socket.FileDescriptor) throws {
-        guard existing[socket] != events else { return }
+    mutating func setEvents(
+        _ events: Socket.Events,
+        for socket: Socket.FileDescriptor,
+        evenIfRecorded: Bool = false
+    ) throws {
+        guard evenIfRecorded || existing[socket] != events else { return }
         var event = CSystemLinux.epoll_event()
         event.events = events.epollEvents(triggerMode: triggerMode).rawValue
         event.data.fd = socket.rawValue
 
-        if existing[socket] != nil {
-            if events.isEmpty {
-                guard epoll_ctl(file.rawValue, EPOLL_CTL_DEL, socket.rawValue, &event) != -1 else {
-                    throw SocketError.makeFailed("epoll_ctl EPOLL_CTL_DEL")
-                }
-            } else {
-                guard epoll_ctl(file.rawValue, EPOLL_CTL_MOD, socket.rawValue, &event) != -1 else {
-                    throw SocketError.makeFailed("epoll_ctl EPOLL_CTL_MOD")
-                }
-            }
-        } else if !events.isEmpty {
-            guard epoll_ctl(file.rawValue, EPOLL_CTL_ADD, socket.rawValue, &event) != -1 else {
-                throw SocketError.makeFailed("epoll_ctl EPOLL_CTL_ADD")
-            }
-        }
+        let recorded = existing[socket] != nil
+        // Record the change before making it. If it fails because the descriptor was closed
+        // while registered, the epoll set has dropped it already.
+        existing[socket] = events.isEmpty ? nil : events
 
         if events.isEmpty {
-            existing[socket] = nil
-        } else {
-            existing[socket] = events
+            guard recorded else { return }
+            guard epoll_ctl(file.rawValue, EPOLL_CTL_DEL, socket.rawValue, &event) != -1 else {
+                throw SocketError.makeFailed("epoll_ctl EPOLL_CTL_DEL")
+            }
+        } else if epoll_ctl(file.rawValue, recorded ? EPOLL_CTL_MOD : EPOLL_CTL_ADD, socket.rawValue, &event) == -1 {
+            // The epoll set can disagree with the record once a descriptor closed while
+            // registered is reused: modifying it then fails with ENOENT, adding with EEXIST.
+            var retry: Int32?
+            if recorded, errno == ENOENT {
+                retry = EPOLL_CTL_ADD
+            } else if !recorded, errno == EEXIST {
+                retry = EPOLL_CTL_MOD
+            }
+            guard let retry, epoll_ctl(file.rawValue, retry, socket.rawValue, &event) != -1 else {
+                existing[socket] = nil
+                throw SocketError.makeFailed(recorded ? "epoll_ctl EPOLL_CTL_MOD" : "epoll_ctl EPOLL_CTL_ADD")
+            }
         }
     }
 

FlyingSocks/Sources/SocketPool+kQueue.swift

@@ -76,12 +76,14 @@ public struct kQueue: EventQueue {
 
     public mutating func addEvents(_ events: Socket.Events, for socket: Socket.FileDescriptor) throws {
         for event in events {
+            // Add the event even when it is recorded as added: EV_ADD of a filter the kqueue
+            // has just updates it, but a descriptor closed while registered has left the
+            // kqueue, and skipping the add would leave a socket since given the descriptor
+            // never woken.
+            try addEvent(event, for: socket)
             var socketEvents = existing[socket] ?? []
-            if !socketEvents.contains(event) {
-                try addEvent(event, for: socket)
-                socketEvents.insert(event)
-                existing[socket] = socketEvents
-            }
+            socketEvents.insert(event)
+            existing[socket] = socketEvents
         }
     }
 
@@ -101,17 +103,13 @@ public struct kQueue: EventQueue {
 
     public mutating func removeEvents(_ events: Socket.Events, for socket: Socket.FileDescriptor) throws {
         for event in events {
-            if var entries = existing[socket] {
-                if entries.contains(event) {
-                    try removeEvent(event, for: socket)
-                    entries.remove(event)
-                    if entries.isEmpty {
-                        existing[socket] = nil
-                    } else {
-                        existing[socket] = entries
-                    }
-                }
-            }
+            guard var entries = existing[socket], entries.contains(event) else { continue }
+            // Forget the event before removing it: a socket closed while registered has left
+            // the kqueue already, so the removal fails, and remembering the event would stop
+            // it being added for a socket later given the same descriptor.
+            entries.remove(event)
+            existing[socket] = entries.isEmpty ? nil : entries
+            try removeEvent(event, for: socket)
         }
     }
 

FlyingSocks/Tests/SocketPoolTests.swift

@@ -175,6 +175,65 @@ struct SocketPoolTests {
         }
     }
 
+#if canImport(Darwin) || canImport(CSystemLinux)
+    @Test
+    func reusedDescriptor_IsWoken_AfterTheSuspendedSocketIsReplacedAndCancelled() async throws {
+        let pool = SocketPool.make(maxEvents: 4)
+        try await pool.prepare()
+        let run = Task { try await pool.run() }
+        defer { run.cancel() }
+
+        let (old, oldPeer) = try Socket.makeNonBlockingPair()
+        defer { try? oldPeer.close() }
+        let suspended = Task { try await pool.suspendSocket(old, untilReadyFor: .read) }
+        try await Task.sleep(seconds: 0.1)
+
+        // closing the socket drops it from the queue; cancelling its suspension then fails
+        // to remove it, and a socket given the same descriptor must still be woken
+        let (reused, peer) = try Socket.makeNonBlockingPair(replacing: old)
+        defer {
+            try? reused.close()
+            try? peer.close()
+        }
+        suspended.cancel()
+        _ = try? await suspended.value
+
+        _ = try peer.write(Data([1]))
+        let woken = Task(timeout: 1) { try await pool.suspendSocket(reused, untilReadyFor: .read) }
+        await #expect(throws: Never.self) {
+            try await woken.value
+        }
+    }
+
+    @Test
+    func reusedDescriptor_IsWoken_WhileTheReplacedSocketIsStillSuspended() async throws {
+        let pool = SocketPool.make(maxEvents: 4)
+        try await pool.prepare()
+        let run = Task { try await pool.run() }
+        defer { run.cancel() }
+
+        let (old, oldPeer) = try Socket.makeNonBlockingPair()
+        defer { try? oldPeer.close() }
+        let suspended = Task { try await pool.suspendSocket(old, untilReadyFor: .read) }
+        defer { suspended.cancel() }
+        try await Task.sleep(seconds: 0.1)
+
+        // the socket is closed and its descriptor given to another before its suspension
+        // ends; the new socket's own suspension must still register with the queue
+        let (reused, peer) = try Socket.makeNonBlockingPair(replacing: old)
+        defer {
+            try? reused.close()
+            try? peer.close()
+        }
+
+        _ = try peer.write(Data([1]))
+        let woken = Task(timeout: 1) { try await pool.suspendSocket(reused, untilReadyFor: .read) }
+        await #expect(throws: Never.self) {
+            try await woken.value
+        }
+    }
+#endif
+
     @Test
     func waiting_IsEmpty() async {
         let cn = await Continuation.make()
@@ -253,6 +312,22 @@ struct SocketPoolTests {
     }
 }
 
+#if canImport(Darwin) || canImport(CSystemLinux)
+private extension Socket {
+    /// A connected pair whose first socket takes over `socket`'s descriptor, closing it, as
+    /// a new socket can be given the descriptor of one just closed. `dup2` closes and
+    /// replaces in one step, so no other socket can be given the descriptor in between.
+    static func makeNonBlockingPair(replacing socket: Socket) throws -> (Socket, Socket) {
+        let (first, second) = try Socket.makeNonBlockingPair()
+        guard dup2(first.file.rawValue, socket.file.rawValue) != -1 else {
+            throw SocketError.makeFailed("dup2")
+        }
+        try first.close()
+        return (Socket(file: socket.file), second)
+    }
+}
+#endif
+
 private extension SocketPool where Queue == MockEventQueue  {
     static func make() -> Self {
         .init(queue: MockEventQueue())