fix(DataStore): model reconciliation queue publisher (#756)

Author: wooj2

AmplifyPlugins/DataStore/AWSDataStoreCategoryPlugin/Sync/SubscriptionSync/ReconcileAndLocalSave/AWSModelReconciliationQueue.swift

@@ -64,7 +64,7 @@ final class AWSModelReconciliationQueue: ModelReconciliationQueue {
     private let modelName: String
 
     private var incomingEventsSink: AnyCancellable?
-    private var reconcileAndLocalSaveOperationSink: AnyCancellable?
+    private var reconcileAndLocalSaveOperationSinks: Set<AnyCancellable?>
 
     private let modelReconciliationQueueSubject: PassthroughSubject<ModelReconciliationQueueEvent, DataStoreError>
     var publisher: AnyPublisher<ModelReconciliationQueueEvent, DataStoreError> {
@@ -98,7 +98,7 @@ final class AWSModelReconciliationQueue: ModelReconciliationQueue {
         let resolvedIncomingSubscriptionEvents = incomingSubscriptionEvents ??
             AWSIncomingSubscriptionEventPublisher(modelType: modelType, api: api, auth: auth)
         self.incomingSubscriptionEvents = resolvedIncomingSubscriptionEvents
-
+        self.reconcileAndLocalSaveOperationSinks = Set<AnyCancellable?>()
         self.incomingEventsSink = resolvedIncomingSubscriptionEvents
             .publisher
             .sink(receiveCompletion: { [weak self] completion in
@@ -133,11 +133,17 @@ final class AWSModelReconciliationQueue: ModelReconciliationQueue {
     func enqueue(_ remoteModel: MutationSync<AnyModel>) {
         let reconcileOp = ReconcileAndLocalSaveOperation(remoteModel: remoteModel,
                                                          storageAdapter: storageAdapter)
-        reconcileAndLocalSaveOperationSink = reconcileOp.publisher.sink(receiveCompletion: { error in
-            self.modelReconciliationQueueSubject.send(completion: error)
+        var reconcileAndLocalSaveOperationSink: AnyCancellable?
+
+        reconcileAndLocalSaveOperationSink = reconcileOp.publisher.sink(receiveCompletion: { completion in
+            self.reconcileAndLocalSaveOperationSinks.remove(reconcileAndLocalSaveOperationSink)
+            if case .failure = completion {
+                self.modelReconciliationQueueSubject.send(completion: completion)
+            }
         }, receiveValue: { mutationEvent in
             self.modelReconciliationQueueSubject.send(.mutationEvent(mutationEvent))
         })
+        reconcileAndLocalSaveOperationSinks.insert(reconcileAndLocalSaveOperationSink)
         reconcileAndSaveQueue.addOperation(reconcileOp)
     }
 

AmplifyPlugins/DataStore/AWSDataStoreCategoryPlugin/Sync/SubscriptionSync/ReconcileAndLocalSave/ReconcileAndLocalSaveOperation.swift

@@ -96,10 +96,12 @@ class ReconcileAndLocalSaveOperation: AsynchronousOperation {
         case .inError(let error):
             // Maybe we have to notify the Hub?
             log.error(error: error)
+            notifyFinished()
             finish()
 
         case .finished:
             // Maybe we have to notify the Hub?
+            notifyFinished()
             finish()
 
         }
@@ -301,6 +303,10 @@ class ReconcileAndLocalSaveOperation: AsynchronousOperation {
         stateMachine.notify(action: .notified)
     }
 
+    private func notifyFinished() {
+        mutationEventPublisher.send(completion: .finished)
+    }
+
     private func getPendingMutations(forModelId modelId: Model.Identifier) -> DataStoreResult<[MutationEvent]> {
         guard let storageAdapter = storageAdapter else {
             return .failure(DataStoreError.nilStorageAdapter())

AmplifyPlugins/DataStore/AWSDataStoreCategoryPluginTests/Sync/SubscriptionSync/ModelReconciliationQueueBehaviorTests.swift

@@ -89,9 +89,34 @@ class ModelReconciliationQueueBehaviorTests: ReconciliationQueueTestBase {
             subscriptionEventsSubject.send(.mutationEvent(mutationSync))
         }
 
+        let eventsSentViaPublisher1 = expectation(description: "id-1 sent via publisher")
+        let eventsSentViaPublisher2 = expectation(description: "id-2 sent via publisher")
+        let eventsSentViaPublisher3 = expectation(description: "id-3 sent via publisher")
+        let queueSink = queue.publisher.sink(receiveCompletion: { _ in
+            XCTFail("Not expecting a call to completion")
+        }, receiveValue: { event in
+            if case let .mutationEvent(mutationEvent) = event {
+                switch mutationEvent.modelId {
+                case "id-1":
+                    eventsSentViaPublisher1.fulfill()
+                case "id-2":
+                    eventsSentViaPublisher2.fulfill()
+                case "id-3":
+                    eventsSentViaPublisher3.fulfill()
+                default:
+                    XCTFail("Not expecting a call to default")
+                }
+            }
+        })
+
         queue.start()
 
-        wait(for: [event1Saved, event2Saved, event3Saved], timeout: 5.0, enforceOrder: true)
+        wait(for: [event1Saved,
+                   event2Saved,
+                   event3Saved], timeout: 5.0, enforceOrder: true)
+        wait(for: [eventsSentViaPublisher1,
+            eventsSentViaPublisher2,
+            eventsSentViaPublisher3], timeout: 2.0)
     }
 
     /// - Given: An AWSModelReconciliationQueue that has been buffering events
@@ -156,9 +181,33 @@ class ModelReconciliationQueueBehaviorTests: ReconciliationQueueTestBase {
             subscriptionEventsSubject.send(.mutationEvent(mutationSync))
         }
 
+        let eventsSentViaPublisher1 = expectation(description: "id-1 sent via publisher")
+        let eventsSentViaPublisher2 = expectation(description: "id-2 sent via publisher")
+        let eventsSentViaPublisher3 = expectation(description: "id-3 sent via publisher")
+
+        let queueSink = queue.publisher.sink(receiveCompletion: { _ in
+            XCTFail("Not expecting a call to completion")
+        }, receiveValue: { event in
+            if case let .mutationEvent(mutationEvent) = event {
+                switch mutationEvent.modelId {
+                case "id-1":
+                    eventsSentViaPublisher1.fulfill()
+                case "id-2":
+                    eventsSentViaPublisher2.fulfill()
+                case "id-3":
+                    eventsSentViaPublisher3.fulfill()
+                default:
+                    break
+                }
+            }
+        })
+
         queue.start()
 
-        wait(for: [allEventsProcessed], timeout: 5.0)
+        wait(for: [allEventsProcessed,
+                   eventsSentViaPublisher1,
+                   eventsSentViaPublisher2,
+                   eventsSentViaPublisher3], timeout: 5.0)
     }
 
     /// - Given: A started AWSModelReconciliationQueue with no pending events
@@ -202,15 +251,35 @@ class ModelReconciliationQueueBehaviorTests: ReconciliationQueueTestBase {
             subscriptionEventsSubject.send(.mutationEvent(mutationSync))
         }
 
+        let eventsSentViaPublisher1 = expectation(description: "id-1 sent via publisher")
+        let eventsSentViaPublisher2 = expectation(description: "id-2 sent via publisher")
+        var queueSink = queue.publisher.sink(receiveCompletion: { _ in
+            XCTFail("Not expecting a call to completion")
+        }, receiveValue: { event in
+            if case let .mutationEvent(mutationEvent) = event {
+                switch mutationEvent.modelId {
+                case "id-1":
+                    eventsSentViaPublisher1.fulfill()
+                case "id-2":
+                    eventsSentViaPublisher2.fulfill()
+                default:
+                    XCTFail("Not expecting a call to default")
+                }
+            }
+        })
+
         queue.start()
 
-        wait(for: [event1ShouldBeProcessed, event2ShouldBeProcessed], timeout: 1.0)
+        wait(for: [event1ShouldBeProcessed,
+                   event2ShouldBeProcessed,
+                   eventsSentViaPublisher1,
+                   eventsSentViaPublisher2], timeout: 1.0)
 
         let event1ShouldNotBeProcessed = expectation(description: "Event 1 should not be processed")
         event1ShouldNotBeProcessed.isInverted = true
         let event2ShouldNotBeProcessed = expectation(description: "Event 2 should not be processed")
         event2ShouldNotBeProcessed.isInverted = true
-        let event3ShouldBeProcessed = expectation(description: "Event 3 should not be processed")
+        let event3ShouldBeProcessed = expectation(description: "Event 3 should be processed")
         storageAdapter.responders[.saveModelCompletion] =
             SaveModelCompletionResponder<MutationSyncMetadata> { model, completion in
                 switch model.id {
@@ -226,6 +295,17 @@ class ModelReconciliationQueueBehaviorTests: ReconciliationQueueTestBase {
                 completion(.success(model))
         }
 
+        let eventsSentViaPublisher3 = expectation(description: "id-3 sent via publisher")
+        queueSink = queue.publisher.sink(receiveCompletion: { _ in
+            XCTFail("Not expecting a call to completion")
+        }, receiveValue: { event in
+            if case let .mutationEvent(mutationEvent) = event {
+                if mutationEvent.modelId == "id-3" {
+                    eventsSentViaPublisher3.fulfill()
+                }
+            }
+        })
+
         let model = try MockSynced(id: "id-3").eraseToAnyModel()
         let syncMetadata = MutationSyncMetadata(id: model.id,
                                                 deleted: false,
@@ -234,8 +314,10 @@ class ModelReconciliationQueueBehaviorTests: ReconciliationQueueTestBase {
         let mutationSync = MutationSync(model: model, syncMetadata: syncMetadata)
         subscriptionEventsSubject.send(.mutationEvent(mutationSync))
 
-        wait(for: [event1ShouldNotBeProcessed, event2ShouldNotBeProcessed, event3ShouldBeProcessed], timeout: 1.0)
-
+        wait(for: [event1ShouldNotBeProcessed,
+                   event2ShouldNotBeProcessed,
+                   event3ShouldBeProcessed,
+                   eventsSentViaPublisher3], timeout: 1.0)
     }
 
 }