Merge branch 'background-fetch-cancellation'

4 files changed, 106 insertions, 38 deletions

diff --git a/ios/MullvadVPN/AppDelegate.swift b/ios/MullvadVPN/AppDelegate.swift
index 3a673b17aa..aeb4dde33e 100644
--- a/ios/MullvadVPN/AppDelegate.swift
+++ b/ios/MullvadVPN/AppDelegate.swift
@@ -184,14 +184,21 @@ class AppDelegate: UIResponder, UIApplicationDelegate {
         let operationQueue = OperationQueue()
 
         let updateAddressCacheOperation = AsyncBlockOperation { operation in
-            _ = self.addressCacheTracker.updateEndpoints { result in
+            let handle = self.addressCacheTracker.updateEndpoints { result in
                 addressCacheFetchResult = result.backgroundFetchResult
                 operation.finish()
             }
+
+            operation.addCancellationBlock {
+                handle.cancel()
+            }
         }
 
         let updateRelaysOperation = AsyncBlockOperation { operation in
+            var cancellationToken: PromiseCancellationToken?
+
             RelayCache.Tracker.shared.updateRelays()
+                .storeCancellationToken(in: &cancellationToken)
                 .observe { completion in
                     switch completion.unwrappedValue {
                     case .success(let result):
@@ -206,15 +213,14 @@ class AppDelegate: UIResponder, UIApplicationDelegate {
 
                     operation.finish()
                 }
-        }
 
-        let rotatePrivateKeyOperation = AsyncBlockOperation { operation in
-            guard !operation.isCancelled else {
-                operation.finish()
-                return
+            operation.addCancellationBlock {
+                cancellationToken?.cancel()
             }
+        }
 
-            TunnelManager.shared.rotatePrivateKey { rotationResult, error in
+        let rotatePrivateKeyOperation = AsyncBlockOperation { operation in
+            let handle = TunnelManager.shared.rotatePrivateKey { rotationResult, error in
                 if let error = error {
                     self.logger?.error(chainedError: error, message: "Failed to rotate the key")
 
@@ -233,32 +239,38 @@ class AppDelegate: UIResponder, UIApplicationDelegate {
 
                 operation.finish()
             }
+
+            operation.addCancellationBlock {
+                handle.cancel()
+            }
         }
 
-        rotatePrivateKeyOperation.addDependency(updateRelaysOperation)
-        rotatePrivateKeyOperation.addDependency(updateAddressCacheOperation)
+        rotatePrivateKeyOperation.addDependencies([updateRelaysOperation, updateAddressCacheOperation])
 
         let backgroundTaskIdentifier = UIApplication.shared.beginBackgroundTask(withName: "AppDelegate.performFetch") {
             operationQueue.cancelAllOperations()
         }
 
-        rotatePrivateKeyOperation.completionBlock = {
-            DispatchQueue.main.async {
-                let operationResults = [addressCacheFetchResult, relaysFetchResult, rotatePrivateKeyFetchResult].compactMap { $0 }
-                let initialResult = operationResults.first ?? .failed
-                let backgroundFetchResult = operationResults.reduce(initialResult) { partialResult, other in
-                    return partialResult.combine(with: other)
-                }
+        let fetchOperations = [updateAddressCacheOperation, updateRelaysOperation, rotatePrivateKeyOperation]
 
-                self.logger?.info("Finish background refresh with \(backgroundFetchResult)")
+        let completionOperation = BlockOperation {
+            let operationResults = [addressCacheFetchResult, relaysFetchResult, rotatePrivateKeyFetchResult].compactMap { $0 }
+            let initialResult = operationResults.first ?? .failed
+            let backgroundFetchResult = operationResults.reduce(initialResult) { partialResult, other in
+                return partialResult.combine(with: other)
+            }
 
-                completionHandler(backgroundFetchResult)
+            self.logger?.info("Finish background refresh with \(backgroundFetchResult)")
 
-                UIApplication.shared.endBackgroundTask(backgroundTaskIdentifier)
-            }
+            completionHandler(backgroundFetchResult)
+
+            UIApplication.shared.endBackgroundTask(backgroundTaskIdentifier)
         }
 
-        operationQueue.addOperations([updateAddressCacheOperation, updateRelaysOperation, rotatePrivateKeyOperation], waitUntilFinished: false)
+        completionOperation.addDependencies(fetchOperations)
+
+        operationQueue.addOperations(fetchOperations, waitUntilFinished: false)
+        OperationQueue.main.addOperation(completionOperation)
     }
 
     // MARK: - Private
diff --git a/ios/MullvadVPN/Operations/AsyncBlockOperation.swift b/ios/MullvadVPN/Operations/AsyncBlockOperation.swift
index a26e4e249f..c02577471f 100644
--- a/ios/MullvadVPN/Operations/AsyncBlockOperation.swift
+++ b/ios/MullvadVPN/Operations/AsyncBlockOperation.swift
@@ -10,13 +10,48 @@ import Foundation
 
 /// Asynchronous block operation
 class AsyncBlockOperation: AsyncOperation {
-    private let block: ((AsyncBlockOperation) -> Void)
+    private let stateLock = NSLock()
+
+    private var executionBlock: ((AsyncBlockOperation) -> Void)?
+    private var cancellationBlocks: [() -> Void] = []
 
     init(block: @escaping (AsyncBlockOperation) -> Void) {
-        self.block = block
+        executionBlock = block
     }
 
     override func main() {
-        block(self)
+        executionBlock?(self)
+        executionBlock = nil
+    }
+
+    override func cancel() {
+        super.cancel()
+
+        stateLock.lock()
+        let blocks = cancellationBlocks
+        cancellationBlocks.removeAll()
+        stateLock.unlock()
+
+        for block in blocks {
+            block()
+        }
+    }
+
+    func addCancellationBlock(_ block: @escaping () -> Void) {
+        stateLock.lock()
+
+        if isCancelled {
+            stateLock.unlock()
+            block()
+        } else {
+            cancellationBlocks.append(block)
+            stateLock.unlock()
+        }
+    }
+
+    override func operationDidFinish() {
+        stateLock.lock()
+        cancellationBlocks.removeAll()
+        stateLock.unlock()
     }
 }
diff --git a/ios/MullvadVPN/Operations/AsyncOperation.swift b/ios/MullvadVPN/Operations/AsyncOperation.swift
index e8572ca56e..c4f628434d 100644
--- a/ios/MullvadVPN/Operations/AsyncOperation.swift
+++ b/ios/MullvadVPN/Operations/AsyncOperation.swift
@@ -10,7 +10,6 @@ import Foundation
 
 /// A base implementation of an asynchronous operation
 class AsyncOperation: Operation {
-
     /// A state lock used for manipulating the operation state flags in a thread safe fashion.
     private let stateLock = NSRecursiveLock()
 
@@ -58,22 +57,40 @@ class AsyncOperation: Operation {
     }
 
     final func finish() {
-        stateLock.withCriticalBlock {
-            if _isExecuting {
-               setExecuting(false)
-            }
+        stateLock.lock()
 
-            if !_isFinished {
-                willChangeValue(for: \.isFinished)
-                _isFinished = true
-                didChangeValue(for: \.isFinished)
-            }
+        if _isExecuting {
+           setExecuting(false)
+        }
+
+        if !_isFinished {
+            willChangeValue(for: \.isFinished)
+            _isFinished = true
+            didChangeValue(for: \.isFinished)
+
+            stateLock.unlock()
+
+            operationDidFinish()
+        } else {
+            stateLock.unlock()
         }
     }
 
+    func operationDidFinish() {
+        // Override in subclasses
+    }
+
     private func setExecuting(_ value: Bool) {
         willChangeValue(for: \.isExecuting)
         _isExecuting = value
         didChangeValue(for: \.isExecuting)
     }
 }
+
+extension Operation {
+    func addDependencies(_ dependencies: [Operation]) {
+        for dependency in dependencies {
+            addDependency(dependency)
+        }
+    }
+}
diff --git a/ios/MullvadVPN/TunnelManager/TunnelManager.swift b/ios/MullvadVPN/TunnelManager/TunnelManager.swift
index c2fe96adcd..a885bf56d4 100644
--- a/ios/MullvadVPN/TunnelManager/TunnelManager.swift
+++ b/ios/MullvadVPN/TunnelManager/TunnelManager.swift
@@ -124,7 +124,7 @@ class TunnelManager: TunnelManagerStateDelegate
         timer.setEventHandler { [weak self] in
             guard let self = self else { return }
 
-            self.rotatePrivateKey { rotationResult, error in
+            _ = self.rotatePrivateKey { rotationResult, error in
                 self.stateQueue.async {
                     if let scheduleDate = self.handlePrivateKeyRotationCompletion(result: rotationResult, error: error) {
                         guard self.isRunningPeriodicPrivateKeyRotation else { return }
@@ -352,7 +352,7 @@ class TunnelManager: TunnelManagerStateDelegate
         operationQueue.addOperation(operation)
     }
 
-    func rotatePrivateKey(completionHandler: @escaping (KeyRotationResult?, TunnelManager.Error?) -> Void) {
+    func rotatePrivateKey(completionHandler: @escaping (KeyRotationResult?, TunnelManager.Error?) -> Void) -> Cancellable {
         let operation = ReplaceKeyOperation.operationForKeyRotation(
             queue: stateQueue,
             state: state,
@@ -399,6 +399,10 @@ class TunnelManager: TunnelManagerStateDelegate
         exclusivityController.addOperation(operation, categories: [OperationCategory.changeTunnelSettings])
 
         operationQueue.addOperation(operation)
+
+        return AnyCancellable {
+            operation.cancel()
+        }
     }
 
     func setRelayConstraints(_ newConstraints: RelayConstraints, completionHandler: @escaping (TunnelManager.Error?) -> Void) {
@@ -666,7 +670,7 @@ extension TunnelManager {
     private func handleBackgroundTask(_ task: BGProcessingTask) {
         logger.debug("Start private key rotation task")
 
-        rotatePrivateKey { rotationResult, error in
+        let request = rotatePrivateKey { rotationResult, error in
             if let scheduleDate = self.handlePrivateKeyRotationCompletion(result: rotationResult, error: error) {
                 // Schedule next background task
                 switch self.submitBackgroundTask(at: scheduleDate) {
@@ -683,7 +687,7 @@ extension TunnelManager {
         }
 
         task.expirationHandler = {
-            // TODO: handle cancellation?
+            request.cancel()
         }
     }
 }