Add async URLSession methods (swiftlang#4970)

* Add `data(from:delegate:)` method.

* Add async URLSession methods

---------

Co-authored-by: ichiho <[email protected]>

Author: jrflat

Sources/FoundationNetworking/DataURLProtocol.swift

@@ -91,8 +91,7 @@ internal class _DataURLProtocol: URLProtocol {
             urlClient.urlProtocolDidFinishLoading(self)
         } else {
             let error = NSError(domain: NSURLErrorDomain, code: NSURLErrorBadURL)
-            if let session = self.task?.session as? URLSession, let delegate = session.delegate as? URLSessionTaskDelegate,
-                let task = self.task {
+            if let task = self.task, let session = task.actualSession, let delegate = task.delegate {
                 delegate.urlSession(session, task: task, didCompleteWithError: error)
             }
         }

Sources/FoundationNetworking/URLSession/FTP/FTPURLProtocol.swift

@@ -119,7 +119,7 @@ internal extension _FTPURLProtocol {
         switch session.behaviour(for: self.task!) {
         case .noDelegate:
             break
-        case .taskDelegate:
+        case .taskDelegate, .dataCompletionHandlerWithTaskDelegate, .downloadCompletionHandlerWithTaskDelegate:
             self.client?.urlProtocol(self, didReceive: response, cacheStoragePolicy: .notAllowed)
         case .dataCompletionHandler:
             break

Sources/FoundationNetworking/URLSession/HTTP/HTTPURLProtocol.swift

@@ -475,7 +475,7 @@ internal class _HTTPURLProtocol: _NativeProtocol {
 
         guard let session = task?.session as? URLSession else { fatalError() }
 
-        if let delegate = session.delegate as? URLSessionTaskDelegate {
+        if let delegate = task?.delegate {
             // At this point we need to change the internal state to note
             // that we're waiting for the delegate to call the completion
             // handler. Then we'll call the delegate callback
@@ -524,7 +524,9 @@ internal class _HTTPURLProtocol: _NativeProtocol {
         switch session.behaviour(for: self.task!) {
         case .noDelegate:
             break
-        case .taskDelegate:
+        case .taskDelegate,
+             .dataCompletionHandlerWithTaskDelegate,
+             .downloadCompletionHandlerWithTaskDelegate:
             //TODO: There's a problem with libcurl / with how we're using it.
             // We're currently unable to pause the transfer / the easy handle:
             // https://curl.haxx.se/mail/lib-2016-03/0222.html

Sources/FoundationNetworking/URLSession/NativeProtocol.swift

@@ -129,43 +129,59 @@ internal class _NativeProtocol: URLProtocol, _EasyHandleDelegate {
     }
 
     fileprivate func notifyDelegate(aboutReceivedData data: Data) {
-        guard let t = self.task else {
+        guard let task = self.task, let session = task.session as? URLSession else {
             fatalError("Cannot notify")
         }
-        if case .taskDelegate(let delegate) = t.session.behaviour(for: self.task!),
-            let dataDelegate = delegate as? URLSessionDataDelegate,
-            let task = self.task as? URLSessionDataTask {
-            // Forward to the delegate:
-            guard let s = self.task?.session as? URLSession else {
-                fatalError()
-            }
-            s.delegateQueue.addOperation {
-                dataDelegate.urlSession(s, dataTask: task, didReceive: data)
-            }
-        } else if case .taskDelegate(let delegate) = t.session.behaviour(for: self.task!),
-            let downloadDelegate = delegate as? URLSessionDownloadDelegate,
-            let task = self.task as? URLSessionDownloadTask {
-            guard let s = self.task?.session as? URLSession else {
-                fatalError()
-            }
-            let fileHandle = try! FileHandle(forWritingTo: self.tempFileURL)
-            _ = fileHandle.seekToEndOfFile()
-            fileHandle.write(data)
-            task.countOfBytesReceived  += Int64(data.count)
-            s.delegateQueue.addOperation {
-                downloadDelegate.urlSession(s, downloadTask: task, didWriteData: Int64(data.count), totalBytesWritten: task.countOfBytesReceived,
-                                            totalBytesExpectedToWrite: task.countOfBytesExpectedToReceive)
+        switch task.session.behaviour(for: task) {
+        case .taskDelegate(let delegate),
+             .dataCompletionHandlerWithTaskDelegate(_, let delegate),
+             .downloadCompletionHandlerWithTaskDelegate(_, let delegate):
+            if let dataDelegate = delegate as? URLSessionDataDelegate,
+               let dataTask = task as? URLSessionDataTask {
+                session.delegateQueue.addOperation {
+                    dataDelegate.urlSession(session, dataTask: dataTask, didReceive: data)
+                }
+            } else if let downloadDelegate = delegate as? URLSessionDownloadDelegate,
+                      let downloadTask = task as? URLSessionDownloadTask {
+                let fileHandle = try! FileHandle(forWritingTo: self.tempFileURL)
+                _ = fileHandle.seekToEndOfFile()
+                fileHandle.write(data)
+                task.countOfBytesReceived  += Int64(data.count)
+                session.delegateQueue.addOperation {
+                    downloadDelegate.urlSession(
+                        session,
+                        downloadTask: downloadTask,
+                        didWriteData: Int64(data.count),
+                        totalBytesWritten: task.countOfBytesReceived,
+                        totalBytesExpectedToWrite: task.countOfBytesExpectedToReceive
+                    )
+                }
             }
+        default:
+            break
         }
     }
 
     fileprivate func notifyDelegate(aboutUploadedData count: Int64) {
-        guard let task = self.task, let session = task.session as? URLSession,
-            case .taskDelegate(let delegate) = session.behaviour(for: task) else { return }
-        task.countOfBytesSent += count
-        session.delegateQueue.addOperation {
-            delegate.urlSession(session, task: task, didSendBodyData: count,
-                                totalBytesSent: task.countOfBytesSent, totalBytesExpectedToSend: task.countOfBytesExpectedToSend)
+        guard let task = self.task, let session = task.session as? URLSession else {
+            return
+        }
+        switch session.behaviour(for: task) {
+        case .taskDelegate(let delegate),
+             .dataCompletionHandlerWithTaskDelegate(_, let delegate),
+             .downloadCompletionHandlerWithTaskDelegate(_, let delegate):
+            task.countOfBytesSent += count
+            session.delegateQueue.addOperation {
+                delegate.urlSession(
+                    session,
+                    task: task,
+                    didSendBodyData: count,
+                    totalBytesSent: task.countOfBytesSent,
+                    totalBytesExpectedToSend: task.countOfBytesExpectedToSend
+                )
+            }
+        default:
+            break
         }
     }
 
@@ -284,7 +300,7 @@ internal class _NativeProtocol: URLProtocol, _EasyHandleDelegate {
 
         var currentInputStream: InputStream?
 
-        if let delegate = session.delegate as? URLSessionTaskDelegate {
+        if let delegate = task?.delegate {
             let dispatchGroup = DispatchGroup()
             dispatchGroup.enter()
 
@@ -338,11 +354,13 @@ internal class _NativeProtocol: URLProtocol, _EasyHandleDelegate {
             // Data will be forwarded to the delegate as we receive it, we don't
             // need to do anything about it.
             return .ignore
-        case .dataCompletionHandler:
+        case .dataCompletionHandler,
+             .dataCompletionHandlerWithTaskDelegate:
             // Data needs to be concatenated in-memory such that we can pass it
             // to the completion handler upon completion.
             return .inMemory(nil)
-        case .downloadCompletionHandler:
+        case .downloadCompletionHandler,
+             .downloadCompletionHandlerWithTaskDelegate:
             // Data needs to be written to a file (i.e. a download task).
             let fileHandle = try! FileHandle(forWritingTo: self.tempFileURL)
             return .toFile(self.tempFileURL, fileHandle)

Sources/FoundationNetworking/URLSession/TaskRegistry.swift

@@ -45,8 +45,12 @@ extension URLSession {
             case callDelegate
             /// Default action for all events, except for completion.
             case dataCompletionHandler(DataTaskCompletion)
+            /// Default action for all asynchronous events.
+            case dataCompletionHandlerWithTaskDelegate(DataTaskCompletion, URLSessionTaskDelegate?)
             /// Default action for all events, except for completion.
             case downloadCompletionHandler(DownloadTaskCompletion)
+            /// Default action for all asynchronous events.
+            case downloadCompletionHandlerWithTaskDelegate(DownloadTaskCompletion, URLSessionTaskDelegate?)
         }
 
         fileprivate var tasks: [Int: URLSessionTask] = [:]