Skip to content

improve X-Amz-Trace-Id for APIGW NG #11327

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 2 commits into from
Aug 8, 2024
Merged

improve X-Amz-Trace-Id for APIGW NG #11327

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
34d1ff7
add trace id to APIGW NG
bentsku Aug 7, 2024
735d404
implement trace-id from incoming
bentsku Aug 7, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 3 additions & 0 deletions localstack-core/localstack/services/apigateway/next_gen/execute_api/context.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -88,6 +88,8 @@ class RestApiInvocationContext(RequestContext):
"""The region the REST API is living in."""
account_id: Optional[str]
"""The account the REST API is living in."""
trace_id: Optional[str]
"""The X-Ray trace ID for the request."""
resource: Optional[Resource]
"""The resource the invocation matched"""
resource_method: Optional[Method]
Expand Down Expand Up @@ -126,3 +128,4 @@ def __init__(self, request: Request):
self.integration_request = None
self.endpoint_response = None
self.invocation_response = None
self.trace_id = None
4 changes: 2 additions & 2 deletions ...-core/localstack/services/apigateway/next_gen/execute_api/handlers/integration_request.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,7 @@
from localstack.constants import APPLICATION_JSON
from localstack.http import Request, Response
from localstack.utils.collections import merge_recursive
from localstack.utils.strings import short_uid, to_bytes, to_str
from localstack.utils.strings import to_bytes, to_str

from ..api import RestApiGatewayHandler, RestApiGatewayHandlerChain
from ..context import IntegrationRequest, InvocationRequest, RestApiInvocationContext
Expand Down Expand Up @@ -282,7 +282,7 @@ def _apply_header_transforms(
default_headers["Content-Type"] = content_type

set_default_headers(headers, default_headers)
headers.set("X-Amzn-Trace-Id", short_uid()) # TODO
headers.set("X-Amzn-Trace-Id", context.trace_id)
if integration_type not in (IntegrationType.AWS_PROXY, IntegrationType.AWS):
headers.set("X-Amzn-Apigateway-Api-Id", context.api_id)

Expand Down
15 changes: 15 additions & 0 deletions localstack-core/localstack/services/apigateway/next_gen/execute_api/handlers/parse.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,7 @@
from ..api import RestApiGatewayHandler, RestApiGatewayHandlerChain
from ..context import InvocationRequest, RestApiInvocationContext
from ..header_utils import should_drop_header_from_invocation
from ..helpers import generate_trace_id, generate_trace_parent, parse_trace_id
from ..moto_helpers import get_stage_variables
from ..variables import ContextVariables, ContextVarsIdentity

Expand Down Expand Up @@ -44,6 +45,8 @@ def parse_and_enrich(self, context: RestApiInvocationContext):
context.stage_variables = self.fetch_stage_variables(context)
LOG.debug("Initializing $stageVariables='%s'", context.stage_variables)

context.trace_id = self.populate_trace_id(context.request.headers)

def create_invocation_request(self, context: RestApiInvocationContext) -> InvocationRequest:
request = context.request
params, multi_value_params = self._get_single_and_multi_values_from_multidict(request.args)
Expand Down Expand Up @@ -166,3 +169,15 @@ def fetch_stage_variables(context: RestApiInvocationContext) -> Optional[dict[st
return None

return stage_variables

@staticmethod
def populate_trace_id(headers: Headers) -> str:
incoming_trace = parse_trace_id(headers.get("x-amzn-trace-id", ""))
# parse_trace_id always return capitalized keys

trace = incoming_trace.get("Root", generate_trace_id())
incoming_parent = incoming_trace.get("Parent")
parent = incoming_parent or generate_trace_parent()
sampled = incoming_trace.get("Sampled", "1" if incoming_parent else "0")
# TODO: lineage? not sure what it related to
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Lineage is related to detecting recursive Lambda loops (July 2023): https://aws.amazon.com/blogs/compute/detecting-and-stopping-recursive-loops-in-aws-lambda-functions/

They extended support for S3, SQS, and SNS in Aug 2024: https://aws.amazon.com/blogs/compute/aws-lambda-introduces-recursive-loop-detection-apis/

return f"Root={trace};Parent={parent};Sampled={sampled}"
5 changes: 3 additions & 2 deletions ...ck-core/localstack/services/apigateway/next_gen/execute_api/handlers/response_enricher.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -23,7 +23,8 @@ def __call__(
headers.set("x-amz-apigw-id", short_uid() + "=")
if (
context.integration
and context.integration["type"] != IntegrationType.HTTP_PROXY
and context.integration["type"]
not in (IntegrationType.HTTP_PROXY, IntegrationType.MOCK)
and not context.context_variables.get("error")
):
headers.set("X-Amzn-Trace-Id", short_uid()) # TODO
headers.set("X-Amzn-Trace-Id", context.trace_id)
26 changes: 26 additions & 0 deletions localstack-core/localstack/services/apigateway/next_gen/execute_api/helpers.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,6 +1,8 @@
import copy
import logging
import re
import time
from secrets import token_hex
from typing import Type, TypedDict

from moto.apigateway.models import RestAPI as MotoRestAPI
Expand Down Expand Up @@ -113,3 +115,27 @@ def validate_sub_dict_of_typed_dict(typed_dict: Type[TypedDict], obj: dict) -> b
typed_dict_keys = {*typed_dict.__required_keys__, *typed_dict.__optional_keys__}

return not bool(set(obj) - typed_dict_keys)


def generate_trace_id():
"""https://docs.aws.amazon.com/xray/latest/devguide/xray-api-sendingdata.html#xray-api-traceids"""
original_request_epoch = int(time.time())
timestamp_hex = hex(original_request_epoch)[2:]
version_number = "1"
unique_id = token_hex(12)
return f"{version_number}-{timestamp_hex}-{unique_id}"


def generate_trace_parent():
return token_hex(8)


def parse_trace_id(trace_id: str) -> dict[str, str]:
split_trace = trace_id.split(";")
trace_values = {}
for trace_part in split_trace:
key_value = trace_part.split("=")
if len(key_value) == 2:
trace_values[key_value[0].capitalize()] = key_value[1]

return trace_values
Comment on lines +120 to +141
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Question: Do you think we should put those method in a shared utils for more services to have access to? I feel like they will mostly be the same for everywhere?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

That's a good point! Not sure where it should live for now 🤔 we might keep it here for now, and if we hear about more x-ray traces, move it somewhere else?

I believe lambda might have a similar implementation..

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think Lambda is only using long_uid(), unless I missed it.

I was thinking that it could make sense in localstack-core/utils maybe create a trace.py but not super convinced either.

I think it is fine with it here for now, but we should stay alert to prevent this behavior to be implemented many times in the code base!

Loading
Loading