diff --git a/ddtrace/__init__.py b/ddtrace/__init__.py index 3f493ce122a..f3b9aab1ebc 100644 --- a/ddtrace/__init__.py +++ b/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__ = [ diff --git a/ddtrace/internal/runtime/__init__.py b/ddtrace/internal/runtime/__init__.py index cdfe5888cbb..c16b100a9d4 100644 --- a/ddtrace/internal/runtime/__init__.py +++ b/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. diff --git a/releasenotes/notes/aws-lambda-microvm-identity-refresh-3a672cd6bcbad16d.yaml b/releasenotes/notes/aws-lambda-microvm-identity-refresh-3a672cd6bcbad16d.yaml new file mode 100644 index 00000000000..125be739a8c --- /dev/null +++ b/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. diff --git a/tests/contrib/flask/test_microvm_identity_refresh.py b/tests/contrib/flask/test_microvm_identity_refresh.py index d9643480a0a..7bd91810f69 100644 --- a/tests/contrib/flask/test_microvm_identity_refresh.py +++ b/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 diff --git a/tests/tracer/runtime/test_runtime_id.py b/tests/tracer/runtime/test_runtime_id.py index 0549461e3af..2df50d0cc75 100644 --- a/tests/tracer/runtime/test_runtime_id.py +++ b/tests/tracer/runtime/test_runtime_id.py @@ -1,3 +1,5 @@ +import os + import pytest @@ -305,3 +307,221 @@ def on_change(self, new_id): """ _, err, status, _ = run_python_code_in_subprocess(code) assert status == 0, err + + +@pytest.mark.parametrize("auto_enable_crashtracking", [False]) +def test_listen_for_identity_refresh_hooks_noop_does_not_import_core(monkeypatch, auto_enable_crashtracking): + import builtins + + import ddtrace.internal.runtime as runtime + + monkeypatch.setattr(runtime, "in_aws_lambda_microvm", lambda: False) + + real_import = builtins.__import__ + + def fail_core_import(name, *args, **kwargs): + fromlist = kwargs.get("fromlist", ()) + if len(args) >= 3: + fromlist = args[2] + if name == "ddtrace.internal" and "core" in fromlist: + raise AssertionError("listen_for_identity_refresh_hooks() imported core outside a MicroVM") + return real_import(name, *args, **kwargs) + + monkeypatch.setattr(builtins, "__import__", fail_core_import) + + runtime.listen_for_identity_refresh_hooks(lambda _event, _callback: None) + + +def test_web_request_starting_event_name_is_shared_with_contrib_event(): + from ddtrace.contrib._events.web_framework import WebFrameworkEvents + + assert "web.request.starting" == WebFrameworkEvents.WEB_REQUEST_STARTING.value + + +@pytest.mark.subprocess(env={"AWS_LAMBDA_MICROVM_IMAGE_ARN": "arn:aws:lambda:us-east-1::runtime:python3.12"}, err=None) +def test_import_ddtrace_in_microvm_environment(): + """MicroVM hook registration must not import contrib before ddtrace.config exists.""" + import ddtrace + + assert ddtrace.config is not None + + +@pytest.mark.subprocess(env={"AWS_LAMBDA_MICROVM_IMAGE_ARN": "arn:aws:lambda:us-east-1::runtime:python3.12"}, err=None) +def test_maybe_refresh_identity_matches_microvm_run_hook(): + """Only the exact AWS Lambda MicroVM /run hook request triggers a refresh.""" + import ddtrace.internal.runtime as runtime + + runtime_id = runtime.get_runtime_id() + + runtime.maybe_refresh_identity(runtime.MICROVM_RUN_HOOK_METHOD, runtime.MICROVM_RUN_HOOK_PATH) + + refreshed_runtime_id = runtime.get_runtime_id() + assert refreshed_runtime_id != runtime_id + + runtime.maybe_refresh_identity(runtime.MICROVM_RUN_HOOK_METHOD, runtime.MICROVM_RUN_HOOK_PATH) + + assert runtime.get_runtime_id() == refreshed_runtime_id + + +@pytest.mark.subprocess(env={"AWS_LAMBDA_MICROVM_IMAGE_ARN": "arn:aws:lambda:us-east-1::runtime:python3.12"}, err=None) +def test_identity_refresh_hook_runs_before_root_span_creation(): + """The pre-request hook must refresh runtime-id before a web root span reads it.""" + from ddtrace import tracer + from ddtrace.contrib._events.web_framework import WebFrameworkEvents + from ddtrace.internal import core + import ddtrace.internal.runtime as runtime + + runtime_id = runtime.get_runtime_id() + core.dispatch( + WebFrameworkEvents.WEB_REQUEST_STARTING.value, (runtime.MICROVM_RUN_HOOK_METHOD, runtime.MICROVM_RUN_HOOK_PATH) + ) + + refreshed_runtime_id = runtime.get_runtime_id() + assert refreshed_runtime_id != runtime_id + + with tracer.trace("web.request") as span: + assert span.get_tag("runtime-id") == refreshed_runtime_id + + core.dispatch( + WebFrameworkEvents.WEB_REQUEST_STARTING.value, (runtime.MICROVM_RUN_HOOK_METHOD, runtime.MICROVM_RUN_HOOK_PATH) + ) + + assert runtime.get_runtime_id() == refreshed_runtime_id + + +@pytest.mark.subprocess(env={"AWS_LAMBDA_MICROVM_IMAGE_ARN": "arn:aws:lambda:us-east-1::runtime:python3.12"}, err=None) +def test_maybe_refresh_identity_is_thread_safe(): + """Concurrent observations of the same MicroVM /run hook refresh identity once.""" + import threading + import time + + import ddtrace.internal.runtime as runtime + + calls = [] + + def refresh_identity(): + calls.append(1) + time.sleep(0.01) + + runtime.refresh_identity = refresh_identity + + workers = 16 + barrier = threading.Barrier(workers) + errors = [] + threads = [] + + def refresh_from_request_layer(): + try: + barrier.wait() + runtime.maybe_refresh_identity(runtime.MICROVM_RUN_HOOK_METHOD, runtime.MICROVM_RUN_HOOK_PATH) + except Exception as e: + errors.append(e) + + for _ in range(workers): + thread = threading.Thread(target=refresh_from_request_layer) + thread.start() + threads.append(thread) + + for thread in threads: + thread.join() + + assert errors == [] + assert len(calls) == 1 + + +@pytest.mark.subprocess(env={"AWS_LAMBDA_MICROVM_IMAGE_ARN": "arn:aws:lambda:us-east-1::runtime:python3.12"}, err=None) +def test_maybe_refresh_identity_retries_after_refresh_failure(): + """A failed identity refresh must not make later /run hooks no-op.""" + import pytest + + import ddtrace.internal.runtime as runtime + + calls = [] + + def refresh_identity(): + calls.append(1) + if len(calls) == 1: + raise RuntimeError("identity refresh failed") + + runtime.refresh_identity = refresh_identity + + with pytest.raises(RuntimeError, match="identity refresh failed"): + runtime.maybe_refresh_identity(runtime.MICROVM_RUN_HOOK_METHOD, runtime.MICROVM_RUN_HOOK_PATH) + + runtime.maybe_refresh_identity(runtime.MICROVM_RUN_HOOK_METHOD, runtime.MICROVM_RUN_HOOK_PATH) + + assert len(calls) == 2 + + +@pytest.mark.subprocess(env={"AWS_LAMBDA_MICROVM_IMAGE_ARN": "arn:aws:lambda:us-east-1::runtime:python3.12"}, err=None) +def test_microvm_refresh_guard_resets_after_fork(): + """A child process must not inherit a locked or already-refreshed MicroVM guard.""" + from ddtrace.internal import forksafe + import ddtrace.internal.runtime as runtime + + runtime._IDENTITY_REFRESH_HOOK_REFRESH_LOCK.acquire() + runtime._IDENTITY_REFRESH_HOOK_REFRESHED.set() + + forksafe.ddtrace_after_in_child() + + assert runtime._IDENTITY_REFRESH_HOOK_REFRESH_LOCK.acquire(False) + assert not runtime._IDENTITY_REFRESH_HOOK_REFRESHED.is_set() + + +@pytest.mark.subprocess(env={"AWS_LAMBDA_MICROVM_IMAGE_ARN": "arn:aws:lambda:us-east-1::runtime:python3.12"}, err=None) +def test_maybe_refresh_identity_ignores_other_requests(): + """A different method/path, or the /resume hook, must not trigger a refresh.""" + import ddtrace.internal.runtime as runtime + + runtime_id = runtime.get_runtime_id() + + runtime.maybe_refresh_identity("GET", runtime.MICROVM_RUN_HOOK_PATH) + runtime.maybe_refresh_identity(runtime.MICROVM_RUN_HOOK_METHOD, "/aws/lambda-microvms/runtime/v1/resume") + runtime.maybe_refresh_identity(runtime.MICROVM_RUN_HOOK_METHOD, "/some/other/path") + runtime.maybe_refresh_identity(None, runtime.MICROVM_RUN_HOOK_PATH) + runtime.maybe_refresh_identity(runtime.MICROVM_RUN_HOOK_METHOD, None) + + assert runtime.get_runtime_id() == runtime_id + + +@pytest.mark.subprocess(env={"AWS_LAMBDA_MICROVM_IMAGE_ARN": None}, err=None) +def test_listen_for_identity_refresh_hooks_noop_outside_microvm(): + """Outside a MicroVM, do not register the request-event listener.""" + from ddtrace.contrib._events.web_framework import WebFrameworkEvents + from ddtrace.internal import core + import ddtrace.internal.runtime as runtime + + core.reset_listeners(WebFrameworkEvents.WEB_REQUEST_STARTING.value) + runtime.listen_for_identity_refresh_hooks(core.on) + + runtime_id = runtime.get_runtime_id() + + core.dispatch( + WebFrameworkEvents.WEB_REQUEST_STARTING.value, (runtime.MICROVM_RUN_HOOK_METHOD, runtime.MICROVM_RUN_HOOK_PATH) + ) + + assert runtime.get_runtime_id() == runtime_id + + +@pytest.mark.parametrize("microvm_image_arn", ["", " "]) +def test_listen_for_identity_refresh_hooks_noop_for_blank_microvm_env(run_python_code_in_subprocess, microvm_image_arn): + """Blank MicroVM image ARN env values must not enable hook registration.""" + code = """ +from ddtrace.contrib._events.web_framework import WebFrameworkEvents +from ddtrace.internal import core +import ddtrace.internal.runtime as runtime + +core.reset_listeners(WebFrameworkEvents.WEB_REQUEST_STARTING.value) +runtime.listen_for_identity_refresh_hooks(core.on) + +runtime_id = runtime.get_runtime_id() +core.dispatch( + WebFrameworkEvents.WEB_REQUEST_STARTING.value, (runtime.MICROVM_RUN_HOOK_METHOD, runtime.MICROVM_RUN_HOOK_PATH) +) + +assert runtime.get_runtime_id() == runtime_id +""" + env = os.environ.copy() + env["AWS_LAMBDA_MICROVM_IMAGE_ARN"] = microvm_image_arn + _, err, status, _ = run_python_code_in_subprocess(code, env=env) + assert status == 0, err