diff --git a/ddtrace/internal/telemetry/writer.py b/ddtrace/internal/telemetry/writer.py index e5ab691ed75..863da58eafb 100644 --- a/ddtrace/internal/telemetry/writer.py +++ b/ddtrace/internal/telemetry/writer.py @@ -22,6 +22,7 @@ from ..runtime import get_ancestor_runtime_id from ..runtime import get_parent_runtime_id from ..runtime import get_runtime_id +from ..runtime import on_runtime_id_change from ..utils.formats import get_test_session_token from ..utils.version import version as tracer_version from .constants import TELEMETRY_APM_PRODUCT @@ -203,6 +204,9 @@ def __init__(self, agentless: Optional[bool] = None) -> None: # runtime's after_fork_child hook, which get_native_runtime() registered # during enable(), so the shared runtime is restarted before we rebuild). forksafe.register(self._fork_writer) + # Same rebuild, triggered by an explicit identity refresh (e.g. an AWS Lambda + # MicroVM /run hook) rather than an actual fork. + on_runtime_id_change(self._on_identity_refresh) get_logger("ddtrace").addHandler(DDTelemetryErrorHandler(self)) def _build_worker(self) -> "TelemetryWorker": @@ -915,6 +919,18 @@ def _fork_writer(self) -> None: # Re-discover dependencies from scratch so the child reports its own imports. self._dependency_tracker.reset() + def _on_identity_refresh(self, new_runtime_id: str) -> None: + # Same rebuild as _fork_writer(): the native worker bakes in get_runtime_id() at + # construction, so it must be dropped and lazily rebuilt on the next telemetry call. + # Unlike after a fork, the worker is still alive here -- it must be stopped (not just + # dropped) or it keeps heartbeating with the stale runtime ID until process shutdown. + if self._worker is not None: + try: + self._worker.stop(send_app_closing=False) + except Exception: + log.debug("Failed to stop the native telemetry worker during identity refresh", exc_info=True) + self._fork_writer() + def _telemetry_excepthook(self, tp, value, root_traceback) -> None: if root_traceback is not None: # Get the frame which raised the exception diff --git a/tests/telemetry/test_writer.py b/tests/telemetry/test_writer.py index 7bbede8e4c6..ce5b71a4c98 100644 --- a/tests/telemetry/test_writer.py +++ b/tests/telemetry/test_writer.py @@ -570,6 +570,61 @@ def test_telemetry_writer_agent_setup(): assert new_telemetry_writer._agentless is False +def test_identity_refresh_rebuilds_native_worker(): + """Same rebuild as after a fork: the native worker bakes in get_runtime_id() at construction.""" + with override_global_config( + {"_dd_site": "datad0g.com", "_dd_api_key": "foobarkey", "_ci_visibility_agentless_enabled": False} + ): + writer = ddtrace.internal.telemetry.TelemetryWriter(agentless=False) + assert writer._worker is not None + + writer._on_identity_refresh("some-new-runtime-id") + + assert writer._worker is None + assert writer.started is False + + +def test_identity_refresh_stops_live_worker_before_dropping(): + """Unlike after a fork, the worker is still alive here and must be explicitly stopped, or it + keeps heartbeating with the stale runtime ID until process shutdown. + """ + with override_global_config( + {"_dd_site": "datad0g.com", "_dd_api_key": "foobarkey", "_ci_visibility_agentless_enabled": False} + ): + writer = ddtrace.internal.telemetry.TelemetryWriter(agentless=False) + assert writer._worker is not None + + # TelemetryWorker is a native extension type -- its methods can't be patched in place, + # so swap in a mock to observe the stop() call instead. + fake_worker = mock.Mock() + writer._worker = fake_worker + + writer._on_identity_refresh("some-new-runtime-id") + + fake_worker.stop.assert_called_once_with(send_app_closing=False) + assert writer._worker is None + + +@pytest.mark.subprocess( + env={"DD_SITE": "datad0g.com", "DD_API_KEY": "foobarkey", "DD_CIVISIBILITY_AGENTLESS_ENABLED": "false"} +) +def test_identity_refresh_wired_to_runtime_id_change(): + """Drives the refresh through runtime.refresh_identity() instead of calling + _on_identity_refresh directly (as the test above does), so a dropped + on_runtime_id_change() subscription would actually fail this. + """ + from ddtrace.internal import runtime + import ddtrace.internal.telemetry + + writer = ddtrace.internal.telemetry.TelemetryWriter(agentless=False) + assert writer._worker is not None + + runtime.refresh_identity() + + assert writer._worker is None + assert writer.started is False + + @pytest.mark.parametrize( "env_agentless,arg_agentless", [