Infer API Gateway spans

README.md

@@ -87,6 +87,17 @@ Initialize the Datadog tracer when set to `true`. Defaults to `false`.
 
 Set to `true` to merge the X-Ray trace and the Datadog trace, when using both the X-Ray and Datadog tracing. Defaults to `false`.
 
+### DD_INFERRED_SPANS (experimental)
+
+Inferred Spans are spans that Datadog can create based on incoming event metadata. 
+Set `DD_INFERRED_SPANS` to `true` to infer spans based on Lambda events.
+Inferring upstream spans is only supported if you are using the [Datadog Lambda Extension](https://docs.datadoghq.com/serverless/libraries_integrations/extension/). 
+Defaults to `false`.
+Infers spans for:
+- API Gateway REST events
+- API Gateway websocket events
+- HTTP API events
+
 ## Opening Issues
 
 If you encounter a bug with this package, we want to hear about it. Before opening a new issue, search the existing issues to avoid duplicates.

datadog_lambda/tracing.py

@@ -6,6 +6,7 @@
 import logging
 import os
 import json
+from enum import Enum
 
 from datadog_lambda.constants import (
     SamplingPriority,
@@ -19,17 +20,48 @@
 )
 from ddtrace import tracer, patch
 from ddtrace import __version__ as ddtrace_version
+from ddtrace.filters import TraceFilter
 from ddtrace.propagation.http import HTTPPropagator
 from datadog_lambda import __version__ as datadog_lambda_version
 
 logger = logging.getLogger(__name__)
 
+
+SPAN_TYPE_TAG = "_dd.span_type"
+SPAN_TYPE_INFERRED = "inferred"
+
+
+class InferredSpanFilter(TraceFilter):
+    def process_trace(self, trace):
+        logger.debug("InferredSpanFilter got a trace of length {}".format(len(trace)))
+        trace_to_send = []
+        for span in trace:
+            if span.get_tag(SPAN_TYPE_TAG) == SPAN_TYPE_INFERRED and len(trace) > 1:
+                logger.debug(
+                    "Found an inferred span. Filtering it out and writing it separately."
+                )
+                tracer.write([span])
+            else:
+                logger.debug("Appending a span to the returned trace")
+                trace_to_send.append(span)
+        return trace_to_send
+
+
+tracer.configure(settings={"FILTERS": [InferredSpanFilter()]})
 dd_trace_context = {}
 dd_tracing_enabled = os.environ.get("DD_TRACE_ENABLED", "false").lower() == "true"
 
 propagator = HTTPPropagator()
 
 
+class ManagedService(Enum):
+    UNKNOWN = 0
+    API_GATEWAY = 1
+    API_GATEWAY_WEBSOCKET = 2
+    HTTP_API = 3
+    APPSYNC = 4
+
+
 def _convert_xray_trace_id(xray_trace_id):
     """
     Convert X-Ray trace id (hex)'s last 63 bits to a Datadog trace id (int).
@@ -377,13 +409,125 @@ def set_dd_trace_py_root(trace_context_source, merge_xray_traces):
         )
 
 
+def create_inferred_span(event, context, function_name):
+    managed_service = detect_inferrable_span_type(event)
+    try:
+        if managed_service == ManagedService.API_GATEWAY:
+            logger.debug("API Gateway event detected. Inferring a span")
+            return create_inferred_span_from_api_gateway_event(event, context)
+        elif managed_service == ManagedService.HTTP_API:
+            logger.debug("HTTP API event detected. Inferring a span")
+            return create_inferred_span_from_http_api_event(event, context)
+        elif managed_service == ManagedService.API_GATEWAY_WEBSOCKET:
+            logger.debug("API Gateway Websocket event detected. Inferring a span")
+            return create_inferred_span_from_api_gateway_websocket_event(event, context)
+    except Exception as e:
+        logger.debug(
+            "Unable to infer span. Detected type: {}. Reason: {}", managed_service, e
+        )
+        return None
+    logger.debug("Unable to infer a span: unknown event type")
+    return None
+
+
+def detect_inferrable_span_type(event):
+    if "httpMethod" in event:  # likely some kind of API Gateway event
+        return ManagedService.API_GATEWAY
+    if "routeKey" in event:  # likely HTTP API
+        return ManagedService.HTTP_API
+    if (
+        "requestContext" in event and "messageDirection" in event["requestContext"]
+    ):  # likely a websocket API
+        return ManagedService.API_GATEWAY_WEBSOCKET
+    return ManagedService.UNKNOWN
+
+
+def create_inferred_span_from_api_gateway_websocket_event(event, context):
+    domain = event["requestContext"]["domainName"]
+    endpoint = event["requestContext"]["routeKey"]
+    tags = {
+        "operation_name": "aws.apigateway.websocket",
+        "service.name": domain,
+        "url": domain + endpoint,
+        "endpoint": endpoint,
+        "resource_name": domain + endpoint,
+        "request_id": context.aws_request_id,
+        "connection_id": event["requestContext"]["connectionId"],
+        SPAN_TYPE_TAG: SPAN_TYPE_INFERRED,
+    }
+    request_time_epoch = event["requestContext"]["requestTimeEpoch"]
+    args = {
+        "resource": domain + endpoint,
+        "span_type": "web",
+    }
+    tracer.set_tags({"_dd.origin": "lambda"})
+    span = tracer.trace("aws.apigateway.websocket", **args)
+    if span:
+        span.set_tags(tags)
+    span.start = request_time_epoch / 1000
+    return span
+
+
+def create_inferred_span_from_api_gateway_event(event, context):
+    domain = event["requestContext"]["domainName"]
+    path = event["path"]
+    tags = {
+        "operation_name": "aws.apigateway.rest",
+        "service.name": domain,
+        "url": domain + path,
+        "endpoint": path,
+        "http.method": event["httpMethod"],
+        "resource_name": domain + path,
+        "request_id": context.aws_request_id,
+        SPAN_TYPE_TAG: SPAN_TYPE_INFERRED,
+    }
+    request_time_epoch = event["requestContext"]["requestTimeEpoch"]
+    args = {
+        "resource": domain + path,
+        "span_type": "http",
+    }
+    tracer.set_tags({"_dd.origin": "lambda"})
+    span = tracer.trace("aws.apigateway", **args)
+    if span:
+        span.set_tags(tags)
+    span.start = request_time_epoch / 1000
+    return span
+
+
+def create_inferred_span_from_http_api_event(event, context):
+    domain = event["requestContext"]["domainName"]
+    path = event["rawPath"]
+    tags = {
+        "operation_name": "aws.httpapi",
+        "service.name": domain,
+        "url": domain + path,
+        "endpoint": path,
+        "http.method": event["requestContext"]["http"]["method"],
+        "resource_name": domain + path,
+        "request_id": context.aws_request_id,
+        SPAN_TYPE_TAG: SPAN_TYPE_INFERRED,
+    }
+    request_time_epoch = event["requestContext"]["timeEpoch"]
+    args = {
+        "resource": domain + path,
+        "span_type": "http",
+    }
+    tracer.set_tags({"_dd.origin": "lambda"})
+    span = tracer.trace("aws.httpapi", **args)
+    if span:
+        span.set_tags(tags)
+    span.start = request_time_epoch / 1000
+    return span
+
+
 def create_function_execution_span(
     context,
     function_name,
     is_cold_start,
     trace_context_source,
     merge_xray_traces,
     trigger_tags,
+    upstream=None,
 ):
     tags = {}
     if context:
@@ -402,6 +546,7 @@ def create_function_execution_span(
             else None,
             "datadog_lambda": datadog_lambda_version,
             "dd_trace": ddtrace_version,
+            "span.name": "aws.lambda",
         }
     if trace_context_source == TraceContextSource.XRAY and merge_xray_traces:
         tags["_dd.parent_source"] = trace_context_source
@@ -415,4 +560,6 @@ def create_function_execution_span(
     span = tracer.trace("aws.lambda", **args)
     if span:
         span.set_tags(tags)
+    if upstream:
+        span.parent_id = upstream.span_id
     return span

datadog_lambda/wrapper.py

@@ -26,12 +26,12 @@
     set_correlation_ids,
     set_dd_trace_py_root,
     create_function_execution_span,
+    create_inferred_span,
 )
 from datadog_lambda.trigger import extract_trigger_tags, extract_http_status_code_tag
 
 logger = logging.getLogger(__name__)
 
-
 """
 Usage:
 
@@ -96,6 +96,11 @@ def __init__(self, func):
             self.extractor_env = os.environ.get("DD_TRACE_EXTRACTOR", None)
             self.trace_extractor = None
             self.span = None
+            self.inferred_span = None
+            self.make_inferred_span = (
+                os.environ.get("DD_INFERRED_SPANS", "false").lower() == "true"
+                and should_use_extension
+            )
             self.response = None
 
             if self.extractor_env:
@@ -138,6 +143,7 @@ def __call__(self, event, context, **kwargs):
 
     def _before(self, event, context):
         try:
+
             set_cold_start()
             submit_invocations_metric(context)
             self.trigger_tags = extract_trigger_tags(event, context)
@@ -153,13 +159,18 @@ def _before(self, event, context):
 
             if dd_tracing_enabled:
                 set_dd_trace_py_root(trace_context_source, self.merge_xray_traces)
+                if self.make_inferred_span:
+                    self.inferred_span = create_inferred_span(
+                        event, context, self.function_name
+                    )
                 self.span = create_function_execution_span(
                     context,
                     self.function_name,
                     is_cold_start(),
                     trace_context_source,
                     self.merge_xray_traces,
                     self.trigger_tags,
+                    upstream=self.inferred_span,
                 )
             else:
                 set_correlation_ids()
@@ -180,16 +191,20 @@ def _after(self, event, context):
                     self.trigger_tags, XraySubsegment.LAMBDA_FUNCTION_TAGS_KEY
                 )
 
-            if not self.flush_to_log or should_use_extension:
-                flush_stats()
-            if should_use_extension:
-                flush_extension()
-
             if self.span:
                 if status_code:
                     self.span.set_tag("http.status_code", status_code)
                 self.span.finish()
+            if self.inferred_span:
+                if status_code:
+                    self.inferred_span.set_tag("http.status_code", status_code)
+                self.inferred_span.finish()
             logger.debug("datadog_lambda_wrapper _after() done")
+
+            if not self.flush_to_log or should_use_extension:
+                flush_stats()
+            if should_use_extension:
+                flush_extension()
         except Exception:
             traceback.print_exc()
 

scripts/run_integration_tests.sh

@@ -208,9 +208,9 @@ for handler_name in "${LAMBDA_HANDLERS[@]}"; do
                 sed -E "s/(\"span_id\"\: \")[A-Z0-9\.\-]+/\1XXXX/g" |
                 sed -E "s/(\"parent_id\"\: \")[A-Z0-9\.\-]+/\1XXXX/g" |
                 sed -E "s/(\"request_id\"\: \")[a-z0-9\.\-]+/\1XXXX/g" |
-                sed -E "s/(\"duration\"\: )[0-9\.\-]+/\1XXXX/g" |
-                sed -E "s/(\"start\"\: )[0-9\.\-]+/\1XXXX/g" |
-                sed -E "s/(\"system\.pid\"\: )[0-9\.\-]+/\1XXXX/g" |
+                sed -E "s/(\"duration\"\: )[0-9\.\-]+/\1\"XXXX\"/g" |
+                sed -E "s/(\"start\"\: )[0-9\.\-]+/\1\"XXXX\"/g" |
+                sed -E "s/(\"system\.pid\"\: )[0-9\.\-]+/\1\"XXXX\"/g" |
                 sed -E "s/(\"runtime-id\"\: \")[a-z0-9\.\-]+/\1XXXX/g" |
                 sed -E "s/(\"datadog_lambda\"\: \")([0-9]+\.[0-9]+\.[0-9])/\1X.X.X/g" |
                 sed -E "s/(\"dd_trace\"\: \")([0-9]+\.[0-9]+\.[0-9])/\1X.X.X/g"