disposable notification

NotificationCenter+Concurrency.swift

public extension NotificationCenter {
    private class FirstReceiver: NSObject {

        let notificationName: Notification.Name

        init(_ name: Notification.Name) {
            notificationName = name
            super.init()
        }

        var continuation: CheckedContinuation<[AnyHashable: Any]?, Error>?

        func setReceive(_ continuation: CheckedContinuation<[AnyHashable: Any]?, Error>) {
            self.continuation = continuation
            NotificationCenter.default.addObserver(self, selector: #selector(receiveUserInfo), name: notificationName, object: nil)
        }

        @objc
        func receiveUserInfo(_ notification: Notification) {
            let continuation = self.continuation
            self.continuation = nil
            NotificationCenter.default.removeObserver(self, name: notificationName, object: nil)
            continuation?.resume(returning: notification.userInfo)
        }

        func cancel() {
            continuation?.resume(throwing: CancellationError())
            NotificationCenter.default.removeObserver(self, name: notificationName, object: nil)
        }
    }

    @discardableResult
    func receiveFirst(forName name: Notification.Name, timeout: TimeInterval = 0) async throws -> [AnyHashable: Any]? {
        if #available(iOS 15, *) {
            return try await withThrowingTaskGroup(of: [AnyHashable: Any]?.self, body: { group in
                group.addTask {
                    let notifications = NotificationCenter.default.notifications(named: name)
                    for await notification in notifications {
                        return notification.userInfo
                    }
                    try Task.checkCancellation()
                    return nil
                }
                if timeout > 0 {
                    group.addTask {
                        try await Task.sleep(nanoseconds: UInt64(timeout * Double(1_000_000_000)))
                        throw CancellationError()
                    }
                }
                do {
                    for try await value in group {
                        return value
                    }
                    return nil
                } catch {
                    throw error
                }
            })
        } else {
            let receiver = FirstReceiver(name)
            return try await withThrowingTaskGroup(of: [AnyHashable: Any]?.self, body: { group in
                group.addTask {
                    try await withTaskCancellationHandler(operation: {
                        try await withCheckedThrowingContinuation({ continuation in
                            receiver.setReceive(continuation)
                        })
                    }, onCancel: {
                        receiver.cancel()
                    })
                }
                if timeout > 0 {
                    group.addTask { [weak receiver] in
                        try await Task.sleep(nanoseconds: UInt64(timeout * Double(1_000_000_000)))
                        receiver?.cancel()
                        throw CancellationError()
                    }
                }

                for try await value in group {
                    return value
                }
                return nil
            })
        }
    }
}