udpate message buffer and ut

Author: xingsy97

Sources/SignalRClient/ConnectionProtocol.swift

@@ -5,7 +5,6 @@ enum ConnectionFeature: String, CaseIterable {
     case Reconnect = "reconnect"
     case Resend = "resend"
     case Disconnected = "disconnected"
-    // Add more feature keys as needed
 }
 
 protocol ConnectionProtocol: AnyObject, Sendable {

Sources/SignalRClient/HttpConnection.swift

@@ -385,13 +385,11 @@ actor HttpConnection: ConnectionProtocol {
                 }
             }   
         }
-
-        do {
-            try await transport!.connect(url: url, transferFormat: transferFormat)
-        } catch {
-            await transport!.onReceive(nil)
-            await transport!.onClose(nil)
-            throw error
+        else { 
+            await transport!.onClose { [weak self] error in
+                guard let self = self else { return }
+                await self.handleConnectionClose(error: error)
+            }
         }
     }
 

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,17 +40,16 @@ public actor HubConnection {
     internal init(connection: ConnectionProtocol,
                   logger: Logger,
                   hubProtocol: HubProtocol,
-                  reconnectPolicy: RetryPolicy,
+                  retryPolicy: RetryPolicy,
                   serverTimeout: TimeInterval?,
                   keepAliveInterval: TimeInterval?,
                   statefulReconnectBufferSize: Int?) {
         self.serverTimeout = serverTimeout ?? HubConnection.defaultTimeout
         self.keepAliveInterval = keepAliveInterval ?? HubConnection.defaultPingInterval
-        self.statefulReconnectBufferSize =
-            statefulReconnectBufferSize ?? HubConnection.defaultStatefulReconnectBufferSize
+        self.statefulReconnectBufferSize = statefulReconnectBufferSize ?? HubConnection.defaultStatefulReconnectBufferSize
 
         self.logger = logger
-        self.reconnectPolicy = reconnectPolicy
+        self.retryPolicy = retryPolicy
 
         self.connection = connection
         self.hubProtocol = hubProtocol
@@ -64,7 +63,7 @@ public actor HubConnection {
     }
 
     public func start() async throws {
-        if connectionStatus != .Stopped {
+        if (connectionStatus != .Stopped) {
             throw SignalRError.invalidOperation("Start client while not in a stopped state.")
         }
 
@@ -345,7 +344,7 @@ public actor HubConnection {
         // 2. Connected: In this case, we should reconnect
         // 3. Reconnecting: In this case, we're in the control of previous reconnect(), let that function handle the reconnection
 
-        if connectionStatus == .Connected {
+        if (connectionStatus == .Connected) {
             do {
                 try await reconnect(error: error)
             } catch {
@@ -361,13 +360,11 @@ public actor HubConnection {
         var lastError: Error? = error
 
         // reconnect
-        while let interval = reconnectPolicy.nextRetryInterval(
-            retryContext: RetryContext(
-                retryCount: retryCount,
-                elapsed: elapsed,
-                retryReason: lastError
-            ))
-        {
+        while let interval = retryPolicy.nextRetryInterval(retryContext: RetryContext(
+            retryCount: retryCount,
+            elapsed: elapsed,
+            retryReason: lastError
+        )) {
             try Task.checkCancellation()
             if connectionStatus == .Stopping {
                 break
@@ -436,16 +433,11 @@ public actor HubConnection {
         do {
             let hubMessage = try hubProtocol.parseMessages(input: data!, binder: invocationBinder)
             for message in hubMessage {
-                do {
-                    if let messageBuffer = self.messageBuffer {
-                        let shouldProcess = try await messageBuffer.shouldProcessMessage(message)
-                        if !shouldProcess {
-                            // Don't process the message, we are either waiting for a SequenceMessage or received a duplicate message
-                            continue
-                        }
+                if let messageBuffer = self.messageBuffer {
+                    if !(try await messageBuffer.shouldProcessMessage(message)) {
+                        // Don't process the message, we are either waiting for a SequenceMessage or received a duplicate message
+                        continue
                     }
-                } catch {
-                    logger.log(level: .error, message: "Error parsing messages: \(error)")
                 }
                 await dispatchMessage(message)
             }
@@ -482,11 +474,18 @@ public actor HubConnection {
         case _ as CloseMessage:
             // Close
             break
-        case _ as AckMessage:
-            // TODO: In stateful reconnect
+        case let message as AckMessage:
+            let result = await self.messageBuffer?.ack(sequenceId: message.sequenceId);
+            if (result == false) {
+                logger.log(level: .warning, message: "Ack message received for sequenceId: \(message.sequenceId), but failed.")
+            }
             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)")
@@ -554,7 +553,7 @@ public actor HubConnection {
             logger.log(level: .error, message: "Unsupported handshake version: \(version)")
             throw SignalRError.unsupportedHandshakeVersion
         }
-        // TODO: enable version 2 when stateful reconnect is done
+        // TODO: enable version 2 when stateful reconnect is ready
 
         receivedHandshakeResponse = false
         let handshakeRequest = HandshakeRequestMessage(protocol: hubProtocol.name, version: version)
@@ -593,17 +592,19 @@ 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
-                    _ = try? await self?.messageBuffer?.disconnected()
+                        _ = await self?.messageBuffer?.disconnected()
                 })
                 await self.connection.setFeature(
                     feature: ConnectionFeature.Resend, 
                     value: { [weak self] () async -> Any? in
-                    return try? await self?.messageBuffer?.resend()
+                        return try? await self?.messageBuffer?.resend()
                 })
             }
 

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

@@ -4,38 +4,131 @@
 import Foundation
 
 actor MessageBuffer {
-    private var hubProtocol: HubProtocol;
-    private var connection: ConnectionProtocol;
-
     private var maxBufferSize: Int
     private var messages: [BufferedItem] = []
     private var bufferedByteCount: Int = 0
     private var totalMessageCount: Int = 0
     private var lastSendSequenceId: Int = 0
     private var nextSendIdx = 0
+    private var nextReceivingIdx: Int64 = 1
+    private var lastReceivedSequenceId: Int64 = 0
     private var dequeueContinuations: [CheckedContinuation<Bool, Never>] = []
     private var closed: Bool = false
+    private var reconnectInprogress: Bool = false;
+    private var waitForSequenceMessage: Bool = false;
+
+    private var ackTimerHandle: DispatchWorkItem?
+
+    private var hubProtocol: HubProtocol;
+    private var connection: ConnectionProtocol;
 
-    init(hubProtocol: HubProtocol, connection: ConnectionProtocol, bufferSize: Int) {
+    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() -> 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)
     }
 
-    public func disconnected() async throws -> Void {
-        throw SignalRError.invalidOperation("Disconnected is not implemented")
+    private func performScheduledAck() async {
+        defer {
+            // 在方法结束时清理定时器
+            ackTimerHandle = nil
+        }
+        
+        do {
+            if !reconnectInprogress {
+                let ackMessage = AckMessage(
+                    sequenceId: lastReceivedSequenceId
+                )
+                
+                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.lastReceivedSequenceId{
+            if currentId == self.lastReceivedSequenceId{
+                // 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.lastReceivedSequenceId = 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 {
@@ -71,7 +164,7 @@ actor MessageBuffer {
         }
     }
 
-    public func ack(sequenceId: Int) throws -> Bool {
+    public func ack(sequenceId: Int64) -> Bool {
         // It might be wrong ack or the ack of previous connection
         if (sequenceId <= 0 || sequenceId > lastSendSequenceId) {
             return false
@@ -137,6 +230,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 {