[SignalR][TS client] Stateful reconnect (#49940)

Author: BrennanConroy

src/SignalR/clients/ts/signalr-protocol-msgpack/src/MessagePackHubProtocol.ts

@@ -6,8 +6,9 @@ import { Encoder, Decoder } from "@msgpack/msgpack";
 import { MessagePackOptions } from "./MessagePackOptions";
 
 import {
+    AckMessage,
     CancelInvocationMessage, CompletionMessage, HubMessage, IHubProtocol, ILogger, InvocationMessage,
-    LogLevel, MessageHeaders, MessageType, NullLogger, StreamInvocationMessage, StreamItemMessage, TransferFormat,
+    LogLevel, MessageHeaders, MessageType, NullLogger, SequenceMessage, StreamInvocationMessage, StreamItemMessage, TransferFormat,
 } from "@microsoft/signalr";
 
 import { BinaryMessageFormat } from "./BinaryMessageFormat";
@@ -114,6 +115,10 @@ export class MessagePackHubProtocol implements IHubProtocol {
                 return this._writeCancelInvocation(message as CancelInvocationMessage);
             case MessageType.Close:
                 return this._writeClose();
+            case MessageType.Ack:
+                return this._writeAck(message as AckMessage);
+            case MessageType.Sequence:
+                return this._writeSequence(message as SequenceMessage);
             default:
                 throw new Error("Invalid message type.");
         }
@@ -142,6 +147,10 @@ export class MessagePackHubProtocol implements IHubProtocol {
                 return this._createPingMessage(properties);
             case MessageType.Close:
                 return this._createCloseMessage(properties);
+            case MessageType.Ack:
+                return this._createAckMessage(properties);
+            case MessageType.Sequence:
+                return this._createSequenceMessage(properties);
             default:
                 // Future protocol changes can add message types, old clients can ignore them
                 logger.log(LogLevel.Information, "Unknown message type '" + messageType + "' ignored.");
@@ -252,6 +261,30 @@ export class MessagePackHubProtocol implements IHubProtocol {
         return completionMessage;
     }
 
+    private _createAckMessage(properties: any[]): HubMessage {
+        // check minimum length to allow protocol to add items to the end of objects in future releases
+        if (properties.length < 1) {
+            throw new Error("Invalid payload for Ack message.");
+        }
+
+        return {
+            sequenceId: properties[1],
+            type: MessageType.Ack,
+        } as HubMessage;
+    }
+
+    private _createSequenceMessage(properties: any[]): HubMessage {
+        // check minimum length to allow protocol to add items to the end of objects in future releases
+        if (properties.length < 1) {
+            throw new Error("Invalid payload for Sequence message.");
+        }
+
+        return {
+            sequenceId: properties[1],
+            type: MessageType.Sequence,
+        } as HubMessage;
+    }
+
     private _writeInvocation(invocationMessage: InvocationMessage): ArrayBuffer {
         let payload: any;
         if (invocationMessage.streamIds) {
@@ -317,6 +350,18 @@ export class MessagePackHubProtocol implements IHubProtocol {
         return BinaryMessageFormat.write(payload.slice());
     }
 
+    private _writeAck(ackMessage: AckMessage): ArrayBuffer {
+        const payload = this._encoder.encode([MessageType.Ack, ackMessage.sequenceId]);
+
+        return BinaryMessageFormat.write(payload.slice());
+    }
+
+    private _writeSequence(sequenceMessage: SequenceMessage): ArrayBuffer {
+        const payload = this._encoder.encode([MessageType.Sequence, sequenceMessage.sequenceId]);
+
+        return BinaryMessageFormat.write(payload.slice());
+    }
+
     private _readHeaders(properties: any): MessageHeaders {
         const headers: MessageHeaders = properties[1] as MessageHeaders;
         if (typeof headers !== "object") {

src/SignalR/clients/ts/signalr-protocol-msgpack/tests/MessagePackHubProtocol.test.ts

@@ -1,7 +1,7 @@
 // Licensed to the .NET Foundation under one or more agreements.
 // The .NET Foundation licenses this file to you under the MIT license.
 
-import { CloseMessage, CompletionMessage, InvocationMessage, MessageType, NullLogger, StreamItemMessage } from "@microsoft/signalr";
+import { AckMessage, CloseMessage, CompletionMessage, InvocationMessage, MessageType, NullLogger, SequenceMessage, StreamItemMessage } from "@microsoft/signalr";
 import { MessagePackHubProtocol } from "../src/MessagePackHubProtocol";
 
 describe("MessagePackHubProtocol", () => {
@@ -284,4 +284,35 @@ describe("MessagePackHubProtocol", () => {
         expect(parsedMessages.length).toEqual(1);
         expect(parsedMessages[0].type).toEqual(MessageType.Close);
     });
+
+    it("can write/read Ack message", () => {
+        const ackMessage = {
+            sequenceId: 13,
+            type: MessageType.Ack,
+        } as AckMessage;
+
+        const protocol = new MessagePackHubProtocol();
+        const parsedMessages = protocol.parseMessages(protocol.writeMessage(ackMessage), NullLogger.instance);
+        expect(parsedMessages.length).toEqual(1);
+        expect(parsedMessages[0].type).toEqual(MessageType.Ack);
+        expect(parsedMessages[0]).toEqual({
+            sequenceId: 13,
+            type: MessageType.Ack
+        });
+    });
+
+    it("can write/read Sequence message", () => {
+        const sequenceMessage = {
+            sequenceId: 24,
+            type: MessageType.Sequence,
+        } as SequenceMessage;
+
+        const protocol = new MessagePackHubProtocol();
+        const parsedMessages = protocol.parseMessages(protocol.writeMessage(sequenceMessage), NullLogger.instance);
+        expect(parsedMessages.length).toEqual(1);
+        expect(parsedMessages[0]).toEqual({
+            sequenceId: 24,
+            type: MessageType.Sequence
+        });
+    });
 });

src/SignalR/clients/ts/signalr/src/HttpConnection.ts

@@ -31,6 +31,7 @@ export interface INegotiateResponse {
     url?: string;
     accessToken?: string;
     error?: string;
+    useAck?: boolean;
 }
 
 /** @private */
@@ -331,6 +332,11 @@ export class HttpConnection implements IConnection {
                 // So we set it equal to connectionId so all our logic can use connectionToken without being aware of the negotiate version
                 negotiateResponse.connectionToken = negotiateResponse.connectionId;
             }
+
+            if (negotiateResponse.useAck && this._options._useStatefulReconnect !== true) {
+                return Promise.reject(new FailedToNegotiateWithServerError("Client didn't negotiate Stateful Reconnect but the server did."));
+            }
+
             return negotiateResponse;
         } catch (e) {
             let errorMessage = "Failed to complete negotiation with the server: " + e;
@@ -368,7 +374,7 @@ export class HttpConnection implements IConnection {
         const transports = negotiateResponse.availableTransports || [];
         let negotiate: INegotiateResponse | undefined = negotiateResponse;
         for (const endpoint of transports) {
-            const transportOrError = this._resolveTransportOrError(endpoint, requestedTransport, requestedTransferFormat);
+            const transportOrError = this._resolveTransportOrError(endpoint, requestedTransport, requestedTransferFormat, negotiate?.useAck === true);
             if (transportOrError instanceof Error) {
                 // Store the error and continue, we don't want to cause a re-negotiate in these cases
                 transportExceptions.push(`${endpoint.transport} failed:`);
@@ -413,7 +419,8 @@ export class HttpConnection implements IConnection {
                 if (!this._options.WebSocket) {
                     throw new Error("'WebSocket' is not supported in your environment.");
                 }
-                return new WebSocketTransport(this._httpClient, this._accessTokenFactory, this._logger, this._options.logMessageContent!, this._options.WebSocket, this._options.headers || {});
+                return new WebSocketTransport(this._httpClient, this._accessTokenFactory, this._logger, this._options.logMessageContent!,
+                    this._options.WebSocket, this._options.headers || {});
             case HttpTransportType.ServerSentEvents:
                 if (!this._options.EventSource) {
                     throw new Error("'EventSource' is not supported in your environment.");
@@ -428,11 +435,34 @@ export class HttpConnection implements IConnection {
 
     private _startTransport(url: string, transferFormat: TransferFormat): Promise<void> {
         this.transport!.onreceive = this.onreceive;
-        this.transport!.onclose = (e) => this._stopConnection(e);
+        if (this.features.reconnect) {
+            this.transport!.onclose = async (e) => {
+                let callStop = false;
+                if (this.features.resend) {
+                    try {
+                        this.features.disconnected();
+                        await this.transport!.connect(url, transferFormat);
+                        await this.features.resend();
+                    } catch {
+                        callStop = true;
+                    }
+                } else {
+                    this._stopConnection(e);
+                    return;
+                }
+
+                if (callStop) {
+                    this._stopConnection(e);
+                }
+            };
+        } else {
+            this.transport!.onclose = (e) => this._stopConnection(e);
+        }
         return this.transport!.connect(url, transferFormat);
     }
 
-    private _resolveTransportOrError(endpoint: IAvailableTransport, requestedTransport: HttpTransportType | undefined, requestedTransferFormat: TransferFormat): ITransport | Error | unknown {
+    private _resolveTransportOrError(endpoint: IAvailableTransport, requestedTransport: HttpTransportType | undefined,
+        requestedTransferFormat: TransferFormat, useAcks: boolean): ITransport | Error | unknown {
         const transport = HttpTransportType[endpoint.transport];
         if (transport === null || transport === undefined) {
             this._logger.log(LogLevel.Debug, `Skipping transport '${endpoint.transport}' because it is not supported by this client.`);
@@ -448,6 +478,7 @@ export class HttpConnection implements IConnection {
                     } else {
                         this._logger.log(LogLevel.Debug, `Selecting transport '${HttpTransportType[transport]}'.`);
                         try {
+                            this.features.reconnect = transport === HttpTransportType.WebSockets ? useAcks : undefined;
                             return this._constructTransport(transport);
                         } catch (ex) {
                             return ex;
@@ -544,19 +575,30 @@ export class HttpConnection implements IConnection {
     }
 
     private _resolveNegotiateUrl(url: string): string {
-        const index = url.indexOf("?");
-        let negotiateUrl = url.substring(0, index === -1 ? url.length : index);
-        if (negotiateUrl[negotiateUrl.length - 1] !== "/") {
-            negotiateUrl += "/";
+        const negotiateUrl = new URL(url);
+
+        if (negotiateUrl.pathname.endsWith('/')) {
+            negotiateUrl.pathname += "negotiate";
+        } else {
+            negotiateUrl.pathname += "/negotiate";
         }
-        negotiateUrl += "negotiate";
-        negotiateUrl += index === -1 ? "" : url.substring(index);
+        const searchParams = new URLSearchParams(negotiateUrl.searchParams);
 
-        if (negotiateUrl.indexOf("negotiateVersion") === -1) {
-            negotiateUrl += index === -1 ? "?" : "&";
-            negotiateUrl += "negotiateVersion=" + this._negotiateVersion;
+        if (!searchParams.has("negotiateVersion")) {
+            searchParams.append("negotiateVersion", this._negotiateVersion.toString());
         }
-        return negotiateUrl;
+
+        if (searchParams.has("useAck")) {
+            if (searchParams.get("useAck") === "true") {
+                this._options._useStatefulReconnect = true;
+            }
+        } else if (this._options._useStatefulReconnect === true) {
+            searchParams.append("useAck", "true");
+        }
+
+        negotiateUrl.search = searchParams.toString();
+
+        return negotiateUrl.toString();
     }
 }
 

src/SignalR/clients/ts/signalr/src/HubConnection.ts

@@ -10,9 +10,11 @@ import { IRetryPolicy } from "./IRetryPolicy";
 import { IStreamResult } from "./Stream";
 import { Subject } from "./Subject";
 import { Arg, getErrorString, Platform } from "./Utils";
+import { MessageBuffer } from "./MessageBuffer";
 
 const DEFAULT_TIMEOUT_IN_MS: number = 30 * 1000;
 const DEFAULT_PING_INTERVAL_IN_MS: number = 15 * 1000;
+const DEFAULT_STATEFUL_RECONNECT_BUFFER_SIZE = 100_000;
 
 /** Describes the current state of the {@link HubConnection} to the server. */
 export enum HubConnectionState {
@@ -36,11 +38,13 @@ export class HubConnection {
     private readonly connection: IConnection;
     private readonly _logger: ILogger;
     private readonly _reconnectPolicy?: IRetryPolicy;
+    private readonly _statefulReconnectBufferSize: number;
     private _protocol: IHubProtocol;
     private _handshakeProtocol: HandshakeProtocol;
     private _callbacks: { [invocationId: string]: (invocationEvent: StreamItemMessage | CompletionMessage | null, error?: Error) => void };
     private _methods: { [name: string]: (((...args: any[]) => void) | ((...args: any[]) => any))[] };
     private _invocationId: number;
+    private _messageBuffer?: MessageBuffer;
 
     private _closedCallbacks: ((error?: Error) => void)[];
     private _reconnectingCallbacks: ((error?: Error) => void)[];
@@ -98,8 +102,10 @@ export class HubConnection {
         protocol: IHubProtocol,
         reconnectPolicy?: IRetryPolicy,
         serverTimeoutInMilliseconds?: number,
-        keepAliveIntervalInMilliseconds?: number): HubConnection {
-        return new HubConnection(connection, logger, protocol, reconnectPolicy, serverTimeoutInMilliseconds, keepAliveIntervalInMilliseconds);
+        keepAliveIntervalInMilliseconds?: number,
+        statefulReconnectBufferSize?: number): HubConnection {
+        return new HubConnection(connection, logger, protocol, reconnectPolicy,
+            serverTimeoutInMilliseconds, keepAliveIntervalInMilliseconds, statefulReconnectBufferSize);
     }
 
     private constructor(
@@ -108,14 +114,17 @@ export class HubConnection {
         protocol: IHubProtocol,
         reconnectPolicy?: IRetryPolicy,
         serverTimeoutInMilliseconds?: number,
-        keepAliveIntervalInMilliseconds?: number) {
+        keepAliveIntervalInMilliseconds?: number,
+        statefulReconnectBufferSize?: number) {
         Arg.isRequired(connection, "connection");
         Arg.isRequired(logger, "logger");
         Arg.isRequired(protocol, "protocol");
 
         this.serverTimeoutInMilliseconds = serverTimeoutInMilliseconds ?? DEFAULT_TIMEOUT_IN_MS;
         this.keepAliveIntervalInMilliseconds = keepAliveIntervalInMilliseconds ?? DEFAULT_PING_INTERVAL_IN_MS;
 
+        this._statefulReconnectBufferSize = statefulReconnectBufferSize ?? DEFAULT_STATEFUL_RECONNECT_BUFFER_SIZE;
+
         this._logger = logger;
         this._protocol = protocol;
         this.connection = connection;
@@ -248,6 +257,17 @@ export class HubConnection {
                 throw this._stopDuringStartError;
             }
 
+            const useAck = this.connection.features.reconnect || false;
+            if (useAck) {
+                this._messageBuffer = new MessageBuffer(this._protocol, this.connection, this._statefulReconnectBufferSize);
+                this.connection.features.disconnected = this._messageBuffer._disconnected.bind(this._messageBuffer);
+                this.connection.features.resend = () => {
+                    if (this._messageBuffer) {
+                        return this._messageBuffer._resend();
+                    }
+                }
+            }
+
             if (!this.connection.features.inherentKeepAlive) {
                 await this._sendMessage(this._cachedPingMessage);
             }
@@ -271,6 +291,7 @@ export class HubConnection {
     public async stop(): Promise<void> {
         // Capture the start promise before the connection might be restarted in an onclose callback.
         const startPromise = this._startPromise;
+        this.connection.features.reconnect = false;
 
         this._stopPromise = this._stopInternal();
         await this._stopPromise;
@@ -399,7 +420,11 @@ export class HubConnection {
      * @param message The js object to serialize and send.
      */
     private _sendWithProtocol(message: any) {
-        return this._sendMessage(this._protocol.writeMessage(message));
+        if (this._messageBuffer) {
+            return this._messageBuffer._send(message);
+        } else {
+            return this._sendMessage(this._protocol.writeMessage(message));
+        }
     }
 
     /** Invokes a hub method on the server using the specified name and arguments. Does not wait for a response from the receiver.
@@ -575,6 +600,11 @@ export class HubConnection {
             const messages = this._protocol.parseMessages(data, this._logger);
 
             for (const message of messages) {
+                if (this._messageBuffer && !this._messageBuffer._shouldProcessMessage(message)) {
+                    // Don't process the message, we are either waiting for a SequenceMessage or received a duplicate message
+                    continue;
+                }
+
                 switch (message.type) {
                     case MessageType.Invocation:
                         // eslint-disable-next-line @typescript-eslint/no-floating-promises
@@ -616,6 +646,16 @@ export class HubConnection {
 
                         break;
                     }
+                    case MessageType.Ack:
+                        if (this._messageBuffer) {
+                            this._messageBuffer._ack(message);
+                        }
+                        break;
+                    case MessageType.Sequence:
+                        if (this._messageBuffer) {
+                            this._messageBuffer._resetSequence(message);
+                        }
+                        break;
                     default:
                         this._logger.log(LogLevel.Warning, `Invalid message type: ${message.type}.`);
                         break;
@@ -800,6 +840,10 @@ export class HubConnection {
         if (this._connectionStarted) {
             this._connectionState = HubConnectionState.Disconnected;
             this._connectionStarted = false;
+            if (this._messageBuffer) {
+                this._messageBuffer._dispose(error ?? new Error("Connection closed."));
+                this._messageBuffer = undefined;
+            }
 
             if (Platform.isBrowser) {
                 window.document.removeEventListener("freeze", this._freezeEventListener);