From ed0fceca629abd1146e570ebafbb505cafd4f101 Mon Sep 17 00:00:00 2001 From: austinpower1258 Date: Sun, 20 Sep 2026 17:50:16 -0700 Subject: [PATCH 1/7] fix: wait for control socket source cancellation before close --- .../Server/ControlClientAsyncTransport.swift | 48 +++++++++++++++++-- .../ControlClientSourceLifecycleTests.swift | 41 ++++++++++++++++ Sources/TerminalController.swift | 7 ++- 3 files changed, 91 insertions(+), 5 deletions(-) create mode 100644 Packages/macOS/CmuxControlSocket/Tests/CmuxControlSocketTests/ControlClientSourceLifecycleTests.swift diff --git a/Packages/macOS/CmuxControlSocket/Sources/CmuxControlSocket/Server/ControlClientAsyncTransport.swift b/Packages/macOS/CmuxControlSocket/Sources/CmuxControlSocket/Server/ControlClientAsyncTransport.swift index c2838f39c33b..4ad6cb41988e 100644 --- a/Packages/macOS/CmuxControlSocket/Sources/CmuxControlSocket/Server/ControlClientAsyncTransport.swift +++ b/Packages/macOS/CmuxControlSocket/Sources/CmuxControlSocket/Server/ControlClientAsyncTransport.swift @@ -46,6 +46,10 @@ public final class ControlClientAsyncLineReader: @unchecked Sendable { private var iterator: AsyncStream.Iterator private let source: any DispatchSourceRead private let revocationSource: (any DispatchSourceRead)? + /// Dispatch source cancellation is asynchronous. The connection owner + /// must wait for both cancellation handlers before closing a borrowed + /// descriptor, otherwise libdispatch can observe a vanished kevent. + private let sourceCancellationGroup: DispatchGroup /// 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 +98,8 @@ public final class ControlClientAsyncLineReader: @unchecked Sendable { ) self.bufferedByteAccounting = bufferedByteAccounting self.limitsActive = OSAllocatedUnfairLock(initialState: initialLimits != nil) + let sourceCancellationGroup = DispatchGroup() + self.sourceCancellationGroup = sourceCancellationGroup self.monotonicNowNanoseconds = monotonicNowNanoseconds ?? { DispatchTime.now().uptimeNanoseconds } @@ -117,6 +123,7 @@ public final class ControlClientAsyncLineReader: @unchecked Sendable { fileDescriptor: socket, queue: DispatchQueue.global(qos: .utility) ) + sourceCancellationGroup.enter() let byteCapForDrain = self.maximumBufferedBytes readSource.setEventHandler { [streamContinuation, sourceBox, bufferedByteAccounting] in Self.drain( @@ -127,9 +134,10 @@ public final class ControlClientAsyncLineReader: @unchecked Sendable { maximumBufferedBytes: byteCapForDrain ) } - readSource.setCancelHandler { [streamContinuation, sourceBox] in + readSource.setCancelHandler { [streamContinuation, sourceBox, sourceCancellationGroup] in sourceBox.source = nil streamContinuation.finish() + sourceCancellationGroup.leave() } sourceBox.source = readSource source = readSource @@ -152,13 +160,15 @@ public final class ControlClientAsyncLineReader: @unchecked Sendable { fileDescriptor: duplicate, queue: DispatchQueue.global(qos: .utility) ) + sourceCancellationGroup.enter() revocation.setEventHandler { [streamContinuation, revocationBox] in streamContinuation.finish() revocationBox.source?.cancel() } - revocation.setCancelHandler { [revocationBox] in + revocation.setCancelHandler { [revocationBox, sourceCancellationGroup] in revocationBox.source = nil close(duplicate) + sourceCancellationGroup.leave() } revocationBox.source = revocation revocationSource = revocation @@ -272,6 +282,18 @@ public final class ControlClientAsyncLineReader: @unchecked Sendable { 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() { + cancel() + sourceCancellationGroup.wait() + } + private var deadlineHasNotExpired: Bool { guard let deadlineUptimeNanoseconds else { return true } return monotonicNowNanoseconds() < deadlineUptimeNanoseconds @@ -441,6 +463,9 @@ public final class ControlClientAsyncWriter: @unchecked Sendable { } private let socket: Int32 + /// One-shot writable sources must finish cancellation before the owner + /// closes the shared socket descriptor. + private let sourceCancellationGroup = DispatchGroup() /// Creates a writer over a non-blocking descriptor. public init(socket: Int32) { @@ -481,6 +506,11 @@ public final class ControlClientAsyncWriter: @unchecked Sendable { // cancellation. There is no persistent writable source to suspend. } + /// Waits for any one-shot writable source to finish cancellation. + public func cancelAndWait() { + sourceCancellationGroup.wait() + } + private static func makeNonBlocking(_ socket: Int32) -> Int32? { let flags = fcntl(socket, F_GETFL, 0) guard flags >= 0 else { return errno } @@ -496,13 +526,15 @@ public final class ControlClientAsyncWriter: @unchecked Sendable { fileDescriptor: socket, queue: DispatchQueue.global(qos: .utility) ) + sourceCancellationGroup.enter() writeSource.setEventHandler { [streamContinuation, sourceBox] in streamContinuation.yield(()) streamContinuation.finish() sourceBox.source?.cancel() } - writeSource.setCancelHandler { + writeSource.setCancelHandler { [sourceCancellationGroup] in streamContinuation.finish() + sourceCancellationGroup.leave() } sourceBox.source = writeSource writeSource.activate() @@ -514,6 +546,16 @@ public final class ControlClientAsyncWriter: @unchecked Sendable { sourceBox.source?.cancel() streamContinuation.finish() } + // The event/cancellation handler may have resumed the iterator before + // libdispatch ran the source's cancellation handler. Wait here while + // the source is still retained, so a subsequent socket close cannot + // race the kevent teardown. + writeSource.cancel() + waitForSourceCancellation() return writable != nil } + + private func waitForSourceCancellation() { + sourceCancellationGroup.wait() + } } diff --git a/Packages/macOS/CmuxControlSocket/Tests/CmuxControlSocketTests/ControlClientSourceLifecycleTests.swift b/Packages/macOS/CmuxControlSocket/Tests/CmuxControlSocketTests/ControlClientSourceLifecycleTests.swift new file mode 100644 index 000000000000..452d1a33cecb --- /dev/null +++ b/Packages/macOS/CmuxControlSocket/Tests/CmuxControlSocketTests/ControlClientSourceLifecycleTests.swift @@ -0,0 +1,41 @@ +import Darwin +import Foundation +import Testing +@testable import CmuxControlSocket + +/// Exercises the same cancellation followed by close used by a CLI connection. +@Suite(.serialized) +struct ControlClientSourceLifecycleTests { + @Test(.timeLimit(.minutes(1))) + func repeatedReaderCancellationBeforeOwnerClose() async throws { + try await withThrowingTaskGroup(of: Int.self) { group in + for _ in 0..<4 { + group.addTask { + for _ in 0..<256 { + let pair = try UnixSocketFixture.makeSocketPair() + let reader = ControlClientAsyncLineReader(socket: pair.reader) + defer { + // This is TerminalController.handleClientAsync's + // teardown order. The reader must wait for + // libdispatch to unregister its borrowed + // descriptor before the owner closes it. + reader.cancelAndWait() + shutdown(pair.reader, SHUT_RDWR) + close(pair.reader) + close(pair.writer) + } + var line: [UInt8] = [112, 105, 110, 103, 10] + #expect(write(pair.writer, &line, line.count) == line.count) + // A delivered line proves the read source was armed; + // do not substitute a delay before teardown. + #expect(await reader.nextLine { true } == "ping") + } + return 256 + } + } + var completed = 0 + for try await count in group { completed += count } + #expect(completed == 1_024) + } + } +} diff --git a/Sources/TerminalController.swift b/Sources/TerminalController.swift index 4d7e00fc9d23..daa2321347a9 100644 --- a/Sources/TerminalController.swift +++ b/Sources/TerminalController.swift @@ -2052,8 +2052,11 @@ class TerminalController { let writer = ControlClientAsyncWriter(socket: socket) let rateLimiter = ControlClientRateLimiter() defer { - lineReader.cancel() - writer.cancel() + // Dispatch source cancellation is asynchronous. Wait for every + // borrowed socket source to unregister before the outer defer + // shuts down and closes the descriptor. + lineReader.cancelAndWait() + writer.cancelAndWait() } while let line = await lineReader.nextLine(shouldContinueReading: { self.socketServer.isConnectionAuthorizationCurrent(authorizationGeneration) From e13d7f30a5cbf60676da0e1e119fe88d1cd7ecd8 Mon Sep 17 00:00:00 2001 From: austinpower1258 Date: Sun, 20 Sep 2026 18:29:55 -0700 Subject: [PATCH 2/7] fix: await control socket source teardown asynchronously --- .../Server/ControlClientAsyncTransport.swift | 106 +++++++++++------- .../ControlClientSourceLifecycleTests.swift | 69 ++++++++++-- ...ocketControlServerManagedPolicyTests.swift | 2 +- Sources/TerminalController.swift | 57 ++++++---- 4 files changed, 165 insertions(+), 69 deletions(-) diff --git a/Packages/macOS/CmuxControlSocket/Sources/CmuxControlSocket/Server/ControlClientAsyncTransport.swift b/Packages/macOS/CmuxControlSocket/Sources/CmuxControlSocket/Server/ControlClientAsyncTransport.swift index 4ad6cb41988e..a01f6e937da4 100644 --- a/Packages/macOS/CmuxControlSocket/Sources/CmuxControlSocket/Server/ControlClientAsyncTransport.swift +++ b/Packages/macOS/CmuxControlSocket/Sources/CmuxControlSocket/Server/ControlClientAsyncTransport.swift @@ -3,6 +3,49 @@ public import Foundation internal import Darwin internal import os +/// Completes asynchronous teardown after every registered dispatch source has +/// run its cancellation handler. +private final class DispatchSourceCancellationBarrier: @unchecked Sendable { + private let lock = NSLock() + private var registrations = 0 + private var waiters: [CheckedContinuation] = [] + + /// 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) in + lock.lock() + guard registrations > 0 else { + lock.unlock() + continuation.resume() + return + } + waiters.append(continuation) + lock.unlock() + } + } +} + /// Bridges one non-blocking control-socket descriptor into an async line /// reader. The descriptor is never closed by this type; its connection owner /// retains that responsibility. @@ -47,9 +90,8 @@ public final class ControlClientAsyncLineReader: @unchecked Sendable { private let source: any DispatchSourceRead private let revocationSource: (any DispatchSourceRead)? /// Dispatch source cancellation is asynchronous. The connection owner - /// must wait for both cancellation handlers before closing a borrowed - /// descriptor, otherwise libdispatch can observe a vanished kevent. - private let sourceCancellationGroup: DispatchGroup + /// 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)? @@ -98,8 +140,8 @@ public final class ControlClientAsyncLineReader: @unchecked Sendable { ) self.bufferedByteAccounting = bufferedByteAccounting self.limitsActive = OSAllocatedUnfairLock(initialState: initialLimits != nil) - let sourceCancellationGroup = DispatchGroup() - self.sourceCancellationGroup = sourceCancellationGroup + let sourceCancellationBarrier = DispatchSourceCancellationBarrier() + self.sourceCancellationBarrier = sourceCancellationBarrier self.monotonicNowNanoseconds = monotonicNowNanoseconds ?? { DispatchTime.now().uptimeNanoseconds } @@ -123,7 +165,7 @@ public final class ControlClientAsyncLineReader: @unchecked Sendable { fileDescriptor: socket, queue: DispatchQueue.global(qos: .utility) ) - sourceCancellationGroup.enter() + sourceCancellationBarrier.register() let byteCapForDrain = self.maximumBufferedBytes readSource.setEventHandler { [streamContinuation, sourceBox, bufferedByteAccounting] in Self.drain( @@ -134,10 +176,10 @@ public final class ControlClientAsyncLineReader: @unchecked Sendable { maximumBufferedBytes: byteCapForDrain ) } - readSource.setCancelHandler { [streamContinuation, sourceBox, sourceCancellationGroup] in + readSource.setCancelHandler { [streamContinuation, sourceBox, sourceCancellationBarrier] in sourceBox.source = nil streamContinuation.finish() - sourceCancellationGroup.leave() + sourceCancellationBarrier.complete() } sourceBox.source = readSource source = readSource @@ -160,15 +202,15 @@ public final class ControlClientAsyncLineReader: @unchecked Sendable { fileDescriptor: duplicate, queue: DispatchQueue.global(qos: .utility) ) - sourceCancellationGroup.enter() + sourceCancellationBarrier.register() revocation.setEventHandler { [streamContinuation, revocationBox] in streamContinuation.finish() revocationBox.source?.cancel() } - revocation.setCancelHandler { [revocationBox, sourceCancellationGroup] in + revocation.setCancelHandler { [revocationBox, sourceCancellationBarrier] in revocationBox.source = nil close(duplicate) - sourceCancellationGroup.leave() + sourceCancellationBarrier.complete() } revocationBox.source = revocation revocationSource = revocation @@ -204,8 +246,7 @@ public final class ControlClientAsyncLineReader: @unchecked Sendable { deinit { deadlineTask?.cancel() - source.cancel() - revocationSource?.cancel() + cancelSources() idleReadTimer?.cancel() continuation.finish() } @@ -249,8 +290,7 @@ public final class ControlClientAsyncLineReader: @unchecked Sendable { let nextChunk = await withTaskCancellationHandler { await iterator.next() } onCancel: { - source.cancel() - revocationSource?.cancel() + cancelSources() idleReadTimer?.cancel() continuation.finish() } @@ -274,8 +314,8 @@ 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() @@ -289,9 +329,9 @@ public final class ControlClientAsyncLineReader: @unchecked Sendable { /// or closing that descriptor. The cancellation handlers are responsible /// for completing the wait; the descriptor itself remains owned by the /// caller. - public func cancelAndWait() { - cancel() - sourceCancellationGroup.wait() + public func cancelAndWait() async { + cancelSources() + await sourceCancellationBarrier.wait() } private var deadlineHasNotExpired: Bool { @@ -465,7 +505,7 @@ public final class ControlClientAsyncWriter: @unchecked Sendable { private let socket: Int32 /// One-shot writable sources must finish cancellation before the owner /// closes the shared socket descriptor. - private let sourceCancellationGroup = DispatchGroup() + private let sourceCancellationBarrier = DispatchSourceCancellationBarrier() /// Creates a writer over a non-blocking descriptor. public init(socket: Int32) { @@ -500,15 +540,9 @@ public final class ControlClientAsyncWriter: @unchecked Sendable { 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. - } - /// Waits for any one-shot writable source to finish cancellation. - public func cancelAndWait() { - sourceCancellationGroup.wait() + public func cancelAndWait() async { + await sourceCancellationBarrier.wait() } private static func makeNonBlocking(_ socket: Int32) -> Int32? { @@ -526,15 +560,15 @@ public final class ControlClientAsyncWriter: @unchecked Sendable { fileDescriptor: socket, queue: DispatchQueue.global(qos: .utility) ) - sourceCancellationGroup.enter() + sourceCancellationBarrier.register() writeSource.setEventHandler { [streamContinuation, sourceBox] in streamContinuation.yield(()) streamContinuation.finish() sourceBox.source?.cancel() } - writeSource.setCancelHandler { [sourceCancellationGroup] in + writeSource.setCancelHandler { [sourceCancellationBarrier] in streamContinuation.finish() - sourceCancellationGroup.leave() + sourceCancellationBarrier.complete() } sourceBox.source = writeSource writeSource.activate() @@ -547,15 +581,11 @@ public final class ControlClientAsyncWriter: @unchecked Sendable { streamContinuation.finish() } // The event/cancellation handler may have resumed the iterator before - // libdispatch ran the source's cancellation handler. Wait here while + // 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() - waitForSourceCancellation() + await sourceCancellationBarrier.wait() return writable != nil } - - private func waitForSourceCancellation() { - sourceCancellationGroup.wait() - } } diff --git a/Packages/macOS/CmuxControlSocket/Tests/CmuxControlSocketTests/ControlClientSourceLifecycleTests.swift b/Packages/macOS/CmuxControlSocket/Tests/CmuxControlSocketTests/ControlClientSourceLifecycleTests.swift index 452d1a33cecb..5f024d07182e 100644 --- a/Packages/macOS/CmuxControlSocket/Tests/CmuxControlSocketTests/ControlClientSourceLifecycleTests.swift +++ b/Packages/macOS/CmuxControlSocket/Tests/CmuxControlSocketTests/ControlClientSourceLifecycleTests.swift @@ -13,22 +13,24 @@ struct ControlClientSourceLifecycleTests { group.addTask { for _ in 0..<256 { let pair = try UnixSocketFixture.makeSocketPair() - let reader = ControlClientAsyncLineReader(socket: pair.reader) - defer { - // This is TerminalController.handleClientAsync's - // teardown order. The reader must wait for - // libdispatch to unregister its borrowed - // descriptor before the owner closes it. - reader.cancelAndWait() - shutdown(pair.reader, SHUT_RDWR) - close(pair.reader) - close(pair.writer) - } + let signal = SocketAuthorizationRevocationSignal() + let reader = ControlClientAsyncLineReader( + socket: pair.reader, + authorizationRevocationSignal: signal + ) var line: [UInt8] = [112, 105, 110, 103, 10] #expect(write(pair.writer, &line, line.count) == line.count) // A delivered line proves the read source was armed; // do not substitute a delay before teardown. #expect(await reader.nextLine { true } == "ping") + // This is TerminalController.handleClientAsync's + // teardown order. The reader must wait for + // libdispatch to unregister its borrowed descriptor + // before the owner closes it. + await reader.cancelAndWait() + shutdown(pair.reader, SHUT_RDWR) + close(pair.reader) + close(pair.writer) } return 256 } @@ -38,4 +40,49 @@ struct ControlClientSourceLifecycleTests { #expect(completed == 1_024) } } + + @Test(.timeLimit(.minutes(1))) + func revocationSourceCancelsBeforeOwnerClose() async throws { + let pair = try UnixSocketFixture.makeSocketPair() + let signal = SocketAuthorizationRevocationSignal() + let reader = ControlClientAsyncLineReader( + socket: pair.reader, + authorizationRevocationSignal: signal + ) + let pending = Task { + await reader.nextLine { true } + } + signal.revoke() + #expect(await pending.value == nil) + await reader.cancelAndWait() + shutdown(pair.reader, SHUT_RDWR) + close(pair.reader) + close(pair.writer) + } + + @Test(.timeLimit(.minutes(1))) + func writableSourceCancelsBeforeOwnerClose() async throws { + let pair = try UnixSocketFixture.makeSocketPair() + var sendBuffer = 4 * 1024 + #expect( + setsockopt( + pair.writer, + SOL_SOCKET, + SO_SNDBUF, + &sendBuffer, + socklen_t(MemoryLayout.size) + ) == 0 + ) + let writer = ControlClientAsyncWriter(socket: pair.writer) + let pending = Task { + await writer.writeAll(Data(repeating: 0x58, count: 16 * 1024 * 1024)) + } + try await Task.sleep(nanoseconds: 20_000_000) + pending.cancel() + #expect(await pending.value == false) + await writer.cancelAndWait() + shutdown(pair.writer, SHUT_RDWR) + close(pair.writer) + close(pair.reader) + } } diff --git a/Packages/macOS/CmuxControlSocket/Tests/CmuxControlSocketTests/SocketControlServerManagedPolicyTests.swift b/Packages/macOS/CmuxControlSocket/Tests/CmuxControlSocketTests/SocketControlServerManagedPolicyTests.swift index 60d93ca57e42..f97784e152ba 100644 --- a/Packages/macOS/CmuxControlSocket/Tests/CmuxControlSocketTests/SocketControlServerManagedPolicyTests.swift +++ b/Packages/macOS/CmuxControlSocket/Tests/CmuxControlSocketTests/SocketControlServerManagedPolicyTests.swift @@ -124,7 +124,7 @@ struct SocketControlServerManagedPolicyTests { } try #require(server.reconfigure(accessMode: .cmuxOnly)) #expect(await reading.value == nil) - reader.cancel() + await reader.cancelAndWait() // The connection owns its fd; neither generation revocation nor the // reader may close it and accidentally affect a reused descriptor. #expect(fcntl(connection.socket, F_GETFD) >= 0) diff --git a/Sources/TerminalController.swift b/Sources/TerminalController.swift index dc4bb71e2444..47412e12f081 100644 --- a/Sources/TerminalController.swift +++ b/Sources/TerminalController.swift @@ -2030,14 +2030,46 @@ class TerminalController { initialReadLimits: ControlClientLineReadLimits? = nil, holdsPreauthorizationSlot initialSlotHeld: Bool = false ) async { - // Shut down before close so a DispatchSource callback racing - // cancellation observes EOF rather than a recycled descriptor number. - defer { - shutdown(socket, SHUT_RDWR) - close(socket) - } let pid = peerPid ?? transport.peerProcessID(of: socket) let peerHasSameUID = transport.peerHasSameUID(socket) + let lineReader = ControlClientAsyncLineReader( + socket: socket, + initialLimits: initialReadLimits, + authorizationRevocationSignal: authorizationRevocationSignal + ) + let writer = ControlClientAsyncWriter(socket: socket) + + await handleClientLoop( + socket: socket, + pid: pid, + peerHasSameUID: peerHasSameUID, + authorizationGeneration: authorizationGeneration, + authorizationRevocationSignal: authorizationRevocationSignal, + initialSlotHeld: initialSlotHeld, + lineReader: lineReader, + writer: writer + ) + + // Dispatch source cancellation is asynchronous. Await every borrowed + // socket source before shutting down and closing the descriptor. + await lineReader.cancelAndWait() + await writer.cancelAndWait() + shutdown(socket, SHUT_RDWR) + close(socket) + } + + /// Runs the admitted command loop while retaining ownership of its async + /// socket readers and writer until the caller joins source cancellation. + private nonisolated func handleClientLoop( + socket: Int32, + pid: pid_t?, + peerHasSameUID: Bool, + authorizationGeneration: UInt64, + authorizationRevocationSignal: SocketAuthorizationRevocationSignal, + initialSlotHeld: Bool, + lineReader: ControlClientAsyncLineReader, + writer: ControlClientAsyncWriter + ) async { let preauthorizationLimiter = socketClientPreauthorizationLimiter var holdsPreauthorizationSlot = initialSlotHeld defer { @@ -2046,20 +2078,7 @@ class TerminalController { } } var passwordAuthorization = SocketPasswordAuthorization() - let lineReader = ControlClientAsyncLineReader( - socket: socket, - initialLimits: initialReadLimits, - authorizationRevocationSignal: authorizationRevocationSignal - ) - let writer = ControlClientAsyncWriter(socket: socket) let rateLimiter = ControlClientRateLimiter() - defer { - // Dispatch source cancellation is asynchronous. Wait for every - // borrowed socket source to unregister before the outer defer - // shuts down and closes the descriptor. - lineReader.cancelAndWait() - writer.cancelAndWait() - } while let line = await lineReader.nextLine(shouldContinueReading: { self.socketServer.isConnectionAuthorizationCurrent(authorizationGeneration) }) { From 87e772040cf59f4d5c4f3855b0849908dde2bef4 Mon Sep 17 00:00:00 2001 From: austinpower1258 Date: Sun, 20 Sep 2026 18:35:31 -0700 Subject: [PATCH 3/7] test: make writer cancellation coverage deterministic --- .../ControlClientSourceLifecycleTests.swift | 18 +++++++++++++++++- 1 file changed, 17 insertions(+), 1 deletion(-) diff --git a/Packages/macOS/CmuxControlSocket/Tests/CmuxControlSocketTests/ControlClientSourceLifecycleTests.swift b/Packages/macOS/CmuxControlSocket/Tests/CmuxControlSocketTests/ControlClientSourceLifecycleTests.swift index 5f024d07182e..0f2a206b12f0 100644 --- a/Packages/macOS/CmuxControlSocket/Tests/CmuxControlSocketTests/ControlClientSourceLifecycleTests.swift +++ b/Packages/macOS/CmuxControlSocket/Tests/CmuxControlSocketTests/ControlClientSourceLifecycleTests.swift @@ -73,11 +73,27 @@ struct ControlClientSourceLifecycleTests { socklen_t(MemoryLayout.size) ) == 0 ) + #expect(fcntl(pair.writer, F_SETFL, O_NONBLOCK) == 0) + let fill = [UInt8](repeating: 0x46, count: 64 * 1024) + var reachedWouldBlock = false + while !reachedWouldBlock { + let written = fill.withUnsafeBytes { + Darwin.write(pair.writer, $0.baseAddress, $0.count) + } + if written < 0, errno == EAGAIN || errno == EWOULDBLOCK { + reachedWouldBlock = true + } else { + #expect(written > 0) + } + } + #expect(reachedWouldBlock) let writer = ControlClientAsyncWriter(socket: pair.writer) let pending = Task { await writer.writeAll(Data(repeating: 0x58, count: 16 * 1024 * 1024)) } - try await Task.sleep(nanoseconds: 20_000_000) + for _ in 0..<16 { + await Task.yield() + } pending.cancel() #expect(await pending.value == false) await writer.cancelAndWait() From 0a31fad6069794edd2b345f2ca609ee4eab44a55 Mon Sep 17 00:00:00 2001 From: austinpower1258 Date: Sun, 20 Sep 2026 19:32:36 -0700 Subject: [PATCH 4/7] fix: complete async source teardown review --- .../Server/ControlClientAsyncTransport.swift | 140 ------------------ .../Server/ControlClientAsyncWriter.swift | 121 +++++++++++++++ .../DispatchSourceCancellationBarrier.swift | 52 +++++++ .../ControlClientSourceLifecycleTests.swift | 73 ++++++--- Sources/TerminalController.swift | 11 ++ 5 files changed, 237 insertions(+), 160 deletions(-) create mode 100644 Packages/macOS/CmuxControlSocket/Sources/CmuxControlSocket/Server/ControlClientAsyncWriter.swift create mode 100644 Packages/macOS/CmuxControlSocket/Sources/CmuxControlSocket/Server/DispatchSourceCancellationBarrier.swift diff --git a/Packages/macOS/CmuxControlSocket/Sources/CmuxControlSocket/Server/ControlClientAsyncTransport.swift b/Packages/macOS/CmuxControlSocket/Sources/CmuxControlSocket/Server/ControlClientAsyncTransport.swift index a01f6e937da4..01e590c78e77 100644 --- a/Packages/macOS/CmuxControlSocket/Sources/CmuxControlSocket/Server/ControlClientAsyncTransport.swift +++ b/Packages/macOS/CmuxControlSocket/Sources/CmuxControlSocket/Server/ControlClientAsyncTransport.swift @@ -3,49 +3,6 @@ public import Foundation internal import Darwin internal import os -/// Completes asynchronous teardown after every registered dispatch source has -/// run its cancellation handler. -private final class DispatchSourceCancellationBarrier: @unchecked Sendable { - private let lock = NSLock() - private var registrations = 0 - private var waiters: [CheckedContinuation] = [] - - /// 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) in - lock.lock() - guard registrations > 0 else { - lock.unlock() - continuation.resume() - return - } - waiters.append(continuation) - lock.unlock() - } - } -} - /// Bridges one non-blocking control-socket descriptor into an async line /// reader. The descriptor is never closed by this type; its connection owner /// retains that responsibility. @@ -492,100 +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 - /// 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. - 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 - } - - /// Waits for any one-shot writable source to finish cancellation. - public func cancelAndWait() async { - await sourceCancellationBarrier.wait() - } - - 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.makeStream(bufferingPolicy: .bufferingNewest(1)) - let streamContinuation = stream.continuation - let sourceBox = SourceBox() - let writeSource = DispatchSource.makeWriteSource( - fileDescriptor: socket, - queue: DispatchQueue.global(qos: .utility) - ) - 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 - } -} diff --git a/Packages/macOS/CmuxControlSocket/Sources/CmuxControlSocket/Server/ControlClientAsyncWriter.swift b/Packages/macOS/CmuxControlSocket/Sources/CmuxControlSocket/Server/ControlClientAsyncWriter.swift new file mode 100644 index 000000000000..de3073d6e948 --- /dev/null +++ b/Packages/macOS/CmuxControlSocket/Sources/CmuxControlSocket/Server/ControlClientAsyncWriter.swift @@ -0,0 +1,121 @@ +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 + private let makeWritableSource: @Sendable (Int32) -> any DispatchSourceWrite + /// 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 convenience init(socket: Int32) { + self.init(socket: socket, makeWritableSource: { descriptor in + DispatchSource.makeWriteSource( + fileDescriptor: descriptor, + queue: DispatchQueue.global(qos: .utility) + ) + }) + } + + /// Injects the OS event-source factory independently of response framing. + /// The factory must return an inactive source for the supplied descriptor; + /// the writer installs event/cancel handlers and owns activation/cancellation. + init( + socket: Int32, + makeWritableSource: @escaping @Sendable (Int32) -> any DispatchSourceWrite + ) { + self.socket = socket + self.makeWritableSource = makeWritableSource + _ = 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 + } + + /// 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. + private func waitForWritable() async -> Bool { + let stream = AsyncStream.makeStream(bufferingPolicy: .bufferingNewest(1)) + let streamContinuation = stream.continuation + let sourceBox = SourceBox() + let writeSource = makeWritableSource(socket) + 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 + } +} diff --git a/Packages/macOS/CmuxControlSocket/Sources/CmuxControlSocket/Server/DispatchSourceCancellationBarrier.swift b/Packages/macOS/CmuxControlSocket/Sources/CmuxControlSocket/Server/DispatchSourceCancellationBarrier.swift new file mode 100644 index 000000000000..7be322c33335 --- /dev/null +++ b/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] = [] + + /// 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) in + lock.lock() + guard registrations > 0 else { + lock.unlock() + continuation.resume() + return + } + waiters.append(continuation) + lock.unlock() + } + } +} + diff --git a/Packages/macOS/CmuxControlSocket/Tests/CmuxControlSocketTests/ControlClientSourceLifecycleTests.swift b/Packages/macOS/CmuxControlSocket/Tests/CmuxControlSocketTests/ControlClientSourceLifecycleTests.swift index 0f2a206b12f0..28c0ac86d902 100644 --- a/Packages/macOS/CmuxControlSocket/Tests/CmuxControlSocketTests/ControlClientSourceLifecycleTests.swift +++ b/Packages/macOS/CmuxControlSocket/Tests/CmuxControlSocketTests/ControlClientSourceLifecycleTests.swift @@ -1,4 +1,5 @@ import Darwin +import Dispatch import Foundation import Testing @testable import CmuxControlSocket @@ -6,6 +7,7 @@ import Testing /// Exercises the same cancellation followed by close used by a CLI connection. @Suite(.serialized) struct ControlClientSourceLifecycleTests { + /// Stresses reuse after read and revocation sources complete cancellation. @Test(.timeLimit(.minutes(1))) func repeatedReaderCancellationBeforeOwnerClose() async throws { try await withThrowingTaskGroup(of: Int.self) { group in @@ -41,6 +43,7 @@ struct ControlClientSourceLifecycleTests { } } + /// Revoking an idle connection must still join both descriptor sources. @Test(.timeLimit(.minutes(1))) func revocationSourceCancelsBeforeOwnerClose() async throws { let pair = try UnixSocketFixture.makeSocketPair() @@ -60,45 +63,75 @@ struct ControlClientSourceLifecycleTests { close(pair.writer) } + /// Cancels only after libdispatch registers the real backpressured source. @Test(.timeLimit(.minutes(1))) func writableSourceCancelsBeforeOwnerClose() async throws { let pair = try UnixSocketFixture.makeSocketPair() - var sendBuffer = 4 * 1024 - #expect( + defer { + close(pair.writer) + close(pair.reader) + } + var sendBuffer: Int32 = 4 * 1024 + try #require( setsockopt( pair.writer, SOL_SOCKET, SO_SNDBUF, &sendBuffer, - socklen_t(MemoryLayout.size) + socklen_t(MemoryLayout.size) ) == 0 ) - #expect(fcntl(pair.writer, F_SETFL, O_NONBLOCK) == 0) - let fill = [UInt8](repeating: 0x46, count: 64 * 1024) - var reachedWouldBlock = false - while !reachedWouldBlock { + try #require(fcntl(pair.writer, F_SETFL, O_NONBLOCK) == 0) + let fill = [UInt8](repeating: 0x46, count: 4 * 1024) + var filledBytes = 0 + while true { let written = fill.withUnsafeBytes { Darwin.write(pair.writer, $0.baseAddress, $0.count) } - if written < 0, errno == EAGAIN || errno == EWOULDBLOCK { - reachedWouldBlock = true - } else { - #expect(written > 0) - } + if written < 0, errno == EINTR { continue } + if written < 0, errno == EAGAIN || errno == EWOULDBLOCK { break } + try #require(written > 0) + filledBytes += written + try #require(filledBytes <= 1024 * 1024, "Socket never applied backpressure") } - #expect(reachedWouldBlock) - let writer = ControlClientAsyncWriter(socket: pair.writer) + try #require(filledBytes > 0) + + let registrations = AsyncStream.makeStream(bufferingPolicy: .bufferingNewest(1)) + defer { registrations.continuation.finish() } + let writer = ControlClientAsyncWriter(socket: pair.writer, makeWritableSource: { descriptor in + let source = DispatchSource.makeWriteSource( + fileDescriptor: descriptor, + queue: .global(qos: .utility) + ) + source.setRegistrationHandler { + registrations.continuation.yield(()) + } + return source + }) + // Exercise combined connection teardown: a pending read, generation + // revocation, and a backpressured write all share the accepted socket. + let signal = SocketAuthorizationRevocationSignal() + try #require(signal.readFileDescriptor >= 0) + let reader = ControlClientAsyncLineReader( + socket: pair.writer, + authorizationRevocationSignal: signal + ) + let reading = Task { await reader.nextLine { true } } let pending = Task { - await writer.writeAll(Data(repeating: 0x58, count: 16 * 1024 * 1024)) - } - for _ in 0..<16 { - await Task.yield() + await writer.writeAll(Data([0x58])) } + var registration = registrations.stream.makeAsyncIterator() + let registered = await registration.next() + // This event comes from libdispatch after activation, not from a yield + // count or elapsed time. An early-cancel-only writer cannot satisfy it. pending.cancel() + signal.revoke() #expect(await pending.value == false) + #expect(await reading.value == nil) + await reader.cancelAndWait() await writer.cancelAndWait() + #expect(registered != nil) + #expect(fcntl(pair.writer, F_GETFD) >= 0) shutdown(pair.writer, SHUT_RDWR) - close(pair.writer) - close(pair.reader) } } diff --git a/Sources/TerminalController.swift b/Sources/TerminalController.swift index 47412e12f081..19f2b7c611d5 100644 --- a/Sources/TerminalController.swift +++ b/Sources/TerminalController.swift @@ -2022,6 +2022,12 @@ class TerminalController { guard submission != .rejected else { return } } + /// Owns the accepted socket until the command loop and source teardown finish. + #if compiler(>=6.2) + @concurrent + #else + @Sendable + #endif private nonisolated func handleClientAsync( _ socket: Int32, peerPid: pid_t? = nil, @@ -2060,6 +2066,11 @@ class TerminalController { /// Runs the admitted command loop while retaining ownership of its async /// socket readers and writer until the caller joins source cancellation. + #if compiler(>=6.2) + @concurrent + #else + @Sendable + #endif private nonisolated func handleClientLoop( socket: Int32, pid: pid_t?, From ee5eef59fe09c85708426e91dbbdb9526bfc18fe Mon Sep 17 00:00:00 2001 From: austinpower1258 Date: Sun, 20 Sep 2026 19:42:20 -0700 Subject: [PATCH 5/7] test: observe writable registration through the runtime operation --- .../Server/ControlClientAsyncTransport.swift | 2 +- .../Server/ControlClientAsyncWriter.swift | 30 ++++++------------- .../ControlClientSourceLifecycleTests.swift | 26 ++++++++-------- 3 files changed, 23 insertions(+), 35 deletions(-) diff --git a/Packages/macOS/CmuxControlSocket/Sources/CmuxControlSocket/Server/ControlClientAsyncTransport.swift b/Packages/macOS/CmuxControlSocket/Sources/CmuxControlSocket/Server/ControlClientAsyncTransport.swift index 01e590c78e77..cf57d44ccf22 100644 --- a/Packages/macOS/CmuxControlSocket/Sources/CmuxControlSocket/Server/ControlClientAsyncTransport.swift +++ b/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 diff --git a/Packages/macOS/CmuxControlSocket/Sources/CmuxControlSocket/Server/ControlClientAsyncWriter.swift b/Packages/macOS/CmuxControlSocket/Sources/CmuxControlSocket/Server/ControlClientAsyncWriter.swift index de3073d6e948..25e3cf01e4c7 100644 --- a/Packages/macOS/CmuxControlSocket/Sources/CmuxControlSocket/Server/ControlClientAsyncWriter.swift +++ b/Packages/macOS/CmuxControlSocket/Sources/CmuxControlSocket/Server/ControlClientAsyncWriter.swift @@ -12,7 +12,6 @@ public final class ControlClientAsyncWriter: @unchecked Sendable { } private let socket: Int32 - private let makeWritableSource: @Sendable (Int32) -> any DispatchSourceWrite /// One-shot writable sources must finish cancellation before the owner /// closes the shared socket descriptor. private let sourceCancellationBarrier = DispatchSourceCancellationBarrier() @@ -20,24 +19,8 @@ public final class ControlClientAsyncWriter: @unchecked Sendable { /// Creates a writer over a non-blocking descriptor. /// /// - Parameter socket: A borrowed descriptor retained by the connection owner. - public convenience init(socket: Int32) { - self.init(socket: socket, makeWritableSource: { descriptor in - DispatchSource.makeWriteSource( - fileDescriptor: descriptor, - queue: DispatchQueue.global(qos: .utility) - ) - }) - } - - /// Injects the OS event-source factory independently of response framing. - /// The factory must return an inactive source for the supplied descriptor; - /// the writer installs event/cancel handlers and owns activation/cancellation. - init( - socket: Int32, - makeWritableSource: @escaping @Sendable (Int32) -> any DispatchSourceWrite - ) { + public init(socket: Int32) { self.socket = socket - self.makeWritableSource = makeWritableSource _ = Self.makeNonBlocking(socket) } @@ -60,7 +43,11 @@ public final class ControlClientAsyncWriter: @unchecked Sendable { } if written < 0, errno == EINTR { continue } if written < 0, errno == EAGAIN || errno == EWOULDBLOCK { - guard await waitForWritable() else { return false } + let source = DispatchSource.makeWriteSource( + fileDescriptor: socket, + queue: DispatchQueue.global(qos: .utility) + ) + guard await waitForWritable(source: source) else { return false } continue } return false @@ -85,11 +72,12 @@ public final class ControlClientAsyncWriter: @unchecked Sendable { } /// Joins cancellation of a one-shot source on readiness or task cancellation. - private func waitForWritable() async -> Bool { + /// - 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.makeStream(bufferingPolicy: .bufferingNewest(1)) let streamContinuation = stream.continuation let sourceBox = SourceBox() - let writeSource = makeWritableSource(socket) sourceCancellationBarrier.register() writeSource.setEventHandler { [streamContinuation, sourceBox] in streamContinuation.yield(()) diff --git a/Packages/macOS/CmuxControlSocket/Tests/CmuxControlSocketTests/ControlClientSourceLifecycleTests.swift b/Packages/macOS/CmuxControlSocket/Tests/CmuxControlSocketTests/ControlClientSourceLifecycleTests.swift index 28c0ac86d902..6bffac7a56c2 100644 --- a/Packages/macOS/CmuxControlSocket/Tests/CmuxControlSocketTests/ControlClientSourceLifecycleTests.swift +++ b/Packages/macOS/CmuxControlSocket/Tests/CmuxControlSocketTests/ControlClientSourceLifecycleTests.swift @@ -98,16 +98,14 @@ struct ControlClientSourceLifecycleTests { let registrations = AsyncStream.makeStream(bufferingPolicy: .bufferingNewest(1)) defer { registrations.continuation.finish() } - let writer = ControlClientAsyncWriter(socket: pair.writer, makeWritableSource: { descriptor in - let source = DispatchSource.makeWriteSource( - fileDescriptor: descriptor, - queue: .global(qos: .utility) - ) - source.setRegistrationHandler { - registrations.continuation.yield(()) - } - return source - }) + let writer = ControlClientAsyncWriter(socket: pair.writer) + let source = DispatchSource.makeWriteSource( + fileDescriptor: pair.writer, + queue: .global(qos: .utility) + ) + source.setRegistrationHandler { + registrations.continuation.yield(()) + } // Exercise combined connection teardown: a pending read, generation // revocation, and a backpressured write all share the accepted socket. let signal = SocketAuthorizationRevocationSignal() @@ -118,12 +116,14 @@ struct ControlClientSourceLifecycleTests { ) let reading = Task { await reader.nextLine { true } } let pending = Task { - await writer.writeAll(Data([0x58])) + // Call the same runtime operation used by writeAll after EAGAIN. + // Source observation stays in the test target; no factory hook. + await writer.waitForWritable(source: source) } var registration = registrations.stream.makeAsyncIterator() - let registered = await registration.next() + let registered: Void? = await registration.next() // This event comes from libdispatch after activation, not from a yield - // count or elapsed time. An early-cancel-only writer cannot satisfy it. + // count or elapsed time. The runtime wait must activate this source. pending.cancel() signal.revoke() #expect(await pending.value == false) From cb453431fb24a2844a190bea7c9296cc19614981 Mon Sep 17 00:00:00 2001 From: austinpower1258 Date: Sun, 20 Sep 2026 19:48:19 -0700 Subject: [PATCH 6/7] test: drive real writer backpressure through an observed executor --- .../Server/ControlClientAsyncWriter.swift | 16 ++++----- .../DispatchSourceCancellationBarrier.swift | 6 ++-- .../ControlClientSourceLifecycleTests.swift | 33 +++++++++++-------- .../SocketWriterTestExecutor.swift | 20 +++++++++++ 4 files changed, 50 insertions(+), 25 deletions(-) create mode 100644 Packages/macOS/CmuxControlSocket/Tests/CmuxControlSocketTests/SocketWriterTestExecutor.swift diff --git a/Packages/macOS/CmuxControlSocket/Sources/CmuxControlSocket/Server/ControlClientAsyncWriter.swift b/Packages/macOS/CmuxControlSocket/Sources/CmuxControlSocket/Server/ControlClientAsyncWriter.swift index 25e3cf01e4c7..aa3b170dfd0f 100644 --- a/Packages/macOS/CmuxControlSocket/Sources/CmuxControlSocket/Server/ControlClientAsyncWriter.swift +++ b/Packages/macOS/CmuxControlSocket/Sources/CmuxControlSocket/Server/ControlClientAsyncWriter.swift @@ -14,7 +14,7 @@ public final class ControlClientAsyncWriter: @unchecked Sendable { private let socket: Int32 /// One-shot writable sources must finish cancellation before the owner /// closes the shared socket descriptor. - private let sourceCancellationBarrier = DispatchSourceCancellationBarrier() + let sourceCancellationBarrier = DispatchSourceCancellationBarrier() /// Creates a writer over a non-blocking descriptor. /// @@ -43,11 +43,7 @@ public final class ControlClientAsyncWriter: @unchecked Sendable { } 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 } + guard await waitForWritable() else { return false } continue } return false @@ -72,12 +68,14 @@ public final class ControlClientAsyncWriter: @unchecked Sendable { } /// 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 { + private func waitForWritable() async -> Bool { let stream = AsyncStream.makeStream(bufferingPolicy: .bufferingNewest(1)) let streamContinuation = stream.continuation let sourceBox = SourceBox() + let writeSource = DispatchSource.makeWriteSource( + fileDescriptor: socket, + queue: DispatchQueue.global(qos: .utility) + ) sourceCancellationBarrier.register() writeSource.setEventHandler { [streamContinuation, sourceBox] in streamContinuation.yield(()) diff --git a/Packages/macOS/CmuxControlSocket/Sources/CmuxControlSocket/Server/DispatchSourceCancellationBarrier.swift b/Packages/macOS/CmuxControlSocket/Sources/CmuxControlSocket/Server/DispatchSourceCancellationBarrier.swift index 7be322c33335..b6bdddc7ac59 100644 --- a/Packages/macOS/CmuxControlSocket/Sources/CmuxControlSocket/Server/DispatchSourceCancellationBarrier.swift +++ b/Packages/macOS/CmuxControlSocket/Sources/CmuxControlSocket/Server/DispatchSourceCancellationBarrier.swift @@ -10,8 +10,9 @@ final class DispatchSourceCancellationBarrier: @unchecked Sendable { // 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 + let lock = NSLock() + /// The source-of-truth count; readers must hold ``lock``. + internal private(set) var registrations = 0 private var waiters: [CheckedContinuation] = [] /// Registers one source before it is activated. @@ -49,4 +50,3 @@ final class DispatchSourceCancellationBarrier: @unchecked Sendable { } } } - diff --git a/Packages/macOS/CmuxControlSocket/Tests/CmuxControlSocketTests/ControlClientSourceLifecycleTests.swift b/Packages/macOS/CmuxControlSocket/Tests/CmuxControlSocketTests/ControlClientSourceLifecycleTests.swift index 6bffac7a56c2..1651bc6d4398 100644 --- a/Packages/macOS/CmuxControlSocket/Tests/CmuxControlSocketTests/ControlClientSourceLifecycleTests.swift +++ b/Packages/macOS/CmuxControlSocket/Tests/CmuxControlSocketTests/ControlClientSourceLifecycleTests.swift @@ -63,7 +63,16 @@ struct ControlClientSourceLifecycleTests { close(pair.writer) } - /// Cancels only after libdispatch registers the real backpressured source. + /// Reads the runtime count under the barrier's existing synchronization. + private func sourceCount(_ writer: ControlClientAsyncWriter) -> Int { + let barrier = writer.sourceCancellationBarrier + barrier.lock.lock() + defer { barrier.lock.unlock() } + return barrier.registrations + } + + /// Cancels only after writeAll suspends with its real EAGAIN source active. + @available(macOS 15.0, *) @Test(.timeLimit(.minutes(1))) func writableSourceCancelsBeforeOwnerClose() async throws { let pair = try UnixSocketFixture.makeSocketPair() @@ -99,12 +108,10 @@ struct ControlClientSourceLifecycleTests { let registrations = AsyncStream.makeStream(bufferingPolicy: .bufferingNewest(1)) defer { registrations.continuation.finish() } let writer = ControlClientAsyncWriter(socket: pair.writer) - let source = DispatchSource.makeWriteSource( - fileDescriptor: pair.writer, - queue: .global(qos: .utility) - ) - source.setRegistrationHandler { - registrations.continuation.yield(()) + let executor = SocketWriterTestExecutor { + if sourceCount(writer) == 1 { + registrations.continuation.yield(()) + } } // Exercise combined connection teardown: a pending read, generation // revocation, and a backpressured write all share the accepted socket. @@ -115,15 +122,14 @@ struct ControlClientSourceLifecycleTests { authorizationRevocationSignal: signal ) let reading = Task { await reader.nextLine { true } } - let pending = Task { - // Call the same runtime operation used by writeAll after EAGAIN. - // Source observation stays in the test target; no factory hook. - await writer.waitForWritable(source: source) + let pending = Task(executorPreference: executor) { + await writer.writeAll(Data(repeating: 0x58, count: 64 * 1024)) } var registration = registrations.stream.makeAsyncIterator() let registered: Void? = await registration.next() - // This event comes from libdispatch after activation, not from a yield - // count or elapsed time. The runtime wait must activate this source. + // The executor reports runtime state only after writeAll suspends. + // A full buffer prevents readiness, so a count of one proves EAGAIN, + // source creation, activation, and suspension all occurred before cancel. pending.cancel() signal.revoke() #expect(await pending.value == false) @@ -131,6 +137,7 @@ struct ControlClientSourceLifecycleTests { await reader.cancelAndWait() await writer.cancelAndWait() #expect(registered != nil) + #expect(sourceCount(writer) == 0) #expect(fcntl(pair.writer, F_GETFD) >= 0) shutdown(pair.writer, SHUT_RDWR) } diff --git a/Packages/macOS/CmuxControlSocket/Tests/CmuxControlSocketTests/SocketWriterTestExecutor.swift b/Packages/macOS/CmuxControlSocket/Tests/CmuxControlSocketTests/SocketWriterTestExecutor.swift new file mode 100644 index 000000000000..985cdcb42d46 --- /dev/null +++ b/Packages/macOS/CmuxControlSocket/Tests/CmuxControlSocketTests/SocketWriterTestExecutor.swift @@ -0,0 +1,20 @@ +import Dispatch + +/// Observes the real writer task at suspension boundaries, never on a timer. +@available(macOS 15.0, *) +final class SocketWriterTestExecutor: TaskExecutor { + private let queue = DispatchQueue(label: "cmux.tests.socket-writer-executor") + private let didRunJob: @Sendable () -> Void + + init(didRunJob: @escaping @Sendable () -> Void) { + self.didRunJob = didRunJob + } + + func enqueue(_ job: consuming ExecutorJob) { + let job = UnownedJob(job) + queue.async { [self] in + job.runSynchronously(on: asUnownedTaskExecutor()) + didRunJob() + } + } +} From d4bad807fa520f7785fa509a805eb7d33b60c29f Mon Sep 17 00:00:00 2001 From: austinpower1258 Date: Sun, 20 Sep 2026 21:20:17 -0700 Subject: [PATCH 7/7] style: add explicit deinitializers to socket lifecycle types --- .../CmuxControlSocket/Server/ControlClientAsyncWriter.swift | 4 ++++ .../Server/DispatchSourceCancellationBarrier.swift | 2 ++ .../CmuxControlSocketTests/SocketWriterTestExecutor.swift | 2 ++ 3 files changed, 8 insertions(+) diff --git a/Packages/macOS/CmuxControlSocket/Sources/CmuxControlSocket/Server/ControlClientAsyncWriter.swift b/Packages/macOS/CmuxControlSocket/Sources/CmuxControlSocket/Server/ControlClientAsyncWriter.swift index aa3b170dfd0f..1e3c9296d136 100644 --- a/Packages/macOS/CmuxControlSocket/Sources/CmuxControlSocket/Server/ControlClientAsyncWriter.swift +++ b/Packages/macOS/CmuxControlSocket/Sources/CmuxControlSocket/Server/ControlClientAsyncWriter.swift @@ -9,6 +9,8 @@ internal import Darwin public final class ControlClientAsyncWriter: @unchecked Sendable { private final class SourceBox: @unchecked Sendable { var source: (any DispatchSourceWrite)? + + deinit {} } private let socket: Int32 @@ -16,6 +18,8 @@ public final class ControlClientAsyncWriter: @unchecked Sendable { /// closes the shared socket descriptor. let sourceCancellationBarrier = DispatchSourceCancellationBarrier() + deinit {} + /// Creates a writer over a non-blocking descriptor. /// /// - Parameter socket: A borrowed descriptor retained by the connection owner. diff --git a/Packages/macOS/CmuxControlSocket/Sources/CmuxControlSocket/Server/DispatchSourceCancellationBarrier.swift b/Packages/macOS/CmuxControlSocket/Sources/CmuxControlSocket/Server/DispatchSourceCancellationBarrier.swift index b6bdddc7ac59..455571774d5b 100644 --- a/Packages/macOS/CmuxControlSocket/Sources/CmuxControlSocket/Server/DispatchSourceCancellationBarrier.swift +++ b/Packages/macOS/CmuxControlSocket/Sources/CmuxControlSocket/Server/DispatchSourceCancellationBarrier.swift @@ -15,6 +15,8 @@ final class DispatchSourceCancellationBarrier: @unchecked Sendable { internal private(set) var registrations = 0 private var waiters: [CheckedContinuation] = [] + deinit {} + /// Registers one source before it is activated. func register() { lock.lock() diff --git a/Packages/macOS/CmuxControlSocket/Tests/CmuxControlSocketTests/SocketWriterTestExecutor.swift b/Packages/macOS/CmuxControlSocket/Tests/CmuxControlSocketTests/SocketWriterTestExecutor.swift index 985cdcb42d46..7103ab681f43 100644 --- a/Packages/macOS/CmuxControlSocket/Tests/CmuxControlSocketTests/SocketWriterTestExecutor.swift +++ b/Packages/macOS/CmuxControlSocket/Tests/CmuxControlSocketTests/SocketWriterTestExecutor.swift @@ -6,6 +6,8 @@ final class SocketWriterTestExecutor: TaskExecutor { private let queue = DispatchQueue(label: "cmux.tests.socket-writer-executor") private let didRunJob: @Sendable () -> Void + deinit {} + init(didRunJob: @escaping @Sendable () -> Void) { self.didRunJob = didRunJob }