Merge pull request #1558 from onevcat/fix/lock-task-cancellation

Avoid racing for accessing originalRequest

Sources/Networking/ImageDownloader.swift

@@ -186,7 +186,7 @@ open class ImageDownloader {
             }
         }
         sessionDelegate.onDidDownloadData.delegate(on: self) { (self, task) in
-            guard let url = task.task.originalRequest?.url else {
+            guard let url = task.originalURL else {
                 return task.mutableData
             }
             return (self.delegate ?? self).imageDownloader(self, didDownload: task.mutableData, for: url)

Sources/Networking/SessionDataTask.swift

@@ -41,6 +41,10 @@ public class SessionDataTask {
     /// Downloaded raw data of current task.
     public private(set) var mutableData: Data
 
+    // This is a copy of `task.originalRequest?.url`. It is for getting a race-safe behavior for a pitfall on iOS 13.
+    // Ref: https://github.com/onevcat/Kingfisher/issues/1511
+    let originalURL: URL?
+
     /// The underlying download task. It is only for debugging purpose when you encountered an error. You should not
     /// modify the content of this task or start it yourself.
     public let task: URLSessionDataTask
@@ -70,6 +74,7 @@ public class SessionDataTask {
 
     init(task: URLSessionDataTask) {
         self.task = task
+        self.originalURL = task.originalRequest?.url
         mutableData = Data()
     }
 

Sources/Networking/SessionDelegate.swift

@@ -73,8 +73,9 @@ class SessionDelegate: NSObject {
             // No other callbacks waiting, we can clear the task now.
             if !task.containsCallbacks {
                 let dataTask = task.task
-                dataTask.cancel()
-                self.remove(dataTask)
+
+                self.cancelTask(dataTask)
+                self.remove(task)
             }
         }
         let token = task.addCallback(callback)
@@ -82,6 +83,12 @@ class SessionDelegate: NSObject {
         return DownloadTask(sessionTask: task, cancelToken: token)
     }
 
+    private func cancelTask(_ dataTask: URLSessionDataTask) {
+        lock.lock()
+        defer { lock.unlock() }
+        dataTask.cancel()
+    }
+
     func append(
         _ task: SessionDataTask,
         url: URL,
@@ -91,23 +98,23 @@ class SessionDelegate: NSObject {
         return DownloadTask(sessionTask: task, cancelToken: token)
     }
 
-    private func remove(_ task: URLSessionTask) {
-        guard let url = task.originalRequest?.url else {
+    private func remove(_ task: SessionDataTask) {
+        lock.lock()
+        defer { lock.unlock() }
+
+        guard let url = task.originalURL else {
             return
         }
-        lock.lock()
-        defer {lock.unlock()}
         tasks[url] = nil
     }
 
     private func task(for task: URLSessionTask) -> SessionDataTask? {
+        lock.lock()
+        defer { lock.unlock() }
 
         guard let url = task.originalRequest?.url else {
             return nil
         }
-
-        lock.lock()
-        defer { lock.unlock() }
         guard let sessionTask = tasks[url] else {
             return nil
         }
@@ -182,7 +189,7 @@ extension SessionDelegate: URLSessionDataDelegate {
     func urlSession(_ session: URLSession, task: URLSessionTask, didCompleteWithError error: Error?) {
         guard let sessionTask = self.task(for: task) else { return }
 
-        if let url = task.originalRequest?.url {
+        if let url = sessionTask.originalURL {
             let result: Result<URLResponse, KingfisherError>
             if let error = error {
                 result = .failure(KingfisherError.responseError(reason: .URLSessionError(error: error)))
@@ -249,7 +256,7 @@ extension SessionDelegate: URLSessionDataDelegate {
         guard let sessionTask = self.task(for: task) else {
             return
         }
-        remove(task)
+        remove(sessionTask)
         sessionTask.onTaskDone.call((result, sessionTask.callbacks))
     }
 }

Tests/KingfisherTests/RetryStrategyTests.swift

@@ -146,11 +146,13 @@ class RetryStrategyTests: XCTestCase {
         var blockCalled: [Bool] = []
         let source = Source.network(URL(string: "url")!)
         let retry = DelayRetryStrategy(maxRetryCount: 3, retryInterval: .seconds(0))
+
+        let task = URLSession.shared.dataTask(with: URL(string: "url")!)
+
         let context1 = RetryContext(
             source: source,
-            error: .requestError(reason: .taskCancelled(task: .init(task: .init()), token: .init()))
+            error: .requestError(reason: .taskCancelled(task: .init(task: task), token: .init()))
         )
-
         retry.retry(context: context1) { decision in
             guard case RetryDecision.stop = decision else {
                 XCTFail("The decision should be `stop` if user cancelled the task.")