diff --git a/Packages/macOS/CmuxControlSocket/Sources/CmuxControlSocket/Server/ControlClientAsyncTransport.swift b/Packages/macOS/CmuxControlSocket/Sources/CmuxControlSocket/Server/ControlClientAsyncTransport.swift index c2838f39c33b..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 @@ -46,6 +46,9 @@ 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 + /// 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.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 - } -} 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..1e3c9296d136 --- /dev/null +++ b/Packages/macOS/CmuxControlSocket/Sources/CmuxControlSocket/Server/ControlClientAsyncWriter.swift @@ -0,0 +1,111 @@ +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)? + + deinit {} + } + + private let socket: Int32 + /// One-shot writable sources must finish cancellation before the owner + /// 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. + 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. + /// + /// 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 = 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/DispatchSourceCancellationBarrier.swift b/Packages/macOS/CmuxControlSocket/Sources/CmuxControlSocket/Server/DispatchSourceCancellationBarrier.swift new file mode 100644 index 000000000000..455571774d5b --- /dev/null +++ b/Packages/macOS/CmuxControlSocket/Sources/CmuxControlSocket/Server/DispatchSourceCancellationBarrier.swift @@ -0,0 +1,54 @@ +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. + let lock = NSLock() + /// The source-of-truth count; readers must hold ``lock``. + internal private(set) var registrations = 0 + private var waiters: [CheckedContinuation] = [] + + deinit {} + + /// 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 new file mode 100644 index 000000000000..1651bc6d4398 --- /dev/null +++ b/Packages/macOS/CmuxControlSocket/Tests/CmuxControlSocketTests/ControlClientSourceLifecycleTests.swift @@ -0,0 +1,144 @@ +import Darwin +import Dispatch +import Foundation +import Testing +@testable import CmuxControlSocket + +/// 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 + for _ in 0..<4 { + group.addTask { + for _ in 0..<256 { + let pair = try UnixSocketFixture.makeSocketPair() + 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 + } + } + var completed = 0 + for try await count in group { completed += count } + #expect(completed == 1_024) + } + } + + /// Revoking an idle connection must still join both descriptor sources. + @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) + } + + /// 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() + 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) + ) == 0 + ) + 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 == 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") + } + try #require(filledBytes > 0) + + let registrations = AsyncStream.makeStream(bufferingPolicy: .bufferingNewest(1)) + defer { registrations.continuation.finish() } + let writer = ControlClientAsyncWriter(socket: pair.writer) + 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. + 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(executorPreference: executor) { + await writer.writeAll(Data(repeating: 0x58, count: 64 * 1024)) + } + var registration = registrations.stream.makeAsyncIterator() + let registered: Void? = await registration.next() + // 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) + #expect(await reading.value == nil) + 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/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/Packages/macOS/CmuxControlSocket/Tests/CmuxControlSocketTests/SocketWriterTestExecutor.swift b/Packages/macOS/CmuxControlSocket/Tests/CmuxControlSocketTests/SocketWriterTestExecutor.swift new file mode 100644 index 000000000000..7103ab681f43 --- /dev/null +++ b/Packages/macOS/CmuxControlSocket/Tests/CmuxControlSocketTests/SocketWriterTestExecutor.swift @@ -0,0 +1,22 @@ +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 + + deinit {} + + 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() + } + } +} diff --git a/Sources/TerminalController.swift b/Sources/TerminalController.swift index 5878e025296f..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, @@ -2030,14 +2036,51 @@ 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. + #if compiler(>=6.2) + @concurrent + #else + @Sendable + #endif + 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,17 +2089,7 @@ class TerminalController { } } var passwordAuthorization = SocketPasswordAuthorization() - let lineReader = ControlClientAsyncLineReader( - socket: socket, - initialLimits: initialReadLimits, - authorizationRevocationSignal: authorizationRevocationSignal - ) - let writer = ControlClientAsyncWriter(socket: socket) let rateLimiter = ControlClientRateLimiter() - defer { - lineReader.cancel() - writer.cancel() - } while let line = await lineReader.nextLine(shouldContinueReading: { self.socketServer.isConnectionAuthorizationCurrent(authorizationGeneration) }) {