Add a shared file IO thread pool per HTTPClient

Author: dnadoba

Sources/AsyncHTTPClient/FileDownloadDelegate.swift

@@ -30,18 +30,10 @@ public final class FileDownloadDelegate: HTTPClientResponseDelegate {
     public typealias Response = Progress
 
     private let filePath: String
-    private let io: NonBlockingFileIO
+    private var io: NonBlockingFileIO?
     private let reportHead: ((HTTPResponseHead) -> Void)?
     private let reportProgress: ((Progress) -> Void)?
 
-    private enum ThreadPool {
-        case unowned
-        // if we own the thread pool we also need to shut it down
-        case owned(NIOThreadPool)
-    }
-
-    private let threadPool: ThreadPool
-
     private var fileHandleFuture: EventLoopFuture<NIOFileHandle>?
     private var writeFuture: EventLoopFuture<Void>?
 
@@ -61,7 +53,7 @@ public final class FileDownloadDelegate: HTTPClientResponseDelegate {
         reportHead: ((HTTPResponseHead) -> Void)? = nil,
         reportProgress: ((Progress) -> Void)? = nil
     ) throws {
-        try self.init(path: path, sharedThreadPool: pool, reportHead: reportHead, reportProgress: reportProgress)
+        try self.init(path: path, pool: .some(pool), reportHead: reportHead, reportProgress: reportProgress)
     }
 
     /// Initializes a new file download delegate and spawns a new thread for file I/O.
@@ -78,31 +70,36 @@ public final class FileDownloadDelegate: HTTPClientResponseDelegate {
         reportHead: ((HTTPResponseHead) -> Void)? = nil,
         reportProgress: ((Progress) -> Void)? = nil
     ) throws {
-        try self.init(path: path, sharedThreadPool: nil, reportHead: reportHead, reportProgress: reportProgress)
+        try self.init(path: path, pool: nil, reportHead: reportHead, reportProgress: reportProgress)
     }
 
     private init(
         path: String,
-        sharedThreadPool: NIOThreadPool?,
+        pool: NIOThreadPool?,
         reportHead: ((HTTPResponseHead) -> Void)? = nil,
         reportProgress: ((Progress) -> Void)? = nil
     ) throws {
-        let pool: NIOThreadPool
-        if let sharedThreadPool = sharedThreadPool {
-            pool = sharedThreadPool
-            self.threadPool = .unowned
+        if let pool = pool {
+            self.io = NonBlockingFileIO(threadPool: pool)
         } else {
-            pool = NIOThreadPool(numberOfThreads: 1)
-            self.threadPool = .owned(pool)
+            // we should use the shared thread pool from the HTTPClient which
+            // we will get shortly through a call to provideSharedThreadPool(fileIOPool:)
+            self.io = nil
         }
-
-        pool.start()
-        self.io = NonBlockingFileIO(threadPool: pool)
+        
         self.filePath = path
 
         self.reportHead = reportHead
         self.reportProgress = reportProgress
     }
+    
+    public func provideSharedThreadPool(fileIOPool: NIOThreadPool) {
+        guard self.io == nil else {
+            // user has provided their own thread pool
+            return
+        }
+        self.io = NonBlockingFileIO(threadPool: fileIOPool)
+    }
 
     public func didReceiveHead(
         task: HTTPClient.Task<Response>,
@@ -122,24 +119,27 @@ public final class FileDownloadDelegate: HTTPClientResponseDelegate {
         task: HTTPClient.Task<Response>,
         _ buffer: ByteBuffer
     ) -> EventLoopFuture<Void> {
+        guard let io = io else {
+            preconditionFailure("thread pool not provided by HTTPClient before calling \(#function)")
+        }
         self.progress.receivedBytes += buffer.readableBytes
         self.reportProgress?(self.progress)
 
         let writeFuture: EventLoopFuture<Void>
         if let fileHandleFuture = self.fileHandleFuture {
             writeFuture = fileHandleFuture.flatMap {
-                self.io.write(fileHandle: $0, buffer: buffer, eventLoop: task.eventLoop)
+                io.write(fileHandle: $0, buffer: buffer, eventLoop: task.eventLoop)
             }
         } else {
-            let fileHandleFuture = self.io.openFile(
+            let fileHandleFuture = io.openFile(
                 path: self.filePath,
                 mode: .write,
                 flags: .allowFileCreation(),
                 eventLoop: task.eventLoop
             )
             self.fileHandleFuture = fileHandleFuture
             writeFuture = fileHandleFuture.flatMap {
-                self.io.write(fileHandle: $0, buffer: buffer, eventLoop: task.eventLoop)
+                io.write(fileHandle: $0, buffer: buffer, eventLoop: task.eventLoop)
             }
         }
 
@@ -150,12 +150,6 @@ public final class FileDownloadDelegate: HTTPClientResponseDelegate {
     private func close(fileHandle: NIOFileHandle) {
         try! fileHandle.close()
         self.fileHandleFuture = nil
-        switch self.threadPool {
-        case .unowned:
-            break
-        case .owned(let pool):
-            try! pool.syncShutdownGracefully()
-        }
     }
 
     private func finalize() {
@@ -177,14 +171,4 @@ public final class FileDownloadDelegate: HTTPClientResponseDelegate {
         self.finalize()
         return self.progress
     }
-
-    deinit {
-        switch threadPool {
-        case .unowned:
-            break
-        case .owned(let pool):
-            // if the delegate is unused we still need to shutdown the thread pool
-            try! pool.syncShutdownGracefully()
-        }
-    }
 }

Sources/AsyncHTTPClient/HTTPClient.swift

@@ -72,6 +72,9 @@ public class HTTPClient {
     let eventLoopGroupProvider: EventLoopGroupProvider
     let configuration: Configuration
     let poolManager: HTTPConnectionPool.Manager
+    
+    /// Shared thread pool used for file IO. It is given to the user through ``HTTPClientResponseDelegate/provideSharedThreadPool(fileIOPool:)-6phmu``
+    private let fileIOThreadPool = NIOThreadPool(numberOfThreads: 1)
     private var state: State
     private let stateLock = Lock()
 
@@ -97,6 +100,7 @@ public class HTTPClient {
     public required init(eventLoopGroupProvider: EventLoopGroupProvider,
                          configuration: Configuration = Configuration(),
                          backgroundActivityLogger: Logger) {
+        self.fileIOThreadPool.start()
         self.eventLoopGroupProvider = eventLoopGroupProvider
         switch self.eventLoopGroupProvider {
         case .shared(let group):
@@ -241,11 +245,13 @@ public class HTTPClient {
                     let error: Error? = (requiresClean && unclean) ? HTTPClientError.uncleanShutdown : nil
                     return (callback, error)
                 }
-
-                self.shutdownEventLoop(queue: queue) { error in
-                    let reportedError = error ?? uncleanError
-                    callback(reportedError)
+                self.fileIOThreadPool.shutdownGracefully(queue: queue) { ioThreadPoolError in
+                    self.shutdownEventLoop(queue: queue) { error in
+                        let reportedError = error ?? ioThreadPoolError ?? uncleanError
+                        callback(reportedError)
+                    }
                 }
+                
             }
         }
     }
@@ -562,9 +568,12 @@ public class HTTPClient {
         case .testOnly_exact(_, delegateOn: let delegateEL):
             taskEL = delegateEL
         }
+        
         logger.trace("selected EventLoop for task given the preference",
                      metadata: ["ahc-eventloop": "\(taskEL)",
                                 "ahc-el-preference": "\(eventLoopPreference)"])
+        
+        delegate.provideSharedThreadPool(fileIOPool: self.fileIOThreadPool)
 
         let failedTask: Task<Delegate.Response>? = self.stateLock.withLock {
             switch state {

Sources/AsyncHTTPClient/HTTPHandler.swift

@@ -16,6 +16,7 @@ import Foundation
 import Logging
 import NIOConcurrencyHelpers
 import NIOCore
+import NIOPosix
 import NIOHTTP1
 import NIOSSL
 
@@ -430,7 +431,14 @@ public class ResponseAccumulator: HTTPClientResponseDelegate {
 ///          object that implements this protocol, but may do so if needed.
 public protocol HTTPClientResponseDelegate: AnyObject {
     associatedtype Response
-
+    
+    /// Called exactly once before any other methods of this delegate are called.
+    /// It is used to give access to the shared thread pool of the ``HTTPClient`` the request is executed on.
+    /// Use this thread pool to do file IO associated with the request e.g. to save a response to disk.
+    ///
+    /// - Parameter fileIOPool: File IO Pool
+    func provideSharedThreadPool(fileIOPool: NIOThreadPool)
+    
     /// Called when the request head is sent. Will be called once.
     ///
     /// - parameters:
@@ -502,7 +510,13 @@ public protocol HTTPClientResponseDelegate: AnyObject {
 }
 
 extension HTTPClientResponseDelegate {
-    /// Default implementation of ``HTTPClientResponseDelegate/didSendRequestHead(task:_:)-6khai``.
+    
+    /// Default implementation of ``HTTPClientResponseDelegate/provideSharedThreadPool(fileIOPool:)-8y1b``
+    ///
+    /// By default, this does nothing.
+    public func provideSharedThreadPool(fileIOPool: NIOThreadPool) {}
+    
+    /// Default implementation of ``HTTPClientResponseDelegate/didSendRequest(task:)-9od5p``.
     ///
     /// By default, this does nothing.
     public func didSendRequestHead(task: HTTPClient.Task<Response>, _ head: HTTPRequestHead) {}