Counter implementation in metrics (#495)

* Counter implementation in metrics * Update vaas/vaas/cluster/cluster.py Co-authored-by: Jan Taczanowski <jan.taczanowski@allegro.pl> * Update vaas/vaas/cluster/cluster.py Co-authored-by: Jan Taczanowski <jan.taczanowski@allegro.pl> --------- Co-authored-by: Jan Taczanowski <jan.taczanowski@allegro.pl>
allegro · Sep 1, 2023 · 6c83c9b · 6c83c9b
1 parent 6551dab
commit 6c83c9b
Show file tree

Hide file tree

Showing 8 changed files with 56 additions and 19 deletions.
diff --git a/vaas/vaas/cluster/cluster.py b/vaas/vaas/cluster/cluster.py
@@ -29,11 +29,11 @@ def load_vcl_task(self, emmit_time, cluster_ids):
     ).prefetch_related('varnishserver_set')
     if len(clusters) > 0:
         varnish_cluster_load_vcl = VarnishCluster().load_vcl(start_processing_time, clusters)
-        metrics.gauge('events_with_change', 1)
+        metrics.counter('events_with_change')
         metrics.time('total_time_of_processing_vcl_task_with_change', time.perf_counter() - emmit_time_aware)
         return varnish_cluster_load_vcl
 
-    metrics.gauge('events_without_change', 1)
+    metrics.counter('events_without_change')
     metrics.time('total_time_of_processing_vcl_task_without_change', time.perf_counter() - emmit_time_aware)
 
     return True
@@ -115,9 +115,9 @@ def load_vcl(self, start_processing_time, clusters):
         else:
             result = parallel_loader.use_vcl_list(start_processing_time, loaded_vcl_list)
             if result is False:
-                metrics.gauge('successful_reload_vcl', 0)
+                metrics.counter('unsuccessful_reload_vcl')
             else:
-                metrics.gauge('successful_reload_vcl', 1)
+                metrics.counter('successful_reload_vcl')
             return result
         finally:
             for phase, processing in processing_stats.items():
@@ -139,7 +139,7 @@ def _update_vcl_versions(self, clusters, start_processing_time, vcl_list):
     @collect_processing
     def _handle_load_error(self, e, clusters, start_processing_time):
         self.logger.error('Loading error: {} - rendered vcl-s not used'.format(e))
-        metrics.gauge('successful_reload_vcl', 0)
+        metrics.counter('unsuccessful_reload_vcl')
 
         for cluster in clusters:
             cluster.error_timestamp = start_processing_time

diff --git a/vaas/vaas/metrics/handler.py b/vaas/vaas/metrics/handler.py
@@ -9,6 +9,10 @@ class MetricsHandler:
     def __init__(self, protocols: List[MetricsProtocol]):
         self.protocols: List[MetricsProtocol] = protocols
 
+    def counter(self, metric_name: str) -> None:
+        for protocol in self.protocols:
+            protocol.counter(metric_name)
+
     def time(self, metric_name: str, value: float) -> None:
         for protocol in self.protocols:
             protocol.sum(metric_name, value)

diff --git a/vaas/vaas/metrics/models.py b/vaas/vaas/metrics/models.py
@@ -7,6 +7,9 @@ def __init__(self, client: Any) -> None:
 
 
 class MetricsProtocol(Protocol):
+    def counter(self, metric_name: str) -> None:
+        ...
+
     def sum(self, metric_name: str, value: float) -> None:
         ...
 

diff --git a/vaas/vaas/metrics/prometheus.py b/vaas/vaas/metrics/prometheus.py
@@ -1,25 +1,27 @@
 import logging
 
-from typing import Dict, Type, Union
+from typing import Dict, Optional, Type, Union
 from django.conf import settings
-from prometheus_client import CollectorRegistry, Gauge, Summary, push_to_gateway
+from prometheus_client import Counter, CollectorRegistry, Gauge, Summary, push_to_gateway
 
 from vaas.metrics.models import Metrics
 
 logger = logging.getLogger(__name__)
 
+Kind = Union[Counter, Summary, Gauge]
+
 
 class PrometheusClient:
     def __init__(self) -> None:
         self.host: str = f'{settings.PROMETHEUS_GATEWAY_HOST}:{settings.PROMETHEUS_GATEWAY_PORT}'
         self.job: str = settings.PROMETHEUS_GATEWAY_JOB
-        self.metrics_bucket: Dict[str, Union[Summary, Gauge]] = {}
+        self.metrics_bucket: Dict[str, Kind] = {}
         self.registry: CollectorRegistry = CollectorRegistry()
 
-    def get_or_create(self, name: str, kind: Type[Union[Summary, Gauge]]) -> Union[Summary, Gauge]:
-        metric: Union[Summary, Gauge, None] = self.metrics_bucket.get(name)
+    def get_or_create(self, name: str, kind: Type[Kind]) -> Kind:
+        metric: Optional[Kind] = self.metrics_bucket.get(name)
         if not metric:
-            new_metrics: Union[Summary, Gauge] = kind(name, name, registry=self.registry)
+            new_metrics: Kind = kind(name, name, registry=self.registry)
             self.metrics_bucket[name] = new_metrics
             return new_metrics
         return metric
@@ -36,6 +38,11 @@ def __init__(self) -> None:
         super().__init__(
             client=PrometheusClient())
 
+    def counter(self, metric_name: str) -> None:
+        if metric_name:
+            self.client.get_or_create(metric_name, Counter).inc()
+            self.client.push()
+
     def sum(self, metric_name: str, value: float) -> None:
         if metric_name:
             self.client.get_or_create(metric_name, Summary).observe(value)

diff --git a/vaas/vaas/metrics/statsd.py b/vaas/vaas/metrics/statsd.py
@@ -11,6 +11,10 @@ def __init__(self) -> None:
             client=StatsClient(host=settings.STATSD_HOST, port=settings.STATSD_PORT, prefix=settings.STATSD_PREFIX)
         )
 
+    def counter(self, metric_name: str) -> None:
+        if metric_name:
+            self.client.incr(metric_name)
+
     def sum(self, metric_name: str, value: float) -> None:
         if metric_name:
             self.client.timing(metric_name, value)

diff --git a/vaas/vaas/metrics/tests/test_handler.py b/vaas/vaas/metrics/tests/test_handler.py
@@ -18,3 +18,7 @@ def test_should_send_metrics_for_defined_protocols(self, mock_statsd_metrics, mo
         metrics.gauge('test', 1)
         mock_statsd_metrics.gauge.assert_called_with('test', 1)
         mock_prometheus_metrics.gauge.assert_called_with('test', 1)
+
+        metrics.counter('test')
+        mock_statsd_metrics.counter.assert_called_with('test')
+        mock_prometheus_metrics.counter.assert_called_with('test')
diff --git a/vaas/vaas/metrics/tests/test_prometheus.py b/vaas/vaas/metrics/tests/test_prometheus.py
@@ -1,6 +1,6 @@
 from unittest.mock import Mock, call, patch
 from django.test import SimpleTestCase, override_settings
-from prometheus_client import Gauge, Summary
+from prometheus_client import Counter, Gauge, Summary
 
 from vaas.metrics.prometheus import PrometheusClient, PrometheusMetrics
 
@@ -46,8 +46,10 @@ def test_should_continue_after_push_to_gateway_raise_an_exception(self):
 class TestPrometheusMetrics(SimpleTestCase):
     def test_should_create_metrics_and_push_them_to_gateway(self):
         with patch('vaas.metrics.prometheus.PrometheusClient') as mock_prometheus_client:
-            mock_summary_observe, mock_gauge_set, mock_prometheus_client.push = Mock(), Mock(), Mock()
-            mock_prometheus_client.get_or_create.return_value = Mock(observe=mock_summary_observe, set=mock_gauge_set)
+            mock_summary_observe, mock_gauge_set, mock_gauge_counter, mock_prometheus_client.push = \
+                Mock(), Mock(), Mock(), Mock()
+            mock_prometheus_client.get_or_create.return_value = Mock(
+                observe=mock_summary_observe, set=mock_gauge_set, inc=mock_gauge_counter)
             prometheus_metrics = PrometheusMetrics()
             prometheus_metrics.client = mock_prometheus_client
 
@@ -61,7 +63,11 @@ def test_should_create_metrics_and_push_them_to_gateway(self):
             mock_prometheus_client.get_or_create.assert_called_with('test', Gauge)
             mock_gauge_set.assert_called_with(1)
 
-            mock_prometheus_client.push.assert_has_calls([call(), call()])
+            prometheus_metrics.counter('test')
+            mock_prometheus_client.get_or_create.assert_called_with('test', Counter)
+            mock_gauge_counter.assert_called_once()
+
+            mock_prometheus_client.push.assert_has_calls([call(), call(), call()])
 
     def test_should_not_create_and_push_metrics_with_empty_names(self):
         with patch('vaas.metrics.prometheus.PrometheusClient') as mock_prometheus_client:
@@ -77,4 +83,7 @@ def test_should_not_create_and_push_metrics_with_empty_names(self):
             prometheus_metrics.gauge('', 1)
             mock_prometheus_client.get_or_create.assert_not_called()
 
+            prometheus_metrics.counter('')
+            mock_prometheus_client.get_or_create.assert_not_called()
+
             mock_prometheus_client.push.assert_not_called()
diff --git a/vaas/vaas/metrics/tests/test_statsd.py b/vaas/vaas/metrics/tests/test_statsd.py
@@ -7,8 +7,8 @@
 class TestStatsdMetrics(SimpleTestCase):
     def test_should_send_defined_metrics(self):
         with patch('statsd.StatsClient') as mock_statsd_client:
-            mock_statsd_timing, mock_statsd_gauge = Mock(), Mock()
-            mock_statsd_client = Mock(timing=mock_statsd_timing, gauge=mock_statsd_gauge)
+            mock_statsd_timing, mock_statsd_gauge, mock_statsd_counter = Mock(), Mock(), Mock()
+            mock_statsd_client = Mock(timing=mock_statsd_timing, gauge=mock_statsd_gauge, incr=mock_statsd_counter)
 
             statsd_metrics = StatsdMetrics()
             statsd_metrics.client = mock_statsd_client
@@ -19,10 +19,13 @@ def test_should_send_defined_metrics(self):
             statsd_metrics.gauge('test', 1)
             mock_statsd_gauge.assert_called_once_with('test', 1)
 
+            statsd_metrics.counter('test')
+            mock_statsd_counter.assert_called_once_with('test')
+
     def test_should_not_send_metric_with_empty_name(self):
         with patch('statsd.StatsClient') as mock_statsd_client:
-            mock_statsd_timing, mock_statsd_gauge = Mock(), Mock()
-            mock_statsd_client = Mock(timing=mock_statsd_timing, gauge=mock_statsd_gauge)
+            mock_statsd_timing, mock_statsd_gauge, mock_statsd_counter = Mock(), Mock(), Mock()
+            mock_statsd_client = Mock(timing=mock_statsd_timing, gauge=mock_statsd_gauge, counter=mock_statsd_counter)
 
             statsd_metrics = StatsdMetrics()
             statsd_metrics.client = mock_statsd_client
@@ -32,3 +35,6 @@ def test_should_not_send_metric_with_empty_name(self):
 
             statsd_metrics.gauge('', 1)
             mock_statsd_gauge.assert_not_called()
+
+            statsd_metrics.counter('')
+            mock_statsd_counter.assert_not_called()