From 5b7839d41d0df5fdd89d07c930cf0a196317e970 Mon Sep 17 00:00:00 2001
From: DarcyRaynerDD <darcy.rayner@datadoghq.com>
Date: Tue, 10 Nov 2020 11:33:21 -0500
Subject: [PATCH 1/2] Fix bug with closing xray socket too early

---
 src/trace/context.ts | 4 +---
 1 file changed, 1 insertion(+), 3 deletions(-)

diff --git a/src/trace/context.ts b/src/trace/context.ts
index 1e36a84a..e3c0340e 100644
--- a/src/trace/context.ts
+++ b/src/trace/context.ts
@@ -135,13 +135,11 @@ export function sendXraySubsegment(segment: string) {
     client = createSocket("udp4");
     // Send segment asynchronously to xray daemon
     client.send(message, 0, message.length, port, address, (error, bytes) => {
+      client?.close();
       logDebug(`Xray daemon received metadata payload`, { error, bytes });
     });
   } catch (error) {
     logDebug("Error occurred submitting to xray daemon", { error });
-  } finally {
-    // Cleanup socket
-    client?.close();
   }
 }
 

From f182d0e17dc2265def004dc8f596211164e2c22a Mon Sep 17 00:00:00 2001
From: DarcyRaynerDD <darcy.rayner@datadoghq.com>
Date: Tue, 10 Nov 2020 14:01:13 -0500
Subject: [PATCH 2/2] Fix unit tests

---
 src/trace/context.spec.ts | 11 ++++++++++-
 src/trace/context.ts      |  1 +
 2 files changed, 11 insertions(+), 1 deletion(-)

diff --git a/src/trace/context.spec.ts b/src/trace/context.spec.ts
index 7674d095..e547851f 100644
--- a/src/trace/context.spec.ts
+++ b/src/trace/context.spec.ts
@@ -25,8 +25,16 @@ jest.mock("dgram", () => {
   return {
     createSocket: () => {
       return {
-        send: (message: string) => {
+        send: (
+          message: string,
+          start: number,
+          length: number,
+          port: number,
+          address: string,
+          callback: (error: string | undefined, bytes: number) => void,
+        ) => {
           sentSegment = message;
+          callback(undefined, 1);
         },
         close: () => {
           closedSocket = true;
@@ -475,6 +483,7 @@ describe("extractTraceContext", () => {
     extractTraceContext(stepFunctionEvent);
 
     expect(sentSegment instanceof Buffer).toBeTruthy();
+
     expect(closedSocket).toBeTruthy();
 
     const sentMessage = sentSegment.toString();
diff --git a/src/trace/context.ts b/src/trace/context.ts
index e3c0340e..c73d6e0b 100644
--- a/src/trace/context.ts
+++ b/src/trace/context.ts
@@ -139,6 +139,7 @@ export function sendXraySubsegment(segment: string) {
       logDebug(`Xray daemon received metadata payload`, { error, bytes });
     });
   } catch (error) {
+    client?.close();
     logDebug("Error occurred submitting to xray daemon", { error });
   }
 }