@jr200-labs/xstate-nats 0.9.13 → 0.9.14

package/src/machines/kv.ts

@@ -3,9 +3,20 @@ import { jetstream } from '@nats-io/jetstream'
 import { KvEntry, Kvm, KvOptions, KvStatus, KvWatchEntry } from '@nats-io/kv'
 import { assign, sendParent, setup } from 'xstate'
 import { KvSubscriptionKey, KvSubscriptionConfig, kvConsolidateState } from '../actions/kv'
+import {
+  byteLength,
+  createEmptyTrafficMetrics,
+  NatsTrafficMetrics,
+  NatsTrafficRecordEvent,
+  NatsTrafficResetEvent,
+  recordTrafficMetric,
+  resetTrafficAll,
+  resetTrafficDownload,
+  resetTrafficUpload,
+} from '../traffic'
 
 // internal events and events from nats connection
-type InternalEvents = { type: 'ERROR'; error: Error }
+type InternalEvents = { type: 'ERROR'; error: Error } | NatsTrafficRecordEvent
 
 // events which can be sent to the machine from the user
 export type ExternalEvents =
@@ -49,6 +60,7 @@ export type ExternalEvents =
   | { type: 'KV.SUBSCRIBE'; config: KvSubscriptionConfig }
   | { type: 'KV.UNSUBSCRIBE'; bucket: string; key: string }
   | { type: 'KV.UNSUBSCRIBE_ALL' }
+  | NatsTrafficResetEvent
 
 export type Events = InternalEvents | ExternalEvents
 
@@ -60,6 +72,7 @@ export interface Context {
   subscriptions: Map<string, QueuedIterator<KvWatchEntry>>
   subscriptionConfigs: Map<string, KvSubscriptionConfig>
   syncRequired: number
+  trafficMetrics: NatsTrafficMetrics
   error?: Error
 }
 export const kvManagerLogic = setup({
@@ -86,19 +99,53 @@ export const kvManagerLogic = setup({
     subscriptions: new Map<string, QueuedIterator<KvWatchEntry>>(),
     subscriptionConfigs: new Map<string, KvSubscriptionConfig>(),
     syncRequired: 0,
+    trafficMetrics: createEmptyTrafficMetrics(),
   },
   on: {
+    'METRICS.RECORD': {
+      actions: assign({
+        trafficMetrics: ({ context, event }) => recordTrafficMetric(context.trafficMetrics, event),
+      }),
+    },
+    'METRICS.RESET_UPLOAD': {
+      actions: assign({
+        trafficMetrics: ({ context, event }) =>
+          resetTrafficUpload(context.trafficMetrics, event.source),
+      }),
+    },
+    'METRICS.RESET_DOWNLOAD': {
+      actions: assign({
+        trafficMetrics: ({ context, event }) =>
+          resetTrafficDownload(context.trafficMetrics, event.source),
+      }),
+    },
+    'METRICS.RESET_ALL': {
+      actions: assign({
+        trafficMetrics: ({ context, event }) =>
+          resetTrafficAll(context.trafficMetrics, event.source),
+      }),
+    },
     'KV.SUBSCRIBE': {
       actions: [
-        assign({
-          subscriptionConfigs: ({ context, event }) => {
-            const { config } = event
-            const newConfigs = new Map(context.subscriptionConfigs)
-            const newKvKey = KvSubscriptionKey.key(config.bucket, config.key)
-            newConfigs.set(newKvKey, config)
-            return newConfigs
-          },
-          syncRequired: ({ context }) => context.syncRequired + 1,
+        assign(({ context, event, self }) => {
+          if (event.type !== 'KV.SUBSCRIBE') return {}
+
+          const config: KvSubscriptionConfig = {
+            ...event.config,
+            onDownloadBytes: (bytes) =>
+              self.send({
+                type: 'METRICS.RECORD',
+                source: 'nats.kv',
+                downloadBytes: bytes,
+              }),
+          }
+          const newConfigs = new Map(context.subscriptionConfigs)
+          const newKvKey = KvSubscriptionKey.key(config.bucket, config.key)
+          newConfigs.set(newKvKey, config)
+          return {
+            subscriptionConfigs: newConfigs,
+            syncRequired: context.syncRequired + 1,
+          }
         }),
       ],
     },
@@ -252,7 +299,7 @@ export const kvManagerLogic = setup({
           },
         },
         'KV.PUT': {
-          actions: async ({ context, event }) => {
+          actions: async ({ context, event, self }) => {
             try {
               const kv = await context.cachedKvm?.open(event.bucket)
               if (!kv) {
@@ -260,6 +307,11 @@ export const kvManagerLogic = setup({
                 return
               }
 
+              self.send({
+                type: 'METRICS.RECORD',
+                source: 'nats.kv',
+                uploadBytes: byteLength(event.value),
+              })
               await kv.put(event.key, event.value)
               event.onResult({ ok: true })
             } catch (error) {

package/src/machines/root.ts

@@ -6,6 +6,7 @@ import { connectToNats, disconnectNats } from '../actions/connection'
 import { type AuthConfig } from '../actions/types'
 import { InternalStatusEvents as NatsStatusEvents } from '../actions/connection'
 import { withSpan } from '../telemetry'
+import { NatsTrafficResetEvent } from '../traffic'
 
 export interface NatsDiagnosticsConfig {
   lifecycle?: boolean
@@ -44,6 +45,7 @@ export type ExternalEvents =
   | { type: 'CONNECT' }
   | { type: 'DISCONNECT' }
   | { type: 'RESET' }
+  | NatsTrafficResetEvent
   | SubjectExternalEvents
   | KvExternalEvents
 
@@ -195,6 +197,15 @@ export const natsMachine = setup({
         },
       ],
     },
+    'METRICS.RESET_UPLOAD': {
+      actions: [sendTo('subject', ({ event }) => event), sendTo('kv', ({ event }) => event)],
+    },
+    'METRICS.RESET_DOWNLOAD': {
+      actions: [sendTo('subject', ({ event }) => event), sendTo('kv', ({ event }) => event)],
+    },
+    'METRICS.RESET_ALL': {
+      actions: [sendTo('subject', ({ event }) => event), sendTo('kv', ({ event }) => event)],
+    },
   },
   states: {
     not_configured: {

package/src/machines/subject.test.ts

@@ -18,7 +18,9 @@ function createMockConnection() {
         next: () => new Promise(() => {}),
       }),
     }),
-    request: vi.fn().mockResolvedValue({ json: () => ({}), string: () => '{}' }),
+    request: vi
+      .fn()
+      .mockResolvedValue({ data: new Uint8Array([1, 2, 3]), json: () => ({}), string: () => '{}' }),
     publish: vi.fn(),
   } as any
 }
@@ -58,6 +60,9 @@ function createParentMachine() {
           assign({ childState: 'disconnected' }),
         ],
       },
+      'METRICS.*': {
+        actions: sendTo('subject', ({ event }: any) => event),
+      },
     },
     states: {
       active: {},
@@ -241,6 +246,26 @@ describe('subjectManagerLogic', () => {
     parentActor.stop()
   })
 
+  it('should count publish upload bytes', async () => {
+    const parentActor = createActor(createParentMachine())
+    parentActor.start()
+
+    parentActor.send({ type: 'SUBJECT.CONNECT', connection: createMockConnection() })
+    parentActor.send({
+      type: 'SUBJECT.PUBLISH',
+      subject: 'test.pub',
+      payload: { msg: 'hello' },
+    })
+
+    await vi.waitFor(() => {
+      const metrics = parentActor.getSnapshot().children.subject?.getSnapshot()
+        ?.context.trafficMetrics
+      expect(metrics?.uploadBytes).toBe(15)
+      expect(metrics?.bySource['nats.publish']).toEqual({ uploadBytes: 15, downloadBytes: 0 })
+    })
+    parentActor.stop()
+  })
+
   it('should handle SUBJECT.REQUEST when connected', () => {
     const parentActor = createActor(createParentMachine())
     parentActor.start()
@@ -264,6 +289,66 @@ describe('subjectManagerLogic', () => {
     parentActor.stop()
   })
 
+  it('should count request upload and reply download bytes', async () => {
+    const parentActor = createActor(createParentMachine())
+    parentActor.start()
+
+    parentActor.send({ type: 'SUBJECT.CONNECT', connection: createMockConnection() })
+    parentActor.send({
+      type: 'SUBJECT.REQUEST',
+      subject: 'test.req',
+      payload: { data: 1 },
+      callback: vi.fn(),
+    })
+
+    await vi.waitFor(() => {
+      const metrics = parentActor.getSnapshot().children.subject?.getSnapshot()
+        ?.context.trafficMetrics
+      expect(metrics?.uploadBytes).toBe(10)
+      expect(metrics?.downloadBytes).toBe(3)
+      expect(metrics?.bySource['nats.request']).toEqual({ uploadBytes: 10, downloadBytes: 3 })
+    })
+    parentActor.stop()
+  })
+
+  it('should reset upload and download metrics without disconnecting', async () => {
+    const parentActor = createActor(createParentMachine())
+    parentActor.start()
+
+    const connection = createMockConnection()
+    parentActor.send({ type: 'SUBJECT.CONNECT', connection })
+    parentActor.send({
+      type: 'SUBJECT.REQUEST',
+      subject: 'test.req',
+      payload: { data: 1 },
+      callback: vi.fn(),
+    })
+
+    await vi.waitFor(() => {
+      const metrics = parentActor.getSnapshot().children.subject?.getSnapshot()
+        ?.context.trafficMetrics
+      expect(metrics?.downloadBytes).toBe(3)
+    })
+
+    parentActor.send({ type: 'METRICS.RESET_UPLOAD' })
+    await vi.waitFor(() => {
+      const snap = parentActor.getSnapshot().children.subject?.getSnapshot()
+      expect(snap?.value).toBe('subject_connected')
+      expect(snap?.context.cachedConnection).toBe(connection)
+      expect(snap?.context.trafficMetrics.uploadBytes).toBe(0)
+      expect(snap?.context.trafficMetrics.downloadBytes).toBe(3)
+    })
+
+    parentActor.send({ type: 'METRICS.RESET_DOWNLOAD' })
+    await vi.waitFor(() => {
+      const metrics = parentActor.getSnapshot().children.subject?.getSnapshot()
+        ?.context.trafficMetrics
+      expect(metrics?.uploadBytes).toBe(0)
+      expect(metrics?.downloadBytes).toBe(0)
+    })
+    parentActor.stop()
+  })
+
   it('should accept SUBJECT.SUBSCRIBE in idle state (global handler)', () => {
     const parentActor = createActor(createParentMachine())
     parentActor.start()

package/src/machines/subject.ts

@@ -7,9 +7,20 @@ import {
   subjectRequest,
   subjectPublish,
 } from '../actions/subject'
+import {
+  byteLength,
+  createEmptyTrafficMetrics,
+  NatsTrafficMetrics,
+  NatsTrafficRecordEvent,
+  NatsTrafficResetEvent,
+  recordTrafficMetric,
+  resetTrafficAll,
+  resetTrafficDownload,
+  resetTrafficUpload,
+} from '../traffic'
 
 // internal events and events from nats connection
-type InternalEvents = { type: 'ERROR'; error: Error }
+type InternalEvents = { type: 'ERROR'; error: Error } | NatsTrafficRecordEvent
 
 // events which can be sent to the machine from the user
 export type ExternalEvents =
@@ -34,6 +45,7 @@ export type ExternalEvents =
   | { type: 'SUBJECT.SUBSCRIBE'; config: SubjectSubscriptionConfig }
   | { type: 'SUBJECT.UNSUBSCRIBE'; subject: string }
   | { type: 'SUBJECT.UNSUBSCRIBE_ALL' }
+  | NatsTrafficResetEvent
 
 export type Events = InternalEvents | ExternalEvents
 
@@ -43,6 +55,7 @@ export interface Context {
   subscriptions: Map<string, Subscription>
   subscriptionConfigs: Map<string, SubjectSubscriptionConfig>
   syncRequired: number
+  trafficMetrics: NatsTrafficMetrics
   error?: Error
 }
 
@@ -65,19 +78,53 @@ export const subjectManagerLogic = setup({
     subscriptionConfigs: new Map<string, SubjectSubscriptionConfig>(),
     cachedConnection: null,
     syncRequired: 0,
+    trafficMetrics: createEmptyTrafficMetrics(),
     error: undefined,
   },
   on: {
+    'METRICS.RECORD': {
+      actions: assign({
+        trafficMetrics: ({ context, event }) => recordTrafficMetric(context.trafficMetrics, event),
+      }),
+    },
+    'METRICS.RESET_UPLOAD': {
+      actions: assign({
+        trafficMetrics: ({ context, event }) =>
+          resetTrafficUpload(context.trafficMetrics, event.source),
+      }),
+    },
+    'METRICS.RESET_DOWNLOAD': {
+      actions: assign({
+        trafficMetrics: ({ context, event }) =>
+          resetTrafficDownload(context.trafficMetrics, event.source),
+      }),
+    },
+    'METRICS.RESET_ALL': {
+      actions: assign({
+        trafficMetrics: ({ context, event }) =>
+          resetTrafficAll(context.trafficMetrics, event.source),
+      }),
+    },
     'SUBJECT.SUBSCRIBE': {
       actions: [
-        assign({
-          subscriptionConfigs: ({ context, event }) => {
-            const { config } = event
-            const newConfigs = new Map(context.subscriptionConfigs)
-            newConfigs.set(config.subject, config)
-            return newConfigs
-          },
-          syncRequired: ({ context }) => context.syncRequired + 1,
+        assign(({ context, event, self }) => {
+          if (event.type !== 'SUBJECT.SUBSCRIBE') return {}
+
+          const config: SubjectSubscriptionConfig = {
+            ...event.config,
+            onDownloadBytes: (bytes) =>
+              self.send({
+                type: 'METRICS.RECORD',
+                source: 'nats.subscribe',
+                downloadBytes: bytes,
+              }),
+          }
+          const newConfigs = new Map(context.subscriptionConfigs)
+          newConfigs.set(config.subject, config)
+          return {
+            subscriptionConfigs: newConfigs,
+            syncRequired: context.syncRequired + 1,
+          }
         }),
       ],
     },
@@ -150,7 +197,12 @@ export const subjectManagerLogic = setup({
           target: 'subject_disconnecting',
         },
         'SUBJECT.REQUEST': {
-          actions: assign(({ event, context }) => {
+          actions: assign(({ event, context, self }) => {
+            self.send({
+              type: 'METRICS.RECORD',
+              source: 'nats.request',
+              uploadBytes: byteLength(event.payload),
+            })
             subjectRequest({
               input: {
                 connection: context.cachedConnection!,
@@ -159,6 +211,12 @@ export const subjectManagerLogic = setup({
                 opts: event.opts,
                 callback: event.callback,
                 onRequestResult: event.onRequestResult,
+                onDownloadBytes: (bytes) =>
+                  self.send({
+                    type: 'METRICS.RECORD',
+                    source: 'nats.request',
+                    downloadBytes: bytes,
+                  }),
               },
             })
             return {}
@@ -166,7 +224,12 @@ export const subjectManagerLogic = setup({
         },
         'SUBJECT.PUBLISH': {
           actions: [
-            ({ context, event }) => {
+            ({ context, event, self }) => {
+              self.send({
+                type: 'METRICS.RECORD',
+                source: 'nats.publish',
+                uploadBytes: byteLength(event.payload),
+              })
               subjectPublish({
                 input: {
                   connection: context.cachedConnection!,

package/src/traffic.test.ts

@@ -0,0 +1,62 @@
+import { describe, expect, it } from 'vitest'
+import {
+  byteLength,
+  createEmptyTrafficMetrics,
+  recordTrafficMetric,
+  resetTrafficAll,
+  resetTrafficDownload,
+  resetTrafficUpload,
+} from './traffic'
+
+describe('traffic metrics', () => {
+  it('counts UTF-8 payload bytes', () => {
+    expect(byteLength('hello')).toBe(5)
+    expect(byteLength('£')).toBe(2)
+    expect(byteLength({ msg: 'hello' })).toBe(15)
+    expect(byteLength(new Uint8Array([1, 2, 3]))).toBe(3)
+  })
+
+  it('records upload and download bytes by source', () => {
+    const metrics = recordTrafficMetric(createEmptyTrafficMetrics(), {
+      type: 'METRICS.RECORD',
+      source: 'nats.request',
+      uploadBytes: 10,
+      downloadBytes: 25,
+    })
+
+    expect(metrics.uploadBytes).toBe(10)
+    expect(metrics.downloadBytes).toBe(25)
+    expect(metrics.bySource['nats.request']).toEqual({ uploadBytes: 10, downloadBytes: 25 })
+  })
+
+  it('resets upload and download independently', () => {
+    const metrics = recordTrafficMetric(
+      recordTrafficMetric(createEmptyTrafficMetrics(), {
+        type: 'METRICS.RECORD',
+        source: 'nats.request',
+        uploadBytes: 10,
+        downloadBytes: 25,
+      }),
+      {
+        type: 'METRICS.RECORD',
+        source: 'nats.subscribe',
+        downloadBytes: 7,
+      },
+    )
+
+    const uploadReset = resetTrafficUpload(metrics)
+    expect(uploadReset.uploadBytes).toBe(0)
+    expect(uploadReset.downloadBytes).toBe(32)
+    expect(uploadReset.bySource['nats.request']).toEqual({ uploadBytes: 0, downloadBytes: 25 })
+
+    const requestDownloadReset = resetTrafficDownload(metrics, 'nats.request')
+    expect(requestDownloadReset.uploadBytes).toBe(10)
+    expect(requestDownloadReset.downloadBytes).toBe(7)
+    expect(requestDownloadReset.bySource['nats.request']).toEqual({
+      uploadBytes: 10,
+      downloadBytes: 0,
+    })
+
+    expect(resetTrafficAll(metrics)).toEqual(createEmptyTrafficMetrics())
+  })
+})

package/src/traffic.ts

@@ -0,0 +1,150 @@
+export type NatsTrafficMetricSource = 'nats.publish' | 'nats.request' | 'nats.subscribe' | 'nats.kv'
+
+export type NatsTrafficDirection = 'upload' | 'download'
+
+export interface NatsTrafficMetricBucket {
+  uploadBytes: number
+  downloadBytes: number
+}
+
+export interface NatsTrafficMetrics extends NatsTrafficMetricBucket {
+  bySource: Partial<Record<NatsTrafficMetricSource, NatsTrafficMetricBucket>>
+}
+
+export type NatsTrafficRecordEvent = {
+  type: 'METRICS.RECORD'
+  source: NatsTrafficMetricSource
+  uploadBytes?: number
+  downloadBytes?: number
+}
+
+export type NatsTrafficResetEvent =
+  | { type: 'METRICS.RESET_UPLOAD'; source?: NatsTrafficMetricSource }
+  | { type: 'METRICS.RESET_DOWNLOAD'; source?: NatsTrafficMetricSource }
+  | { type: 'METRICS.RESET_ALL'; source?: NatsTrafficMetricSource }
+
+export type NatsTrafficEvent = NatsTrafficRecordEvent | NatsTrafficResetEvent
+
+const textEncoder = new TextEncoder()
+
+export function createEmptyTrafficMetrics(): NatsTrafficMetrics {
+  return {
+    uploadBytes: 0,
+    downloadBytes: 0,
+    bySource: {},
+  }
+}
+
+export function byteLength(value: unknown): number {
+  if (value == null) return 0
+  if (value instanceof Uint8Array) return value.byteLength
+  if (value instanceof ArrayBuffer) return value.byteLength
+  if (ArrayBuffer.isView(value)) return value.byteLength
+  if (typeof value === 'string') return textEncoder.encode(value).byteLength
+
+  try {
+    return textEncoder.encode(JSON.stringify(value)).byteLength
+  } catch {
+    return 0
+  }
+}
+
+export function recordTrafficMetric(
+  metrics: NatsTrafficMetrics,
+  event: NatsTrafficRecordEvent,
+): NatsTrafficMetrics {
+  const uploadBytes = normaliseByteCount(event.uploadBytes)
+  const downloadBytes = normaliseByteCount(event.downloadBytes)
+
+  if (uploadBytes === 0 && downloadBytes === 0) return metrics
+
+  const currentSource = metrics.bySource[event.source] ?? { uploadBytes: 0, downloadBytes: 0 }
+  return {
+    uploadBytes: metrics.uploadBytes + uploadBytes,
+    downloadBytes: metrics.downloadBytes + downloadBytes,
+    bySource: {
+      ...metrics.bySource,
+      [event.source]: {
+        uploadBytes: currentSource.uploadBytes + uploadBytes,
+        downloadBytes: currentSource.downloadBytes + downloadBytes,
+      },
+    },
+  }
+}
+
+export function resetTrafficUpload(
+  metrics: NatsTrafficMetrics,
+  source?: NatsTrafficMetricSource,
+): NatsTrafficMetrics {
+  return resetTrafficDirection(metrics, 'upload', source)
+}
+
+export function resetTrafficDownload(
+  metrics: NatsTrafficMetrics,
+  source?: NatsTrafficMetricSource,
+): NatsTrafficMetrics {
+  return resetTrafficDirection(metrics, 'download', source)
+}
+
+export function resetTrafficAll(
+  metrics: NatsTrafficMetrics,
+  source?: NatsTrafficMetricSource,
+): NatsTrafficMetrics {
+  if (!source) return createEmptyTrafficMetrics()
+
+  const sourceMetrics = metrics.bySource[source] ?? { uploadBytes: 0, downloadBytes: 0 }
+  const bySource = { ...metrics.bySource }
+  delete bySource[source]
+
+  return {
+    uploadBytes: metrics.uploadBytes - sourceMetrics.uploadBytes,
+    downloadBytes: metrics.downloadBytes - sourceMetrics.downloadBytes,
+    bySource,
+  }
+}
+
+function resetTrafficDirection(
+  metrics: NatsTrafficMetrics,
+  direction: NatsTrafficDirection,
+  source?: NatsTrafficMetricSource,
+): NatsTrafficMetrics {
+  const key = direction === 'upload' ? 'uploadBytes' : 'downloadBytes'
+
+  if (!source) {
+    const bySource = Object.fromEntries(
+      Object.entries(metrics.bySource).map(([sourceName, sourceMetrics]) => [
+        sourceName,
+        {
+          ...sourceMetrics,
+          [key]: 0,
+        },
+      ]),
+    ) as NatsTrafficMetrics['bySource']
+
+    return {
+      ...metrics,
+      [key]: 0,
+      bySource,
+    }
+  }
+
+  const sourceMetrics = metrics.bySource[source]
+  if (!sourceMetrics) return metrics
+
+  return {
+    ...metrics,
+    [key]: metrics[key] - sourceMetrics[key],
+    bySource: {
+      ...metrics.bySource,
+      [source]: {
+        ...sourceMetrics,
+        [key]: 0,
+      },
+    },
+  }
+}
+
+function normaliseByteCount(value: number | undefined): number {
+  if (!Number.isFinite(value) || value === undefined || value <= 0) return 0
+  return Math.floor(value)
+}