Skip to content

Commit

Permalink
feat: add support for payload sizelimiting for aws-lambda
Browse files Browse the repository at this point in the history
  • Loading branch information
@povilasv
povilasv committed Oct 19, 2023
1 parent df8a761 commit e13560d
Showing 1 changed file with 35 additions and 20 deletions.
55 changes: 35 additions & 20 deletions ...metry-instrumentation-aws-lambda/src/opentelemetry/instrumentation/aws_lambda/__init__.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -220,7 +220,7 @@ def _determine_parent_context(


def _set_api_gateway_v1_proxy_attributes(
lambda_event: Any, span: Span
lambda_event: Any, span: Span, payload_size_limit: int
) -> Span:
"""Sets HTTP attributes for REST APIs and v1 HTTP APIs
Expand All @@ -234,7 +234,7 @@ def _set_api_gateway_v1_proxy_attributes(
if lambda_event.get("body"):
span.set_attribute(
"http.request.body",
lambda_event.get("body"),
limit_string_size(payload_size_limit, lambda_event.get("body")),
)

if lambda_event.get("headers"):
Expand Down Expand Up @@ -270,7 +270,7 @@ def _set_api_gateway_v1_proxy_attributes(


def _set_api_gateway_v2_proxy_attributes(
lambda_event: Any, span: Span
lambda_event: Any, span: Span, payload_size_limit: int
) -> Span:
"""Sets HTTP attributes for v2 HTTP APIs
Expand All @@ -286,7 +286,7 @@ def _set_api_gateway_v2_proxy_attributes(
if lambda_event.get("body"):
span.set_attribute(
"http.request.body",
lambda_event.get("body"),
limit_string_size(payload_size_limit, lambda_event.get("body")),
)

span.set_attribute(
Expand Down Expand Up @@ -340,6 +340,14 @@ def _instrumented_lambda_handler_call( # noqa pylint: disable=too-many-branches
[wrapped_module_name, wrapped_function_name]
)

payload_size_limit = 51200
try:
payload_size_limit = int(os.environ.get("OTEL_PAYLOAD_SIZE_LIMIT", 51200))
except ValueError:
logger.error(
"OTEL_PAYLOAD_SIZE_LIMIT is not a number"
)

lambda_event = args[0]

parent_context = _determine_parent_context(
Expand Down Expand Up @@ -414,7 +422,7 @@ def _instrumented_lambda_handler_call( # noqa pylint: disable=too-many-branches
if lambda_event["Records"][0].get("s3"):
s3TriggerSpan.set_attribute(
"rpc.request.body",
json.dumps(lambda_event["Records"][0].get("s3")),
limit_string_size(payload_size_limit, json.dumps(lambda_event["Records"][0].get("s3"))),
)
except Exception as ex:
pass
Expand All @@ -441,7 +449,7 @@ def _instrumented_lambda_handler_call( # noqa pylint: disable=too-many-branches
if lambda_event["Records"][0].get("body"):
sqsTriggerSpan.set_attribute(
"rpc.request.body",
lambda_event["Records"][0].get("body"),
limit_string_size(payload_size_limit, lambda_event["Records"][0].get("body")),
)

except Exception as ex:
Expand All @@ -461,7 +469,7 @@ def _instrumented_lambda_handler_call( # noqa pylint: disable=too-many-branches
if lambda_event["Records"][0]["Sns"] and lambda_event["Records"][0]["Sns"].get("Message"):
snsTriggerSpan.set_attribute(
"rpc.request.body",
lambda_event["Records"][0]["Sns"].get("Message"),
limit_string_size(payload_size_limit, lambda_event["Records"][0]["Sns"].get("Message"),
)
except Exception as ex:
pass
Expand All @@ -482,7 +490,7 @@ def _instrumented_lambda_handler_call( # noqa pylint: disable=too-many-branches
if lambda_event["Records"][0].get("dynamodb"):
dynamoTriggerSpan.set_attribute(
"rpc.request.body",
json.dumps(lambda_event["Records"][0].get("dynamodb")),
limit_string_size(payload_size_limit, json.dumps(lambda_event["Records"][0].get("dynamodb"))),
)
except Exception as ex:
pass
Expand All @@ -503,7 +511,7 @@ def _instrumented_lambda_handler_call( # noqa pylint: disable=too-many-branches
if lambda_event["datasetRecords"]:
cognitoTriggerSpan.set_attribute(
"rpc.request.body",
json.dumps(lambda_event["datasetRecords"]),
limit_string_size(payload_size_limit, json.dumps(lambda_event["datasetRecords"])),
)
except Exception as ex:
pass
Expand All @@ -528,7 +536,7 @@ def _instrumented_lambda_handler_call( # noqa pylint: disable=too-many-branches

eventBridgeTriggerSpan.set_attribute(
"rpc.request.body",
json.dumps(lambda_event),
limit_string_size(payload_size_limit, json.dumps(lambda_event)),
)
except Exception as ex:
pass
Expand Down Expand Up @@ -566,9 +574,9 @@ def _instrumented_lambda_handler_call( # noqa pylint: disable=too-many-branches
apiGwSpan.set_attribute(SpanAttributes.FAAS_TRIGGER, "http")

if lambda_event.get("version") == "2.0":
_set_api_gateway_v2_proxy_attributes(lambda_event, apiGwSpan)
_set_api_gateway_v2_proxy_attributes(lambda_event, apiGwSpan, payload_size_limit)
else:
_set_api_gateway_v1_proxy_attributes(lambda_event, apiGwSpan)
_set_api_gateway_v1_proxy_attributes(lambda_event, apiGwSpan, payload_size_limit)

if isinstance(result, dict) and result.get("statusCode"):
apiGwSpan.set_attribute(
Expand All @@ -578,7 +586,7 @@ def _instrumented_lambda_handler_call( # noqa pylint: disable=too-many-branches
if isinstance(result, dict) and result.get("body"):
apiGwSpan.set_attribute(
"http.response.body",
result.get("body"),
limit_string_size(payload_size_limit, result.get("body")),
)
if lambda_event.get("headers"):
for key, value in lambda_event.get("headers").items():
Expand All @@ -594,7 +602,7 @@ def _instrumented_lambda_handler_call( # noqa pylint: disable=too-many-branches
if lambda_event["Records"][0]["eventSource"] == "aws:sqs":
span.set_attribute(SpanAttributes.FAAS_TRIGGER, "pubsub")
span.set_attribute("messaging.message",
lambda_event["Records"])
limit_string_size(payload_size_limit, lambda_event["Records"]))
except Exception:
pass
except Exception:
Expand All @@ -618,7 +626,7 @@ def _instrumented_lambda_handler_call( # noqa pylint: disable=too-many-branches
if isinstance(result, dict) and result.get("body"):
s3TriggerSpan.set_attribute(
"rpc.response.body",
result.get("body"),
limit_string_size(payload_size_limit, result.get("body")),
)
except Exception:
pass
Expand All @@ -635,7 +643,7 @@ def _instrumented_lambda_handler_call( # noqa pylint: disable=too-many-branches
if isinstance(result, dict) and result.get("body"):
sqsTriggerSpan.set_attribute(
"rpc.response.body",
result.get("body"),
limit_string_size(payload_size_limit, result.get("body")),
)
except Exception:
pass
Expand All @@ -651,7 +659,7 @@ def _instrumented_lambda_handler_call( # noqa pylint: disable=too-many-branches
if isinstance(result, dict) and result.get("body"):
snsTriggerSpan.set_attribute(
"rpc.response.body",
result.get("body"),
limit_string_size(payload_size_limit, result.get("body")),
)
except Exception:
pass
Expand All @@ -668,7 +676,7 @@ def _instrumented_lambda_handler_call( # noqa pylint: disable=too-many-branches
if isinstance(result, dict) and result.get("body"):
dynamoTriggerSpan.set_attribute(
"rpc.response.body",
result.get("body"),
limit_string_size(payload_size_limit, result.get("body")),
)
except Exception:
pass
Expand All @@ -684,7 +692,7 @@ def _instrumented_lambda_handler_call( # noqa pylint: disable=too-many-branches
if isinstance(result, dict) and result.get("body"):
cognitoTriggerSpan.set_attribute(
"rpc.response.body",
result.get("body"),
limit_string_size(payload_size_limit, result.get("body")),
)
except Exception:
pass
Expand All @@ -700,7 +708,7 @@ def _instrumented_lambda_handler_call( # noqa pylint: disable=too-many-branches
if isinstance(result, dict) and result.get("body"):
eventBridgeTriggerSpan.set_attribute(
"rpc.response.body",
result.get("body"),
limit_string_size(payload_size_limit, result.get("body")),
)
except Exception:
pass
Expand Down Expand Up @@ -867,3 +875,10 @@ def keys(
) -> typing.List[str]:
"""Keys implementation that returns all keys from a dictionary."""
return list(carrier.keys())

def limit_string_size(max_size: int, s: str) -> str:
if len(s) > max_size:
return s[:max_size]
else:
return s

0 comments on commit e13560d

Please sign in to comment.