Operations: let operation subclasses handle early cancellation within main()

6 files changed, 70 insertions, 35 deletions

diff --git a/ios/MullvadVPN/Operations/AsyncBlockOperation.swift b/ios/MullvadVPN/Operations/AsyncBlockOperation.swift
index a7f0ffdd29..a26e4e249f 100644
--- a/ios/MullvadVPN/Operations/AsyncBlockOperation.swift
+++ b/ios/MullvadVPN/Operations/AsyncBlockOperation.swift
@@ -10,16 +10,13 @@ import Foundation
 
 /// Asynchronous block operation
 class AsyncBlockOperation: AsyncOperation {
-    private let block: ((@escaping () -> Void) -> Void)
+    private let block: ((AsyncBlockOperation) -> Void)
 
-    init(_ block: @escaping (@escaping () -> Void) -> Void) {
+    init(block: @escaping (AsyncBlockOperation) -> Void) {
         self.block = block
-        super.init()
     }
 
     override func main() {
-        block { [weak self] in
-            self?.finish()
-        }
+        block(self)
     }
 }
diff --git a/ios/MullvadVPN/Operations/AsyncOperation.swift b/ios/MullvadVPN/Operations/AsyncOperation.swift
index 2d1819ee30..e8572ca56e 100644
--- a/ios/MullvadVPN/Operations/AsyncOperation.swift
+++ b/ios/MullvadVPN/Operations/AsyncOperation.swift
@@ -36,15 +36,10 @@ class AsyncOperation: Operation {
     }
 
     final override func start() {
-        stateLock.lock()
-        if _isCancelled {
-            finish()
-            stateLock.unlock()
-        } else {
+        stateLock.withCriticalBlock {
             setExecuting(true)
-            stateLock.unlock()
-            main()
         }
+        main()
     }
 
     override func main() {
diff --git a/ios/MullvadVPN/Operations/PresentAlertOperation.swift b/ios/MullvadVPN/Operations/PresentAlertOperation.swift
index 0bc420f770..7ed1efcc3f 100644
--- a/ios/MullvadVPN/Operations/PresentAlertOperation.swift
+++ b/ios/MullvadVPN/Operations/PresentAlertOperation.swift
@@ -26,6 +26,9 @@ class PresentAlertOperation: AsyncOperation {
             // Guard against executing cancellation more than once.
             guard !self.isCancelled else { return }
 
+            // Call super implementation to toggle isCancelled flag
+            super.cancel()
+
             // Guard against trying to dismiss the alert when operation hasn't started yet.
             guard self.isExecuting else { return }
 
@@ -33,14 +36,16 @@ class PresentAlertOperation: AsyncOperation {
             if !self.alertController.isBeingPresented && !self.alertController.isBeingDismissed {
                 self.dismissAndFinish()
             }
-
-            // Call super implementation to toggle isCancelled flag
-            super.cancel()
         }
     }
 
     override func main() {
         DispatchQueue.main.async {
+            guard !self.isCancelled else {
+                self.finish()
+                return
+            }
+
             NotificationCenter.default.addObserver(
                 self,
                 selector: #selector(self.alertControllerDidDismiss(_:)),
diff --git a/ios/MullvadVPN/Operations/ProductsRequestOperation.swift b/ios/MullvadVPN/Operations/ProductsRequestOperation.swift
index e1f8c99136..6e34e05977 100644
--- a/ios/MullvadVPN/Operations/ProductsRequestOperation.swift
+++ b/ios/MullvadVPN/Operations/ProductsRequestOperation.swift
@@ -11,7 +11,7 @@ import StoreKit
 
 class ProductsRequestOperation: AsyncOperation, SKProductsRequestDelegate {
     private let productIdentifiers: Set<String>
-    private let completionHandler: (Result<SKProductsResponse, Error>) -> Void
+    private var completionHandler: ((Result<SKProductsResponse, Error>) -> Void)?
 
     private let maxRetryCount = 10
     private let retryDelay: DispatchTimeInterval = .seconds(2)
@@ -20,6 +20,12 @@ class ProductsRequestOperation: AsyncOperation, SKProductsRequestDelegate {
     private var retryTimer: DispatchSourceTimer?
     private var request: SKProductsRequest?
 
+    struct OperationCancelledError: LocalizedError {
+        var errorDescription: String? {
+            return "Operation is cancelled"
+        }
+    }
+
     init(productIdentifiers: Set<String>, completionHandler: @escaping (Result<SKProductsResponse, Error>) -> Void) {
         self.productIdentifiers = productIdentifiers
         self.completionHandler = completionHandler
@@ -29,6 +35,11 @@ class ProductsRequestOperation: AsyncOperation, SKProductsRequestDelegate {
 
     override func main() {
         DispatchQueue.main.async {
+            guard !self.isCancelled else {
+                self.finish(with: .failure(OperationCancelledError()))
+                return
+            }
+
             self.startRequest()
         }
     }
@@ -89,7 +100,11 @@ class ProductsRequestOperation: AsyncOperation, SKProductsRequestDelegate {
     }
 
     private func finish(with result: Result<SKProductsResponse, Error>) {
-        completionHandler(result)
+        assert(Thread.isMainThread)
+
+        completionHandler?(result)
+        completionHandler = nil
+
         finish()
     }
 }
diff --git a/ios/MullvadVPN/Operations/ReceiptRefreshOperation.swift b/ios/MullvadVPN/Operations/ReceiptRefreshOperation.swift
index c386fc1cb1..61d9857904 100644
--- a/ios/MullvadVPN/Operations/ReceiptRefreshOperation.swift
+++ b/ios/MullvadVPN/Operations/ReceiptRefreshOperation.swift
@@ -11,36 +11,59 @@ import StoreKit
 
 class ReceiptRefreshOperation: AsyncOperation, SKRequestDelegate {
     private let request: SKReceiptRefreshRequest
-    private let completionHandler: (Result<(), Error>) -> Void
+    private var completionHandler: ((Result<(), Error>) -> Void)?
 
-    init(receiptProperties: [String: Any]?, completionHandler: @escaping (Result<(), Error>) -> Void) {
-        request = SKReceiptRefreshRequest(receiptProperties: receiptProperties)
-        self.completionHandler = completionHandler
-
-        super.init()
+    struct OperationCancelledError: LocalizedError {
+        var errorDescription: String? {
+            return "Operation is cancelled"
+        }
+    }
 
-        request.delegate = self
+    init(receiptProperties: [String: Any]?, completionHandler completion: @escaping (Result<(), Error>) -> Void) {
+        request = SKReceiptRefreshRequest(receiptProperties: receiptProperties)
+        completionHandler = completion
     }
 
     override func main() {
-        request.start()
+        DispatchQueue.main.async {
+            guard !self.isCancelled else {
+                self.finish(with: .failure(OperationCancelledError()))
+                return
+            }
+
+            self.request.delegate = self
+            self.request.start()
+        }
     }
 
     override func cancel() {
-        super.cancel()
+        DispatchQueue.main.async {
+            super.cancel()
 
-        request.cancel()
+            self.request.cancel()
+        }
     }
 
     // - MARK: SKRequestDelegate
 
     func requestDidFinish(_ request: SKRequest) {
-        completionHandler(.success(()))
-        finish()
+        DispatchQueue.main.async {
+            self.finish(with: .success(()))
+        }
     }
 
     func request(_ request: SKRequest, didFailWithError error: Error) {
-        completionHandler(.failure(error))
+        DispatchQueue.main.async {
+            self.finish(with: .failure(error))
+        }
+    }
+
+    private func finish(with result: Result<(), Error>) {
+        assert(Thread.isMainThread)
+
+        completionHandler?(result)
+        completionHandler = nil
+
         finish()
     }
 }
diff --git a/ios/MullvadVPN/Promise/Promise+OperationQueue.swift b/ios/MullvadVPN/Promise/Promise+OperationQueue.swift
index be8c4e5166..1ee55db770 100644
--- a/ios/MullvadVPN/Promise/Promise+OperationQueue.swift
+++ b/ios/MullvadVPN/Promise/Promise+OperationQueue.swift
@@ -13,10 +13,10 @@ extension Promise {
     /// Returns a promise that adds operation that finishes along with the upstream.
     func run(on operationQueue: OperationQueue) -> Promise<Value> {
         return Promise { resolver in
-            let operation = AsyncBlockOperation { finish in
+            let operation = AsyncBlockOperation { operation in
                 self.observe { completion in
                     resolver.resolve(completion: completion)
-                    finish()
+                    operation.finish()
                 }
             }
 
@@ -31,10 +31,10 @@ extension Promise {
     /// Returns a promise that adds a mutually exclusive operation that finishes along with the upstream.
     func run(on operationQueue: OperationQueue, categories: [String]) -> Promise<Value> {
         return Promise { resolver in
-            let operation = AsyncBlockOperation { finish in
+            let operation = AsyncBlockOperation { operation in
                 self.observe { completion in
                     resolver.resolve(completion: completion)
-                    finish()
+                    operation.finish()
                 }
             }