fix: take a better approach when identifying subscription responses - use correlation

Author: maximovicd

src/connection/listeners.ts

@@ -1,4 +1,4 @@
-import { OpenfeedGatewayMessage, Result, SubscriptionType } from "@gen/openfeed_api";
+import { OpenfeedGatewayMessage, Result } from "@gen/openfeed_api";
 import { InstrumentDefinition } from "@gen/openfeed_instrument";
 import Long from "long";
 import { ActionType, HeartBeat } from "@gen/openfeed";
@@ -11,9 +11,9 @@ const IDGetters: ((msg: OpenfeedGatewayMessage) => Long | undefined)[] = [
     (msg) => msg.volumeAtPrice?.marketId,
 ];
 export class OpenFeedListeners {
-    private readonly instrumentByMarketId: Map<string, [InstrumentDefinition?, [string, SubscriptionType][]?]> = new Map<
+    private readonly instrumentByMarketId: Map<string, [InstrumentDefinition?, [string, string][]?]> = new Map<
         string,
-        [InstrumentDefinition?, [string, SubscriptionType][]?]
+        [InstrumentDefinition?, [string, string][]?]
     >();
 
     constructor() {
@@ -22,39 +22,44 @@ export class OpenFeedListeners {
 
     private addDetails = (message: OpenfeedGatewayMessage) => {
         let def: InstrumentDefinition | undefined;
-        let symbols: [string, SubscriptionType][] | undefined;
+        let symbols: [string, string][] | undefined;
 
         const getInstrumentDefinition = (marketId: Long) => {
             const res = this.instrumentByMarketId.get(marketId.toString());
 
             return res ?? [undefined, undefined];
         };
 
-        const includesSymbolSubscription = (arr: [string, SubscriptionType][], item: [string, SubscriptionType]) => {
-            return arr.some(([symbol, type]) => symbol === item[0] && type === item[1]);
+        const includesSymbolSubscription = (arr: [string, string][], item: [string, string]) => {
+            return arr.some(([symbol, correlationId]) => symbol === item[0] && correlationId === item[1]);
         };
 
         if (message.subscriptionResponse) {
-            const { marketId, symbol, unsubscribe, status, subscriptionType } = message.subscriptionResponse;
+            const { marketId, symbol, unsubscribe, status, correlationId } = message.subscriptionResponse;
+            const corIdStr = correlationId.toString();
             if (marketId !== Long.ZERO) {
                 [def, symbols] = getInstrumentDefinition(marketId);
                 if (status?.result === Result.SUCCESS) {
-                    const currentEntry: [string, SubscriptionType] = [symbol, subscriptionType];
+                    const currentEntry: [string, string] = [symbol, corIdStr];
+
                     if (!unsubscribe) {
                         if (!symbols) {
                             symbols = [currentEntry];
                         } else if (!includesSymbolSubscription(symbols, currentEntry)) {
                             symbols = [...symbols, currentEntry];
                         }
+                        this.instrumentByMarketId.set(marketId.toString(), [def, symbols]);
                     } else {
+                        let symbolsToSave: [string, string][] | undefined;
                         if (symbols) {
-                            symbols = symbols.filter(([s, t]) => !(s === symbol && t === subscriptionType));
+                            symbolsToSave = symbols.filter(([s, t]) => !(s === symbol && t === corIdStr));
                         }
-                        if (!symbols) {
+                        if (!symbolsToSave?.length) {
                             this.instrumentByMarketId.delete(marketId.toString());
+                        } else {
+                            this.instrumentByMarketId.set(marketId.toString(), [def, symbolsToSave]);
                         }
                     }
-                    this.instrumentByMarketId.set(marketId.toString(), [def, symbols]);
                 }
             }
         } else if (message.instrumentDefinition) {

test/listeners.test.ts

@@ -10,6 +10,7 @@ const ALIAS_REGULAR_NEW = "SYM*5";
 const MESSAGES = {
     SUBSCRIPTION_RESPONSE_MESSAGE_OLD: {
         subscriptionResponse: {
+            correlationId: Long.fromNumber(1),
             marketId: MARKET_ID,
             symbol: ALIAS_REGULAR_OLD,
             unsubscribe: false,
@@ -21,6 +22,7 @@ const MESSAGES = {
     } as Partial<OpenfeedGatewayMessage> as OpenfeedGatewayMessage,
     SUBSCRIPTION_RESPONSE_DEPTH_OLD: {
         subscriptionResponse: {
+            correlationId: Long.fromNumber(2),
             marketId: MARKET_ID,
             symbol: ALIAS_REGULAR_OLD,
             unsubscribe: false,
@@ -51,6 +53,7 @@ const MESSAGES = {
     } as Partial<OpenfeedGatewayMessage> as OpenfeedGatewayMessage,
     SUBSCRIPTION_RESPONSE_MESSAGE_NEW: {
         subscriptionResponse: {
+            correlationId: Long.fromNumber(3),
             marketId: MARKET_ID,
             symbol: ALIAS_REGULAR_NEW,
             unsubscribe: false,
@@ -83,6 +86,7 @@ const MESSAGES = {
     } as Partial<OpenfeedGatewayMessage> as OpenfeedGatewayMessage,
     DUMMY_SUBSCRIPTION_UNAUTHORIZED: {
         subscriptionResponse: {
+            correlationId: Long.fromNumber(1),
             marketId: MARKET_ID,
             symbol: ALIAS_REGULAR_OLD,
             unsubscribe: false,
@@ -94,6 +98,7 @@ const MESSAGES = {
     } as Partial<OpenfeedGatewayMessage> as OpenfeedGatewayMessage,
     DUMMY_UNSUBSCRIPTION_UNSUBSCRIBED: {
         subscriptionResponse: {
+            correlationId: Long.fromNumber(1),
             marketId: MARKET_ID,
             symbol: ALIAS_REGULAR_OLD,
             unsubscribe: true,
@@ -105,6 +110,7 @@ const MESSAGES = {
     } as Partial<OpenfeedGatewayMessage> as OpenfeedGatewayMessage,
     DUMMY_UNSUBSCRIPTION_DEPTH: {
         subscriptionResponse: {
+            correlationId: Long.fromNumber(2),
             marketId: MARKET_ID,
             symbol: ALIAS_REGULAR_OLD,
             unsubscribe: true,
@@ -236,6 +242,23 @@ describe("OpenFeedListeners", () => {
         );
     });
 
+    it("should handle unsubscription messages correctly", () => {
+        // We require the subscriptionResponse and instrumentDefinition message to arrive first
+        listeners.onMessage(MESSAGES.SUBSCRIPTION_RESPONSE_DEPTH_OLD);
+        listeners.onMessage(MESSAGES.INSTRUMENT_DEFINITION_MESSAGE_OLD);
+        listeners.onMessage(MESSAGES.DUMMY_UNSUBSCRIPTION_DEPTH);
+        listeners.onMessage(MESSAGES.DUMMY_SNAPSHOT_MESSAGE);
+
+        expect(onMessageWithMetadataSpy).toHaveBeenCalledWith(MESSAGES.SUBSCRIPTION_RESPONSE_DEPTH_OLD, [ALIAS_REGULAR_OLD], undefined);
+        expect(onMessageWithMetadataSpy).toHaveBeenCalledWith(MESSAGES.INSTRUMENT_DEFINITION_MESSAGE_OLD, [ALIAS_REGULAR_OLD], undefined);
+        expect(onMessageWithMetadataSpy).toHaveBeenCalledWith(
+            MESSAGES.DUMMY_UNSUBSCRIPTION_DEPTH,
+            [ALIAS_REGULAR_OLD],
+            MESSAGES.INSTRUMENT_DEFINITION_MESSAGE_OLD.instrumentDefinition
+        );
+        expect(onMessageWithMetadataSpy).toHaveBeenCalledWith(MESSAGES.DUMMY_SNAPSHOT_MESSAGE, [], undefined);
+    });
+
     it("should handle alias_changed instrument action correctly", () => {
         // We expect the messages to arrive in the following order - subscriptionResponse, instrumentDefinition, marketSnapshot
         // after an alias change - we expect the subscriptionResponse and instrumentDefinition to be repeated so that we can reconnect everything