chore(flask): emit MicroVM request-starting event

ddtrace/contrib/_events/web_framework.py

@@ -13,6 +13,7 @@
 
 class WebFrameworkEvents(str, Enum):
     WEB_REQUEST = "web.request"
+    WEB_REQUEST_STARTING = "web.request.starting"
 
 
 @dataclass

ddtrace/contrib/internal/flask/patch.py

@@ -7,6 +7,7 @@
 from werkzeug.exceptions import NotFound
 
 from ddtrace.contrib import trace_utils
+from ddtrace.contrib._events.web_framework import WebFrameworkEvents
 from ddtrace.ext import SpanTypes
 from ddtrace.internal import core
 from ddtrace.internal.constants import COMPONENT
@@ -15,6 +16,7 @@
 from ddtrace.internal.schema import schematize_service_name
 from ddtrace.internal.schema import schematize_url_operation
 from ddtrace.internal.schema.span_attribute_schema import SpanDirection
+from ddtrace.internal.serverless import in_aws_lambda_microvm
 from ddtrace.internal.settings.appsec_telemetry import config as appsec_telemetry_config
 from ddtrace.internal.span_bus import span_from_context
 from ddtrace.internal.utils import get_blocked
@@ -391,8 +393,14 @@ def unpatch():
 
 def patched_wsgi_app(wrapped, instance, args, kwargs):
     environ, start_response = args
+    script_name = (environ.get("SCRIPT_NAME") or "").rstrip("/")
+    if in_aws_lambda_microvm():
+        path_info = environ.get("PATH_INFO") or ""
+        core.dispatch(
+            WebFrameworkEvents.WEB_REQUEST_STARTING.value, (environ.get("REQUEST_METHOD"), script_name + path_info)
+        )
     # Registration is gated on asm_config, not tracing — keep this above the tracing short-circuit.
-    _collect_routes_once(instance, environ.get("SCRIPT_NAME") or "")
+    _collect_routes_once(instance, script_name)
     if not is_tracing_enabled():
         return wrapped(*args, **kwargs)
     middleware = _FlaskWSGIMiddleware(wrapped, None, config.flask)

ddtrace/internal/serverless/__init__.py

@@ -20,6 +20,12 @@ def has_aws_lambda_agent_extension():
     return path.exists("/opt/extensions/datadog-agent")
 
 
+def in_aws_lambda_microvm():
+    # type: () -> bool
+    """Returns whether the environment is an AWS Lambda MicroVM."""
+    return bool(env.get("AWS_LAMBDA_MICROVM_IMAGE_ARN", "").strip())
+
+
 def in_gcp_function():
     # type: () -> bool
     """Returns whether the environment is a GCP Function.

tests/contrib/flask/test_microvm_identity_refresh.py

@@ -0,0 +1,93 @@
+import mock
+
+from ddtrace.contrib._events.web_framework import WebFrameworkEvents
+from ddtrace.contrib.internal.flask.patch import patched_wsgi_app
+from ddtrace.internal import core
+
+from . import BaseFlaskTestCase
+
+
+REQUEST_STARTING_PATH = "/web-request-starting"
+MICROVM_RUNTIME_PREFIX = "/aws/lambda-microvms/runtime/v1"
+MICROVM_RUN_HOOK_PATH = "/run"
+
+
+class FlaskMicrovmIdentityRefreshTestCase(BaseFlaskTestCase):
+    """patched_wsgi_app() must dispatch every request's method/path before tracing starts.
+
+    The matching logic itself is tested in tests/tracer/runtime/test_runtime_id.py.
+    """
+
+    def test_microvm_run_hook_request(self):
+        """No route is registered at the hook path: wsgi_app() runs before routing, so it
+        must fire even on a 404 -- the stronger, more general form of this check (whether the
+        route matches doesn't change what gets dispatched).
+        """
+        with (
+            mock.patch("ddtrace.contrib.internal.flask.patch.in_aws_lambda_microvm", return_value=True),
+            mock.patch("ddtrace.contrib.internal.flask.patch.core.dispatch", wraps=core.dispatch) as m,
+        ):
+            res = self.client.post(REQUEST_STARTING_PATH)
+
+        self.assertEqual(res.status_code, 404)
+        m.assert_any_call(WebFrameworkEvents.WEB_REQUEST_STARTING.value, ("POST", REQUEST_STARTING_PATH))
+
+    def test_other_request_does_not_dispatch_outside_microvm(self):
+        @self.app.route("/")
+        def index():
+            return "ok", 200
+
+        with (
+            mock.patch("ddtrace.contrib.internal.flask.patch.in_aws_lambda_microvm", return_value=False),
+            mock.patch("ddtrace.contrib.internal.flask.patch.core.dispatch", wraps=core.dispatch) as m,
+        ):
+            res = self.client.get("/")
+
+        self.assertEqual(res.status_code, 200)
+        assert all(call.args[0] != WebFrameworkEvents.WEB_REQUEST_STARTING.value for call in m.call_args_list)
+
+    def test_dispatches_script_name_prefixed_request_path(self):
+        with (
+            mock.patch("ddtrace.contrib.internal.flask.patch.in_aws_lambda_microvm", return_value=True),
+            mock.patch("ddtrace.contrib.internal.flask.patch.core.dispatch", wraps=core.dispatch) as m,
+        ):
+            res = self.client.post(MICROVM_RUN_HOOK_PATH, environ_overrides={"SCRIPT_NAME": MICROVM_RUNTIME_PREFIX})
+
+        self.assertEqual(res.status_code, 404)
+        m.assert_any_call(
+            WebFrameworkEvents.WEB_REQUEST_STARTING.value,
+            ("POST", MICROVM_RUNTIME_PREFIX + MICROVM_RUN_HOOK_PATH),
+        )
+
+    def test_pre_request_event_dispatches_before_wsgi_middleware(self):
+        events = []
+        environ = {"REQUEST_METHOD": "POST", "PATH_INFO": REQUEST_STARTING_PATH, "SCRIPT_NAME": ""}
+
+        def start_response(status, headers, exc_info=None):
+            pass
+
+        def wrapped(environ, start_response):
+            return []
+
+        def dispatch(name, args):
+            if name == WebFrameworkEvents.WEB_REQUEST_STARTING.value:
+                events.append("starting")
+
+        class WSGIMiddleware:
+            def __init__(self, app, tracer, integration_config):
+                pass
+
+            def __call__(self, environ, start_response):
+                events.append("middleware")
+                return []
+
+        with (
+            mock.patch("ddtrace.contrib.internal.flask.patch.in_aws_lambda_microvm", return_value=True),
+            mock.patch("ddtrace.contrib.internal.flask.patch.core.dispatch", side_effect=dispatch),
+            mock.patch("ddtrace.contrib.internal.flask.patch._collect_routes_once"),
+            mock.patch("ddtrace.contrib.internal.flask.patch.is_tracing_enabled", return_value=True),
+            mock.patch("ddtrace.contrib.internal.flask.patch._FlaskWSGIMiddleware", WSGIMiddleware),
+        ):
+            patched_wsgi_app(wrapped, self.app, (environ, start_response), {})
+
+        assert events == ["starting", "middleware"]

tests/internal/test_serverless.py

@@ -1,5 +1,6 @@
 import pytest
 
+from ddtrace.internal.serverless import in_aws_lambda_microvm
 from ddtrace.internal.serverless import in_azure_function
 from ddtrace.internal.serverless import in_gcp_function
 from tests.utils import override_env
@@ -19,6 +20,21 @@ def test_not_gcp_function():
     assert in_gcp_function() is False
 
 
+def test_is_aws_lambda_microvm():
+    with override_env(dict(AWS_LAMBDA_MICROVM_IMAGE_ARN="arn:aws:lambda:us-east-1::runtime:python3.12")):
+        assert in_aws_lambda_microvm() is True
+
+
+def test_not_aws_lambda_microvm():
+    assert in_aws_lambda_microvm() is False
+
+
+@pytest.mark.parametrize("image_arn", ["", "   "])
+def test_blank_aws_lambda_microvm_image_arn(image_arn):
+    with override_env(dict(AWS_LAMBDA_MICROVM_IMAGE_ARN=image_arn)):
+        assert in_aws_lambda_microvm() is False
+
+
 def test_is_azure_function():
     with override_env(dict(FUNCTIONS_WORKER_RUNTIME="python", FUNCTIONS_EXTENSION_VERSION="2")):
         assert in_azure_function() is True