Merge pull request #2549 from onevcat/fix-force-cancel-race

Fix data race on SessionDataTask.started and tighten concurrency tests

Author: onevcat

Sources/Networking/SessionDataTask.swift

@@ -98,7 +98,13 @@ public class SessionDataTask: @unchecked Sendable {
     let onTaskDone = Delegate<(Result<(Data, URLResponse?), KingfisherError>, [TaskCallback]), Void>()
     let onCallbackCancelled = Delegate<(CancelToken, TaskCallback), Void>()
 
-    var started = false
+    private var _started = false
+    var started: Bool {
+        lock.lock()
+        defer { lock.unlock() }
+        return _started
+    }
+
     var containsCallbacks: Bool {
         // We should be able to use `task.state != .running` to check it.
         // However, in some rare cases, cancelling the task does not change
@@ -154,8 +160,14 @@ public class SessionDataTask: @unchecked Sendable {
     }
 
     func resume() {
-        guard !started else { return }
-        started = true
+        // Atomic check-and-set; `task.resume()` is called outside the lock.
+        lock.lock()
+        guard !_started else {
+            lock.unlock()
+            return
+        }
+        _started = true
+        lock.unlock()
         task.resume()
     }
 
@@ -167,16 +179,9 @@ public class SessionDataTask: @unchecked Sendable {
     }
 
     func forceCancel() {
-        // `callbacksStore` is protected by `lock` for every other access. Snapshot the
-        // tokens under the lock before cancelling, for two reasons:
-        // 1. `forceCancel()` is reachable from the public `ImageDownloader.cancelAll()` and
-        //    `cancel(url:)` on arbitrary caller threads, while the session delegate queue may
-        //    concurrently mutate the store via `addCallback`/`completeAndRemoveAllCallbacks`.
-        //    Reading the live `keys` view here would be a data race on the dictionary.
-        // 2. `cancel(token:)` removes the token via `removeCallback`, so iterating the live
-        //    `keys` view would mutate the dictionary mid-iteration.
-        // Iterating an immutable snapshot avoids both. The lock is released before calling
-        // `cancel(token:)`, which re-acquires it (the non-recursive `lock` would otherwise deadlock).
+        // Snapshot the tokens under the lock, then cancel outside of it: `forceCancel` can run on
+        // any thread while `callbacksStore` is being mutated, and `cancel(token:)` re-acquires the
+        // non-recurrent lock.
         lock.lock()
         let tokens = Array(callbacksStore.keys)
         lock.unlock()

Tests/KingfisherTests/ImageDownloaderTests.swift

@@ -505,6 +505,38 @@ class ImageDownloaderTests: XCTestCase {
         }
     }
 
+    // Hammers all mutating and reading surfaces of `SessionDataTask` concurrently, including the
+    // `started` flag which `resume()` writes while `ImageDownloader.startDownloadTask` reads it
+    // from other threads. Meaningful mainly under Thread Sanitizer.
+    func testSessionDataTaskConcurrentAccessIsThreadSafe() {
+        let url = URL(string: "https://example.com/concurrent-access")!
+        // `resume()` resumes the underlying URLSessionDataTask; stub the URL so no real request leaves.
+        stub(url, data: Data())
+        let options = KingfisherParsedOptionsInfo(nil)
+
+        for _ in 0..<50 {
+            let task = SessionDataTask(task: URLSession.shared.dataTask(with: url))
+            let group = DispatchGroup()
+
+            func hammer(_ body: @escaping () -> Void) {
+                group.enter()
+                DispatchQueue.global().async {
+                    body()
+                    group.leave()
+                }
+            }
+
+            hammer { for _ in 0..<20 { _ = task.addCallback(.init(onCompleted: nil, options: options)) } }
+            hammer { for _ in 0..<5 { task.forceCancel() } }
+            hammer { for _ in 0..<20 { task.resume() } }
+            hammer { for _ in 0..<20 { _ = task.started; _ = task.containsCallbacks } }
+            hammer { for _ in 0..<20 { task.didReceiveData(Data([0x01])); _ = task.mutableDataCount } }
+            hammer { _ = task.completeAndRemoveAllCallbacks() }
+
+            group.wait()
+        }
+    }
+
     // Issue 532 https://github.com/onevcat/Kingfisher/issues/532#issuecomment-305644311
     func testCancelThenRestartSameDownload() {
         let exp = expectation(description: #function)