feat: add prom-clilent, metrics handler & server

packages/server/metrics.ts

@@ -0,0 +1,93 @@
+import {collectDefaultMetrics, Counter, Gauge, Histogram, register} from 'prom-client'
+import {HttpResponse} from 'uWebSockets.js'
+
+// Default Node.js metrics collection
+collectDefaultMetrics()
+
+// Caching for metrics
+let cachedMetrics = ''
+let lastMetricsUpdate = 0
+const METRICS_CACHE_DURATION = 5000 // 5 seconds
+
+async function getMetrics(): Promise<string> {
+  const now = Date.now()
+  if (now - lastMetricsUpdate > METRICS_CACHE_DURATION) {
+    cachedMetrics = await register.metrics()
+    lastMetricsUpdate = now
+  }
+  return cachedMetrics
+}
+
+export async function handleMetricsRequest(res: HttpResponse): Promise<void> {
+  let aborted = false
+  res.onAborted(() => (aborted = true))
+
+  try {
+    const metrics = await getMetrics()
+    if (aborted) return
+
+    res.cork(() => {
+      res.writeHeader('Content-Type', register.contentType)
+      res.end(metrics)
+    })
+  } catch (error) {
+    console.error('Error retrieving metrics:', error)
+    if (!aborted) {
+      res.cork(() => {
+        res.writeHeader('Content-Type', 'text/plain')
+        res.end('Error retrieving metrics')
+      })
+    }
+  }
+}
+
+// Websocket metrics
+export const websocketConnections = new Gauge({
+  name: 'websocket_active_connections',
+  help: 'Number of active WebSocket connections'
+})
+
+export const websocketMessagesSent = new Counter({
+  name: 'websocket_messages_sent_total',
+  help: 'Total number of WebSocket messages sent'
+})
+
+export const websocketMessagesReceived = new Counter({
+  name: 'websocket_messages_received_total',
+  help: 'Total number of WebSocket messages received'
+})
+
+export const websocketConnectionDuration = new Histogram({
+  name: 'websocket_connection_duration_seconds',
+  help: 'Duration of WebSocket connections',
+  buckets: [1, 5, 10, 30, 60, 120, 300] // Buckets for duration
+})
+
+const connectionStartTimes = new Map()
+
+export function incrementWebSocketConnections() {
+  websocketConnections.inc()
+}
+
+export function decrementWebSocketConnections(ws: any) {
+  websocketConnections.dec()
+
+  // Track connection duration
+  const startTime = connectionStartTimes.get(ws)
+  if (startTime) {
+    websocketConnectionDuration.observe((Date.now() - startTime) / 1000)
+    connectionStartTimes.delete(ws)
+  }
+}
+
+export function trackWebSocketOpen(ws: any) {
+  connectionStartTimes.set(ws, Date.now())
+}
+
+export function trackWsMessageSent() {
+  websocketMessagesSent.inc()
+}
+
+export function trackWsMessageReceived() {
+  websocketMessagesReceived.inc()
+}

packages/server/package.json

@@ -128,6 +128,7 @@
     "oy-vey": "^0.12.1",
     "parabol-client": "8.25.4",
     "pg": "^8.5.1",
+    "prom-client": "^15.1.3",
     "react": "^17.0.2",
     "react-dom": "^17.0.2",
     "relay-runtime": "^14.1.0",

packages/server/server.ts

@@ -1,5 +1,5 @@
 import tracer from 'dd-trace'
-import uws, {SHARED_COMPRESSOR} from 'uWebSockets.js'
+import uws, {SHARED_COMPRESSOR, WebSocket} from 'uWebSockets.js'
 import sleep from '../client/utils/sleep'
 import ICSHandler from './ICSHandler'
 import PWAHandler from './PWAHandler'
@@ -12,6 +12,13 @@ import './initSentry'
 import mattermostWebhookHandler from './integrations/mattermost/mattermostWebhookHandler'
 import jiraImagesHandler from './jiraImagesHandler'
 import listenHandler from './listenHandler'
+import {
+  decrementWebSocketConnections,
+  handleMetricsRequest,
+  incrementWebSocketConnections,
+  trackWebSocketOpen,
+  trackWsMessageReceived
+} from './metrics'
 import './monkeyPatchFetch'
 import selfHostedHandler from './selfHostedHandler'
 import handleClose from './socketHandlers/handleClose'
@@ -53,6 +60,9 @@ process.on('SIGTERM', async (signal) => {
 })
 
 const PORT = Number(__PRODUCTION__ ? process.env.PORT : process.env.SOCKET_PORT)
+const METRICS_PORT = Number(process.env.METRICS_PORT || 9100) // Default to 9100 if not specified
+
+// Main App
 uws
   .App()
   .get('/favicon.ico', PWAHandler)
@@ -76,10 +86,31 @@ uws
     idleTimeout: 0,
     maxPayloadLength: 5 * 2 ** 20,
     upgrade: handleUpgrade,
-    open: handleOpen,
-    message: handleMessage,
-    // today, we don't send folks enough data to worry about backpressure
-    close: handleClose
+    open: (ws: WebSocket<any>) => {
+      incrementWebSocketConnections()
+      trackWebSocketOpen(ws)
+      if (handleOpen) handleOpen(ws)
+    },
+    message: (ws: WebSocket<any>, message: ArrayBuffer, isBinary: boolean) => {
+      trackWsMessageReceived()
+      if (handleMessage) handleMessage(ws, message, isBinary)
+    },
+    close: (ws: WebSocket<any>) => {
+      decrementWebSocketConnections(ws)
+      handleClose(ws)
+    }
   })
   .any('/*', createSSR)
   .listen(PORT, listenHandler)
+
+// Metrics App
+uws
+  .App()
+  .get('/metrics', handleMetricsRequest)
+  .listen(METRICS_PORT, (listenSocket) => {
+    if (listenSocket) {
+      Logger.log(`Metrics server listening on port ${METRICS_PORT}`)
+    } else {
+      Logger.error('Failed to start metrics server')
+    }
+  })