Fix observer deregistration on deinit

Author: kylehughes

Sources/PersistentKeyValueKit/Property Wrapper/PersistentKeyUIObservableObject.swift

@@ -28,6 +28,9 @@ public final class PersistentKeyUIObservableObject<Key>: ObservableObject where
 
     /// The store that the key is being observed in.
     private var _store: (any PersistentKeyValueStore)?
+
+    /// The registration lifetime for the current store.
+    private var observation: StoreObservation<Key>?
 
     /// The object used for key-value and `NotificationCenter` observation.
     private var observer: Observer<Key>!
@@ -45,8 +48,6 @@ public final class PersistentKeyUIObservableObject<Key>: ObservableObject where
 
         observer = Observer(keyID: key.id) { [weak self] in
             self?.objectWillChange.send()
-        } deregister: { @MainActor [weak self] in
-            self?.deregisterObserver(on: store)
         }
 
         registerObserver(on: store)
@@ -60,14 +61,10 @@ public final class PersistentKeyUIObservableObject<Key>: ObservableObject where
             _store
         }
         set {
-            deregisterObserver(on: _store)
+            observation = nil
 
             _store = newValue
 
-            observer.deregister = { @MainActor [weak self] in
-                self?.deregisterObserver(on: newValue)
-            }
-            
             registerObserver(on: newValue)
         }
     }
@@ -94,18 +91,50 @@ public final class PersistentKeyUIObservableObject<Key>: ObservableObject where
         assert(store != nil, "Store should always be set on initialization or immediately after.")
     }
 
-    private func deregisterObserver(on store: (any PersistentKeyValueStore)?) {
-        store?.deregister(observer, for: key, context: nil)
-    }
-    
     private func registerObserver(on store: (any PersistentKeyValueStore)?) {
-        store?.register(
+        guard let store else {
+            return
+        }
+
+        observation = StoreObservation(
+            store: store,
+            key: key,
+            observer: observer
+        )
+    }
+}
+
+// MARK: - StoreObservation Definition
+
+private final class StoreObservation<Key> where Key: PersistentKeyProtocol {
+    private let key: Key
+    private let observer: Observer<Key>
+    private let store: any PersistentKeyValueStore
+
+    // MARK: Internal Initialization
+
+    internal init(
+        store: any PersistentKeyValueStore,
+        key: Key,
+        observer: Observer<Key>
+    ) {
+        self.key = key
+        self.observer = observer
+        self.store = store
+
+        store.register(
             observer: observer,
             for: key,
             with: nil,
             and: #selector(Observer<Key>.didReceive(_:))
         )
     }
+
+    // MARK: Deinitialization
+
+    deinit {
+        store.deregister(observer, for: key, context: nil)
+    }
 }
 
 // MARK: - Observer Definition
@@ -114,24 +143,14 @@ private final class Observer<Key>: NSObject, ObservableObject where Key: Persist
     internal let keyID: String
     internal let objectWillChange: @MainActor () -> Void
 
-    internal var deregister: @MainActor () -> Void
-    
     // MARK: Internal Initialization
 
     internal init(
         keyID: String,
-        objectWillChange: @escaping @MainActor () -> Void,
-        deregister: @escaping @MainActor () -> Void
+        objectWillChange: @escaping @MainActor () -> Void
     ) {
         self.keyID = keyID
         self.objectWillChange = objectWillChange
-        self.deregister = deregister
-    }
-    
-    // MARK: Deinitialization
-    
-    deinit {
-        performOnMainIfNecessary(deregister)
     }
 
     // MARK: NSObject Implementation

Tests/PersistentKeyValueKitTests/Tests/Persistent Key/PersistentKeyObserverTests.swift

@@ -136,4 +136,93 @@ extension PersistentKeyObserverTests {
 
         XCTAssertNil(weakObserver, "Observer should be deallocated")
     }
+
+    @MainActor
+    func test_deinitialization_deregistersFromStore() {
+        let store = ObservablePersistentKeyValueStore()
+        let key = PersistentKey("key:\(#function)", defaultValue: "defaultValue")
+        var observer: PersistentKeyUIObservableObject? = PersistentKeyUIObservableObject(store: store, key: key)
+
+        XCTAssertEqual(store.registrationCount, 1)
+
+        weak let weakObserver = observer
+        observer = nil
+
+        XCTAssertNil(weakObserver)
+        XCTAssertEqual(store.deregistrationCount, 1)
+    }
+
+    @MainActor
+    func test_deinitialization_deregistersCurrentStoreAfterStoreChange() {
+        let oldStore = ObservablePersistentKeyValueStore()
+        let newStore = ObservablePersistentKeyValueStore()
+        let key = PersistentKey("key:\(#function)", defaultValue: "defaultValue")
+        var observer: PersistentKeyUIObservableObject? = PersistentKeyUIObservableObject(store: oldStore, key: key)
+
+        XCTAssertEqual(oldStore.registrationCount, 1)
+        XCTAssertEqual(newStore.registrationCount, 0)
+
+        observer?.store = newStore
+
+        XCTAssertEqual(oldStore.deregistrationCount, 1)
+        XCTAssertEqual(newStore.registrationCount, 1)
+
+        observer = nil
+
+        XCTAssertEqual(oldStore.deregistrationCount, 1)
+        XCTAssertEqual(newStore.deregistrationCount, 1)
+    }
+
+    @MainActor
+    func test_deinitialization_deregistersStoreSetAfterNilInitialization() {
+        let store = ObservablePersistentKeyValueStore()
+        let key = PersistentKey("key:\(#function)", defaultValue: "defaultValue")
+        var observer: PersistentKeyUIObservableObject? = PersistentKeyUIObservableObject(store: nil, key: key)
+
+        XCTAssertEqual(store.registrationCount, 0)
+
+        observer?.store = store
+
+        XCTAssertEqual(store.registrationCount, 1)
+        XCTAssertEqual(store.deregistrationCount, 0)
+
+        observer = nil
+
+        XCTAssertEqual(store.deregistrationCount, 1)
+    }
+
+    @MainActor
+    func test_deinitialization_doesNotDeregisterOldStoreAgainAfterStoreChangesToNil() {
+        let store = ObservablePersistentKeyValueStore()
+        let key = PersistentKey("key:\(#function)", defaultValue: "defaultValue")
+        var observer: PersistentKeyUIObservableObject? = PersistentKeyUIObservableObject(store: store, key: key)
+
+        XCTAssertEqual(store.registrationCount, 1)
+
+        observer?.store = nil
+
+        XCTAssertEqual(store.deregistrationCount, 1)
+
+        observer = nil
+
+        XCTAssertEqual(store.deregistrationCount, 1)
+    }
+
+    @MainActor
+    func test_deinitialization_deregistersCurrentRegistrationAfterStoreChangesToSameStore() {
+        let store = ObservablePersistentKeyValueStore()
+        let key = PersistentKey("key:\(#function)", defaultValue: "defaultValue")
+        var observer: PersistentKeyUIObservableObject? = PersistentKeyUIObservableObject(store: store, key: key)
+
+        XCTAssertEqual(store.registrationCount, 1)
+
+        observer?.store = store
+
+        XCTAssertEqual(store.registrationCount, 2)
+        XCTAssertEqual(store.deregistrationCount, 1)
+
+        observer = nil
+
+        XCTAssertEqual(store.deregistrationCount, 2)
+    }
 }