chore(runtime): add explicit identity refresh

ddtrace/internal/_runtime_id.py

@@ -1,3 +1,4 @@
+import logging
 import typing as t
 import uuid
 
@@ -6,12 +7,16 @@
 from . import forksafe
 
 
+log = logging.getLogger(__name__)
+
+
 __all__ = [
     "get_ancestor_runtime_id",
     "get_process_role",
     "get_runtime_id",
     "get_parent_runtime_id",
     "get_runtime_propagation_envs",
+    "refresh_identity",
 ]
 
 
@@ -31,6 +36,7 @@ def _generate_runtime_id() -> str:
 # IMPORTANT: Do not change t.Set to set until minimum Python version is 3.11+
 # Module-level set[...] in Python 3.10 affects import timing. See packages.py for details.
 _ON_RUNTIME_ID_CHANGE: t.Set[t.Callable[[str], None]] = set()  # noqa: UP006
+_ON_RUNTIME_IDENTITY_REFRESH: t.Set[t.Callable[[str], None]] = set()  # noqa: UP006
 
 
 def on_runtime_id_change(cb: t.Callable[[str], None]) -> None:
@@ -42,6 +48,45 @@ def on_runtime_id_change(cb: t.Callable[[str], None]) -> None:
     _ON_RUNTIME_ID_CHANGE.add(cb)
 
 
+def on_runtime_identity_refresh(cb: t.Callable[[str], None]) -> None:
+    """Register a callback to be called after an explicit runtime identity refresh.
+
+    This is separate from the fork callback because a logical runtime replacement
+    must not be treated as a child process or update the fork lineage.
+    """
+    global _ON_RUNTIME_IDENTITY_REFRESH
+    _ON_RUNTIME_IDENTITY_REFRESH.add(cb)
+
+
+def _notify_runtime_id_callbacks(callbacks: t.Set[t.Callable[[str], None]]) -> None:  # noqa: UP006
+    for cb in list(callbacks):
+        try:
+            cb(_RUNTIME_ID)
+        except Exception:
+            log.exception("Exception ignored in runtime ID callback %r", cb)
+
+
+def _notify_runtime_identity_refresh_callbacks(*, raise_on_error: bool = False) -> None:  # noqa: UP006
+    # Direct refresh callers keep subscriber failures isolated so every component gets
+    # a chance to rebuild. The MicroVM coordinator opts into propagation so a failed
+    # rebuild leaves its completion guard unset and the same identity can be retried.
+    for cb in list(_ON_RUNTIME_IDENTITY_REFRESH):
+        if raise_on_error:
+            cb(_RUNTIME_ID)
+            continue
+        try:
+            cb(_RUNTIME_ID)
+        except Exception:
+            log.exception("Exception ignored in runtime ID callback %r", cb)
+
+
+def _refresh_runtime_id() -> None:
+    global _RUNTIME_ID
+
+    _RUNTIME_ID = _generate_runtime_id()
+    _notify_runtime_id_callbacks(_ON_RUNTIME_ID_CHANGE)
+
+
 @forksafe.register
 def _set_runtime_id() -> None:
     global _RUNTIME_ID, _ANCESTOR_RUNTIME_ID, _PARENT_RUNTIME_ID
@@ -51,9 +96,25 @@ def _set_runtime_id() -> None:
         _ANCESTOR_RUNTIME_ID = _RUNTIME_ID
 
     _PARENT_RUNTIME_ID = _RUNTIME_ID
-    _RUNTIME_ID = _generate_runtime_id()
-    for cb in _ON_RUNTIME_ID_CHANGE:
-        cb(_RUNTIME_ID)
+    _refresh_runtime_id()
+
+
+def refresh_identity(raise_on_error: bool = False) -> None:
+    """Regenerate the runtime ID without recording fork lineage.
+
+    Unlike a fork, this does not update _PARENT_RUNTIME_ID / _ANCESTOR_RUNTIME_ID:
+    the previous runtime ID was not a real parent process, so recording it there
+    would make get_process_role() and friends misreport a fork lineage that never
+    existed. Use this when a new logical process instance is created by a mechanism
+    other than fork().
+    """
+    # This is an opt-in lifecycle path: non-MicroVM processes do not call it from
+    # ordinary request handling, so their runtime-ID and fork behavior is unchanged.
+    # Notify consumers that only need the new ID first. The explicit refresh
+    # callbacks below are for components that must rebuild restore-sensitive
+    # state, which is different from the fork handling in _set_runtime_id().
+    _refresh_runtime_id()
+    _notify_runtime_identity_refresh_callbacks(raise_on_error=raise_on_error)
 
 
 def get_runtime_id() -> str:

ddtrace/internal/runtime/__init__.py

@@ -4,6 +4,8 @@
 from ddtrace.internal._runtime_id import get_runtime_id
 from ddtrace.internal._runtime_id import get_runtime_propagation_envs
 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
 
 
 __all__ = [
@@ -13,4 +15,6 @@
     "get_parent_runtime_id",
     "get_runtime_propagation_envs",
     "on_runtime_id_change",
+    "on_runtime_identity_refresh",
+    "refresh_identity",
 ]