Fix leaks

Sources/Realtime/Channel.swift

@@ -281,9 +281,9 @@ public final class RealtimeChannelV2: @unchecked Sendable {
       continuation.yield($0)
     }
 
-    continuation.onTermination = { _ in
+    continuation.onTermination = { [weak callbackManager] _ in
       debug("Removing presence callback with id: \(id)")
-      self.callbackManager.removeCallback(id: id)
+      callbackManager?.removeCallback(id: id)
     }
 
     return stream
@@ -321,9 +321,9 @@ public final class RealtimeChannelV2: @unchecked Sendable {
       }
     }
 
-    continuation.onTermination = { _ in
+    continuation.onTermination = { [weak callbackManager] _ in
       debug("Removing postgres callback with id: \(id)")
-      self.callbackManager.removeCallback(id: id)
+      callbackManager?.removeCallback(id: id)
     }
 
     return stream
@@ -338,9 +338,9 @@ public final class RealtimeChannelV2: @unchecked Sendable {
       continuation.yield($0)
     }
 
-    continuation.onTermination = { _ in
+    continuation.onTermination = { [weak callbackManager] _ in
       debug("Removing broadcast callback with id: \(id)")
-      self.callbackManager.removeCallback(id: id)
+      callbackManager?.removeCallback(id: id)
     }
 
     return stream

Sources/Realtime/RealtimeChannel.swift

@@ -336,9 +336,9 @@ public class RealtimeChannel {
   ///
   /// - parameter msg: The Message received by the client from the server
   /// - return: Must return the message, modified or unmodified
-//  public var onMessage: (_ message: RealtimeMessage) -> RealtimeMessage = { message in
-//    message
-//  }
+  public var onMessage: (_ message: RealtimeMessage) -> RealtimeMessage = { message in
+    message
+  }
 
   /// Joins the channel
   ///
@@ -852,9 +852,9 @@ public class RealtimeChannel {
   /// - parameter payload: The payload for the message
   /// - parameter ref: The reference of the message
   /// - return: Must return the payload, modified or unmodified
-//  public func onMessage(callback: @escaping (RealtimeMessage) -> RealtimeMessage) {
-//    onMessage = callback
-//  }
+  public func onMessage(callback: @escaping (RealtimeMessage) -> RealtimeMessage) {
+    onMessage = callback
+  }
 
   // ----------------------------------------------------------------------
 

Sources/Realtime/RealtimeClient.swift

@@ -54,6 +54,7 @@ struct StateChangeCallbacks {
 /// The `RealtimeClient` constructor takes the mount point of the socket,
 /// the authentication params, as well as options that can be found in
 /// the Socket docs, such as configuring the heartbeat.
+@available(*, deprecated, message: "Use new Realtime class instead.")
 public class RealtimeClient: PhoenixTransportDelegate {
   // ----------------------------------------------------------------------
 

Sources/Realtime/WebSocketClient.swift

@@ -63,7 +63,10 @@ final class WebSocketClient: NSObject, URLSessionWebSocketDelegate, WebSocketCli
   }
 
   func cancel() {
-    mutableState.task?.cancel()
+    mutableState.withValue {
+      $0.task?.cancel()
+      $0.statusContinuation?.finish()
+    }
   }
 
   func urlSession(
@@ -74,7 +77,7 @@ final class WebSocketClient: NSObject, URLSessionWebSocketDelegate, WebSocketCli
     mutableState.statusContinuation?.yield(.open)
   }
 
-  nonisolated func urlSession(
+  func urlSession(
     _: URLSession,
     webSocketTask _: URLSessionWebSocketTask,
     didCloseWith _: URLSessionWebSocketTask.CloseCode,
@@ -83,7 +86,7 @@ final class WebSocketClient: NSObject, URLSessionWebSocketDelegate, WebSocketCli
     mutableState.statusContinuation?.yield(.close)
   }
 
-  nonisolated func urlSession(
+  func urlSession(
     _: URLSession,
     task _: URLSessionTask,
     didCompleteWithError error: Error?

Tests/RealtimeTests/CallbackManagerTests.swift

@@ -14,6 +14,8 @@ import XCTest
 final class CallbackManagerTests: XCTestCase {
   func testIntegration() {
     let callbackManager = CallbackManager()
+    XCTAssertNoLeak(callbackManager)
+
     let filter = PostgresJoinConfig(
       event: .update,
       schema: "public",
@@ -48,6 +50,8 @@ final class CallbackManagerTests: XCTestCase {
 
   func testSetServerChanges() {
     let callbackManager = CallbackManager()
+    XCTAssertNoLeak(callbackManager)
+
     let changes = [PostgresJoinConfig(
       event: .update,
       schema: "public",
@@ -63,6 +67,8 @@ final class CallbackManagerTests: XCTestCase {
 
   func testTriggerPostgresChanges() {
     let callbackManager = CallbackManager()
+    XCTAssertNoLeak(callbackManager)
+
     let updateUsersFilter = PostgresJoinConfig(
       event: .update,
       schema: "public",
@@ -165,6 +171,8 @@ final class CallbackManagerTests: XCTestCase {
 
   func testTriggerBroadcast() {
     let callbackManager = CallbackManager()
+    XCTAssertNoLeak(callbackManager)
+
     let event = "new_user"
     let message = RealtimeMessageV2(
       joinRef: nil,
@@ -212,3 +220,11 @@ final class CallbackManagerTests: XCTestCase {
     XCTAssertEqual(receivedAction.value?.leaves.count, 1)
   }
 }
+
+extension XCTestCase {
+  func XCTAssertNoLeak(_ object: AnyObject, file: StaticString = #file, line: UInt = #line) {
+    addTeardownBlock { [weak object] in
+      XCTAssertNil(object, file: file, line: line)
+    }
+  }
+}

Tests/RealtimeTests/RealtimeTests.swift

@@ -22,6 +22,8 @@ final class RealtimeTests: XCTestCase {
       makeWebSocketClient: { _ in mock }
     )
 
+    XCTAssertNoLeak(realtime)
+
     try await realtime.connect()
 
     XCTAssertEqual(realtime.status, .connected)
@@ -34,8 +36,10 @@ final class RealtimeTests: XCTestCase {
       config: Realtime.Configuration(url: url, apiKey: apiKey, authTokenProvider: nil),
       makeWebSocketClient: { _ in mock }
     )
+    XCTAssertNoLeak(realtime)
 
     let channel = realtime.channel("users")
+    XCTAssertNoLeak(channel)
 
     let changes = channel.postgresChange(
       AnyAction.self,
@@ -126,5 +130,13 @@ final class RealtimeTests: XCTestCase {
 
     let receivedChange = await receivedPostgresChangeTask.value
     XCTAssertNoDifference(receivedChange, action)
+
+    try await channel.unsubscribe()
+
+    mock.mockReceive(
+      RealtimeMessageV2(joinRef: nil, ref: nil, topic: "realtime:users", event: ChannelEvent.leave, payload: [:])
+    )
+
+    await Task.megaYield()
   }
 }