fix: Error: WebSocket is not open: readyState 2 (CLOSING) (#5846)

Author: paveltiunov

packages/cubejs-backend-shared/src/env.ts

@@ -1315,6 +1315,9 @@ const variables: Record<string, (...args: any) => any> = {
     .asString(),
   cubeStorePass: () => get('CUBEJS_CUBESTORE_PASS')
     .asString(),
+  cubeStoreMaxConnectRetries: () => get('CUBEJS_CUBESTORE_MAX_CONNECT_RETRIES')
+    .default('5')
+    .asInt(),
 
   // Redis
   redisPoolMin: () => get('CUBEJS_REDIS_POOL_MIN')

packages/cubejs-cubestore-driver/src/WebSocketConnection.ts

@@ -1,6 +1,7 @@
 import WebSocket from 'ws';
 import { flatbuffers } from 'flatbuffers';
 import { InlineTable } from '@cubejs-backend/base-driver';
+import { getEnv } from '@cubejs-backend/shared';
 import {
   HttpCommand,
   HttpError,
@@ -13,13 +14,19 @@ import {
 export class WebSocketConnection {
   protected messageCounter: number;
 
+  protected maxConnectRetries: number;
+
+  protected currentConnectionTry: number;
+
   protected webSocket: any;
 
   private url: string;
 
   public constructor(url: string) {
     this.url = url;
     this.messageCounter = 1;
+    this.maxConnectRetries = getEnv('cubeStoreMaxConnectRetries');
+    this.currentConnectionTry = 0;
   }
 
   protected async initWebSocket() {
@@ -38,22 +45,33 @@ export class WebSocketConnection {
         }, 5000);
 
         webSocket.sendAsync = async (message) => new Promise<void>((resolveSend, rejectSend) => {
-          webSocket.send(message, (err) => {
-            if (err) {
-              rejectSend(err);
-            } else {
-              resolveSend();
-            }
-          });
+          // If socket is closing this message should be resent
+          if (webSocket.readyState === WebSocket.OPEN) {
+            webSocket.send(message, (err) => {
+              if (err) {
+                rejectSend(err);
+              } else {
+                resolveSend();
+              }
+            });
+          }
         });
         webSocket.on('open', () => resolve(webSocket));
         webSocket.on('error', (err) => {
-          reject(err);
+          this.currentConnectionTry += 1;
+          if (this.currentConnectionTry < this.maxConnectRetries) {
+            setTimeout(async () => {
+              resolve(this.initWebSocket());
+            }, this.retryWaitTime());
+          } else {
+            reject(err);
+          }
           if (webSocket === this.webSocket) {
             this.webSocket = undefined;
           }
         });
         webSocket.on('pong', () => {
+          this.currentConnectionTry = 0;
           webSocket.lastHeartBeat = new Date();
         });
         webSocket.on('close', () => {
@@ -74,7 +92,7 @@ export class WebSocketConnection {
                   webSocket.sentMessages[key].reject(e);
                 }
               }
-            }, 1000);
+            }, this.retryWaitTime());
           }
 
           if (webSocket === this.webSocket) {
@@ -136,20 +154,28 @@ export class WebSocketConnection {
     return this.webSocket.readyPromise;
   }
 
+  private retryWaitTime() {
+    return 1000 * this.currentConnectionTry;
+  }
+
   private async sendMessage(messageId: number, buffer: Uint8Array): Promise<any> {
     const socket = await this.initWebSocket();
     return new Promise((resolve, reject) => {
-      socket.send(buffer, (err) => {
-        if (err) {
-          delete socket.sentMessages[messageId];
-          reject(err);
-        }
-      });
-      socket.sentMessages[messageId] = {
-        resolve,
-        reject,
-        buffer
-      };
+      if (socket.readyState === WebSocket.OPEN) {
+        socket.send(buffer, (err) => {
+          if (err) {
+            delete socket.sentMessages[messageId];
+            reject(err);
+          }
+        });
+        socket.sentMessages[messageId] = {
+          resolve,
+          reject,
+          buffer
+        };
+      } else {
+        resolve(this.sendMessage(messageId, buffer));
+      }
     });
   }