feat: add multiprocessing support

add multiprocessing support by ensuring ControllerClient is spawned for each pid

Author: Kl0ven

sentry_dynamic_sampling_lib/__init__.py

@@ -52,7 +52,7 @@ def init_wrapper():
         app_key = build_app_key(client.options)
         controller_endpoint = urljoin(controller_host, controller_path)
         metric_endpoint = urljoin(controller_host, metric_path)
-        print("Sentry Wrapper: Injecting TracesSampler")
+        print(f"Sentry Wrapper: Injecting TracesSampler. App Key : {app_key}")
         client.options["traces_sampler"] = TraceSampler(
             poll_interval=poll_interval,
             metric_interval=metric_interval,

sentry_dynamic_sampling_lib/sampler.py

@@ -1,7 +1,9 @@
 import logging
+import os
 import signal
 from threading import Event, Thread
 from time import sleep
+from typing import Optional
 
 import schedule
 from requests.exceptions import RequestException
@@ -25,16 +27,28 @@ def on_exit(*args, **kwargs):
 
 
 class ControllerClient(Thread):
-    def __init__(self, stop, config, metric, *args, **kwargs) -> None:
-        self.poll_interval = kwargs.pop("poll_interval")
-        self.metric_interval = kwargs.pop("metric_interval")
-        self.controller_endpoint = kwargs.pop("controller_endpoint")
-        self.metric_endpoint = kwargs.pop("metric_endpoint")
-        self.app_key = kwargs.pop("app_key")
-        self.stop: Event = stop
-        self.config: Config = config
-        self.metrics: Metric = metric
+    def __init__(
+        self,
+        *args,
+        poll_interval=None,
+        metric_interval=None,
+        controller_endpoint=None,
+        metric_endpoint=None,
+        app_key=None,
+        **kwargs
+    ) -> None:
+        self.poll_interval = poll_interval
+        self.metric_interval = metric_interval
+        self.controller_endpoint = controller_endpoint
+        self.metric_endpoint = metric_endpoint
+        self.app_key = app_key
+
+        self._stop = Event()
+        self.config = Config()
+        self.metrics = Metric()
         self.session = CachedSession(backend="memory", cache_control=True)
+
+        LOGGER.debug("ControllerClient Initialized")
         super().__init__(*args, name="SentryControllerClient", **kwargs)
 
     def run(self):
@@ -44,10 +58,16 @@ def run(self):
         sleep(5)
         schedule.every(self.poll_interval).seconds.do(self.update_config)
         schedule.every(self.metric_interval).seconds.do(self.update_metrics)
-        while not self.stop.is_set():
+        LOGGER.debug("ControllerClient Started")
+        while not self._stop.is_set():
             schedule.run_pending()
             sleep(1)
 
+    def kill(self):
+        self._stop.set()
+        if self.is_alive():
+            self.join()
+
     def update_config(self):
         try:
             resp = self.session.get(
@@ -59,8 +79,10 @@ def update_config(self):
             return
 
         if resp.from_cache:
+            LOGGER.debug("Config Polled from cache")
             return
 
+        LOGGER.debug("Config Polled")
         data = resp.json()
         self.config.update(data)
         self.metrics.set_mode(MetricType.CELERY, data["celery_collect_metrics"])
@@ -71,6 +93,7 @@ def update_metrics(self):
             # check if metric is enable
             mode = self.metrics.get_mode(metric_type)
             if not mode:
+                LOGGER.debug("Metric %s disabled", metric_type.value)
                 continue
 
             counter = self.metrics.get_and_reset(metric_type)
@@ -86,36 +109,59 @@ def update_metrics(self):
                     self.metric_endpoint.format(self.app_key, metric_type.value),
                     json=data,
                 )
+                LOGGER.debug("Metric %s pushed", metric_type.value)
             except RequestException as err:
                 LOGGER.warning("Metric Request Failed: %s", err)
                 return
 
 
 class TraceSampler(metaclass=Singleton):
     def __init__(self, *args, **kwargs) -> None:
-        self.stop = Event()
-        self.config = Config()
-        self.metrics = Metric()
-        self.controller = ControllerClient(
-            *args, self.stop, self.config, self.metrics, **kwargs
-        )
-        self.controller.start()
+        self.params = (args, kwargs)
+        self._controller: Optional[ControllerClient] = None
+        self._tread_for_pid: Optional[int] = None
 
         signal.signal(signal.SIGINT, on_exit)
-
         # HACK: Celery has a built in signal mechanism
         # so we use it
         if worker_shutdown:
             worker_shutdown.connect(on_exit)
 
-    def __del__(self):
-        on_exit(self.stop, self.controller)
+    @property
+    def has_running_controller(self):
+        if self._tread_for_pid != os.getpid():
+            return False
+        if not self._controller:
+            return None
+        return self._controller.is_alive()
+
+    @property
+    def config(self) -> Config:
+        return self._controller.config
+
+    @property
+    def metrics(self) -> Metric:
+        return self._controller.metrics
 
     def kill(self):
-        self.stop.set()
-        self.controller.join()
+        if self._controller:
+            self._controller.kill()
+
+    def _start_controller(self):
+        args, kwargs = self.params
+        self._controller = ControllerClient(*args, **kwargs)
+        self._controller.start()
+        self._tread_for_pid = os.getpid()
+
+    def _ensure_controller(self):
+        if not self.has_running_controller:
+            self._start_controller()
+
+    def __del__(self):
+        self.kill()
 
     def __call__(self, sampling_context):
+        self._ensure_controller()
         if sampling_context:
             if "wsgi_environ" in sampling_context:
                 path = sampling_context["wsgi_environ"].get("PATH_INFO", "")