Promise: make receive(on:) calls cancellable and add timer type

1 files changed, 29 insertions, 5 deletions

diff --git a/ios/MullvadVPN/Promise/Promise+ReceiveOn.swift b/ios/MullvadVPN/Promise/Promise+ReceiveOn.swift
index a7ce8f039b..6979d5e3fb 100644
--- a/ios/MullvadVPN/Promise/Promise+ReceiveOn.swift
+++ b/ios/MullvadVPN/Promise/Promise+ReceiveOn.swift
@@ -9,24 +9,48 @@
 import Foundation
 
 extension Promise {
+    /// A type of timer.
+    enum TimerType {
+        case deadline
+        case walltime
+    }
+
     /// Dispatch the upstream value on another queue.
     func receive(on queue: DispatchQueue) -> Promise<Value> {
         return Promise<Value> { resolver in
-            _ = self.observe { completion in
-                queue.async {
+            self.observe { completion in
+                let work = DispatchWorkItem {
                     resolver.resolve(completion: completion, queue: queue)
                 }
+
+                resolver.setCancelHandler {
+                    work.cancel()
+                }
+
+                queue.async(execute: work)
             }
         }
     }
 
     /// Dispatch the upstream value on another queue after delay.
-    func receive(on queue: DispatchQueue, after deadline: DispatchTime) -> Promise<Value> {
+    func receive(on queue: DispatchQueue, after timeInterval: DispatchTimeInterval, timerType: TimerType) -> Promise<Value> {
         return Promise<Value> { resolver in
-            _ = self.observe { completion in
-                queue.asyncAfter(deadline: deadline) {
+            self.observe { completion in
+                let work = DispatchWorkItem {
                     resolver.resolve(completion: completion, queue: queue)
                 }
+
+                resolver.setCancelHandler {
+                    work.cancel()
+                }
+
+                switch timerType {
+                case .deadline:
+                    queue.asyncAfter(deadline: .now() + timeInterval, execute: work)
+
+                case .walltime:
+                    queue.asyncAfter(wallDeadline: .now() + timeInterval, execute: work)
+                }
             }
         }
     }