UpdateAddressCacheOperation: subclass from ResultOperation

3 files changed, 64 insertions, 73 deletions

diff --git a/ios/MullvadVPN/AddressCache/AddressCacheTracker.swift b/ios/MullvadVPN/AddressCache/AddressCacheTracker.swift
index 8502d12ffc..de84265b1e 100644
--- a/ios/MullvadVPN/AddressCache/AddressCacheTracker.swift
+++ b/ios/MullvadVPN/AddressCache/AddressCacheTracker.swift
@@ -83,16 +83,16 @@ extension AddressCache {
             }
         }
 
-        func updateEndpoints(completionHandler: ((_ result: CacheUpdateResult) -> Void)? = nil) -> Cancellable {
+        func updateEndpoints(completionHandler: ((_ completion: OperationCompletion<CacheUpdateResult, Error>) -> Void)? = nil) -> Cancellable {
             let operation = UpdateAddressCacheOperation(
                 queue: stateQueue,
                 restClient: restClient,
                 store: store,
                 updateInterval: Self.updateInterval,
-                completionHandler: { [weak self] result in
-                    self?.handleCacheUpdateResult(result)
+                completionHandler: { [weak self] completion in
+                    self?.handleCacheUpdateCompletion(completion)
 
-                    completionHandler?(result)
+                    completionHandler?(completion)
                 }
             )
 
@@ -144,20 +144,22 @@ extension AddressCache {
             }
         }
 
-        private func handleCacheUpdateResult(_ result: AddressCache.CacheUpdateResult) {
-            switch result {
-            case .success:
-                logger.debug("Finished updating address cache")
+        private func handleCacheUpdateCompletion(_ completion: OperationCompletion<AddressCache.CacheUpdateResult, Error>) {
+            switch completion {
+            case .success(let updateResult):
+                switch updateResult {
+                case .finished:
+                    logger.debug("Finished updating address cache")
+                case .throttled:
+                    logger.debug("Address cache update was throttled")
+                }
+
                 lastFailureAttemptDate = nil
 
             case .failure(let error):
                 logger.error(chainedError: AnyChainedError(error), message: "Failed to update address cache")
                 lastFailureAttemptDate = Date()
 
-            case .throttled:
-                logger.debug("Address cache update was throttled")
-                lastFailureAttemptDate = nil
-
             case .cancelled:
                 logger.debug("Address cache update was cancelled")
                 lastFailureAttemptDate = Date()
diff --git a/ios/MullvadVPN/AddressCache/UpdateAddressCacheOperation.swift b/ios/MullvadVPN/AddressCache/UpdateAddressCacheOperation.swift
index 8fcd41b7dd..b66f89bb3b 100644
--- a/ios/MullvadVPN/AddressCache/UpdateAddressCacheOperation.swift
+++ b/ios/MullvadVPN/AddressCache/UpdateAddressCacheOperation.swift
@@ -11,63 +11,48 @@ import Foundation
 extension AddressCache {
 
     enum CacheUpdateResult {
-        /// Operation was cancelled.
-        case cancelled
-
         /// Address cache update was throttled as it was requested too early.
         case throttled(_ lastUpdateDate: Date)
 
-        /// Failure to update address cache.
-        case failure(Error)
-
         /// Address cache is successfully updated.
-        case success
-
-        var isTaskCompleted: Bool {
-            switch self {
-            case .cancelled, .failure:
-                return false
-            case .success, .throttled:
-                return true
-            }
-        }
+        case finished
     }
 
-    class UpdateAddressCacheOperation: AsyncOperation {
-        typealias CompletionHandler = (_ result: CacheUpdateResult) -> Void
-
+    class UpdateAddressCacheOperation: ResultOperation<CacheUpdateResult, Error> {
         private let queue: DispatchQueue
         private let restClient: REST.Client
         private let store: AddressCache.Store
         private let updateInterval: TimeInterval
 
-        private var completionHandler: CompletionHandler?
-        private var restCancellationHandle: Cancellable?
+        private var requestTask: Cancellable?
 
         init(queue: DispatchQueue, restClient: REST.Client, store: AddressCache.Store, updateInterval: TimeInterval, completionHandler: CompletionHandler?) {
             self.queue = queue
             self.restClient = restClient
             self.store = store
             self.updateInterval = updateInterval
-            self.completionHandler = completionHandler
+
+            super.init(completionQueue: queue, completionHandler: completionHandler)
         }
 
-        override func cancel() {
+        override func main() {
             queue.async {
-                super.cancel()
-                self.restCancellationHandle?.cancel()
+                self.startUpdate()
             }
         }
 
-        override func main() {
+        override func cancel() {
+            super.cancel()
+
             queue.async {
-                self.startUpdate()
+                self.requestTask?.cancel()
+                self.requestTask = nil
             }
         }
 
         private func startUpdate() {
             guard !isCancelled else {
-                completeOperation(with: .cancelled)
+                finish(completion: .cancelled)
                 return
             }
 
@@ -75,40 +60,46 @@ extension AddressCache {
             let nextUpdate = Date(timeInterval: updateInterval, since: lastUpdate)
 
             guard nextUpdate <= Date() else {
-                completeOperation(with: .throttled(lastUpdate))
+                finish(completion: .success(.throttled(lastUpdate)))
                 return
             }
 
-            restCancellationHandle = restClient.getAddressList(retryStrategy: .default) { restResult in
-                    self.queue.async {
-                        switch restResult {
-                        case .success(let newEndpoints):
-                            self.store.setEndpoints(newEndpoints) { error in
-                                self.queue.async {
-                                    if let error = error {
-                                        self.completeOperation(with: .failure(error))
-                                    } else {
-                                        self.completeOperation(with: .success)
-                                    }
-                                }
-                            }
+            requestTask = restClient.getAddressList(retryStrategy: .default) { result in
+                self.queue.async {
+                    self.handleResponse(result)
+                }
+            }
+        }
 
-                        case .failure(let error):
-                            if case URLError.cancelled = error {
-                                self.completeOperation(with: .cancelled)
-                            } else {
-                                self.completeOperation(with: .failure(error))
-                            }
-                        }
+        private func handleResponse(_ result: Result<[AnyIPEndpoint], REST.Error>) {
+            switch result {
+            case .success(let newEndpoints):
+                self.store.setEndpoints(newEndpoints) { error in
+                    if let error = error {
+                        self.finish(completion: .failure(error))
+                    } else {
+                        self.finish(completion: .success(.finished))
                     }
                 }
-        }
 
-        private func completeOperation(with result: CacheUpdateResult) {
-            completionHandler?(result)
-            completionHandler = nil
+            case .failure(let error):
+                if case URLError.cancelled = error {
+                    self.finish(completion: .cancelled)
+                } else {
+                    self.finish(completion: .failure(error))
+                }
+            }
+        }
+    }
+}
 
-            finish()
+extension OperationCompletion where Success == AddressCache.CacheUpdateResult {
+    var isTaskCompleted: Bool {
+        switch self {
+        case .success:
+            return true
+        case .cancelled, .failure:
+            return false
         }
     }
 }
diff --git a/ios/MullvadVPN/Result+UIBackgroundFetchResult.swift b/ios/MullvadVPN/Result+UIBackgroundFetchResult.swift
index 0727d961cd..40aad31c70 100644
--- a/ios/MullvadVPN/Result+UIBackgroundFetchResult.swift
+++ b/ios/MullvadVPN/Result+UIBackgroundFetchResult.swift
@@ -8,17 +8,15 @@
 
 import UIKit
 
-extension AddressCache.CacheUpdateResult {
+extension OperationCompletion where Success == AddressCache.CacheUpdateResult {
     var backgroundFetchResult: UIBackgroundFetchResult {
         switch self {
-        case .failure:
-            return .failed
-        case .throttled:
-            return .noData
-        case .success:
+        case .success(.finished):
             return .newData
-        case .cancelled:
+        case .success(.throttled), .cancelled:
             return .noData
+        case .failure:
+            return .failed
         }
     }
 }