From 179445f088c4f7dd80415a53a1423460540ebb37 Mon Sep 17 00:00:00 2001 From: Alex Hoppen Date: Tue, 19 Mar 2024 15:58:44 +0100 Subject: [PATCH] Fix a memory leak in `DownloadTaskManager` and `DataTaskManager` MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `DownloadTaskManager` had a strong reference to `URLSession`, which retained its delegate, which was the `DownloadTaskManager`. This formed a retain cycle. Make the reference from `URLSession` to the delegate weak and only keep `DownloadTaskManager` alive as long as `DownloadTask`s need it. This causes the `DownloadTaskManager` to be deallocated once nobody holds a reference to it anymore and all the tasks it manages are finished. Finally, we need to call `invalidate` on the `URLSession` to tell it that we’re done and that it should release its delegate (which is now the `WeakDownloadTaskManager`wrapper). rdar://125012584 --- .../HTTPClient/URLSessionHTTPClient.swift | 93 +++++++++++++++++-- 1 file changed, 86 insertions(+), 7 deletions(-) diff --git a/Sources/Basics/HTTPClient/URLSessionHTTPClient.swift b/Sources/Basics/HTTPClient/URLSessionHTTPClient.swift index d6302003d2a..2d253705170 100644 --- a/Sources/Basics/HTTPClient/URLSessionHTTPClient.swift +++ b/Sources/Basics/HTTPClient/URLSessionHTTPClient.swift @@ -87,7 +87,35 @@ final class URLSessionHTTPClient { } } -private class DataTaskManager: NSObject, URLSessionDataDelegate { +/// A weak wrapper around `DataTaskManager` that conforms to `URLSessionDataDelegate`. +/// +/// This ensures that we don't get a retain cycle between `DataTaskManager.session` -> `URLSession.delegate` -> `DataTaskManager`. +/// +/// The `DataTaskManager` is being kept alive by a reference from all `DataTask`s that it manages. Once all the +/// `DataTasks` have finished and are deallocated, `DataTaskManager` will get deinitialized, which invalidates the +/// session, which then lets go of `WeakDataTaskManager`. +private class WeakDataTaskManager: NSObject, URLSessionDataDelegate { + private weak var dataTaskManager: DataTaskManager? + + init(_ dataTaskManager: DataTaskManager? = nil) { + self.dataTaskManager = dataTaskManager + } + + func urlSession( + _ session: URLSession, + dataTask: URLSessionDataTask, + didReceive response: URLResponse, + completionHandler: @escaping @Sendable (URLSession.ResponseDisposition) -> Void + ) { + dataTaskManager?.urlSession(session, dataTask: dataTask, didReceive: response, completionHandler: completionHandler) + } + + func urlSession(_ session: URLSession, dataTask: URLSessionDataTask, didReceive data: Data) { + dataTaskManager?.urlSession(session, dataTask: dataTask, didReceive: data) + } +} + +private class DataTaskManager { private var tasks = ThreadSafeKeyValueStore() private let delegateQueue: OperationQueue private var session: URLSession! @@ -96,8 +124,11 @@ private class DataTaskManager: NSObject, URLSessionDataDelegate { self.delegateQueue = OperationQueue() self.delegateQueue.name = "org.swift.swiftpm.urlsession-http-client-data-delegate" self.delegateQueue.maxConcurrentOperationCount = 1 - super.init() - self.session = URLSession(configuration: configuration, delegate: self, delegateQueue: self.delegateQueue) + self.session = URLSession(configuration: configuration, delegate: WeakDataTaskManager(self), delegateQueue: self.delegateQueue) + } + + deinit { + session.finishTasksAndInvalidate() } func makeTask( @@ -110,6 +141,7 @@ private class DataTaskManager: NSObject, URLSessionDataDelegate { self.tasks[task.taskIdentifier] = DataTask( task: task, progressHandler: progress, + dataTaskManager: self, completionHandler: completion, authorizationProvider: authorizationProvider ) @@ -192,6 +224,11 @@ private class DataTaskManager: NSObject, URLSessionDataDelegate { class DataTask { let task: URLSessionDataTask let completionHandler: LegacyHTTPClient.CompletionHandler + /// A strong reference to keep the `DataTaskManager` alive so it can handle the callbacks from the + /// `URLSession`. + /// + /// See comment on `WeakDataTaskManager`. + let dataTaskManager: DataTaskManager let progressHandler: LegacyHTTPClient.ProgressHandler? let authorizationProvider: LegacyHTTPClientConfiguration.AuthorizationProvider? @@ -202,18 +239,49 @@ private class DataTaskManager: NSObject, URLSessionDataDelegate { init( task: URLSessionDataTask, progressHandler: LegacyHTTPClient.ProgressHandler?, + dataTaskManager: DataTaskManager, completionHandler: @escaping LegacyHTTPClient.CompletionHandler, authorizationProvider: LegacyHTTPClientConfiguration.AuthorizationProvider? ) { self.task = task self.progressHandler = progressHandler + self.dataTaskManager = dataTaskManager self.completionHandler = completionHandler self.authorizationProvider = authorizationProvider } } } -private class DownloadTaskManager: NSObject, URLSessionDownloadDelegate { +/// This uses the same pattern as `WeakDataTaskManager`. See comment on that type. +private class WeakDownloadTaskManager: NSObject, URLSessionDownloadDelegate { + private weak var downloadTaskManager: DownloadTaskManager? + + init(_ downloadTaskManager: DownloadTaskManager? = nil) { + self.downloadTaskManager = downloadTaskManager + } + + func urlSession(_ session: URLSession, downloadTask: URLSessionDownloadTask, didFinishDownloadingTo location: URL) { + downloadTaskManager?.urlSession(session, downloadTask: downloadTask, didFinishDownloadingTo: location) + } + + func urlSession( + _ session: URLSession, + downloadTask: URLSessionDownloadTask, + didWriteData bytesWritten: Int64, + totalBytesWritten: Int64, + totalBytesExpectedToWrite: Int64 + ) { + downloadTaskManager?.urlSession( + session, + downloadTask: downloadTask, + didWriteData: bytesWritten, + totalBytesWritten: totalBytesWritten, + totalBytesExpectedToWrite: totalBytesExpectedToWrite + ) + } +} + +private class DownloadTaskManager { private var tasks = ThreadSafeKeyValueStore() private let delegateQueue: OperationQueue private var session: URLSession! @@ -222,8 +290,11 @@ private class DownloadTaskManager: NSObject, URLSessionDownloadDelegate { self.delegateQueue = OperationQueue() self.delegateQueue.name = "org.swift.swiftpm.urlsession-http-client-download-delegate" self.delegateQueue.maxConcurrentOperationCount = 1 - super.init() - self.session = URLSession(configuration: configuration, delegate: self, delegateQueue: self.delegateQueue) + self.session = URLSession(configuration: configuration, delegate: WeakDownloadTaskManager(self), delegateQueue: self.delegateQueue) + } + + deinit { + session.finishTasksAndInvalidate() } func makeTask( @@ -238,6 +309,7 @@ private class DownloadTaskManager: NSObject, URLSessionDownloadDelegate { task: task, fileSystem: fileSystem, destination: destination, + downloadTaskManager: self, progressHandler: progress, completionHandler: completion ) @@ -314,8 +386,13 @@ private class DownloadTaskManager: NSObject, URLSessionDownloadDelegate { let task: URLSessionDownloadTask let fileSystem: FileSystem let destination: AbsolutePath - let completionHandler: LegacyHTTPClient.CompletionHandler + /// A strong reference to keep the `DownloadTaskManager` alive so it can handle the callbacks from the + /// `URLSession`. + /// + /// See comment on `WeakDownloadTaskManager`. + private let downloadTaskManager: DownloadTaskManager let progressHandler: LegacyHTTPClient.ProgressHandler? + let completionHandler: LegacyHTTPClient.CompletionHandler var moveFileError: Error? @@ -323,12 +400,14 @@ private class DownloadTaskManager: NSObject, URLSessionDownloadDelegate { task: URLSessionDownloadTask, fileSystem: FileSystem, destination: AbsolutePath, + downloadTaskManager: DownloadTaskManager, progressHandler: LegacyHTTPClient.ProgressHandler?, completionHandler: @escaping LegacyHTTPClient.CompletionHandler ) { self.task = task self.fileSystem = fileSystem self.destination = destination + self.downloadTaskManager = downloadTaskManager self.progressHandler = progressHandler self.completionHandler = completionHandler }