Simplify reusing http request Tasks in ImageDownloader

WordPress/Classes/Utility/Media/ImageDownloader.swift

@@ -30,7 +30,7 @@ actor ImageDownloader {
         )
     }
 
-    private var tasks: [String: ImageDataTask] = [:]
+    private var tasks: [String: Task<Data, any Error>] = [:]
 
     init(cache: MemoryCacheProtocol = MemoryCache.shared) {
         self.cache = cache
@@ -115,32 +115,22 @@ actor ImageDownloader {
 
     private func data(for request: URLRequest, options: ImageRequestOptions) async throws -> Data {
         let requestKey = request.urlRequest?.url?.absoluteString ?? ""
-        let task = tasks[requestKey] ?? ImageDataTask(key: requestKey, Task {
-            try await self._data(for: request, options: options, key: requestKey)
-        })
-        task.downloader = self
 
-        let subscriptionID = UUID()
-        task.subscriptions.insert(subscriptionID)
-        tasks[requestKey] = task
-
-        return try await task.getData(subscriptionID: subscriptionID)
-    }
-
-    fileprivate nonisolated func unsubscribe(_ subscriptionID: UUID, key: String) {
-        Task {
-            await _unsubscribe(subscriptionID, key: key)
+        let theTask: Task<Data, any Error>
+        if let task = tasks[requestKey] {
+            theTask = task
+        } else {
+            theTask = Task {
+                try await self._data(for: request, options: options, key: requestKey)
+            }
+            tasks[requestKey] = theTask
         }
-    }
 
-    private func _unsubscribe(_ subscriptionID: UUID, key: String) {
-        guard let task = tasks[key],
-              task.subscriptions.remove(subscriptionID) != nil,
-              task.subscriptions.isEmpty else {
-            return
+        return try await withTaskCancellationHandler {
+            try await theTask.value
+        } onCancel: {
+            theTask.cancel()
         }
-        task.task.cancel()
-        tasks[key] = nil
     }
 
     private func _data(for request: URLRequest, options: ImageRequestOptions, key: String) async throws -> Data {
@@ -161,27 +151,6 @@ actor ImageDownloader {
     }
 }
 
-private final class ImageDataTask {
-    let key: String
-    var subscriptions = Set<UUID>()
-    let task: Task<Data, Error>
-    weak var downloader: ImageDownloader?
-
-    init(key: String, _ task: Task<Data, Error>) {
-        self.key = key
-        self.task = task
-    }
-
-    func getData(subscriptionID: UUID) async throws -> Data {
-        try await withTaskCancellationHandler {
-            try await task.value
-        } onCancel: { [weak self] in
-            guard let self else { return }
-            self.downloader?.unsubscribe(subscriptionID, key: self.key)
-        }
-    }
-}
-
 // MARK: - ImageDownloader (Closures)
 
 extension ImageDownloader {

WordPress/WordPressTest/ImageDownloaderTests.swift

@@ -125,6 +125,40 @@ class ImageDownloaderTests: CoreDataTestCase {
         XCTAssertEqual(image.size, CGSize(width: 1024, height: 680))
     }
 
+    func testReuseHTTPRequest() async throws {
+        // GIVEN
+        let httpRequestReceived = self.expectation(description: "HTTP request received")
+        let imageURL = try XCTUnwrap(URL(string: "https://example.files.wordpress.com/2023/09/image.jpg"))
+        stub(condition: { _ in true }, response: { _ in
+            httpRequestReceived.fulfill()
+
+            guard let sourceURL = try? XCTUnwrap(Bundle.test.url(forResource: "test-image", withExtension: "jpg")),
+                  let data = try? Data(contentsOf: sourceURL) else {
+                return HTTPStubsResponse(error: URLError(.unknown))
+            }
+
+            return HTTPStubsResponse(data: data, statusCode: 200, headers: nil)
+                .responseTime(0.3)
+        })
+
+        // WHEN
+        let taskCompleted = self.expectation(description: "Image downloaded")
+        taskCompleted.expectedFulfillmentCount = 3
+        for _ in 1...3 {
+            try await Task.sleep(for: .milliseconds(50))
+            Task.detached {
+                do {
+                    _ = try await self.sut.image(from: imageURL)
+                } catch {
+                    XCTFail("Unexpected error: \(error)")
+                }
+                taskCompleted.fulfill()
+            }
+        }
+
+        await fulfillment(of: [httpRequestReceived, taskCompleted], timeout: 0.5)
+    }
+
     // MARK: - Helpers
 
     /// `Media` is hardcoded to work with a specific direcoty URL managed by `MediaFileManager`