Operations: add convenience initializer that returns cancellable object

3 files changed, 59 insertions, 143 deletions

diff --git a/ios/Operations/AsyncBlockOperation.swift b/ios/Operations/AsyncBlockOperation.swift
index 43be9b8275..c8889bc9b5 100644
--- a/ios/Operations/AsyncBlockOperation.swift
+++ b/ios/Operations/AsyncBlockOperation.swift
@@ -7,30 +7,34 @@
 //
 
 import Foundation
+import protocol MullvadTypes.Cancellable
 
 /// Asynchronous block operation
 public class AsyncBlockOperation: AsyncOperation {
-    private var executionBlock: ((AsyncBlockOperation) -> Void)?
-    private var cancellationBlocks: [() -> Void] = []
+    public typealias ExecutionBlock = (AsyncBlockOperation) -> Void
 
-    override public init(dispatchQueue: DispatchQueue? = nil) {
-        super.init(dispatchQueue: dispatchQueue)
-    }
+    private var executionBlock: ExecutionBlock?
+    private var cancellableTask: Cancellable?
 
-    public init(
-        dispatchQueue: DispatchQueue? = nil,
-        block: @escaping (AsyncBlockOperation) -> Void
-    ) {
+    public init(dispatchQueue: DispatchQueue? = nil, block: @escaping ExecutionBlock) {
         executionBlock = block
         super.init(dispatchQueue: dispatchQueue)
     }
 
-    public init(dispatchQueue: DispatchQueue? = nil, block: @escaping () -> Void) {
-        executionBlock = { operation in
+    public convenience init(dispatchQueue: DispatchQueue? = nil, block: @escaping () -> Void) {
+        self.init(dispatchQueue: dispatchQueue, block: { operation in
             block()
             operation.finish()
-        }
-        super.init(dispatchQueue: dispatchQueue)
+        })
+    }
+
+    public convenience init(
+        dispatchQueue: DispatchQueue? = nil,
+        cancellableTask: @escaping (AsyncBlockOperation) -> Cancellable
+    ) {
+        self.init(dispatchQueue: dispatchQueue, block: { operation in
+            operation.cancellableTask = cancellableTask(operation)
+        })
     }
 
     override public func main() {
@@ -45,40 +49,11 @@ public class AsyncBlockOperation: AsyncOperation {
     }
 
     override public func operationDidCancel() {
-        let blocks = cancellationBlocks
-        cancellationBlocks.removeAll()
-
-        for block in blocks {
-            block()
-        }
+        cancellableTask?.cancel()
     }
 
     override public func operationDidFinish() {
-        cancellationBlocks.removeAll()
         executionBlock = nil
-    }
-
-    public func setExecutionBlock(_ block: @escaping (AsyncBlockOperation) -> Void) {
-        dispatchQueue.async {
-            assert(!self.isExecuting && !self.isFinished)
-            self.executionBlock = block
-        }
-    }
-
-    public func setExecutionBlock(_ block: @escaping () -> Void) {
-        setExecutionBlock { operation in
-            block()
-            operation.finish()
-        }
-    }
-
-    public func addCancellationBlock(_ block: @escaping () -> Void) {
-        dispatchQueue.async {
-            if self.isCancelled, self.isExecuting {
-                block()
-            } else {
-                self.cancellationBlocks.append(block)
-            }
-        }
+        cancellableTask = nil
     }
 }
diff --git a/ios/Operations/ResultBlockOperation.swift b/ios/Operations/ResultBlockOperation.swift
index 5d848b438d..e1a037e4db 100644
--- a/ios/Operations/ResultBlockOperation.swift
+++ b/ios/Operations/ResultBlockOperation.swift
@@ -7,18 +7,15 @@
 //
 
 import Foundation
+import protocol MullvadTypes.Cancellable
 
 public final class ResultBlockOperation<Success>: ResultOperation<Success> {
     public typealias ExecutionBlock = (ResultBlockOperation<Success>) -> Void
-    public typealias ThrowingExecutionBlock = () throws -> Success
 
     private var executionBlock: ExecutionBlock?
-    private var cancellationBlocks: [() -> Void] = []
+    private var cancellableTask: Cancellable?
 
-    public convenience init(
-        dispatchQueue: DispatchQueue? = nil,
-        executionBlock: ExecutionBlock? = nil
-    ) {
+    public convenience init(dispatchQueue: DispatchQueue? = nil, executionBlock: @escaping ExecutionBlock) {
         self.init(
             dispatchQueue: dispatchQueue,
             executionBlock: executionBlock,
@@ -27,13 +24,26 @@ public final class ResultBlockOperation<Success>: ResultOperation<Success> {
         )
     }
 
+    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(
         dispatchQueue: DispatchQueue? = nil,
-        executionBlock: @escaping ThrowingExecutionBlock
+        cancellableTask: @escaping (ResultBlockOperation<Success>) -> Cancellable
     ) {
         self.init(
             dispatchQueue: dispatchQueue,
-            executionBlock: Self.wrapThrowingBlock(executionBlock),
+            executionBlock: { operation in
+                operation.cancellableTask = cancellableTask(operation)
+            },
             completionQueue: nil,
             completionHandler: nil
         )
@@ -41,7 +51,7 @@ public final class ResultBlockOperation<Success>: ResultOperation<Success> {
 
     public init(
         dispatchQueue: DispatchQueue?,
-        executionBlock: ExecutionBlock?,
+        executionBlock: @escaping ExecutionBlock,
         completionQueue: DispatchQueue?,
         completionHandler: CompletionHandler?
     ) {
@@ -62,47 +72,11 @@ public final class ResultBlockOperation<Success>: ResultOperation<Success> {
     }
 
     override public func operationDidCancel() {
-        let blocks = cancellationBlocks
-        cancellationBlocks.removeAll()
-
-        for block in blocks {
-            block()
-        }
+        cancellableTask?.cancel()
     }
 
     override public func operationDidFinish() {
-        cancellationBlocks.removeAll()
         executionBlock = nil
-    }
-
-    public func setExecutionBlock(_ block: @escaping ExecutionBlock) {
-        dispatchQueue.async {
-            assert(!self.isExecuting && !self.isFinished)
-            self.executionBlock = block
-        }
-    }
-
-    public func setExecutionBlock(_ block: @escaping ThrowingExecutionBlock) {
-        setExecutionBlock(Self.wrapThrowingBlock(block))
-    }
-
-    public func addCancellationBlock(_ block: @escaping () -> Void) {
-        dispatchQueue.async {
-            if self.isCancelled, self.isExecuting {
-                block()
-            } else {
-                self.cancellationBlocks.append(block)
-            }
-        }
-    }
-
-    private class func wrapThrowingBlock(_ executionBlock: @escaping ThrowingExecutionBlock)
-        -> ExecutionBlock
-    {
-        return { operation in
-            let result = Result { try executionBlock() }
-
-            operation.finish(result: result)
-        }
+        cancellableTask = nil
     }
 }
diff --git a/ios/Operations/TransformOperation.swift b/ios/Operations/TransformOperation.swift
index 609e07896c..49361ac641 100644
--- a/ios/Operations/TransformOperation.swift
+++ b/ios/Operations/TransformOperation.swift
@@ -7,10 +7,10 @@
 //
 
 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 ThrowingExecutionBlock = (Input) throws -> Output
     public typealias InputBlock = () -> Input?
 
     private let nslock = NSLock()
@@ -36,28 +36,33 @@ public final class TransformOperation<Input, Output>: ResultOperation<Output>, I
     private var inputBlock: InputBlock?
 
     private var executionBlock: ExecutionBlock?
-    private var cancellationBlocks: [() -> Void] = []
+    private var cancellableTask: Cancellable?
 
-    public init(
-        dispatchQueue: DispatchQueue? = nil,
-        input: Input? = nil,
-        block: ExecutionBlock? = nil
-    ) {
+    public init(dispatchQueue: DispatchQueue? = nil, input: Input? = nil, block: ExecutionBlock? = nil) {
         __input = input
         executionBlock = block
 
         super.init(dispatchQueue: dispatchQueue)
     }
 
-    public init(
+    public convenience init(
         dispatchQueue: DispatchQueue? = nil,
         input: Input? = nil,
-        throwingBlock: @escaping ThrowingExecutionBlock
+        throwingBlock: @escaping (Input) throws -> Output
     ) {
-        __input = input
-        executionBlock = Self.wrapThrowingBlock(throwingBlock)
+        self.init(dispatchQueue: dispatchQueue, input: input, block: { input, operation in
+            operation.finish(result: Result { try throwingBlock(input) })
+        })
+    }
 
-        super.init(dispatchQueue: dispatchQueue)
+    public convenience init(
+        dispatchQueue: DispatchQueue? = nil,
+        input: Input? = nil,
+        cancellableTask: @escaping (Input, TransformOperation<Input, Output>) -> Cancellable
+    ) {
+        self.init(dispatchQueue: dispatchQueue, input: input, block: { input, operation in
+            operation.cancellableTask = cancellableTask(input, operation)
+        })
     }
 
     override public func main() {
@@ -74,40 +79,12 @@ public final class TransformOperation<Input, Output>: ResultOperation<Output>, I
     }
 
     override public func operationDidCancel() {
-        let blocks = cancellationBlocks
-        cancellationBlocks.removeAll()
-
-        for block in blocks {
-            block()
-        }
+        cancellableTask?.cancel()
     }
 
     override public func operationDidFinish() {
-        cancellationBlocks.removeAll()
         executionBlock = nil
-    }
-
-    // MARK: - Block handlers
-
-    public func setExecutionBlock(_ block: @escaping ExecutionBlock) {
-        dispatchQueue.async {
-            assert(!self.isExecuting && !self.isFinished)
-            self.executionBlock = block
-        }
-    }
-
-    public func setExecutionBlock(_ block: @escaping ThrowingExecutionBlock) {
-        setExecutionBlock(Self.wrapThrowingBlock(block))
-    }
-
-    public func addCancellationBlock(_ block: @escaping () -> Void) {
-        dispatchQueue.async {
-            if self.isCancelled, self.isExecuting {
-                block()
-            } else {
-                self.cancellationBlocks.append(block)
-            }
-        }
+        cancellableTask = nil
     }
 
     // MARK: - Input injection
@@ -117,14 +94,4 @@ public final class TransformOperation<Input, Output>: ResultOperation<Output>, I
             self.inputBlock = block
         }
     }
-
-    private class func wrapThrowingBlock(_ executionBlock: @escaping ThrowingExecutionBlock)
-        -> ExecutionBlock
-    {
-        return { input, operation in
-            let result = Result { try executionBlock(input) }
-
-            operation.finish(result: result)
-        }
-    }
 }