Concurrency Improvements in ImageDownloadService (#435)

Automattic · Sep 30, 2024 · 517cd73 · 517cd73
2 parents c996f42 + 6380bde
commit 517cd73
Show file tree

Hide file tree

Showing 8 changed files with 213 additions and 104 deletions.
diff --git a/Demo/Demo/Gravatar-UIKit-Demo/DemoUIImageViewExtensionViewController.swift b/Demo/Demo/Gravatar-UIKit-Demo/DemoUIImageViewExtensionViewController.swift
@@ -135,11 +135,12 @@ class DemoUIImageViewExtensionViewController: UIViewController {
 
         present(controller, animated: true)
     }
+    var task: Task<Void, Never>?
 
     @objc private func fetchAvatarButtonHandler() {
         let options = setupOptions()
         let placeholderImage: UIImage? = showPlaceholderSwitchWithLabel.isOn ? UIImage(named: "placeholder") : nil
-        Task {
+        task = Task {
             do {
                 let result = try await avatarImageView.gravatar.setImage(
                     avatarID: .email(emailInputField.text ?? ""),
@@ -150,17 +151,14 @@ class DemoUIImageViewExtensionViewController: UIViewController {
                 print("success!")
                 print("result url: \(result.sourceURL)")
                 print("retrived Image point size: \(result.image.size)")
-
             } catch {
                 print(error)
             }
         }
     }
 
     @objc private func cancelOngoingButtonHandler() {
-        Task {
-            await avatarImageView.gravatar.cancelImageDownload()
-        }
+        task?.cancel()
     }
 
     private func setupOptions() -> [ImageSettingOption] {

diff --git a/Sources/Gravatar/Network/Services/ImageDownloadService.swift b/Sources/Gravatar/Network/Services/ImageDownloadService.swift
@@ -4,7 +4,7 @@ import UIKit
 ///
 /// This is the default type which implements ``ImageDownloader``..
 /// Unless specified otherwise, `ImageDownloadService` will use a `URLSession` based `HTTPClient`, and a in-memory image cache.
-public struct ImageDownloadService: ImageDownloader, Sendable {
+public actor ImageDownloadService: ImageDownloader, Sendable {
     private let client: HTTPClient
     let imageCache: ImageCaching
 
@@ -29,12 +29,20 @@ public struct ImageDownloadService: ImageDownloader, Sendable {
         let request = URLRequest.imageRequest(url: url, forceRefresh: forceRefresh)
 
         if !forceRefresh, let image = try await cachedImage(for: url) {
+            try Task.checkCancellation()
             return ImageDownloadResult(image: image, sourceURL: url)
         }
+
         let task = Task<UIImage, Error> {
             try await fetchAndProcessImage(request: request, processor: processingMethod.processor)
         }
-        let image = try await awaitAndCacheImage(from: task, cacheKey: url.absoluteString)
+
+        // Create `.inProgress` entry before we await to prevent re-entrancy issues
+        let cacheKey = url.absoluteString
+        imageCache.setEntry(.inProgress(task), for: cacheKey)
+
+        let image = try await awaitAndCacheImage(from: task, cacheKey: cacheKey)
+        try Task.checkCancellation()
         return ImageDownloadResult(image: image, sourceURL: url)
     }
 
@@ -43,17 +51,18 @@ public struct ImageDownloadService: ImageDownloader, Sendable {
         switch entry {
         case .inProgress(let task):
             let image = try await task.value
+            try Task.checkCancellation()
             return image
         case .ready(let image):
             return image
         }
     }
 
     private func awaitAndCacheImage(from task: Task<UIImage, Error>, cacheKey key: String) async throws -> UIImage {
-        imageCache.setEntry(.inProgress(task), for: key)
         let image: UIImage
         do {
             image = try await task.value
+            try Task.checkCancellation()
         } catch {
             imageCache.setEntry(nil, for: key)
             throw error
@@ -65,6 +74,7 @@ public struct ImageDownloadService: ImageDownloader, Sendable {
     private func fetchAndProcessImage(request: URLRequest, processor: ImageProcessor) async throws -> UIImage {
         do {
             let (data, _) = try await client.fetchData(with: request)
+            try Task.checkCancellation()
             guard let image = processor.process(data) else {
                 throw ImageFetchingError.imageProcessorFailed
             }
@@ -77,20 +87,6 @@ public struct ImageDownloadService: ImageDownloader, Sendable {
             throw ImageFetchingError.responseError(reason: .unexpected(error))
         }
     }
-
-    public func cancelTask(for url: URL) {
-        if let entry = imageCache.getEntry(with: url.absoluteString) {
-            switch entry {
-            case .inProgress(let task):
-                if !task.isCancelled {
-                    task.cancel()
-                    imageCache.setEntry(nil, for: url.absoluteString)
-                }
-            default:
-                break
-            }
-        }
-    }
 }
 
 extension URLRequest {

diff --git a/Sources/Gravatar/Network/Services/ImageDownloader.swift b/Sources/Gravatar/Network/Services/ImageDownloader.swift
@@ -14,8 +14,4 @@ public protocol ImageDownloader: Sendable {
         forceRefresh: Bool,
         processingMethod: ImageProcessingMethod
     ) async throws -> ImageDownloadResult
-
-    /// Cancels the download task for the given `URL`.
-    /// - Parameter url: `URL` of the download task.
-    func cancelTask(for url: URL)
 }
diff --git a/Sources/GravatarUI/GravatarCompatibleUI/UIImageView+Gravatar.swift b/Sources/GravatarUI/GravatarCompatibleUI/UIImageView+Gravatar.swift
@@ -124,12 +124,6 @@ extension GravatarWrapper where Component: UIImageView {
         }
     }
 
-    public func cancelImageDownload() {
-        if let sourceURL {
-            imageDownloader?.cancelTask(for: sourceURL)
-        }
-    }
-
     /// Downloads the Gravatar profile image and sets it to this UIImageView. Throws ``ImageFetchingComponentError``.
     ///
     /// - Parameters:

diff --git a/Sources/TestHelpers/TestImageCache.swift b/Sources/TestHelpers/TestImageCache.swift
@@ -1,31 +1,63 @@
 import Gravatar
 import UIKit
 
-package class TestImageCache: ImageCaching, @unchecked Sendable {
+package final class TestImageCache: ImageCaching, @unchecked Sendable {
     let imageCache = ImageCache()
 
-    package private(set) var getImageCallsCount = 0
-    package private(set) var setImageCallsCount = 0
-    package private(set) var setTaskCallsCount = 0
+    package typealias CacheMessage = (operation: CacheMessageType, key: String)
+    private var cacheMessages = [CacheMessage]()
+
+    package enum CacheMessageType {
+        case setToNil
+        case inProgress
+        case ready
+        case get
+    }
+
+    package var getImageCallsCount: Int { messageCount(type: .get) }
+    package var setImageCallsCount: Int { messageCount(type: .ready) }
+    package var setTaskCallsCount: Int { messageCount(type: .inProgress) }
+
+    // Serial queue to synchronize access to shared mutable state
+    private let accessQueue = DispatchQueue(label: "com.testImageCache.accessQueue")
 
     package init() {}
 
     package func setEntry(_ entry: Gravatar.CacheEntry?, for key: String) {
-        guard let entry else {
-            imageCache.setEntry(nil, for: key)
-            return
-        }
-        switch entry {
-        case .inProgress:
-            setTaskCallsCount += 1
-        case .ready:
-            setImageCallsCount += 1
+        accessQueue.sync {
+            var message: CacheMessage
+            defer { cacheMessages.append(message) }
+            guard let entry else {
+                imageCache.setEntry(nil, for: key)
+                message = (operation: .setToNil, key: key)
+                return
+            }
+            switch entry {
+            case .inProgress:
+                message = (operation: .inProgress, key: key)
+            case .ready:
+                message = (operation: .ready, key: key)
+            }
+            imageCache.setEntry(entry, for: key)
         }
-        imageCache.setEntry(entry, for: key)
     }
 
     package func getEntry(with key: String) -> Gravatar.CacheEntry? {
-        getImageCallsCount += 1
-        return imageCache.getEntry(with: key)
+        accessQueue.sync {
+            cacheMessages.append(CacheMessage(operation: .get, key: key))
+            return imageCache.getEntry(with: key)
+        }
+    }
+
+    package func messageCount(type: CacheMessageType) -> Int {
+        accessQueue.sync {
+            cacheMessages.filter { $0.operation == type }.count
+        }
+    }
+
+    package func messageCount(type: CacheMessageType, forKey key: String) -> Int {
+        accessQueue.sync {
+            cacheMessages.filter { $0.operation == type && $0.key == key }.count
+        }
     }
 }
diff --git a/Sources/TestHelpers/URLSessionMock.swift b/Sources/TestHelpers/URLSessionMock.swift
@@ -11,7 +11,7 @@ package actor URLSessionMock: URLSessionProtocol {
 
     let returnData: Data
     let response: HTTPURLResponse
-    let error: NSError?
+    private(set) var error: NSError?
     private(set) var isCancellable: Bool = false
     private(set) var maxDurationSeconds: Double = 2
     package private(set) var callsCount = 0
@@ -63,6 +63,10 @@ package actor URLSessionMock: URLSessionProtocol {
         self.request = request
     }
 
+    package func update(error: NSError?) async {
+        self.error = error
+    }
+
     package func update(isCancellable: Bool) async {
         self.isCancellable = isCancellable
     }