Pass closure to block operations instead of the entire operation

Block operations receive (Error?) -> Void
Result block operations receive (Result<Output, Error>) -> Void
Transform operations receive (Input, (Result<Output, Error>) -> Void) -> Void

17 files changed, 223 insertions, 232 deletions

diff --git a/ios/MullvadREST/RESTAccessTokenManager.swift b/ios/MullvadREST/RESTAccessTokenManager.swift
index a0fe79ba75..8b374b88d5 100644
--- a/ios/MullvadREST/RESTAccessTokenManager.swift
+++ b/ios/MullvadREST/RESTAccessTokenManager.swift
@@ -27,12 +27,11 @@ extension REST {
             accountNumber: String,
             completionHandler: @escaping (Result<REST.AccessTokenData, Swift.Error>) -> Void
         ) -> Cancellable {
-            let operation = ResultBlockOperation<REST.AccessTokenData>(
-                dispatchQueue: dispatchQueue,
-                cancellableTask: { operation -> Cancellable in
+            let operation =
+                ResultBlockOperation<REST.AccessTokenData>(dispatchQueue: dispatchQueue) { finish -> Cancellable in
                     if let tokenData = self.tokens[accountNumber], tokenData.expiry > Date() {
-                        operation.finish(result: .success(tokenData))
-                        return AnyCancellable {}
+                        finish(.success(tokenData))
+                        return AnyCancellable()
                     }
 
                     return self.proxy.getAccessToken(accountNumber: accountNumber, retryStrategy: .noRetry) { result in
@@ -51,11 +50,10 @@ extension REST {
                                 break
                             }
 
-                            operation.finish(result: result)
+                            finish(result)
                         }
                     }
                 }
-            )
 
             operation.completionQueue = .main
             operation.completionHandler = completionHandler
diff --git a/ios/MullvadTypes/Cancellable.swift b/ios/MullvadTypes/Cancellable.swift
index 1a8db756f6..8f658a6da1 100644
--- a/ios/MullvadTypes/Cancellable.swift
+++ b/ios/MullvadTypes/Cancellable.swift
@@ -14,14 +14,21 @@ public protocol Cancellable {
 
 extension Operation: Cancellable {}
 
+/// An object representing a cancellation token.
 public final class AnyCancellable: Cancellable {
-    private let block: () -> Void
+    private let block: (() -> Void)?
 
+    /// Create cancellation token with block handler.
     public init(block: @escaping () -> Void) {
         self.block = block
     }
 
+    /// Create empty cancellation token.
+    public init() {
+        block = nil
+    }
+
     public func cancel() {
-        block()
+        block?()
     }
 }
diff --git a/ios/MullvadVPN/AddressCacheTracker/AddressCacheTracker.swift b/ios/MullvadVPN/AddressCacheTracker/AddressCacheTracker.swift
index 3d281bf869..d7c402c802 100644
--- a/ios/MullvadVPN/AddressCacheTracker/AddressCacheTracker.swift
+++ b/ios/MullvadVPN/AddressCacheTracker/AddressCacheTracker.swift
@@ -85,18 +85,18 @@ final class AddressCacheTracker {
     }
 
     func updateEndpoints(completionHandler: ((Result<Bool, Error>) -> Void)? = nil) -> Cancellable {
-        let operation = ResultBlockOperation<Bool>(cancellableTask: { operation -> Cancellable in
+        let operation = ResultBlockOperation<Bool> { finish -> Cancellable in
             guard self.nextScheduleDate() <= Date() else {
-                operation.finish(result: .success(false))
-                return AnyCancellable {}
+                finish(.success(false))
+                return AnyCancellable()
             }
 
             return self.apiProxy.getAddressList(retryStrategy: .default) { result in
                 self.setEndpoints(from: result)
 
-                operation.finish(result: result.map { _ in true })
+                finish(result.map { _ in true })
             }
-        })
+        }
 
         operation.completionQueue = .main
         operation.completionHandler = completionHandler
diff --git a/ios/MullvadVPN/AppDelegate.swift b/ios/MullvadVPN/AppDelegate.swift
index 640d7fc86c..233c17d73f 100644
--- a/ios/MullvadVPN/AppDelegate.swift
+++ b/ios/MullvadVPN/AppDelegate.swift
@@ -370,7 +370,7 @@ class AppDelegate: UIResponder, UIApplicationDelegate, UNUserNotificationCenterD
     private func startInitialization(application: UIApplication) {
         let wipeSettingsOperation = getWipeSettingsOperation()
 
-        let loadTunnelStoreOperation = AsyncBlockOperation(dispatchQueue: .main) { operation in
+        let loadTunnelStoreOperation = AsyncBlockOperation(dispatchQueue: .main) { finish in
             self.tunnelStore.loadPersistentTunnels { error in
                 if let error = error {
                     self.logger.error(
@@ -378,16 +378,14 @@ class AppDelegate: UIResponder, UIApplicationDelegate, UNUserNotificationCenterD
                         message: "Failed to load persistent tunnels."
                     )
                 }
-                operation.finish()
+                finish(nil)
             }
         }
 
-        let migrateSettingsOperation = ResultBlockOperation<SettingsMigrationResult>(
-            dispatchQueue: .main
-        ) { operation in
+        let migrateSettingsOperation = ResultBlockOperation<SettingsMigrationResult>(dispatchQueue: .main) { finish in
             SettingsManager.migrateStore(with: self.proxyFactory) { migrationResult in
                 let finishHandler = {
-                    operation.finish(result: .success(migrationResult))
+                    finish(.success(migrationResult))
                 }
 
                 guard case let .failure(error) = migrationResult,
@@ -404,7 +402,7 @@ class AppDelegate: UIResponder, UIApplicationDelegate, UNUserNotificationCenterD
         }
         migrateSettingsOperation.addDependencies([wipeSettingsOperation, loadTunnelStoreOperation])
 
-        let initTunnelManagerOperation = AsyncBlockOperation(dispatchQueue: .main) { operation in
+        let initTunnelManagerOperation = AsyncBlockOperation(dispatchQueue: .main) { finish in
             self.tunnelManager.loadConfiguration { error in
                 // TODO: avoid throwing fatal error and show the problem report UI instead.
                 if let error = error {
@@ -416,7 +414,7 @@ class AppDelegate: UIResponder, UIApplicationDelegate, UNUserNotificationCenterD
                 NotificationManager.shared.updateNotifications()
                 self.storePaymentManager.startPaymentQueueMonitoring()
 
-                operation.finish()
+                finish(nil)
             }
         }
         initTunnelManagerOperation.addDependency(migrateSettingsOperation)
diff --git a/ios/MullvadVPN/RelayCacheTracker/RelayCacheTracker.swift b/ios/MullvadVPN/RelayCacheTracker/RelayCacheTracker.swift
index afc761ebf5..402063052d 100644
--- a/ios/MullvadVPN/RelayCacheTracker/RelayCacheTracker.swift
+++ b/ios/MullvadVPN/RelayCacheTracker/RelayCacheTracker.swift
@@ -97,18 +97,18 @@ final class RelayCacheTracker {
     func updateRelays(completionHandler: ((Result<RelaysFetchResult, Error>) -> Void)? = nil)
         -> Cancellable
     {
-        let operation = ResultBlockOperation<RelaysFetchResult>(cancellableTask: { operation in
+        let operation = ResultBlockOperation<RelaysFetchResult> { finish in
             let cachedRelays = try? self.getCachedRelays()
 
             if self.getNextUpdateDate() > Date() {
-                operation.finish(result: .success(.throttled))
-                return AnyCancellable {}
+                finish(.success(.throttled))
+                return AnyCancellable()
             }
 
             return self.apiProxy.getRelays(etag: cachedRelays?.etag, retryStrategy: .noRetry) { result in
-                operation.finish(result: self.handleResponse(result: result))
+                finish(self.handleResponse(result: result))
             }
-        })
+        }
 
         operation.addObserver(
             BackgroundObserver(
diff --git a/ios/MullvadVPN/StorePaymentManager/StorePaymentManager.swift b/ios/MullvadVPN/StorePaymentManager/StorePaymentManager.swift
index 7e8545bea5..62ff2f11f9 100644
--- a/ios/MullvadVPN/StorePaymentManager/StorePaymentManager.swift
+++ b/ios/MullvadVPN/StorePaymentManager/StorePaymentManager.swift
@@ -181,10 +181,12 @@ final class StorePaymentManager: NSObject, SKPaymentTransactionObserver {
         accountNumber: String,
         completionHandler: @escaping (StorePaymentManagerError?) -> Void
     ) {
-        let accountOperation = ResultBlockOperation<REST.AccountData>(dispatchQueue: .main) { op in
-            return self.accountsProxy.getAccountData(accountNumber: accountNumber, retryStrategy: .default) { result in
-                op.finish(result: result)
-            }
+        let accountOperation = ResultBlockOperation<REST.AccountData>(dispatchQueue: .main) { finish in
+            return self.accountsProxy.getAccountData(
+                accountNumber: accountNumber,
+                retryStrategy: .default,
+                completion: finish
+            )
         }
 
         accountOperation.addObserver(BackgroundObserver(
diff --git a/ios/MullvadVPN/TunnelManager/SetAccountOperation.swift b/ios/MullvadVPN/TunnelManager/SetAccountOperation.swift
index fc518500ce..e13e4373fe 100644
--- a/ios/MullvadVPN/TunnelManager/SetAccountOperation.swift
+++ b/ios/MullvadVPN/TunnelManager/SetAccountOperation.swift
@@ -184,7 +184,7 @@ class SetAccountOperation: ResultOperation<StoredAccountData?> {
     }
 
     private func getCreateAccountOperation() -> ResultBlockOperation<StoredAccountData> {
-        return ResultBlockOperation<StoredAccountData>(dispatchQueue: dispatchQueue, cancellableTask: { operation in
+        return ResultBlockOperation<StoredAccountData>(dispatchQueue: dispatchQueue) { finish -> Cancellable in
             self.logger.debug("Create new account...")
 
             return self.accountsProxy.createAccount(retryStrategy: .default) { result in
@@ -205,40 +205,37 @@ class SetAccountOperation: ResultOperation<StoredAccountData?> {
                     )
                 }
 
-                operation.finish(result: result)
+                finish(result)
             }
-        })
+        }
     }
 
     private func getExistingAccountOperation(accountNumber: String) -> ResultOperation<StoredAccountData> {
-        return ResultBlockOperation<StoredAccountData>(
-            dispatchQueue: dispatchQueue,
-            cancellableTask: { operation -> Cancellable in
-                self.logger.debug("Request account data...")
-
-                return self.accountsProxy
-                    .getAccountData(accountNumber: accountNumber, retryStrategy: .default) { result in
-                        let result = result.inspectError { error in
-                            guard !error.isOperationCancellationError else { return }
+        return ResultBlockOperation<StoredAccountData>(dispatchQueue: dispatchQueue) { finish -> Cancellable in
+            self.logger.debug("Request account data...")
 
-                            self.logger.error(
-                                error: error,
-                                message: "Failed to receive account data."
-                            )
-                        }.map { accountData -> StoredAccountData in
-                            self.logger.debug("Received account data.")
+            return self.accountsProxy
+                .getAccountData(accountNumber: accountNumber, retryStrategy: .default) { result in
+                    let result = result.inspectError { error in
+                        guard !error.isOperationCancellationError else { return }
 
-                            return StoredAccountData(
-                                identifier: accountData.id,
-                                number: accountNumber,
-                                expiry: accountData.expiry
-                            )
-                        }
+                        self.logger.error(
+                            error: error,
+                            message: "Failed to receive account data."
+                        )
+                    }.map { accountData -> StoredAccountData in
+                        self.logger.debug("Received account data.")
 
-                        operation.finish(result: result)
+                        return StoredAccountData(
+                            identifier: accountData.id,
+                            number: accountNumber,
+                            expiry: accountData.expiry
+                        )
                     }
-            }
-        )
+
+                    finish(result)
+                }
+        }
     }
 
     private func getDeleteDeviceOperation() -> AsyncBlockOperation? {
@@ -246,7 +243,7 @@ class SetAccountOperation: ResultOperation<StoredAccountData?> {
             return nil
         }
 
-        let operation = AsyncBlockOperation(dispatchQueue: dispatchQueue, cancellableTask: { operation -> Cancellable in
+        let operation = AsyncBlockOperation(dispatchQueue: dispatchQueue) { finish -> Cancellable in
             self.logger.debug("Delete current device...")
 
             return self.devicesProxy.deleteDevice(
@@ -268,15 +265,15 @@ class SetAccountOperation: ResultOperation<StoredAccountData?> {
                     break
                 }
 
-                operation.finish(error: result.error)
+                finish(result.error)
             }
-        })
+        }
 
         return operation
     }
 
     private func getUnsetDeviceStateOperation() -> AsyncBlockOperation {
-        return AsyncBlockOperation(dispatchQueue: dispatchQueue, block: { operation in
+        return AsyncBlockOperation(dispatchQueue: dispatchQueue) { finish in
             // Tell the caller to unsubscribe from VPN status notifications.
             self.interactor.prepareForVPNConfigurationDeletion()
 
@@ -289,7 +286,7 @@ class SetAccountOperation: ResultOperation<StoredAccountData?> {
 
             // Finish immediately if tunnel provider is not set.
             guard let tunnel = self.interactor.tunnel else {
-                operation.finish()
+                finish(nil)
                 return
             }
 
@@ -306,55 +303,53 @@ class SetAccountOperation: ResultOperation<StoredAccountData?> {
 
                     self.interactor.setTunnel(nil, shouldRefreshTunnelState: false)
 
-                    operation.finish()
+                    finish(nil)
                 }
             }
-        })
+        }
     }
 
     private func getCreateDeviceOperation() -> TransformOperation<StoredAccountData, (PrivateKey, REST.Device)> {
-        let createDeviceOperation = TransformOperation<StoredAccountData, (PrivateKey, REST.Device)>(
-            dispatchQueue: dispatchQueue,
-            cancellableTask: { storedAccountData, operation -> Cancellable in
-                self.logger.debug("Store last used account.")
+        return TransformOperation<StoredAccountData, (
+            PrivateKey,
+            REST.Device
+        )>(dispatchQueue: dispatchQueue) { storedAccountData, finish -> Cancellable in
+            self.logger.debug("Store last used account.")
 
-                do {
-                    try SettingsManager.setLastUsedAccount(storedAccountData.number)
-                } catch {
-                    self.logger.error(
-                        error: error,
-                        message: "Failed to store last used account number."
-                    )
-                }
+            do {
+                try SettingsManager.setLastUsedAccount(storedAccountData.number)
+            } catch {
+                self.logger.error(
+                    error: error,
+                    message: "Failed to store last used account number."
+                )
+            }
 
-                self.logger.debug("Create device...")
+            self.logger.debug("Create device...")
 
-                let privateKey = PrivateKey()
+            let privateKey = PrivateKey()
 
-                let request = REST.CreateDeviceRequest(
-                    publicKey: privateKey.publicKey,
-                    hijackDNS: false
-                )
+            let request = REST.CreateDeviceRequest(
+                publicKey: privateKey.publicKey,
+                hijackDNS: false
+            )
 
-                return self.devicesProxy.createDevice(
-                    accountNumber: storedAccountData.number,
-                    request: request,
-                    retryStrategy: .default
-                ) { result in
-                    let result = result
-                        .map { device in
-                            return (privateKey, device)
-                        }
-                        .inspectError { error in
-                            self.logger.error(error: error, message: "Failed to create device.")
-                        }
+            return self.devicesProxy.createDevice(
+                accountNumber: storedAccountData.number,
+                request: request,
+                retryStrategy: .default
+            ) { result in
+                let result = result
+                    .map { device in
+                        return (privateKey, device)
+                    }
+                    .inspectError { error in
+                        self.logger.error(error: error, message: "Failed to create device.")
+                    }
 
-                    operation.finish(result: result)
-                }
+                finish(result)
             }
-        )
-
-        return createDeviceOperation
+        }
     }
 
     private func getSaveSettingsOperation() -> TransformOperation<SetAccountResult, StoredAccountData> {
diff --git a/ios/MullvadVPN/TunnelManager/TunnelManager.swift b/ios/MullvadVPN/TunnelManager/TunnelManager.swift
index 5bf37d8574..b2644dcec0 100644
--- a/ios/MullvadVPN/TunnelManager/TunnelManager.swift
+++ b/ios/MullvadVPN/TunnelManager/TunnelManager.swift
@@ -268,24 +268,23 @@ final class TunnelManager: StorePaymentObserver {
     }
 
     func reconnectTunnel(selectNewRelay: Bool, completionHandler: ((Error?) -> Void)? = nil) {
-        let operation = AsyncBlockOperation(dispatchQueue: internalQueue, cancellableTask: { operation -> Cancellable in
-            guard let tunnel = self.tunnel else {
-                operation.finish(error: UnsetTunnelError())
-                return AnyCancellable {}
-            }
-
+        let operation = AsyncBlockOperation(dispatchQueue: internalQueue) { finish -> Cancellable in
             do {
+                guard let tunnel = self.tunnel else {
+                    throw UnsetTunnelError()
+                }
+
                 let selectorResult = selectNewRelay ? try self.selectRelay() : nil
 
                 return tunnel.reconnectTunnel(relaySelectorResult: selectorResult) { result in
-                    operation.finish(error: result.error)
+                    finish(result.error)
                 }
             } catch {
-                operation.finish(error: error)
+                finish(error)
 
-                return AnyCancellable {}
+                return AnyCancellable()
             }
-        })
+        }
 
         operation.completionBlock = {
             DispatchQueue.main.async {
diff --git a/ios/MullvadVPN/View controllers/Tunnel/MapViewController.swift b/ios/MullvadVPN/View controllers/Tunnel/MapViewController.swift
index 243672597a..bff8888133 100644
--- a/ios/MullvadVPN/View controllers/Tunnel/MapViewController.swift
+++ b/ios/MullvadVPN/View controllers/Tunnel/MapViewController.swift
@@ -234,9 +234,9 @@ final class MapViewController: UIViewController, MKMapViewDelegate {
         cancelOtherAnimations: Bool,
         block: @escaping (_ finish: @escaping () -> Void) -> Void
     ) {
-        let operation = AsyncBlockOperation(dispatchQueue: .main) { operation in
+        let operation = AsyncBlockOperation(dispatchQueue: .main) { finish in
             block {
-                operation.finish()
+                finish(nil)
             }
         }
 
diff --git a/ios/Operations/AsyncBlockOperation.swift b/ios/Operations/AsyncBlockOperation.swift
index 1a17683cca..062fc528f5 100644
--- a/ios/Operations/AsyncBlockOperation.swift
+++ b/ios/Operations/AsyncBlockOperation.swift
@@ -11,38 +11,41 @@ import protocol MullvadTypes.Cancellable
 
 /// Asynchronous block operation
 public class AsyncBlockOperation: AsyncOperation {
-    public typealias ExecutionBlock = (AsyncBlockOperation) -> Void
-
-    private var executionBlock: ExecutionBlock?
+    private var executor: ((@escaping (Error?) -> Void) -> Cancellable?)?
     private var cancellableTask: Cancellable?
 
-    public init(dispatchQueue: DispatchQueue? = nil, block: @escaping ExecutionBlock) {
-        executionBlock = block
+    public init(dispatchQueue: DispatchQueue? = nil, block: @escaping (@escaping (Error?) -> Void) -> Void) {
         super.init(dispatchQueue: dispatchQueue)
+        executor = { finish in
+            block(finish)
+            return nil
+        }
     }
 
-    public convenience init(dispatchQueue: DispatchQueue? = nil, block: @escaping () -> Void) {
-        self.init(dispatchQueue: dispatchQueue, block: { operation in
+    public init(dispatchQueue: DispatchQueue? = nil, block: @escaping () -> Void) {
+        super.init(dispatchQueue: dispatchQueue)
+        executor = { finish in
             block()
-            operation.finish()
-        })
+            finish(nil)
+            return nil
+        }
     }
 
-    public convenience init(
+    public init(
         dispatchQueue: DispatchQueue? = nil,
-        cancellableTask: @escaping (AsyncBlockOperation) -> Cancellable
+        cancellableTask: @escaping (@escaping (Error?) -> Void) -> Cancellable
     ) {
-        self.init(dispatchQueue: dispatchQueue, block: { operation in
-            operation.cancellableTask = cancellableTask(operation)
-        })
+        super.init(dispatchQueue: dispatchQueue)
+        executor = { cancellableTask($0) }
     }
 
     override public func main() {
-        let block = executionBlock
-        executionBlock = nil
+        let executor = executor
+        self.executor = nil
+
+        assert(executor != nil)
 
-        assert(block != nil)
-        block?(self)
+        cancellableTask = executor?(self.finish)
     }
 
     override public func operationDidCancel() {
@@ -50,7 +53,7 @@ public class AsyncBlockOperation: AsyncOperation {
     }
 
     override public func operationDidFinish() {
-        executionBlock = nil
+        executor = nil
         cancellableTask = nil
     }
 }
diff --git a/ios/Operations/ResultBlockOperation.swift b/ios/Operations/ResultBlockOperation.swift
index 3972262980..867dee45cf 100644
--- a/ios/Operations/ResultBlockOperation.swift
+++ b/ios/Operations/ResultBlockOperation.swift
@@ -10,66 +10,42 @@ import Foundation
 import protocol MullvadTypes.Cancellable
 
 public final class ResultBlockOperation<Success>: ResultOperation<Success> {
-    public typealias ExecutionBlock = (ResultBlockOperation<Success>) -> Void
-
-    private var executionBlock: ExecutionBlock?
+    private var executor: ((@escaping (Result<Success, Error>) -> Void) -> Cancellable?)?
     private var cancellableTask: Cancellable?
 
-    public convenience init(dispatchQueue: DispatchQueue? = nil, executionBlock: @escaping ExecutionBlock) {
-        self.init(
-            dispatchQueue: dispatchQueue,
-            executionBlock: executionBlock,
-            completionQueue: nil,
-            completionHandler: nil
-        )
-    }
-
-    public convenience init(dispatchQueue: DispatchQueue? = nil, executionBlock: @escaping () throws -> Success) {
-        self.init(
-            dispatchQueue: dispatchQueue,
-            executionBlock: { operation in
-                operation.finish(result: Result { try executionBlock() })
-            },
-            completionQueue: nil,
-            completionHandler: nil
-        )
-    }
-
-    public convenience init(
+    public init(
         dispatchQueue: DispatchQueue? = nil,
-        cancellableTask: @escaping (ResultBlockOperation<Success>) -> Cancellable
+        executionBlock: @escaping (_ finish: @escaping (Result<Success, Error>) -> Void) -> Void
     ) {
-        self.init(
-            dispatchQueue: dispatchQueue,
-            executionBlock: { operation in
-                operation.cancellableTask = cancellableTask(operation)
-            },
-            completionQueue: nil,
-            completionHandler: nil
-        )
+        super.init(dispatchQueue: dispatchQueue)
+        executor = { finish in
+            executionBlock(finish)
+            return nil
+        }
+    }
+
+    public init(dispatchQueue: DispatchQueue? = nil, executionBlock: @escaping () throws -> Success) {
+        super.init(dispatchQueue: dispatchQueue)
+        executor = { finish in
+            finish(Result { try executionBlock() })
+            return nil
+        }
     }
 
     public init(
-        dispatchQueue: DispatchQueue?,
-        executionBlock: @escaping ExecutionBlock,
-        completionQueue: DispatchQueue?,
-        completionHandler: CompletionHandler?
+        dispatchQueue: DispatchQueue? = nil,
+        cancellableTask: @escaping (_ finish: @escaping (Result<Success, Error>) -> Void) -> Cancellable
     ) {
-        self.executionBlock = executionBlock
-
-        super.init(
-            dispatchQueue: dispatchQueue,
-            completionQueue: completionQueue,
-            completionHandler: completionHandler
-        )
+        super.init(dispatchQueue: dispatchQueue)
+        executor = { cancellableTask($0) }
     }
 
     override public func main() {
-        let block = executionBlock
-        executionBlock = nil
+        let executor = executor
+        self.executor = nil
 
-        assert(block != nil)
-        block?(self)
+        assert(executor != nil)
+        cancellableTask = executor?(self.finish)
     }
 
     override public func operationDidCancel() {
@@ -77,7 +53,7 @@ public final class ResultBlockOperation<Success>: ResultOperation<Success> {
     }
 
     override public func operationDidFinish() {
-        executionBlock = nil
+        executor = nil
         cancellableTask = nil
     }
 }
diff --git a/ios/Operations/TransformOperation.swift b/ios/Operations/TransformOperation.swift
index 51484563c1..b5cc36adad 100644
--- a/ios/Operations/TransformOperation.swift
+++ b/ios/Operations/TransformOperation.swift
@@ -10,7 +10,6 @@ import Foundation
 import protocol MullvadTypes.Cancellable
 
 public final class TransformOperation<Input, Output>: ResultOperation<Output>, InputOperation {
-    public typealias ExecutionBlock = (Input, TransformOperation<Input, Output>) -> Void
     public typealias InputBlock = () -> Input?
 
     private let nslock = NSLock()
@@ -35,34 +34,43 @@ public final class TransformOperation<Input, Output>: ResultOperation<Output>, I
 
     private var inputBlock: InputBlock?
 
-    private var executionBlock: ExecutionBlock?
+    private var executor: ((Input, @escaping (Result<Output, Error>) -> Void) -> Cancellable?)?
     private var cancellableTask: Cancellable?
 
-    public init(dispatchQueue: DispatchQueue? = nil, input: Input? = nil, block: @escaping ExecutionBlock) {
-        __input = input
-        executionBlock = block
-
+    public init(
+        dispatchQueue: DispatchQueue? = nil,
+        input: Input? = nil,
+        block: @escaping (_ input: Input, _ finish: @escaping (Result<Output, Error>) -> Void) -> Void
+    ) {
         super.init(dispatchQueue: dispatchQueue)
+        __input = input
+        executor = { input, finish in
+            block(input, finish)
+            return nil
+        }
     }
 
-    public convenience init(
+    public init(
         dispatchQueue: DispatchQueue? = nil,
         input: Input? = nil,
-        throwingBlock: @escaping (Input) throws -> Output
+        throwingBlock: @escaping (_ input: Input) throws -> Output
     ) {
-        self.init(dispatchQueue: dispatchQueue, input: input, block: { input, operation in
-            operation.finish(result: Result { try throwingBlock(input) })
-        })
+        super.init(dispatchQueue: dispatchQueue)
+        __input = input
+        executor = { input, finish in
+            finish(Result { try throwingBlock(input) })
+            return nil
+        }
     }
 
-    public convenience init(
+    public init(
         dispatchQueue: DispatchQueue? = nil,
         input: Input? = nil,
-        cancellableTask: @escaping (Input, TransformOperation<Input, Output>) -> Cancellable
+        cancellableTask: @escaping (_ input: Input, _ finish: @escaping (Result<Output, Error>) -> Void) -> Cancellable
     ) {
-        self.init(dispatchQueue: dispatchQueue, input: input, block: { input, operation in
-            operation.cancellableTask = cancellableTask(input, operation)
-        })
+        super.init(dispatchQueue: dispatchQueue)
+        __input = input
+        executor = cancellableTask
     }
 
     override public func main() {
@@ -70,12 +78,17 @@ public final class TransformOperation<Input, Output>: ResultOperation<Output>, I
             _input = inputBlock()
         }
 
-        guard let inputValue = _input, let executionBlock = executionBlock else {
+        guard let inputValue = _input else {
             finish(result: .failure(OperationError.unsatisfiedRequirement))
             return
         }
 
-        executionBlock(inputValue, self)
+        let executor = executor
+        self.executor = nil
+
+        assert(executor != nil)
+
+        cancellableTask = executor?(inputValue, self.finish)
     }
 
     override public func operationDidCancel() {
@@ -83,7 +96,7 @@ public final class TransformOperation<Input, Output>: ResultOperation<Output>, I
     }
 
     override public func operationDidFinish() {
-        executionBlock = nil
+        executor = nil
         cancellableTask = nil
     }
 
diff --git a/ios/OperationsTests/AsyncBlockOperationTests.swift b/ios/OperationsTests/AsyncBlockOperationTests.swift
index 026f8db911..07a65d0ccd 100644
--- a/ios/OperationsTests/AsyncBlockOperationTests.swift
+++ b/ios/OperationsTests/AsyncBlockOperationTests.swift
@@ -17,9 +17,9 @@ final class AsyncBlockOperationTests: XCTestCase {
         let executionExpectation = expectation(description: "Should execute")
         let finishExpectation = expectation(description: "Should finish")
 
-        let operation = AsyncBlockOperation(block: { op in
+        let operation = AsyncBlockOperation(block: { finish in
             executionExpectation.fulfill()
-            op.finish()
+            finish(nil)
         })
 
         operation.completionBlock = {
@@ -53,15 +53,14 @@ final class AsyncBlockOperationTests: XCTestCase {
         let cancelExpectation = expectation(description: "Should cancel")
         let finishExpectation = expectation(description: "Should finish")
 
-        let operation = AsyncBlockOperation(cancellableTask: { op in
+        let operation = AsyncBlockOperation { finish -> Cancellable in
             executionExpectation.fulfill()
 
             return AnyCancellable {
                 cancelExpectation.fulfill()
-
-                op.finish()
+                finish(nil)
             }
-        })
+        }
 
         operation.completionBlock = {
             finishExpectation.fulfill()
diff --git a/ios/OperationsTests/AsyncResultBlockOperationTests.swift b/ios/OperationsTests/AsyncResultBlockOperationTests.swift
index 160e7b2580..5bae7f7d9c 100644
--- a/ios/OperationsTests/AsyncResultBlockOperationTests.swift
+++ b/ios/OperationsTests/AsyncResultBlockOperationTests.swift
@@ -16,8 +16,8 @@ final class AsyncResultBlockOperationTests: XCTestCase {
     func testBlockOperation() {
         let expectation = expectation(description: "Should finish")
 
-        let operation = ResultBlockOperation<Bool> { op in
-            op.finish(result: .success(true))
+        let operation = ResultBlockOperation<Bool> { finish in
+            finish(.success(true))
         }
 
         operation.onFinish { op, error in
@@ -52,11 +52,11 @@ final class AsyncResultBlockOperationTests: XCTestCase {
     func testCancellableTaskOperation() {
         let expectation = expectation(description: "Should finish")
 
-        let operation = ResultBlockOperation<Bool>(cancellableTask: { op in
+        let operation = ResultBlockOperation<Bool> { finish -> Cancellable in
             return AnyCancellable {
-                op.finish(result: .failure(URLError(.cancelled)))
+                finish(.failure(URLError(.cancelled)))
             }
-        })
+        }
 
         operation.onStart { op in
             op.cancel()
diff --git a/ios/OperationsTests/OperationConditionTests.swift b/ios/OperationsTests/OperationConditionTests.swift
index b35e687319..27898ea424 100644
--- a/ios/OperationsTests/OperationConditionTests.swift
+++ b/ios/OperationsTests/OperationConditionTests.swift
@@ -123,10 +123,10 @@ class OperationConditionTests: XCTestCase {
         let exclusiveCategory = "exclusiveOperations"
         let operationQueue = AsyncOperationQueue()
 
-        let firstOperation = AsyncBlockOperation { op in
+        let firstOperation = AsyncBlockOperation { finish in
             DispatchQueue.main.asyncAfter(deadline: .now() + .seconds(1)) {
                 expectFirstOperationExecution.fulfill()
-                op.finish()
+                finish(nil)
             }
         }
         firstOperation.addCondition(MutuallyExclusive(category: exclusiveCategory))
diff --git a/ios/OperationsTests/TransformOperationTests.swift b/ios/OperationsTests/TransformOperationTests.swift
index 19e932602b..c16bcd3be1 100644
--- a/ios/OperationsTests/TransformOperationTests.swift
+++ b/ios/OperationsTests/TransformOperationTests.swift
@@ -16,8 +16,8 @@ final class TransformOperationTests: XCTestCase {
     func testBlockTransformOperation() {
         let finishExpectation = expectation(description: "Should finish")
 
-        let transform = TransformOperation(input: Int.zero) { input, op in
-            op.finish(result: .success(input + 1))
+        let transform = TransformOperation(input: Int.zero) { input, finish in
+            finish(.success(input + 1))
         }
 
         transform.onFinish { op, error in
@@ -52,11 +52,11 @@ final class TransformOperationTests: XCTestCase {
     func testCancellableTaskBlockTransformOperation() {
         let finishExpectation = expectation(description: "Should finish")
 
-        let transform = TransformOperation<Int, Int>(input: Int.zero, cancellableTask: { _, op in
+        let transform = TransformOperation<Int, Int>(input: Int.zero) { _, finish -> Cancellable in
             return AnyCancellable {
-                op.finish(result: .failure(URLError(.cancelled)))
+                finish(.failure(URLError(.cancelled)))
             }
-        })
+        }
 
         transform.onStart { op in
             op.cancel()
@@ -76,8 +76,8 @@ final class TransformOperationTests: XCTestCase {
     func testShouldFailWithUnsatisfiedRequirement() {
         let finishExpectation = expectation(description: "Should finish")
 
-        let transform = TransformOperation<Int, Int> { input, op in
-            op.finish(result: .success(input))
+        let transform = TransformOperation<Int, Int> { input, finish in
+            finish(.success(input))
         }
 
         transform.onFinish { _, error in
diff --git a/ios/PacketTunnel/PacketTunnelProvider.swift b/ios/PacketTunnel/PacketTunnelProvider.swift
index 614626b2e8..c70d5e0d35 100644
--- a/ios/PacketTunnel/PacketTunnelProvider.swift
+++ b/ios/PacketTunnel/PacketTunnelProvider.swift
@@ -585,14 +585,14 @@ class PacketTunnelProvider: NEPacketTunnelProvider, TunnelMonitorDelegate {
         // Ignore all requests to reconnect once tunnel is preparing to stop.
         guard !isStopping else { return }
 
-        let blockOperation = AsyncBlockOperation(dispatchQueue: dispatchQueue, block: { operation in
+        let blockOperation = AsyncBlockOperation(dispatchQueue: dispatchQueue, block: { finish in
             if shouldStopTunnelMonitor {
                 self.tunnelMonitor.stop()
             }
 
             self.reconnectTunnelInner(to: nextRelay) { error in
                 completionHandler?(error)
-                operation.finish()
+                finish(nil)
             }
         })
 
@@ -780,26 +780,27 @@ class PacketTunnelProvider: NEPacketTunnelProvider, TunnelMonitorDelegate {
     }
 
     private func createGetAccountDataOperation(accountNumber: String) -> ResultOperation<REST.AccountData> {
-        return ResultBlockOperation<REST.AccountData>(dispatchQueue: dispatchQueue, cancellableTask: { operation in
-            return self.accountsProxy.getAccountData(accountNumber: accountNumber, retryStrategy: .noRetry) { result in
-                operation.finish(result: result)
-            }
-        })
+        return ResultBlockOperation<REST.AccountData>(dispatchQueue: dispatchQueue) { finish -> Cancellable in
+            return self.accountsProxy.getAccountData(
+                accountNumber: accountNumber,
+                retryStrategy: .noRetry,
+                completion: finish
+            )
+        }
     }
 
     private func createGetDeviceDataOperation(
         accountNumber: String,
         identifier: String
     ) -> ResultOperation<REST.Device> {
-        return ResultBlockOperation<REST.Device>(dispatchQueue: dispatchQueue, cancellableTask: { operation in
+        return ResultBlockOperation<REST.Device>(dispatchQueue: dispatchQueue) { finish -> Cancellable in
             return self.devicesProxy.getDevice(
                 accountNumber: accountNumber,
                 identifier: identifier,
-                retryStrategy: .noRetry
-            ) { result in
-                operation.finish(result: result)
-            }
-        })
+                retryStrategy: .noRetry,
+                completion: finish
+            )
+        }
     }
 }