Comments

Author: mxiao-cll

src/transports/metrics.ts

@@ -29,26 +29,33 @@ export const messageSubsLabels = <T extends TransportGenerics>(
   }
 }
 
+export const recordWsMessageSentMetrics = <T extends TransportGenerics>(
+  context: EndpointContext<T>,
+  subscribes: TypeFromDefinition<T['Parameters']>[],
+  unsubscribes: TypeFromDefinition<T['Parameters']>[],
+) => {
+  for (const params of [...subscribes, ...unsubscribes]) {
+    const baseLabels = messageSubsLabels(context, params)
+
+    // Record total number of ws messages sent
+    metrics
+      .get('wsMessageTotal')
+      .labels({ ...baseLabels, direction: 'sent' })
+      .inc()
+  }
+}
+
 // Record WS message and subscription metrics
 // Recalculate cacheKey and feedId for metrics
 // since avoiding storing extra info in expiring sorted set
-export const recordWsMessageMetrics = <T extends TransportGenerics>(
+export const recordWsMessageSubMetrics = <T extends TransportGenerics>(
   context: EndpointContext<T>,
   subscribes: TypeFromDefinition<T['Parameters']>[],
   unsubscribes: TypeFromDefinition<T['Parameters']>[],
-  wsMessageTotal: boolean = true,
 ): void => {
   const recordMetrics = (params: TypeFromDefinition<T['Parameters']>, type: 'sub' | 'unsub') => {
     const baseLabels = messageSubsLabels(context, params)
 
-    // Record total number of ws messages sent
-    if (wsMessageTotal) {
-      metrics
-        .get('wsMessageTotal')
-        .labels({ ...baseLabels, direction: 'sent' })
-        .inc()
-    }
-
     // Record total number of subscriptions made
     if (type === 'sub') {
       metrics.get('wsSubscriptionTotal').labels(baseLabels).inc()

src/transports/websocket.ts

@@ -6,7 +6,11 @@ import { PartialSuccessfulResponse, ProviderResult, TimestampedProviderResult }
 import { TypeFromDefinition } from '../validation/input-params'
 import { TransportGenerics } from './'
 import { StreamingTransport, SubscriptionDeltas } from './abstract/streaming'
-import { connectionErrorLabels, recordWsMessageMetrics } from './metrics'
+import {
+  connectionErrorLabels,
+  recordWsMessageSentMetrics,
+  recordWsMessageSubMetrics,
+} from './metrics'
 
 // Aliasing type for use at adapter level
 export { WebSocket, RawData as WebSocketRawData }
@@ -453,8 +457,8 @@ export class WebSocketTransport<
       connectionClosed = true
 
       if (subscriptions.desired.length) {
-        // Clear wsSubscriptionActive metrics for all active subscriptions
-        recordWsMessageMetrics(context, [], subscriptions.desired, false)
+        // Clear subscription metrics for all active subscriptions
+        recordWsMessageSubMetrics(context, [], subscriptions.desired)
         censorLogs(() =>
           logger.trace(
             `Connection will be reopened and will subscribe to new and resubscribe to existing: ${JSON.stringify(
@@ -505,15 +509,16 @@ export class WebSocketTransport<
             ? subscriptions.stale.map((sub) => unsubscribeMessage(sub, context))
             : subscriptions.stale,
         )
+        recordWsMessageSentMetrics(context, subscriptions.new, subscriptions.stale)
       } else {
         logger.trace(
           "This ws transport has no builders configured, so we're not sending any messages",
         )
       }
     }
 
-    // Record WS message and subscription metrics
-    recordWsMessageMetrics(context, subscriptions.new, subscriptions.stale)
+    // Record WS subscription metrics
+    recordWsMessageSubMetrics(context, subscriptions.new, subscriptions.stale)
 
     // The background execute loop no longer sleeps between executions, so we have to do it here
     logger.trace(

test/transports/websocket.test.ts

@@ -341,6 +341,11 @@ test.serial('reconnects if connection becomes unresponsive', async (t) => {
   let metrics = await testAdapter.getMetrics()
   metrics.assert(t, { name: 'ws_subscription_active', labels, expectedValue: 1 })
   metrics.assert(t, { name: 'ws_subscription_total', labels, expectedValue: 1 })
+  metrics.assert(t, {
+    name: 'ws_message_total',
+    labels: { ...labels, direction: 'sent' },
+    expectedValue: 1,
+  })
 
   // Advance to next cycle where connection is unhealthy and reconnect
   await runAllUntilTime(t.context.clock, BACKGROUND_EXECUTE_MS_WS + 100)