initial commit

Author: joeyzhao2018

datadog_lambda/patch.py

@@ -31,14 +31,16 @@ def patch_all():
 
     if config.trace_enabled:
         patch_all_dd()
-        # Todo: remove this for PR. This is just a testing helper
-        # Manually patch the durable execution integration since it may not
-        # be registered in the PyPI ddtrace's _monkey.py yet.
+        # Todo: remove this for PR. This is just a testing helper.
+        # Call the aws_durable_execution_sdk_python integration's patch()
+        # directly because PyPI ddtrace's _monkey.py doesn't know about it yet,
+        # so ddtrace.patch(aws_durable_execution_sdk_python=True) would be a no-op.
         try:
-            from ddtrace import patch as _patch_dd
-            _patch_dd(aws_durable_execution_sdk_python=True)
-        except Exception:
-            pass
+            from ddtrace.contrib.internal.aws_durable_execution_sdk_python.patch import patch as _patch_ade
+            _patch_ade()
+            print("[DD-DURABLE] aws_durable_execution_sdk_python integration patched")
+        except Exception as e:
+            print(f"[DD-DURABLE] Failed to patch aws_durable_execution_sdk_python: {e}")
     else:
         _patch_http()
         _ensure_patch_requests()

datadog_lambda/tracing.py

@@ -659,19 +659,19 @@ def is_durable_execution_replay(event):
     return is_replay
 
 
+_TRACE_CHECKPOINT_PREFIX = "_dd_trace_context_"
+
+
 def extract_context_from_durable_execution(event, lambda_context):
     """
     Extract Datadog trace context from AWS Lambda Durable Execution event.
 
-    Looks for extra trace context checkpoints created by the dd-trace plugin.
-    These are STEP operations with Name="_dd_trace_context" that store trace
-    headers in their StepDetails.Result payload. Customer operation payloads
-    are never read or modified.
-
-    Scans operations in reverse to find the LAST trace checkpoint, which
-    corresponds to the most recently completed customer operation. This gives
-    proper parent chaining: each invocation's root span is parented to the
-    last operation span from the previous invocation.
+    Looks for trace-context checkpoints created by the dd-trace-py integration
+    on previous invocations.  Checkpoints are STEP operations named
+    ``_dd_trace_context_{N}`` with trace headers stored as their StepDetails.Result
+    payload.  The one with the highest ``{N}`` wins — it corresponds to the
+    latest trace-context state from the previous invocation.  Customer
+    operation payloads are never read or modified.
     """
     try:
         if not isinstance(event, dict):
@@ -687,54 +687,59 @@ def extract_context_from_durable_execution(event, lambda_context):
 
         print(f"[DD-DURABLE] Found {len(operations)} operations in InitialExecutionState")
 
-        # Scan in reverse to find the LAST trace context checkpoint
-        # (corresponds to the most recently completed customer operation)
-        for idx in range(len(operations) - 1, -1, -1):
-            operation = operations[idx]
+        # Collect all _dd_trace_context_{N} checkpoints, then pick the highest N
+        candidates = []  # list of (number, operation)
+        for operation in operations:
             op_name = operation.get("Name")
-
-            if op_name != "_dd_trace_context":
+            if not op_name or not op_name.startswith(_TRACE_CHECKPOINT_PREFIX):
                 continue
+            suffix = op_name[len(_TRACE_CHECKPOINT_PREFIX):]
+            try:
+                number = int(suffix)
+            except ValueError:
+                continue
+            candidates.append((number, operation))
 
-            operation_id = operation.get("Id")
-            print(f"[DD-DURABLE] Found trace checkpoint: id={operation_id}, index={idx}")
+        if not candidates:
+            print("[DD-DURABLE] No trace context checkpoints found in operations")
+            return None
 
-            # Trace context is in StepDetails.Result (standard STEP format)
-            step_details = operation.get("StepDetails", {})
-            payload_str = step_details.get("Result")
+        candidates.sort(key=lambda t: t[0])
+        number, operation = candidates[-1]
+        operation_id = operation.get("Id")
+        op_name = operation.get("Name")
+        print(f"[DD-DURABLE] Using latest trace checkpoint: name={op_name}, id={operation_id}")
 
-            if not payload_str:
-                print(f"[DD-DURABLE] Trace checkpoint {operation_id} has no Result, skipping")
-                continue
+        step_details = operation.get("StepDetails", {})
+        payload_str = step_details.get("Result")
+        if not payload_str:
+            print(f"[DD-DURABLE] Trace checkpoint {op_name} has no Result, skipping")
+            return None
 
-            try:
-                payload = json.loads(payload_str)
-                if not isinstance(payload, dict):
-                    print(f"[DD-DURABLE] Trace checkpoint payload is not a dict: {type(payload)}")
-                    continue
-
-                trace_id = payload.get("x-datadog-trace-id")
-                span_id = payload.get("x-datadog-parent-id")
-
-                if trace_id and span_id:
-                    # Use HTTPPropagator to restore full context including
-                    # baggage, _dd.p.* tags, origin, and sampling priority
-                    context = propagator.extract(payload)
-                    if context and context.trace_id:
-                        print(f"[DD-DURABLE] Extracted trace context from trace checkpoint {operation_id}")
-                        print(f"[DD-DURABLE]   trace_id={trace_id}, span_id={span_id}, headers={list(payload.keys())}")
-                        logger.debug(
-                            "Extracted Datadog trace context from trace checkpoint %s: %s",
-                            operation_id,
-                            context,
-                        )
-                        return context
-            except (json.JSONDecodeError, TypeError, ValueError) as e:
-                print(f"[DD-DURABLE] Failed to parse trace checkpoint payload: {e}")
-                logger.debug("Failed to parse trace checkpoint payload: %s", e)
-                continue
+        try:
+            payload = json.loads(payload_str)
+        except (json.JSONDecodeError, TypeError, ValueError) as e:
+            print(f"[DD-DURABLE] Failed to parse trace checkpoint payload: {e}")
+            logger.debug("Failed to parse trace checkpoint payload: %s", e)
+            return None
+
+        if not isinstance(payload, dict):
+            print(f"[DD-DURABLE] Trace checkpoint payload is not a dict: {type(payload)}")
+            return None
 
-        print("[DD-DURABLE] No trace context checkpoints found in operations")
+        context = propagator.extract(payload)
+        if context and context.trace_id:
+            print(
+                f"[DD-DURABLE] Extracted trace context from {op_name}: "
+                f"trace_id={context.trace_id}, span_id={context.span_id}, "
+                f"headers={list(payload.keys())}"
+            )
+            logger.debug(
+                "Extracted Datadog trace context from trace checkpoint %s: %s",
+                op_name,
+                context,
+            )
+            return context
     except Exception as e:
         logger.debug("Failed to extract trace context from durable execution: %s", e)
 
@@ -1161,12 +1166,16 @@ def process_injected_data(event, request_time_epoch_ms, args, tags):
             start_time_ns = int(
                 injected_authorizer_data.get(Headers.Parent_Span_Finish_Time)
             )
-            integration_latency = int(
-                event["requestContext"]["authorizer"].get("integrationLatency", 0)
-            )
-            finish_time_ns = max(
-                start_time_ns, (request_time_epoch_ms + integration_latency) * 1e6
-            )
+            finish_time_ns = (
+                request_time_epoch_ms
+                + (
+                    int(
+                        event["requestContext"]["authorizer"].get(
+                            "integrationLatency", 0
+                        )
+                    )
+                )
+            ) * 1e6
             upstream_authorizer_span = insert_upstream_authorizer_span(
                 args, tags, start_time_ns, finish_time_ns
             )
@@ -1629,9 +1638,9 @@ def create_function_execution_span(
     trace_context_source,
     merge_xray_traces,
     trigger_tags,
-    durable_function_tags=None,
     parent_span=None,
     span_pointers=None,
+    durable_function_tags=None,
 ):
     tags = None
     if context:
@@ -1640,7 +1649,6 @@ def create_function_execution_span(
         function_arn = ":".join(tk[0:7]) if len(tk) > 7 else function_arn
         function_version = tk[7] if len(tk) > 7 else "$LATEST"
         tags = {
-            "span.kind": "server",
             "cold_start": str(is_cold_start).lower(),
             "function_arn": function_arn,
             "function_version": function_version,

datadog_lambda/wrapper.py

@@ -153,6 +153,7 @@ def __init__(self, func):
             self.blocking_response = None
             self.durable_root_span = None
             self.is_durable = False
+            self.durable_status = None
 
             if config.profiling_enabled and profiler:
                 self.prof = profiler.Profiler(env=config.env, service=config.service)
@@ -307,14 +308,6 @@ def _before(self, event, context):
                     # Component 4: On replays, returns None (context from checkpoint)
                     self.durable_root_span = create_durable_execution_root_span(event)
 
-                    # Store root span reference for Component 2 (checkpoint save)
-                    if self.durable_root_span:
-                        try:
-                            from ddtrace.contrib.internal.aws_durable_execution_sdk_python.patch import set_durable_root_span
-                            set_durable_root_span(self.durable_root_span)
-                        except Exception:
-                            pass
-
                 # Create aws.lambda span — child of root durable span (first invocation)
                 # or child of checkpoint context (replay), or normal parent otherwise
                 self.span = create_function_execution_span(
@@ -333,6 +326,7 @@ def _before(self, event, context):
                 if config.appsec_enabled:
                     asm_start_request(self.span, event, event_source, self.trigger_tags)
                     self.blocking_response = get_asm_blocked_response(self.event_source)
+
             else:
                 set_correlation_ids()
             if config.profiling_enabled and profiler and is_new_sandbox():
@@ -359,6 +353,11 @@ def _after(self, event, context):
             if should_trace_cold_start:
                 trace_ctx = tracer.current_trace_context()
 
+            if self.is_durable:
+                self.durable_status = extract_durable_execution_status(
+                    self.response, event
+                )
+
             if self.span:
                 if config.appsec_enabled and not self.blocking_response:
                     asm_start_response(
@@ -384,45 +383,18 @@ def _after(self, event, context):
                 if status_code:
                     self.span.set_tag("http.status_code", status_code)
 
-                durable_status = extract_durable_execution_status(self.response, event)
-                if durable_status:
+                if self.durable_status:
                     self.span.set_tag(
                         "aws_lambda.durable_function.execution_status",
-                        durable_status,
+                        self.durable_status,
                     )
 
                 self.span.finish()
 
-            # Component 3: After aws.lambda closes but BEFORE root span closes,
-            # check if trace context was enriched during this invocation.
-            # The context is still alive because _reactivate=True on the parent context.
-            # This is best-effort for the end-of-invocation case; the piggyback in
-            # _patched_create_checkpoint handles saves during handler execution.
-            if self.is_durable:
-                try:
-                    from ddtrace.contrib.internal.aws_durable_execution_sdk_python.patch import (
-                        _has_context_changed, save_updated_trace_checkpoint,
-                        _get_current_execution_state, _inject_current_context,
-                    )
-                    state = _get_current_execution_state()
-                    if state:
-                        before_headers = getattr(state, "_dd_before_trace_headers", None)
-                        if before_headers is not None and _has_context_changed(before_headers):
-                            print("[DD-DURABLE] Context changed at end of invocation, saving updated checkpoint")
-                            state._dd_saving_trace_checkpoint = True
-                            try:
-                                save_updated_trace_checkpoint(state)
-                            finally:
-                                state._dd_saving_trace_checkpoint = False
-                            current = _inject_current_context()
-                            if current:
-                                state._dd_before_trace_headers = current
-                except Exception as e:
-                    print(f"[DD-DURABLE] Best-effort context change check in _after: {e}")
-
-            # Finish durable execution root span LAST (Component 1)
-            # Component 2 (root checkpoint) is handled in _patched_create_checkpoint
-            # piggybacked on the first operation's checkpoint call.
+            # Finish durable execution root span LAST, after aws.lambda.
+            # The trace-context checkpoint (for cross-invocation continuity)
+            # is saved by dd-trace-py's aws_durable_execution_sdk_python
+            # integration when the aws.durable_execution.execute span closes.
             if self.durable_root_span:
                 self.durable_root_span.finish()
                 print(f"[DD-DURABLE] Finished root span: trace_id={self.durable_root_span.trace_id}, span_id={self.durable_root_span.span_id}")