feat: add xray to analytics scheduler and websockets

.pnp.cjs

@@ -5740,6 +5740,7 @@ const RAW_RUNTIME_STATE =
           ["@types/node", "npm:20.5.7"],\
           ["@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"],\
           ["dayjs", "npm:1.11.7"],\
           ["dotenv", "npm:16.1.3"],\
           ["eslint", "npm:8.41.0"],\
@@ -6235,6 +6236,7 @@ const RAW_RUNTIME_STATE =
           ["@types/node", "npm:20.5.7"],\
           ["@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"],\
           ["dayjs", "npm:1.11.7"],\
           ["dotenv", "npm:16.1.3"],\
           ["eslint", "npm:8.41.0"],\
@@ -6472,6 +6474,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"],\
           ["axios", "npm:1.4.0"],\
           ["cors", "npm:2.8.5"],\
           ["dotenv", "npm:16.1.3"],\

packages/analytics/bin/worker.ts

@@ -1,6 +1,7 @@
 import 'reflect-metadata'
 
 import { Logger } from 'winston'
+import * as AWSXRay from 'aws-xray-sdk'
 import { DomainEventSubscriberFactoryInterface } from '@standardnotes/domain-events'
 import * as dayjs from 'dayjs'
 import * as utc from 'dayjs/plugin/utc'
@@ -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.ECSPlugin])
+  }
+
   const logger: Logger = container.get(TYPES.Logger)
 
   logger.info('Starting worker...')

packages/analytics/package.json

@@ -46,6 +46,7 @@
     "@standardnotes/domain-events": "workspace:*",
     "@standardnotes/domain-events-infra": "workspace:*",
     "@standardnotes/time": "workspace:*",
+    "aws-xray-sdk": "^3.5.2",
     "dayjs": "^1.11.6",
     "dotenv": "^16.0.1",
     "inversify": "^6.0.1",

packages/analytics/src/Bootstrap/Container.ts

@@ -6,7 +6,7 @@ import {
   DomainEventMessageHandlerInterface,
   DomainEventSubscriberFactoryInterface,
 } from '@standardnotes/domain-events'
-import { MapperInterface } from '@standardnotes/domain-core'
+import { MapperInterface, ServiceIdentifier } from '@standardnotes/domain-core'
 // eslint-disable-next-line @typescript-eslint/no-var-requires
 const Mixpanel = require('mixpanel')
 
@@ -18,7 +18,7 @@ import {
   SNSDomainEventPublisher,
   SQSDomainEventSubscriberFactory,
   SQSEventMessageHandler,
-  SQSNewRelicEventMessageHandler,
+  SQSXRayEventMessageHandler,
 } from '@standardnotes/domain-events-infra'
 import { Timer, TimerInterface } from '@standardnotes/time'
 import { PeriodKeyGeneratorInterface } from '../Domain/Time/PeriodKeyGeneratorInterface'
@@ -57,6 +57,7 @@ import { SNSClient, SNSClientConfig } from '@aws-sdk/client-sns'
 import { SQSClient, SQSClientConfig } from '@aws-sdk/client-sqs'
 import { SessionCreatedEventHandler } from '../Domain/Handler/SessionCreatedEventHandler'
 import { SessionRefreshedEventHandler } from '../Domain/Handler/SessionRefreshedEventHandler'
+import { captureAWSv3Client } from 'aws-xray-sdk'
 
 export class ContainerConfigLoader {
   async load(): Promise<Container> {
@@ -107,7 +108,7 @@ export class ContainerConfigLoader {
         secretAccessKey: env.get('SNS_SECRET_ACCESS_KEY', true),
       }
     }
-    container.bind<SNSClient>(TYPES.SNS).toConstantValue(new SNSClient(snsConfig))
+    container.bind<SNSClient>(TYPES.SNS).toConstantValue(captureAWSv3Client(new SNSClient(snsConfig)))
 
     if (env.get('SQS_QUEUE_URL', true)) {
       const sqsConfig: SQSClientConfig = {
@@ -122,7 +123,7 @@ export class ContainerConfigLoader {
           secretAccessKey: env.get('SQS_SECRET_ACCESS_KEY', true),
         }
       }
-      container.bind<SQSClient>(TYPES.SQS).toConstantValue(new SQSClient(sqsConfig))
+      container.bind<SQSClient>(TYPES.SQS).toConstantValue(captureAWSv3Client(new SQSClient(sqsConfig)))
     }
 
     // env vars
@@ -246,7 +247,11 @@ export class ContainerConfigLoader {
       .bind<DomainEventMessageHandlerInterface>(TYPES.DomainEventMessageHandler)
       .toConstantValue(
         env.get('NEW_RELIC_ENABLED', true) === 'true'
-          ? new SQSNewRelicEventMessageHandler(eventHandlers, container.get(TYPES.Logger))
+          ? new SQSXRayEventMessageHandler(
+              ServiceIdentifier.NAMES.AnalyticsWorker,
+              eventHandlers,
+              container.get(TYPES.Logger),
+            )
           : new SQSEventMessageHandler(eventHandlers, container.get(TYPES.Logger)),
       )
     container

packages/domain-core/src/Domain/Service/ServiceIdentifier.ts

@@ -4,6 +4,7 @@ import { ServiceIdentifierProps } from './ServiceIdentifierProps'
 
 export class ServiceIdentifier extends ValueObject<ServiceIdentifierProps> {
   static readonly NAMES = {
+    AnalyticsWorker: 'AnalyticsWorker',
     ApiGateway: 'ApiGateway',
     Auth: 'Auth',
     AuthWorker: 'AuthWorker',
@@ -13,6 +14,11 @@ export class ServiceIdentifier extends ValueObject<ServiceIdentifierProps> {
     RevisionsWorker: 'RevisionsWorker',
     Files: 'Files',
     FilesWorker: 'FilesWorker',
+    SchedulerWorker: 'SchedulerWorker',
+    Email: 'Email',
+    EmailWorker: 'EmailWorker',
+    Websockets: 'Websockets',
+    WebsocketsWorker: 'WebsocketsWorker',
   }
 
   get value(): string {

packages/scheduler/bin/worker.ts

@@ -1,6 +1,7 @@
 import 'reflect-metadata'
 
 import { Logger } from 'winston'
+import * as AWSXRay from 'aws-xray-sdk'
 import { DomainEventSubscriberFactoryInterface } from '@standardnotes/domain-events'
 import * as dayjs from 'dayjs'
 import * as utc from 'dayjs/plugin/utc'
@@ -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.ECSPlugin])
+  }
+
   const logger: Logger = container.get(TYPES.Logger)
 
   logger.info('Starting worker...')

packages/scheduler/package.json

@@ -30,6 +30,7 @@
     "@standardnotes/domain-events-infra": "workspace:*",
     "@standardnotes/predicates": "workspace:*",
     "@standardnotes/time": "workspace:*",
+    "aws-xray-sdk": "^3.5.2",
     "dayjs": "^1.11.6",
     "dotenv": "^16.0.1",
     "inversify": "^6.0.1",

packages/scheduler/src/Bootstrap/Container.ts

@@ -17,7 +17,7 @@ import {
   SNSDomainEventPublisher,
   SQSDomainEventSubscriberFactory,
   SQSEventMessageHandler,
-  SQSNewRelicEventMessageHandler,
+  SQSXRayEventMessageHandler,
 } from '@standardnotes/domain-events-infra'
 import { Timer, TimerInterface } from '@standardnotes/time'
 import { PredicateRepositoryInterface } from '../Domain/Predicate/PredicateRepositoryInterface'
@@ -35,6 +35,8 @@ import { VerifyPredicates } from '../Domain/UseCase/VerifyPredicates/VerifyPredi
 import { UserRegisteredEventHandler } from '../Domain/Handler/UserRegisteredEventHandler'
 import { SubscriptionCancelledEventHandler } from '../Domain/Handler/SubscriptionCancelledEventHandler'
 import { ExitDiscountAppliedEventHandler } from '../Domain/Handler/ExitDiscountAppliedEventHandler'
+import { ServiceIdentifier } from '@standardnotes/domain-core'
+import { captureAWSv3Client } from 'aws-xray-sdk'
 
 export class ContainerConfigLoader {
   async load(): Promise<Container> {
@@ -86,7 +88,7 @@ export class ContainerConfigLoader {
           secretAccessKey: env.get('SNS_SECRET_ACCESS_KEY', true),
         }
       }
-      container.bind<SNSClient>(TYPES.SNS).toConstantValue(new SNSClient(snsConfig))
+      container.bind<SNSClient>(TYPES.SNS).toConstantValue(captureAWSv3Client(new SNSClient(snsConfig)))
     }
 
     if (env.get('SQS_QUEUE_URL', true)) {
@@ -102,7 +104,7 @@ export class ContainerConfigLoader {
           secretAccessKey: env.get('SQS_SECRET_ACCESS_KEY', true),
         }
       }
-      container.bind<SQSClient>(TYPES.SQS).toConstantValue(new SQSClient(sqsConfig))
+      container.bind<SQSClient>(TYPES.SQS).toConstantValue(captureAWSv3Client(new SQSClient(sqsConfig)))
     }
 
     // env vars
@@ -156,7 +158,11 @@ export class ContainerConfigLoader {
       .bind<DomainEventMessageHandlerInterface>(TYPES.DomainEventMessageHandler)
       .toConstantValue(
         env.get('NEW_RELIC_ENABLED', true) === 'true'
-          ? new SQSNewRelicEventMessageHandler(eventHandlers, container.get(TYPES.Logger))
+          ? new SQSXRayEventMessageHandler(
+              ServiceIdentifier.NAMES.SchedulerWorker,
+              eventHandlers,
+              container.get(TYPES.Logger),
+            )
           : new SQSEventMessageHandler(eventHandlers, container.get(TYPES.Logger)),
       )
     container

packages/websockets/bin/server.ts

@@ -4,6 +4,7 @@ import '../src/Infra/InversifyExpressUtils/AnnotatedHealthCheckController'
 import '../src/Infra/InversifyExpressUtils/AnnotatedWebSocketsController'
 
 import * as cors from 'cors'
+import * as AWSXRay from 'aws-xray-sdk'
 import { urlencoded, json, Request, Response, NextFunction } from 'express'
 import * as winston from 'winston'
 
@@ -11,15 +12,20 @@ import { InversifyExpressServer } from 'inversify-express-utils'
 import { ContainerConfigLoader } from '../src/Bootstrap/Container'
 import TYPES from '../src/Bootstrap/Types'
 import { Env } from '../src/Bootstrap/Env'
+import { ServiceIdentifier } from '@standardnotes/domain-core'
 
 const container = new ContainerConfigLoader()
 void container.load().then((container) => {
   const env: Env = new Env()
   env.load()
 
+  AWSXRay.config([AWSXRay.plugins.ECSPlugin])
+
   const server = new InversifyExpressServer(container)
 
   server.setConfig((app) => {
+    app.use(AWSXRay.express.openSegment(ServiceIdentifier.NAMES.Websockets))
+
     app.use((_request: Request, response: Response, next: NextFunction) => {
       response.setHeader('X-Websockets-Version', container.get(TYPES.VERSION))
       next()
@@ -46,6 +52,8 @@ void container.load().then((container) => {
 
   const serverInstance = server.build()
 
+  serverInstance.use(AWSXRay.express.closeSegment())
+
   serverInstance.listen(env.get('PORT'))
 
   logger.info(`Server started on port ${process.env.PORT}`)

packages/websockets/bin/worker.ts

@@ -1,6 +1,7 @@
 import 'reflect-metadata'
 
 import { Logger } from 'winston'
+import * as AWSXRay from 'aws-xray-sdk'
 
 import { ContainerConfigLoader } from '../src/Bootstrap/Container'
 import TYPES from '../src/Bootstrap/Types'
@@ -12,6 +13,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.ECSPlugin])
+  }
+
   const logger: Logger = container.get(TYPES.Logger)
 
   logger.info('Starting worker...')

packages/websockets/package.json

@@ -31,6 +31,7 @@
     "@standardnotes/domain-events-infra": "workspace:^",
     "@standardnotes/responses": "^1.13.27",
     "@standardnotes/security": "workspace:^",
+    "aws-xray-sdk": "^3.5.2",
     "axios": "^1.1.3",
     "cors": "2.8.5",
     "dotenv": "^16.0.1",

packages/websockets/src/Bootstrap/Container.ts

@@ -21,7 +21,7 @@ import { WebSocketsClientMessenger } from '../Infra/WebSockets/WebSocketsClientM
 import {
   SQSDomainEventSubscriberFactory,
   SQSEventMessageHandler,
-  SQSNewRelicEventMessageHandler,
+  SQSXRayEventMessageHandler,
 } from '@standardnotes/domain-events-infra'
 import { ApiGatewayAuthMiddleware } from '../Controller/ApiGatewayAuthMiddleware'
 
@@ -38,6 +38,8 @@ import { WebSocketsController } from '../Controller/WebSocketsController'
 import { WebSocketServerInterface } from '@standardnotes/api'
 import { ClientMessengerInterface } from '../Client/ClientMessengerInterface'
 import { WebSocketMessageRequestedEventHandler } from '../Domain/Handler/WebSocketMessageRequestedEventHandler'
+import { ServiceIdentifier } from '@standardnotes/domain-core'
+import { captureAWSv3Client } from 'aws-xray-sdk'
 
 export class ContainerConfigLoader {
   async load(): Promise<Container> {
@@ -86,16 +88,18 @@ export class ContainerConfigLoader {
           secretAccessKey: env.get('SQS_SECRET_ACCESS_KEY', true),
         }
       }
-      container.bind<SQSClient>(TYPES.SQS).toConstantValue(new SQSClient(sqsConfig))
+      container.bind<SQSClient>(TYPES.SQS).toConstantValue(captureAWSv3Client(new SQSClient(sqsConfig)))
     }
 
     container.bind(TYPES.WEBSOCKETS_API_URL).toConstantValue(env.get('WEBSOCKETS_API_URL', true))
 
     container.bind<ApiGatewayManagementApiClient>(TYPES.WebSockets_ApiGatewayManagementApiClient).toConstantValue(
-      new ApiGatewayManagementApiClient({
-        endpoint: container.get(TYPES.WEBSOCKETS_API_URL),
-        region: env.get('API_GATEWAY_AWS_REGION', true) ?? 'us-east-1',
-      }),
+      captureAWSv3Client(
+        new ApiGatewayManagementApiClient({
+          endpoint: container.get(TYPES.WEBSOCKETS_API_URL),
+          region: env.get('API_GATEWAY_AWS_REGION', true) ?? 'us-east-1',
+        }),
+      ),
     )
 
     // Controller
@@ -154,7 +158,11 @@ export class ContainerConfigLoader {
       .bind<DomainEventMessageHandlerInterface>(TYPES.DomainEventMessageHandler)
       .toConstantValue(
         env.get('NEW_RELIC_ENABLED', true) === 'true'
-          ? new SQSNewRelicEventMessageHandler(eventHandlers, container.get(TYPES.Logger))
+          ? new SQSXRayEventMessageHandler(
+              ServiceIdentifier.NAMES.WebsocketsWorker,
+              eventHandlers,
+              container.get(TYPES.Logger),
+            )
           : new SQSEventMessageHandler(eventHandlers, container.get(TYPES.Logger)),
       )
     container

yarn.lock

@@ -4621,6 +4621,7 @@ __metadata:
     "@types/node": "npm:^20.5.7"
     "@typescript-eslint/eslint-plugin": "npm:^6.5.0"
     "@typescript-eslint/parser": "npm:^6.5.0"
+    aws-xray-sdk: "npm:^3.5.2"
     dayjs: "npm:^1.11.6"
     dotenv: "npm:^16.0.1"
     eslint: "npm:^8.39.0"
@@ -5131,6 +5132,7 @@ __metadata:
     "@types/node": "npm:^20.5.7"
     "@typescript-eslint/eslint-plugin": "npm:^6.5.0"
     "@typescript-eslint/parser": "npm:^6.5.0"
+    aws-xray-sdk: "npm:^3.5.2"
     dayjs: "npm:^1.11.6"
     dotenv: "npm:^16.0.1"
     eslint: "npm:^8.39.0"
@@ -5371,6 +5373,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"
     axios: "npm:^1.1.3"
     cors: "npm:2.8.5"
     dotenv: "npm:^16.0.1"