chore(telemetry): rebuild worker on identity refresh

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

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",
     [