diff --git a/ddtrace/internal/symbol_db/symbols.py b/ddtrace/internal/symbol_db/symbols.py index b8b9daa35e6..5073ae6cd7c 100644 --- a/ddtrace/internal/symbol_db/symbols.py +++ b/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: diff --git a/tests/internal/symbol_db/test_symbols.py b/tests/internal/symbol_db/test_symbols.py index c1c5b075312..8b64d7ecf7e 100644 --- a/tests/internal/symbol_db/test_symbols.py +++ b/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(): """