UserDefaults KVO observation with AsyncSequence/AsyncStream

UserDefaultsAsyncSequence.swift

// UserDefaults KVO observation with AsyncSequence/AsyncStream
// Ole Begemann, 2023-04
// Updated for Swift 6, 2024-11
// https://gist.github.com/ole/fc5c1f4c763d28d9ba70940512e81916

import Foundation

// This is ugly, but UserDefaults is documented to be thread-safe, so this
// should be OK.
extension UserDefaults: @retroactive @unchecked Sendable {}

extension UserDefaults {
    func observeKey<Value: Sendable>(_ key: String, valueType _: Value.Type) -> AsyncStream<Value?> {
        let (stream, continuation) = AsyncStream.makeStream(of: Value?.self)
        let observer = KVOObserver { newValue in
            continuation.yield(newValue)
        }
        continuation.onTermination = { [weak self] termination in
            print("UserDefaults.observeKey('\(key)') sequence terminated. Reason: \(termination)")
            guard let self else { return }
            // Referencing observer here retains it.
            self.removeObserver(observer, forKeyPath: key)
        }
        self.addObserver(observer, forKeyPath: key, options: [.initial, .new], context: nil)
        return stream
    }
}

private final class KVOObserver<Value: Sendable>: NSObject, Sendable {
    let send: @Sendable (Value?) -> Void

    init(send: @escaping @Sendable (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: Int].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": 1, "age": 23] as [String: Int], forKey: "user")
UserDefaults.standard.set(["name": 2] as [String: Int], forKey: "user")
UserDefaults.standard.set(["name": 3, "age": 42] as [String: Int], 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": 4] as [String: Int], 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.

What's the best way to make this work in Swift 6? So far, what I found is conforming KVOObserver to @unchecked Sendable, and constraining Value to be Sendable, too.

I guess having a Value: Sendable is logical and necessary, but what is an actually good way to get rid of the Sendable warning around KVOObserver?

extension UserDefaults {

    func observeKey<Value: Sendable>(_ 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, @unchecked Sendable {
    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)")
        }
    }
}

@Gernot I updated the code so that it compiles in Swift 6 language mode. Here's the diff: https://gist.github.com/ole/fc5c1f4c763d28d9ba70940512e81916/revisions#diff-236817b08adf0885e4b9e978670cb741cc51be6f1cbe4f9e18b32daaa260b5aa

Notes:

• The example dictionary makes a bit less sense because I changed it from [String: Any] to [String: Int] in order to satisfy the Value: Sendable constraint. But as you said, I think it makes sense that the value must be Sendable.
• I made class KVOObserver immutable, which allows me to make it Sendable.
• extension UserDefaults: @unchecked Sendable {} is still required. I don't know how to get rid of it because the UserDefaults object must cross isolation domains if we want to use it in the continuation's termination handler. It should be OK if we trust the docs that UserDefaults is thread-safe, but who knows.

I'd love to see a better solution.

Ah thanks! I came almost as far, but didn't dare to make the KVOObserver.send immutable, because of these lines that are now removed:

// Break retain cycle (is there one?)
observer.send = nil

I'm always afraid if someone writes that there might be a retain cycle ;-)

I'm always afraid if someone writes that there might be a retain cycle ;-)

Yeah, I don't know. I guess the observer retains the continuation and the continuation's onTermination handler retains the observer, so we do have a cycle. But the docs for onTermination say:

After reaching a terminal state as a result of cancellation, the AsyncStream sets the callback to nil.

So I guess that resolves the cycle. At least in my little usage example, the observer's deinit is being called. I haven't tested it any further.

I don't think this pattern will work in Swift 6:

        var continuation: AsyncStream<Value?>.Continuation? = nil
        let stream = AsyncStream(Value?.self) {
            continuation = $0
        }

You'll need to switch to the original version of the AsyncStream API that creates a lifetime for your object to exist in, e.g.

  return AsyncStream(Value?.self) { continuation in
        // init the observer here and it will be retained while the stream is working
            let keyValueObservation = object.observe(...) {
                continuation.(...
            }
       }

You'll also need to create your stream inside a static func, usually inside a new struct type. Take a look at CLLocationUpdate.liveUpdates for an example.

If you use Swift's observe API then the observation will be cancelled automatically when the NSKeyValueObservation handle is deinit, which will happen automatically when the stream is cancelled and the handle goes out of scope. So if done right, I believe you won't even need an onTermination handler.