From 912093245df3b1a97b61b9138214ca92da8d402f Mon Sep 17 00:00:00 2001
From: tiagosiebler <tiagosiebler@github.com>
Date: Mon, 17 Feb 2025 11:00:20 +0000
Subject: [PATCH] feat(): optional promisified subscribe requests

---
 src/types/websockets/ws-general.ts |  9 +++++++++
 src/util/BaseWSClient.ts           | 32 ++++++++++++++++++++++++------
 src/util/websockets/WsStore.ts     |  3 +++
 3 files changed, 38 insertions(+), 6 deletions(-)

diff --git a/src/types/websockets/ws-general.ts b/src/types/websockets/ws-general.ts
index bd4d52e..7a61f3d 100644
--- a/src/types/websockets/ws-general.ts
+++ b/src/types/websockets/ws-general.ts
@@ -123,6 +123,15 @@ export interface WSClientConfigurableOptions {
 
   wsUrl?: string;
 
+  /**
+   * Default: false.
+   *
+   * When enabled, any calls to the subscribe method will return a promise.
+   * Note: internally, subscription requests are sent in batches. This may not behave as expected when
+   * subscribing to a large number of topics, especially if you are not yet connected when subscribing.
+   */
+  promiseSubscribeRequests?: boolean;
+
   /**
    * Allows you to provide a custom "signMessage" function, e.g. to use node's much faster createHmac method
    *
diff --git a/src/util/BaseWSClient.ts b/src/util/BaseWSClient.ts
index f344907..24b7f76 100644
--- a/src/util/BaseWSClient.ts
+++ b/src/util/BaseWSClient.ts
@@ -163,10 +163,15 @@ export abstract class BaseWebsocketClient<
       reconnectTimeout: 500,
       recvWindow: 5000,
 
+      // Calls to subscribeV5() are wrapped in a promise, allowing you to await a subscription request.
+      // Note: due to internal complexity, it's only recommended if you connect before subscribing.
+      promiseSubscribeRequests: false,
+
       // Automatically send an authentication op/request after a connection opens, for private connections.
       authPrivateConnectionsOnConnect: true,
       // Individual requests do not require a signature, so this is disabled.
       authPrivateRequests: false,
+
       ...options,
     };
   }
@@ -305,6 +310,9 @@ export abstract class BaseWebsocketClient<
 
       for (const requestKey in pendingSubReqs) {
         const request = pendingSubReqs[requestKey];
+        this.logger.trace(
+          `clearTopicsPendingSubscriptions(${wsKey}, ${rejectAll}, ${rejectReason}, ${requestKey}): rejecting promise for: ${JSON.stringify(request?.requestData || {})}`,
+        );
         request?.rejector(request.requestData, rejectReason);
       }
     }
@@ -854,9 +862,11 @@ export abstract class BaseWebsocketClient<
     for (const midflightRequest of subscribeWsMessages) {
       const wsMessage = midflightRequest.requestEvent;
 
-      promises.push(
-        this.upsertPendingTopicSubscribeRequests(wsKey, midflightRequest),
-      );
+      if (this.options.promiseSubscribeRequests) {
+        promises.push(
+          this.upsertPendingTopicSubscribeRequests(wsKey, midflightRequest),
+        );
+      }
 
       this.logger.trace(
         `Sending batch via message: "${JSON.stringify(wsMessage)}"`,
@@ -899,9 +909,11 @@ export abstract class BaseWebsocketClient<
     for (const midflightRequest of subscribeWsMessages) {
       const wsMessage = midflightRequest.requestEvent;
 
-      promises.push(
-        this.upsertPendingTopicSubscribeRequests(wsKey, midflightRequest),
-      );
+      if (this.options.promiseSubscribeRequests) {
+        promises.push(
+          this.upsertPendingTopicSubscribeRequests(wsKey, midflightRequest),
+        );
+      }
 
       this.logger.trace(`Sending batch via message: "${wsMessage}"`);
       this.tryWsSend(wsKey, JSON.stringify(wsMessage));
@@ -1004,6 +1016,7 @@ export abstract class BaseWebsocketClient<
     } catch (e) {
       this.logger.error(
         'Exception trying to resolve "connectionInProgress" promise',
+        e,
       );
     }
 
@@ -1073,6 +1086,7 @@ export abstract class BaseWebsocketClient<
     } catch (e) {
       this.logger.error(
         'Exception trying to resolve "connectionInProgress" promise',
+        e,
       );
     }
 
@@ -1216,6 +1230,9 @@ export abstract class BaseWebsocketClient<
     if (
       this.wsStore.getConnectionState(wsKey) !== WsConnectionStateEnum.CLOSING
     ) {
+      this.logger.trace(
+        `onWsClose(${wsKey}): rejecting all deferred promises...`,
+      );
       // clean up any pending promises for this connection
       this.getWsStore().rejectAllDeferredPromises(
         wsKey,
@@ -1230,6 +1247,9 @@ export abstract class BaseWebsocketClient<
       this.emit('reconnect', { wsKey, event });
     } else {
       // clean up any pending promises for this connection
+      this.logger.trace(
+        `onWsClose(${wsKey}): rejecting all deferred promises...`,
+      );
       this.getWsStore().rejectAllDeferredPromises(wsKey, 'disconnected');
       this.setWsState(wsKey, WsConnectionStateEnum.INITIAL);
       this.emit('close', { wsKey, event });
diff --git a/src/util/websockets/WsStore.ts b/src/util/websockets/WsStore.ts
index 2744deb..6b42b2b 100644
--- a/src/util/websockets/WsStore.ts
+++ b/src/util/websockets/WsStore.ts
@@ -205,6 +205,9 @@ export class WsStore<
     const promise = this.getDeferredPromise(wsKey, promiseRef);
 
     if (promise?.reject) {
+      this.logger.trace(
+        `rejectDeferredPromise(): rejecting ${wsKey}/${promiseRef}/${value}`,
+      );
       promise.reject(value);
     }