Prevents data races when creating assembles concurrently (#40)

* Thread safety for assembly creation

* Thread safety for callbacks

* Fix testing of concurrent assembly creation

* Bump CI to Swift 6 toolchain

* Temporarily turn off concurrent test and bump TUSKit version

* re-enable concurrency test with lower concurrent operation count

Author: donnywals

.github/workflows/transloaditkit-ci.yml

@@ -6,7 +6,7 @@ jobs:
     strategy:
       matrix:
         os: ["macos-latest"]
-        swift: ["5.10"]
+        swift: ["6.0.2"]
     runs-on: ${{ matrix.os }}
     steps:
       - name: Extract Branch Name

.gitignore

@@ -121,7 +121,9 @@ playground.xcworkspace
 # Package.resolved
 .build/
 # Add this line if you want to avoid checking in Xcode SPM integration.
-# .swiftpm/xcode
+.swiftpm/xcode
+
+.vscode/
 
 # CocoaPods
 # We recommend against adding the Pods directory to your .gitignore. However

Package.resolved

@@ -14,7 +14,7 @@ let package = Package(
     ],
     dependencies: [
         // Dependencies declare other packages that this package depends on.
-        .package(name: "TUSKit", url: "https://github.com/tus/TUSKit", from: "3.3.0")
+        .package(name: "TUSKit", url: "https://github.com/tus/TUSKit", from: "3.4.2")
     ],
     targets: [
         // Targets are the basic building blocks of a package. A target can define a module or a test suite.

Package.swift

@@ -51,7 +51,7 @@ public final class Transloadit {
     }
 
     typealias FileId = UUID
-    var pollers = [[URL]: TransloaditPoller]()
+    let pollers = TransloaditPollers()
 
     // "private" -- only exposed for unit testing
     let api: TransloaditAPI
@@ -201,10 +201,10 @@ public final class Transloadit {
 
         let poller = TransloaditPoller(transloadit: self, didFinish: { [weak self] in
             guard let self = self else { return }
-            self.pollers[files] = nil
+            self.pollers.remove(for: files)
         })
 
-        if let existingPoller = self.pollers[files], existingPoller === poller {
+        if let existingPoller = self.pollers.get(for: files), existingPoller === poller {
             assertionFailure("Transloadit: Somehow already got a poller for this url and these files")
         }
 
@@ -234,7 +234,8 @@ public final class Transloadit {
             }
         })
 
-        pollers[files] = poller
+        pollers.register(poller, for: files)
+
         return poller
     }
 
@@ -273,10 +274,10 @@ public final class Transloadit {
 
         let poller = TransloaditPoller(transloadit: self, didFinish: { [weak self] in
             guard let self = self else { return }
-            self.pollers[files] = nil
+            self.pollers.remove(for: files)
         })
 
-        if let existingPoller = self.pollers[files], existingPoller === poller {
+        if let existingPoller = self.pollers.get(for: files), existingPoller === poller {
             assertionFailure("Transloadit: Somehow already got a poller for this url and these files")
         }
 
@@ -306,7 +307,7 @@ public final class Transloadit {
                 }
             })
 
-        pollers[files] = poller
+        pollers.register(poller, for: files)
         return poller
     }
 
@@ -436,6 +437,29 @@ extension Transloadit: TUSClientDelegate {
     }
 }
 
+class TransloaditPollers {
+    private var pollers = [[URL]: TransloaditPoller]()
+    private var syncQueue = DispatchQueue(label: "com.transloadit.pollers")
+
+    func register(_ poller: TransloaditPoller, for files: [URL]) {
+        syncQueue.sync {
+            self.pollers[files] = poller
+        }
+    }
+
+    func remove(for files: [URL]) {
+        syncQueue.sync {
+            self.pollers[files] = nil
+        }
+    }
+
+    func get(for files: [URL]) -> TransloaditPoller? {
+        return syncQueue.sync {
+            return self.pollers[files]
+        }
+    }
+}
+
 
 /// Small helper function to get an `Assembly` out of a context from TUSKit.
 /// - Parameter context: A dictionary that's passed when uploading a file via TUSKit

Sources/TransloaditKit/Transloadit.swift

@@ -2,11 +2,11 @@ import Foundation
 
 extension TransloaditAPI: URLSessionDataDelegate {
     public func urlSession(_ session: URLSession, task: URLSessionTask, didCompleteWithError error: Error?) {
-        guard let completionHandler = callbacks[task] else {
+        guard let completionHandler = callbacks.get(for: task) else {
             return
         }
 
-        defer { callbacks[task] = nil }
+        defer { callbacks.remove(for: task) }
 
         if let error {
             completionHandler.callback(.failure(error))
@@ -22,6 +22,6 @@ extension TransloaditAPI: URLSessionDataDelegate {
     }
 
     public func urlSession(_ session: URLSession, dataTask: URLSessionDataTask, didReceive data: Data) {
-        callbacks[dataTask]?.data.append(data)
+        callbacks.get(for: dataTask)?.data.append(data)
     }
 }

Sources/TransloaditKit/TransloaditAPI+URLSessionDelegate.swift

@@ -41,7 +41,7 @@ final class TransloaditAPI: NSObject {
     }()
 
     private let credentials: Transloadit.Credentials
-    var callbacks: [URLSessionTask: URLSessionCompletionHandler] = [:]
+    let callbacks = TransloaditCallbacks()
 
     init(credentials: Transloadit.Credentials, session: URLSession) {
         self.credentials = credentials
@@ -76,7 +76,7 @@ final class TransloaditAPI: NSObject {
         }
 
         let task = session.uploadTask(with: request.request, fromFile: request.httpBody)
-        callbacks[task] = URLSessionCompletionHandler(callback: { result in
+        callbacks.register(URLSessionCompletionHandler(callback: { result in
             switch result {
             case .failure(let error):
                 completion(.failure(.couldNotCreateAssembly(error)))
@@ -95,7 +95,7 @@ final class TransloaditAPI: NSObject {
                     completion(.failure(TransloaditAPIError.couldNotCreateAssembly(error)))
                 }
             }
-        })
+        }), for: task)
         task.resume()
     }
 
@@ -118,7 +118,7 @@ final class TransloaditAPI: NSObject {
         }
 
         let task = session.uploadTask(with: request.request, fromFile: request.httpBody)
-        callbacks[task] = URLSessionCompletionHandler(callback: { result in
+        callbacks.register(URLSessionCompletionHandler(callback: { result in
             switch result {
             case .failure(let error):
                 completion(.failure(.couldNotCreateAssembly(error)))
@@ -137,7 +137,7 @@ final class TransloaditAPI: NSObject {
                     completion(.failure(TransloaditAPIError.couldNotCreateAssembly(error)))
                 }
             }
-        })
+        }), for: task)
         task.resume()
     }
 
@@ -304,7 +304,7 @@ final class TransloaditAPI: NSObject {
         }
 
         let task = session.dataTask(with: makeRequest())
-        callbacks[task] = URLSessionCompletionHandler(callback: { result in
+        callbacks.register(URLSessionCompletionHandler(callback: { result in
             switch result {
             case .failure:
                 completion(.failure(.couldNotFetchStatus))
@@ -318,7 +318,7 @@ final class TransloaditAPI: NSObject {
                     completion(.failure(.couldNotFetchStatus))
                 }
             }
-        })
+        }), for: task)
         task.resume()
     }
 
@@ -330,7 +330,7 @@ final class TransloaditAPI: NSObject {
         }
 
         let task = session.dataTask(with: makeRequest())
-        callbacks[task] = URLSessionCompletionHandler(callback: { result in
+        callbacks.register(URLSessionCompletionHandler(callback: { result in
             switch result {
             case .failure:
                 completion(.failure(.couldNotFetchStatus))
@@ -344,11 +344,33 @@ final class TransloaditAPI: NSObject {
                     completion(.failure(.couldNotFetchStatus))
                 }
             }
-        })
+        }), for: task)
         task.resume()
     }
 }
 
+class TransloaditCallbacks {
+    private var callbacks = [URLSessionTask: URLSessionCompletionHandler]()
+    private let syncQueue = DispatchQueue(label: "com.transloadit.callbacks")
+
+    func register(_ callback: URLSessionCompletionHandler, for task: URLSessionTask) {
+        syncQueue.sync {
+            callbacks[task] = callback
+        }
+    }
+
+    func remove(for task: URLSessionTask) {
+        syncQueue.sync {
+            callbacks[task] = nil
+        }
+    }
+
+    func get(for task: URLSessionTask) -> URLSessionCompletionHandler? {
+        return syncQueue.sync {
+            return callbacks[task]
+        }
+    }
+}
 
 extension String {
 

Sources/TransloaditKit/TransloaditAPI.swift

@@ -87,6 +87,27 @@ class TransloaditKitTests: XCTestCase {
         XCTAssertEqual(numFiles, fileDelegate.finishedUploads.count)
         XCTAssertEqual(numFiles, fileDelegate.startedUploads.count)
     }
+
+    func testConcurrentAssemblyCreation() throws {
+        let expect = expectation(description: "Wait for all assemblies to be created")
+        expect.expectedFulfillmentCount = 3
+        
+        DispatchQueue.concurrentPerform(iterations: 3) { _ in
+            do {
+                let (files, serverAssembly) = try Network.prepareForUploadingFiles(data: data)
+                let numFiles = files.count
+                let _ = createAssembly(files) { _ in
+                    expect.fulfill()
+                }
+            } catch {
+                XCTFail("Failed with error \(error)")
+            }
+        }
+        
+        wait(for: [expect], timeout: 10)
+        
+        try transloadit.reset()
+    }
 
     func testCanReset() throws {
         XCTAssertEqual(0, transloadit.remainingUploads)

Tests/TransloaditKitTests/TransloaditKitTests.swift

@@ -16,7 +16,7 @@ final class MyUploader: ObservableObject {
     let transloadit: Transloadit
 
     init(backgroundUploader: Bool = false) {
-        let credentials = Transloadit.Credentials(key: "OsCOAe4ro8CyNsHTp8pdhSiyEzuqwBue", secret: "jB5gZqmkiu2sdSwc7pko8iajD9ailws1eYUtwoKj")
+        let credentials = Transloadit.Credentials(key: "nOumdyAazaiVnkLDLBiERm0gEmIeeIeW", secret: "FfL2HOBEBxyXIp7zsXZguCMjuaIwhWkVqshF48lB")
 
         if backgroundUploader {
             self.transloadit = Transloadit(credentials: credentials, sessionConfiguration: .background(withIdentifier: "com.transloadit.bg_sample"))