Add instance id to socks metrics (#2137)

(cherry picked from commit 76bd1c8)

indexer/Dockerfile.bazooka.remote

@@ -11,6 +11,7 @@ COPY ./patches ./patches
 
 # Copy bazooka and imported packages
 COPY ./packages/base/ ./packages/base/
+COPY ./packages/dev/ ./packages/dev/
 COPY ./packages/postgres/ ./packages/postgres/
 COPY ./packages/v4-protos/ ./packages/v4-protos/
 COPY ./packages/kafka/ ./packages/kafka/

indexer/packages/base/__tests__/instance-id.test.ts

@@ -0,0 +1,87 @@
+import { setInstanceId, getInstanceId, resetForTests } from '../src/instance-id';
+import { axiosRequest } from '../src/axios';
+import { asMock } from '@dydxprotocol-indexer/dev';
+import logger from '../src/logger';
+import config from '../src/config';
+
+jest.mock('../src/axios', () => ({
+  ...(jest.requireActual('../src/axios') as object),
+  axiosRequest: jest.fn(),
+}));
+
+describe('instance-id', () => {
+  describe('setInstanceId', () => {
+    const defaultTaskArn = 'defaultTaskArn';
+    const defaultResponse = {
+      TaskARN: defaultTaskArn,
+    };
+    const ecsUrl = config.ECS_CONTAINER_METADATA_URI_V4;
+
+    beforeEach(() => {
+      config.ECS_CONTAINER_METADATA_URI_V4 = ecsUrl;
+      resetForTests();
+      jest.resetAllMocks();
+      jest.restoreAllMocks();
+      asMock(axiosRequest).mockResolvedValue(defaultResponse);
+    });
+
+    afterAll(() => {
+      jest.clearAllMocks();
+      jest.restoreAllMocks();
+    });
+
+    it('should set instance id to task ARN in staging', async () => {
+      jest.spyOn(config, 'isStaging').mockReturnValueOnce(true);
+      config.ECS_CONTAINER_METADATA_URI_V4 = 'url';
+      await setInstanceId();
+
+      expect(getInstanceId()).toEqual(defaultTaskArn);
+    });
+
+    it('should set instance id to task ARN in production', async () => {
+      jest.spyOn(config, 'isProduction').mockReturnValueOnce(true);
+      config.ECS_CONTAINER_METADATA_URI_V4 = 'url';
+      await setInstanceId();
+
+      expect(getInstanceId()).toEqual(defaultTaskArn);
+    });
+
+    it('should not call metadata endpoint if not production or staging', async () => {
+      config.ECS_CONTAINER_METADATA_URI_V4 = 'url';
+      await setInstanceId();
+
+      expect(getInstanceId()).not.toEqual(defaultTaskArn);
+      expect(asMock(axiosRequest)).not.toHaveBeenCalled();
+    });
+
+    it('should not set instance id if already set', async () => {
+      jest.spyOn(config, 'isStaging').mockReturnValue(true);
+      config.ECS_CONTAINER_METADATA_URI_V4 = 'url';
+      await setInstanceId();
+      const instanceId = getInstanceId();
+      await setInstanceId();
+
+      expect(getInstanceId()).toEqual(instanceId);
+      expect(axiosRequest).toHaveBeenCalledTimes(1);
+    });
+
+    it('should log error and set instance id to uuid if request errors', async () => {
+      jest.spyOn(config, 'isStaging').mockReturnValue(true);
+      config.ECS_CONTAINER_METADATA_URI_V4 = 'url';
+      const loggerErrorSpy = jest.spyOn(logger, 'error');
+      const emptyInstanceId = getInstanceId();
+      asMock(axiosRequest).mockRejectedValueOnce(new Error());
+      await setInstanceId();
+
+      expect(loggerErrorSpy).toHaveBeenCalledTimes(1);
+      expect(getInstanceId()).not.toEqual(emptyInstanceId);
+    });
+
+    it('should not call metadata endpoint if url is empty', async () => {
+      jest.spyOn(config, 'isStaging').mockReturnValue(true);
+      await setInstanceId();
+
+      expect(axiosRequest).not.toHaveBeenCalled();
+    });
+  });
+});

indexer/packages/base/package.json

@@ -2,14 +2,15 @@
   "name": "@dydxprotocol-indexer/base",
   "version": "0.0.1",
   "description": "",
-  "main": "build/index.js",
+  "main": "build/src/index.js",
   "devDependencies": {
     "@dydxprotocol-indexer/dev": "workspace:^0.0.1",
     "@types/big.js": "^6.1.5",
     "@types/express": "^4.17.13",
     "@types/lodash": "^4.14.182",
     "@types/traverse": "^0.6.32",
     "express": "^4.18.1",
+    "jest": "^28.1.2",
     "typescript": "^4.7.4"
   },
   "scripts": {
@@ -18,7 +19,7 @@
     "build": "rm -rf build/ && tsc",
     "build:prod": "pnpm run build",
     "build:watch": "pnpm run build -- --watch",
-    "test": "echo \"Error: no test specified\""
+    "test": "NODE_ENV=test jest --runInBand --forceExit"
   },
   "repository": {
     "type": "git",
@@ -38,6 +39,7 @@
     "hot-shots": "^9.1.0",
     "lodash": "^4.17.21",
     "traverse": "^0.6.6",
+    "uuid": "^8.3.2",
     "winston": "^3.8.1",
     "winston-transport": "^4.5.0",
     "@bugsnag/core": "^7.18.0",

indexer/packages/base/src/config.ts

@@ -30,14 +30,15 @@ export const baseConfigSchema = {
   SEND_BUGSNAG_ERRORS: parseBoolean({
     default: true,
   }),
-  SERVICE_NAME: parseString(),
+  SERVICE_NAME: parseString({ default: '' }),
 
   // Optional environment variables.
   NODE_ENV: parseString({ default: null }),
   ENABLE_LOGS_IN_TEST: parseBoolean({ default: false }),
   STATSD_HOST: parseString({ default: 'localhost' }),
   STATSD_PORT: parseInteger({ default: 8125 }),
   LOG_LEVEL: parseString({ default: 'debug' }),
+  ECS_CONTAINER_METADATA_URI_V4: parseString({ default: '' }),
 };
 
 export default parseSchema(baseConfigSchema);

indexer/packages/base/src/index.ts

@@ -13,6 +13,7 @@ export * from './constants';
 export * from './bugsnag';
 export * from './stats-util';
 export * from './date-helpers';
+export * from './instance-id';
 
 // Do this outside logger.ts to avoid a dependency cycle with logger transports that may trigger
 // additional logging.

indexer/packages/base/src/instance-id.ts

@@ -0,0 +1,51 @@
+import { v4 as uuidv4 } from 'uuid';
+
+import { axiosRequest } from './axios';
+import config from './config';
+import logger from './logger';
+
+let INSTANCE_ID: string = '';
+
+export function getInstanceId(): string {
+  return INSTANCE_ID;
+}
+
+export async function setInstanceId(): Promise<void> {
+  if (INSTANCE_ID !== '') {
+    return;
+  }
+  if (config.ECS_CONTAINER_METADATA_URI_V4 !== '' &&
+      (
+        config.isProduction() || config.isStaging()
+      )
+  ) {
+    // https://docs.aws.amazon.com/AmazonECS/latest/developerguide/task-metadata-endpoint-v4.html
+    const taskUrl = `${config.ECS_CONTAINER_METADATA_URI_V4}/task`;
+    try {
+      const response = await axiosRequest({
+        method: 'GET',
+        url: taskUrl,
+      }) as { TaskARN: string };
+      INSTANCE_ID = response.TaskARN;
+    } catch (error) {
+      logger.error({
+        at: 'instance-id#setInstanceId',
+        message: 'Failed to retrieve task arn from metadata endpoint. Falling back to uuid.',
+        error,
+        taskUrl,
+      });
+      INSTANCE_ID = uuidv4();
+    }
+  } else {
+    INSTANCE_ID = uuidv4();
+
+  }
+}
+
+// Exported for tests
+export function resetForTests(): void {
+  if (!config.isTest()) {
+    throw new Error(`resetForTests() cannot be called for env: ${config.NODE_ENV}`);
+  }
+  INSTANCE_ID = '';
+}

indexer/packages/base/tsconfig.json

@@ -4,6 +4,7 @@
     "outDir": "build"
   },
   "include": [
-    "src"
+    "src",
+    "__tests__"
   ]
 }

indexer/services/socks/src/helpers/wss.ts

@@ -1,4 +1,4 @@
-import { stats, logger } from '@dydxprotocol-indexer/base';
+import { stats, getInstanceId, logger } from '@dydxprotocol-indexer/base';
 import WebSocket from 'ws';
 
 import config from '../config';
@@ -105,6 +105,7 @@ export function sendMessageString(
       1,
       config.MESSAGE_FORWARDER_STATSD_SAMPLE_RATE,
       {
+        instance: getInstanceId(),
         reason: WEBSOCKET_NOT_OPEN,
         readyState: ws.readyState.toString(),
       },
@@ -118,7 +119,10 @@ export function sendMessageString(
         `${config.SERVICE_NAME}.ws_send.error`,
         1,
         config.MESSAGE_FORWARDER_STATSD_SAMPLE_RATE,
-        { code: (error as WssError)?.code },
+        {
+          instance: getInstanceId(),
+          code: (error as WssError)?.code,
+        },
       );
       const errorLog = { // type is InfoObject in node-service-base
         at: 'wss#sendMessage',
@@ -148,7 +152,10 @@ export function sendMessageString(
           stats.increment(
             `${config.SERVICE_NAME}.ws_send.stream_destroyed_errors`,
             1,
-            { action: 'close' },
+            {
+              action: 'close',
+              instance: getInstanceId(),
+            },
           );
         } else {
           logger.error(closeErrorLog);

indexer/services/socks/src/index.ts

@@ -1,4 +1,6 @@
-import { logger, startBugsnag, wrapBackgroundTask } from '@dydxprotocol-indexer/base';
+import {
+  logger, setInstanceId, getInstanceId, startBugsnag, wrapBackgroundTask,
+} from '@dydxprotocol-indexer/base';
 import { startConsumer } from '@dydxprotocol-indexer/kafka';
 import { blockHeightRefresher, perpetualMarketRefresher } from '@dydxprotocol-indexer/postgres';
 
@@ -39,6 +41,18 @@ async function start(): Promise<void> {
 
   startBugsnag();
 
+  logger.info({
+    at: 'index#start',
+    message: 'Getting instance id...',
+  });
+
+  await setInstanceId();
+
+  logger.info({
+    at: 'index#start',
+    message: `Got instance id ${getInstanceId()}.`,
+  });
+
   // Initialize PerpetualMarkets and BlockHeight cache
   await Promise.all([
     blockHeightRefresher.updateBlockHeight(),