From d6a74b0232192cb4daa35256f371ed5d90ea65af Mon Sep 17 00:00:00 2001
From: Aleksey Gaponov <aleks.gapp@gmail.com>
Date: Thu, 29 Sep 2016 14:13:17 +0200
Subject: [PATCH] [SR-2798] Fix deadlock in URLSession when calling
 finishTasksAndInvalidate().

---
 Foundation/NSURLSession/NSURLSession.swift | 22 ++++++++++--------
 Foundation/NSURLSession/TaskRegistry.swift | 10 ++++----
 TestFoundation/TestNSURLSession.swift      | 27 +++++++++++++++++++++-
 3 files changed, 43 insertions(+), 16 deletions(-)

diff --git a/Foundation/NSURLSession/NSURLSession.swift b/Foundation/NSURLSession/NSURLSession.swift
index f564b0970b..ab6bd1d217 100644
--- a/Foundation/NSURLSession/NSURLSession.swift
+++ b/Foundation/NSURLSession/NSURLSession.swift
@@ -265,18 +265,20 @@ open class URLSession : NSObject {
            //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()
+           let invalidateSessionCallback = { [weak self] in
+               //invoke the delegate method and break the delegate link
+               guard let `self` = self, let sessionDelegate = self.delegate else { return }
+               self.delegateQueue.addOperation {
+                   sessionDelegate.urlSession(self, didBecomeInvalidWithError: nil)
+                   self.delegate = nil
+               }
            }
 
-           //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
+           //wait for running tasks to finish
+           if !self.taskRegistry.isEmpty {
+               self.taskRegistry.notify(on: invalidateSessionCallback)
+           } else {
+               invalidateSessionCallback()
            }
        }
     }
diff --git a/Foundation/NSURLSession/TaskRegistry.swift b/Foundation/NSURLSession/TaskRegistry.swift
index 9b1809ce32..e772eb2b67 100644
--- a/Foundation/NSURLSession/TaskRegistry.swift
+++ b/Foundation/NSURLSession/TaskRegistry.swift
@@ -45,7 +45,7 @@ extension URLSession {
         
         fileprivate var tasks: [Int: URLSessionTask] = [:]
         fileprivate var behaviours: [Int: _Behaviour] = [:]
-        fileprivate var tasksFinished: DispatchSemaphore?
+        fileprivate var tasksFinishedCallback: (() -> ())?
     }
 }
 
@@ -81,14 +81,14 @@ extension URLSession._TaskRegistry {
         }
         behaviours.remove(at: behaviourIdx)
 
-        guard let allTasksFinished = tasksFinished else { return }
+        guard let allTasksFinished = tasksFinishedCallback else { return }
         if self.isEmpty {
-            allTasksFinished.signal()
+            allTasksFinished()
         }
     }
 
-    func notify(on semaphore: DispatchSemaphore) {
-        tasksFinished = semaphore
+    func notify(on tasksCompetion: @escaping () -> ()) {
+        tasksFinishedCallback = tasksCompetion
     }
 
     var isEmpty: Bool {
diff --git a/TestFoundation/TestNSURLSession.swift b/TestFoundation/TestNSURLSession.swift
index 487389900e..e3b721bc9f 100644
--- a/TestFoundation/TestNSURLSession.swift
+++ b/TestFoundation/TestNSURLSession.swift
@@ -29,7 +29,7 @@ class TestURLSession : XCTestCase {
             ("test_downloadTaskWithURLRequest", test_downloadTaskWithURLRequest),
             ("test_downloadTaskWithRequestAndHandler", test_downloadTaskWithRequestAndHandler),
             ("test_downloadTaskWithURLAndHandler", test_downloadTaskWithURLAndHandler),
-            
+            ("test_finishTaskAndInvalidate", test_finishTasksAndInvalidate)
         ]
     }
 
@@ -247,6 +247,31 @@ class TestURLSession : XCTestCase {
         task.resume()
         waitForExpectations(timeout: 12)
     }
+    
+    func test_finishTasksAndInvalidate() {
+        let invalidateExpectation = expectation(description: "URLSession wasn't invalidated")
+        let delegate = SessionDelegate(invalidateExpectation: invalidateExpectation)
+        let url = URL(string: "http://127.0.0.1:\(serverPort)/Nepal")!
+        let session = URLSession(configuration: URLSessionConfiguration.default,
+                                 delegate: delegate, delegateQueue: nil)
+        let completionExpectation = expectation(description: "dataTask completion block wasn't called")
+        let task = session.dataTask(with: url) { _ in
+            completionExpectation.fulfill()
+        }
+        task.resume()
+        session.finishTasksAndInvalidate()
+        waitForExpectations(timeout: 12)
+    }
+}
+
+class SessionDelegate: NSObject, URLSessionDelegate {
+    let invalidateExpectation: XCTestExpectation
+    init(invalidateExpectation: XCTestExpectation){
+        self.invalidateExpectation = invalidateExpectation
+    }
+    func urlSession(_ session: URLSession, didBecomeInvalidWithError error: Error?) {
+        invalidateExpectation.fulfill()
+    }
 }
 
 class DataTask : NSObject {