diff --git a/datadog_lambda/wrapper.py b/datadog_lambda/wrapper.py index 8dbd7e35..31609c76 100644 --- a/datadog_lambda/wrapper.py +++ b/datadog_lambda/wrapper.py @@ -177,11 +177,13 @@ def __init__(self, func): def __call__(self, event, context, **kwargs): """Executes when the wrapped function gets called""" - self._before(event, context) + actual_event, custom_probes = self._extract_di_request(event) + self._before(actual_event, context, custom_probes) try: if self.blocking_response: return self.blocking_response - self.response = self.func(event, context, **kwargs) + # Pass unwrapped event to user's handler + self.response = self.func(actual_event, context, **kwargs) return self.response except BlockingException: self.blocking_response = get_asm_blocked_response(self.event_source) @@ -194,10 +196,109 @@ def __call__(self, event, context, **kwargs): self.span.set_traceback() raise finally: - self._after(event, context) + self._after(actual_event, context, custom_probes) if self.blocking_response: return self.blocking_response + def _enable_dynamic_instrumentation(self): + from ddtrace.debugging import DynamicInstrumentation + DynamicInstrumentation.enable() + logger.debug("DI framework enabled for this invocation") + + def _disable_dynamic_instrumentation(self): + try: + from ddtrace.debugging import DynamicInstrumentation + DynamicInstrumentation.disable() + logger.debug("DI framework disabled after invocation") + except Exception as e: + logger.error(f"Failed to disable DI framework: {e}") + + def _inject_probes(self, probe_definitions): + if not probe_definitions: + return + + try: + from ddtrace.debugging._debugger import Debugger + from ddtrace.debugging._probe.remoteconfig import build_probe, ProbePollerEvent + + # Get the debugger instance + debugger = Debugger._instance + if debugger is None: + logger.error("Debugger instance not available, cannot inject probes") + return + + # Convert probe definitions to Probe objects + probes = [build_probe(p) for p in probe_definitions] + + # Track probe IDs for cleanup + for probe in probes: + self._active_probe_ids.append(probe.probe_id) + + # Inject probes using the debugger's configuration method + debugger._on_configuration(ProbePollerEvent.NEW_PROBES, probes) + logger.debug(f"Injected {len(probes)} probes: {[p.probe_id for p in probes]}") + + except Exception as e: + logger.error(f"Failed to inject custom probes: {e}") + import traceback + logger.error(f"Traceback: {traceback.format_exc()}") + + def _remove_probes(self): + try: + from ddtrace.debugging._debugger import Debugger + from ddtrace.debugging._probe.remoteconfig import ProbePollerEvent + + debugger = Debugger._instance + if debugger is None: + logger.debug("Debugger instance not available, cannot remove probes") + return + + probes_to_remove = [] + for probe_id in self._active_probe_ids: + probe = debugger._probe_registry.get(probe_id) + if probe is not None: + probes_to_remove.append(probe) + logger.debug(f"Marking probe for removal: {probe_id}") + else: + logger.debug(f"Probe {probe_id} not found in registry") + if probes_to_remove: + debugger._on_configuration(ProbePollerEvent.DELETED_PROBES, probes_to_remove) + logger.debug(f"Removed {len(probes_to_remove)} probes") + + self._active_probe_ids = [] + + except Exception as e: + logger.error(f"Failed to remove probes: {e}") + import traceback + logger.error(f"Traceback: {traceback.format_exc()}") + + def _extract_di_request(self, event): + """ + Check if event requests Dynamic Instrumentation and extract actual payload and probes. + + Supported formats: + 1. Regular payload: {"key1": "value1"} + -> (event, False, None) + DI is NOT enabled + + 2. DI with custom probes: {"probes": [...], "payload": {...}} + -> (unwrapped_payload, True, [probe_definitions]) + DI is enabled with custom probes, both disabled/removed after + Note: Presence of "probes" implies enableDI=true + + Returns: (actual_payload, probes) + """ + if isinstance(event, dict): + has_probes = "probes" in event + + if has_probes: + probes = event.get("probes", None) + actual_payload = event.get("payload", event) + logger.debug(f"DI request detected: custom_probes={probes is not None}, probe_count={len(probes) if probes else 0}") + return actual_payload, probes + + return event, [] + def _inject_authorizer_span_headers(self, request_id): reference_span = self.inferred_span if self.inferred_span else self.span assert reference_span.finished @@ -226,12 +327,18 @@ def _inject_authorizer_span_headers(self, request_id): self.response.setdefault("context", {}) self.response["context"]["_datadog"] = datadog_data - def _before(self, event, context): + def _before(self, event, context, custom_probes=None): try: self.response = None self.blocking_response = None + self._active_probe_ids = [] set_cold_start(init_timestamp_ns) + if custom_probes: + logger.debug(f"Injecting {len(custom_probes)} custom probes for this invocation") + self._enable_dynamic_instrumentation() + self._inject_probes(custom_probes) + if not should_use_extension: from datadog_lambda.metric import submit_invocations_metric @@ -289,8 +396,12 @@ def _before(self, event, context): except Exception as e: logger.error(format_err_with_traceback(e)) - def _after(self, event, context): + def _after(self, event, context, custom_probes=None): try: + if custom_probes: + self._remove_probes() + self._disable_dynamic_instrumentation() + from datadog_lambda.metric import submit_batch_item_failures_metric submit_batch_item_failures_metric(self.response, context) diff --git a/scripts/publish_layers.sh b/scripts/publish_layers.sh index 8c78093f..e42d0fc1 100755 --- a/scripts/publish_layers.sh +++ b/scripts/publish_layers.sh @@ -42,18 +42,18 @@ LAYER_PATHS=( ".layers/datadog_lambda_py-arm64-3.13.zip" ) AVAILABLE_LAYERS=( - "Datadog-Python38" - "Datadog-Python38-ARM" - "Datadog-Python39" - "Datadog-Python39-ARM" - "Datadog-Python310" - "Datadog-Python310-ARM" - "Datadog-Python311" - "Datadog-Python311-ARM" - "Datadog-Python312" - "Datadog-Python312-ARM" - "Datadog-Python313" - "Datadog-Python313-ARM" + "John-Datadog-Python38" + "John-Datadog-Python38-ARM" + "John-Datadog-Python39" + "John-Datadog-Python39-ARM" + "John-Datadog-Python310" + "John-Datadog-Python310-ARM" + "John-Datadog-Python311" + "John-Datadog-Python311-ARM" + "John-Datadog-Python312" + "John-Datadog-Python312-ARM" + "John-Datadog-Python313" + "John-Datadog-Python313-ARM" ) AVAILABLE_REGIONS=$(aws ec2 describe-regions | jq -r '.[] | .[] | .RegionName') @@ -63,15 +63,6 @@ PIDS=() # Makes sure any subprocesses will be terminated with this process trap "pkill -P $$; exit 1;" INT -# Check that the layer files exist -for layer_file in "${LAYER_PATHS[@]}" -do - if [ ! -f $layer_file ]; then - echo "Could not find $layer_file." - exit 1 - fi -done - # Determine the target regions if [ -z "$REGIONS" ]; then echo "Region not specified, running for all available regions." @@ -110,11 +101,7 @@ else echo "Layer version specified: $VERSION" fi -read -p "Ready to publish version $VERSION of layers ${LAYERS[*]} to regions ${REGIONS[*]} (y/n)?" CONT -if [ "$CONT" != "y" ]; then - echo "Exiting" - exit 1 -fi +echo "Publishing version $VERSION of layers ${LAYERS[*]} to regions ${REGIONS[*]}" index_of_layer() { layer_name=$1 diff --git a/scripts/publish_sandbox.sh b/scripts/publish_sandbox.sh index 8f7d1de5..4027fe22 100755 --- a/scripts/publish_sandbox.sh +++ b/scripts/publish_sandbox.sh @@ -3,9 +3,12 @@ # Usage: VERSION=5 ./scripts/publish_sandbox.sh set -e -./scripts/build_layers.sh -aws-vault exec sso-serverless-sandbox-account-admin -- ./scripts/sign_layers.sh sandbox -aws-vault exec sso-serverless-sandbox-account-admin -- ./scripts/publish_layers.sh +# Build only ARM64 layers for Python 3.12 +ARCH=arm64 PYTHON_VERSION=3.12 ./scripts/build_layers.sh +# Signing is commented out for sandbox - not needed for internal testing +## aws-vault exec sso-serverless-sandbox-account-admin -- ./scripts/sign_layers.sh sandbox +# Publish to us-east-1 only +LAYERS=John-Datadog-Python312-ARM VERSION=$VERSION REGIONS=us-east-1 aws-vault exec sso-serverless-sandbox-account-admin -- ./scripts/publish_layers.sh # Automatically create PR against github.com/DataDog/documentation # If you'd like to test, please uncomment the below line diff --git a/venv/pyvenv.cfg b/venv/pyvenv.cfg new file mode 100644 index 00000000..48b1e427 --- /dev/null +++ b/venv/pyvenv.cfg @@ -0,0 +1,3 @@ +home = /Users/john.chrostek/.pyenv/versions/3.10.13/bin +include-system-site-packages = false +version = 3.10.13