RD-4242 - handle malformed trace id (#167)

* handle malformed trace id

Author: saartochner-lumigo

src/lumigo_tracer/spans_container.py

@@ -31,6 +31,7 @@
 _VERSION_PATH = os.path.join(os.path.dirname(__file__), "VERSION")
 MAX_LAMBDA_TIME = 15 * 60 * 1000
 FUNCTION_TYPE = "function"
+MALFORMED_TXID = "000000000000000000000000"
 
 
 class SpansContainer:
@@ -62,6 +63,10 @@ def __init__(
         self.region = region
         self.trace_root = trace_root
         self.trace_id_suffix = trace_id_suffix
+        malformed_txid = False
+        if transaction_id == MALFORMED_TXID:
+            transaction_id = os.urandom(12).hex()
+            malformed_txid = True
         self.transaction_id = transaction_id
         self.max_finish_time = max_finish_time
         self.base_msg = {
@@ -88,6 +93,7 @@ def __init__(
                     "logGroupName": log_group_name,
                     **(trigger_by or {}),
                 },
+                "isMalformedTransactionId": malformed_txid,
                 EXECUTION_TAGS_KEY: [],
             },
             self.base_msg,
@@ -260,8 +266,13 @@ def _set_error_extra_data(self, event):
             )
 
     def get_patched_root(self):
+        """
+        We're changing the root in order to pass/share the transaction id. More info:
+        https://docs.aws.amazon.com/xray/latest/devguide/xray-api-sendingdata.html#xray-api-traceids
+        """
+        current_time = int(time.time())
         root = safe_split_get(self.trace_root, "-", 0)
-        return f"Root={root}-0000{os.urandom(2).hex()}-{self.transaction_id}{self.trace_id_suffix}"
+        return f"Root={root}-{hex(current_time)[2:]}-{self.transaction_id}{self.trace_id_suffix}"
 
     @classmethod
     def get_span(cls) -> "SpansContainer":

src/test/unit/test_spans_container.py

@@ -1,11 +1,17 @@
 import copy
 import inspect
 import json
+from datetime import datetime
 
 import pytest
 
 from lumigo_tracer.wrappers.http.http_parser import HTTP_TYPE
-from lumigo_tracer.spans_container import SpansContainer, TimeoutMechanism, FUNCTION_TYPE
+from lumigo_tracer.spans_container import (
+    SpansContainer,
+    TimeoutMechanism,
+    FUNCTION_TYPE,
+    MALFORMED_TXID,
+)
 from lumigo_tracer.lumigo_utils import Configuration, EXECUTION_TAGS_KEY
 
 
@@ -179,3 +185,32 @@ def test_get_span_by_id():
     container.add_span({"id": 3, "extra": "c"})
     assert SpansContainer.get_span().get_span_by_id(2)["extra"] == "b"
     assert SpansContainer.get_span().get_span_by_id(5) is None
+
+
+def test_get_patched_root(monkeypatch, context):
+    monkeypatch.setenv(
+        "_X_AMZN_TRACE_ID",
+        "Root=1-5fd891b8-252f5de90a085ae04267aa4e;Parent=0a885f800de045d4;Sampled=0",
+    )
+    SpansContainer.create_span({}, context)
+    result = SpansContainer.get_span().get_patched_root()
+    root = result.split(";")[0].split("=")[1]
+    one, current_time, txid = root.split("-")
+
+    result_time = datetime.fromtimestamp(int(current_time, 16))
+    assert one == "1"
+    assert (result_time - datetime.now()).total_seconds() < 5
+    assert txid == "252f5de90a085ae04267aa4e"
+
+
+def test_malformed_txid(monkeypatch, context):
+    monkeypatch.setenv(
+        "_X_AMZN_TRACE_ID", f"Root=1-5fd891b8-{MALFORMED_TXID};Parent=0a885f800de045d4;Sampled=0"
+    )
+    SpansContainer.create_span({}, context)
+
+    assert SpansContainer.get_span().transaction_id != MALFORMED_TXID
+    assert SpansContainer.get_span().function_span["isMalformedTransactionId"]
+    result = SpansContainer.get_span().get_patched_root()
+    output_trace_id = result.split(";")[0].split("=")[1].split("-")[2]
+    assert output_trace_id == SpansContainer.get_span().transaction_id