fix(datastore): mutation event got ignored while executing parallel saving (#2781)

* fix(datastore): mutation event got ignored while executing parallel saving

* fix(datastore): refactor with extension peek method

* test(core): add test cases for extension method peek on Optional

* fix(core): rename optional peek method to ifSome

Author: Di Wu

Amplify/Core/Support/Optional+Extension.swift

@@ -0,0 +1,22 @@
+//
+// Copyright Amazon.com Inc. or its affiliates.
+// All Rights Reserved.
+//
+// SPDX-License-Identifier: Apache-2.0
+//
+		
+
+import Foundation
+
+extension Optional {
+    ///
+    /// Performing side effect function when data is exist
+    /// - parameters:
+    ///     - then: a closure that takes wrapped data as a parameter
+    @_spi(OptionalExtension)
+    public func ifSome(_ then: (Wrapped) throws -> Void) rethrows {
+        if case .some(let wrapped) = self {
+            try then(wrapped)
+        }
+    }
+}

AmplifyPlugins/DataStore/Sources/AWSDataStorePlugin/Sync/MutationSync/AWSMutationDatabaseAdapter/AWSMutationDatabaseAdapter+MutationEventIngester.swift

@@ -5,7 +5,7 @@
 // SPDX-License-Identifier: Apache-2.0
 //
 
-import Amplify
+@_spi(OptionalExtension) import Amplify
 import Combine
 import Foundation
 
@@ -215,24 +215,31 @@ extension AWSMutationDatabaseAdapter: MutationEventIngester {
 
     /// Saves the deconflicted mutationEvent, invokes `nextEventPromise` if it exists, and the save was successful,
     /// and finally invokes the completion promise from the future of the original invocation of `submit`
-    func save(mutationEvent: MutationEvent,
-              storageAdapter: StorageEngineAdapter,
-              completionPromise: @escaping Future<MutationEvent, DataStoreError>.Promise) {
-
+    func save(
+        mutationEvent: MutationEvent,
+        storageAdapter: StorageEngineAdapter,
+        completionPromise: @escaping Future<MutationEvent, DataStoreError>.Promise
+    ) {
         log.verbose("\(#function) mutationEvent: \(mutationEvent)")
+        let nextEventPromise = self.nextEventPromise.getAndSet(nil)
         var eventToPersist = mutationEvent
-        if nextEventPromise.get() != nil {
+        if nextEventPromise != nil {
             eventToPersist.inProcess = true
         }
+
         storageAdapter.save(eventToPersist, condition: nil, eagerLoad: true) { result in
             switch result {
             case .failure(let dataStoreError):
                 self.log.verbose("\(#function): Error saving mutation event: \(dataStoreError)")
+                // restore the `nextEventPromise` value when failed to save mutation event
+                // as nextEventPromise is expecting to hanlde error of querying unprocessed mutaiton events
+                // not the failure of saving mutaiton event operation
+                nextEventPromise.ifSome(self.nextEventPromise.set(_:))
             case .success(let savedMutationEvent):
                 self.log.verbose("\(#function): saved \(savedMutationEvent)")
-                if let nextEventPromise = self.nextEventPromise.getAndSet(nil) {
+                nextEventPromise.ifSome {
                     self.log.verbose("\(#function): invoking nextEventPromise with \(savedMutationEvent)")
-                    nextEventPromise(.success(savedMutationEvent))
+                    $0(.success(savedMutationEvent))
                 }
             }
             self.log.verbose("\(#function): invoking completionPromise with \(result)")

AmplifyPlugins/DataStore/Tests/DataStoreHostApp/AWSDataStorePluginIntegrationTests/DataStoreEndToEndTests.swift

@@ -479,6 +479,64 @@ class DataStoreEndToEndTests: SyncEngineIntegrationTestBase {
         try await Amplify.DataStore.start()
     }
 
+    ///
+    /// - Given: DataStore with clean state
+    /// - When:
+    ///     - do some mutaions to ensure MutationEventPublisher works fine
+    ///     - wait 1 second for OutgoingMutationQueue stateMachine transform to waitingForEventToProcess state
+    ///     - do some mutations in parallel
+    /// - Then: verify no mutaiton loss
+    func testParallelMutations_whenWaitingForEventToProcess_noMutationLoss() async throws {
+        await setUp(withModels: TestModelRegistration())
+        try await startAmplifyAndWaitForSync()
+        let parallelSize = 100
+        let initExpectation = expectation(description: "expect MutationEventPublisher works fine")
+        let parallelExpectation = expectation(description: "expect parallel processing no data loss")
+
+        let newPost = Post(title: UUID().uuidString, content: UUID().uuidString, createdAt: .now())
+
+        let titlePrefix = UUID().uuidString
+        let posts = (0..<parallelSize).map { Post(title: "\(titlePrefix)-\($0)", content: UUID().uuidString, createdAt: .now()) }
+        var expectedResult = Set<String>()
+        let cancellable = Amplify.Hub.publisher(for: .dataStore)
+            .sink { payload in
+                let event = DataStoreHubEvent(payload: payload)
+                switch event {
+                case .outboxMutationProcessed(let mutationEvent):
+                    guard mutationEvent.modelName == Post.modelName,
+                          let post = mutationEvent.element.model as? Post
+                    else
+                    { break }
+
+                    if post.title == newPost.title {
+                        initExpectation.fulfill()
+                    }
+
+                    if post.title.hasPrefix(titlePrefix) {
+                        expectedResult.insert(post.title)
+                    }
+
+                    if expectedResult.count == parallelSize {
+                        parallelExpectation.fulfill()
+                    }
+                default: break
+                }
+            }
+        try await Amplify.DataStore.save(newPost)
+        wait(for: [initExpectation], timeout: 5)
+        try await Task.sleep(seconds: 1)
+
+        for post in posts {
+            Task {
+                try? await Amplify.DataStore.save(post)
+            }
+        }
+
+        wait(for: [parallelExpectation], timeout: Double(parallelSize))
+        cancellable.cancel()
+        XCTAssertEqual(expectedResult, Set(posts.map(\.title)))
+    }
+
     // MARK: - Helpers
 
     func validateSavePost() async throws {

AmplifyTests/CoreTests/Optional+ExtentionTests.swift

@@ -0,0 +1,62 @@
+//
+// Copyright Amazon.com Inc. or its affiliates.
+// All Rights Reserved.
+//
+// SPDX-License-Identifier: Apache-2.0
+//
+		
+@_spi(OptionalExtension) import Amplify
+import XCTest
+
+class OptionalExtensionTests: XCTestCase {
+
+    /// - Given:
+    ///     Optional of integer with none
+    /// - When:
+    ///     apply a function to increase a captured local integer variable with value 0
+    /// - Then:
+    ///     - the ifSome function on optional should not be applied
+    ///     - local integer variable value stays 0
+    func testIfSome_withNone_doNothing() {
+        var sideEffect = 0
+        let optional: Int? = .none
+        optional.ifSome { _ in sideEffect += 1 }
+        XCTAssertEqual(0, sideEffect)
+    }
+
+    /// - Given:
+    ///     Optional of integer with value 10
+    /// - When:
+    ///     apply a function to increase a captured local integer variable with value 0
+    /// - Then:
+    ///     - the ifSome function on optioanl should be applied
+    ///     - capture local integer value equals 1
+    func testIfSome_withValue_applyFunction() {
+        var sideEffect = 0
+        let optional: Int? = .some(10)
+        optional.ifSome { _ in sideEffect += 1 }
+        XCTAssertEqual(1, sideEffect)
+    }
+
+    /// - Given:
+    ///     Optional of integer with value 10
+    /// - When:
+    ///     apply a function that throw error
+    /// - Then:
+    ///     - the ifSome function on optioanl should be applied
+    ///     - the error is rethrowed
+    func testIfSome_withValue_applyFunctionRethrowError() {
+        let optional: Int? = .some(10)
+        let expectedError = TestRuntimeError()
+        XCTAssertThrowsError(try optional.ifSome {_ in
+            throw expectedError
+        }) { error in
+            XCTAssertEqual(expectedError, error as? TestRuntimeError)
+        }
+    }
+
+}
+
+fileprivate struct TestRuntimeError: Error, Equatable {
+    let id = UUID()
+}