feat: lambda support for DSM

datadog_lambda/tracing.py

@@ -67,6 +67,24 @@
 LOWER_64_BITS = "LOWER_64_BITS"
 
 
+def _dsm_set_checkpoint(context_json, event_type, arn):
+    if not config.data_streams_enabled:
+        return
+
+    if not arn:
+        return
+
+    try:
+        from ddtrace.data_streams import set_consume_checkpoint
+
+        carrier_get = lambda k: context_json and context_json.get(k)  # noqa: E731
+        set_consume_checkpoint(event_type, arn, carrier_get, manual_checkpoint=False)
+    except Exception as e:
+        logger.debug(
+            f"DSM:Failed to set consume checkpoint for {event_type} {arn}: {e}"
+        )
+
+
 def _convert_xray_trace_id(xray_trace_id):
     """
     Convert X-Ray trace id (hex)'s last 63 bits to a Datadog trace id (int).
@@ -202,7 +220,9 @@ def create_sns_event(message):
     }
 
 
-def extract_context_from_sqs_or_sns_event_or_context(event, lambda_context):
+def extract_context_from_sqs_or_sns_event_or_context(
+    event, lambda_context, event_source
+):
     """
     Extract Datadog trace context from an SQS event.
 
@@ -214,7 +234,10 @@ def extract_context_from_sqs_or_sns_event_or_context(event, lambda_context):
     Lambda Context.
 
     Falls back to lambda context if no trace data is found in the SQS message attributes.
+    Set a DSM checkpoint if DSM is enabled and the method for context propagation is supported.
     """
+    source_arn = ""
+    event_type = "sqs" if event_source.equals(EventTypes.SQS) else "sns"
 
     # EventBridge => SQS
     try:
@@ -226,6 +249,7 @@ def extract_context_from_sqs_or_sns_event_or_context(event, lambda_context):
 
     try:
         first_record = event.get("Records")[0]
+        source_arn = first_record.get("eventSourceARN", "")
 
         # logic to deal with SNS => SQS event
         if "body" in first_record:
@@ -241,6 +265,9 @@ def extract_context_from_sqs_or_sns_event_or_context(event, lambda_context):
         msg_attributes = first_record.get("messageAttributes")
         if msg_attributes is None:
             sns_record = first_record.get("Sns") or {}
+            # SNS->SQS event would extract SNS arn without this check
+            if event_source.equals(EventTypes.SNS):
+                source_arn = sns_record.get("TopicArn", "")
             msg_attributes = sns_record.get("MessageAttributes") or {}
         dd_payload = msg_attributes.get("_datadog")
         if dd_payload:
@@ -272,8 +299,9 @@ def extract_context_from_sqs_or_sns_event_or_context(event, lambda_context):
                         logger.debug(
                             "Failed to extract Step Functions context from SQS/SNS event."
                         )
-
-                return propagator.extract(dd_data)
+                context = propagator.extract(dd_data)
+                _dsm_set_checkpoint(dd_data, event_type, source_arn)
+                return context
         else:
             # Handle case where trace context is injected into attributes.AWSTraceHeader
             # example: Root=1-654321ab-000000001234567890abcdef;Parent=0123456789abcdef;Sampled=1
@@ -296,9 +324,13 @@ def extract_context_from_sqs_or_sns_event_or_context(event, lambda_context):
                             span_id=int(x_ray_context["parent_id"], 16),
                             sampling_priority=float(x_ray_context["sampled"]),
                         )
+        # Still want to set a DSM checkpoint even if DSM context not propagated
+        _dsm_set_checkpoint(None, event_type, source_arn)
         return extract_context_from_lambda_context(lambda_context)
     except Exception as e:
         logger.debug("The trace extractor returned with error %s", e)
+        # Still want to set a DSM checkpoint even if DSM context not propagated
+        _dsm_set_checkpoint(None, event_type, source_arn)
         return extract_context_from_lambda_context(lambda_context)
 
 
@@ -357,9 +389,12 @@ def extract_context_from_eventbridge_event(event, lambda_context):
 def extract_context_from_kinesis_event(event, lambda_context):
     """
     Extract datadog trace context from a Kinesis Stream's base64 encoded data string
+    Set a DSM checkpoint if DSM is enabled and the method for context propagation is supported.
     """
+    source_arn = ""
     try:
         record = get_first_record(event)
+        source_arn = record.get("eventSourceARN", "")
         kinesis = record.get("kinesis")
         if not kinesis:
             return extract_context_from_lambda_context(lambda_context)
@@ -373,10 +408,13 @@ def extract_context_from_kinesis_event(event, lambda_context):
             data_obj = json.loads(data_str)
             dd_ctx = data_obj.get("_datadog")
             if dd_ctx:
-                return propagator.extract(dd_ctx)
+                context = propagator.extract(dd_ctx)
+                _dsm_set_checkpoint(dd_ctx, "kinesis", source_arn)
+                return context
     except Exception as e:
         logger.debug("The trace extractor returned with error %s", e)
-
+    # Still want to set a DSM checkpoint even if DSM context not propagated
+    _dsm_set_checkpoint(None, "kinesis", source_arn)
     return extract_context_from_lambda_context(lambda_context)
 
 
@@ -594,7 +632,7 @@ def extract_dd_trace_context(
         )
     elif event_source.equals(EventTypes.SNS) or event_source.equals(EventTypes.SQS):
         context = extract_context_from_sqs_or_sns_event_or_context(
-            event, lambda_context
+            event, lambda_context, event_source
         )
     elif event_source.equals(EventTypes.EVENTBRIDGE):
         context = extract_context_from_eventbridge_event(event, lambda_context)

datadog_lambda/wrapper.py

@@ -10,7 +10,6 @@
 from time import time_ns
 
 from datadog_lambda.asm import asm_start_response, asm_start_request
-from datadog_lambda.dsm import set_dsm_context
 from datadog_lambda.extension import should_use_extension, flush_extension
 from datadog_lambda.cold_start import (
     set_cold_start,
@@ -237,8 +236,6 @@ def _before(self, event, context):
                     self.inferred_span = create_inferred_span(
                         event, context, event_source, config.decode_authorizer_context
                     )
-                if config.data_streams_enabled:
-                    set_dsm_context(event, event_source)
                 self.span = create_function_execution_span(
                     context=context,
                     function_name=config.function_name,