add progress tracking to http client (#3255)

motivation:
* safer handling of large responses
* prepare to consolidate Downloader from TSC and http client from Basics

changes:
* add progress support to HTTPclient::execute and URLSession based impl
* add HTTPClientRequest::Options::maxResponseSize and logic to fail the request if response is too large
* add authorization support to HTTPClientRequest::Options and HTTPClient::Configuration
* refactor call-sites to use new ways
* add and adopt tests

Author: tomerd

Sources/Basics/ConcurrencyHelpers.swift

@@ -6,8 +6,8 @@
  See http://swift.org/CONTRIBUTORS.txt for Swift project authors
  */
 
-import TSCBasic
 import class Foundation.ProcessInfo
+import TSCBasic
 
 /// Thread-safe dictionary like structure
 public final class ThreadSafeKeyValueStore<Key, Value> where Key: Hashable {
@@ -37,6 +37,12 @@ public final class ThreadSafeKeyValueStore<Key, Value> where Key: Hashable {
         }
     }
 
+    public func removeValue(forKey key: Key) -> Value? {
+        self.lock.withLock {
+            self.underlying.removeValue(forKey: key)
+        }
+    }
+
     public func clear() {
         self.lock.withLock {
             self.underlying.removeAll()

Sources/Basics/HTPClient+URLSession.swift

@@ -1,4 +1,12 @@
+/*
+ This source file is part of the Swift.org open source project
 
+ Copyright (c) 2020 Apple Inc. and the Swift project authors
+ Licensed under Apache License v2.0 with Runtime Library Exception
+
+ See http://swift.org/LICENSE.txt for license information
+ See http://swift.org/CONTRIBUTORS.txt for Swift project authors
+ */
 
 import Foundation
 import struct TSCUtility.Versioning
@@ -8,28 +16,84 @@ import struct TSCUtility.Versioning
 import FoundationNetworking
 #endif
 
-public struct URLSessionHTTPClient: HTTPClientProtocol {
+public final class URLSessionHTTPClient: NSObject, HTTPClientProtocol {
     private let configuration: URLSessionConfiguration
+    private let delegateQueue: OperationQueue
+    private var session: URLSession!
+    private var tasks = ThreadSafeKeyValueStore<Int, DataTask>()
 
     public init(configuration: URLSessionConfiguration = .default) {
         self.configuration = configuration
+        self.delegateQueue = OperationQueue()
+        self.delegateQueue.name = "org.swift.swiftpm.urlsession-http-client"
+        self.delegateQueue.maxConcurrentOperationCount = 1
+        super.init()
+        self.session = URLSession(configuration: self.configuration, delegate: self, delegateQueue: self.delegateQueue)
     }
 
-    public func execute(_ request: HTTPClient.Request, callback: @escaping (Result<HTTPClient.Response, Error>) -> Void) {
-        let session = URLSession(configuration: self.configuration)
-        let task = session.dataTask(with: request.urlRequest()) { data, response, error in
-            if let error = error {
-                callback(.failure(error))
-            } else if let response = response as? HTTPURLResponse {
-                callback(.success(response.response(body: data)))
-            } else {
-                callback(.failure(HTTPClientError.invalidResponse))
-            }
-        }
+    public func execute(_ request: HTTPClient.Request, progress: ProgressHandler?, completion: @escaping CompletionHandler) {
+        let task = self.session.dataTask(with: request.urlRequest())
+        self.tasks[task.taskIdentifier] = DataTask(task: task, progressHandler: progress, completionHandler: completion)
         task.resume()
     }
 }
 
+extension URLSessionHTTPClient: URLSessionDataDelegate {
+    public func urlSession(_ session: URLSession,
+                           dataTask: URLSessionDataTask,
+                           didReceive response: URLResponse,
+                           completionHandler: @escaping (URLSession.ResponseDisposition) -> Void) {
+        guard let task = self.tasks[dataTask.taskIdentifier] else {
+            return completionHandler(.cancel)
+        }
+        task.response = response as? HTTPURLResponse
+        task.expectedContentLength = response.expectedContentLength
+        task.progressHandler?(0, response.expectedContentLength)
+        completionHandler(.allow)
+    }
+
+    public func urlSession(_ session: URLSession, dataTask: URLSessionDataTask, didReceive data: Data) {
+        guard let task = self.tasks[dataTask.taskIdentifier] else {
+            return
+        }
+        if task.buffer != nil {
+            task.buffer?.append(data)
+        } else {
+            task.buffer = data
+        }
+        task.progressHandler?(Int64(task.buffer?.count ?? 0), task.expectedContentLength) // safe since created in the line above
+    }
+
+    public func urlSession(_ session: URLSession, task: URLSessionTask, didCompleteWithError error: Error?) {
+        guard let task = self.tasks.removeValue(forKey: task.taskIdentifier) else {
+            return
+        }
+        if let error = error {
+            task.completionHandler(.failure(error))
+        } else if let response = task.response {
+            task.completionHandler(.success(response.response(body: task.buffer)))
+        } else {
+            task.completionHandler(.failure(HTTPClientError.invalidResponse))
+        }
+    }
+
+    class DataTask {
+        let task: URLSessionDataTask
+        let completionHandler: CompletionHandler
+        let progressHandler: ProgressHandler?
+
+        var response: HTTPURLResponse?
+        var expectedContentLength: Int64?
+        var buffer: Data?
+
+        init(task: URLSessionDataTask, progressHandler: ProgressHandler?, completionHandler: @escaping CompletionHandler) {
+            self.task = task
+            self.progressHandler = progressHandler
+            self.completionHandler = completionHandler
+        }
+    }
+}
+
 extension HTTPClient.Request {
     func urlRequest() -> URLRequest {
         var request = URLRequest(url: self.url)

Sources/Basics/HTTPClient.swift

@@ -25,13 +25,24 @@ import CRT
 #endif
 
 public protocol HTTPClientProtocol {
-    func execute(_ request: HTTPClientRequest, callback: @escaping (Result<HTTPClientResponse, Error>) -> Void)
+    typealias ProgressHandler = (_ bytesReceived: Int64, _ totalBytes: Int64?) -> Void
+    typealias CompletionHandler = (Result<HTTPClientResponse, Error>) -> Void
+
+    /// Execute an HTTP request asynchronously
+    ///
+    /// - Parameters:
+    ///   - request: The `HTTPClientRequest` to perform.
+    ///   - callback: A closure to be notified of the completion of the request.
+    func execute(_ request: HTTPClientRequest,
+                 progress: ProgressHandler?,
+                 completion: @escaping CompletionHandler)
 }
 
 public enum HTTPClientError: Error, Equatable {
     case invalidResponse
     case badResponseStatusCode(Int)
     case circuitBreakerTriggered
+    case responseTooLarge(Int64)
 }
 
 // MARK: - HTTPClient
@@ -40,7 +51,7 @@ public struct HTTPClient: HTTPClientProtocol {
     public typealias Configuration = HTTPClientConfiguration
     public typealias Request = HTTPClientRequest
     public typealias Response = HTTPClientResponse
-    public typealias Handler = (Request, @escaping (Result<Response, Error>) -> Void) -> Void
+    public typealias Handler = (Request, ProgressHandler?, @escaping (Result<Response, Error>) -> Void) -> Void
 
     public var configuration: HTTPClientConfiguration
     private let diagnosticsEngine: DiagnosticsEngine?
@@ -57,7 +68,7 @@ public struct HTTPClient: HTTPClientProtocol {
         self.underlying = handler ?? URLSessionHTTPClient().execute
     }
 
-    public func execute(_ request: Request, callback: @escaping (Result<Response, Error>) -> Void) {
+    public func execute(_ request: Request, progress: ProgressHandler? = nil, completion: @escaping CompletionHandler) {
         // merge configuration
         var request = request
         if request.options.callbackQueue == nil {
@@ -72,6 +83,9 @@ public struct HTTPClient: HTTPClientProtocol {
         if request.options.timeout == nil {
             request.options.timeout = self.configuration.requestTimeout
         }
+        if request.options.authorizationProvider == nil {
+            request.options.authorizationProvider = self.configuration.authorizationProvider
+        }
         // add additional headers
         if let additionalHeaders = self.configuration.requestHeaders {
             additionalHeaders.forEach {
@@ -81,43 +95,68 @@ public struct HTTPClient: HTTPClientProtocol {
         if request.options.addUserAgent, !request.headers.contains("User-Agent") {
             request.headers.add(name: "User-Agent", value: "SwiftPackageManager/\(SwiftVersion.currentVersion.displayString)")
         }
+        if let authorization = request.options.authorizationProvider?(request.url), !request.headers.contains("Authorization") {
+            request.headers.add(name: "Authorization", value: authorization)
+        }
         // execute
-        self._execute(request: request, requestNumber: 0) { result in
-            let callbackQueue = request.options.callbackQueue ?? self.configuration.callbackQueue
-            callbackQueue.async {
-                callback(result)
+        let callbackQueue = request.options.callbackQueue ?? self.configuration.callbackQueue
+        self._execute(
+            request: request, requestNumber: 0,
+            progress: progress.map { handler in
+                { received, expected in
+                    callbackQueue.async {
+                        handler(received, expected)
+                    }
+                }
+            },
+            completion: { result in
+                callbackQueue.async {
+                    completion(result)
+                }
             }
-        }
+        )
     }
 
-    private func _execute(request: Request, requestNumber: Int, callback: @escaping (Result<Response, Error>) -> Void) {
+    private func _execute(request: Request, requestNumber: Int, progress: ProgressHandler?, completion: @escaping CompletionHandler) {
         if self.shouldCircuitBreak(request: request) {
             diagnosticsEngine?.emit(warning: "Circuit breaker triggered for \(request.url)")
-            return callback(.failure(HTTPClientError.circuitBreakerTriggered))
+            return completion(.failure(HTTPClientError.circuitBreakerTriggered))
         }
 
-        self.underlying(request) { result in
-            switch result {
-            case .failure(let error):
-                callback(.failure(error))
-            case .success(let response):
-                // record host errors for circuit breaker
-                self.recordErrorIfNecessary(response: response, request: request)
-                // handle retry strategy
-                if let retryDelay = self.shouldRetry(response: response, request: request, requestNumber: requestNumber) {
-                    self.diagnosticsEngine?.emit(warning: "\(request.url) failed, retrying in \(retryDelay)")
-                    // TODO: dedicated retry queue?
-                    return self.configuration.callbackQueue.asyncAfter(deadline: .now() + retryDelay) {
-                        self._execute(request: request, requestNumber: requestNumber + 1, callback: callback)
+        self.underlying(
+            request,
+            { received, expected in
+                if let max = request.options.maximumResponseSizeInBytes {
+                    guard received < max else {
+                        // FIXME: cancel the request?
+                        return completion(.failure(HTTPClientError.responseTooLarge(received)))
                     }
                 }
-                // check for valid response codes
-                if let validResponseCodes = request.options.validResponseCodes, !validResponseCodes.contains(response.statusCode) {
-                    return callback(.failure(HTTPClientError.badResponseStatusCode(response.statusCode)))
+                progress?(received, expected)
+            },
+            { result in
+                switch result {
+                case .failure(let error):
+                    completion(.failure(error))
+                case .success(let response):
+                    // record host errors for circuit breaker
+                    self.recordErrorIfNecessary(response: response, request: request)
+                    // handle retry strategy
+                    if let retryDelay = self.shouldRetry(response: response, request: request, requestNumber: requestNumber) {
+                        self.diagnosticsEngine?.emit(warning: "\(request.url) failed, retrying in \(retryDelay)")
+                        // TODO: dedicated retry queue?
+                        return self.configuration.callbackQueue.asyncAfter(deadline: .now() + retryDelay) {
+                            self._execute(request: request, requestNumber: requestNumber + 1, progress: progress, completion: completion)
+                        }
+                    }
+                    // check for valid response codes
+                    if let validResponseCodes = request.options.validResponseCodes, !validResponseCodes.contains(response.statusCode) {
+                        return completion(.failure(HTTPClientError.badResponseStatusCode(response.statusCode)))
+                    }
+                    completion(.success(response))
                 }
-                callback(.success(response))
             }
-        }
+        )
     }
 
     private func shouldRetry(response: Response, request: Request, requestNumber: Int) -> DispatchTimeInterval? {
@@ -179,39 +218,43 @@ public struct HTTPClient: HTTPClientProtocol {
 }
 
 public extension HTTPClient {
-    func head(_ url: URL, headers: HTTPClientHeaders = .init(), options: Request.Options = .init(), callback: @escaping (Result<Response, Error>) -> Void) {
-        self.execute(Request(method: .head, url: url, headers: headers, body: nil, options: options), callback: callback)
+    func head(_ url: URL, headers: HTTPClientHeaders = .init(), options: Request.Options = .init(), completion: @escaping (Result<Response, Error>) -> Void) {
+        self.execute(Request(method: .head, url: url, headers: headers, body: nil, options: options), completion: completion)
     }
 
-    func get(_ url: URL, headers: HTTPClientHeaders = .init(), options: Request.Options = .init(), callback: @escaping (Result<Response, Error>) -> Void) {
-        self.execute(Request(method: .get, url: url, headers: headers, body: nil, options: options), callback: callback)
+    func get(_ url: URL, headers: HTTPClientHeaders = .init(), options: Request.Options = .init(), completion: @escaping (Result<Response, Error>) -> Void) {
+        self.execute(Request(method: .get, url: url, headers: headers, body: nil, options: options), completion: completion)
     }
 
-    func put(_ url: URL, body: Data?, headers: HTTPClientHeaders = .init(), options: Request.Options = .init(), callback: @escaping (Result<Response, Error>) -> Void) {
-        self.execute(Request(method: .put, url: url, headers: headers, body: body, options: options), callback: callback)
+    func put(_ url: URL, body: Data?, headers: HTTPClientHeaders = .init(), options: Request.Options = .init(), completion: @escaping (Result<Response, Error>) -> Void) {
+        self.execute(Request(method: .put, url: url, headers: headers, body: body, options: options), completion: completion)
     }
 
-    func post(_ url: URL, body: Data?, headers: HTTPClientHeaders = .init(), options: Request.Options = .init(), callback: @escaping (Result<Response, Error>) -> Void) {
-        self.execute(Request(method: .post, url: url, headers: headers, body: body, options: options), callback: callback)
+    func post(_ url: URL, body: Data?, headers: HTTPClientHeaders = .init(), options: Request.Options = .init(), completion: @escaping (Result<Response, Error>) -> Void) {
+        self.execute(Request(method: .post, url: url, headers: headers, body: body, options: options), completion: completion)
     }
 
-    func delete(_ url: URL, headers: HTTPClientHeaders = .init(), options: Request.Options = .init(), callback: @escaping (Result<Response, Error>) -> Void) {
-        self.execute(Request(method: .delete, url: url, headers: headers, body: nil, options: options), callback: callback)
+    func delete(_ url: URL, headers: HTTPClientHeaders = .init(), options: Request.Options = .init(), completion: @escaping (Result<Response, Error>) -> Void) {
+        self.execute(Request(method: .delete, url: url, headers: headers, body: nil, options: options), completion: completion)
     }
 }
 
 // MARK: - HTTPClientConfiguration
 
+public typealias HTTPClientAuthorizationProvider = (URL) -> String?
+
 public struct HTTPClientConfiguration {
     public var requestHeaders: HTTPClientHeaders?
     public var requestTimeout: DispatchTimeInterval?
+    public var authorizationProvider: HTTPClientAuthorizationProvider?
     public var retryStrategy: HTTPClientRetryStrategy?
     public var circuitBreakerStrategy: HTTPClientCircuitBreakerStrategy?
     public var callbackQueue: DispatchQueue
 
     public init() {
         self.requestHeaders = .none
         self.requestTimeout = .none
+        self.authorizationProvider = .none
         self.retryStrategy = .none
         self.circuitBreakerStrategy = .none
         self.callbackQueue = .global()
@@ -259,6 +302,8 @@ public struct HTTPClientRequest {
         public var addUserAgent: Bool
         public var validResponseCodes: [Int]?
         public var timeout: DispatchTimeInterval?
+        public var maximumResponseSizeInBytes: Int64?
+        public var authorizationProvider: HTTPClientAuthorizationProvider?
         public var retryStrategy: HTTPClientRetryStrategy?
         public var circuitBreakerStrategy: HTTPClientCircuitBreakerStrategy?
         public var callbackQueue: DispatchQueue?
@@ -267,6 +312,8 @@ public struct HTTPClientRequest {
             self.addUserAgent = true
             self.validResponseCodes = .none
             self.timeout = .none
+            self.maximumResponseSizeInBytes = .none
+            self.authorizationProvider = .none
             self.retryStrategy = .none
             self.circuitBreakerStrategy = .none
             self.callbackQueue = .none