Avoid deadlock when accessing @Shared in a @Dependency

...by not holding locks while accessing `Dependency` logic with locks of its own.

Also simplify `PersistentReferences` implementation by maintaining weak references to
`_PersistentReference<Key>` objects directly, without custom reference counting.

Fixes #149.

Author: pyrtsa

Sources/Sharing/Internal/PersistentReferences.swift

@@ -6,9 +6,8 @@ final class PersistentReferences: @unchecked Sendable, DependencyKey {
   static var liveValue: PersistentReferences { PersistentReferences() }
   static var testValue: PersistentReferences { PersistentReferences() }
 
-  struct Pair<Key: SharedReaderKey> {
-    var cachedValue: Key.Value
-    var reference: _PersistentReference<Key>?
+  struct Weak<Key: SharedReaderKey> {
+    weak var reference: _PersistentReference<Key>?
   }
 
   private var storage: [AnyHashable: Any] = [:]
@@ -18,37 +17,35 @@ final class PersistentReferences: @unchecked Sendable, DependencyKey {
     forKey key: Key,
     default value: @autoclosure () throws -> Key.Value,
     skipInitialLoad: Bool
-  ) rethrows -> _ManagedReference<Key> {
-    try lock.withLock {
-      guard var pair = storage[key.id] as? Pair<Key> else {
-        let value = try value()
-        let persistentReference = _PersistentReference(
+  ) rethrows -> _PersistentReference<Key> {
+    guard let reference = lock.withLock({ (storage[key.id] as? Weak<Key>)?.reference }) else {
+      let value = try value()
+      return withExtendedLifetime(
+        _PersistentReference(
           key: key,
           value: value,
           skipInitialLoad: skipInitialLoad
         )
-        storage[key.id] = Pair(cachedValue: value, reference: persistentReference)
-        return _ManagedReference(persistentReference)
+      ) { reference in
+        lock.withLock {
+          if let reference = (storage[key.id] as? Weak<Key>)?.reference {
+            return reference
+          } else {
+            storage[key.id] = Weak(reference: reference)
+            reference.onDeinit = { [self] in
+              removeReference(forKey: key)
+            }
+            return reference
+          }
+        }
       }
-      guard let persistentReference = pair.reference else {
-        let persistentReference = _PersistentReference(
-          key: key,
-          value: skipInitialLoad ? (try? value()) ?? pair.cachedValue : pair.cachedValue,
-          skipInitialLoad: skipInitialLoad
-        )
-        pair.reference = persistentReference
-        storage[key.id] = pair
-        return _ManagedReference(persistentReference)
-      }
-      return _ManagedReference(persistentReference)
     }
+    return reference
   }
 
   func removeReference<Key: SharedReaderKey>(forKey key: Key) {
     lock.withLock {
-      guard var pair = storage[key.id] as? Pair<Key> else { return }
-      pair.reference = nil
-      storage[key.id] = pair
+      _ = storage.removeValue(forKey: key.id)
     }
   }
 }

Sources/Sharing/Internal/Reference.swift

@@ -196,6 +196,7 @@ final class _PersistentReference<Key: SharedReaderKey>:
   private var _saveError: (any Error)?
   private var _referenceCount = 0
   private var subscription: SharedSubscription?
+  internal var onDeinit: (() -> Void)?
 
   init(key: Key, value initialValue: Key.Value, skipInitialLoad: Bool) {
     self.key = key
@@ -231,6 +232,10 @@ final class _PersistentReference<Key: SharedReaderKey>:
     )
   }
 
+  deinit {
+    onDeinit?()
+  }
+
   var id: ObjectIdentifier { ObjectIdentifier(self) }
 
   var isLoading: Bool {
@@ -313,16 +318,6 @@ final class _PersistentReference<Key: SharedReaderKey>:
     lock.withLock { _referenceCount += 1 }
   }
 
-  func release() {
-    let shouldRelease = lock.withLock {
-      _referenceCount -= 1
-      return _referenceCount <= 0
-    }
-    guard shouldRelease else { return }
-    @Dependency(PersistentReferences.self) var persistentReferences
-    persistentReferences.removeReference(forKey: key)
-  }
-
   func access<Member>(
     keyPath: KeyPath<_PersistentReference, Member>,
     fileID: StaticString = #fileID,
@@ -454,85 +449,6 @@ extension _PersistentReference: MutableReference, Equatable where Key: SharedKey
   }
 }
 
-final class _ManagedReference<Key: SharedReaderKey>: Reference, Observable {
-  private let base: _PersistentReference<Key>
-
-  init(_ base: _PersistentReference<Key>) {
-    base.retain()
-    self.base = base
-  }
-
-  deinit {
-    base.release()
-  }
-
-  var id: ObjectIdentifier {
-    base.id
-  }
-
-  var isLoading: Bool {
-    base.isLoading
-  }
-
-  var loadError: (any Error)? {
-    base.loadError
-  }
-
-  var wrappedValue: Key.Value {
-    base.wrappedValue
-  }
-
-  func load() async throws {
-    try await base.load()
-  }
-
-  func touch() {
-    base.touch()
-  }
-
-  #if canImport(Combine)
-    var publisher: any Publisher<Key.Value, Never> {
-      base.publisher
-    }
-  #endif
-
-  var description: String {
-    base.description
-  }
-}
-
-extension _ManagedReference: MutableReference, Equatable where Key: SharedKey {
-  var saveError: (any Error)? {
-    base.saveError
-  }
-
-  var snapshot: Key.Value? {
-    base.snapshot
-  }
-
-  func takeSnapshot(
-    _ value: Key.Value,
-    fileID: StaticString,
-    filePath: StaticString,
-    line: UInt,
-    column: UInt
-  ) {
-    base.takeSnapshot(value, fileID: fileID, filePath: filePath, line: line, column: column)
-  }
-
-  func withLock<R>(_ body: (inout Key.Value) throws -> R) rethrows -> R {
-    try base.withLock(body)
-  }
-
-  func save() async throws {
-    try await base.save()
-  }
-
-  static func == (lhs: _ManagedReference, rhs: _ManagedReference) -> Bool {
-    lhs.base == rhs.base
-  }
-}
-
 final class _AppendKeyPathReference<
   Base: Reference, Value, Path: KeyPath<Base.Value, Value> & Sendable
 >: Reference, Observable {

Sources/Sharing/SharedContinuations.swift

@@ -229,7 +229,7 @@ private final class ContinuationBox<Value>: Sendable {
         invoked. This will cause tasks waiting on it to resume immediately.
         """
       )
-      callback.withLock { $0?(.success(nil)) }
+      callback.withLock(\.self)?(.success(nil))
     }
   }
 
@@ -246,9 +246,10 @@ private final class ContinuationBox<Value>: Sendable {
       )
       return
     }
-    callback.withLock { callback in
+    let callback = callback.withLock { callback in
       defer { callback = nil }
-      callback?(result)
+      return callback
     }
+    callback?(result)
   }
 }