From a96d2ca597aa6bde29be17538473a6aafd079204 Mon Sep 17 00:00:00 2001
From: Doug <douglase@element.io>
Date: Wed, 27 Nov 2024 13:57:59 +0000
Subject: [PATCH] Wait until the sync has stopped before marking the task as
 complete.

---
 .../Sources/Application/AppCoordinator.swift  | 29 ++++++++-------
 .../Mocks/Generated/GeneratedMocks.swift      | 35 +++++++++++++++++++
 .../Sources/Services/Client/ClientProxy.swift |  2 +-
 .../Services/Client/ClientProxyProtocol.swift |  1 +
 4 files changed, 54 insertions(+), 13 deletions(-)

diff --git a/ElementX/Sources/Application/AppCoordinator.swift b/ElementX/Sources/Application/AppCoordinator.swift
index 111a913405..74f446003f 100644
--- a/ElementX/Sources/Application/AppCoordinator.swift
+++ b/ElementX/Sources/Application/AppCoordinator.swift
@@ -882,12 +882,12 @@ class AppCoordinator: AppCoordinatorProtocol, AuthenticationFlowCoordinatorDeleg
 
     // MARK: - Application State
 
-    private func stopSync(isBackgroundTask: Bool) {
+    private func stopSync(isBackgroundTask: Bool, completion: (() -> Void)? = nil) {
         if isBackgroundTask, UIApplication.shared.applicationState == .active {
             // Attempt to stop the background task sync loop cleanly, only if the app not already running
             return
         }
-        userSession?.clientProxy.stopSync()
+        userSession?.clientProxy.stopSync(completion: completion)
         clientProxyObserver = nil
     }
 
@@ -968,9 +968,11 @@ class AppCoordinator: AppCoordinatorProtocol, AuthenticationFlowCoordinatorDeleg
         backgroundTask = appMediator.beginBackgroundTask { [weak self] in
             guard let self else { return }
             
-            stopSync(isBackgroundTask: true)
-            
-            if let backgroundTask {
+            MXLog.info("Background task is about to expire.")
+            stopSync(isBackgroundTask: true) { [weak self] in
+                guard let self, let backgroundTask else { return }
+                
+                MXLog.info("Ending background task.")
                 appMediator.endBackgroundTask(backgroundTask)
                 self.backgroundTask = nil
             }
@@ -1026,10 +1028,12 @@ class AppCoordinator: AppCoordinatorProtocol, AuthenticationFlowCoordinatorDeleg
         scheduleBackgroundAppRefresh()
         
         task.expirationHandler = { [weak self] in
-            self?.stopSync(isBackgroundTask: true)
+            MXLog.info("Background app refresh task is about to expire.")
             
-            MXLog.info("Background app refresh task expired")
-            task.setTaskCompleted(success: true)
+            self?.stopSync(isBackgroundTask: true) {
+                MXLog.info("Marking Background app refresh task as complete.")
+                task.setTaskCompleted(success: true)
+            }
         }
         
         guard let userSession else {
@@ -1047,13 +1051,14 @@ class AppCoordinator: AppCoordinatorProtocol, AuthenticationFlowCoordinatorDeleg
             .sink(receiveValue: { [weak self] _ in
                 guard let self else { return }
                 MXLog.info("Background app refresh finished")
+                backgroundRefreshSyncObserver?.cancel()
                 
                 // Make sure we stop the sync loop, otherwise the ongoing request is immediately
                 // handled the next time the app refreshes, which can trigger timeout failures.
-                stopSync(isBackgroundTask: true)
-                backgroundRefreshSyncObserver?.cancel()
-                
-                task.setTaskCompleted(success: true)
+                stopSync(isBackgroundTask: true) {
+                    MXLog.info("Marking Background app refresh task as complete.")
+                    task.setTaskCompleted(success: true)
+                }
             })
     }
 }
diff --git a/ElementX/Sources/Mocks/Generated/GeneratedMocks.swift b/ElementX/Sources/Mocks/Generated/GeneratedMocks.swift
index 45e2fd5d50..99b5ae7cd6 100644
--- a/ElementX/Sources/Mocks/Generated/GeneratedMocks.swift
+++ b/ElementX/Sources/Mocks/Generated/GeneratedMocks.swift
@@ -2346,6 +2346,41 @@ class ClientProxyMock: ClientProxyProtocol {
         stopSyncCallsCount += 1
         stopSyncClosure?()
     }
+    //MARK: - stopSync
+
+    var stopSyncCompletionUnderlyingCallsCount = 0
+    var stopSyncCompletionCallsCount: Int {
+        get {
+            if Thread.isMainThread {
+                return stopSyncCompletionUnderlyingCallsCount
+            } else {
+                var returnValue: Int? = nil
+                DispatchQueue.main.sync {
+                    returnValue = stopSyncCompletionUnderlyingCallsCount
+                }
+
+                return returnValue!
+            }
+        }
+        set {
+            if Thread.isMainThread {
+                stopSyncCompletionUnderlyingCallsCount = newValue
+            } else {
+                DispatchQueue.main.sync {
+                    stopSyncCompletionUnderlyingCallsCount = newValue
+                }
+            }
+        }
+    }
+    var stopSyncCompletionCalled: Bool {
+        return stopSyncCompletionCallsCount > 0
+    }
+    var stopSyncCompletionClosure: (((() -> Void)?) -> Void)?
+
+    func stopSync(completion: (() -> Void)?) {
+        stopSyncCompletionCallsCount += 1
+        stopSyncCompletionClosure?(completion)
+    }
     //MARK: - accountURL
 
     var accountURLActionUnderlyingCallsCount = 0
diff --git a/ElementX/Sources/Services/Client/ClientProxy.swift b/ElementX/Sources/Services/Client/ClientProxy.swift
index 61c0aa4888..41d264eda0 100644
--- a/ElementX/Sources/Services/Client/ClientProxy.swift
+++ b/ElementX/Sources/Services/Client/ClientProxy.swift
@@ -305,7 +305,7 @@ class ClientProxy: ClientProxyProtocol {
         stopSync(completion: nil)
     }
     
-    private func stopSync(completion: (() -> Void)?) {
+    func stopSync(completion: (() -> Void)?) {
         MXLog.info("Stopping sync")
         
         if restartTask != nil {
diff --git a/ElementX/Sources/Services/Client/ClientProxyProtocol.swift b/ElementX/Sources/Services/Client/ClientProxyProtocol.swift
index 0feae1793e..562e3fa902 100644
--- a/ElementX/Sources/Services/Client/ClientProxyProtocol.swift
+++ b/ElementX/Sources/Services/Client/ClientProxyProtocol.swift
@@ -123,6 +123,7 @@ protocol ClientProxyProtocol: AnyObject, MediaLoaderProtocol {
     func startSync()
 
     func stopSync()
+    func stopSync(completion: (() -> Void)?) // Hopefully this will become async once we get SE-0371.
     
     func accountURL(action: AccountManagementAction) async -> URL?