refactor: metrics

refactor metric to simplify use

Author: Kl0ven

sentry_dynamic_sampling_lib/sampler.py

@@ -23,6 +23,7 @@
 def on_exit(*args, **kwargs):
     ts = TraceSampler()
     ts.kill()
+    LOGGER.debug("ControllerClient Killed")
     raise KeyboardInterrupt
 
 
@@ -89,20 +90,11 @@ def update_config(self):
         self.metrics.set_mode(MetricType.WSGI, data["wsgi_collect_metrics"])
 
     def update_metrics(self):
-        for metric_type in MetricType:
-            # 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)
-            if len(counter) == 0:
-                return
+        for metric_type, data in self.metrics:
             data = {
                 "app": self.app_key,
                 "type": metric_type.value,
-                "data": dict(counter.most_common(10)),
+                "data": dict(data.most_common(10)),
             }
             try:
                 self.session.post(
@@ -112,7 +104,7 @@ def update_metrics(self):
                 LOGGER.debug("Metric %s pushed", metric_type.value)
             except RequestException as err:
                 LOGGER.warning("Metric Request Failed: %s", err)
-                return
+                continue
 
 
 class TraceSampler(metaclass=Singleton):

sentry_dynamic_sampling_lib/shared.py

@@ -1,3 +1,4 @@
+import logging
 from collections import Counter
 from enum import Enum
 from threading import RLock
@@ -9,6 +10,8 @@
 )
 from sentry_dynamic_sampling_lib.utils import synchronized
 
+LOGGER = logging.getLogger("SentryWrapper")
+
 
 class Config:
     def __init__(self) -> None:
@@ -62,33 +65,47 @@ class MetricType(Enum):
 class Metric:
     def __init__(self) -> None:
         self._lock = RLock()
-        self._activate = {
-            MetricType.WSGI: False,
-            MetricType.CELERY: False,
-        }
-        self._counters = {
-            MetricType.WSGI: Counter(),
-            MetricType.CELERY: Counter(),
+        self._metrics = {
+            MetricType.WSGI: {"activated": False, "data": Counter()},
+            MetricType.CELERY: {"activated": False, "data": Counter()},
         }
 
     def set_mode(self, _type, mode):
-        self._activate[_type] = mode
+        self._metrics[_type]["activated"] = mode
 
     def get_mode(self, _type):
-        return self._activate[_type]
+        return self._metrics[_type]["activated"]
 
     @synchronized
     def count_path(self, path):
-        if self._activate[MetricType.WSGI]:
-            self._counters[MetricType.WSGI][path] += 1
+        metric = self._metrics[MetricType.WSGI]
+        if metric["activated"]:
+            metric["data"][path] += 1
 
     @synchronized
     def count_task(self, path):
-        if self._activate[MetricType.CELERY]:
-            self._counters[MetricType.CELERY][path] += 1
-
-    @synchronized
-    def get_and_reset(self, _type):
-        counter = self._counters[_type]
-        self._counters[_type] = Counter()
-        return counter
+        metric = self._metrics[MetricType.CELERY]
+        if metric["activated"]:
+            metric["data"][path] += 1
+
+    def __iter__(self):
+        """
+        List activated non-empty metrics
+
+        Yields:
+            Tuple(MetricType, Counter): the activated non-empty metrics
+        """
+        for metric_type, metric in self._metrics.items():
+            # check if metric is activated
+            if not metric["activated"]:
+                LOGGER.debug("Metric %s disabled", metric_type.value)
+            with self._lock:
+                # check im metric is empty
+                if len(metric["data"]) == 0:
+                    LOGGER.debug("Metric %s is empty", metric_type.value)
+                    continue
+                data = metric["data"]
+                metric["data"] = Counter()
+
+            # yield outside of the lock to not block write while callee execute
+            yield metric_type, data