Merge pull request #2105 from karthikkeyan/karthik/SR-10428-gettaskswithcompletion

Author: swift-ci

Foundation/URLSession/URLSession.swift

@@ -345,14 +345,38 @@ open class URLSession : NSObject {
     open func reset(completionHandler: @escaping () -> Void) { NSUnimplemented() } /* empty all cookies, cache and credential stores, removes disk files, issues -flushWithCompletionHandler:. Invokes completionHandler() on the delegate queue if not nil. */
 
     open func flush(completionHandler: @escaping () -> Void)  { NSUnimplemented() }/* flush storage to disk and clear transient network caches.  Invokes completionHandler() on the delegate queue if not nil. */
-    
-    open func getTasksWithCompletionHandler(completionHandler: @escaping ([URLSessionDataTask], [URLSessionUploadTask], [URLSessionDownloadTask]) -> Void)  { NSUnimplemented() }/* invokes completionHandler with outstanding data, upload and download tasks. */
+
+    /* invokes completionHandler with outstanding data, upload and download tasks. */
+    open func getTasksWithCompletionHandler(completionHandler: @escaping ([URLSessionDataTask], [URLSessionUploadTask], [URLSessionDownloadTask]) -> Void)  {
+        workQueue.async {
+            self.delegateQueue.addOperation {
+                var dataTasks = [URLSessionDataTask]()
+                var uploadTasks = [URLSessionUploadTask]()
+                var downloadTasks = [URLSessionDownloadTask]()
+
+                for task in self.taskRegistry.allTasks {
+                    guard task.state == .running || task.isSuspendedAfterResume else { continue }
+
+                    if let uploadTask = task as? URLSessionUploadTask {
+                        uploadTasks.append(uploadTask)
+                    } else if let dataTask = task as? URLSessionDataTask {
+                        dataTasks.append(dataTask)
+                    } else if let downloadTask = task as? URLSessionDownloadTask {
+                        downloadTasks.append(downloadTask)
+                    } else {
+                        // Above three are the only required tasks to be returned from this API, so we can ignore any other types of tasks.
+                    }
+                }
+                completionHandler(dataTasks, uploadTasks, downloadTasks)
+            }
+        }
+    }
 
     /* invokes completionHandler with all outstanding tasks. */
     open func getAllTasks(completionHandler: @escaping ([URLSessionTask]) -> Void)  {
         workQueue.async {
             self.delegateQueue.addOperation {
-                completionHandler(self.taskRegistry.allTasks.filter { $0.state == .running || $0.state == .suspended })
+                completionHandler(self.taskRegistry.allTasks.filter { $0.state == .running || $0.isSuspendedAfterResume })
             }
         }
     }

Foundation/URLSession/URLSessionTask.swift

@@ -35,6 +35,10 @@ open class URLSessionTask : NSObject, NSCopying {
     internal let body: _Body
     fileprivate var _protocol: URLProtocol? = nil
     private let syncQ = DispatchQueue(label: "org.swift.URLSessionTask.SyncQ")
+    private var hasTriggeredResume: Bool = false
+    internal var isSuspendedAfterResume: Bool {
+        return self.syncQ.sync { return self.hasTriggeredResume } && self.state == .suspended
+    }
 
     /// All operations must run on this queue.
     internal let workQueue: DispatchQueue 
@@ -258,6 +262,7 @@ open class URLSessionTask : NSObject, NSCopying {
             guard 0 <= self.suspendCount else { fatalError("Resuming a task that's not suspended. Calls to resume() / suspend() need to be matched.") }
             self.updateTaskState()
             if self.suspendCount == 0 {
+                self.hasTriggeredResume = true
                 self.workQueue.async {
                     if let _protocol = self._protocol {
                         _protocol.startLoading()

TestFoundation/TestURLSession.swift

@@ -53,6 +53,8 @@ class TestURLSession : LoopbackServerTest {
             ("test_checkErrorTypeAfterInvalidateAndCancel", test_checkErrorTypeAfterInvalidateAndCancel),
             ("test_taskCountAfterInvalidateAndCancel", test_taskCountAfterInvalidateAndCancel),
             ("test_sessionDelegateAfterInvalidateAndCancel", test_sessionDelegateAfterInvalidateAndCancel),
+            ("test_getAllTasks", test_getAllTasks),
+            ("test_getTasksWithCompletion", test_getTasksWithCompletion),
         ]
     }
 
@@ -819,32 +821,31 @@ class TestURLSession : LoopbackServerTest {
         session.invalidateAndCancel()
         waitForExpectations(timeout: 5)
     }
-    
+
     func test_taskCountAfterInvalidateAndCancel() {
         let expect = expectation(description: "Check task count after invalidateAndCancel")
-        
+
         let session = URLSession(configuration: .default)
         let task1 = session.dataTask(with: URL(string: "https://www.apple.com")!)
         let task2 = session.dataTask(with: URL(string: "https://developer.apple.com")!)
         let task3 = session.dataTask(with: URL(string: "https://developer.apple.com/swift")!)
-        
+
         task1.resume()
         task2.resume()
         session.invalidateAndCancel()
         Thread.sleep(forTimeInterval: 1)
-        
+
         session.getAllTasks { tasksBeforeResume in
             XCTAssertEqual(tasksBeforeResume.count, 0)
-            
+
             // Resume a task after invalidating a session shouldn't change the task's status
             task3.resume()
-            
+
             session.getAllTasks { tasksAfterResume in
                 XCTAssertEqual(tasksAfterResume.count, 0)
                 expect.fulfill()
             }
         }
-        
         waitForExpectations(timeout: 5)
     }
 
@@ -856,6 +857,91 @@ class TestURLSession : LoopbackServerTest {
         XCTAssertNil(session.delegate)
     }
 
+    func test_getAllTasks() {
+        let expect = expectation(description: "Tasks URLSession.getAllTasks")
+
+        let session = URLSession(configuration: .default)
+        let dataTask1 = session.dataTask(with: URL(string: "https://www.apple.com")!)
+        let dataTask2 = session.dataTask(with: URL(string: "https://developer.apple.com")!)
+        let dataTask3 = session.dataTask(with: URL(string: "https://developer.apple.com/swift")!)
+
+        session.getAllTasks { (tasksBeforeResume) in
+            XCTAssertEqual(tasksBeforeResume.count, 0)
+
+            dataTask1.cancel()
+
+            dataTask2.resume()
+            dataTask2.suspend()
+            // dataTask3 is suspended even before it was resumed, so the next call to `getAllTasks` should not include this tasks
+            dataTask3.suspend()
+            session.getAllTasks { (tasksAfterCancel) in
+                // tasksAfterCancel should only contain dataTask2
+                XCTAssertEqual(tasksAfterCancel.count, 1)
+
+                // A task will in be in suspended state when it was created.
+                // Given that, dataTask3 was suspended once again earlier above, so it should receive `resume()` twice in order to be executed
+                // Calling `getAllTasks` next time should not include dataTask3
+                dataTask3.resume()
+
+                session.getAllTasks { (tasksAfterFirstResume) in
+                    // tasksAfterFirstResume should only contain dataTask2
+                    XCTAssertEqual(tasksAfterFirstResume.count, 1)
+
+                    // Now dataTask3 received `resume()` twice, this time `getAllTasks` should include
+                    dataTask3.resume()
+                    session.getAllTasks { (tasksAfterSecondResume) in
+                        // tasksAfterSecondResume should contain dataTask2 and dataTask2 this time
+                        XCTAssertEqual(tasksAfterSecondResume.count, 2)
+                        expect.fulfill()
+                    }
+                }
+            }
+        }
+
+        waitForExpectations(timeout: 20)
+    }
+
+    func test_getTasksWithCompletion() {
+        let expect = expectation(description: "Test URLSession.getTasksWithCompletion")
+
+        let session = URLSession(configuration: .default)
+        let dataTask1 = session.dataTask(with: URL(string: "https://www.apple.com")!)
+        let dataTask2 = session.dataTask(with: URL(string: "https://developer.apple.com")!)
+        let dataTask3 = session.dataTask(with: URL(string: "https://developer.apple.com/swift")!)
+
+        let uploadTask1 = session.uploadTask(with: URLRequest(url: URL(string: "https://developer.apple.com")!), from: Data())
+        let uploadTask2 = session.uploadTask(with: URLRequest(url: URL(string: "https://developer.apple.com/swift")!), from: Data())
+
+        let downloadTask1 = session.downloadTask(with: URL(string: "https://developer.apple.com/assets/elements/icons/brandmark/apple-developer-brandmark.svg")!)
+
+        session.getTasksWithCompletionHandler { (dataTasksBeforeCancel, uploadTasksBeforeCancel, downloadTasksBeforeCancel) in
+            XCTAssertEqual(dataTasksBeforeCancel.count, 0)
+            XCTAssertEqual(uploadTasksBeforeCancel.count, 0)
+            XCTAssertEqual(downloadTasksBeforeCancel.count, 0)
+
+            dataTask1.cancel()
+            dataTask2.resume()
+            // dataTask3 is resumed and suspended, so this task should be a part of `getTasksWithCompletionHandler` response
+            dataTask3.resume()
+            dataTask3.suspend()
+
+            // uploadTask1 suspended even before it was resumed, so this task shouldn't be a part of `getTasksWithCompletionHandler` response
+            uploadTask1.suspend()
+            uploadTask2.resume()
+
+            downloadTask1.cancel()
+
+            session.getTasksWithCompletionHandler{ (dataTasksAfterCancel, uploadTasksAfterCancel, downloadTasksAfterCancel) in
+                XCTAssertEqual(dataTasksAfterCancel.count, 2)
+                XCTAssertEqual(uploadTasksAfterCancel.count, 1)
+                XCTAssertEqual(downloadTasksAfterCancel.count, 0)
+                expect.fulfill()
+            }
+        }
+
+        waitForExpectations(timeout: 20)
+    }
+
 }
 
 class SharedDelegate: NSObject {