ResultOperation: add block based implementation

1 files changed, 109 insertions, 8 deletions

diff --git a/ios/MullvadVPN/Operations/ResultOperation.swift b/ios/MullvadVPN/Operations/ResultOperation.swift
index a0cab6dcc9..fcaeec7091 100644
--- a/ios/MullvadVPN/Operations/ResultOperation.swift
+++ b/ios/MullvadVPN/Operations/ResultOperation.swift
@@ -13,10 +13,10 @@ class ResultOperation<Success, Failure: Error>: AsyncOperation {
     typealias Completion = OperationCompletion<Success, Failure>
     typealias CompletionHandler = (Completion) -> Void
 
-    private let stateLock = NSLock()
+    fileprivate let stateLock = NSLock()
     private var completionValue: Completion?
-    private let completionQueue: DispatchQueue?
-    private var completionHandler: CompletionHandler?
+    private var _completionQueue: DispatchQueue?
+    private var _completionHandler: CompletionHandler?
     private var pendingFinish = false
 
     var completion: Completion? {
@@ -25,9 +25,39 @@ class ResultOperation<Success, Failure: Error>: AsyncOperation {
         return completionValue
     }
 
+    var completionQueue: DispatchQueue? {
+        get {
+            stateLock.lock()
+            defer { stateLock.unlock() }
+
+            return _completionQueue
+        }
+        set {
+            stateLock.lock()
+            _completionQueue = newValue
+            stateLock.unlock()
+        }
+    }
+
+    var completionHandler: CompletionHandler? {
+        get {
+            stateLock.lock()
+            defer { stateLock.unlock() }
+
+            return _completionHandler
+        }
+        set {
+            stateLock.lock()
+            defer { stateLock.unlock() }
+            if !pendingFinish {
+                _completionHandler = newValue
+            }
+        }
+    }
+
     init(completionQueue: DispatchQueue?, completionHandler: CompletionHandler?) {
-        self.completionQueue = completionQueue
-        self.completionHandler = completionHandler
+        _completionQueue = completionQueue
+        _completionHandler = completionHandler
 
         super.init()
     }
@@ -47,7 +77,7 @@ class ResultOperation<Success, Failure: Error>: AsyncOperation {
         _finish()
     }
 
-    private func _finish() {
+    fileprivate func _finish() {
         stateLock.lock()
         // Bail if operation is already finishing.
         guard !pendingFinish else {
@@ -59,13 +89,16 @@ class ResultOperation<Success, Failure: Error>: AsyncOperation {
         pendingFinish = true
 
         // Copy completion handler.
-        let completionHandler = self.completionHandler
+        let completionHandler = _completionHandler
 
         // Unset completion handler.
-        self.completionHandler = nil
+        _completionHandler = nil
 
         // Copy completion value.
         let completion = completionValue ?? .cancelled
+
+        // Copy completion queue.
+        let completionQueue = _completionQueue
         stateLock.unlock()
 
         let block = {
@@ -83,3 +116,71 @@ class ResultOperation<Success, Failure: Error>: AsyncOperation {
         }
     }
 }
+
+class ResultBlockOperation<Success, Failure: Error>: ResultOperation<Success, Failure> {
+    typealias ExecutionBlock = (ResultBlockOperation<Success, Failure>) -> Void
+
+    private var executionBlock: ExecutionBlock?
+    private var cancellationBlocks: [() -> Void] = []
+
+    convenience init(executionBlock: @escaping ExecutionBlock) {
+        self.init(
+            executionBlock: executionBlock,
+            completionQueue: nil,
+            completionHandler: nil
+        )
+    }
+
+    init(
+        executionBlock: @escaping ExecutionBlock,
+        completionQueue: DispatchQueue?,
+        completionHandler: CompletionHandler?
+    )
+    {
+        self.executionBlock = executionBlock
+        super.init(completionQueue: completionQueue, completionHandler: completionHandler)
+    }
+
+    override func main() {
+        stateLock.lock()
+        let block = executionBlock
+        executionBlock = nil
+        stateLock.unlock()
+
+        block?(self)
+    }
+
+    override func cancel() {
+        super.cancel()
+
+        stateLock.lock()
+        let blocks = cancellationBlocks
+        cancellationBlocks.removeAll()
+        stateLock.unlock()
+
+        for block in blocks {
+            block()
+        }
+    }
+
+    override func _finish() {
+        stateLock.lock()
+        cancellationBlocks.removeAll()
+        executionBlock = nil
+        stateLock.unlock()
+
+        super._finish()
+    }
+
+    func addCancellationBlock(_ block: @escaping () -> Void) {
+        stateLock.lock()
+        if isCancelled {
+            stateLock.unlock()
+            block()
+        } else {
+            cancellationBlocks.append(block)
+            stateLock.unlock()
+        }
+    }
+}
+