fix: awaitReconciliation and tests

spotify · Dec 9, 2024 · 0a9657e · 0a9657e
1 parent a7f666d
commit 0a9657e
Show file tree

Hide file tree

Showing 2 changed files with 78 additions and 12 deletions.
diff --git a/Sources/Confidence/Confidence.swift b/Sources/Confidence/Confidence.swift
@@ -23,7 +23,13 @@ public class Confidence: ConfidenceEventSender {
     // Synchronization and task management resources
     private var cancellables = Set<AnyCancellable>()
     private let cacheQueue = DispatchQueue(label: "com.confidence.queue.cache")
-    private var currentFetchTask: Task<(), Never>?
+    private var currentFetchTask: Task<(), Never>? {
+        didSet {
+            if let oldTask = oldValue {
+                oldTask.cancel()
+            }
+        }
+    }
 
     // Internal for testing
     internal let remoteFlagResolver: ConfidenceResolveClient
@@ -155,7 +161,6 @@ public class Confidence: ConfidenceEventSender {
     }
 
     public func putContextAndWait(key: String, value: ConfidenceValue) async {
-        self.currentFetchTask?.cancel()
         self.currentFetchTask = Task {
             let newContext = contextManager.updateContext(withValues: [key: value], removedKeys: [])
             do {
@@ -169,7 +174,6 @@ public class Confidence: ConfidenceEventSender {
     }
 
     public func putContextAndWait(context: ConfidenceStruct, removedKeys: [String] = []) async {
-        self.currentFetchTask?.cancel()
         self.currentFetchTask = Task {
             let newContext = contextManager.updateContext(withValues: context, removedKeys: removedKeys)
             do {
@@ -183,7 +187,6 @@ public class Confidence: ConfidenceEventSender {
     }
 
     public func putContextAndWait(context: ConfidenceStruct) async {
-        self.currentFetchTask?.cancel()
         self.currentFetchTask = Task {
             let newContext = contextManager.updateContext(withValues: context, removedKeys: [])
             do {
@@ -201,7 +204,6 @@ public class Confidence: ConfidenceEventSender {
     }
 
     public func removeContextAndWait(key: String) async {
-        self.currentFetchTask?.cancel()
         self.currentFetchTask = Task {
             let newContext = contextManager.updateContext(withValues: [:], removedKeys: [key])
             do {
@@ -229,35 +231,30 @@ public class Confidence: ConfidenceEventSender {
     }
 
     public func putContext(key: String, value: ConfidenceValue) {
-        self.currentFetchTask?.cancel()
         self.currentFetchTask = Task {
             await putContextAndWait(key: key, value: value)
         }
     }
 
     public func putContext(context: ConfidenceStruct) {
-        self.currentFetchTask?.cancel()
         self.currentFetchTask = Task {
             await putContextAndWait(context: context)
         }
     }
 
     public func putContext(context: ConfidenceStruct, removeKeys removedKeys: [String] = []) {
-        self.currentFetchTask?.cancel()
         self.currentFetchTask = Task {
             await putContextAndWait(context: context, removedKeys: removedKeys)
         }
     }
 
     public func removeContext(key: String) {
-        self.currentFetchTask?.cancel()
         self.currentFetchTask = Task {
             await removeContextAndWait(key: key)
         }
     }
 
     public func putContext(context: ConfidenceStruct, removedKeys: [String]) {
-        self.currentFetchTask?.cancel()
         self.currentFetchTask = Task {
             let newContext = contextManager.updateContext(withValues: context, removedKeys: removedKeys)
             do {
@@ -277,8 +274,19 @@ public class Confidence: ConfidenceEventSender {
     Ensures all the already-started context changes prior to this function have been reconciliated
     */
     public func awaitReconciliation() async {
-        if let task = self.currentFetchTask {
+        while let task = self.currentFetchTask {
+            // If current task is cancelled return
+            if task.isCancelled {
+                return
+            }
+            // else wait for result
             await task.value
+            if task.isCancelled {
+                continue
+            }
+            if self.currentFetchTask == task {
+                return
+            }
         }
     }
 

diff --git a/Tests/ConfidenceTests/ConfidenceTest.swift b/Tests/ConfidenceTests/ConfidenceTest.swift
@@ -107,8 +107,8 @@ class ConfidenceTest: XCTestCase {
         XCTAssertEqual(2, client.resolveContexts.count)
         XCTAssertEqual(confidence.getContext(), client.resolveContexts[1])
     }
+
     // swiftlint:enable function_body_length
-
     func testRefresh() async throws {
         class FakeClient: ConfidenceResolveClient {
             var resolveStats: Int = 0
@@ -490,6 +490,64 @@ class ConfidenceTest: XCTestCase {
         XCTAssertEqual(flagApplier.applyCallCount, 1)
     }
 
+    func testAwaitReconciliationFailingTask() async throws {
+        class FakeClient: XCTestCase, ConfidenceResolveClient {
+            var resolveStats: Int = 0
+            var resolvedValues: [ResolvedValue] = []
+
+            func resolve(ctx: ConfidenceStruct) async throws -> ResolvesResult {
+                self.resolveStats += 1
+                if resolveStats == 1 {
+                    try await Task.sleep(nanoseconds: 2_000_000_000) // Sleep for 3 seconds
+                    return .init(resolvedValues: [], resolveToken: "token")
+                } else {
+                    try await Task.sleep(nanoseconds: 3_000_000_000) // Sleep for 1 second
+                    if ctx["hello"] == .init(string: "world") {
+                        return .init(resolvedValues: resolvedValues, resolveToken: "token")
+                    } else {
+                        return .init(resolvedValues: [], resolveToken: "token")
+                    }
+                }
+            }
+        }
+
+        let client = FakeClient()
+        client.resolvedValues = [
+            ResolvedValue(
+                variant: "control",
+                value: .init(structure: ["size": .init(integer: 3)]),
+                flag: "flag",
+                resolveReason: .match)
+        ]
+
+        let confidence = Confidence.Builder(clientSecret: "test")
+            .withContext(initialContext: ["targeting_key": .init(string: "user2")])
+            .withFlagResolverClient(flagResolver: client)
+            .withFlagApplier(flagApplier: flagApplier)
+            .withStorage(storage: storage)
+            .build()
+
+        confidence.putContext(context: ["hello": .init(string: "not-world")])
+        Task {
+            confidence.putContext(context: ["hello": .init(string: "world")])
+        }
+        try await Task.sleep(nanoseconds: 1_000_000_000) // Sleep for 1 second
+        await confidence.awaitReconciliation()
+        let evaluation = confidence.getEvaluation(
+            key: "flag.size",
+            defaultValue: 0
+        )
+
+        XCTAssertEqual(client.resolveStats, 2)
+        XCTAssertEqual(evaluation.value, 3)
+        XCTAssertNil(evaluation.errorCode)
+        XCTAssertNil(evaluation.errorMessage)
+        XCTAssertEqual(evaluation.reason, .match)
+        XCTAssertEqual(evaluation.variant, "control")
+        await fulfillment(of: [flagApplier.applyExpectation], timeout: 1)
+        XCTAssertEqual(flagApplier.applyCallCount, 1)
+    }
+
     func testResolveBooleanFlag() async throws {
         class FakeClient: ConfidenceResolveClient {
             var resolveStats: Int = 0