UserDefaults KVO observation with AsyncSequence/AsyncStream

UserDefaultsAsyncSequence.swift

// UserDefaults KVO observation with AsyncSequence/AsyncStream
// Ole Begemann, 2023-04
// https://gist.github.com/ole/fc5c1f4c763d28d9ba70940512e81916

import Foundation

extension UserDefaults {
    func observeKey<Value>(_ key: String, valueType _: Value.Type) -> AsyncStream<Value?> {
        var continuation: AsyncStream<Value?>.Continuation? = nil
        let stream = AsyncStream(Value?.self) {
            continuation = $0
        }
        let observer = KVOObserver(send: { continuation!.yield($0) })
        continuation!.onTermination = { [weak self] termination in
            print("UserDefaults.observeKey('\(key)') sequence terminated. Reason: \(termination)")
            // Warning: Capture of 'self' with non-sendable type 'UserDefaults?' in a `@Sendable` closure.
            // UserDefaults is documented to be thread-safe, so this should be ok.
            guard let self else { return }
            // Referencing observer here retains it.
            self.removeObserver(observer, forKeyPath: key)
            // Break retain cycle (is there one?)
            observer.send = nil
        }
        self.addObserver(observer, forKeyPath: key, options: [.initial, .new], context: nil)
        return stream
    }
}

private final class KVOObserver<Value>: NSObject {
    var send: Optional<(Value?) -> Void>

    init(send: @escaping (Value?) -> Void) {
        self.send = send
    }

    deinit {
        print("KVOObserver deinit")
    }

    override func observeValue(forKeyPath keyPath: String?, of object: Any?, change: [NSKeyValueChangeKey: Any]?, context: UnsafeMutableRawPointer?) {
        let newValue = change![.newKey]!
        switch newValue {
        case let typed as Value:
            send?(typed)
        case nil as Value?:
            send?(nil)
        default:
            assertionFailure("UserDefaults value at keyPath '\(keyPath!)' has unexpected type \(type(of: newValue)), expected \(Value.self)")
        }
    }
}

// MARK: - Usage

let observationTask = Task<Void, Never> {
    for await value in UserDefaults.standard.observeKey("user", valueType: [String: Any].self) {
        print("KVO: \(value?.description ?? "nil")")
    }
}

// Give observation task an opportunity to run
try? await Task.sleep(for: .seconds(0.1))

// These trigger the for loop in observationTask
UserDefaults.standard.set(["name": "Alice", "age": 23] as [String: Any], forKey: "user")
UserDefaults.standard.set(["name": "Bob"] as [String: Any], forKey: "user")
UserDefaults.standard.set(["name": "Charlie", "age": 42] as [String: Any], forKey: "user")
UserDefaults.standard.removeObject(forKey: "user")

// Cancel observation
try? await Task.sleep(for: .seconds(1))
print("Canceling UserDefaults observation")
observationTask.cancel()

// These don't print anything because observationTask has been canceled.
UserDefaults.standard.set(["name": "Danny"] as [String: Any], forKey: "user")
UserDefaults.standard.removeObject(forKey: "user")

try? await Task.sleep(for: .seconds(1))

Isn't it easier to just use AsyncPublisher?

extension UserDefaults {
    typealias AsyncValues<T> = AsyncPublisher<AnyPublisher<T, Never>>
    func observeKey<T>(at path: KeyPath<UserDefaults, T>) -> AsyncValues<T> {
        return self.publisher(for: path, options: [.initial, .new])
            .eraseToAnyPublisher()
            .values
    }
}

KVO-compliant properties already have the .publisher(for:) method to create a KVOPublisher, and then any publisher can be converted into an AsyncPublisher on their .values property. And AsyncPublisher conforms to AsyncSequence.

I think going through Combine is unnecessary overhead. However, perhaps the code could be simplified by replacing KVOObserver with Swift's built-in NSKeyValueObservation as demonstrated in Using Key-Value Observing in Swift. It could also be improved by using the continuation inside the AsyncStream's init closure instead of bouncing it outside which will likely lead to a memory leak, like the onTermination [weak self] leak which is explained in this similar example.