Skip to content

Commit c6a1a3d

Browse files
committed
Fix thread leak in FileDownloadDelegate
1 parent 5e3e58d commit c6a1a3d

File tree

1 file changed

+54
-4
lines changed

1 file changed

+54
-4
lines changed

Sources/AsyncHTTPClient/FileDownloadDelegate.swift

Lines changed: 54 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -33,6 +33,13 @@ public final class FileDownloadDelegate: HTTPClientResponseDelegate {
3333
private let io: NonBlockingFileIO
3434
private let reportHead: ((HTTPResponseHead) -> Void)?
3535
private let reportProgress: ((Progress) -> Void)?
36+
37+
private enum ThreadPool {
38+
case unowned
39+
// if we own the thread pool we also need to shut it down
40+
case owned(NIOThreadPool)
41+
}
42+
private let threadPool: ThreadPool
3643

3744
private var fileHandleFuture: EventLoopFuture<NIOFileHandle>?
3845
private var writeFuture: EventLoopFuture<Void>?
@@ -47,19 +54,46 @@ public final class FileDownloadDelegate: HTTPClientResponseDelegate {
4754
/// the total byte count and download byte count passed to it as arguments. The callbacks
4855
/// will be invoked in the same threading context that the delegate itself is invoked,
4956
/// as controlled by `EventLoopPreference`.
50-
public init(
57+
public convenience init(
58+
path: String,
59+
pool: NIOThreadPool,
60+
reportHead: ((HTTPResponseHead) -> Void)? = nil,
61+
reportProgress: ((Progress) -> Void)? = nil
62+
) throws {
63+
try self.init(path: path, pool: .some(pool), reportHead: reportHead, reportProgress: reportProgress)
64+
}
65+
66+
public convenience init(
67+
path: String,
68+
reportHead: ((HTTPResponseHead) -> Void)? = nil,
69+
reportProgress: ((Progress) -> Void)? = nil
70+
) throws {
71+
try self.init(path: path, pool: nil, reportHead: reportHead, reportProgress: reportProgress)
72+
}
73+
74+
private init(
5175
path: String,
52-
pool: NIOThreadPool = NIOThreadPool(numberOfThreads: 1),
76+
pool: NIOThreadPool?,
5377
reportHead: ((HTTPResponseHead) -> Void)? = nil,
5478
reportProgress: ((Progress) -> Void)? = nil
5579
) throws {
56-
pool.start()
57-
self.io = NonBlockingFileIO(threadPool: pool)
80+
let unwrappedPool: NIOThreadPool
81+
if let pool = pool {
82+
unwrappedPool = pool
83+
self.threadPool = .unowned
84+
} else {
85+
unwrappedPool = NIOThreadPool(numberOfThreads: 1)
86+
self.threadPool = .owned(unwrappedPool)
87+
}
88+
89+
unwrappedPool.start()
90+
self.io = NonBlockingFileIO(threadPool: unwrappedPool)
5891
self.filePath = path
5992

6093
self.reportHead = reportHead
6194
self.reportProgress = reportProgress
6295
}
96+
6397

6498
public func didReceiveHead(
6599
task: HTTPClient.Task<Response>,
@@ -107,6 +141,12 @@ public final class FileDownloadDelegate: HTTPClientResponseDelegate {
107141
private func close(fileHandle: NIOFileHandle) {
108142
try! fileHandle.close()
109143
self.fileHandleFuture = nil
144+
switch threadPool {
145+
case .unowned:
146+
break
147+
case .owned(let pool):
148+
try! pool.syncShutdownGracefully()
149+
}
110150
}
111151

112152
private func finalize() {
@@ -128,4 +168,14 @@ public final class FileDownloadDelegate: HTTPClientResponseDelegate {
128168
self.finalize()
129169
return self.progress
130170
}
171+
172+
deinit {
173+
switch threadPool {
174+
case .unowned:
175+
break
176+
case .owned(let pool):
177+
// if the delegate is unused we still need to shutdown the thread pool
178+
try! pool.syncShutdownGracefully()
179+
}
180+
}
131181
}

0 commit comments

Comments
 (0)