AsyncOperation: add shortcuts for cancellation and completion events observation

4 files changed, 27 insertions, 22 deletions

diff --git a/ios/MullvadVPN/TunnelManager/SetAccountOperation.swift b/ios/MullvadVPN/TunnelManager/SetAccountOperation.swift
index fd8e7008ec..4f99c86bdc 100644
--- a/ios/MullvadVPN/TunnelManager/SetAccountOperation.swift
+++ b/ios/MullvadVPN/TunnelManager/SetAccountOperation.swift
@@ -122,11 +122,9 @@ class SetAccountOperation: ResultOperation<StoredAccountData?> {
                     })
                     .reduce()
 
-                saveSettingsOperation.addBlockObserver(
-                    OperationBlockObserver(didFinish: { operation, error in
-                        self.completeOperation(accountData: operation.output)
-                    })
-                )
+                saveSettingsOperation.onFinish { operation, error in
+                    self.completeOperation(accountData: operation.output)
+                }
 
                 return [accountOperation, createDeviceOperation, saveSettingsOperation]
             } ?? []
diff --git a/ios/MullvadVPN/TunnelManager/Tunnel+Messaging.swift b/ios/MullvadVPN/TunnelManager/Tunnel+Messaging.swift
index 3b69dd34eb..f0902e4e29 100644
--- a/ios/MullvadVPN/TunnelManager/Tunnel+Messaging.swift
+++ b/ios/MullvadVPN/TunnelManager/Tunnel+Messaging.swift
@@ -73,21 +73,19 @@ extension Tunnel {
             completionHandler: completionHandler
         )
 
-        operation.addBlockObserver(
-            OperationBlockObserver(didCancel: { [weak self] _ in
-                guard let self = self else { return }
+        operation.onCancel { [weak self] _ in
+            guard let self = self else { return }
 
-                let cancelOperation = SendTunnelProviderMessageOperation(
-                    dispatchQueue: dispatchQueue,
-                    application: .shared,
-                    tunnel: self,
-                    message: .cancelURLRequest(proxyRequest.id),
-                    completionHandler: nil
-                )
+            let cancelOperation = SendTunnelProviderMessageOperation(
+                dispatchQueue: dispatchQueue,
+                application: .shared,
+                tunnel: self,
+                message: .cancelURLRequest(proxyRequest.id),
+                completionHandler: nil
+            )
 
-                operationQueue.addOperation(cancelOperation)
-            })
-        )
+            operationQueue.addOperation(cancelOperation)
+        }
 
         operationQueue.addOperation(operation)
 
diff --git a/ios/Operations/AsyncOperation.swift b/ios/Operations/AsyncOperation.swift
index 2939f0708e..78ae2fa627 100644
--- a/ios/Operations/AsyncOperation.swift
+++ b/ios/Operations/AsyncOperation.swift
@@ -422,6 +422,17 @@ public protocol OperationBlockObserverSupport {}
 extension AsyncOperation: OperationBlockObserverSupport {}
 
 extension OperationBlockObserverSupport where Self: AsyncOperation {
+    /// Add observer responding to cancellation event.
+    public func onCancel(_ fn: @escaping (Self) -> Void) {
+        addBlockObserver(OperationBlockObserver(didCancel: fn))
+    }
+
+    /// Add observer responding to finish event.
+    public func onFinish(_ fn: @escaping (Self, Error?) -> Void) {
+        addBlockObserver(OperationBlockObserver(didFinish: fn))
+    }
+
+    /// Add block-based observer.
     public func addBlockObserver(_ observer: OperationBlockObserver<Self>) {
         addObserver(observer)
     }
diff --git a/ios/Operations/AsyncOperationQueue.swift b/ios/Operations/AsyncOperationQueue.swift
index 4579546084..d1269a8175 100644
--- a/ios/Operations/AsyncOperationQueue.swift
+++ b/ios/Operations/AsyncOperationQueue.swift
@@ -73,11 +73,9 @@ private final class ExclusivityManager {
 
             operationsByCategory[category] = operations
 
-            let blockObserver = OperationBlockObserver(didFinish: { [weak self] op, error in
+            operation.onFinish { [weak self] op, error in
                 self?.removeOperation(op, categories: categories)
-            })
-
-            operation.addObserver(blockObserver)
+            }
         }
     }