fix(rcm): update the forking behavior in remote config [backport 2.2]

ddtrace/appsec/_remoteconfiguration.py

@@ -13,6 +13,7 @@
 from ddtrace.appsec._constants import PRODUCTS
 from ddtrace.appsec._utils import _appsec_rc_features_is_enabled
 from ddtrace.constants import APPSEC_ENV
+from ddtrace.internal import forksafe
 from ddtrace.internal.logger import get_logger
 from ddtrace.internal.remoteconfig._connectors import PublisherSubscriberConnector
 from ddtrace.internal.remoteconfig._publishers import RemoteConfigPublisherMergeDicts
@@ -25,6 +26,8 @@
 
 log = get_logger(__name__)
 
+APPSEC_PRODUCTS = [PRODUCTS.ASM_FEATURES, PRODUCTS.ASM, PRODUCTS.ASM_DATA, PRODUCTS.ASM_DD]
+
 
 class AppSecRC(PubSub):
     __subscriber_class__ = RemoteConfigSubscriber
@@ -36,6 +39,10 @@ def __init__(self, _preprocess_results, callback):
         self._subscriber = self.__subscriber_class__(self.__shared_data__, callback, "ASM")
 
 
+def _forksafe_appsec_rc():
+    remoteconfig_poller.start_subscribers_by_product(APPSEC_PRODUCTS)
+
+
 def enable_appsec_rc(test_tracer: Optional[Tracer] = None) -> None:
     """Remote config will be used by ASM libraries to receive four different updates from the backend.
     Each update has it’s own product:
@@ -70,14 +77,13 @@ def enable_appsec_rc(test_tracer: Optional[Tracer] = None) -> None:
         remoteconfig_poller.register(PRODUCTS.ASM, asm_callback)  # Exclusion Filters & Custom Rules
         remoteconfig_poller.register(PRODUCTS.ASM_DD, asm_callback)  # DD Rules
 
+    forksafe.register(_forksafe_appsec_rc)
+
 
 def disable_appsec_rc():
     # only used to avoid data leaks between tests
-
-    remoteconfig_poller.unregister(PRODUCTS.ASM_FEATURES)
-    remoteconfig_poller.unregister(PRODUCTS.ASM_DATA)
-    remoteconfig_poller.unregister(PRODUCTS.ASM)
-    remoteconfig_poller.unregister(PRODUCTS.ASM_DD)
+    for product_name in APPSEC_PRODUCTS:
+        remoteconfig_poller.unregister(product_name)
 
 
 def _add_rules_to_list(features: Mapping[str, Any], feature: str, message: str, ruleset: Dict[str, Any]) -> None:

ddtrace/internal/remoteconfig/client.py

@@ -244,6 +244,13 @@ def update_product_callback(self, product_name, callback):
             return True
         return False
 
+    def start_products(self, products_list):
+        # type: (list) -> None
+        for product_name in products_list:
+            pubsub_instance = self._products.get(product_name)
+            if pubsub_instance:
+                pubsub_instance.restart_subscriber()
+
     def unregister_product(self, product_name):
         # type: (str) -> None
         self._products.pop(product_name, None)

ddtrace/internal/remoteconfig/worker.py

@@ -1,4 +1,5 @@
 import os
+from typing import List
 
 from ddtrace.internal import agent
 from ddtrace.internal import atexit
@@ -80,19 +81,21 @@ def enable(self):
                 return True
 
             self.start()
-            forksafe.register(self.start_subscribers)
+            forksafe.register(self.reset_at_fork)
             atexit.register(self.disable)
             return True
         return False
 
-    def start_subscribers(self):
+    def reset_at_fork(self):
         # type: () -> None
-        """Subscribers need to be restarted when application forks"""
+        """Client Id needs to be refreshed when application forks"""
         self._enable = False
-        log.debug("[%s][P: %s] Remote Config Poller fork. Starting Pubsub services", os.getpid(), os.getppid())
+        log.debug("[%s][P: %s] Remote Config Poller fork. Refreshing state", os.getpid(), os.getppid())
         self._client.renew_id()
-        for pubsub in self._client.get_pubsubs():
-            pubsub.restart_subscriber()
+
+    def start_subscribers_by_product(self, products_list):
+        # type: (List[str]) -> None
+        self._client.start_products(products_list)
 
     def _poll_data(self, test_tracer=None):
         """Force subscribers to poll new data. This function is only used in tests"""
@@ -121,7 +124,7 @@ def disable(self, join=False):
         if self.status == ServiceStatus.STOPPED:
             return
 
-        forksafe.unregister(self.start_subscribers)
+        forksafe.unregister(self.reset_at_fork)
         atexit.unregister(self.disable)
 
         self.stop(join=join)
@@ -147,14 +150,12 @@ def register(self, product, pubsub_instance, skip_enabled=False):
         try:
             # By enabling on registration we ensure we start the RCM client only
             # if there is at least one registered product.
-            enabled = True
             if not skip_enabled:
-                enabled = self.enable()
+                self.enable()
 
-            if enabled:
-                self._client.register_product(product, pubsub_instance)
-                if not self._client.is_subscriber_running(pubsub_instance):
-                    pubsub_instance.start_subscriber()
+            self._client.register_product(product, pubsub_instance)
+            if not self._client.is_subscriber_running(pubsub_instance):
+                pubsub_instance.start_subscriber()
         except Exception:
             log.debug("error starting the RCM client", exc_info=True)
 

releasenotes/notes/fix-debugger-rc-instance-954539a2e57b7887.yaml

@@ -0,0 +1,6 @@
+---
+fixes:
+  - |
+    dynamic instrumentation: Needs to update the pubsub instance when the application forks because the probe 
+    mechanism should run in the child process. For that, DI needs the callback as the method of an instance of 
+    Debugger, which lives in the child process.

tests/internal/remoteconfig/test_remoteconfig.py

@@ -146,7 +146,7 @@ def get_mock_encoded_msg_with_signed_errors(msg, path, signed_errors):
 
 def test_remote_config_register_auto_enable(remote_config_worker):
     # ASM_FEATURES product is enabled by default, but LIVE_DEBUGGER isn't
-    class MockPubsub:
+    class MockPubsub(PubSub):
         def stop(self, *args, **kwargs):
             pass
 
@@ -165,10 +165,15 @@ def stop(self, *args, **kwargs):
 
 def test_remote_config_register_validate_rc_disabled(remote_config_worker):
     remoteconfig_poller.disable()
+
+    class MockPubsub(PubSub):
+        def stop(self, *args, **kwargs):
+            pass
+
     assert remoteconfig_poller.status == ServiceStatus.STOPPED
 
     with override_global_config(dict(_remote_config_enabled=False)):
-        remoteconfig_poller.register("LIVE_DEBUGGER", lambda m, c: None)
+        remoteconfig_poller.register("LIVE_DEBUGGER", MockPubsub())
 
         assert remoteconfig_poller.status == ServiceStatus.STOPPED