fix(telemetry): rebuild worker on identity refresh

ddtrace/internal/_runtime_id.py

    _refresh_runtime_id()
    _refresh_runtime_id()
@@ -62,6 +62,11 @@ def on_runtime_identity_refresh(cb: t.Callable[[str], None]) -> None:
     _ON_RUNTIME_IDENTITY_REFRESH.add(cb)
 
 
+def remove_runtime_identity_refresh(cb: t.Callable[[str], None]) -> None:
+    """Unregister a callback for explicit runtime identity refreshes."""
+    _ON_RUNTIME_IDENTITY_REFRESH.discard(cb)
+
+
 def _notify_runtime_id_callbacks(callbacks: t.Set[t.Callable[[str], None]]) -> None:  # noqa: UP006
     for cb in list(callbacks):
         try:

    _refresh_runtime_id()

ddtrace/internal/runtime/__init__.py

@@ -7,6 +7,7 @@
 from ddtrace.internal._runtime_id import on_runtime_id_change
 from ddtrace.internal._runtime_id import on_runtime_identity_refresh
 from ddtrace.internal._runtime_id import refresh_identity
+from ddtrace.internal._runtime_id import remove_runtime_identity_refresh
 
 
 __all__ = [
@@ -17,6 +18,7 @@
     "get_runtime_propagation_envs",
     "on_runtime_id_change",
     "on_runtime_identity_refresh",
+    "remove_runtime_identity_refresh",
     "maybe_refresh_identity",
     "refresh_identity",
 ]

ddtrace/internal/telemetry/dependency.py

@@ -130,6 +130,10 @@ def mark_all_metadata_sent(self) -> None:
         for m in self.metadata:
             m._mark_sent()
 
+    def reset_for_refresh(self) -> None:
+        """Mark this dependency for reporting to a new worker."""
+        self._initial_report_sent = False
+
     def add_metadata(self, cve_id: str, path: str = "", symbol: str = "", line: int = 0) -> bool:
         """Add or update reachability metadata for a CVE.
 

ddtrace/internal/telemetry/dependency_tracker.py

@@ -57,6 +57,7 @@ class DependencyTracker:
 
     def __init__(self) -> None:
         self._imported_dependencies: dict[str, DependencyEntry] = {}
+        self._report_all = False
         self._modules_already_imported: set[str] = set()
         self._lock = Lock()
 
@@ -77,15 +78,12 @@ def collect_report(self) -> Optional[list[dict[str, Any]]]:
             new_keys = {_normalize_dep_name(d["name"]) for d in new_deps}
             self._mark_sent(new_keys)
 
-            # Skip the re-report scan when SCA is disabled.
-            # Without SCA, no entry will ever have unsent metadata, so the
-            # scan over all _imported_dependencies is pure overhead (~887us
-            # at 10K deps).  Only entries created by the SCA hook or with
-            # metadata attached can trigger needs_report() after initial send.
-            if not appsec_telemetry_config.SCA_ENABLED:
+            # Skip re-report scanning when SCA is disabled, except after identity refresh.
+            if not appsec_telemetry_config.SCA_ENABLED and not self._report_all:
                 return new_deps if new_deps else None
 
             re_report_deps = self._collect_rereports(new_keys)
+            self._report_all = False
             all_deps = new_deps + re_report_deps
             return all_deps if all_deps else None
 
@@ -187,11 +185,19 @@ def enable_sca_metadata(self) -> None:
                 if entry.metadata is None:
                     entry.metadata = []
 
+    def refresh(self) -> None:
+        """Preserve dependency metadata while scheduling a full report for a new worker."""
+        with self._lock:
+            for dependency in self._imported_dependencies.values():
+                dependency.reset_for_refresh()
+            self._report_all = True
+
     def reset(self) -> None:
         """Reset all state (used on fork / queue reset)."""
         with self._lock:
             self._imported_dependencies = {}
             self._modules_already_imported = set()
+            self._report_all = False
 
 
 def update_imported_dependencies(

ddtrace/internal/telemetry/noop_writer.py

@@ -108,13 +108,19 @@ def set_test_session_token(self, token: Optional[str]) -> None:
     def _restart_sequence(self) -> None:
         pass
 
+    def _refresh_runtime_identity(self, _runtime_id: str) -> None:
+        pass
+
     def _fork_writer(self) -> None:
         pass
 
     def _report_dependencies(self) -> Optional[list[dict[str, Any]]]:
         return None
 
-    def _subscribe_worker_changes(self, callback: Any) -> None:
+    def _subscribe_worker_changes(self, callback: Any, expected_worker: Any) -> None:
+        pass
+
+    def _unsubscribe_worker_changes(self, callback: Any) -> None:
         pass
 
     def periodic(self, force_flush: bool = False) -> None:

ddtrace/internal/writer/writer.py

                    self._exporter = new_exporter
            old_exporter = self._exporter
            new_exporter = self._create_exporter()
            with self._exporter_lock:
                # A refresh may have dropped the exporter since the check above; a replacement
                # installed now would be skipped by on_shutdown() and leak.
                discarded = self._discarded_by_refresh
                if not discarded:
                    self._exporter = new_exporter
            with self._exporter_lock:
                old_exporter = self._exporter
            new_exporter = self._create_exporter()
            with self._exporter_lock:
                # A refresh may have dropped the exporter since the check above; a replacement
                # installed now would be skipped by on_shutdown() and leak.
                discarded = self._discarded_by_refresh
                if not discarded:
                    self._exporter = new_exporter
@@ -2,6 +2,7 @@
 import binascii
 from collections import defaultdict
 from collections.abc import Sequence
+from functools import partial
 import gzip
 import os
 import socket
@@ -74,6 +75,7 @@
 
 if TYPE_CHECKING:  # pragma: no cover
     from ddtrace.internal.http import HTTPConnection  # noqa:F401
+    from ddtrace.internal.telemetry.writer import TelemetryWriter
     from ddtrace.vendor.dogstatsd import DogStatsd
 
 
@@ -899,6 +901,12 @@ def __init__(
         self._accepting_writes = True
         self._exporter_dropped = False
         self._owner_pid = os.getpid()
+        self._telemetry_writer: Optional[TelemetryWriter] = None
+        self._telemetry_worker_subscribed = False
+        # A MicroVM worker change that arrived while a send held the exporter lock.
+        self._pending_telemetry_worker: Optional[tuple] = None
+        # Serializes publishing a pending worker with a send consuming one, so the newer is not lost.
+        self._pending_telemetry_worker_lock = forksafe.Lock()
 
         # Native exporter methods require exclusive access because PyO3 rejects
         # overlapping mutable borrows.
@@ -911,6 +919,7 @@ def __del__(self) -> None:
         try:
             if getattr(self, "_owner_pid", None) != os.getpid():
                 return
+            self._unsubscribe_telemetry_worker()
             exporter = getattr(self, "_exporter", None)
             if exporter is not None and not getattr(self, "_exporter_dropped", False):
                 if self._discarded_by_refresh:
@@ -1013,14 +1022,57 @@ def _create_exporter(self) -> native.TraceExporter:
         exporter = builder.build(get_native_runtime())
         if shared_worker is not None:
             exporter.set_telemetry_handle(shared_worker)
-            telemetry_writer._subscribe_worker_changes(self._on_telemetry_worker_changed)
+            late_callback = (
+                partial(self._on_telemetry_worker_changed, exporter=exporter) if telemetry_writer._is_microvm else None
+            )
+            if telemetry_writer._is_microvm:
+                # Publish before subscribing: worker-change notifications re-point self._exporter,
+                # so a refresh before the caller assigns this result would otherwise miss it.
+                self._exporter = exporter
+                telemetry_writer._subscribe_worker_changes(
+                    self._on_telemetry_worker_changed, shared_worker, late_callback
+                )
+            else:
+                telemetry_writer._subscribe_worker_changes(self._on_telemetry_worker_changed, shared_worker)
+            self._telemetry_writer = telemetry_writer
+            self._telemetry_worker_subscribed = True
         return exporter
 
-    def _on_telemetry_worker_changed(self, worker: "Optional[native.TelemetryWorker]") -> None:
+    def _unsubscribe_telemetry_worker(self) -> None:
+        if not self._telemetry_worker_subscribed:
+            return
+        telemetry_writer = self._telemetry_writer
+        self._telemetry_writer = None
+        self._telemetry_worker_subscribed = False
+        if telemetry_writer is not None:
+            telemetry_writer._unsubscribe_worker_changes(self._on_telemetry_worker_changed)
+
+    def _on_telemetry_worker_changed(
+        self,
+        worker: "Optional[native.TelemetryWorker]",
+        exporter: "Optional[native.TraceExporter]" = None,
+    ) -> None:
         """Follow the telemetry writer onto a rebuilt worker (or off a stopped one)."""
+        # A send holds the exporter lock across its retries, and MicroVM identity refresh rebuilds
+        # the telemetry worker on the /run thread: leave the worker for the next send to apply.
+        if not self._exporter_lock.acquire(blocking=self._writer_lock is None):
+            with self._pending_telemetry_worker_lock:
+                self._pending_telemetry_worker = (worker, exporter)
+            return
         try:
-            with self._exporter_lock:
-                self._exporter.set_telemetry_handle(worker)
+            self._pending_telemetry_worker = None
+            self._set_telemetry_handle(worker, exporter)
+        finally:
+            self._exporter_lock.release()
+
+    def _set_telemetry_handle(
+        self,
+        worker: "Optional[native.TelemetryWorker]",
+        exporter: "Optional[native.TraceExporter]" = None,
+    ) -> None:
+        # Caller holds the exporter lock.
+        try:
+            (exporter if exporter is not None else self._exporter).set_telemetry_handle(worker)
         except Exception:
             log.debug("Failed to re-point the trace exporter at the telemetry worker", exc_info=True)
 
@@ -1053,9 +1105,11 @@ def set_test_session_token(self, token: Optional[str]) -> None:
 
     def shutdown_exporter(self) -> None:
         """Tear down the native exporter without going through ``stop()``."""
+        self._unsubscribe_telemetry_worker()
         self._shutdown_exporter(self._exporter)
 
     def _drop_exporter(self) -> None:
+        self._unsubscribe_telemetry_worker()
         with self._exporter_lock:
             # The refresh, the finishing sender, and on_shutdown() may all get here; only one drops.
             if not self._exporter_dropped:
@@ -1116,13 +1170,15 @@ def _downgrade(self, status, client):
             self._api_version = "v0.4"
             # Built outside the exporter lock: _create_exporter() can take the telemetry enable lock,
             # whose holder notifies _on_telemetry_worker_changed(), which takes the exporter lock.
+            # Read the old exporter first: in MicroVMs _create_exporter() publishes its result.
+            old_exporter = self._exporter
             new_exporter = self._create_exporter()
             with self._exporter_lock:
                 # A refresh may have dropped the exporter since the check above; a replacement
                 # installed now would be skipped by on_shutdown() and leak.
                 discarded = self._discarded_by_refresh
                 if not discarded:
-                    old_exporter, self._exporter = self._exporter, new_exporter
+                    self._exporter = new_exporter
             if discarded:
                 new_exporter.drop()
                 return
@@ -1191,6 +1247,10 @@ def _send_payload(self, payload: bytes, count: int, client: WriterClientBase):
                 # starts after a refresh and none reaches a dropped exporter.
                 if self._discarded_by_refresh:
                     return
+                if self._pending_telemetry_worker is not None:
+                    with self._pending_telemetry_worker_lock:
+                        pending, self._pending_telemetry_worker = self._pending_telemetry_worker, None
+                    self._set_telemetry_handle(*pending)
                 response_body = self._exporter.send(payload)
         except native.RequestError as e:
             try:
@@ -1396,6 +1456,7 @@ def _stop_service(
 
     def on_shutdown(self):
         if self._exporter_dropped:
+            self._unsubscribe_telemetry_worker()
             return
         if self._discarded_by_refresh:
             # Finish the refresh's discard without the final flush. drop() avoids shutdown()'s stats flush but
@@ -1406,6 +1467,7 @@ def on_shutdown(self):
         try:
             self.periodic()
         finally:
+            self._unsubscribe_telemetry_worker()
             self._shutdown_exporter(self._exporter)
 
 

            old_exporter = self._exporter
            new_exporter = self._create_exporter()
            with self._exporter_lock:
                # A refresh may have dropped the exporter since the check above; a replacement
                # installed now would be skipped by on_shutdown() and leak.
                discarded = self._discarded_by_refresh
                if not discarded:
                    self._exporter = new_exporter

            with self._exporter_lock:
                old_exporter = self._exporter
            new_exporter = self._create_exporter()
            with self._exporter_lock:
                # A refresh may have dropped the exporter since the check above; a replacement
                # installed now would be skipped by on_shutdown() and leak.
                discarded = self._discarded_by_refresh
                if not discarded:
                    self._exporter = new_exporter

releasenotes/notes/fix-telemetry-worker-identity-refresh-19fe1775be48b696.yaml

@@ -0,0 +1,5 @@
+---
+fixes:
+  - |
+    telemetry: Fixes an issue where telemetry events can be associated with a previous
+    AWS Lambda MicroVM invocation after the runtime identity is refreshed.

tests/appsec/sca/test_telemetry.py

@@ -282,6 +282,7 @@ def _make_writer_and_tracker(sca_enabled=False, deps=None, enabled=True):
     appsec_telemetry_config.SCA_ENABLED = sca_enabled
     writer = TelemetryWriter.__new__(TelemetryWriter)
     writer._metric_lock = MagicMock()
+    writer._worker_access_lock = MagicMock()
     writer._enabled = enabled
     # The native worker is mocked: _report_dependencies() forwards to worker.add_dependency and
     # returns the reported records, which is what these tests assert on.

tests/telemetry/test_dependency.py

@@ -452,9 +452,57 @@ def test_report_dependencies_no_rereport_without_new_metadata(self):
 
             result = tracker.collect_report()
 
-        # No new deps, no new metadata -> None
         assert result is None
 
+    def test_refresh_preserves_and_rereports_dependency_metadata(self):
+        from unittest.mock import patch
+
+        _, tracker = _make_writer_and_tracker(sca_enabled=True)
+        entry = DependencyEntry(name="requests", version="2.28.0", metadata=[])
+        entry.add_metadata("CVE-1", "requests.sessions", "send", 10)
+        entry.mark_initial_sent()
+        entry.mark_all_metadata_sent()
+        tracker._imported_dependencies["requests"] = entry
+
+        tracker.refresh()
+
+        with (
+            patch("ddtrace.internal.telemetry.dependency_tracker.modules") as mock_modules,
+            patch("ddtrace.internal.telemetry.dependency_tracker.telemetry_config") as mock_config,
+        ):
+            mock_config.DEPENDENCY_COLLECTION = True
+            mock_modules.get_newly_imported_modules.return_value = set()
+
+            result = tracker.collect_report()
+
+        assert result is not None
+        assert result[0]["name"] == "requests"
+        assert len(result[0]["metadata"]) == 1
+        assert json.loads(result[0]["metadata"][0]["value"])["id"] == "CVE-1"
+
+    def test_refresh_rereports_dependencies_when_sca_disabled(self):
+        from unittest.mock import patch
+
+        _, tracker = _make_writer_and_tracker(sca_enabled=False)
+        entry = DependencyEntry(name="requests", version="2.28.0")
+        entry.mark_initial_sent()
+        tracker._imported_dependencies["requests"] = entry
+
+        tracker.refresh()
+
+        with (
+            patch("ddtrace.internal.telemetry.dependency_tracker.modules") as mock_modules,
+            patch("ddtrace.internal.telemetry.dependency_tracker.telemetry_config") as mock_config,
+        ):
+            mock_config.DEPENDENCY_COLLECTION = True
+            mock_modules.get_newly_imported_modules.return_value = set()
+
+            result = tracker.collect_report()
+            second_result = tracker.collect_report()
+
+        assert result == [{"name": "requests", "version": "2.28.0"}]
+        assert second_result is None
+
     def test_rereport_includes_all_metadata_per_rfc(self):
         """Re-report includes ALL metadata (sent + unsent) per RFC."""
         from unittest.mock import patch