Fix crash when accessing setValue concurrently (follow-up)

Sources/CacheContainer.swift

@@ -1,23 +1,23 @@
 import Foundation
 
-public final class CacheContainer {
+public final class CacheContainer: Sendable {
     static let shared = CacheContainer()
 
     public static func clearAll() {
-        shared.storage.forEach { (key, _) in
-            shared.storage.removeValue(forKey: key)
+        shared.storage.keys.forEach {
+            shared.storage.removeValue(forKey: $0)
         }
     }
 
     public static func clearAll(where shouldBeCleared: (CacheKey) -> Bool) {
-        shared.storage.forEach { (key, _) in
-            if shouldBeCleared(key) {
-                shared.storage.removeValue(forKey: key)
+        shared.storage.keys.forEach {
+            if shouldBeCleared($0) {
+                shared.storage.removeValue(forKey: $0)
             }
         }
     }
 
-    var storage: [CacheKey : Data] = [:]
+    let storage: ThreadSafeDictionary<CacheKey, Data> = .init()
 
     private let decoder = JSONDecoder()
     private let encoder = JSONEncoder()
@@ -46,3 +46,53 @@ public final class CacheContainer {
         }
     }
 }
+
+final class ThreadSafeDictionary<Key: Hashable & Sendable, Value: Sendable>: Collection, @unchecked Sendable {
+    private var dictionary: [Key : Value]
+    private let concurrentQueue = DispatchQueue(label: "Dictionary Barrier Queue", attributes: .concurrent)
+
+    var keys: Dictionary<Key, Value>.Keys {
+        concurrentQueue.sync { dictionary.keys }
+    }
+
+    var values: Dictionary<Key, Value>.Values {
+        concurrentQueue.sync { dictionary.values }
+    }
+
+    var startIndex: Dictionary<Key, Value>.Index {
+        concurrentQueue.sync { dictionary.startIndex }
+    }
+
+    var endIndex: Dictionary<Key, Value>.Index {
+        concurrentQueue.sync { dictionary.endIndex }
+    }
+
+    init(dictionary: [Key : Value] = [:]) {
+        self.dictionary = dictionary
+    }
+
+    func index(after i: Dictionary<Key, Value>.Index) -> Dictionary<Key, Value>.Index {
+        concurrentQueue.sync { dictionary.index(after: i) }
+    }
+
+    subscript(key: Key) -> Value? {
+        set(newValue) {
+            concurrentQueue.async(flags: .barrier) { [weak self] in
+                self?.dictionary[key] = newValue
+            }
+        }
+        get {
+            concurrentQueue.sync { dictionary[key] }
+        }
+    }
+
+    subscript(index: Dictionary<Key, Value>.Index) -> Dictionary<Key, Value>.Element {
+        concurrentQueue.sync { dictionary[index] }
+    }
+
+    func removeValue(forKey key: Key) {
+        concurrentQueue.async(flags: .barrier) { [weak self] in
+            self?.dictionary.removeValue(forKey: key)
+        }
+    }
+}

Tests/CacheContainerTests.swift

@@ -38,6 +38,17 @@ final class CacheContainerTests: XCTestCase {
         let data: CacheContainer.CacheData<String>? = sut.value(forKey: "key.value")
         XCTAssertEqual(data, .init(value: "value", cacheDate: cacheDate))
     }
+
+    func test_setValue_concurrently() async throws {
+        let cacheDate = try Date("2023-01-01T00:00:00Z", strategy: .iso8601)
+        (0..<100).forEach { i in
+            Task.detached { [sut] in
+                sut.setValue("value\(i)", forKey: "key.value\(i)", cacheDate: cacheDate)
+            }
+        }
+        await Task.megaYield()
+        XCTAssertEqual(sut.storage.count, 100)
+    }
 }
 
 extension CacheContainer.CacheData<String>: Equatable {
@@ -46,3 +57,10 @@ extension CacheContainer.CacheData<String>: Equatable {
     }
 }
 
+private extension Task where Success == Never, Failure == Never {
+    static func megaYield(count: Int = 10) async {
+        for _ in 0..<count {
+            await Task<Void, Never>.detached(priority: .background) { await Task.yield() }.value
+        }
+    }
+}