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))

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.