Fix synchronization issue in DataLoader

kean · Sep 21, 2022 · 2684188 · 2684188
1 parent f0795a2
commit 2684188
Showing 1 changed file with 34 additions and 16 deletions.
diff --git a/Sources/Get/DataLoader.swift b/Sources/Get/DataLoader.swift
@@ -9,7 +9,7 @@ import FoundationNetworking
 
 // A simple URLSession wrapper adding async/await APIs compatible with older platforms.
 final class DataLoader: NSObject, URLSessionDataDelegate, URLSessionDownloadDelegate, @unchecked Sendable {
-    private var handlers = [URLSessionTask: TaskHandler]()
+    private let handlers = TaskHandlersDictionary()
 
     var userSessionDelegate: URLSessionDelegate? {
         didSet {
@@ -31,11 +31,10 @@ final class DataLoader: NSObject, URLSessionDataDelegate, URLSessionDownloadDele
     func startDataTask(_ task: URLSessionDataTask, session: URLSession, delegate: URLSessionDataDelegate?) async throws -> Response<Data> {
         try await withTaskCancellationHandler(handler: { task.cancel() }) {
             try await withUnsafeThrowingContinuation { continuation in
-                session.delegateQueue.addOperation {
-                    let handler = DataTaskHandler(delegate: delegate)
-                    handler.completion = continuation.resume(with:)
-                    self.handlers[task] = handler
-                }
+                let handler = DataTaskHandler(delegate: delegate)
+                handler.completion = continuation.resume(with:)
+                self.handlers[task] = handler
+
                 task.resume()
             }
         }
@@ -44,11 +43,10 @@ final class DataLoader: NSObject, URLSessionDataDelegate, URLSessionDownloadDele
     func startDownloadTask(_ task: URLSessionDownloadTask, session: URLSession, delegate: URLSessionDownloadDelegate?) async throws -> Response<URL> {
         try await withTaskCancellationHandler(handler: { task.cancel() }) {
             try await withUnsafeThrowingContinuation { continuation in
-                session.delegateQueue.addOperation {
-                    let handler = DownloadTaskHandler(delegate: delegate)
-                    handler.completion = continuation.resume(with:)
-                    self.handlers[task] = handler
-                }
+                let handler = DownloadTaskHandler(delegate: delegate)
+                handler.completion = continuation.resume(with:)
+                self.handlers[task] = handler
+
                 task.resume()
             }
         }
@@ -57,11 +55,10 @@ final class DataLoader: NSObject, URLSessionDataDelegate, URLSessionDownloadDele
     func startUploadTask(_ task: URLSessionUploadTask, session: URLSession, delegate: URLSessionTaskDelegate?) async throws -> Response<Data> {
         try await withTaskCancellationHandler(handler: { task.cancel() }) {
             try await withUnsafeThrowingContinuation { continuation in
-                session.delegateQueue.addOperation {
-                    let handler = DataTaskHandler(delegate: delegate)
-                    handler.completion = continuation.resume(with:)
-                    self.handlers[task] = handler
-                }
+                let handler = DataTaskHandler(delegate: delegate)
+                handler.completion = continuation.resume(with:)
+                self.handlers[task] = handler
+
                 task.resume()
             }
         }
@@ -359,3 +356,24 @@ func decode<T: Decodable>(_ data: Data, using decoder: JSONDecoder) async throws
         }.value
     }
 }
+
+/// With iOS 16, there is now a delegate method (`didCreateTask`) that gets
+/// called outside of the session's delegate queue, which means that the access
+/// needs to be synchronized.
+private final class TaskHandlersDictionary {
+    private let lock = NSLock()
+    private var handlers = [URLSessionTask: TaskHandler]()
+
+    subscript(task: URLSessionTask) -> TaskHandler? {
+        get {
+            lock.lock()
+            defer { lock.unlock() }
+            return handlers[task]
+        }
+        set {
+            lock.lock()
+            defer { lock.unlock() }
+            handlers[task] = newValue
+        }
+    }
+}