chore(remoteconfig): refresh client identity on runtime refresh

ddtrace/internal/remoteconfig/client.py

@@ -5,18 +5,21 @@
 from typing import Optional
 from typing import Sequence
 import uuid
+import weakref
 
 import ddtrace
+from ddtrace.internal import forksafe
 from ddtrace.internal import gitmetadata
 from ddtrace.internal import process_tags
-from ddtrace.internal import runtime
 from ddtrace.internal.hostname import get_hostname
 from ddtrace.internal.logger import get_logger
 from ddtrace.internal.packages import is_distribution_available
 from ddtrace.internal.remoteconfig import ConfigMetadata
 from ddtrace.internal.remoteconfig import Payload
 from ddtrace.internal.remoteconfig import PayloadType
 from ddtrace.internal.remoteconfig import RCCallback
+from ddtrace.internal.runtime import get_runtime_id
+from ddtrace.internal.runtime import on_runtime_id_change
 from ddtrace.internal.settings._agent import config as agent_config
 from ddtrace.internal.settings._core import DDConfig
 from ddtrace.internal.telemetry import telemetry_writer
@@ -98,6 +101,17 @@ def __init__(self) -> None:
         self._native: Optional[Any] = None
         self._reader: Optional[Any] = None
 
+        client_ref = weakref.ref(self)
+
+        # Runtime keeps callbacks in a module-level set, so registering a bound
+        # method would keep this client alive after callers drop their reference.
+        def _on_identity_refresh(new_runtime_id: str) -> None:
+            client = client_ref()
+            if client is not None:
+                client._on_identity_refresh(new_runtime_id)
+
+        on_runtime_id_change(_on_identity_refresh)
+
     def ensure_native(self) -> Any:
         if self._native is None:
             from ddtrace.internal.native import RemoteConfigClient as _NativeClient
@@ -109,7 +123,7 @@ def ensure_native(self) -> Any:
                 agent_url=str(self.agent_url),
                 tracer_version=tracer_version,
                 client_id=self.id,
-                runtime_id=runtime.get_runtime_id(),
+                runtime_id=get_runtime_id(),
                 service=ddtrace.config.service or "",
                 env=ddtrace.config.env or "",
                 app_version=ddtrace.config.version or "",
@@ -125,6 +139,25 @@ def ensure_native(self) -> Any:
     def renew_id(self) -> None:
         self.id = str(uuid.uuid4())
 
+    def _on_identity_refresh(self, new_runtime_id: str) -> None:
+        # Regenerate the client id and drop the native client, which bakes both ids in
+        # as immutable constructor arguments (get_client_id() is documented "stable for
+        # the process lifetime"). The next ensure_native() call rebuilds it bound to the
+        # fresh ids. Safe across threads: request() captures self._native into a local
+        # before calling .poll(), so an in-flight poll on the old client is unaffected.
+        # After fork, keep a reader over inherited SHM before dropping the inherited
+        # native client; otherwise forked consumers lose the only handles they can read.
+        if forksafe.is_fork_child():
+            if self._native is not None and self._reader is None:
+                try:
+                    self._reader = self._native.make_reader()
+                except Exception:
+                    log.debug("failed to create remote config reader after fork", exc_info=True)
+        else:
+            self._reader = None
+        self.renew_id()
+        self._native = None
+
     def register_callback(self, product_name: "RemoteConfigProduct", callback: RCCallback) -> None:
         self._product_callbacks[product_name] = callback
         log.debug("[%s][P: %s] Registered callback for product %s", os.getpid(), os.getppid(), product_name)

tests/internal/remoteconfig/test_remoteconfig_native.py

@@ -418,3 +418,68 @@ def test_enable_builds_native_runtime_before_registering_fork_hook(monkeypatch):
         assert poller.enable() is True
 
     assert order == ["native", "before_fork", "start"], order
+
+
+def test_identity_refresh_renews_client_id_and_drops_native():
+    # get_client_id() on the native client is documented "stable for the process lifetime",
+    # so refreshing must drop it (not mutate it in place) for the id to actually change.
+    client = RemoteConfigClient()
+    old_id = client.id
+    client.ensure_native()
+    assert client._native is not None
+
+    client._on_identity_refresh("some-new-runtime-id")
+
+    assert client.id != old_id
+    assert client._native is None
+
+
+def test_identity_refresh_drops_cached_reader_outside_fork():
+    client = RemoteConfigClient()
+    client._reader = object()
+
+    client._on_identity_refresh("some-new-runtime-id")
+
+    assert client._reader is None
+
+
+def test_identity_refresh_callback_does_not_keep_client_alive():
+    import gc
+    import weakref
+
+    client = RemoteConfigClient()
+    client_ref = weakref.ref(client)
+
+    del client
+    gc.collect()
+
+    assert client_ref() is None
+
+
+def test_identity_refresh_rebuilds_native_client_with_fresh_id():
+    client = RemoteConfigClient()
+    native_before = client.ensure_native()
+    old_native_client_id = native_before.get_client_id()
+
+    client._on_identity_refresh("some-new-runtime-id")
+    native_after = client.ensure_native()
+
+    assert native_after is not native_before
+    assert native_after.get_client_id() == client.id
+    assert native_after.get_client_id() != old_native_client_id
+
+
+@pytest.mark.subprocess
+def test_identity_refresh_wired_to_runtime_id_change():
+    """A RemoteConfigClient subscribes itself at construction; refresh_identity() reaches it."""
+    from ddtrace.internal import runtime
+    from ddtrace.internal.remoteconfig.client import RemoteConfigClient
+
+    client = RemoteConfigClient()
+    old_id = client.id
+    client.ensure_native()
+
+    runtime.refresh_identity()
+
+    assert client.id != old_id
+    assert client._native is None