diff --git a/Sources/PostgresNIO/New/NotificationListener.swift b/Sources/PostgresNIO/New/NotificationListener.swift index 708b4ff0..58ba8757 100644 --- a/Sources/PostgresNIO/New/NotificationListener.swift +++ b/Sources/PostgresNIO/New/NotificationListener.swift @@ -1,24 +1,29 @@ import NIOCore -// This object is @unchecked Sendable, since we synchronize state on the EL -final class NotificationListener: @unchecked Sendable { +// Thread safety is guaranteed in the NotificationListener by dispatching all access to the shared +// state onto the underlying EventLoop. +final class NotificationListener: Sendable { let eventLoop: any EventLoop let channel: String let id: Int - private var state: State + private let stateBox: NIOLoopBoundBox enum State { case streamInitialized(CheckedContinuation) case streamListening(AsyncThrowingStream.Continuation) - case closure(PostgresListenContext, (PostgresListenContext, PostgresMessage.NotificationResponse) -> Void) + case closure(PostgresListenContext, @Sendable (PostgresListenContext, PostgresMessage.NotificationResponse) -> Void) case done } deinit { - switch self.state { + // The state may only be inspected on the EventLoop. If we are deinitialized somewhere else, + // we can't validate that the listener has been used correctly. + guard self.eventLoop.inEventLoop else { return } + + switch self.stateBox.value { case .streamInitialized: preconditionFailure("Notification continuation had not been used") case .closure: @@ -37,7 +42,10 @@ final class NotificationListener: @unchecked Sendable { self.channel = channel self.id = id self.eventLoop = eventLoop - self.state = .streamInitialized(checkedContinuation) + self.stateBox = NIOLoopBoundBox.makeBoxSendingValue( + .streamInitialized(checkedContinuation), + eventLoop: eventLoop + ) } init( @@ -50,104 +58,176 @@ final class NotificationListener: @unchecked Sendable { self.channel = channel self.id = id self.eventLoop = eventLoop - self.state = .closure(context, closure) + self.stateBox = NIOLoopBoundBox.makeBoxSendingValue( + .closure(context, closure), + eventLoop: eventLoop + ) + } + + // Every state modification below returns the side effects it wants to have run as an action. The + // actions are then run *after* the modifying closure has returned. This is important, since all + // side effects (resuming continuations, yielding into the stream, invoking the user supplied + // closure, ...) may reenter this type. Running them while we hold exclusive access to the state + // would violate the exclusivity of that access. + + private enum StartListeningSucceededAction { + case none + case resumeContinuation( + CheckedContinuation, + PostgresNotificationSequence + ) } func startListeningSucceeded(handler: PostgresChannelHandler) { self.eventLoop.preconditionInEventLoop() let handlerLoopBound = NIOLoopBound(handler, eventLoop: self.eventLoop) - - switch self.state { - case .streamInitialized(let checkedContinuation): - let (stream, continuation) = AsyncThrowingStream.makeStream(of: PostgresNotification.self) - let eventLoop = self.eventLoop - let channel = self.channel - let listenerID = self.id - continuation.onTermination = { reason in - switch reason { - case .cancelled: - eventLoop.execute { - handlerLoopBound.value.cancelNotificationListener(channel: channel, id: listenerID) + let eventLoop = self.eventLoop + let channel = self.channel + let listenerID = self.id + + let action = self.stateBox.withValue { state -> StartListeningSucceededAction in + switch state { + case .streamInitialized(let checkedContinuation): + let (stream, continuation) = AsyncThrowingStream.makeStream(of: PostgresNotification.self) + continuation.onTermination = { reason in + switch reason { + case .cancelled: + eventLoop.execute { + handlerLoopBound.value.cancelNotificationListener(channel: channel, id: listenerID) + } + + case .finished: + break + + @unknown default: + break } + } + state = .streamListening(continuation) - case .finished: - break + return .resumeContinuation(checkedContinuation, PostgresNotificationSequence(base: stream)) - @unknown default: - break - } + case .streamListening, .done: + fatalError("Invalid state: \(state)") + + case .closure: + return .none // ignore } - self.state = .streamListening(continuation) + } - let notificationSequence = PostgresNotificationSequence(base: stream) + switch action { + case .none: + break + case .resumeContinuation(let checkedContinuation, let notificationSequence): checkedContinuation.resume(returning: notificationSequence) - - case .streamListening, .done: - fatalError("Invalid state: \(self.state)") - - case .closure: - break // ignore } } + private enum NotificationReceivedAction { + case yield(AsyncThrowingStream.Continuation, PostgresNotification) + case invokeClosure( + PostgresListenContext, + @Sendable (PostgresListenContext, PostgresMessage.NotificationResponse) -> Void, + PostgresMessage.NotificationResponse + ) + } + func notificationReceived(_ backendMessage: PostgresBackendMessage.NotificationResponse) { self.eventLoop.preconditionInEventLoop() - switch self.state { - case .streamInitialized, .done: - fatalError("Invalid state: \(self.state)") - case .streamListening(let continuation): - continuation.yield(.init(payload: backendMessage.payload)) - - case .closure(let postgresListenContext, let closure): - let message = PostgresMessage.NotificationResponse( - backendPID: backendMessage.backendPID, - channel: backendMessage.channel, - payload: backendMessage.payload - ) + let action = self.stateBox.withValue { state -> NotificationReceivedAction in + switch state { + case .streamInitialized, .done: + fatalError("Invalid state: \(state)") + + case .streamListening(let continuation): + return .yield(continuation, .init(payload: backendMessage.payload)) + + case .closure(let postgresListenContext, let closure): + let message = PostgresMessage.NotificationResponse( + backendPID: backendMessage.backendPID, + channel: backendMessage.channel, + payload: backendMessage.payload + ) + return .invokeClosure(postgresListenContext, closure, message) + } + } + + switch action { + case .yield(let continuation, let notification): + continuation.yield(notification) + case .invokeClosure(let postgresListenContext, let closure, let message): closure(postgresListenContext, message) } } + private enum EndAction { + case none + case failContinuation(CheckedContinuation, any Error) + case finishStream(AsyncThrowingStream.Continuation, (any Error)?) + case cancelListenContext(PostgresListenContext) + } + func failed(_ error: any Error) { self.eventLoop.preconditionInEventLoop() - switch self.state { - case .streamInitialized(let checkedContinuation): - self.state = .done - checkedContinuation.resume(throwing: error) + let action = self.stateBox.withValue { state -> EndAction in + switch state { + case .streamInitialized(let checkedContinuation): + state = .done + return .failContinuation(checkedContinuation, error) - case .streamListening(let continuation): - self.state = .done - continuation.finish(throwing: error) + case .streamListening(let continuation): + state = .done + return .finishStream(continuation, error) - case .closure(let postgresListenContext, _): - self.state = .done - postgresListenContext.cancel() + case .closure(let postgresListenContext, _): + state = .done + return .cancelListenContext(postgresListenContext) - case .done: - break // ignore + case .done: + return .none // ignore + } } + + self.run(action) } func cancelled() { self.eventLoop.preconditionInEventLoop() - switch self.state { - case .streamInitialized(let checkedContinuation): - self.state = .done - checkedContinuation.resume(throwing: PSQLError(code: .queryCancelled)) + let action = self.stateBox.withValue { state -> EndAction in + switch state { + case .streamInitialized(let checkedContinuation): + state = .done + return .failContinuation(checkedContinuation, PSQLError(code: .queryCancelled)) - case .streamListening(let continuation): - self.state = .done - continuation.finish() + case .streamListening(let continuation): + state = .done + return .finishStream(continuation, nil) - case .closure(let postgresListenContext, _): - self.state = .done - postgresListenContext.cancel() + case .closure(let postgresListenContext, _): + state = .done + return .cancelListenContext(postgresListenContext) - case .done: - break // ignore + case .done: + return .none // ignore + } + } + + self.run(action) + } + + private func run(_ action: EndAction) { + switch action { + case .none: + break + case .failContinuation(let checkedContinuation, let error): + checkedContinuation.resume(throwing: error) + case .finishStream(let continuation, let error): + continuation.finish(throwing: error) + case .cancelListenContext(let postgresListenContext): + postgresListenContext.cancel() } } }