fix: wait for control socket dispatch sources before close

Packages/macOS/CmuxControlSocket/Sources/CmuxControlSocket/Server/ControlClientAsyncTransport.swift

@@ -1,5 +1,5 @@
 internal import Dispatch
-public import Foundation
+internal import Foundation
 internal import Darwin
 internal import os
 
@@ -46,6 +46,9 @@ public final class ControlClientAsyncLineReader: @unchecked Sendable {
     private var iterator: AsyncStream<Data>.Iterator
     private let source: any DispatchSourceRead
     private let revocationSource: (any DispatchSourceRead)?
+    /// Dispatch source cancellation is asynchronous. The connection owner
+    /// awaits this barrier before closing the borrowed descriptor.
+    private let sourceCancellationBarrier: DispatchSourceCancellationBarrier
     /// One-shot idle deadline carried over from SO_RCVTIMEO. DispatchSource
     /// is the low-level descriptor/timer bridge; it never blocks the task.
     private let idleReadTimer: (any DispatchSourceTimer)?
@@ -94,6 +97,8 @@ public final class ControlClientAsyncLineReader: @unchecked Sendable {
         )
         self.bufferedByteAccounting = bufferedByteAccounting
         self.limitsActive = OSAllocatedUnfairLock(initialState: initialLimits != nil)
+        let sourceCancellationBarrier = DispatchSourceCancellationBarrier()
+        self.sourceCancellationBarrier = sourceCancellationBarrier
         self.monotonicNowNanoseconds = monotonicNowNanoseconds ?? {
             DispatchTime.now().uptimeNanoseconds
         }
@@ -117,6 +122,7 @@ public final class ControlClientAsyncLineReader: @unchecked Sendable {
             fileDescriptor: socket,
             queue: DispatchQueue.global(qos: .utility)
         )
+        sourceCancellationBarrier.register()
         let byteCapForDrain = self.maximumBufferedBytes
         readSource.setEventHandler { [streamContinuation, sourceBox, bufferedByteAccounting] in
             Self.drain(
@@ -127,9 +133,10 @@ public final class ControlClientAsyncLineReader: @unchecked Sendable {
                 maximumBufferedBytes: byteCapForDrain
             )
         }
-        readSource.setCancelHandler { [streamContinuation, sourceBox] in
+        readSource.setCancelHandler { [streamContinuation, sourceBox, sourceCancellationBarrier] in
             sourceBox.source = nil
             streamContinuation.finish()
+            sourceCancellationBarrier.complete()
         }
         sourceBox.source = readSource
         source = readSource
@@ -152,13 +159,15 @@ public final class ControlClientAsyncLineReader: @unchecked Sendable {
                     fileDescriptor: duplicate,
                     queue: DispatchQueue.global(qos: .utility)
                 )
+                sourceCancellationBarrier.register()
                 revocation.setEventHandler { [streamContinuation, revocationBox] in
                     streamContinuation.finish()
                     revocationBox.source?.cancel()
                 }
-                revocation.setCancelHandler { [revocationBox] in
+                revocation.setCancelHandler { [revocationBox, sourceCancellationBarrier] in
                     revocationBox.source = nil
                     close(duplicate)
+                    sourceCancellationBarrier.complete()
                 }
                 revocationBox.source = revocation
                 revocationSource = revocation
@@ -194,8 +203,7 @@ public final class ControlClientAsyncLineReader: @unchecked Sendable {
 
     deinit {
         deadlineTask?.cancel()
-        source.cancel()
-        revocationSource?.cancel()
+        cancelSources()
         idleReadTimer?.cancel()
         continuation.finish()
     }
@@ -239,8 +247,7 @@ public final class ControlClientAsyncLineReader: @unchecked Sendable {
             let nextChunk = await withTaskCancellationHandler {
                 await iterator.next()
             } onCancel: {
-                source.cancel()
-                revocationSource?.cancel()
+                cancelSources()
                 idleReadTimer?.cancel()
                 continuation.finish()
             }
@@ -264,14 +271,26 @@ public final class ControlClientAsyncLineReader: @unchecked Sendable {
         return nil
     }
 
-    /// Explicitly terminates the reader's event sources.
-    public func cancel() {
+    /// Cancels the reader's event sources without waiting for their handlers.
+    private func cancelSources() {
         source.cancel()
         revocationSource?.cancel()
         idleReadTimer?.cancel()
         continuation.finish()
     }
 
+    /// Terminates the reader and waits for libdispatch to unregister every
+    /// descriptor source before returning to the connection owner.
+    ///
+    /// Callers that own ``socket`` must use this method before shutting down
+    /// or closing that descriptor. The cancellation handlers are responsible
+    /// for completing the wait; the descriptor itself remains owned by the
+    /// caller.
+    public func cancelAndWait() async {
+        cancelSources()
+        await sourceCancellationBarrier.wait()
+    }
+
     private var deadlineHasNotExpired: Bool {
         guard let deadlineUptimeNanoseconds else { return true }
         return monotonicNowNanoseconds() < deadlineUptimeNanoseconds
@@ -430,90 +449,3 @@ public final class ControlClientAsyncLineReader: @unchecked Sendable {
         }
     }
 }
-
-/// Asynchronously writes one connection's responses without parking a thread
-/// when a client applies backpressure. One writer task owns the instance.
-// @unchecked Sendable is safe because one connection task performs all writes;
-// the one-shot writable source communicates only through its continuation.
-public final class ControlClientAsyncWriter: @unchecked Sendable {
-    private final class SourceBox: @unchecked Sendable {
-        var source: (any DispatchSourceWrite)?
-    }
-
-    private let socket: Int32
-
-    /// Creates a writer over a non-blocking descriptor.
-    public init(socket: Int32) {
-        self.socket = socket
-        _ = Self.makeNonBlocking(socket)
-    }
-
-    /// Writes all bytes, suspending on `EAGAIN`; returns false after EOF,
-    /// cancellation, or a non-retryable write error.
-    public func writeAll(_ data: Data) async -> Bool {
-        var offset = 0
-        while offset < data.count, !Task.isCancelled {
-            let written = data.withUnsafeBytes { rawBuffer -> Int in
-                guard let baseAddress = rawBuffer.baseAddress else { return 0 }
-                return Darwin.write(
-                    socket,
-                    baseAddress.advanced(by: offset),
-                    rawBuffer.count - offset
-                )
-            }
-            if written > 0 {
-                offset += written
-                continue
-            }
-            if written < 0, errno == EINTR { continue }
-            if written < 0, errno == EAGAIN || errno == EWOULDBLOCK {
-                guard await waitForWritable() else { return false }
-                continue
-            }
-            return false
-        }
-        return offset == data.count
-    }
-
-    /// Cancels future write notifications.
-    public func cancel() {
-        // Each would-block wait owns a one-shot source and observes task
-        // cancellation. There is no persistent writable source to suspend.
-    }
-
-    private static func makeNonBlocking(_ socket: Int32) -> Int32? {
-        let flags = fcntl(socket, F_GETFL, 0)
-        guard flags >= 0 else { return errno }
-        guard fcntl(socket, F_SETFL, flags | O_NONBLOCK) >= 0 else { return errno }
-        return nil
-    }
-
-    private func waitForWritable() async -> Bool {
-        let stream = AsyncStream<Void>.makeStream(bufferingPolicy: .bufferingNewest(1))
-        let streamContinuation = stream.continuation
-        let sourceBox = SourceBox()
-        let writeSource = DispatchSource.makeWriteSource(
-            fileDescriptor: socket,
-            queue: DispatchQueue.global(qos: .utility)
-        )
-        writeSource.setEventHandler { [streamContinuation, sourceBox] in
-            streamContinuation.yield(())
-            streamContinuation.finish()
-            sourceBox.source?.cancel()
-        }
-        writeSource.setCancelHandler {
-            streamContinuation.finish()
-        }
-        sourceBox.source = writeSource
-        writeSource.activate()
-
-        var iterator = stream.stream.makeAsyncIterator()
-        let writable: Void? = await withTaskCancellationHandler {
-            await iterator.next()
-        } onCancel: {
-            sourceBox.source?.cancel()
-            streamContinuation.finish()
-        }
-        return writable != nil
-    }
-}

Packages/macOS/CmuxControlSocket/Sources/CmuxControlSocket/Server/ControlClientAsyncWriter.swift

@@ -0,0 +1,109 @@
+internal import Dispatch
+public import Foundation
+internal import Darwin
+
+/// Asynchronously writes one connection's responses without parking a thread
+/// when a client applies backpressure. One writer task owns the instance.
+// @unchecked Sendable is safe because one connection task performs all writes;
+// the one-shot writable source communicates only through its continuation.
+public final class ControlClientAsyncWriter: @unchecked Sendable {
+    private final class SourceBox: @unchecked Sendable {
+        var source: (any DispatchSourceWrite)?
+    }
+
+    private let socket: Int32
+    /// One-shot writable sources must finish cancellation before the owner
+    /// closes the shared socket descriptor.
+    private let sourceCancellationBarrier = DispatchSourceCancellationBarrier()
+
+    /// Creates a writer over a non-blocking descriptor.
+    ///
+    /// - Parameter socket: A borrowed descriptor retained by the connection owner.
+    public init(socket: Int32) {
+        self.socket = socket
+        _ = Self.makeNonBlocking(socket)
+    }
+
+    /// Writes all bytes, suspending on `EAGAIN`; returns false after EOF,
+    /// cancellation, or a non-retryable write error.
+    public func writeAll(_ data: Data) async -> Bool {
+        var offset = 0
+        while offset < data.count, !Task.isCancelled {
+            let written = data.withUnsafeBytes { rawBuffer -> Int in
+                guard let baseAddress = rawBuffer.baseAddress else { return 0 }
+                return Darwin.write(
+                    socket,
+                    baseAddress.advanced(by: offset),
+                    rawBuffer.count - offset
+                )
+            }
+            if written > 0 {
+                offset += written
+                continue
+            }
+            if written < 0, errno == EINTR { continue }
+            if written < 0, errno == EAGAIN || errno == EWOULDBLOCK {
+                let source = DispatchSource.makeWriteSource(
+                    fileDescriptor: socket,
+                    queue: DispatchQueue.global(qos: .utility)
+                )
+                guard await waitForWritable(source: source) else { return false }
+                continue
+            }
+            return false
+        }
+        return offset == data.count
+    }
+
+    /// Waits for any one-shot writable source to finish cancellation.
+    ///
+    /// Call after the owning write task has returned. To interrupt a write,
+    /// cancel that task first; this method does not cancel an in-flight write.
+    public func cancelAndWait() async {
+        await sourceCancellationBarrier.wait()
+    }
+
+    /// Enables would-block handling instead of parking the connection task.
+    private static func makeNonBlocking(_ socket: Int32) -> Int32? {
+        let flags = fcntl(socket, F_GETFL, 0)
+        guard flags >= 0 else { return errno }
+        guard fcntl(socket, F_SETFL, flags | O_NONBLOCK) >= 0 else { return errno }
+        return nil
+    }
+
+    /// Joins cancellation of a one-shot source on readiness or task cancellation.
+    /// - Parameter source: The inactive source for this socket; activation and
+    ///   cancellation belong to this operation.
+    func waitForWritable(source writeSource: any DispatchSourceWrite) async -> Bool {
+        let stream = AsyncStream<Void>.makeStream(bufferingPolicy: .bufferingNewest(1))
+        let streamContinuation = stream.continuation
+        let sourceBox = SourceBox()
+        sourceCancellationBarrier.register()
+        writeSource.setEventHandler { [streamContinuation, sourceBox] in
+            streamContinuation.yield(())
+            streamContinuation.finish()
+            sourceBox.source?.cancel()
+        }
+        writeSource.setCancelHandler { [sourceCancellationBarrier] in
+            streamContinuation.finish()
+            sourceCancellationBarrier.complete()
+        }
+        sourceBox.source = writeSource
+        writeSource.activate()
+
+        var iterator = stream.stream.makeAsyncIterator()
+        let writable: Void? = await withTaskCancellationHandler {
+            await iterator.next()
+        } onCancel: {
+            sourceBox.source?.cancel()
+            streamContinuation.finish()
+        }
+        // The event/cancellation handler may have resumed the iterator before
+        // libdispatch ran the source's cancellation handler. Await here while
+        // the source is still retained, so a subsequent socket close cannot
+        // race the kevent teardown.
+        writeSource.cancel()
+        await sourceCancellationBarrier.wait()
+        return writable != nil
+    }
+}

Packages/macOS/CmuxControlSocket/Sources/CmuxControlSocket/Server/DispatchSourceCancellationBarrier.swift

@@ -0,0 +1,52 @@
+internal import Foundation
+
+/// Completes asynchronous teardown after every registered dispatch source has
+/// run its cancellation handler.
+final class DispatchSourceCancellationBarrier: @unchecked Sendable {
+    // Lock carve-out: register() runs synchronously before source activation,
+    // while DispatchSource cancellation callbacks synchronously call complete().
+    // Moving these updates into unstructured actor Tasks could reorder register
+    // and complete relative to wait(), letting teardown observe zero too early.
+    // This lock protects only the count and continuation handoff, never I/O or
+    // an await. Every continuation is removed under the lock and resumed after
+    // unlocking; all mutable state is locked, which justifies Sendable above.
+    private let lock = NSLock()
+    private var registrations = 0
+    private var waiters: [CheckedContinuation<Void, Never>] = []
+
+    /// Registers one source before it is activated.
+    func register() {
+        lock.lock()
+        registrations += 1
+        lock.unlock()
+    }
+
+    /// Marks one source's cancellation handler complete.
+    func complete() {
+        lock.lock()
+        registrations -= 1
+        let waiters = registrations == 0 ? self.waiters : []
+        if registrations == 0 {
+            self.waiters.removeAll(keepingCapacity: false)
+        }
+        lock.unlock()
+        for waiter in waiters {
+            waiter.resume()
+        }
+    }
+
+    /// Suspends without blocking an executor until all sources have cancelled.
+    func wait() async {
+        await withCheckedContinuation { (continuation: CheckedContinuation<Void, Never>) in
+            lock.lock()
+            guard registrations > 0 else {
+                lock.unlock()
+                continuation.resume()
+                return
+            }
+            waiters.append(continuation)
+            lock.unlock()
+        }
+    }
+}
+