fix: update last event id after CC transaction completes - WPB-22219 (#4005)

Author: typfel

WireDomain/Sources/WireDomain/Synchronization/PullPendingUpdateEventsSync.swift

@@ -91,10 +91,11 @@ public struct PullPendingUpdateEventsSync: PullPendingUpdateEventsSyncProtocol {
             // We'll insert new events from this index.
             let currentIndex = try await store.indexOfLastEventEnvelope() + 1
 
+            var lastEnvelopeID: UUID?
+
             // We are decrypting the batch within one core crypto transaction
             try await coreCryptoProvider.coreCrypto().perform { context in
 
-                var lastEnvelopeID: UUID?
                 var decryptedEnvelopes: [UpdateEventEnvelope] = []
 
                 for envelope in envelopes {
@@ -125,15 +126,18 @@ public struct PullPendingUpdateEventsSync: PullPendingUpdateEventsSyncProtocol {
                     decryptedEnvelopes,
                     index: currentIndex
                 )
+            }
 
-                if let lastEnvelopeID {
-                    // We keep track of the last event id so next time we fetch
-                    // only new events. We don't track transient events because
-                    // these events aren't stored in the backend.
-                    WireLogger.sync.debug("storing last event id", attributes: [.eventEnvelopeID: lastEnvelopeID])
-                    store.storeLastEventID(id: lastEnvelopeID)
-                }
-
+            if let lastEnvelopeID {
+                // We keep track of the last event id so next time we fetch
+                // only new events. We don't track transient events because
+                // these events aren't stored in the backend.
+                //
+                // NOTE: it's important the we are updating the last event ID
+                // after the CC transaction has successfully completed,
+                // otherwise we risk data loss in case of a crash.
+                WireLogger.sync.debug("storing last event id", attributes: [.eventEnvelopeID: lastEnvelopeID])
+                store.storeLastEventID(id: lastEnvelopeID)
             }
         }
 

WireDomain/Sources/WireDomain/Synchronization/PullPendingUpdateEventsV2Sync.swift

@@ -146,7 +146,10 @@ public struct PullPendingUpdateEventsSyncV2: PullPendingUpdateEventsSyncV2Protoc
             }
         }
 
-        // ack
+        // ack the decrypted events
+        //
+        // NOTE: it's important that we ack after the CC transaction has succesfully completed,
+        // otherwise we risk data loss in case of a crash.
         if let lastEnvelope = storedEnvelopes.last?.0 {
             try await acknowledgeUntilEnvelope(lastEnvelope, through: pushChannel, batchSize: storedEnvelopes.count)
         }

WireDomain/Tests/WireDomainTests/Synchronization/PullPendingUpdateEventsSyncTests.swift

@@ -120,6 +120,52 @@ final class PullPendingUpdateEventsSyncTests: XCTestCase {
         XCTAssertEqual(journal[.brokenMLSGroupIDs].first, Scaffolding.mlsGroupID)
     }
 
+    func testLastEventIDIsNotPersisted_untilTransactionIsCompleted() async throws {
+        // Mock
+        store.lastEventID_MockValue = Scaffolding.lastEventID
+        store.indexOfLastEventEnvelope_MockValue = Scaffolding.indexOfLastEventEnvelope
+
+        api.getUpdateEventsSelfClientIDSinceEventID_MockValue = PayloadPager(start: "page2") { start in
+            switch start {
+            case "page2":
+                return Scaffolding.page2
+
+            default:
+                throw "unknown page: \(start ?? "nil")"
+            }
+        }
+
+        decryptor.decryptEventsInContext_MockMethod = { envelope, _ in
+            EventDecryptorResult(events: envelope.events, brokenMLSGroupIDs: [Scaffolding.mlsGroupID])
+        }
+
+        store.persistEventEnvelopesIndex_MockMethod = { _, _ in }
+        store.storeLastEventIDId_MockMethod = { _ in }
+        store.storeServerTimeDelta_MockMethod = { _ in }
+
+        coreCrypto.completeTransactionByDefault = false
+
+        // When
+        let pullingEventsTask = Task { [sut] in
+            try await sut.pull()
+        }
+
+        // we wait until the sync tries to commit the batch of decrypted events
+        try await coreCrypto.waitUntilTransactionIsPending()
+
+        // Then
+        try XCTAssertCount(store.storeLastEventIDId_Invocations, count: 0)
+
+        coreCrypto.completeAllTransactions()
+        _ = await pullingEventsTask.result
+
+        // after allowing the transaction to complete we should we see
+        // that the last event ID got persisted
+        let storeLastEventIDInvocations = store.storeLastEventIDId_Invocations
+        try XCTAssertCount(storeLastEventIDInvocations, count: 1)
+        XCTAssertEqual(storeLastEventIDInvocations[0], Scaffolding.envelope4.id)
+    }
+
 }
 
 private enum Scaffolding {

wire-ios-data-model/Support/Sources/MockSafeCoreCrypto.swift

@@ -24,6 +24,9 @@ public class MockSafeCoreCrypto: SafeCoreCryptoProtocol {
 
     var coreCrypto: MockCoreCryptoProtocol
     var coreCryptoContext: MockCoreCryptoContextProtocol
+    var completeTransactionByDefault: Bool = true
+
+    private var transactionContinuations: [CheckedContinuation<Void, Never>] = []
 
     public init(
         coreCrypto: MockCoreCryptoProtocol = .init(),
@@ -48,7 +51,16 @@ public class MockSafeCoreCrypto: SafeCoreCryptoProtocol {
     var performAsyncCount = 0
     public func perform<T>(_ block: (WireCoreCrypto.CoreCryptoContextProtocol) async throws -> T) async rethrows -> T {
         performAsyncCount += 1
-        return try await block(coreCryptoContext)
+
+        let result = try await block(coreCryptoContext)
+
+        if !completeTransactionByDefault {
+            await withCheckedContinuation { continuation in
+                transactionContinuations.append(continuation)
+            }
+        }
+
+        return result
     }
 
     public func configure(block: (any WireCoreCrypto.CoreCryptoProtocol) async throws -> Void) async throws {
@@ -65,6 +77,19 @@ public class MockSafeCoreCrypto: SafeCoreCryptoProtocol {
         try mock(clientID)
     }
 
+    public func waitUntilTransactionIsPending() async throws {
+        while transactionContinuations.isEmpty {
+            try await Task.sleep(nanoseconds: 1_000_000)
+        }
+    }
+
+    public func completeAllTransactions() {
+        transactionContinuations.forEach { cont in
+            cont.resume()
+        }
+
+    }
+
     var tearDownCount = 0
     public func tearDown() throws {
         tearDownCount += 1