only load metrics once in workers module

aarontp · Jun 13, 2024 · a606942 · a606942
1 parent 23d4a8d
commit a606942
Show file tree

Hide file tree

Showing 3 changed files with 29 additions and 21 deletions.
diff --git a/turbinia/worker_test.py b/turbinia/worker_test.py
@@ -55,7 +55,7 @@ def tearDown(self):
 
   def unregisterMetrics(self):
     """Unset all the metrics to avoid duplicated timeseries error."""
-    for collector, names in tuple(REGISTRY._get_names()):
+    for collector, names in tuple(REGISTRY._collector_to_names.items()):
       REGISTRY.unregister(collector)
 
   @mock.patch('turbinia.client.task_manager.CeleryTaskManager._backend_setup')

diff --git a/turbinia/workers/__init__.py b/turbinia/workers/__init__.py
@@ -21,6 +21,7 @@
 from datetime import timedelta
 from enum import IntEnum
 
+from itertools import chain
 import json
 import logging
 import os
@@ -48,6 +49,7 @@
 from turbinia import log_and_report
 
 from celery.exceptions import SoftTimeLimitExceeded
+from prometheus_client import REGISTRY
 
 METRICS = {}
 # Set the maximum size that the report can be before truncating it.  This is a
@@ -59,25 +61,30 @@
 
 log = logging.getLogger(__name__)
 
-turbinia_worker_exception_failure = Counter(
-    'turbinia_worker_exception_failure',
-    'Total number Tasks failed due to uncaught exception')
-turbinia_worker_tasks_started_total = Counter(
-    'turbinia_worker_tasks_started_total',
-    'Total number of started worker tasks')
-turbinia_worker_tasks_completed_total = Counter(
-    'turbinia_worker_tasks_completed_total',
-    'Total number of completed worker tasks')
-turbinia_worker_tasks_queued_total = Counter(
-    'turbinia_worker_tasks_queued_total', 'Total number of queued worker tasks')
-turbinia_worker_tasks_failed_total = Counter(
-    'turbinia_worker_tasks_failed_total', 'Total number of failed worker tasks')
-turbinia_worker_tasks_timeout_total = Counter(
-    'turbinia_worker_tasks_timeout_total',
-    'Total number of worker tasks timed out during dependency execution.')
-turbinia_worker_tasks_timeout_celery_soft = Counter(
-    'turbinia_worker_tasks_timeout_celery_soft',
-    'Total number of Tasks timed out due to Celery soft timeout')
+# Prevent re-registering metrics if module is loaded multiple times.
+metric_names = list(chain.from_iterable(REGISTRY._collector_to_names.values()))
+if 'turbinia_worker_exception_failure' not in metric_names:
+  turbinia_worker_exception_failure = Counter(
+      'turbinia_worker_exception_failure',
+      'Total number Tasks failed due to uncaught exception')
+  turbinia_worker_tasks_started_total = Counter(
+      'turbinia_worker_tasks_started_total',
+      'Total number of started worker tasks')
+  turbinia_worker_tasks_completed_total = Counter(
+      'turbinia_worker_tasks_completed_total',
+      'Total number of completed worker tasks')
+  turbinia_worker_tasks_queued_total = Counter(
+      'turbinia_worker_tasks_queued_total',
+      'Total number of queued worker tasks')
+  turbinia_worker_tasks_failed_total = Counter(
+      'turbinia_worker_tasks_failed_total',
+      'Total number of failed worker tasks')
+  turbinia_worker_tasks_timeout_total = Counter(
+      'turbinia_worker_tasks_timeout_total',
+      'Total number of worker tasks timed out during dependency execution.')
+  turbinia_worker_tasks_timeout_celery_soft = Counter(
+      'turbinia_worker_tasks_timeout_celery_soft',
+      'Total number of Tasks timed out due to Celery soft timeout')
 
 
 class Priority(IntEnum):

diff --git a/turbinia/workers/workers_test.py b/turbinia/workers/workers_test.py
@@ -119,7 +119,8 @@ def setResults(
 
   def unregisterMetrics(self):
     """Unset all the metrics to avoid duplicated timeseries error."""
-    for collector, names in tuple(REGISTRY._get_names()):
+    for collector, names in tuple(REGISTRY._collector_to_names.items()):
+      print(collector, names)
       REGISTRY.unregister(collector)