chore: add some basic instrumentation to Node WebSocket server (apache#14354)

* chore: add some basic instrumentation

* Switch to statsd using hot-shots

* Cleanup a few leftover bits

Author: benjreinhart

superset-websocket/README.md

@@ -93,6 +93,22 @@ GLOBAL_ASYNC_QUERIES_JWT_SECRET
 
 More info on Superset configuration values for async queries: https://github.com/apache/superset/blob/master/CONTRIBUTING.md#async-chart-queries
 
+## StatsD monitoring
+
+The application is tracking a couple of metrics with `statsd` using the [hot-shots](https://www.npmjs.com/package/hot-shots) library, such as the number of connected clients and the number of failed attempts to send a message to a client.
+
+`statsd` can be configured with the `statsd` object in the configuration file. See the [hot-shots](https://www.npmjs.com/package/hot-shots) readme for more info. The default configuration is:
+
+```json
+{
+  "statsd": {
+    "host": "127.0.0.1",
+    "port": 8125,
+    "globalTags": []
+  }
+}
+```
+
 ## Running
 
 Running locally via dev server:

superset-websocket/config.example.json

@@ -3,6 +3,11 @@
   "logLevel": "info",
   "logToFile": false,
   "logFilename": "app.log",
+  "statsd": {
+    "host": "127.0.0.1",
+    "port": 8125,
+    "globalTags": []
+  },
   "redis": {
     "port": 6379,
     "host": "127.0.0.1",

superset-websocket/config.test.json

@@ -6,6 +6,11 @@
     "db": 10,
     "ssl": false
   },
+  "statsd": {
+    "host": "127.0.0.1",
+    "port": 8125,
+    "globalTags": []
+  },
   "redisStreamPrefix": "test-async-events-",
   "jwtSecret": "test123-test123-test123-test123-test123-test123-test123",
   "jwtCookieName": "test-async-token"

superset-websocket/package-lock.json

@@ -16,6 +16,7 @@
   "license": "Apache-2.0",
   "dependencies": {
     "cookie": "^0.4.1",
+    "hot-shots": "^8.3.1",
     "ioredis": "^4.16.1",
     "jsonwebtoken": "^8.5.1",
     "uuid": "^8.3.2",

superset-websocket/package.json

@@ -56,11 +56,19 @@ const streamReturnValue: server.StreamResult[] = [
 ];
 
 import * as server from '../src/index';
+import { statsd } from '../src/index';
 
 describe('server', () => {
+  let statsdIncrementMock: jest.SpyInstance;
+
   beforeEach(() => {
     mockRedisXrange.mockClear();
     server.resetState();
+    statsdIncrementMock = jest.spyOn(statsd, 'increment').mockReturnValue();
+  });
+
+  afterEach(() => {
+    statsdIncrementMock.mockRestore();
   });
 
   describe('HTTP requests', () => {
@@ -90,7 +98,7 @@ describe('server', () => {
       expect(endMock).toHaveBeenLastCalledWith('OK');
     });
 
-    test('reponds with a 404 otherwise', () => {
+    test('reponds with a 404 when not found', () => {
       const endMock = jest.fn();
       const writeHeadMock = jest.fn();
 
@@ -168,9 +176,17 @@ describe('server', () => {
       const ws = new wsMock('localhost');
       const sendMock = jest.spyOn(ws, 'send');
       const socketInstance = { ws: ws, channel: channelId, pongTs: Date.now() };
+
+      expect(statsdIncrementMock).toBeCalledTimes(0);
       server.trackClient(channelId, socketInstance);
+      expect(statsdIncrementMock).toBeCalledTimes(1);
+      expect(statsdIncrementMock).toHaveBeenNthCalledWith(
+        1,
+        'ws_connected_client',
+      );
 
       server.processStreamResults(streamReturnValue);
+      expect(statsdIncrementMock).toBeCalledTimes(1);
 
       const message1 = `{"id":"1615426152415-0","channel_id":"${channelId}","job_id":"c9b99965-8f1e-4ce5-aa43-d6fc94d6a510","user_id":"1","status":"done","errors":[],"result_url":"/superset/explore_json/data/ejr-37281682b1282cdb8f25e0de0339b386"}`;
       const message2 = `{"id":"1615426152516-0","channel_id":"${channelId}","job_id":"f1e5bb1f-f2f1-4f21-9b2f-c9b91dcc9b59","user_id":"1","status":"done","errors":[],"result_url":"/api/v1/chart/data/qc-64e8452dc9907dd77746cb75a19202de"}`;
@@ -182,7 +198,9 @@ describe('server', () => {
       const ws = new wsMock('localhost');
       const sendMock = jest.spyOn(ws, 'send');
 
+      expect(statsdIncrementMock).toBeCalledTimes(0);
       server.processStreamResults(streamReturnValue);
+      expect(statsdIncrementMock).toBeCalledTimes(0);
 
       expect(sendMock).not.toHaveBeenCalled();
     });
@@ -194,9 +212,21 @@ describe('server', () => {
       });
       const cleanChannelMock = jest.spyOn(server, 'cleanChannel');
       const socketInstance = { ws: ws, channel: channelId, pongTs: Date.now() };
+
+      expect(statsdIncrementMock).toBeCalledTimes(0);
       server.trackClient(channelId, socketInstance);
+      expect(statsdIncrementMock).toBeCalledTimes(1);
+      expect(statsdIncrementMock).toHaveBeenNthCalledWith(
+        1,
+        'ws_connected_client',
+      );
 
       server.processStreamResults(streamReturnValue);
+      expect(statsdIncrementMock).toBeCalledTimes(2);
+      expect(statsdIncrementMock).toHaveBeenNthCalledWith(
+        2,
+        'ws_client_send_error',
+      );
 
       expect(sendMock).toHaveBeenCalled();
       expect(cleanChannelMock).toHaveBeenCalledWith(channelId);

superset-websocket/spec/index.test.ts

@@ -23,6 +23,7 @@ import { v4 as uuidv4 } from 'uuid';
 import jwt from 'jsonwebtoken';
 import cookie from 'cookie';
 import Redis from 'ioredis';
+import StatsD from 'hot-shots';
 
 import { createLogger } from './logger';
 
@@ -84,6 +85,11 @@ export const opts = {
   logLevel: 'info',
   logToFile: false,
   logFilename: 'app.log',
+  statsd: {
+    host: '127.0.0.1',
+    port: 8125,
+    globalTags: [],
+  },
   redis: {
     port: 6379,
     host: '127.0.0.1',
@@ -121,6 +127,13 @@ const logger = createLogger({
   logFilename: opts.logFilename,
 });
 
+export const statsd = new StatsD({
+  ...opts.statsd,
+  errorHandler: (e: Error) => {
+    logger.error(e);
+  },
+});
+
 // enforce JWT secret length
 if (startServer && opts.jwtSecret.length < 32)
   throw new Error('Please provide a JWT secret at least 32 bytes long');
@@ -160,6 +173,8 @@ export const trackClient = (
   channel: string,
   socketInstance: SocketInstance,
 ): string => {
+  statsd.increment('ws_connected_client');
+
   const socketId = uuidv4();
   sockets[socketId] = socketInstance;
 
@@ -189,6 +204,7 @@ export const sendToChannel = (channel: string, value: EventValue): void => {
     try {
       socketInstance.ws.send(strData);
     } catch (err) {
+      statsd.increment('ws_client_send_error');
       logger.debug(`Error sending to socket: ${err}`);
       // check that the connection is still active
       cleanChannel(channel);