feat(aws-lambda-microvm): refresh identity on Flask run hook

ddtrace/__init__.py

@@ -6,30 +6,38 @@
 
 # Ensure we capture references to unpatched modules as early as possible
 import ddtrace.internal._unpatched  # noqa
-from ._logger import configure_ddtrace_logger
+from ._logger import configure_ddtrace_logger  # noqa: E402
 
 # configure ddtrace logger before other modules log
 configure_ddtrace_logger()  # noqa: E402
 
 # Enable telemetry writer and excepthook as early as possible to ensure we capture any exceptions from initialization
-import ddtrace.internal.telemetry  # noqa: E402
-from ddtrace.vendor import debtcollector
+# isort: off
+import ddtrace.internal.telemetry  # noqa: F401,E402
+from ddtrace.internal.runtime import listen_for_identity_refresh_hooks  # noqa: E402,I001
+
+# isort: on
+from ddtrace.vendor import debtcollector  # noqa: E402
 
 from ._monkey import patch  # noqa: E402
 from ._monkey import patch_all  # noqa: E402
+from .internal import core as _core  # noqa: E402
 from .internal.compat import PYTHON_VERSION_INFO  # noqa: E402
-from .internal.settings import env
-from .internal.settings._config import config
+from .internal.settings import env  # noqa: E402
+from .internal.settings._config import config  # noqa: E402
 from .internal.utils.deprecations import DDTraceDeprecationWarning  # noqa: E402
-from .version import __version__
+from .version import __version__  # noqa: E402
+
 
+# Register import-time hooks that depend on ddtrace.config being exported.
+listen_for_identity_refresh_hooks(_core.on)
 
 # TODO: Deprecate accessing tracer from ddtrace.__init__ module in v4.0
 if env.get("_DD_GLOBAL_TRACER_INIT", "true").lower() in ("1", "true"):
     from ddtrace.trace import tracer  # noqa: F401
 
 # Initialize DSM support and register DSM handlers (if enabled)
-import ddtrace.internal.datastreams as _  # noqa: F401
+import ddtrace.internal.datastreams as _  # noqa: F401,E402
 
 
 __all__ = [

ddtrace/internal/runtime/__init__.py

@@ -1,6 +1,7 @@
 import typing as t
 import uuid
 
+from ddtrace.internal.serverless import in_aws_lambda_microvm
 from ddtrace.internal.settings import env
 
 from .. import forksafe
@@ -12,6 +13,8 @@
     "get_runtime_id",
     "get_parent_runtime_id",
     "get_runtime_propagation_envs",
+    "listen_for_identity_refresh_hooks",
+    "maybe_refresh_identity",
     "refresh_identity",
 ]
 
@@ -75,6 +78,40 @@ def refresh_identity() -> None:
     _refresh_runtime_id()
 
 
+# Fixed platform request details for the AWS Lambda MicroVM /run lifecycle hook.
+MICROVM_RUN_HOOK_METHOD = "POST"
+MICROVM_RUN_HOOK_PATH = "/aws/lambda-microvms/runtime/v1/run"
+
+# Multiple request layers can observe the same /run hook. Refresh identity once per
+# process so a single logical MicroVM instance gets one runtime-id rotation.
+_IDENTITY_REFRESH_HOOK_REFRESHED = forksafe.Event()
+_IDENTITY_REFRESH_HOOK_REFRESH_LOCK = forksafe.Lock()
+
+
+def listen_for_identity_refresh_hooks(
+    on_event: t.Callable[[str, t.Callable[[t.Optional[str], t.Optional[str]], None]], None],
+) -> None:
+    """Refresh MicroVM identity from request events emitted before root span creation."""
+    if not in_aws_lambda_microvm():
+        return
+
+    on_event("web.request.starting", maybe_refresh_identity)
+
+
+def maybe_refresh_identity(method: t.Optional[str], path: t.Optional[str]) -> None:
+    """Call refresh_identity() if this request is the AWS Lambda MicroVM /run hook."""
+    if not method or not path:
+        return
+    if method != MICROVM_RUN_HOOK_METHOD or path != MICROVM_RUN_HOOK_PATH:
+        return
+
+    with _IDENTITY_REFRESH_HOOK_REFRESH_LOCK:
+        if _IDENTITY_REFRESH_HOOK_REFRESHED.is_set():
+            return
+        refresh_identity()
+        _IDENTITY_REFRESH_HOOK_REFRESHED.set()
+
+
 def get_runtime_id() -> str:
     """Return a unique string identifier for this runtime.
 

releasenotes/notes/aws-lambda-microvm-identity-refresh-3a672cd6bcbad16d.yaml

@@ -0,0 +1,6 @@
+---
+features:
+  - |
+    flask: Adds support for refreshing stable identifiers, including the runtime ID and
+    Remote Config client ID, when Flask applications start in AWS Lambda MicroVM
+    environments.

tests/contrib/flask/test_microvm_identity_refresh.py

@@ -3,17 +3,13 @@
 from ddtrace.contrib._events.web_framework import WebFrameworkEvents
 from ddtrace.contrib.internal.flask.patch import patched_wsgi_app
 from ddtrace.internal import core
+import ddtrace.internal.runtime as runtime
 
 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.
+    """patched_wsgi_app() must dispatch MicroVM requests before tracing starts.
 
     The matching logic itself is tested in tests/tracer/runtime/test_runtime_id.py.
     """
@@ -27,10 +23,38 @@ def test_microvm_run_hook_request(self):
             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)
+            res = self.client.post(runtime.MICROVM_RUN_HOOK_PATH)
 
         self.assertEqual(res.status_code, 404)
-        m.assert_any_call(WebFrameworkEvents.WEB_REQUEST_STARTING.value, ("POST", REQUEST_STARTING_PATH))
+        m.assert_any_call(WebFrameworkEvents.WEB_REQUEST_STARTING.value, ("POST", runtime.MICROVM_RUN_HOOK_PATH))
+
+    def test_microvm_run_hook_refreshes_identity(self):
+        event_name = WebFrameworkEvents.WEB_REQUEST_STARTING.value
+        core.reset_listeners(event_name, runtime.maybe_refresh_identity)
+        runtime._IDENTITY_REFRESH_HOOK_REFRESHED.clear()
+
+        try:
+            with (
+                mock.patch("ddtrace.contrib.internal.flask.patch.in_aws_lambda_microvm", return_value=True),
+                mock.patch("ddtrace.internal.runtime.in_aws_lambda_microvm", return_value=True),
+            ):
+                runtime.listen_for_identity_refresh_hooks(core.on)
+
+                runtime_id = runtime.get_runtime_id()
+
+                res = self.client.post(runtime.MICROVM_RUN_HOOK_PATH)
+
+                self.assertEqual(res.status_code, 404)
+                refreshed_runtime_id = runtime.get_runtime_id()
+                self.assertNotEqual(refreshed_runtime_id, runtime_id)
+
+                res = self.client.post(runtime.MICROVM_RUN_HOOK_PATH)
+
+                self.assertEqual(res.status_code, 404)
+                self.assertEqual(runtime.get_runtime_id(), refreshed_runtime_id)
+        finally:
+            core.reset_listeners(event_name, runtime.maybe_refresh_identity)
+            runtime._IDENTITY_REFRESH_HOOK_REFRESHED.clear()
 
     def test_other_request_does_not_dispatch_outside_microvm(self):
         @self.app.route("/")
@@ -47,21 +71,24 @@ def index():
         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):
+        script_name, hook_name = runtime.MICROVM_RUN_HOOK_PATH.rsplit("/", 1)
+        path_info = "/" + hook_name
+
         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})
+            res = self.client.post(path_info, environ_overrides={"SCRIPT_NAME": script_name})
 
         self.assertEqual(res.status_code, 404)
         m.assert_any_call(
             WebFrameworkEvents.WEB_REQUEST_STARTING.value,
-            ("POST", MICROVM_RUNTIME_PREFIX + MICROVM_RUN_HOOK_PATH),
+            (runtime.MICROVM_RUN_HOOK_METHOD, runtime.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": ""}
+        environ = {"REQUEST_METHOD": "POST", "PATH_INFO": runtime.MICROVM_RUN_HOOK_PATH, "SCRIPT_NAME": ""}
 
         def start_response(status, headers, exc_info=None):
             pass