Cleans up socket after sending payload to x-ray (#114)

* Cleans up socket after sending payload to x-ray

* Fix code formatting

src/trace/context.spec.ts

@@ -19,6 +19,7 @@ import {
 } from "./context";
 
 let sentSegment: any;
+let closedSocket = false;
 
 jest.mock("dgram", () => {
   return {
@@ -27,6 +28,9 @@ jest.mock("dgram", () => {
         send: (message: string) => {
           sentSegment = message;
         },
+        close: () => {
+          closedSocket = true;
+        },
       };
     },
   };
@@ -39,6 +43,7 @@ jest.mock("crypto", () => {
 
 beforeEach(() => {
   sentSegment = undefined;
+  closedSocket = false;
   setLogLevel(LogLevel.NONE);
 });
 
@@ -424,6 +429,7 @@ describe("extractTraceContext", () => {
     });
 
     expect(sentSegment instanceof Buffer).toBeTruthy();
+    expect(closedSocket).toBeTruthy();
     const sentMessage = sentSegment.toString();
     expect(sentMessage).toMatchInlineSnapshot(`
       "{\\"format\\": \\"json\\", \\"version\\": 1}
@@ -469,6 +475,8 @@ describe("extractTraceContext", () => {
     extractTraceContext(stepFunctionEvent);
 
     expect(sentSegment instanceof Buffer).toBeTruthy();
+    expect(closedSocket).toBeTruthy();
+
     const sentMessage = sentSegment.toString();
     expect(sentMessage).toMatchInlineSnapshot(`
       "{\\"format\\": \\"json\\", \\"version\\": 1}

src/trace/context.ts

@@ -1,6 +1,6 @@
 import { BigNumber } from "bignumber.js";
 import { randomBytes } from "crypto";
-import { createSocket } from "dgram";
+import { createSocket, Socket } from "dgram";
 
 import { logDebug, logError } from "../utils";
 import {
@@ -130,14 +130,18 @@ export function sendXraySubsegment(segment: string) {
   const address = parts[0];
 
   const message = new Buffer(`{\"format\": \"json\", \"version\": 1}\n${segment}`);
+  let client: Socket | undefined;
   try {
-    const client = createSocket("udp4");
+    client = createSocket("udp4");
     // Send segment asynchronously to xray daemon
     client.send(message, 0, message.length, port, address, (error, bytes) => {
       logDebug(`Xray daemon received metadata payload`, { error, bytes });
     });
   } catch (error) {
     logDebug("Error occurred submitting to xray daemon", { error });
+  } finally {
+    // Cleanup socket
+    client?.close();
   }
 }