Add WS Bad Closure Failover Mechanism

cawthorne · cawthorne · commit 7536f4268242 · 2026-03-05T13:46:39.000Z
diff --git a/src/transports/websocket.ts b/src/transports/websocket.ts
@@ -289,20 +289,27 @@ export class WebSocketTransport<
 
       // Called when the WS connection closes for any reason
       close: (event: WebSocket.CloseEvent) => {
-        // If the connection closed with 1000, it's a usual closure
-        const level = event.code === 1000 ? 'debug' : 'info'
-        logger[level](
-          `Closed websocket connection. Code: ${event.code} ; reason: ${event.reason?.toString()}`,
-        )
+        const filteredUrl = this.currentUrl.split('?')[0]
+        const isAbnormal = event.code !== 1000
 
-        // Record active ws connections by decrementing count on close
-        // Using URL in label since connection_key is removed from v3
-        metrics.get('wsConnectionActive').dec()
+        if (isAbnormal) {
+          this.streamHandlerInvocationsWithNoConnection += 1
+          logger.warn(
+            `WebSocket closed abnormally (code: ${event.code}, reason: ${event.reason?.toString() || 'none'}). ` +
+              `Failover counter incremented to ${this.streamHandlerInvocationsWithNoConnection}. ` +
+              `URL: ${filteredUrl}`,
+          )
+          metrics
+            .get('wsConnectionFailoverCount')
+            .labels({ transport_name: this.name, url: filteredUrl })
+            .set(this.streamHandlerInvocationsWithNoConnection)
+        } else {
+          logger.debug(
+            `WebSocket closed normally (code: ${event.code}). URL: ${filteredUrl}`,
+          )
+        }
 
-        // Also, register that the connection was closed and the reason why.
-        // We need to filter out query params from the URL to avoid having
-        // the cardinality of the metric go out of control.
-        const filteredUrl = this.currentUrl.split('?')[0]
+        metrics.get('wsConnectionActive').dec()
         metrics.get('wsConnectionClosures').inc({
           code: event.code,
           url: filteredUrl,
diff --git a/test/transports/websocket.test.ts b/test/transports/websocket.test.ts
@@ -1312,3 +1312,227 @@ test.serial('does not heartbeat when handler throws an error', async (t) => {
   mockWsServer.close()
   await t.context.clock.runToLastAsync()
 })
+
+test.serial(
+  'increments failover counter on abnormal closure and passes it to url function',
+  async (t) => {
+    const base = 'ETH'
+    const quote = 'DOGE'
+    process.env['METRICS_ENABLED'] = 'true'
+    eaMetrics.clear()
+
+    const counterValues: number[] = []
+
+    mockWebSocketProvider(WebSocketClassProvider)
+    const mockWsServer = new Server(ENDPOINT_URL, { mock: false })
+    mockWsServer.on('connection', (socket) => {
+      socket.on('message', () => {
+        socket.close({ code: 4000, reason: 'Simulated abnormal closure', wasClean: false })
+      })
+    })
+
+    const transport = new WebSocketTransport<WebSocketTypes>({
+      url: (_context, _desiredSubs, params) => {
+        counterValues.push(params.streamHandlerInvocationsWithNoConnection)
+        return ENDPOINT_URL
+      },
+      handlers: {
+        message(message) {
+          if (!message.pair) {
+            return []
+          }
+          const [curBase, curQuote] = message.pair.split('/')
+          return [
+            {
+              params: { base: curBase, quote: curQuote },
+              response: {
+                data: {
+                  result: message.value,
+                },
+                result: message.value,
+              },
+            },
+          ]
+        },
+      },
+      builders: {
+        subscribeMessage: (params) => ({
+          request: 'subscribe',
+          pair: `${params.base}/${params.quote}`,
+        }),
+        unsubscribeMessage: (params) => ({
+          request: 'unsubscribe',
+          pair: `${params.base}/${params.quote}`,
+        }),
+      },
+    })
+
+    const webSocketEndpoint = new AdapterEndpoint({
+      name: 'TEST',
+      transport: transport,
+      inputParameters,
+    })
+
+    const config = new AdapterConfig(
+      {},
+      {
+        envDefaultOverrides: {
+          BACKGROUND_EXECUTE_MS_WS,
+          WS_SUBSCRIPTION_UNRESPONSIVE_TTL: 180_000,
+        },
+      },
+    )
+
+    const adapter = new Adapter({
+      name: 'TEST',
+      defaultEndpoint: 'test',
+      config,
+      endpoints: [webSocketEndpoint],
+    })
+
+    const testAdapter = await TestAdapter.startWithMockedCache(adapter, t.context)
+
+    await testAdapter.request({ base, quote })
+
+    // Each cycle: connect -> subscribe -> server closes with 4000 -> counter increments
+    await runAllUntilTime(t.context.clock, BACKGROUND_EXECUTE_MS_WS * 5 + 100)
+
+    // The counter should be strictly increasing due to abnormal close incrementing it
+    t.true(counterValues.length >= 3, `Expected at least 3 url calls, got ${counterValues.length}`)
+
+    for (let i = 1; i < counterValues.length; i++) {
+      t.true(
+        counterValues[i] > counterValues[i - 1],
+        `Counter should increase: index ${i} (${counterValues[i]}) should be > index ${i - 1} (${counterValues[i - 1]})`,
+      )
+    }
+
+    process.env['METRICS_ENABLED'] = 'false'
+    await testAdapter.api.close()
+    mockWsServer.close()
+    await t.context.clock.runToLastAsync()
+  },
+)
+
+test.serial(
+  'cycles between primary and secondary URLs on abnormal closure',
+  async (t) => {
+    const base = 'ETH'
+    const quote = 'DOGE'
+
+    const PRIMARY_URL = 'wss://primary.test.com/ws'
+    const SECONDARY_URL = 'wss://secondary.test.com/ws'
+    const urlsConnected: string[] = []
+
+    mockWebSocketProvider(WebSocketClassProvider)
+
+    const mockPrimary = new Server(PRIMARY_URL, { mock: false })
+    mockPrimary.on('connection', (socket) => {
+      socket.on('message', () => {
+        socket.close({ code: 4000, reason: 'Primary abnormal close', wasClean: false })
+      })
+    })
+
+    const mockSecondary = new Server(SECONDARY_URL, { mock: false })
+    mockSecondary.on('connection', (socket) => {
+      socket.on('message', () => {
+        socket.close({ code: 4001, reason: 'Secondary abnormal close', wasClean: false })
+      })
+    })
+
+    // Mimics Tiingo's wsSelectUrl with a 1:1 primary:secondary ratio
+    const transport = new WebSocketTransport<WebSocketTypes>({
+      url: (_context, _desiredSubs, params) => {
+        const counter = params.streamHandlerInvocationsWithNoConnection
+        const zeroIndexed = counter - 1
+        const cyclePos = zeroIndexed % 2
+        const url = cyclePos < 1 ? PRIMARY_URL : SECONDARY_URL
+        urlsConnected.push(url)
+        return url
+      },
+      handlers: {
+        message(message) {
+          if (!message.pair) {
+            return []
+          }
+          const [curBase, curQuote] = message.pair.split('/')
+          return [
+            {
+              params: { base: curBase, quote: curQuote },
+              response: {
+                data: { result: message.value },
+                result: message.value,
+              },
+            },
+          ]
+        },
+      },
+      builders: {
+        subscribeMessage: (params) => ({
+          request: 'subscribe',
+          pair: `${params.base}/${params.quote}`,
+        }),
+        unsubscribeMessage: (params) => ({
+          request: 'unsubscribe',
+          pair: `${params.base}/${params.quote}`,
+        }),
+      },
+    })
+
+    const webSocketEndpoint = new AdapterEndpoint({
+      name: 'TEST',
+      transport: transport,
+      inputParameters,
+    })
+
+    const config = new AdapterConfig(
+      {},
+      {
+        envDefaultOverrides: {
+          BACKGROUND_EXECUTE_MS_WS,
+          WS_SUBSCRIPTION_UNRESPONSIVE_TTL: 180_000,
+        },
+      },
+    )
+
+    const adapter = new Adapter({
+      name: 'TEST',
+      defaultEndpoint: 'test',
+      config,
+      endpoints: [webSocketEndpoint],
+    })
+
+    const testAdapter = await TestAdapter.startWithMockedCache(adapter, t.context)
+
+    await testAdapter.request({ base, quote })
+
+    // Run through enough cycles to see URL cycling:
+    // counter=0 -> primary, counter=1 -> primary, counter=2 -> secondary,
+    // counter=3 -> primary, counter=4 -> secondary, ...
+    await runAllUntilTime(t.context.clock, BACKGROUND_EXECUTE_MS_WS * 6 + 100)
+
+    const primaryCount = urlsConnected.filter((u) => u === PRIMARY_URL).length
+    const secondaryCount = urlsConnected.filter((u) => u === SECONDARY_URL).length
+
+    t.true(primaryCount >= 2, `Expected at least 2 primary connections, got ${primaryCount}`)
+    t.true(secondaryCount >= 1, `Expected at least 1 secondary connection, got ${secondaryCount}`)
+
+    // After hitting secondary, it should cycle back to primary
+    const firstSecondaryIndex = urlsConnected.indexOf(SECONDARY_URL)
+    t.true(firstSecondaryIndex >= 0, 'Should have connected to secondary')
+    t.true(
+      firstSecondaryIndex < urlsConnected.length - 1,
+      'Secondary should not be the last connection (should have returned to primary)',
+    )
+    t.is(
+      urlsConnected[firstSecondaryIndex + 1],
+      PRIMARY_URL,
+      'After secondary, should reconnect to primary',
+    )
+
+    await testAdapter.api.close()
+    mockPrimary.close()
+    mockSecondary.close()
+    await t.context.clock.runToLastAsync()
+  },
+)
