feat: implement failover handling for non-retriable callbacks and enhance tests for disconnect scenarios

Author: qevolg

nodejs/src/client/wsConnector.ts

@@ -286,6 +286,20 @@ export class WebSocketConnector {
         return this._suppressedSockets.has(conn);
     }
 
+    private async failNonRetriableCallbacksOnDisconnect(): Promise<void> {
+        const keepReqIds = new Set<bigint>(
+            this._inflightStore.getRequests().map((req) => req.reqId)
+        );
+        await WsEventCallback.instance().rejectCallbacksExceptReqIds(
+            keepReqIds,
+            new TDWebSocketClientError(
+                ErrorCode.ERR_CONNECTION_CLOSED,
+                "websocket connection closed before response was received"
+            ),
+            this._poolKey
+        );
+    }
+
     private async handleDisconnect(conn: w3cwebsocket, event?: ICloseEvent): Promise<void> {
         if (this.shouldSkipReconnect(conn) || this._isReconnecting) {
             return;
@@ -294,6 +308,7 @@ export class WebSocketConnector {
             logger.info("Websocket closed normally, skipping reconnect.");
             return;
         }
+        await this.failNonRetriableCallbacksOnDisconnect();
         try {
             await this.triggerReconnect();
         } catch (err: unknown) {
@@ -335,10 +350,16 @@ export class WebSocketConnector {
         }
 
         const errObj = err as { code?: unknown; message?: unknown };
-        const code = typeof errObj?.code === "string" ? errObj.code.toLowerCase() : "";
-        if (code.length > 0 && NETWORK_ERROR_CODES.has(code)) {
+        const code = errObj?.code;
+        if (typeof code === "number" && code === ErrorCode.ERR_CONNECTION_CLOSED) {
             return true;
         }
+        if (typeof code === "string") {
+            const loweredCode = code.toLowerCase();
+            if (loweredCode.length > 0 && NETWORK_ERROR_CODES.has(loweredCode)) {
+                return true;
+            }
+        }
 
         const message = err instanceof Error
             ? err.message
@@ -539,6 +560,7 @@ export class WebSocketConnector {
                 action: msg.action,
                 req_id: reqId,
                 timeout: this._timeout,
+                poolKey: this._poolKey,
             },
             resolveResp,
             rejectResp
@@ -644,6 +666,7 @@ export class WebSocketConnector {
                         req_id: reqId,
                         timeout: this._timeout,
                         id: callbackId,
+                        poolKey: this._poolKey,
                     },
                     safeResolve,
                     safeReject

nodejs/src/client/wsEventCallback.ts

@@ -12,6 +12,7 @@ interface MessageId {
     req_id: bigint;
     id?: bigint;
     timeout?: number;
+    poolKey?: string;
 }
 
 interface MessageAction {
@@ -156,4 +157,34 @@ export class WsEventCallback {
             release();
         }
     }
+
+    async rejectCallbacksExceptReqIds(
+        keepReqIds: Set<bigint>,
+        error: Error,
+        poolKey?: string
+    ): Promise<void> {
+        const toReject: Function[] = [];
+        const release = await eventMutex.acquire();
+        try {
+            for (const [k, v] of WsEventCallback._msgActionRegister) {
+                if (poolKey !== undefined && k.poolKey !== poolKey) {
+                    continue;
+                }
+                const keepByReqId = keepReqIds.has(k.req_id);
+                const keepByCallbackId = k.id !== undefined && keepReqIds.has(k.id);
+                if (keepByReqId || keepByCallbackId) {
+                    continue;
+                }
+                clearTimeout(v.timer);
+                WsEventCallback._msgActionRegister.delete(k);
+                toReject.push(v.reject);
+            }
+        } finally {
+            release();
+        }
+
+        for (const reject of toReject) {
+            reject(error);
+        }
+    }
 }

nodejs/test/bulkPulling/stmt2.func.test.ts

@@ -323,9 +323,6 @@ describe("TDWebSocket.Stmt()", () => {
             "('2024-12-19 17:12:47.642', 72.30000, 206, 0.31000) ";
         await wsSql.exec(insertQuery);
 
-        // let result = await wsSql.exec("select * from query_meters")
-        // console.log(result)
-
         let stmt = await wsSql.stmtInit();
         expect(stmt).toBeTruthy();
         expect(stmt).toBeInstanceOf(WsStmt2);

nodejs/test/bulkPulling/wsConnector.failover.test.ts

@@ -173,6 +173,64 @@ describe("WebSocketConnector failover and retry", () => {
         expect(connector.triggerReconnect).toHaveBeenCalledTimes(1);
     });
 
+    test("handleDisconnect rejects non-retriable pending callbacks immediately", async () => {
+        const connector = createBareConnector();
+        connector.triggerReconnect = jest.fn(() => new Promise<void>(() => { }));
+        connector._conn.send = jest.fn(() => { });
+
+        const pending = connector.sendMsg(
+            JSON.stringify({
+                action: "query",
+                args: {
+                    req_id: 1201,
+                },
+            })
+        );
+        void connector.handleDisconnect(
+            connector._conn,
+            { code: 1006, reason: "abnormal close" }
+        );
+
+        const state = await Promise.race([
+            pending.then(() => "resolved").catch(() => "rejected"),
+            delay(40).then(() => "pending"),
+        ]);
+
+        expect(state).toBe("rejected");
+        expect(connector.triggerReconnect).toHaveBeenCalledTimes(1);
+        expect(hasInflightRequest(connector, 1201n)).toBe(false);
+    });
+
+    test("handleDisconnect keeps retriable pending callbacks for replay", async () => {
+        const connector = createBareConnector();
+        connector.triggerReconnect = jest.fn(() => new Promise<void>(() => { }));
+        connector._conn.send = jest.fn(() => { });
+
+        const pending = connector.sendMsg(
+            JSON.stringify({
+                action: "insert",
+                args: {
+                    req_id: 1202,
+                },
+            })
+        );
+        void pending.catch(() => { });
+        void connector.handleDisconnect(
+            connector._conn,
+            { code: 1006, reason: "abnormal close" }
+        );
+
+        const state = await Promise.race([
+            pending.then(() => "resolved").catch(() => "rejected"),
+            delay(40).then(() => "pending"),
+        ]);
+
+        expect(state).toBe("pending");
+        expect(connector.triggerReconnect).toHaveBeenCalledTimes(1);
+        expect(hasInflightRequest(connector, 1202n)).toBe(true);
+        connector.failAllInflightRequests(new Error("cleanup"));
+    });
+
     test("_doReconnect fails all inflight requests when reconnect fails", async () => {
         const connector = createBareConnector();
         const rejectSpy = jest.fn();