fix(realtime): make Push associated to MainActor

Author: grdsdev

Sources/Realtime/PushV2.swift

@@ -15,7 +15,8 @@ public enum PushStatus: String, Sendable {
   case timeout
 }
 
-actor PushV2 {
+@MainActor
+final class PushV2 {
   private weak var channel: RealtimeChannelV2?
   let message: RealtimeMessageV2
 

Sources/Realtime/RealtimeChannelV2.swift

@@ -32,13 +32,15 @@ public final class RealtimeChannelV2: Sendable {
     var pushes: [String: PushV2] = [:]
   }
 
-  private let mutableState = LockIsolated(MutableState())
+  @MainActor
+  private var mutableState = MutableState()
 
   let topic: String
   let config: RealtimeChannelConfig
   let logger: (any SupabaseLogger)?
   let socket: RealtimeClientV2
-  var joinRef: String? { mutableState.joinRef }
+
+  @MainActor var joinRef: String? { mutableState.joinRef }
 
   let callbackManager = CallbackManager()
   private let statusSubject = AsyncValueSubject<RealtimeChannelStatus>(.unsubscribed)
@@ -81,6 +83,7 @@ public final class RealtimeChannelV2: Sendable {
   }
 
   /// Subscribes to the channel
+  @MainActor
   public func subscribe() async {
     if socket.status != .connected {
       if socket.options.connectOnSubscribe != true {
@@ -109,7 +112,7 @@ public final class RealtimeChannelV2: Sendable {
     )
 
     let joinRef = socket.makeRef()
-    mutableState.withValue { $0.joinRef = joinRef }
+    mutableState.joinRef = joinRef
 
     logger?.debug("Subscribing to channel with body: \(joinConfig)")
 
@@ -497,8 +500,8 @@ public final class RealtimeChannelV2: Sendable {
       filter: filter
     )
 
-    mutableState.withValue {
-      $0.clientChanges.append(config)
+    Task { @MainActor in
+      mutableState.clientChanges.append(config)
     }
 
     let id = callbackManager.addPostgresCallback(filter: config, callback: callback)
@@ -538,32 +541,28 @@ public final class RealtimeChannelV2: Sendable {
     self.onSystem { _ in callback() }
   }
 
+  @MainActor
   @discardableResult
   func push(_ event: String, ref: String? = nil, payload: JSONObject = [:]) async -> PushStatus {
-    let push = mutableState.withValue {
-      let message = RealtimeMessageV2(
-        joinRef: $0.joinRef,
-        ref: ref ?? socket.makeRef(),
-        topic: self.topic,
-        event: event,
-        payload: payload
-      )
-
-      let push = PushV2(channel: self, message: message)
-      if let ref = message.ref {
-        $0.pushes[ref] = push
-      }
+    let message = RealtimeMessageV2(
+      joinRef: joinRef,
+      ref: ref ?? socket.makeRef(),
+      topic: self.topic,
+      event: event,
+      payload: payload
+    )
 
-      return push
+    let push = PushV2(channel: self, message: message)
+    if let ref = message.ref {
+      mutableState.pushes[ref] = push
     }
 
     return await push.send()
   }
 
-  private func didReceiveReply(ref: String, status: String) async {
-    let push = mutableState.withValue {
-      $0.pushes.removeValue(forKey: ref)
-    }
-    await push?.didReceive(status: PushStatus(rawValue: status) ?? .ok)
+  @MainActor
+  private func didReceiveReply(ref: String, status: String) {
+    let push = mutableState.pushes.removeValue(forKey: ref)
+    push?.didReceive(status: PushStatus(rawValue: status) ?? .ok)
   }
 }

Tests/RealtimeTests/_PushTests.swift

@@ -12,19 +12,12 @@ import XCTest
 
 @testable import Realtime
 
+@MainActor
 @available(macOS 13.0, iOS 16.0, watchOS 9.0, tvOS 16.0, *)
 final class _PushTests: XCTestCase {
   var ws: FakeWebSocket!
   var socket: RealtimeClientV2!
 
-  #if !os(Windows) && !os(Linux) && !os(Android)
-  override func invokeTest() {
-    withMainSerialExecutor {
-      super.invokeTest()
-    }
-  }
-  #endif
-
   override func setUp() {
     super.setUp()
 
@@ -67,36 +60,35 @@ final class _PushTests: XCTestCase {
     XCTAssertEqual(status, .ok)
   }
 
-  // FIXME: Flaky test, it fails some time due the task scheduling, even tho we're using withMainSerialExecutor.
-  //  func testPushWithAck() async {
-  //    let channel = RealtimeChannelV2(
-  //      topic: "realtime:users",
-  //      config: RealtimeChannelConfig(
-  //        broadcast: .init(acknowledgeBroadcasts: true),
-  //        presence: .init(),
-  //        isPrivate: false
-  //      ),
-  //      socket: Socket(client: socket),
-  //      logger: nil
-  //    )
-  //    let push = PushV2(
-  //      channel: channel,
-  //      message: RealtimeMessageV2(
-  //        joinRef: nil,
-  //        ref: "1",
-  //        topic: "realtime:users",
-  //        event: "broadcast",
-  //        payload: [:]
-  //      )
-  //    )
-  //
-  //    let task = Task {
-  //      await push.send()
-  //    }
-  //    await Task.yield()
-  //    await push.didReceive(status: .ok)
-  //
-  //    let status = await task.value
-  //    XCTAssertEqual(status, .ok)
-  //  }
+  func testPushWithAck() async {
+    let channel = RealtimeChannelV2(
+      topic: "realtime:users",
+      config: RealtimeChannelConfig(
+        broadcast: .init(acknowledgeBroadcasts: true),
+        presence: .init(),
+        isPrivate: false
+      ),
+      socket: socket,
+      logger: nil
+    )
+    let push = PushV2(
+      channel: channel,
+      message: RealtimeMessageV2(
+        joinRef: nil,
+        ref: "1",
+        topic: "realtime:users",
+        event: "broadcast",
+        payload: [:]
+      )
+    )
+
+    let task = Task {
+      await push.send()
+    }
+    await Task.megaYield()
+    push.didReceive(status: .ok)
+
+    let status = await task.value
+    XCTAssertEqual(status, .ok)
+  }
 }