From 2ca81b22bf90faa0d04da5d2f28a44733edd1e7e Mon Sep 17 00:00:00 2001
From: branberry <brandon.ly@mongodb.com>
Date: Mon, 13 Nov 2023 14:08:00 -0600
Subject: [PATCH] Send subsegment to xray daemon

---
 src/enhanced/enhancedApp.ts       |  2 +-
 src/enhanced/utils/queue/index.ts | 34 +++++++++++++++++++++++++++++++
 src/entities/queueMessage.ts      |  1 +
 3 files changed, 36 insertions(+), 1 deletion(-)

diff --git a/src/enhanced/enhancedApp.ts b/src/enhanced/enhancedApp.ts
index a777eadfd..5f5e47808 100644
--- a/src/enhanced/enhancedApp.ts
+++ b/src/enhanced/enhancedApp.ts
@@ -1,5 +1,5 @@
 import AWSXRay from 'aws-xray-sdk';
-// eslint-disable-next-line @typescript-eslint/no-var-requires
+
 
 import mongodb, { MongoClient } from 'mongodb';
 import c from 'config';
diff --git a/src/enhanced/utils/queue/index.ts b/src/enhanced/utils/queue/index.ts
index 9078fdbc8..7e21e4036 100644
--- a/src/enhanced/utils/queue/index.ts
+++ b/src/enhanced/utils/queue/index.ts
@@ -1,5 +1,7 @@
 import { ReceiveMessageCommandInput, SQS } from '@aws-sdk/client-sqs';
+import crypto from 'crypto';
 import config from 'config';
+import dgram from 'dgram';
 import AWSXRay from 'aws-xray-sdk-core';
 import { JobsQueuePayload } from '../../types/job-types';
 import { isJobQueuePayload } from '../../types/utils/type-guards';
@@ -72,6 +74,38 @@ export async function listenToJobQueue(): Promise<JobsQueuePayload> {
 
     const payload = JSON.parse(message.Body);
 
+    const { xrayTraceId } = payload;
+
+    if (xrayTraceId) {
+      console.log('Xray trace id: ', xrayTraceId);
+      const startTime = Date.now();
+      const traceId = xrayTraceId.split(';')[0];
+      const parentSegment = xrayTraceId.split(';')[1].split('=')[1];
+      const segmentId = crypto.randomBytes(16).toString('hex');
+
+      const newSegment = {
+        name: 'Autobuilder',
+        id: segmentId,
+        trace_id: traceId,
+        parent_id: parentSegment,
+        type: 'subsegment',
+        start_time: startTime,
+        end_time: Date.now(),
+      };
+
+      const client = dgram.createSocket('udp4');
+
+      client.send(JSON.stringify(newSegment), 2000, '127.0.0.1', (err) => {
+        if (err) {
+          console.error('Error occurred when sending udp message to xray daemon', err);
+        }
+
+        client.close();
+      });
+    } else {
+      console.log('no trace id found');
+    }
+
     // Use type guard here to validate payload we have received from the queue.
     // This ensures that the `payload` object will be of type `JobQueuePayload` after the if statement.
     if (!isJobQueuePayload(payload)) {
diff --git a/src/entities/queueMessage.ts b/src/entities/queueMessage.ts
index 28aac0a04..2e7b78e19 100644
--- a/src/entities/queueMessage.ts
+++ b/src/entities/queueMessage.ts
@@ -4,6 +4,7 @@ export class JobQueueMessage {
   jobId: string;
   jobStatus: JobStatus;
   tries: number;
+  xrayTraceId?: string = process.env._X_AMZN_TRACE_ID;
   taskId?: string;
 
   constructor(jobId: string, status: JobStatus, tries = 0, taskId?: string) {