feat: add xray segment tracing on auth-worker

.pnp.cjs

@@ -5963,6 +5963,7 @@ const RAW_RUNTIME_STATE =
           ["@types/newrelic", "npm:9.14.0"],\
           ["@typescript-eslint/eslint-plugin", "virtual:fd909b174d079e30b336c4ce72c38a88c1e447767b1a8dd7655e07719a1e31b97807f0931368724fc78897ff15e6a6d00b83316c0f76d11f85111f342e08bb79#npm:6.5.0"],\
           ["@typescript-eslint/parser", "virtual:fd909b174d079e30b336c4ce72c38a88c1e447767b1a8dd7655e07719a1e31b97807f0931368724fc78897ff15e6a6d00b83316c0f76d11f85111f342e08bb79#npm:6.5.0"],\
+          ["aws-xray-sdk", "npm:3.5.2"],\
           ["eslint", "npm:8.41.0"],\
           ["eslint-plugin-prettier", "virtual:fd909b174d079e30b336c4ce72c38a88c1e447767b1a8dd7655e07719a1e31b97807f0931368724fc78897ff15e6a6d00b83316c0f76d11f85111f342e08bb79#npm:5.0.0"],\
           ["ioredis", "npm:5.3.2"],\

packages/auth/bin/worker.ts

@@ -8,6 +8,7 @@ import { Env } from '../src/Bootstrap/Env'
 import { DomainEventSubscriberFactoryInterface } from '@standardnotes/domain-events'
 import * as dayjs from 'dayjs'
 import * as utc from 'dayjs/plugin/utc'
+import * as AWSXRay from 'aws-xray-sdk'
 
 const container = new ContainerConfigLoader('worker')
 void container.load().then((container) => {
@@ -16,6 +17,14 @@ void container.load().then((container) => {
   const env: Env = new Env()
   env.load()
 
+  const isConfiguredForAWSProduction =
+    env.get('MODE', true) !== 'home-server' && env.get('MODE', true) !== 'self-hosted'
+
+  if (isConfiguredForAWSProduction) {
+    AWSXRay.enableManualMode()
+    AWSXRay.config([AWSXRay.plugins.EC2Plugin, AWSXRay.plugins.ECSPlugin])
+  }
+
   const logger: Logger = container.get(TYPES.Auth_Logger)
 
   logger.info('Starting worker...')

packages/auth/src/Bootstrap/Container.ts

@@ -93,7 +93,7 @@ import {
   SNSDomainEventPublisher,
   SQSDomainEventSubscriberFactory,
   SQSEventMessageHandler,
-  SQSNewRelicEventMessageHandler,
+  SQSXRayEventMessageHandler,
 } from '@standardnotes/domain-events-infra'
 import { GetUserSubscription } from '../Domain/UseCase/GetUserSubscription/GetUserSubscription'
 import { ChangeCredentials } from '../Domain/UseCase/ChangeCredentials/ChangeCredentials'
@@ -1236,8 +1236,8 @@ export class ContainerConfigLoader {
       container
         .bind<DomainEventMessageHandlerInterface>(TYPES.Auth_DomainEventMessageHandler)
         .toConstantValue(
-          env.get('NEW_RELIC_ENABLED', true) === 'true'
+          isConfiguredForAWSProduction
-            ? new SQSNewRelicEventMessageHandler(eventHandlers, container.get(TYPES.Auth_Logger))
+            ? new SQSXRayEventMessageHandler(eventHandlers, container.get(TYPES.Auth_Logger))
             : new SQSEventMessageHandler(eventHandlers, container.get(TYPES.Auth_Logger)),
         )
 

packages/domain-events-infra/package.json

@@ -27,6 +27,7 @@
     "@aws-sdk/client-sns": "^3.332.0",
     "@aws-sdk/client-sqs": "^3.332.0",
     "@standardnotes/domain-events": "workspace:*",
+    "aws-xray-sdk": "^3.5.2",
     "ioredis": "^5.2.4",
     "reflect-metadata": "^0.1.13",
     "sqs-consumer": "^6.2.1",

packages/domain-events-infra/src/Infra/SQS/SQSXRayEventMessageHandler.ts

@@ -0,0 +1,60 @@
+import { Logger } from 'winston'
+import * as zlib from 'zlib'
+import { Segment, Subsegment, captureAsyncFunc } from 'aws-xray-sdk'
+
+import {
+  DomainEventHandlerInterface,
+  DomainEventInterface,
+  DomainEventMessageHandlerInterface,
+} from '@standardnotes/domain-events'
+
+export class SQSXRayEventMessageHandler implements DomainEventMessageHandlerInterface {
+  constructor(
+    private handlers: Map<string, DomainEventHandlerInterface>,
+    private logger: Logger,
+  ) {}
+
+  async handleMessage(message: string): Promise<void> {
+    const messageParsed = JSON.parse(message)
+
+    const domainEventJson = zlib.unzipSync(Buffer.from(messageParsed.Message, 'base64')).toString()
+
+    const domainEvent: DomainEventInterface = JSON.parse(domainEventJson)
+
+    domainEvent.createdAt = new Date(domainEvent.createdAt)
+
+    const handler = this.handlers.get(domainEvent.type)
+    if (!handler) {
+      this.logger.debug(`Event handler for event type ${domainEvent.type} does not exist`)
+
+      return
+    }
+
+    this.logger.debug(`Received event: ${domainEvent.type}`)
+
+    const xRaySegment = new Segment(domainEvent.type)
+
+    if (domainEvent.meta.correlation.userIdentifierType === 'uuid') {
+      xRaySegment.setUser(domainEvent.meta.correlation.userIdentifier)
+    }
+
+    await captureAsyncFunc(
+      `${handler.constructor.name}.handle}`,
+      async (subsegment?: Subsegment) => {
+        await handler.handle(domainEvent)
+
+        if (subsegment) {
+          subsegment.close()
+        }
+      },
+      xRaySegment,
+    )
+
+    xRaySegment.close()
+    xRaySegment.flush()
+  }
+
+  async handleError(error: Error): Promise<void> {
+    this.logger.error('Error occured while handling SQS message: %O', error)
+  }
+}

packages/domain-events-infra/src/Infra/index.ts

@@ -12,3 +12,4 @@ export * from './SQS/SQSNewRelicBounceNotificiationHandler'
 export * from './SQS/SQSDomainEventSubscriberFactory'
 export * from './SQS/SQSEventMessageHandler'
 export * from './SQS/SQSNewRelicEventMessageHandler'
+export * from './SQS/SQSXRayEventMessageHandler'

yarn.lock

@@ -4834,6 +4834,7 @@ __metadata:
     "@types/newrelic": "npm:^9.14.0"
     "@typescript-eslint/eslint-plugin": "npm:^6.5.0"
     "@typescript-eslint/parser": "npm:^6.5.0"
+    aws-xray-sdk: "npm:^3.5.2"
     eslint: "npm:^8.39.0"
     eslint-plugin-prettier: "npm:^5.0.0"
     ioredis: "npm:^5.2.4"