[fix]: eliminate race condition in SocketController

This commit is contained in:
Jamie 2025-12-03 20:17:09 -06:00
parent bb0b6d8dc3
commit 8ed9f275af

View File

@ -37,7 +37,15 @@ public struct SocketController {
port = SocketPort(path: path)
fileHandle = FileHandle(fileDescriptor: port.socket, closeOnDealloc: true)
Task { @MainActor [fileHandle, sessionsContinuation, logger] in
for await notification in NotificationCenter.default.notifications(named: .NSFileHandleConnectionAccepted) {
// Create the iterator before triggering the notification to
// ensure it will not be missed.
let connectionAvailableIterator = NotificationCenter.default
.notifications(named: .NSFileHandleConnectionAccepted)
.makeAsyncIterator()
fileHandle.acceptConnectionInBackgroundAndNotify()
while let notification = await connectionAvailableIterator.next() {
logger.debug("Socket controller accepted connection")
guard let new = notification.userInfo?[NSFileHandleNotificationFileHandleItem] as? FileHandle else { continue }
let session = Session(fileHandle: new)
@ -45,7 +53,6 @@ public struct SocketController {
fileHandle.acceptConnectionInBackgroundAndNotify()
}
}
fileHandle.acceptConnectionInBackgroundAndNotify()
logger.debug("Socket listening at \(path)")
}
@ -77,8 +84,16 @@ extension SocketController {
self.fileHandle = fileHandle
provenance = SigningRequestTracer().provenance(from: fileHandle)
(messages, messagesContinuation) = AsyncStream.makeStream()
Task { [messagesContinuation, logger] in
for await _ in NotificationCenter.default.notifications(named: .NSFileHandleDataAvailable, object: fileHandle) {
Task { @MainActor [messagesContinuation, logger] in
// Create the iterator before triggering the notification to
// ensure it will not be missed.
let dataAvailableIterator = NotificationCenter.default
.notifications(named: .NSFileHandleDataAvailable, object: fileHandle)
.makeAsyncIterator()
fileHandle.waitForDataInBackgroundAndNotify()
while let _ = await dataAvailableIterator.next() {
let data = fileHandle.availableData
guard !data.isEmpty else {
logger.debug("Socket controller received empty data, ending continuation.")
@ -90,17 +105,14 @@ extension SocketController {
logger.debug("Socket controller yielded data.")
}
}
fileHandle.waitForDataInBackgroundAndNotify()
}
/// Writes new data to the socket.
/// - Parameter data: The data to write.
public func write(_ data: Data) async throws {
@MainActor public func write(_ data: Data) throws {
try fileHandle.write(contentsOf: data)
await MainActor.run {
fileHandle.waitForDataInBackgroundAndNotify()
}
}
/// Closes the socket and cleans up resources.
public func close() throws {