[SLS-1683] Add sync/async tag and set inferred span end time based on value

datadog_lambda/constants.py

@@ -41,3 +41,6 @@ class XrayDaemon(object):
     XRAY_TRACE_ID_HEADER_NAME = "_X_AMZN_TRACE_ID"
     XRAY_DAEMON_ADDRESS = "AWS_XRAY_DAEMON_ADDRESS"
     FUNCTION_NAME_HEADER_NAME = "AWS_LAMBDA_FUNCTION_NAME"
+
+
+IS_ASYNC_TAG = "is_async"

datadog_lambda/tracing.py

@@ -6,12 +6,14 @@
 import logging
 import os
 import json
+from datetime import datetime, timezone
 
 from datadog_lambda.constants import (
     SamplingPriority,
     TraceHeader,
     TraceContextSource,
     XrayDaemon,
+    IS_ASYNC_TAG,
 )
 from datadog_lambda.xray import (
     send_segment,
@@ -27,7 +29,6 @@
     EventTypes,
     EventSubtypes,
 )
-from datetime import datetime, timezone
 
 logger = logging.getLogger(__name__)
 
@@ -436,6 +437,14 @@ def create_inferred_span(event, context):
     return None
 
 
+def is_api_gateway_invocation_async(event):
+    return (
+        "headers" in event
+        and "X-Amz-Invocation-Type" in event["headers"]
+        and event["headers"]["X-Amz-Invocation-Type"] == "Event"
+    )
+
+
 def create_inferred_span_from_api_gateway_websocket_event(event, context):
     domain = event["requestContext"]["domainName"]
     endpoint = event["requestContext"]["routeKey"]
@@ -448,6 +457,7 @@ def create_inferred_span_from_api_gateway_websocket_event(event, context):
         "request_id": context.aws_request_id,
         "connection_id": event["requestContext"]["connectionId"],
         SPAN_TYPE_TAG: SPAN_TYPE_INFERRED,
+        IS_ASYNC_TAG: is_api_gateway_invocation_async(event),
     }
     request_time_epoch = event["requestContext"]["requestTimeEpoch"]
     args = {
@@ -474,6 +484,7 @@ def create_inferred_span_from_api_gateway_event(event, context):
         "resource_names": domain + path,
         "request_id": context.aws_request_id,
         SPAN_TYPE_TAG: SPAN_TYPE_INFERRED,
+        IS_ASYNC_TAG: is_api_gateway_invocation_async(event),
     }
     request_time_epoch = event["requestContext"]["requestTimeEpoch"]
     args = {
@@ -500,6 +511,7 @@ def create_inferred_span_from_http_api_event(event, context):
         "resource_names": domain + path,
         "request_id": context.aws_request_id,
         SPAN_TYPE_TAG: SPAN_TYPE_INFERRED,
+        IS_ASYNC_TAG: is_api_gateway_invocation_async(event),
     }
     request_time_epoch = event["requestContext"]["timeEpoch"]
     args = {
@@ -522,6 +534,7 @@ def create_inferred_span_from_sqs_event(event, context):
         "service.name": "sqs",
         "resource_names": queue_name,
         SPAN_TYPE_TAG: SPAN_TYPE_INFERRED,
+        IS_ASYNC_TAG: True,
     }
     request_time_epoch = event_record["attributes"]["SentTimestamp"]
     args = {
@@ -544,6 +557,7 @@ def create_inferred_span_from_sns_event(event, context):
         "service.name": "sns",
         "resource_names": topic_name,
         SPAN_TYPE_TAG: SPAN_TYPE_INFERRED,
+        IS_ASYNC_TAG: True,
     }
     sns_dt_format = "%Y-%m-%dT%H:%M:%S.%fZ"
     timestamp = event_record["Sns"]["Timestamp"]
@@ -569,6 +583,7 @@ def create_inferred_span_from_kinesis_event(event, context):
         "service.name": "kinesis",
         "resource_names": stream_name,
         SPAN_TYPE_TAG: SPAN_TYPE_INFERRED,
+        IS_ASYNC_TAG: True,
     }
     request_time_epoch = event_record["kinesis"]["approximateArrivalTimestamp"]
 
@@ -592,6 +607,7 @@ def create_inferred_span_from_dynamodb_event(event, context):
         "service.name": "dynamodb",
         "resource_names": table_name,
         SPAN_TYPE_TAG: SPAN_TYPE_INFERRED,
+        IS_ASYNC_TAG: True,
     }
     request_time_epoch = event_record["dynamodb"]["ApproximateCreationDateTime"]
 
@@ -615,6 +631,7 @@ def create_inferred_span_from_s3_event(event, context):
         "service.name": "s3",
         "resource_names": bucket_name,
         SPAN_TYPE_TAG: SPAN_TYPE_INFERRED,
+        IS_ASYNC_TAG: True,
     }
     dt_format = "%Y-%m-%dT%H:%M:%S.%fZ"
     timestamp = event_record["eventTime"]
@@ -639,6 +656,7 @@ def create_inferred_span_from_eventbridge_event(event, context):
         "service.name": "eventbridge",
         "resource_names": source,
         SPAN_TYPE_TAG: SPAN_TYPE_INFERRED,
+        IS_ASYNC_TAG: True,
     }
     dt_format = "%Y-%m-%dT%H:%M:%SZ"
     timestamp = event["time"]

datadog_lambda/wrapper.py

@@ -10,7 +10,11 @@
 
 from datadog_lambda.extension import should_use_extension, flush_extension
 from datadog_lambda.cold_start import set_cold_start, is_cold_start
-from datadog_lambda.constants import XraySubsegment, TraceContextSource
+from datadog_lambda.constants import (
+    XraySubsegment,
+    TraceContextSource,
+    IS_ASYNC_TAG,
+)
 from datadog_lambda.metric import (
     flush_stats,
     submit_invocations_metric,
@@ -201,10 +205,15 @@ def _after(self, event, context):
                 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()
+
+                if self.inferred_span.get_tag(IS_ASYNC_TAG) == "True" and self.span:
+                    self.inferred_span.finish(finish_time=self.span.start)
+                else:
+                    self.inferred_span.finish()
 
             if not self.flush_to_log or should_use_extension:
                 flush_stats()

tests/event_samples/api-gateway-non-proxy-async.json

@@ -0,0 +1,106 @@
+{
+    "resource": "/http/get",
+    "path": "/http/get",
+    "httpMethod": "GET",
+    "headers": {
+        "Accept": "*/*",
+        "CloudFront-Forwarded-Proto": "https",
+        "CloudFront-Is-Desktop-Viewer": "true",
+        "CloudFront-Is-Mobile-Viewer": "false",
+        "CloudFront-Is-SmartTV-Viewer": "false",
+        "CloudFront-Is-Tablet-Viewer": "false",
+        "CloudFront-Viewer-Country": "US",
+        "Host": "lgxbo6a518.execute-api.sa-east-1.amazonaws.com",
+        "User-Agent": "curl/7.64.1",
+        "Via": "2.0 a1882a601559755135741e91a9f86c28.cloudfront.net (CloudFront)",
+        "X-Amz-Cf-Id": "dHMNSBrMT0Xjg3rtMrI0Ie9BDg3D_OIPoj7m0mTuIOpTujrU0Ob8_A==",
+        "X-Amzn-Trace-Id": "Root=1-613a4da3-5012576973e2e5670d4c549a",
+        "X-Forwarded-For": "38.122.226.210, 70.132.52.143",
+        "X-Forwarded-Port": "443",
+        "X-Forwarded-Proto": "https",
+        "X-Amz-Invocation-Type": "Event"
+    },
+    "multiValueHeaders": {
+        "Accept": [
+            "*/*"
+        ],
+        "CloudFront-Forwarded-Proto": [
+            "https"
+        ],
+        "CloudFront-Is-Desktop-Viewer": [
+            "true"
+        ],
+        "CloudFront-Is-Mobile-Viewer": [
+            "false"
+        ],
+        "CloudFront-Is-SmartTV-Viewer": [
+            "false"
+        ],
+        "CloudFront-Is-Tablet-Viewer": [
+            "false"
+        ],
+        "CloudFront-Viewer-Country": [
+            "US"
+        ],
+        "Host": [
+            "lgxbo6a518.execute-api.sa-east-1.amazonaws.com"
+        ],
+        "User-Agent": [
+            "curl/7.64.1"
+        ],
+        "Via": [
+            "2.0 a1882a601559755135741e91a9f86c28.cloudfront.net (CloudFront)"
+        ],
+        "X-Amz-Cf-Id": [
+            "dHMNSBrMT0Xjg3rtMrI0Ie9BDg3D_OIPoj7m0mTuIOpTujrU0Ob8_A=="
+        ],
+        "X-Amzn-Trace-Id": [
+            "Root=1-613a4da3-5012576973e2e5670d4c549a"
+        ],
+        "X-Forwarded-For": [
+            "38.122.226.210, 70.132.52.143"
+        ],
+        "X-Forwarded-Port": [
+            "443"
+        ],
+        "X-Forwarded-Proto": [
+            "https"
+        ]
+    },
+    "queryStringParameters": null,
+    "multiValueQueryStringParameters": null,
+    "pathParameters": null,
+    "stageVariables": null,
+    "requestContext": {
+        "resourceId": "8ajqil",
+        "resourcePath": "/http/get",
+        "httpMethod": "GET",
+        "extendedRequestId": "FaERiG1RGjQFb6g=",
+        "requestTime": "09/Sep/2021:18:08:35 +0000",
+        "path": "/dev/http/get",
+        "accountId": "601427279990",
+        "protocol": "HTTP/1.1",
+        "stage": "dev",
+        "domainPrefix": "lgxbo6a518",
+        "requestTimeEpoch": 1631210915251,
+        "requestId": "7bf3b161-f698-432c-a639-6fef8b445137",
+        "identity": {
+            "cognitoIdentityPoolId": null,
+            "accountId": null,
+            "cognitoIdentityId": null,
+            "caller": null,
+            "sourceIp": "38.122.226.210",
+            "principalOrgId": null,
+            "accessKey": null,
+            "cognitoAuthenticationType": null,
+            "cognitoAuthenticationProvider": null,
+            "userArn": null,
+            "userAgent": "curl/7.64.1",
+            "user": null
+        },
+        "domainName": "lgxbo6a518.execute-api.sa-east-1.amazonaws.com",
+        "apiId": "lgxbo6a518"
+    },
+    "body": null,
+    "isBase64Encoded": false
+}