Generate root trace id if it doesn't exist

joe4dev · joe4dev · commit d35c6d4635a0 · 2024-10-18T16:17:54.000+02:00
diff --git a/localstack-core/localstack/aws/handlers/tracing.py b/localstack-core/localstack/aws/handlers/tracing.py
@@ -15,6 +15,7 @@ def __call__(self, chain: HandlerChain, context: RequestContext, response: Respo
         trace_header_str = context.request.headers.get("X-Amzn-Trace-Id")
         # Naming aws_trace_header inspired by AWSTraceHeader convention for SQS:
         # https://docs.aws.amazon.com/xray/latest/devguide/xray-services-sqs.html
-        context.trace_context["aws_trace_header"] = TraceHeader.from_header_str(trace_header_str)
+        aws_trace_header = TraceHeader.from_header_str(trace_header_str).ensure_root_exists()
+        context.trace_context["aws_trace_header"] = aws_trace_header
         # NOTE: X-Ray sampling might require service-specific decisions:
         #  https://docs.aws.amazon.com/xray/latest/devguide/xray-console-sampling.html
diff --git a/localstack-core/localstack/services/lambda_/invocation/event_manager.py b/localstack-core/localstack/services/lambda_/invocation/event_manager.py
@@ -59,7 +59,7 @@ class SQSInvocation:
 
     def encode(self) -> str:
         # Encode TraceHeader as string
-        aws_trace_header = self.invocation.trace_context.get("aws_trace_header") or TraceHeader()
+        aws_trace_header = self.invocation.trace_context.get("aws_trace_header")
         aws_trace_header_str = aws_trace_header.to_header_str()
         self.invocation.trace_context["aws_trace_header"] = aws_trace_header_str
         return json.dumps(
diff --git a/localstack-core/localstack/services/lambda_/invocation/execution_environment.py b/localstack-core/localstack/services/lambda_/invocation/execution_environment.py
@@ -8,7 +8,6 @@
 from typing import Callable, Dict, Optional
 
 from localstack import config
-from localstack.aws.api.lambda_ import TracingMode
 from localstack.aws.connect import connect_to
 from localstack.services.lambda_.invocation.lambda_models import (
     Credentials,
@@ -26,7 +25,6 @@
     is_lambda_debug_timeout_enabled_for,
 )
 from localstack.utils.strings import to_str
-from localstack.utils.xray.trace_header import TraceHeader
 
 STARTUP_TIMEOUT_SEC = config.LAMBDA_RUNTIME_ENVIRONMENT_TIMEOUT
 HEX_CHARS = [str(num) for num in range(10)] + ["a", "b", "c", "d", "e", "f"]
@@ -342,11 +340,8 @@ def get_prefixed_logs(self) -> str:
 
     def invoke(self, invocation: Invocation) -> InvocationResult:
         assert self.status == RuntimeStatus.RUNNING
-        aws_trace_header = invocation.trace_context.get("aws_trace_header") or TraceHeader()
-        # TODO: test and implement passive tracing and sampling decisions. Using a simple heuristic for now:
-        #  If TracingMode is "Active", we always sample the request.
-        if self.function_version.config.tracing_config_mode == TracingMode.Active:
-            aws_trace_header.sampled = 1
+        aws_trace_header = invocation.trace_context.get("aws_trace_header")
+        # TODO: test and implement Active and PassThrough tracing and sampling decisions.
         # TODO: implement Lambda lineage: https://docs.aws.amazon.com/lambda/latest/dg/invocation-recursion.html
         invoke_payload = {
             "invoke-id": invocation.request_id,  # TODO: rename to request-id (requires change in lambda-init)
diff --git a/localstack-core/localstack/utils/xray/trace_header.py b/localstack-core/localstack/utils/xray/trace_header.py
@@ -2,10 +2,14 @@
 # It is adapted from aws-xray-sdk-python licensed under the Apache License 2.0.
 # You may obtain a copy of the Apache License 2.0 at http://www.apache.org/licenses/LICENSE-2.0
 # Original source: https://github.com/aws/aws-xray-sdk-python/blob/master/aws_xray_sdk/core/models/trace_header.py
-# Modifications: Add optional lineage field for https://docs.aws.amazon.com/lambda/latest/dg/invocation-recursion.html
+# Modifications:
+# * Add optional lineage field for https://docs.aws.amazon.com/lambda/latest/dg/invocation-recursion.html
+# * Add ensure_root_exists() to generate root trace id
 
 import logging
 
+from localstack.utils.xray.traceid import TraceId
+
 log = logging.getLogger(__name__)
 
 ROOT = "Root"
@@ -103,6 +107,15 @@ def to_header_str(self):
 
         return HEADER_DELIMITER.join(h_parts)
 
+    def ensure_root_exists(self):
+        """
+        Ensures that a root trace id exists by generating one if None.
+        Return self to allow for chaining.
+        """
+        if self._root is None:
+            self._root = TraceId().to_id()
+        return self
+
     @property
     def root(self):
         """
