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

ddtrace/internal/runtime/__init__.py

@@ -1,3 +1,4 @@
+import threading
 import typing as t
 import uuid
 import weakref
@@ -17,6 +18,8 @@
     "get_parent_runtime_id",
     "get_runtime_propagation_envs",
     "refresh_identity",
+    "maybe_refresh_identity",
+    "listen_for_identity_refresh_hooks",
 ]
 
 
@@ -43,14 +46,19 @@ def _generate_runtime_id() -> str:
 
 
 def on_runtime_id_change(cb: t.Callable[[str], None]) -> None:
-    """Register a callback to be called when refresh_identity() runs.
-
-    refresh_identity() is the non-fork trigger for a new logical process
-    instance. It is deliberately not called after a plain fork: forked children
-    already get a fresh runtime ID silently (see _set_runtime_id()), and code
-    that needs to react to a fork specifically should use forksafe.register().
-    Only a weak reference to cb is kept, so the caller must keep it alive for
-    it to keep firing.
+    """Register a callback to be called when ``refresh_identity()`` runs.
+
+    ``refresh_identity()`` is the non-fork trigger (e.g. an AWS Lambda MicroVM
+    ``/run`` hook) for a new logical process instance. It is deliberately not
+    called after a plain fork: forked children already get a fresh runtime ID
+    silently (see ``_set_runtime_id()``), and code that needs to react to a
+    fork specifically should use ``forksafe.register()`` instead -- unlike a
+    fork, no state (spans, SHM-backed native clients, ...) was inherited from
+    a different process here, so subscribers generally only need to rebuild
+    what bakes the runtime/client id in at construction, not reset buffers or
+    handles the way a fork hook would. Only a weak reference to ``cb`` is
+    kept, so the caller must keep it alive (e.g. by registering a bound
+    method of a long-lived object) for it to keep firing.
     """
     global _ON_RUNTIME_ID_CHANGE
     try:
@@ -83,8 +91,8 @@ def _notify_runtime_id_subscribers() -> None:
         try:
             cb(_RUNTIME_ID)
         except Exception:
-            # One broken subscriber must not prevent other subscribers from seeing the
-            # refreshed runtime ID.
+            # This can run on a web framework's request-dispatch path (maybe_refresh_identity());
+            # one broken subscriber must not take down the request or block the others.
             log.debug("Error notifying on_runtime_id_change() subscriber", exc_info=True)
     _ON_RUNTIME_ID_CHANGE -= dead
 
@@ -108,16 +116,72 @@ def _set_runtime_id() -> None:
 def refresh_identity() -> None:
     """Regenerate the runtime ID without recording fork lineage.
 
-    Unlike a fork, this does not update _PARENT_RUNTIME_ID / _ANCESTOR_RUNTIME_ID:
-    the previous runtime ID was not a real parent process, so recording it there
-    would make get_process_role() and friends misreport a fork lineage that never
-    existed. Use this when a new logical process instance is created by a mechanism
-    other than fork().
+    Unlike a fork, this does not update ``_PARENT_RUNTIME_ID`` /
+    ``_ANCESTOR_RUNTIME_ID``: the previous runtime ID was not a real parent
+    process, so recording it there would make ``get_process_role()`` and
+    friends misreport a fork lineage that never existed. Use this when a new
+    logical process instance is created by a mechanism other than ``fork()``
+    -- e.g. an AWS Lambda MicroVM instance launched from a shared image
+    snapshot.
     """
     _regenerate_runtime_id()
     _notify_runtime_id_subscribers()
 
 
+# Fixed platform path for the AWS Lambda MicroVM "/run" lifecycle hook. Never the "/resume"
+# hook path -- see refresh_identity()'s docstring for why. Exported (no leading underscore)
+# so tests can reference the same constant instead of duplicating the literal.
+MICROVM_RUN_HOOK_METHOD = "POST"
+MICROVM_RUN_HOOK_PATH = "/aws/lambda-microvms/runtime/v1/run"
+
+# Same env var #18017 uses to detect a MicroVM at runtime (rand64bits() OS-entropy fallback).
+# Read once at import, like that fix does, so listener registration is skipped outside a
+# MicroVM. This exact method+path is otherwise just an unauthenticated trigger on every
+# ddtrace user's request-dispatch path, MicroVM or not.
+_IS_AWS_LAMBDA_MICROVM = env.get("AWS_LAMBDA_MICROVM_IMAGE_ARN") is not None
+# Multiple instrumented request layers can observe the same MicroVM /run hook
+# (for example, Werkzeug's http.server layer plus Flask). Refresh identity once
+# per process so a single logical MicroVM instance gets one runtime-id rotation.
+_IDENTITY_REFRESH_HOOK_REFRESHED = threading.Event()
+_IDENTITY_REFRESH_HOOK_REFRESH_LOCK = threading.Lock()
+
+
+def listen_for_identity_refresh_hooks() -> None:
+    """Refresh MicroVM identity from request events emitted before root span creation."""
+    if not _IS_AWS_LAMBDA_MICROVM:
+        return
+
+    from ddtrace.internal import core
+
+    core.on(core.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.
+
+    Called from every instrumented web framework's request-dispatch patch in a MicroVM
+    with that request's method and path, so the platform-defined hook path only needs
+    to be matched in one place.
+    ``method``/``path`` may be ``None`` -- some callers read them straight off a raw
+    request/environ mapping (e.g. ``environ.get("REQUEST_METHOD")``) that has no guarantee
+    either key is present.
+    """
+    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
+        _IDENTITY_REFRESH_HOOK_REFRESHED.set()
+
+    refresh_identity()
+
+
+listen_for_identity_refresh_hooks()
+
+
 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,7 @@
+---
+fixes:
+  - |
+    stable id uniqueness: Regenerates stable identifiers, including the runtime ID and
+    Remote Config client ID, when a new AWS Lambda MicroVM instance starts to ensure
+    they remain unique. This behavior only applies to AWS Lambda MicroVM deployments;
+    behavior in other environments is unchanged.

tests/contrib/asgi/test_microvm_identity_refresh.py

@@ -4,13 +4,11 @@
 
 from ddtrace.contrib.internal.asgi.middleware import TraceMiddleware
 from ddtrace.internal import core
+from ddtrace.internal.runtime import MICROVM_RUN_HOOK_PATH
 
 from .test_asgi import basic_app
 
 
-REQUEST_STARTING_PATH = "/web-request-starting"
-
-
 def _scope(method, path):
     return {
         "client": ("127.0.0.1", 32767),
@@ -32,14 +30,14 @@ async def test_microvm_run_hook_request(test_spans):
     three.
     """
     app = TraceMiddleware(basic_app)
-    instance = ApplicationCommunicator(app, _scope("POST", REQUEST_STARTING_PATH))
+    instance = ApplicationCommunicator(app, _scope("POST", MICROVM_RUN_HOOK_PATH))
 
     with mock.patch("ddtrace.contrib.internal.asgi.middleware.core.dispatch", wraps=core.dispatch) as m:
         await instance.send_input({"type": "http.request", "body": b""})
         await instance.receive_output(1)
         await instance.receive_output(1)
 
-    m.assert_any_call(core.WEB_REQUEST_STARTING, ("POST", REQUEST_STARTING_PATH))
+    m.assert_any_call(core.WEB_REQUEST_STARTING, ("POST", MICROVM_RUN_HOOK_PATH))
 
 
 @pytest.mark.asyncio
@@ -61,7 +59,7 @@ async def test_sub_app_does_not_double_refresh(test_spans):
     (matches the existing not-is_subapp guard around route collection/distributed headers).
     """
     app = TraceMiddleware(basic_app)
-    scope = _scope("POST", REQUEST_STARTING_PATH)
+    scope = _scope("POST", MICROVM_RUN_HOOK_PATH)
     # marks this as a sub-app request, per TraceMiddleware.__call__; request_spans matches
     # the shape _on_asgi_request always creates the dict with (ddtrace/_trace/trace_handlers.py)
     scope["datadog"] = {"request_spans": []}

tests/contrib/bottle/test_microvm_identity_refresh.py

@@ -5,12 +5,10 @@
 from ddtrace.contrib.internal.bottle.patch import TracePlugin
 from ddtrace.contrib.internal.bottle.patch import patch
 from ddtrace.internal import core
+from ddtrace.internal.runtime import MICROVM_RUN_HOOK_PATH
 from tests.utils import TracerTestCase
 
 
-REQUEST_STARTING_PATH = "/web-request-starting"
-
-
 class BottleMicrovmIdentityRefreshTestCase(TracerTestCase):
     """traced_wsgi() wraps Bottle.wsgi() -- the WSGI entry point, run before routing -- so it
     emits every request's real method/path before request tracing starts. Bottle has no
@@ -34,10 +32,10 @@ def test_microvm_run_hook_request(self):
         self._trace_app()
 
         with mock.patch("ddtrace.contrib.internal.bottle.trace.core.dispatch", wraps=core.dispatch) as m:
-            resp = self.app.post(REQUEST_STARTING_PATH, expect_errors=True)
+            resp = self.app.post(MICROVM_RUN_HOOK_PATH, expect_errors=True)
 
         assert resp.status_int == 404
-        m.assert_any_call(core.WEB_REQUEST_STARTING, ("POST", REQUEST_STARTING_PATH))
+        m.assert_any_call(core.WEB_REQUEST_STARTING, ("POST", MICROVM_RUN_HOOK_PATH))
 
     def test_other_request(self):
         @self.app.route("/hi/<name>")

tests/contrib/cherrypy/test_microvm_identity_refresh.py

@@ -4,14 +4,12 @@
 
 from ddtrace.contrib.internal.cherrypy.patch import TraceMiddleware
 from ddtrace.internal import core
+from ddtrace.internal.runtime import MICROVM_RUN_HOOK_PATH
 from tests.utils import TracerTestCase
 
 from .web import StubApp
 
 
-REQUEST_STARTING_PATH = "/web-request-starting"
-
-
 class CherrypyMicrovmIdentityRefreshTestCase(TracerTestCase, helper.CPWebCase):
     """TraceTool._on_start_resource() must dispatch method/path before request tracing starts.
 
@@ -38,10 +36,10 @@ def test_microvm_run_hook_request(self):
         404 (see test_404 in test_middleware.py).
         """
         with mock.patch("ddtrace.contrib.internal.cherrypy.patch.core.dispatch", wraps=core.dispatch) as m:
-            self.getPage(REQUEST_STARTING_PATH, method="POST")
+            self.getPage(MICROVM_RUN_HOOK_PATH, method="POST")
 
         self.assertStatus("404 Not Found")
-        m.assert_any_call(core.WEB_REQUEST_STARTING, ("POST", REQUEST_STARTING_PATH))
+        m.assert_any_call(core.WEB_REQUEST_STARTING, ("POST", MICROVM_RUN_HOOK_PATH))
 
     def test_other_request(self):
         with mock.patch("ddtrace.contrib.internal.cherrypy.patch.core.dispatch", wraps=core.dispatch) as m:

tests/contrib/django/test_microvm_identity_refresh.py

@@ -1,8 +1,7 @@
 import mock
 
 from ddtrace.internal import core
-
-REQUEST_STARTING_PATH = "/web-request-starting"
+from ddtrace.internal.runtime import MICROVM_RUN_HOOK_PATH
 
 
 def test_microvm_run_hook_request(client):
@@ -12,10 +11,10 @@ def test_microvm_run_hook_request(client):
     test_django_request_not_found).
     """
     with mock.patch("ddtrace.contrib.internal.django.response.core.dispatch", wraps=core.dispatch) as m:
-        resp = client.post(REQUEST_STARTING_PATH)
+        resp = client.post(MICROVM_RUN_HOOK_PATH)
 
     assert resp.status_code == 404
-    m.assert_any_call(core.WEB_REQUEST_STARTING, ("POST", REQUEST_STARTING_PATH))
+    m.assert_any_call(core.WEB_REQUEST_STARTING, ("POST", MICROVM_RUN_HOOK_PATH))
 
 
 def test_other_request(client):

tests/contrib/falcon/test_microvm_identity_refresh.py

@@ -2,13 +2,11 @@
 import mock
 
 from ddtrace.internal import core
+from ddtrace.internal.runtime import MICROVM_RUN_HOOK_PATH
 
 from .app import get_app
 
 
-REQUEST_STARTING_PATH = "/web-request-starting"
-
-
 def _client():
     return testing.TestClient(get_app())
 
@@ -20,10 +18,10 @@ def test_microvm_run_hook_request():
     fires on the 404 (see test_404 in test_suite.py).
     """
     with mock.patch("ddtrace.contrib.internal.falcon.middleware.core.dispatch", wraps=core.dispatch) as m:
-        response = _client().simulate_post(REQUEST_STARTING_PATH)
+        response = _client().simulate_post(MICROVM_RUN_HOOK_PATH)
 
     assert response.status[:3] == "404"
-    m.assert_any_call(core.WEB_REQUEST_STARTING, ("POST", REQUEST_STARTING_PATH))
+    m.assert_any_call(core.WEB_REQUEST_STARTING, ("POST", MICROVM_RUN_HOOK_PATH))
 
 
 def test_other_request():

tests/contrib/flask/test_microvm_identity_refresh.py

@@ -2,13 +2,11 @@
 
 from ddtrace.contrib.internal.flask.patch import patched_wsgi_app
 from ddtrace.internal import core
+from ddtrace.internal.runtime import MICROVM_RUN_HOOK_PATH
 
 from . import BaseFlaskTestCase
 
 
-REQUEST_STARTING_PATH = "/web-request-starting"
-
-
 class FlaskMicrovmIdentityRefreshTestCase(BaseFlaskTestCase):
     """patched_wsgi_app() must dispatch every request's method/path before tracing starts.
 
@@ -21,10 +19,10 @@ def test_microvm_run_hook_request(self):
         route matches doesn't change what gets dispatched).
         """
         with 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(MICROVM_RUN_HOOK_PATH)
 
         self.assertEqual(res.status_code, 404)
-        m.assert_any_call(core.WEB_REQUEST_STARTING, ("POST", REQUEST_STARTING_PATH))
+        m.assert_any_call(core.WEB_REQUEST_STARTING, ("POST", MICROVM_RUN_HOOK_PATH))
 
     def test_other_request(self):
         @self.app.route("/")
@@ -39,7 +37,7 @@ def index():
 
     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": MICROVM_RUN_HOOK_PATH, "SCRIPT_NAME": ""}
 
         def start_response(status, headers, exc_info=None):
             pass

tests/contrib/http_server/test_microvm_identity_refresh.py

@@ -7,8 +7,7 @@
 from ddtrace.contrib.internal.http_server.patch import patch
 from ddtrace.contrib.internal.http_server.patch import unpatch
 from ddtrace.internal import core
-
-REQUEST_STARTING_PATH = "/web-request-starting"
+from ddtrace.internal.runtime import MICROVM_RUN_HOOK_PATH
 
 
 def _handler_for(method, path):
@@ -35,10 +34,10 @@ def test_microvm_run_hook_request():
     supported web framework.
     """
     with mock.patch("ddtrace.contrib.internal.http_server.patch.core.dispatch", wraps=core.dispatch) as m:
-        parsed = _handler_for("POST", REQUEST_STARTING_PATH).parse_request()
+        parsed = _handler_for("POST", MICROVM_RUN_HOOK_PATH).parse_request()
 
     assert parsed is True
-    m.assert_any_call(core.WEB_REQUEST_STARTING, ("POST", REQUEST_STARTING_PATH))
+    m.assert_any_call(core.WEB_REQUEST_STARTING, ("POST", MICROVM_RUN_HOOK_PATH))
 
 
 def test_other_request():

tests/contrib/molten/test_microvm_identity_refresh.py

@@ -5,12 +5,10 @@
 from ddtrace.contrib.internal.molten.patch import patch
 from ddtrace.contrib.internal.molten.patch import unpatch
 from ddtrace.internal import core
+from ddtrace.internal.runtime import MICROVM_RUN_HOOK_PATH
 from tests.utils import TracerTestCase
 
 
-REQUEST_STARTING_PATH = "/web-request-starting"
-
-
 def greet():
     return "Greetings"
 
@@ -33,10 +31,10 @@ def test_microvm_run_hook_request(self):
         point, ahead of molten's router, so it still fires on the 404.
         """
         with mock.patch("ddtrace.contrib.internal.molten.patch.core.dispatch", wraps=core.dispatch) as m:
-            response = self.client.request("POST", REQUEST_STARTING_PATH)
+            response = self.client.request("POST", MICROVM_RUN_HOOK_PATH)
 
         self.assertEqual(response.status_code, 404)
-        m.assert_any_call(core.WEB_REQUEST_STARTING, ("POST", REQUEST_STARTING_PATH))
+        m.assert_any_call(core.WEB_REQUEST_STARTING, ("POST", MICROVM_RUN_HOOK_PATH))
 
     def test_other_request(self):
         with mock.patch("ddtrace.contrib.internal.molten.patch.core.dispatch", wraps=core.dispatch) as m: