Fix bug with closing xray socket too early (#124)

* Fix bug with closing xray socket too early

* Fix unit tests

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();

src/trace/context.ts

@@ -135,13 +135,12 @@ 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();
+    logDebug("Error occurred submitting to xray daemon", { error });
   }
 }