Skip to content

Commit 1abf3f2

Browse files
chemystery09cursoragent
chemystery09
and
cursoragent
authored
feat: inferred alb spans (#838)
* feat: emit inferred aws.alb span for ALB-triggered Lambdas When a Lambda is invoked behind an Application Load Balancer, create an inferred aws.alb span with HTTP request tags and populate ALB HTTP facet tags for http.url, http.route, and http.useragent. FRSLES-851 / APMSVLS-542 Co-authored-by: Cursor <cursoragent@cursor.com> * test: add unit and e2e coverage for inferred aws.alb span Cover ALB inferred-span creation, service mapping, inbound trace context, HTTP trigger tags, and full wrapper lifecycle including status code and parent/child linkage. FRSLES-851 / APMSVLS-542 Co-authored-by: Cursor <cursoragent@cursor.com> * fix: support ALB multiValueHeaders for inferred aws.alb span Handle both ALB event subtypes via resolve_alb_request_headers, fix inbound trace propagation from multiValueHeaders, and expand unit/wrapper coverage for HTTPS URLs, peer.service, and dd_resource_key. FRSLES-851 Co-authored-by: Cursor <cursoragent@cursor.com> * style: fix black formatting for ALB lint CI Collapse http_url assignment and reformat TestAlbInferredSpanWrapper tests so check_format.sh passes. FRSLES-851 Co-authored-by: Cursor <cursoragent@cursor.com> * fix: extract http.referer from resolved ALB headers Reuse resolve_alb_request_headers for referer tagging so multiValueHeaders events and lowercase referer keys are handled consistently with other ALB HTTP facet tags. Co-authored-by: Cursor <cursoragent@cursor.com> * style: fix black formatting for referer test assertions Co-authored-by: Cursor <cursoragent@cursor.com> * refactor: rename resolve_multivalue_headers and simplify extraction Address PR review feedback: the header resolver is not ALB-specific, and extract_context_from_http_event_or_context can delegate to it directly. Co-authored-by: Cursor <cursoragent@cursor.com> --------- Co-authored-by: Cursor <cursoragent@cursor.com>
1 parent 897fb0d commit 1abf3f2

5 files changed

Lines changed: 485 additions & 9 deletions

File tree

‎datadog_lambda/tracing.py‎

Lines changed: 56 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -43,6 +43,7 @@
4343
is_step_function_event,
4444
EventTypes,
4545
EventSubtypes,
46+
resolve_multivalue_headers,
4647
)
4748
from datadog_lambda.durable import extract_context_from_durable_execution
4849

@@ -196,7 +197,7 @@ def extract_context_from_http_event_or_context(
196197
if _is_context_complete(context):
197198
return context
198199

199-
headers = event.get("headers")
200+
headers = resolve_multivalue_headers(event)
200201
context = propagator.extract(headers)
201202

202203
if not _is_context_complete(context):
@@ -658,7 +659,9 @@ def extract_dd_trace_context(
658659
context = extract_context_from_request_header_or_context(
659660
event, lambda_context, event_source
660661
)
661-
elif isinstance(event, (set, dict)) and "headers" in event:
662+
elif isinstance(event, (set, dict)) and (
663+
"headers" in event or "multiValueHeaders" in event
664+
):
662665
context = extract_context_from_http_event_or_context(
663666
event, lambda_context, event_source, decode_authorizer_context
664667
)
@@ -837,6 +840,9 @@ def create_inferred_span(
837840
elif event_source.equals(EventTypes.LAMBDA_FUNCTION_URL):
838841
logger.debug("Function URL event detected. Inferring a span")
839842
return create_inferred_span_from_lambda_function_url_event(event, context)
843+
elif event_source.event_type == EventTypes.ALB:
844+
logger.debug("ALB event detected. Inferring a span")
845+
return create_inferred_span_from_alb_event(event, context)
840846
elif event_source.equals(
841847
EventTypes.API_GATEWAY, subtype=EventSubtypes.HTTP_API
842848
):
@@ -952,6 +958,54 @@ def create_inferred_span_from_lambda_function_url_event(event, context):
952958
return span
953959

954960

961+
def create_inferred_span_from_alb_event(event, context):
962+
request_context = event.get("requestContext") or {}
963+
elb = request_context.get("elb") or {}
964+
target_group_arn = elb.get("targetGroupArn")
965+
966+
headers = resolve_multivalue_headers(event)
967+
host = headers.get("host")
968+
method = event.get("httpMethod")
969+
path = event.get("path")
970+
proto = headers.get("x-forwarded-proto", "http")
971+
972+
# ALB has no api id; key the service mapping off the load-balancer host and
973+
# fall back to it when DD_TRACE_AWS_SERVICE_REPRESENTATION_ENABLED is on.
974+
service_name = determine_service_name(service_mapping, host, "lambda_alb", host)
975+
976+
http_url = "%s://%s%s" % (proto, host, path) if host and path is not None else None
977+
if method and path is not None:
978+
resource = f"{method} {path}"
979+
else:
980+
resource = method or path
981+
982+
tags = {
983+
"operation_name": "aws.alb",
984+
"span.kind": "server",
985+
"http.method": method,
986+
"http.url": http_url,
987+
"http.useragent": headers.get("user-agent"),
988+
"endpoint": path,
989+
"resource_names": resource,
990+
"request_id": context.aws_request_id,
991+
"target_group_arn": target_group_arn,
992+
}
993+
# Drop tags we couldn't derive so the span never carries malformed values.
994+
tags = {key: value for key, value in tags.items() if value is not None}
995+
996+
tracer.set_tags(_dd_origin)
997+
# ALB events carry no request timestamp (unlike API GW requestTimeEpoch /
998+
# Function URL timeEpoch), so the span starts at handler time.
999+
span = tracer.trace(
1000+
"aws.alb", service=service_name, resource=resource, span_type="http"
1001+
)
1002+
InferredSpanInfo.set_tags(tags, tag_source="self", synchronicity="sync")
1003+
if span:
1004+
span.set_tags(tags)
1005+
span.set_metric(InferredSpanInfo.METRIC, 1.0)
1006+
return span
1007+
1008+
9551009
def is_api_gateway_invocation_async(event):
9561010
hdrs = event.get("headers")
9571011
if not hdrs:

‎datadog_lambda/trigger.py‎

Lines changed: 53 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -293,6 +293,39 @@ def get_event_source_arn(source: _EventSource, event: dict, context: Any) -> str
293293
return event_source_arn
294294

295295

296+
def resolve_multivalue_headers(event):
297+
"""
298+
Resolve request headers from single-value ``headers`` or ``multiValueHeaders``
299+
(first value per key, matching datadog-lambda-js).
300+
"""
301+
headers = event.get("headers")
302+
if isinstance(headers, dict) and headers:
303+
return headers
304+
305+
multi_value = event.get("multiValueHeaders")
306+
if not isinstance(multi_value, dict):
307+
return {}
308+
309+
resolved = {}
310+
for key, value in multi_value.items():
311+
if isinstance(value, list):
312+
if value:
313+
resolved[key] = value[0]
314+
elif isinstance(value, str):
315+
resolved[key] = value
316+
return resolved
317+
318+
319+
def _get_header_case_insensitive(headers, name):
320+
if not isinstance(headers, dict):
321+
return None
322+
name_lower = name.lower()
323+
for key, value in headers.items():
324+
if isinstance(key, str) and key.lower() == name_lower and value:
325+
return value
326+
return None
327+
328+
296329
def extract_http_tags(event):
297330
"""
298331
Extracts HTTP facet tags from the triggering event
@@ -324,19 +357,32 @@ def extract_http_tags(event):
324357
path = apigateway_v2_http.get("path")
325358
method = apigateway_v2_http.get("method")
326359

360+
elif request_context and request_context.get("elb"):
361+
# ALB events have no requestContext.stage; derive the URL from the
362+
# forwarded host/proto headers and the top-level path.
363+
alb_headers = resolve_multivalue_headers(event)
364+
host = alb_headers.get("host")
365+
if host:
366+
proto = alb_headers.get("x-forwarded-proto", "http")
367+
http_tags["http.url"] = proto + "://" + host
368+
369+
user_agent = alb_headers.get("user-agent")
370+
if user_agent:
371+
http_tags["http.useragent"] = user_agent
372+
373+
# ALB carries no route template, so use the request path as the route.
374+
if path:
375+
http_tags["http.route"] = path
376+
327377
if path:
328378
if http_tags.get("http.url"):
329379
http_tags["http.url"] += path
330380
if method:
331381
http_tags["http.method"] = method
332382

333-
# Safely get headers
334-
headers = event.get("headers", {})
335-
if not isinstance(headers, dict):
336-
headers = {}
337-
338-
if headers and headers.get("Referer"):
339-
http_tags["http.referer"] = headers.get("Referer")
383+
referer = _get_header_case_insensitive(resolve_multivalue_headers(event), "referer")
384+
if referer:
385+
http_tags["http.referer"] = referer
340386

341387
# Try to get `routeKey` from API GW v2; otherwise try to get `resource` from API GW v1
342388
route = event.get("routeKey") or event.get("resource")

‎tests/test_tracing.py‎

Lines changed: 176 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -112,6 +112,10 @@ def _wrap(*args, **kwargs):
112112
"api-gateway-websocket-disconnect",
113113
Context(trace_id=12345, span_id=67890, sampling_priority=2),
114114
),
115+
(
116+
"application-load-balancer",
117+
Context(trace_id=12345, span_id=67890, sampling_priority=2),
118+
),
115119
(
116120
"authorizer-request-api-gateway-v1",
117121
Context(
@@ -1953,6 +1957,178 @@ def test_remaps_specific_inferred_span_service_names_from_eventbridge_event(
19531957
self.assertEqual(span2.get_tag("operation_name"), "aws.eventbridge")
19541958
self.assertEqual(span2.service, "different.eventbridge.custom.event.sender")
19551959

1960+
def test_remaps_all_inferred_span_service_names_from_alb_event(self):
1961+
self.set_service_mapping({"lambda_alb": "new-name"})
1962+
with open(f"{event_samples}application-load-balancer.json") as event:
1963+
original_event = json.load(event)
1964+
1965+
ctx = get_mock_context()
1966+
ctx.aws_request_id = "123"
1967+
1968+
span1 = create_inferred_span(original_event, ctx)
1969+
self.assertEqual(span1.get_tag("operation_name"), "aws.alb")
1970+
self.assertEqual(span1.service, "new-name")
1971+
1972+
event2 = copy.deepcopy(original_event)
1973+
event2["headers"]["host"] = "different-alb.us-east-2.elb.amazonaws.com"
1974+
span2 = create_inferred_span(event2, ctx)
1975+
self.assertEqual(span2.get_tag("operation_name"), "aws.alb")
1976+
self.assertEqual(span2.service, "new-name")
1977+
1978+
def test_remaps_specific_inferred_span_service_names_from_alb_event(self):
1979+
host = "lambda-alb-123578498.us-east-2.elb.amazonaws.com"
1980+
self.set_service_mapping({host: "mapped-alb-service"})
1981+
with open(f"{event_samples}application-load-balancer.json") as event:
1982+
original_event = json.load(event)
1983+
1984+
ctx = get_mock_context()
1985+
ctx.aws_request_id = "123"
1986+
1987+
span1 = create_inferred_span(original_event, ctx)
1988+
self.assertEqual(span1.get_tag("operation_name"), "aws.alb")
1989+
self.assertEqual(span1.service, "mapped-alb-service")
1990+
1991+
event2 = copy.deepcopy(original_event)
1992+
event2["headers"]["host"] = "other-alb.us-east-2.elb.amazonaws.com"
1993+
span2 = create_inferred_span(event2, ctx)
1994+
self.assertEqual(span2.get_tag("operation_name"), "aws.alb")
1995+
self.assertEqual(span2.service, "other-alb.us-east-2.elb.amazonaws.com")
1996+
1997+
1998+
class TestAlbInferredSpan(unittest.TestCase):
1999+
ALB_SAMPLE = "application-load-balancer"
2000+
ALB_MULTIVALUE = "application-load-balancer-multivalue-headers"
2001+
ALB_HOST = "lambda-alb-123578498.us-east-2.elb.amazonaws.com"
2002+
ALB_USER_AGENT = (
2003+
"Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 "
2004+
"(KHTML, like Gecko) Chrome/71.0.3578.98 Safari/537.36"
2005+
)
2006+
2007+
def _load_event(self, sample_name):
2008+
with open(f"{event_samples}{sample_name}.json") as event_file:
2009+
return json.load(event_file)
2010+
2011+
def test_create_inferred_span_from_alb_event(self):
2012+
event = self._load_event(self.ALB_SAMPLE)
2013+
ctx = get_mock_context(aws_request_id="123")
2014+
2015+
span = create_inferred_span(event, ctx)
2016+
2017+
self.assertIsNotNone(span)
2018+
self.assertEqual(span.name, "aws.alb")
2019+
self.assertEqual(span.span_type, "http")
2020+
self.assertEqual(span.service, self.ALB_HOST)
2021+
self.assertEqual(span.resource, "GET /lambda")
2022+
self.assertEqual(span.get_tag("operation_name"), "aws.alb")
2023+
self.assertEqual(span.get_tag("span.kind"), "server")
2024+
self.assertEqual(span.get_tag("http.method"), "GET")
2025+
self.assertEqual(
2026+
span.get_tag("http.url"),
2027+
"http://%s/lambda" % self.ALB_HOST,
2028+
)
2029+
self.assertEqual(span.get_tag("http.useragent"), self.ALB_USER_AGENT)
2030+
self.assertEqual(span.get_tag("endpoint"), "/lambda")
2031+
self.assertEqual(span.get_tag("resource_names"), "GET /lambda")
2032+
self.assertEqual(span.get_tag("request_id"), "123")
2033+
self.assertEqual(span.get_tag("_inferred_span.synchronicity"), "sync")
2034+
self.assertEqual(span.get_tag("_inferred_span.tag_source"), "self")
2035+
self.assertEqual(span.get_metric("_dd._inferred_span"), 1.0)
2036+
self.assertEqual(
2037+
span.get_tag("target_group_arn"),
2038+
"arn:aws:elasticloadbalancing:us-east-2:123456789012:targetgroup/lambda-xyz/123abc",
2039+
)
2040+
2041+
def test_create_inferred_span_omits_tags_when_headers_missing(self):
2042+
event = self._load_event(self.ALB_SAMPLE)
2043+
del event["headers"]
2044+
event["httpMethod"] = None
2045+
event["path"] = None
2046+
2047+
span = create_inferred_span(event, get_mock_context())
2048+
2049+
self.assertIsNotNone(span)
2050+
self.assertNotIn("http.url", span.get_tags())
2051+
self.assertNotIn("http.method", span.get_tags())
2052+
self.assertNotIn("http.useragent", span.get_tags())
2053+
2054+
def test_multivalue_headers_subtype_emits_inferred_span(self):
2055+
event = self._load_event(self.ALB_MULTIVALUE)
2056+
span = create_inferred_span(event, get_mock_context())
2057+
self.assertIsNotNone(span)
2058+
self.assertEqual(span.name, "aws.alb")
2059+
self.assertEqual(span.get_tag("http.method"), "GET")
2060+
self.assertEqual(
2061+
span.get_tag("http.url"),
2062+
"http://%s/lambda" % self.ALB_HOST,
2063+
)
2064+
self.assertEqual(span.get_tag("http.useragent"), self.ALB_USER_AGENT)
2065+
2066+
@with_trace_propagation_style("datadog")
2067+
def test_inbound_datadog_context_from_multivalue_headers(self):
2068+
event = self._load_event(self.ALB_MULTIVALUE)
2069+
ctx = get_mock_context()
2070+
2071+
parent_ctx, source, _ = extract_dd_trace_context(event, ctx)
2072+
self.assertIsNotNone(parent_ctx)
2073+
self.assertEqual(parent_ctx.trace_id, 12345)
2074+
self.assertEqual(parent_ctx.span_id, 67890)
2075+
2076+
set_dd_trace_py_root(source, merge_xray_traces=False)
2077+
span = create_inferred_span(event, ctx)
2078+
self.assertEqual(span.trace_id, parent_ctx.trace_id)
2079+
self.assertEqual(span.parent_id, parent_ctx.span_id)
2080+
2081+
@with_trace_propagation_style("datadog")
2082+
def test_inbound_datadog_context_parents_inferred_span(self):
2083+
event = self._load_event(self.ALB_SAMPLE)
2084+
ctx = get_mock_context()
2085+
2086+
parent_ctx, source, _ = extract_dd_trace_context(event, ctx)
2087+
set_dd_trace_py_root(source, merge_xray_traces=False)
2088+
span = create_inferred_span(event, ctx)
2089+
2090+
self.assertEqual(span.trace_id, parent_ctx.trace_id)
2091+
self.assertEqual(span.parent_id, parent_ctx.span_id)
2092+
2093+
def test_inbound_w3c_context_extracted_from_alb_event(self):
2094+
event = self._load_event(self.ALB_SAMPLE)
2095+
event["headers"] = {
2096+
"host": self.ALB_HOST,
2097+
"user-agent": self.ALB_USER_AGENT,
2098+
"x-forwarded-proto": "http",
2099+
"traceparent": "00-0000000000000000000000000000abcd-000000000000004d-01",
2100+
"tracestate": "dd=s:1",
2101+
}
2102+
2103+
ctx, source, _ = extract_dd_trace_context(event, get_mock_context())
2104+
2105+
self.assertIsNotNone(ctx)
2106+
self.assertEqual(source, TraceContextSource.EVENT)
2107+
self.assertEqual(ctx.trace_id, 0xABCD)
2108+
self.assertEqual(ctx.span_id, 0x4D)
2109+
2110+
def test_http_url_uses_https_when_forwarded_proto_is_https(self):
2111+
event = self._load_event(self.ALB_SAMPLE)
2112+
event["headers"]["x-forwarded-proto"] = "https"
2113+
2114+
span = create_inferred_span(event, get_mock_context())
2115+
2116+
self.assertEqual(
2117+
span.get_tag("http.url"),
2118+
"https://%s/lambda" % self.ALB_HOST,
2119+
)
2120+
2121+
def test_http_url_excludes_query_string(self):
2122+
event = self._load_event(self.ALB_SAMPLE)
2123+
2124+
span = create_inferred_span(event, get_mock_context())
2125+
2126+
self.assertEqual(
2127+
span.get_tag("http.url"),
2128+
"http://%s/lambda" % self.ALB_HOST,
2129+
)
2130+
self.assertNotIn("query=", span.get_tag("http.url") or "")
2131+
19562132

19572133
class _Span(object):
19582134
def __init__(self, service, start, span_type, parent_name=None, tags=None):

0 commit comments

Comments
 (0)