Use single socket instance for all xray api calls.

datadog_lambda/xray.py

@@ -10,32 +10,45 @@
 logger = logging.getLogger(__name__)
 
 
-def get_xray_host_port(address):
-    if address == "":
-        logger.debug("X-Ray daemon env var not set, not sending sub-segment")
-        return None
-    parts = address.split(":")
-    if len(parts) <= 1:
-        logger.debug("X-Ray daemon env var not set, not sending sub-segment")
-        return None
-    port = int(parts[1])
-    host = parts[0]
-    return (host, port)
-
-
-def send(host_port_tuple, payload):
-    sock = None
-    try:
-        sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
-        sock.setblocking(0)
-        sock.connect(host_port_tuple)
-        sock.send(payload.encode("utf-8"))
-    except Exception as e_send:
-        logger.error("Error occurred submitting to xray daemon: %s", e_send)
-    try:
-        sock.close()
-    except Exception as e_close:
-        logger.error("Error while closing the socket: %s", e_close)
+class Socket(object):
+    def __init__(self):
+        self.sock = None
+
+    @property
+    def host_port_tuple(self):
+        if not hasattr(self, "_host_port_tuple"):
+            self._host_port_tuple = self._get_xray_host_port(
+                os.environ.get(XrayDaemon.XRAY_DAEMON_ADDRESS, "")
+            )
+        return self._host_port_tuple
+
+    def send(self, payload):
+        if not self.sock:
+            self._connect()
+        try:
+            self.sock.send(payload.encode("utf-8"))
+        except Exception as e_send:
+            logger.error("Error occurred submitting to xray daemon: %s", e_send)
+
+    def _connect(self):
+        self.sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
+        self.sock.setblocking(0)
+        self.sock.connect(self.host_port_tuple)
+
+    def _get_xray_host_port(self, address):
+        if address == "":
+            logger.debug("X-Ray daemon env var not set, not sending sub-segment")
+            return None
+        parts = address.split(":")
+        if len(parts) <= 1:
+            logger.debug("X-Ray daemon env var not set, not sending sub-segment")
+            return None
+        port = int(parts[1])
+        host = parts[0]
+        return (host, port)
+
+
+sock = Socket()
 
 
 def build_segment_payload(payload):
@@ -95,10 +108,7 @@ def build_segment(context, key, metadata):
 
 
 def send_segment(key, metadata):
-    host_port_tuple = get_xray_host_port(
-        os.environ.get(XrayDaemon.XRAY_DAEMON_ADDRESS, "")
-    )
-    if host_port_tuple is None:
+    if sock.host_port_tuple is None:
         return None
     context = parse_xray_header(
         os.environ.get(XrayDaemon.XRAY_TRACE_ID_HEADER_NAME, "")
@@ -115,4 +125,4 @@ def send_segment(key, metadata):
         return None
     segment = build_segment(context, key, metadata)
     segment_payload = build_segment_payload(segment)
-    send(host_port_tuple, segment_payload)
+    sock.send(segment_payload)

tests/test_benchmarks.py

@@ -13,7 +13,7 @@
 
 from datadog_lambda.constants import XrayDaemon, XraySubsegment
 
-from tests.utils import get_mock_context
+from tests.utils import get_mock_context, reset_xray_connection
 
 
 event_samples_dir = "tests/event_samples"
@@ -74,14 +74,24 @@ def test_trigger_extract_trigger_tags(event, benchmark):
 
 
 def test_xray_send_segment(benchmark, monkeypatch):
+    reset_xray_connection()
+
     monkeypatch.setenv(XrayDaemon.XRAY_DAEMON_ADDRESS, "localhost:9000")
     monkeypatch.setenv(
         XrayDaemon.XRAY_TRACE_ID_HEADER_NAME,
         "Root=1-5e272390-8c398be037738dc042009320;Parent=94ae789b969f1cc5;Sampled=1;Lineage=c6c5b1b9:0",
     )
+
+    def socket_send(*a, **k):
+        sends.append(True)
+
+    sends = []
+    monkeypatch.setattr("socket.socket.send", socket_send)
+
     key = {
         "trace-id": "12345678901234567890123456789012",
         "parent-id": "1234567890123456",
         "sampling-priority": "1",
     }
     benchmark(xray.send_segment, XraySubsegment.TRACE_KEY, key)
+    assert sends

tests/test_wrapper.py

@@ -7,12 +7,13 @@
 from datadog_lambda.constants import TraceHeader
 
 import datadog_lambda.wrapper as wrapper
+import datadog_lambda.xray as xray
 from datadog_lambda.metric import lambda_metric
 from datadog_lambda.thread_stats_writer import ThreadStatsWriter
 from ddtrace import Span, tracer
 from ddtrace.internal.constants import MAX_UINT_64BITS
 
-from tests.utils import get_mock_context
+from tests.utils import get_mock_context, reset_xray_connection
 
 
 class TestDatadogLambdaWrapper(unittest.TestCase):
@@ -590,7 +591,9 @@ class TestLambdaWrapperWithTraceContext(unittest.TestCase):
         },
     )
     def test_event_bridge_sqs_payload(self):
-        patcher = patch("datadog_lambda.xray.send")
+        reset_xray_connection()
+
+        patcher = patch("datadog_lambda.xray.sock.send")
         mock_send = patcher.start()
         self.addCleanup(patcher.stop)
 
@@ -623,7 +626,7 @@ def handler(event, context):
         self.assertEqual(result.span_id, aws_lambda_span.span_id)
         self.assertEqual(result.sampling_priority, 1)
         mock_send.assert_called_once()
-        (_, raw_payload), _ = mock_send.call_args
+        (raw_payload,), _ = mock_send.call_args
         payload = json.loads(raw_payload[33:])  # strip formatting prefix
         self.assertEqual(self.xray_root, payload["trace_id"])
         self.assertEqual(self.xray_parent, payload["parent_id"])

tests/test_xray.py

@@ -4,15 +4,14 @@
 
 from unittest.mock import MagicMock, patch
 
-from datadog_lambda.xray import (
-    get_xray_host_port,
-    build_segment_payload,
-    build_segment,
-    send_segment,
-)
+from datadog_lambda.xray import build_segment_payload, build_segment, send_segment, sock
+from tests.utils import reset_xray_connection
 
 
 class TestXRay(unittest.TestCase):
+    def setUp(self):
+        reset_xray_connection()
+
     def tearDown(self):
         if os.environ.get("_X_AMZN_TRACE_ID"):
             os.environ.pop("_X_AMZN_TRACE_ID")
@@ -21,15 +20,15 @@ def tearDown(self):
         return super().tearDown()
 
     def test_get_xray_host_port_empty_(self):
-        result = get_xray_host_port("")
+        result = sock._get_xray_host_port("")
         self.assertIsNone(result)
 
     def test_get_xray_host_port_invalid_value(self):
-        result = get_xray_host_port("myVar")
+        result = sock._get_xray_host_port("myVar")
         self.assertIsNone(result)
 
     def test_get_xray_host_port_success(self):
-        result = get_xray_host_port("mySuperHost:1000")
+        result = sock._get_xray_host_port("mySuperHost:1000")
         self.assertEqual("mySuperHost", result[0])
         self.assertEqual(1000, result[1])
 
@@ -40,7 +39,7 @@ def test_send_segment_sampled_out(self):
         ] = "Root=1-5e272390-8c398be037738dc042009320;Parent=94ae789b969f1cc5;Sampled=0;Lineage=c6c5b1b9:0"
 
         with patch(
-            "datadog_lambda.xray.send", MagicMock(return_value=None)
+            "datadog_lambda.xray.sock.send", MagicMock(return_value=None)
         ) as mock_send:
             # XRay trace won't be sampled according to the trace header.
             send_segment("my_key", {"data": "value"})
@@ -52,7 +51,7 @@ def test_send_segment_sampled(self):
             "_X_AMZN_TRACE_ID"
         ] = "Root=1-5e272390-8c398be037738dc042009320;Parent=94ae789b969f1cc5;Sampled=1;Lineage=c6c5b1b9:0"
         with patch(
-            "datadog_lambda.xray.send", MagicMock(return_value=None)
+            "datadog_lambda.xray.sock.send", MagicMock(return_value=None)
         ) as mock_send:
             # X-Ray trace will be sampled according to the trace header.
             send_segment("my_key", {"data": "value"})

tests/utils.py

@@ -24,3 +24,13 @@ def get_mock_context(
     lambda_context.function_name = function_name
     lambda_context.client_context = ClientContext(custom)
     return lambda_context
+
+
+def reset_xray_connection():
+    from datadog_lambda.xray import sock
+
+    if hasattr(sock, "_host_port_tuple"):
+        del sock._host_port_tuple
+    if sock.sock:
+        sock.sock.close()
+        sock.sock = None