Fix crash in ImageDownloader

Author: kean

WordPress/Classes/Utility/Media/ImageDownloader.swift

@@ -106,23 +106,25 @@ actor ImageDownloader {
 
     private func data(for request: URLRequest, options: ImageRequestOptions) async throws -> Data {
         let requestKey = request.urlRequest?.url?.absoluteString ?? ""
-        let task = tasks[requestKey] ?? ImageDataTask(task: Task {
+        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 withTaskCancellationHandler {
-            try await task.task.value
-        } onCancel: {
-            Task {
-                await self.unsubscribe(subscriptionID, key: requestKey)
-            }
+        return try await task.getData(subscriptionID: subscriptionID)
+    }
+
+    fileprivate nonisolated func unsubscribe(_ subscriptionID: UUID, key: String) {
+        Task {
+            await _unsubscribe(subscriptionID, key: key)
         }
     }
 
-    private func unsubscribe(_ subscriptionID: UUID, key: String) {
+    private func _unsubscribe(_ subscriptionID: UUID, key: String) {
         guard let task = tasks[key],
               task.subscriptions.remove(subscriptionID) != nil,
               task.subscriptions.isEmpty else {
@@ -151,14 +153,24 @@ actor ImageDownloader {
 }
 
 private final class ImageDataTask {
+    let key: String
     var subscriptions = Set<UUID>()
-    var isCancelled = false
-    var task: Task<Data, Error>
+    let task: Task<Data, Error>
+    weak var downloader: ImageDownloader?
 
-    init(subscriptions: Set<UUID> = Set<UUID>(), task: Task<Data, Error>) {
-        self.subscriptions = subscriptions
+    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)