From 69247baf7be2fd6f5820192caef0082d01849cd0 Mon Sep 17 00:00:00 2001
From: Pyry Jahkola <pyry.jahkola@sanoma.com>
Date: Thu, 14 Nov 2024 21:14:20 +0200
Subject: [PATCH] Fix crash when cancelling subscription for custom
 PersistenceKey (#3494)

* Add crashing unit test for AppStorageKey subscription cancellation

* Fix double cancellation of Shared.Subscription
---
 .../SharedState/PersistenceKey.swift             |  5 +++--
 .../SharedTests.swift                            | 16 ++++++++++++++++
 2 files changed, 19 insertions(+), 2 deletions(-)

diff --git a/Sources/ComposableArchitecture/SharedState/PersistenceKey.swift b/Sources/ComposableArchitecture/SharedState/PersistenceKey.swift
index 710b6b4067ca..13d008a580bd 100644
--- a/Sources/ComposableArchitecture/SharedState/PersistenceKey.swift
+++ b/Sources/ComposableArchitecture/SharedState/PersistenceKey.swift
@@ -73,7 +73,7 @@ extension Shared {
   /// This object is returned from ``PersistenceReaderKey/subscribe(initialValue:didSet:)``, which
   /// will feed updates from an external system for its lifetime, or till ``cancel()`` is called.
   public class Subscription {
-    let onCancel: () -> Void
+    var onCancel: (() -> Void)?
 
     /// Initializes the subscription with the given cancel closure.
     ///
@@ -88,7 +88,8 @@ extension Shared {
 
     /// Cancels the subscription.
     public func cancel() {
-      self.onCancel()
+      self.onCancel?()
+      self.onCancel = nil
     }
   }
 }
diff --git a/Tests/ComposableArchitectureTests/SharedTests.swift b/Tests/ComposableArchitectureTests/SharedTests.swift
index fc78cef09979..7aa8ba5c48b2 100644
--- a/Tests/ComposableArchitectureTests/SharedTests.swift
+++ b/Tests/ComposableArchitectureTests/SharedTests.swift
@@ -1084,6 +1084,22 @@ final class SharedTests: XCTestCase {
       }
     }
   }
+
+  func testPersistenceKeySubscription() async throws {
+    let persistenceKey: AppStorageKey<Int> = .appStorage("shared")
+    let changes = LockIsolated<[Int?]>([])
+    var subscription: Optional = persistenceKey.subscribe(initialValue: nil) { value in
+      changes.withValue { $0.append(value) }
+    }
+    @Dependency(\.defaultAppStorage) var userDefaults
+    userDefaults.set(1, forKey: "shared")
+    userDefaults.set(42, forKey: "shared")
+    subscription?.cancel()
+    userDefaults.set(123, forKey: "shared")
+    subscription = nil
+    XCTAssertEqual([1, 42], changes.value)
+    XCTAssertEqual(123, persistenceKey.load(initialValue: nil))
+  }
 }
 
 @globalActor actor GA: GlobalActor {