Accept queue in OperationBlockObserver

2 files changed, 26 insertions, 7 deletions

diff --git a/ios/MullvadVPN/Operations/OperationBlockObserver.swift b/ios/MullvadVPN/Operations/OperationBlockObserver.swift
index 728e9a5a02..eccb7f63ab 100644
--- a/ios/MullvadVPN/Operations/OperationBlockObserver.swift
+++ b/ios/MullvadVPN/Operations/OperationBlockObserver.swift
@@ -12,22 +12,41 @@ class OperationBlockObserver<OperationType: OperationProtocol>: OperationObserve
     private var willFinish: ((OperationType) -> Void)?
     private var didFinish: ((OperationType) -> Void)?
 
-    init(willFinish: ((OperationType) -> Void)? = nil, didFinish: ((OperationType) -> Void)? = nil) {
+    let queue: DispatchQueue?
+
+    init(queue: DispatchQueue? = nil, willFinish: ((OperationType) -> Void)? = nil, didFinish: ((OperationType) -> Void)? = nil) {
+        self.queue = queue
         self.willFinish = willFinish
         self.didFinish = didFinish
     }
 
     func operationWillFinish(_ operation: OperationType) {
-        self.willFinish?(operation)
+        if let willFinish = self.willFinish {
+            scheduleEvent {
+                willFinish(operation)
+            }
+        }
     }
 
     func operationDidFinish(_ operation: OperationType) {
-        self.didFinish?(operation)
+        if let didFinish = self.didFinish {
+            scheduleEvent {
+                didFinish(operation)
+            }
+        }
+    }
+
+    private func scheduleEvent(_ body: @escaping () -> Void) {
+        if let queue = queue {
+            queue.async(execute: body)
+        } else {
+            body()
+        }
     }
 }
 
 extension OperationProtocol {
-    func addDidFinishBlockObserver(_ block: @escaping (Self) -> Void) {
-        addObserver(OperationBlockObserver(didFinish: block))
+    func addDidFinishBlockObserver(queue: DispatchQueue? = nil, _ block: @escaping (Self) -> Void) {
+        addObserver(OperationBlockObserver(queue: queue, didFinish: block))
     }
 }
diff --git a/ios/MullvadVPN/Operations/OutputOperation.swift b/ios/MullvadVPN/Operations/OutputOperation.swift
index 533d5e5151..064d54a2e9 100644
--- a/ios/MullvadVPN/Operations/OutputOperation.swift
+++ b/ios/MullvadVPN/Operations/OutputOperation.swift
@@ -40,8 +40,8 @@ extension OutputOperation where Self: OperationSubclassing {
 }
 
 extension OperationProtocol where Self: OutputOperation {
-    func addDidFinishBlockObserver(_ block: @escaping (Self, Output) -> Void) {
-        addDidFinishBlockObserver { (operation) in
+    func addDidFinishBlockObserver(queue: DispatchQueue? = nil, _ block: @escaping (Self, Output) -> Void) {
+        addDidFinishBlockObserver(queue: queue) { (operation) in
             if let output = operation.output {
                 block(operation, output)
             }