fix cycle

Signed-off-by: Michael Rebello <[email protected]>
connectrpc · Jan 9, 2025 · ff35359 · ff35359
1 parent 20c7871
commit ff35359
Show file tree

Hide file tree

Showing 2 changed files with 73 additions and 15 deletions.
diff --git a/Examples/ElizaSharedSources/AppSources/MenuView.swift b/Examples/ElizaSharedSources/AppSources/MenuView.swift
@@ -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"
@@ -87,7 +87,7 @@ struct MenuView: View {
                             destination: LazyNavigationView {
                                 MessagingView(
                                     viewModel: UnaryMessagingViewModel(
-                                        client: self.createClient(withProtocol: .connect)
+                                        client: Self.createClient(withProtocol: .connect)
                                     )
                                 )
                             }
@@ -100,7 +100,7 @@ struct MenuView: View {
                             destination: LazyNavigationView {
                                 MessagingView(
                                     viewModel: BidirectionalStreamingMessagingViewModel(
-                                        client: self.createClient(withProtocol: .connect)
+                                        client: Self.createClient(withProtocol: .connect)
                                     )
                                 )
                             }
@@ -114,7 +114,7 @@ struct MenuView: View {
                             destination: LazyNavigationView {
                                 MessagingView(
                                     viewModel: UnaryMessagingViewModel(
-                                        client: self.createClient(withProtocol: .grpc)
+                                        client: Self.createClient(withProtocol: .grpc)
                                     )
                                 )
                             }
@@ -128,7 +128,7 @@ struct MenuView: View {
                             destination: LazyNavigationView {
                                 MessagingView(
                                     viewModel: UnaryMessagingViewModel(
-                                        client: self.createClient(withProtocol: .grpcWeb)
+                                        client: Self.createClient(withProtocol: .grpcWeb)
                                     )
                                 )
                             }
@@ -142,7 +142,7 @@ struct MenuView: View {
                             destination: LazyNavigationView {
                                 MessagingView(
                                     viewModel: BidirectionalStreamingMessagingViewModel(
-                                        client: self.createClient(withProtocol: .grpc)
+                                        client: Self.createClient(withProtocol: .grpc)
                                     )
                                 )
                             }
@@ -156,7 +156,7 @@ struct MenuView: View {
                             destination: LazyNavigationView {
                                 MessagingView(
                                     viewModel: BidirectionalStreamingMessagingViewModel(
-                                        client: self.createClient(withProtocol: .grpcWeb)
+                                        client: Self.createClient(withProtocol: .grpcWeb)
                                     )
                                 )
                             }

diff --git a/Libraries/Connect/Public/Implementation/Clients/URLSessionHTTPClient.swift b/Libraries/Connect/Public/Implementation/Clients/URLSessionHTTPClient.swift
@@ -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
@@ -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
@@ -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?
     ) {
@@ -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
+    ) {
+        client?.urlSession(
+            session, dataTask: dataTask, didReceive: response, completionHandler: completionHandler
+        )
+    }
+
+    func urlSession(
+        _ session: URLSession, dataTask: URLSessionDataTask, didReceive data: Data
+    ) {
+        client?.urlSession(session, dataTask: dataTask, didReceive: data)
+    }
+
+    func urlSession(
+        _ session: URLSession, task: URLSessionTask,
+        needNewBodyStream completionHandler: @escaping (InputStream?) -> Void
+    ) {
+        client?.urlSession(session, task: task, needNewBodyStream: completionHandler)
+    }
+}
+
+extension URLSessionDelegateWrapper: URLSessionTaskDelegate {
+    func urlSession(
+        _ session: URLSession, task: URLSessionTask, didCompleteWithError error: Swift.Error?
+    ) {
+        client?.urlSession(session, task: task, didCompleteWithError: error)
+    }
+
+    func urlSession(
+        _ session: URLSession, task: URLSessionTask,
+        didFinishCollecting metrics: URLSessionTaskMetrics
+    ) {
+        client?.urlSession(session, task: task, didFinishCollecting: metrics)
+    }
+}
+
+// MARK: - Extensions
+
 extension HTTPURLResponse {
     func formattedLowercasedHeaders() -> Headers {
         return self.allHeaderFields.reduce(into: Headers()) { headers, current in