chore(symbol-db): refresh metadata after identity refresh

ddtrace/internal/symbol_db/symbols.py

@@ -38,6 +38,7 @@
 from ddtrace.internal.periodic import Timer
 from ddtrace.internal.runtime import get_ancestor_runtime_id
 from ddtrace.internal.runtime import get_runtime_id
+from ddtrace.internal.runtime import on_runtime_id_change
 from ddtrace.internal.safety import _isinstance
 from ddtrace.internal.settings._agent import config as agent_config
 from ddtrace.internal.settings.dynamic_instrumentation import config as di_config
@@ -560,23 +561,46 @@ def __init__(self, scopes: t.Optional[list[Scope]] = None) -> None:
             "final": False,
         }
 
+        self._runtime_id_change_fork_generation = forksafe.get_generation()
         forksafe.register(self._reset_on_fork)
+        # 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)
 
     @cached_property
     def _sender(self) -> SymDBSender:
         # Built on the first upload to reduce startup cost.
         return build_symdb_sender()
 
-    def _reset_on_fork(self) -> None:
-        # Runs after fork in the child. The runtime module's own forksafe
-        # hook regenerates the runtime_id first (registered earlier at
-        # import time), so the calls below return the post-fork values.
+    def _reset_upload_metadata(self) -> None:
         self._upload_id = str(uuid.uuid4())
         self._batch_counter = 0
         self._event_data["uploadId"] = self._upload_id
         self._event_data["runtimeId"] = get_runtime_id()
         self._event_data["parentId"] = get_ancestor_runtime_id()
 
+    def _reset_on_fork(self) -> None:
+        # Runs after fork in the child. The runtime module's own forksafe
+        # hook regenerates the runtime_id first (registered earlier at
+        # import time), so the calls below return the post-fork values.
+        self._runtime_id_change_fork_generation = forksafe.get_generation()
+        self._reset_upload_metadata()
+
+    def _on_identity_refresh(self, new_runtime_id: str) -> None:
+        fork_generation = forksafe.get_generation()
+        if fork_generation != self._runtime_id_change_fork_generation:
+            # Runtime ID callbacks also run from the forksafe hook, before this
+            # context's own fork hook. Skip that notification so fork setup keeps
+            # one reset path; explicit later identity refreshes in the child still
+            # reset below.
+            self._runtime_id_change_fork_generation = fork_generation
+            return
+
+        # Explicit identity refresh runs in a live process and can race with the
+        # timer upload thread, which reads the same upload metadata.
+        with self._scopes_lock:
+            self._reset_upload_metadata()
+
     def _set_timer(self) -> None:
         with self._timer_lock:
             if self._timer is None:

tests/internal/symbol_db/test_symbols.py

@@ -526,6 +526,63 @@ def test_symbols_fork_uploads():
         assert os.WIFEXITED(status) and os.WEXITSTATUS(status) == 0, f"child {pid} exited with status {status}"
 
 
+@pytest.mark.subprocess(ddtrace_run=True, err=None)
+def test_symbols_identity_refresh_updates_runtime_id():
+    """A non-fork identity refresh (e.g. an AWS Lambda MicroVM /run hook) must also refresh the
+    ScopeContext's cached runtimeId, the same way _reset_on_fork() does after an actual fork --
+    otherwise every upload after a MicroVM /run keeps reporting the pre-refresh snapshot's ID.
+    """
+    import typing as t
+
+    import ddtrace.internal.runtime as runtime
+    from ddtrace.internal.symbol_db.symbols import SymbolDatabaseUploader
+
+    SymbolDatabaseUploader.install()
+
+    context = t.cast(SymbolDatabaseUploader, SymbolDatabaseUploader._instance)._context
+    old_runtime_id = runtime.get_runtime_id()
+    old_upload_id = context._upload_id
+    assert context._event_data["runtimeId"] == old_runtime_id
+
+    runtime.refresh_identity()
+
+    assert runtime.get_runtime_id() != old_runtime_id
+    assert context._event_data["runtimeId"] == runtime.get_runtime_id()
+    assert context._upload_id != old_upload_id
+
+
+def test_symbols_identity_refresh_skips_first_new_fork_generation():
+    from unittest import mock
+
+    context = ScopeContext()
+    fork_generation = context._runtime_id_change_fork_generation
+    old_upload_id = context._upload_id
+
+    with mock.patch("ddtrace.internal.symbol_db.symbols.forksafe.get_generation", return_value=fork_generation + 1):
+        context._on_identity_refresh("new-runtime-id")
+
+    assert context._upload_id == old_upload_id
+
+
+def test_symbols_identity_refresh_resets_metadata_under_upload_lock():
+    class RecordingLock:
+        entered = False
+
+        def __enter__(self):
+            self.entered = True
+
+        def __exit__(self, *exc_info):
+            return None
+
+    context = ScopeContext()
+    lock = RecordingLock()
+    context._scopes_lock = lock  # type: ignore[assignment]
+
+    context._on_identity_refresh("new-runtime-id")
+
+    assert lock.entered
+
+
 @pytest.mark.subprocess(ddtrace_run=True, err=None)
 def test_symbols_fork_forces_reenable_and_install():
     """