update messageBuffer and fix test

Author: xingsy97

Sources/SignalRClient/HubConnection.swift

@@ -16,7 +16,7 @@ public actor HubConnection {
     private let logger: Logger
     private let hubProtocol: HubProtocol
     private let connection: ConnectionProtocol
-    private let reconnectPolicy: RetryPolicy
+    private let retryPolicy: RetryPolicy
     private let keepAliveScheduler: TimeScheduler
     private let serverTimeoutScheduler: TimeScheduler
     private let statefulReconnectBufferSize: Int
@@ -40,7 +40,7 @@ public actor HubConnection {
     internal init(connection: ConnectionProtocol,
                   logger: Logger,
                   hubProtocol: HubProtocol,
-                  reconnectPolicy: RetryPolicy,
+                  retryPolicy: RetryPolicy,
                   serverTimeout: TimeInterval?,
                   keepAliveInterval: TimeInterval?,
                   statefulReconnectBufferSize: Int?) {
@@ -50,7 +50,7 @@ public actor HubConnection {
             statefulReconnectBufferSize ?? HubConnection.defaultStatefulReconnectBufferSize
 
         self.logger = logger
-        self.reconnectPolicy = reconnectPolicy
+        self.retryPolicy = retryPolicy
 
         self.connection = connection
         self.hubProtocol = hubProtocol
@@ -361,7 +361,7 @@ public actor HubConnection {
         var lastError: Error? = error
 
         // reconnect
-        while let interval = reconnectPolicy.nextRetryInterval(
+        while let interval = retryPolicy.nextRetryInterval(
             retryContext: RetryContext(
                 retryCount: retryCount,
                 elapsed: elapsed,
@@ -438,16 +438,15 @@ public actor HubConnection {
             for message in hubMessage {
                 do {
                     if let messageBuffer = self.messageBuffer {
-                        let shouldProcess = try await messageBuffer.shouldProcessMessage(message)
-                        if !shouldProcess {
+                        if !(try await messageBuffer.shouldProcessMessage(message)) {
                             // Don't process the message, we are either waiting for a SequenceMessage or received a duplicate message
                             continue
                         }
                     }
+                    await dispatchMessage(message)
                 } catch {
                     logger.log(level: .error, message: "Error parsing messages: \(error)")
                 }
-                await dispatchMessage(message)
             }
         } catch {
             logger.log(level: .error, message: "Error parsing messages: \(error)")
@@ -482,11 +481,23 @@ public actor HubConnection {
         case _ as CloseMessage:
             // Close
             break
-        case _ as AckMessage:
-            // TODO: In stateful reconnect
+        case let message as AckMessage:
+            do {
+                let result = try await self.messageBuffer?.ack(sequenceId: Int(message.sequenceId));
+                if (result == false) {
+                    logger.log(level: .warning, message: "Ack message received for sequenceId: \(message.sequenceId), but failed.")
+                }
+            }
+            catch {
+                logger.log(level: .error, message: "Ack message received for sequenceId: \(message.sequenceId), but failed with error: \(error)")
+            }
             break
-        case _ as SequenceMessage:
-            // TODO: In stateful reconnect
+        case let message as SequenceMessage:
+            if let messageBuffer = self.messageBuffer {
+                await messageBuffer.resetSequenceMessage(message: message)
+            } else {
+                logger.log(level: .warning, message: "Sequence message received but no message buffer is available.")
+            }
             break
         default:
             logger.log(level: .warning, message: "Unknown message type: \(message)")
@@ -593,8 +604,7 @@ public actor HubConnection {
             let useStatefulReconnect = await (self.connection.features[ConnectionFeature.Reconnect] as? Bool) == true
             if useStatefulReconnect {
                 self.messageBuffer = MessageBuffer(
-                    hubProtocol: self.hubProtocol, connection: self.connection,
-                    bufferSize: self.statefulReconnectBufferSize)
+                    bufferSize: self.statefulReconnectBufferSize, hubProtocol: self.hubProtocol, connection: self.connection)
                 await self.connection.setFeature(
                     feature: ConnectionFeature.Disconnected, 
                     value: { [weak self] () async -> Void in

Sources/SignalRClient/HubConnectionBuilder.swift

@@ -103,7 +103,7 @@ public class HubConnectionBuilder {
         return HubConnection(connection: connection,
                              logger: logger,
                              hubProtocol: hubProtocol,
-                             reconnectPolicy: retryPolicy,
+                             retryPolicy: retryPolicy,
                              serverTimeout: serverTimeout,
                              keepAliveInterval: keepAliveInterval,
                              statefulReconnectBufferSize: statefulReconnectBufferSize)

Sources/SignalRClient/MessageBuffer.swift

@@ -13,29 +13,123 @@ actor MessageBuffer {
     private var totalMessageCount: Int = 0
     private var lastSendSequenceId: Int = 0
     private var nextSendIdx = 0
+    private var nextReceivingIdx: Int64 = 1
+    private var latestReceivedSequenceId: Int64 = 0
     private var dequeueContinuations: [CheckedContinuation<Bool, Never>] = []
     private var closed: Bool = false
+    private var reconnectInprogress: Bool = false;
+    private var waitForSequenceMessage: Bool = false;
 
-    init(hubProtocol: HubProtocol, connection: ConnectionProtocol, bufferSize: Int) {
+    private var ackTimerHandle: DispatchWorkItem?
+
+
+    init(bufferSize: Int, hubProtocol: HubProtocol, connection: ConnectionProtocol) {
+        self.maxBufferSize = bufferSize
         self.hubProtocol = hubProtocol
         self.connection = connection
-        self.maxBufferSize = bufferSize
     }
 
     public func send(message: HubMessage) async throws -> Void {
-        throw SignalRError.invalidOperation("Send is not implemented")
+        let serializedMessage = try self.hubProtocol.writeMessage(message: message);
+
+        try await self.enqueue(content: serializedMessage);
+
+        if (!self.reconnectInprogress) {
+            do {
+                try await self.connection.send(serializedMessage);
+            }
+            catch {
+                self.disconnected();
+            }
+        }
     }
 
     public func resend() async throws -> Void {
-        throw SignalRError.invalidOperation("Resend is not implemented")
+        let sequenceId = Int64(self.messages.count > 0 ? self.messages[0].id : self.totalMessageCount + 1);
+        let serializedMessage = try self.hubProtocol.writeMessage(message: SequenceMessage(sequenceId: sequenceId));
+        try await self.connection.send(serializedMessage);
+        
+        let messages = self.messages;
+        for element in messages {
+            try await self.connection.send(element.content);
+        }
+
+        self.reconnectInprogress = false;
     }
 
-    public func disconnected() async throws -> Void {
-        throw SignalRError.invalidOperation("Disconnected is not implemented")
+    public func disconnected() -> Void {
+        self.reconnectInprogress = true;
+        self.waitForSequenceMessage = true;
+    }
+
+    private func ackTimer() {
+        guard ackTimerHandle == nil else {
+            return
+        }
+        
+        let workItem = DispatchWorkItem { [weak self] in
+            guard let self = self else { return }
+            
+            Task {
+                await self.performScheduledAck()
+            }
+        }
+        
+        ackTimerHandle = workItem
+        DispatchQueue.main.asyncAfter(deadline: .now() + 1.0, execute: workItem)
+    }
+
+    private func performScheduledAck() async {
+        defer {
+            // 在方法结束时清理定时器
+            ackTimerHandle = nil
+        }
+        
+        do {
+            if !reconnectInprogress {
+                let ackMessage = AckMessage(
+                    sequenceId: latestReceivedSequenceId
+                )
+                
+                let serializedMessage = try hubProtocol.writeMessage(message: ackMessage)
+                try await connection.send(serializedMessage)
+            }
+        } catch {
+            // 忽略错误，连接关闭时不需要发送ACK
+        }
     }
 
     public func shouldProcessMessage(_ message: HubMessage) throws -> Bool {
-        throw SignalRError.invalidOperation("ShouldProcessMessage is not implemented")
+        if (self.waitForSequenceMessage) {
+            if (message.type != .sequence) {
+                return false;
+            } else {
+                self.waitForSequenceMessage = false;
+                return true;
+            }
+        }
+
+        if !self.isInvocationMessage(message: message) {
+            return true
+        }
+
+        let currentId = self.nextReceivingIdx;
+        self.nextReceivingIdx += 1;
+        if currentId <= self.latestReceivedSequenceId{
+            if currentId == self.latestReceivedSequenceId{
+                // Should only hit this if we just reconnected and the server is sending
+                // Messages it has buffered, which would mean it hasn't seen an Ack for these messages
+                self.ackTimer();
+            }
+            // Ignore, this is a duplicate message
+            return false;
+        }
+        self.latestReceivedSequenceId = currentId;
+
+        // Only start the timer for sending an Ack message when we have a message to ack. This also conveniently solves
+        // timer throttling by not having a recursive timer, and by starting the timer via a network call (recv)
+        self.ackTimer();
+        return true;     
     }
 
     public func enqueue(content: StringOrData) async throws -> Void {
@@ -137,6 +231,17 @@ actor MessageBuffer {
         }
     }
 
+    public func resetSequenceMessage(message: SequenceMessage) async {
+        if message.sequenceId > self.nextReceivingIdx {
+            // do not await stop
+            Task {
+                await self.connection.stop(error: SignalRError.invalidOperation("Received sequence message with sequenceId \(message.sequenceId) greater than nextReceivingIdx \(self.nextReceivingIdx)"))
+            }
+            return 
+        }
+        self.nextReceivingIdx = message.sequenceId;
+    }
+
     private func isInvocationMessage(message: HubMessage) -> Bool {
         switch (message.type) {
             case .invocation, .streamItem, .completion, .streamInvocation, .cancelInvocation:

Tests/SignalRClientTests/HubConnection+OnResultTests.swift

@@ -29,7 +29,7 @@ final class HubConnectionOnResultTests: XCTestCase {
             connection: mockConnection,
             logger: Logger(logLevel: .debug, logHandler: logHandler),
             hubProtocol: hubProtocol,
-            reconnectPolicy: DefaultRetryPolicy(retryDelays: []), // No retry
+            retryPolicy: DefaultRetryPolicy(retryDelays: []), // No retry
             serverTimeout: nil,
             keepAliveInterval: nil,
             statefulReconnectBufferSize: nil

Tests/SignalRClientTests/HubConnection+OnTests.swift

@@ -25,7 +25,7 @@ final class HubConnectionOnTests: XCTestCase {
             connection: mockConnection,
             logger: Logger(logLevel: .debug, logHandler: logHandler),
             hubProtocol: hubProtocol,
-            reconnectPolicy: DefaultRetryPolicy(retryDelays: []), // No retry
+            retryPolicy: DefaultRetryPolicy(retryDelays: []), // No retry
             serverTimeout: nil,
             keepAliveInterval: nil,
             statefulReconnectBufferSize: nil

Tests/SignalRClientTests/HubConnectionTests.swift

@@ -13,6 +13,7 @@ class MockConnection: ConnectionProtocol, @unchecked Sendable {
     var onSend: ((StringOrData) -> Void)?
     var onStart: (() -> Void)?
     var onStop: ((Error?) -> Void)?
+    var features: [ConnectionFeature : Any] = [:]
 
     private(set) var startCalled = false
     private(set) var sendCalled = false
@@ -42,6 +43,10 @@ class MockConnection: ConnectionProtocol, @unchecked Sendable {
     func onClose(_ handler: @escaping @Sendable ((any Error)?) async -> Void) async {
         onClose = handler
     }
+
+    func setFeature(feature: SignalRClient.ConnectionFeature, value: Any) async {
+        features[feature] = value
+    }
 }
 
 final class HubConnectionTests: XCTestCase {

Tests/SignalRClientTests/MessageBufferTests.swift

@@ -3,8 +3,14 @@ import XCTest
 @testable import SignalRClient
 
 class MessageBufferTest: XCTestCase {
+    func getTestMessageBuffer(bufferSize: Int) -> MessageBuffer {
+        return MessageBuffer(bufferSize: bufferSize, 
+                             hubProtocol: JsonHubProtocol(),
+                             connection: MockConnection())
+    }
+
     func testSendWithinBufferSize() async throws {
-        let buffer = MessageBuffer(bufferSize: 100)
+        let buffer = getTestMessageBuffer(bufferSize: 100)
         let expectation = XCTestExpectation(description: "Should enqueue")
         Task {
             try await buffer.enqueue(content: .string("data"))
@@ -14,7 +20,7 @@ class MessageBufferTest: XCTestCase {
     }
 
     func testSendTriggersBackpressure() async throws {
-        let buffer = MessageBuffer(bufferSize: 5)
+        let buffer = getTestMessageBuffer(bufferSize: 5)
         let expectation1 = XCTestExpectation(description: "Should not enqueue")
         expectation1.isInverted = true
         let expectation2 = XCTestExpectation(description: "Should enqueue")
@@ -33,7 +39,7 @@ class MessageBufferTest: XCTestCase {
     }
 
     func testBackPressureAndRelease() async throws {
-        let buffer = MessageBuffer(bufferSize: 10)
+        let buffer = getTestMessageBuffer(bufferSize: 10)
         try await buffer.enqueue(content: .string("1234567890"))
         async let eq1 = buffer.enqueue(content: .string("1"))
         async let eq2 = buffer.enqueue(content: .string("2"))
@@ -51,7 +57,7 @@ class MessageBufferTest: XCTestCase {
     }
 
     func testBackPressureAndRelease2() async throws {
-        let buffer = MessageBuffer(bufferSize: 10)
+        let buffer = getTestMessageBuffer(bufferSize: 10)
         let expect1 = XCTestExpectation(description: "Should not release 1")
         expect1.isInverted = true
         let expect2 = XCTestExpectation(description: "Should not release 2")
@@ -93,7 +99,7 @@ class MessageBufferTest: XCTestCase {
     }
 
     func testAckInvalidSequenceIdIgnored() async throws {
-        let buffer = MessageBuffer(bufferSize: 100)
+        let buffer = getTestMessageBuffer(bufferSize: 100)
         let rst = try await buffer.ack(sequenceId: 1) // without any send
         XCTAssertEqual(false, rst)
 
@@ -104,7 +110,7 @@ class MessageBufferTest: XCTestCase {
     }
 
     func testWaitToDequeueReturnsImmediatelyIfAvailable() async throws {
-        let buffer = MessageBuffer(bufferSize: 100)
+        let buffer = getTestMessageBuffer(bufferSize: 100)
         _ = try await buffer.enqueue(content: .string("msg"))
         let result = try await buffer.WaitToDequeue()
         XCTAssertTrue(result)
@@ -113,7 +119,7 @@ class MessageBufferTest: XCTestCase {
     }
 
     func testWaitToDequeueFirst() async throws {
-        let buffer = MessageBuffer(bufferSize: 100)
+        let buffer = getTestMessageBuffer(bufferSize: 100)
         async let dqueue: Bool = try await buffer.WaitToDequeue()
         try await Task.sleep(for: .milliseconds(10))
 
@@ -127,7 +133,7 @@ class MessageBufferTest: XCTestCase {
     }
 
     func testMultipleDequeueWait() async throws {
-        let buffer = MessageBuffer(bufferSize: 100)
+        let buffer = getTestMessageBuffer(bufferSize: 100)
         async let dqueue1: Bool = try await buffer.WaitToDequeue()
         async let dqueue2: Bool = try await buffer.WaitToDequeue()
         try await Task.sleep(for: .milliseconds(10))
@@ -143,13 +149,13 @@ class MessageBufferTest: XCTestCase {
     }
 
     func testTryDequeueReturnsNilIfEmpty() async throws {
-        let buffer = MessageBuffer(bufferSize: 100)
+        let buffer = getTestMessageBuffer(bufferSize: 100)
         let result = try await buffer.TryDequeue()
         XCTAssertNil(result)
     }
 
     func testResetDequeueResetsCorrectly() async throws {
-        let buffer = MessageBuffer(bufferSize: 100)
+        let buffer = getTestMessageBuffer(bufferSize: 100)
         try await buffer.enqueue(content: .string("test1"))
         try await buffer.enqueue(content: .string("test2"))
         let t1 = try await buffer.TryDequeue()
@@ -172,7 +178,7 @@ class MessageBufferTest: XCTestCase {
     }
 
     func testContinuousBackPressure() async throws {
-        let buffer = MessageBuffer(bufferSize: 5)
+        let buffer = getTestMessageBuffer(bufferSize: 5)
         var tasks: [Task<Void, any Error>] = []
         for i in 0..<100 {
             let task = Task {