Fix NSURLSessionTask ‘state’.

Author: danieleggert

Foundation/NSURLSession.swift

@@ -493,6 +493,9 @@ public class NSURLSessionTask : NSObject, NSCopying {
             if oldValue.isEasyHandlePaused && !internalState.isEasyHandlePaused {
                 fatalError("Need to solve pausing receive.")
             }
+            if case .taskCompleted = internalState {
+                updateTaskState()
+            }
         }
     }
     private let workQueue: dispatch_queue_t
@@ -559,9 +562,7 @@ public class NSURLSessionTask : NSObject, NSCopying {
             dispatch_sync(taskAttributesIsolation) { r = self._currentRequest }
             return r
         }
-        set {
-            dispatch_barrier_async(taskAttributesIsolation) { self._currentRequest = newValue }
-        }
+        set { dispatch_barrier_async(taskAttributesIsolation) { self._currentRequest = newValue } }
     }
     private var _currentRequest: NSURLRequest? = nil
     /*@NSCopying*/ public private(set) var response: NSURLResponse? {
@@ -570,9 +571,7 @@ public class NSURLSessionTask : NSObject, NSCopying {
             dispatch_sync(taskAttributesIsolation) { r = self._response }
             return r
         }
-        set {
-            dispatch_barrier_async(taskAttributesIsolation) { self._response = newValue }
-        }
+        set { dispatch_barrier_async(taskAttributesIsolation) { self._response = newValue } }
     }
     private var _response: NSURLResponse? = nil
 
@@ -588,9 +587,7 @@ public class NSURLSessionTask : NSObject, NSCopying {
             dispatch_sync(taskAttributesIsolation) { r = self._countOfBytesReceived }
             return r
         }
-        set {
-            dispatch_barrier_async(taskAttributesIsolation) { self._countOfBytesReceived = newValue }
-        }
+        set { dispatch_barrier_async(taskAttributesIsolation) { self._countOfBytesReceived = newValue } }
     }
     private var _countOfBytesReceived: Int64 = 0
 
@@ -601,9 +598,7 @@ public class NSURLSessionTask : NSObject, NSCopying {
             dispatch_sync(taskAttributesIsolation) { r = self._countOfBytesSent }
             return r
         }
-        set {
-            dispatch_barrier_async(taskAttributesIsolation) { self._countOfBytesSent = newValue }
-        }
+        set { dispatch_barrier_async(taskAttributesIsolation) { self._countOfBytesSent = newValue } }
     }
     private var _countOfBytesSent: Int64 = 0
 
@@ -629,10 +624,14 @@ public class NSURLSessionTask : NSObject, NSCopying {
      * The current state of the task within the session.
      */
     public var state: NSURLSessionTaskState {
-        //TODO: Make this thread-safe
-        guard suspendCount == 0 else { return .Suspended }
-        NSUnimplemented()
+        get {
+            var r: NSURLSessionTaskState = .Suspended
+            dispatch_sync(taskAttributesIsolation) { r = self._state }
+            return r
+        }
+        set { dispatch_barrier_async(taskAttributesIsolation) { self._state = newValue } }
     }
+    private var _state: NSURLSessionTaskState = .Suspended
 
     /*
      * The error, if any, delivered via -URLSession:task:didCompleteWithError:
@@ -659,25 +658,38 @@ public class NSURLSessionTask : NSObject, NSCopying {
         // 
         // TODO: It may be worth looking into starting over a task that gets
         // resumed. The Darwin Foundation documentation states that that's what
-        // it down for anything but download tasks.
-        //
-        // TODO: is `suspend` / `resume` supposed to be thread safe? If so, we'd
-        // need to use atomic increment / decrement of the suspend count.
-        dispatch_async(workQueue) {
+        // it does for anything but download tasks.
+
+        // We perform the increment and call to `updateTaskState()`
+        // synchronous, to make sure the `state` is updated when this method
+        // returns, but the actual suspend will be done asynchronous to avoid
+        // dead-locks.
+        dispatch_sync(workQueue) {
             self.suspendCount += 1
-            guard self.suspendCount == 1  else { return }
-            self.performSuspend()
+            guard self.suspendCount < Int.max else { fatalError("Task suspended too many times \(Int.max).") }
+            self.updateTaskState()
+            
+            if self.suspendCount == 1 {
+                dispatch_async(self.workQueue) {
+                    self.performSuspend()
+                }
+            }
         }
     }
     /// Resume the task.
     ///
     /// - SeeAlso: `suspend()`
     public func resume() {
-        dispatch_async(workQueue) {
+        dispatch_sync(workQueue) {
             self.suspendCount -= 1
             guard 0 <= self.suspendCount else { fatalError("Resuming a task that's not suspended. Calls to resume() / suspend() need to be matched.") }
-            guard self.suspendCount == 0  else { return }
-            self.performResume()
+            self.updateTaskState()
+            
+            if self.suspendCount == 0 {
+                dispatch_async(self.workQueue) {
+                    self.performResume()
+                }
+            }
         }
     }
 
@@ -709,6 +721,27 @@ public class NSURLSessionTask : NSObject, NSCopying {
     private var _priority: Float = NSURLSessionTaskPriorityDefault
 }
 
+private extension NSURLSessionTask {
+    /// The calls to `suspend` can be nested. This one is only called when the
+    /// task is not suspended and needs to go into suspended state.
+    func performSuspend() {
+        if case .transferInProgress(let transferState) = internalState {
+            internalState = .transferReady(transferState)
+        }
+    }
+    /// The calls to `resume` can be nested. This one is only called when the
+    /// task is suspended and needs to go out of suspended state.
+    func performResume() {
+        if case .initial = internalState {
+            guard let r = originalRequest else { fatalError("Task has no original request.") }
+            startNewTransfer(r)
+        }
+        if case .transferReady(let transferState) = internalState {
+            internalState = .transferInProgress(transferState)
+        }
+    }
+}
+
 internal extension NSURLSessionTask {
     /// The is independent of the public `state: NSURLSessionTaskState`.
     enum InternalState {
@@ -744,28 +777,11 @@ internal extension NSURLSessionTask {
         case waitingForResponseCompletionHandler(TransferState)
         /// The task is completed
         ///
-        /// Contrast this with `.transferComplted`.
+        /// Contrast this with `.transferCompleted`.
         case taskCompleted
     }
-    /// The calls to `suspend` can be nested. This one is only called when the
-    /// task is not suspended and needs to go into suspended state.
-    private func performSuspend() {
-        if case .transferInProgress(let transferState) = internalState {
-            internalState = .transferReady(transferState)
-        }
-    }
-    /// The calls to `resume` can be nested. This one is only called when the
-    /// task is suspended and needs to go out of suspended state.
-    private func performResume() {
-        if case .initial = internalState {
-            guard let r = originalRequest else { fatalError("Task has no original request.") }
-            startNewTransfer(r)
-        }
-        if case .transferReady(let transferState) = internalState {
-            internalState = .transferInProgress(transferState)
-        }
-    }
 }
+
 private extension NSURLSessionTask.InternalState {
     var isEasyHandleAddedToMultiHandle: Bool {
         switch self {
@@ -793,6 +809,25 @@ private extension NSURLSessionTask.InternalState {
     }
 }
 
+internal extension NSURLSessionTask {
+    /// Updates the (public) state based on private / internal state.
+    ///
+    /// - Note: This must be called on the `workQueue`.
+    private func updateTaskState() {
+        func calculateState() -> NSURLSessionTaskState {
+            if case .taskCompleted = internalState {
+                return .Completed
+            }
+            if suspendCount == 0 {
+                return .Running
+            } else {
+                return .Suspended
+            }
+        }
+        state = calculateState()
+    }
+}
+
 private extension NSURLSessionTask {
     enum Body {
         case none

TestFoundation/TestNSURLSession.swift

@@ -208,6 +208,7 @@ extension TestNSURLSession {
                 delegate.completionExpectation = expectationWithDescription("Task did complete")
                 task.resume()
                 waitForExpectationsWithTimeout(0.1, handler: nil)
+                XCTAssertEqual(task.state, NSURLSessionTaskState.Completed)
 
                 guard let error = delegate.completionError else { XCTFail("URL: '\(urlString)'"); return }
                 XCTAssertEqual(error.domain, NSURLErrorDomain, "URL: '\(urlString)'")
@@ -243,6 +244,7 @@ extension TestNSURLSession {
                 delegate.completionExpectation = expectationWithDescription("Task did complete")
                 task.resume()
                 waitForExpectationsWithTimeout(1000, handler: nil)
+                XCTAssertEqual(task.state, NSURLSessionTaskState.Completed)
 
                 guard let error = delegate.completionError else { XCTFail(); return }
                 XCTAssertEqual(error.domain, NSURLErrorDomain)
@@ -279,10 +281,12 @@ extension TestNSURLSession {
                 let task = sut.dataTaskWithURL(originalURL)
 
                 delegate.redirectExpectation = expectationWithDescription("Redirect")
+                XCTAssertEqual(task.state, NSURLSessionTaskState.Suspended)
                 task.resume()
 
                 waitForExpectationsWithTimeout(1.0, handler: nil)
-                
+                XCTAssertEqual(task.state, NSURLSessionTaskState.Running)
+
                 do {
                     XCTAssertEqual(delegate.allResponses.count, 1, "Should only have received the redirect response.")
                     guard let r = delegate.redirect else { XCTFail("No redirect."); return }
@@ -302,6 +306,7 @@ extension TestNSURLSession {
                     r.completionHandler(r.newRequest)
                 }
                 waitForExpectationsWithTimeout(0.1, handler: nil)
+                XCTAssertEqual(task.state, NSURLSessionTaskState.Completed)
 
                 do {
                     XCTAssertEqual(delegate.allResponses.count, 2, "Should have received the redirect response and the final.")
@@ -360,8 +365,11 @@ extension TestNSURLSession {
                     completionExpectation.fulfill()
                 })
 
+                XCTAssertEqual(task.state, NSURLSessionTaskState.Suspended)
                 task.resume()
+                XCTAssertEqual(task.state, NSURLSessionTaskState.Running)
                 waitForExpectationsWithTimeout(0.1, handler: nil)
+                XCTAssertEqual(task.state, NSURLSessionTaskState.Completed)
             }
         }
     }
@@ -371,13 +379,6 @@ extension TestNSURLSession {
     // NSURLSessionUploadTask
     // Requests with a body
     //
-    // public func uploadTaskWithRequest(request: NSURLRequest, fromFile fileURL: NSURL) -> NSURLSessionUploadTask
-    // public func uploadTaskWithRequest(request: NSURLRequest, fromData bodyData: NSData) -> NSURLSessionUploadTask
-    // public func uploadTaskWithStreamedRequest(request: NSURLRequest) -> NSURLSessionUploadTask
-    //
-    // public func uploadTaskWithRequest(request: NSURLRequest, fromFile fileURL: NSURL, completionHandler: (NSData?, NSURLResponse?, NSError?) -> Void) -> NSURLSessionUploadTask
-    // public func uploadTaskWithRequest(request: NSURLRequest, fromData bodyData: NSData?, completionHandler: (NSData?, NSURLResponse?, NSError?) -> Void) -> NSURLSessionUploadTask
-    //
 
     /// Test a task with a completion handler performs a `POST` request.
     func test_functional_POST_fromData() {
@@ -543,6 +544,10 @@ extension TestNSURLSession {
             }
         }
     }
+    
+    //TODO: Uploading data with completion handler
+    // public func uploadTaskWithRequest(request: NSURLRequest, fromFile fileURL: NSURL, completionHandler: (NSData?, NSURLResponse?, NSError?) -> Void) -> NSURLSessionUploadTask
+    // public func uploadTaskWithRequest(request: NSURLRequest, fromData bodyData: NSData?, completionHandler: (NSData?, NSURLResponse?, NSError?) -> Void) -> NSURLSessionUploadTask
 }
 
 private extension HTTPResponse {