implement client subscription id and asset selection

Author: premiumjibles

node/proxy/api/src/coincap.ts

@@ -3,9 +3,10 @@ import WebSocket from 'ws'
 import { BaseWebsocketClient, Args, Options, BaseConnectionHandler } from '@shapeshiftoss/websocket'
 import { MarketDataClient, MarketDataConnectionHandler, MarketDataMessage } from './marketData'
 
-// TODO: track assets and subscriptionId for clients
 export class CoincapWebsocketClient extends BaseWebsocketClient implements MarketDataClient {
   clients = new Map<string, BaseConnectionHandler>()
+  // clientId -> subscriptionId -> assets
+  subscriptions = new Map<string, Map<string, string[]>>()
 
   constructor(url: string, args: Args, opts?: Options) {
     super(url, { logger: args.logger }, opts)
@@ -27,29 +28,71 @@ export class CoincapWebsocketClient extends BaseWebsocketClient implements Marke
   }
 
   subscribe(clientId: string, subscriptionId: string, connection: MarketDataConnectionHandler, assets: Array<string>) {
-    console.log(this.clients.size, { subscriptionId, assets })
     if (!this.clients.size) this.connect()
     this.clients.set(clientId, connection)
+
+    // Get or create the client's subscription map
+    let clientSubscriptions = this.subscriptions.get(clientId)
+    if (!clientSubscriptions) {
+      clientSubscriptions = new Map<string, string[]>()
+      this.subscriptions.set(clientId, clientSubscriptions)
+    }
+
+    // Add this subscription
+    clientSubscriptions.set(subscriptionId, assets)
   }
 
-  unsubscribe(clientId: string, subscriptionId: string, assets: Array<string>) {
-    console.log({ subscriptionId, assets })
-    if (!this.clients.has(clientId)) return
-    this.clients.delete(clientId)
+  unsubscribe(clientId: string, subscriptionId?: string) {
+    const clientSubscriptions = this.subscriptions.get(clientId)
+    if (!clientSubscriptions) return
+
+    if (subscriptionId) {
+      // Remove specific subscription
+      clientSubscriptions.delete(subscriptionId)
+
+      // If client has no more subscriptions, remove them entirely
+      if (clientSubscriptions.size === 0) {
+        this.clients.delete(clientId)
+        this.subscriptions.delete(clientId)
+      }
+    } else {
+      // Remove all subscriptions for this client
+      this.clients.delete(clientId)
+      this.subscriptions.delete(clientId)
+    }
+
+    // Close connection if no more clients
     if (!this.clients.size) this.socket?.close(1000)
   }
 
   private handleMessage(message: Record<string, string>): void {
     for (const [clientId, client] of this.clients) {
       try {
-        const payload: MarketDataMessage = {
-          type: 'price_update',
-          source: 'coincap',
-          data: message,
-          timestamp: Date.now(),
-        }
+        const clientSubscriptions = this.subscriptions.get(clientId)
+        if (!clientSubscriptions) continue
 
-        client.publish(clientId, payload)
+        // Send updates for each subscription
+        for (const [subscriptionId, assets] of clientSubscriptions) {
+          // Filter data to only include assets this subscription requested
+          const filteredData: Record<string, string> = {}
+          for (const asset of assets) {
+            if (message[asset] !== undefined) {
+              filteredData[asset] = message[asset]
+            }
+          }
+
+          // Only send if there's relevant data for this subscription
+          if (Object.keys(filteredData).length > 0) {
+            const payload: MarketDataMessage = {
+              type: 'price_update',
+              source: 'coincap',
+              data: filteredData,
+              timestamp: Date.now(),
+            }
+
+            client.publish(subscriptionId, payload)
+          }
+        }
       } catch (error) {
         this.logger.error({ clientId, error }, 'failed to handle message')
       }

node/proxy/api/src/marketData.ts

@@ -26,7 +26,7 @@ export interface MarketDataClient {
     connection: MarketDataConnectionHandler,
     assets: Array<string>
   ): void
-  unsubscribe(clientId: string, subscriptionId: string, assets: Array<string>): void
+  unsubscribe(clientId: string, subscriptionId?: string): void
 }
 
 export class MarketDataConnectionHandler extends BaseConnectionHandler {
@@ -65,19 +65,14 @@ export class MarketDataConnectionHandler extends BaseConnectionHandler {
 
     if (subscriptionId) {
       this.subscriptionIds.delete(subscriptionId)
-      this.client.unsubscribe(this.clientId, subscriptionId, data.assets)
+      this.client.unsubscribe(this.clientId, subscriptionId)
     } else {
-      for (const subscriptionId of this.subscriptionIds) {
-        this.client.unsubscribe(this.clientId, subscriptionId, [])
-      }
-
       this.subscriptionIds.clear()
+      this.client.unsubscribe(this.clientId)
     }
   }
 
   onClose(): void {
-    for (const subscriptionId of this.subscriptionIds) {
-      this.client.unsubscribe(this.clientId, subscriptionId, [])
-    }
+    this.client.unsubscribe(this.clientId)
   }
 }