Skip to content

Commit

Permalink
Fix retain cycle between URLSessionHTTPClient and `URLSession.deleg…
Browse files Browse the repository at this point in the history 
…ate` (#335)

See #334 for more
information on the underlying issue.

This PR creates a private `URLSessionDelegateWrapper` class to act as an
intermediary between `URLSessionHTTPClient` and `URLSession.delegate`.
From the inline documentation I added:

```
/// This class exists to avoid a retain cycle between `URLSessionHTTPClient` and its underlying
/// `URLSession.delegate`. Since `URLSession` retains its `delegate` strongly, setting the
/// `URLSessionHTTPClient` directly as the `delegate` will cause a retain cycle:
/// https://developer.apple.com/documentation/foundation/urlsession/1411597-init#parameters
///
/// To work around this, `URLSessionDelegateWrapper` maintains a `weak` reference to the
/// `URLSessionHTTPClient` and passes delegate calls through to it, avoiding the retain cycle.
```

The client now also calls `self.session.finishTasksAndInvalidate()` on
`deinit`.

This workaround isn't the _cleanest_, but it does solve the issue and
I'm not sure there's a better way of doing it. Open to suggestions,
though.

I also made the builder function in the example app `static` to avoid
capturing `self` in the builder closures.

With these changes, I was able to confirm the cycle demonstrated in the
above issue is resolved in the memory debugger:

<img width="544" alt="Screenshot 2025-01-09 at 9 56 03 AM"
src="https://github.com/user-attachments/assets/713a59b8-958f-4a61-84fe-4c05eaaba83e"
/>

Note that I avoided changing the `open func` signatures of
`URLSessionHTTPClient` since that would be a breaking API change.

Resolves #334.

---------

Signed-off-by: Michael Rebello <[email protected]>
  • Loading branch information
@rebello95
rebello95 authored Jan 14, 2025
1 parent cee49ee commit f8ae421
Show file tree
Hide file tree
Showing 2 changed files with 73 additions and 15 deletions.
14 changes: 7 additions & 7 deletions Examples/ElizaSharedSources/AppSources/MenuView.swift
View file
Original file line number Diff line number Diff line change
Expand Up @@ -38,7 +38,7 @@ extension MessagingConnectionType: Identifiable {
}

struct MenuView: View {
private func createClient(withProtocol networkProtocol: NetworkProtocol)
private static func createClient(withProtocol networkProtocol: NetworkProtocol)
-> Connectrpc_Eliza_V1_ElizaServiceClient
{
let host = "https://demo.connectrpc.com"
Expand Down Expand Up @@ -87,7 +87,7 @@ struct MenuView: View {
destination: LazyNavigationView {
MessagingView(
viewModel: UnaryMessagingViewModel(
client: self.createClient(withProtocol: .connect)
client: Self.createClient(withProtocol: .connect)
)
)
}
Expand All @@ -100,7 +100,7 @@ struct MenuView: View {
destination: LazyNavigationView {
MessagingView(
viewModel: BidirectionalStreamingMessagingViewModel(
client: self.createClient(withProtocol: .connect)
client: Self.createClient(withProtocol: .connect)
)
)
}
Expand All @@ -114,7 +114,7 @@ struct MenuView: View {
destination: LazyNavigationView {
MessagingView(
viewModel: UnaryMessagingViewModel(
client: self.createClient(withProtocol: .grpc)
client: Self.createClient(withProtocol: .grpc)
)
)
}
Expand All @@ -128,7 +128,7 @@ struct MenuView: View {
destination: LazyNavigationView {
MessagingView(
viewModel: UnaryMessagingViewModel(
client: self.createClient(withProtocol: .grpcWeb)
client: Self.createClient(withProtocol: .grpcWeb)
)
)
}
Expand All @@ -142,7 +142,7 @@ struct MenuView: View {
destination: LazyNavigationView {
MessagingView(
viewModel: BidirectionalStreamingMessagingViewModel(
client: self.createClient(withProtocol: .grpc)
client: Self.createClient(withProtocol: .grpc)
)
)
}
Expand All @@ -156,7 +156,7 @@ struct MenuView: View {
destination: LazyNavigationView {
MessagingView(
viewModel: BidirectionalStreamingMessagingViewModel(
client: self.createClient(withProtocol: .grpcWeb)
client: Self.createClient(withProtocol: .grpcWeb)
)
)
}
Expand Down
74 changes: 66 additions & 8 deletions Libraries/Connect/Public/Implementation/Clients/URLSessionHTTPClient.swift
View file
Original file line number Diff line number Diff line change
Expand Up @@ -25,21 +25,27 @@ open class URLSessionHTTPClient: NSObject, HTTPClientInterface, @unchecked Senda
private let lock = Lock()
/// Closures stored for notifying when metrics are available.
private var metricsClosures = [Int: @Sendable (HTTPMetrics) -> Void]()
/// Force unwrapped to allow using `self` as the delegate.
private var session: URLSession!
/// Session used for performing requests.
private let session: URLSession
/// List of active streams.
/// TODO: Remove in favor of simply setting
/// `URLSessionTask.delegate = <URLSessionStream instance>` once we are able to set iOS 15
/// as the base deployment target.
private var streams = [Int: URLSessionStream]()

public init(configuration: URLSessionConfiguration = .default) {
super.init()
let delegate = URLSessionDelegateWrapper()
self.session = URLSession(
configuration: configuration,
delegate: self,
delegate: delegate,
delegateQueue: .main
)
super.init()
delegate.client = self
}

deinit {
self.session.finishTasksAndInvalidate()
}

@discardableResult
Expand Down Expand Up @@ -124,9 +130,9 @@ open class URLSessionHTTPClient: NSObject, HTTPClientInterface, @unchecked Senda
sendClose: { urlSessionStream.close() }
)
}
}

extension URLSessionHTTPClient: URLSessionDataDelegate {
// MARK: - URLSession delegate functions

open func urlSession(
_ session: URLSession, dataTask: URLSessionDataTask, didReceive response: URLResponse,
completionHandler: @escaping (URLSession.ResponseDisposition) -> Void
Expand Down Expand Up @@ -156,9 +162,7 @@ extension URLSessionHTTPClient: URLSessionDataDelegate {
self.lock.perform { self.streams[task.taskIdentifier]?.requestBodyStream }
)
}
}

extension URLSessionHTTPClient: URLSessionTaskDelegate {
open func urlSession(
_ session: URLSession, task: URLSessionTask, didCompleteWithError error: Swift.Error?
) {
Expand All @@ -178,6 +182,60 @@ extension URLSessionHTTPClient: URLSessionTaskDelegate {
}
}

// MARK: - URLSession delegate wrapper

/// This class exists to avoid a retain cycle between `URLSessionHTTPClient` and its underlying
/// `URLSession.delegate`. Since `URLSession` retains its `delegate` strongly, setting the
/// `URLSessionHTTPClient` directly as the `delegate` will cause a retain cycle:
/// https://developer.apple.com/documentation/foundation/urlsession/1411597-init#parameters
///
/// To work around this, `URLSessionDelegateWrapper` maintains a `weak` reference to the
/// `URLSessionHTTPClient` and passes delegate calls through to it, avoiding the retain cycle.
private final class URLSessionDelegateWrapper: NSObject, @unchecked Sendable {
weak var client: URLSessionHTTPClient?
}

extension URLSessionDelegateWrapper: URLSessionDataDelegate {
func urlSession(
_ session: URLSession, dataTask: URLSessionDataTask, didReceive response: URLResponse,
completionHandler: @escaping (URLSession.ResponseDisposition) -> Void
) {
self.client?.urlSession(
session, dataTask: dataTask, didReceive: response, completionHandler: completionHandler
)
}

func urlSession(
_ session: URLSession, dataTask: URLSessionDataTask, didReceive data: Data
) {
self.client?.urlSession(session, dataTask: dataTask, didReceive: data)
}

func urlSession(
_ session: URLSession, task: URLSessionTask,
needNewBodyStream completionHandler: @escaping (InputStream?) -> Void
) {
self.client?.urlSession(session, task: task, needNewBodyStream: completionHandler)
}
}

extension URLSessionDelegateWrapper: URLSessionTaskDelegate {
func urlSession(
_ session: URLSession, task: URLSessionTask, didCompleteWithError error: Swift.Error?
) {
self.client?.urlSession(session, task: task, didCompleteWithError: error)
}

func urlSession(
_ session: URLSession, task: URLSessionTask,
didFinishCollecting metrics: URLSessionTaskMetrics
) {
self.client?.urlSession(session, task: task, didFinishCollecting: metrics)
}
}

// MARK: - Extensions

extension HTTPURLResponse {
func formattedLowercasedHeaders() -> Headers {
return self.allHeaderFields.reduce(into: Headers()) { headers, current in
Expand Down

0 comments on commit f8ae421

Please sign in to comment.