Merge pull request #640 from pushkarnk/sr2631

Author: swift-ci

Foundation/NSURLSession/NSURLSession.swift

@@ -189,8 +189,9 @@ open class URLSession : NSObject {
     /// This queue is used to make public attributes on `URLSessionTask` instances thread safe.
     /// - Note: It's a **concurrent** queue.
     internal let taskAttributesIsolation: DispatchQueue 
-    fileprivate let taskRegistry = URLSession._TaskRegistry()
+    internal let taskRegistry = URLSession._TaskRegistry()
     fileprivate let identifier: Int32
+    fileprivate var invalidated = false
 
     /*
      * The shared session uses the currently set global NSURLCache,
@@ -237,7 +238,7 @@ open class URLSession : NSObject {
     }
 
     open let delegateQueue: OperationQueue
-    open let delegate: URLSessionDelegate?
+    open var delegate: URLSessionDelegate?
     open let configuration: URLSessionConfiguration
 
     /*
@@ -258,7 +259,27 @@ open class URLSession : NSObject {
      * session with the same identifier until URLSession:didBecomeInvalidWithError: has
      * been issued.
      */
-    open func finishTasksAndInvalidate() { NSUnimplemented() }
+    open func finishTasksAndInvalidate() {
+       //we need to return immediately
+       workQueue.async {
+           //don't allow creation of new tasks from this point onwards
+           self.invalidated = true
+
+           //wait for running tasks to finish
+           if !self.taskRegistry.isEmpty {
+               let tasksCompletion = DispatchSemaphore(value: 0)
+               self.taskRegistry.notify(on: tasksCompletion)
+               tasksCompletion.wait()
+           }
+
+           //invoke the delegate method and break the delegate link
+           guard let sessionDelegate = self.delegate else { return }
+           self.delegateQueue.addOperation {
+               sessionDelegate.urlSession(self, didBecomeInvalidWithError: nil)
+               self.delegate = nil
+           }
+       }
+    }
 
     /* -invalidateAndCancel acts as -finishTasksAndInvalidate, but issues
      * -cancel to all outstanding tasks for this session.  Note task
@@ -367,6 +388,7 @@ fileprivate extension URLSession {
     ///
     /// All public methods funnel into this one.
     func dataTask(with request: _Request, behaviour: _TaskRegistry._Behaviour) -> URLSessionDataTask {
+        guard !self.invalidated else { fatalError("Session invalidated") }
         let r = createConfiguredRequest(from: request)
         let i = createNextTaskIdentifier()
         let task = URLSessionDataTask(session: self, request: r, taskIdentifier: i)
@@ -380,6 +402,7 @@ fileprivate extension URLSession {
     ///
     /// All public methods funnel into this one.
     func uploadTask(with request: _Request, body: URLSessionTask._Body, behaviour: _TaskRegistry._Behaviour) -> URLSessionUploadTask {
+        guard !self.invalidated else { fatalError("Session invalidated") }
         let r = createConfiguredRequest(from: request)
         let i = createNextTaskIdentifier()
         let task = URLSessionUploadTask(session: self, request: r, taskIdentifier: i, body: body)
@@ -391,6 +414,7 @@ fileprivate extension URLSession {
 
     /// Create a download task
     func downloadTask(with request: _Request, behavior: _TaskRegistry._Behaviour) -> URLSessionDownloadTask {
+        guard !self.invalidated else { fatalError("Session invalidated") }
         let r = createConfiguredRequest(from: request)
         let i = createNextTaskIdentifier()
         let task = URLSessionDownloadTask(session: self, request: r, taskIdentifier: i)

Foundation/NSURLSession/NSURLSessionTask.swift

@@ -29,7 +29,7 @@ open class URLSessionTask : NSObject, NSCopying {
     fileprivate var suspendCount = 1
     fileprivate var easyHandle: _EasyHandle!
     fileprivate var totalDownloaded = 0
-    fileprivate unowned let session: URLSessionProtocol
+    fileprivate var session: URLSessionProtocol! //change to nil when task completes
     fileprivate let body: _Body
     fileprivate let tempFileURL: URL
 
@@ -62,6 +62,10 @@ open class URLSessionTask : NSObject, NSCopying {
             }
             if case .taskCompleted = internalState {
                 updateTaskState()
+                guard let s = session as? URLSession else { fatalError() }
+                s.workQueue.async {
+                    s.taskRegistry.remove(self)
+                }
             }
         }
     }
@@ -825,16 +829,19 @@ extension URLSessionTask {
         guard case .transferCompleted(response: let response, bodyDataDrain: let bodyDataDrain) = internalState else {
             fatalError("Trying to complete the task, but its transfer isn't complete.")
         }
-        internalState = .taskCompleted
         self.response = response
+
+        //because we deregister the task with the session on internalState being set to taskCompleted
+        //we need to do the latter after the delegate/handler was notified/invoked
         switch session.behaviour(for: self) {
         case .taskDelegate(let delegate):
             guard let s = session as? URLSession else { fatalError() }
             s.delegateQueue.addOperation {
                 delegate.urlSession(s, task: self, didCompleteWithError: nil)
+                self.internalState = .taskCompleted
             }
         case .noDelegate:
-            break
+            internalState = .taskCompleted
         case .dataCompletionHandler(let completion):
             guard case .inMemory(let bodyData) = bodyDataDrain else {
                 fatalError("Task has data completion handler, but data drain is not in-memory.")
@@ -849,6 +856,8 @@ extension URLSessionTask {
 
             s.delegateQueue.addOperation {
                 completion(data, response, nil)
+                self.internalState = .taskCompleted
+                self.session = nil
             }
         case .downloadCompletionHandler(let completion):
             guard case .toFile(let url, let fileHandle?) = bodyDataDrain else {
@@ -861,6 +870,8 @@ extension URLSessionTask {
 
             s.delegateQueue.addOperation {
                 completion(url, response, nil) 
+                self.internalState = .taskCompleted
+                self.session = nil
             }
 
         }
@@ -869,24 +880,26 @@ extension URLSessionTask {
         guard case .transferFailed = internalState else {
             fatalError("Trying to complete the task, but its transfer isn't complete / failed.")
         }
-        internalState = .taskCompleted
         switch session.behaviour(for: self) {
         case .taskDelegate(let delegate):
             guard let s = session as? URLSession else { fatalError() }
             s.delegateQueue.addOperation {
                 delegate.urlSession(s, task: self, didCompleteWithError: error as Error)
+                self.internalState = .taskCompleted
             }
         case .noDelegate:
-            break
+            internalState = .taskCompleted
         case .dataCompletionHandler(let completion):
             guard let s = session as? URLSession else { fatalError() }
             s.delegateQueue.addOperation {
                 completion(nil, nil, error)
+                self.internalState = .taskCompleted
             }
         case .downloadCompletionHandler(let completion): 
             guard let s = session as? URLSession else { fatalError() }
             s.delegateQueue.addOperation {
                 completion(nil, nil, error)
+                self.internalState = .taskCompleted
             }
         }
     }

Foundation/NSURLSession/TaskRegistry.swift

@@ -17,7 +17,7 @@
 // -----------------------------------------------------------------------------
 
 import CoreFoundation
-
+import Dispatch
 
 extension URLSession {
     /// This helper class keeps track of all tasks, and their behaviours.
@@ -45,6 +45,7 @@ extension URLSession {
 
         fileprivate var tasks: [Int: URLSessionTask] = [:]
         fileprivate var behaviours: [Int: _Behaviour] = [:]
+        fileprivate var tasksFinished: DispatchSemaphore?
     }
 }
 
@@ -79,6 +80,19 @@ extension URLSession._TaskRegistry {
             fatalError("Trying to remove task's behaviour, but it's not in the registry.")
         }
         behaviours.remove(at: behaviourIdx)
+
+        guard let allTasksFinished = tasksFinished else { return }
+        if self.isEmpty {
+            allTasksFinished.signal()
+        }
+    }
+
+    func notify(on semaphore: DispatchSemaphore) {
+        tasksFinished = semaphore
+    }
+
+    var isEmpty: Bool {
+        return tasks.count == 0
     }
 }
 extension URLSession._TaskRegistry {