Repository: google/turbinia

Issue:
Prometheus Duplicated timeseries in CollectorRegistry
The second time a Task runs with the same name, the prometheus code errors out due to duplicate timeseries:

```
2021-03-02 00:57:10 [INFO] Starting Task PartitionEnumerationTask 84181e3126ca4fa5ba2acda57d93322c
[2021-03-02 00:57:10,236: INFO/MainProcess] Lock 140581005659272 released on /evidence/turbinia-worker.lock
[2021-03-02 00:57:10,241: ERROR/MainProcess] Task task_runner[3de967e0-099d-430c-a1a0-1b635681ab78] raised unexpected: ValueError("Duplicated timeseries in CollectorRegistry: {'PartitionEnumerationTask_duration_seconds_created', 'PartitionEnumerationTask_duration_seconds_count', 'PartitionEnumerationTask_duration_seconds_bucket', 'PartitionEnumerationTask_duration_seconds', 'PartitionEnumerationTask_duration_seconds_sum'}",)
Traceback (most recent call last):
  File "/usr/local/lib/python3.6/dist-packages/celery/app/trace.py", line 412, in trace_task
    R = retval = fun(*args, **kwargs)
  File "/usr/local/lib/python3.6/dist-packages/celery/app/trace.py", line 704, in __protected_call__
    return self.run(*args, **kwargs)
  File "/workspaces/turbinia/turbinia/task_manager.py", line 97, in task_runner
    return obj.run_wrapper(*args, **kwargs)
  File "/workspaces/turbinia/turbinia/workers/__init__.py", line 785, in run_wrapper
    'Seconds to run {0:s}'.format(self.name))
  File "/usr/local/lib/python3.6/dist-packages/prometheus_client/metrics.py", line 517, in __init__
    labelvalues=labelvalues,
  File "/usr/local/lib/python3.6/dist-packages/prometheus_client/metrics.py", line 121, in __init__
    registry.register(self)
  File "/usr/local/lib/python3.6/dist-packages/prometheus_client/registry.py", line 31, in register
    duplicates))
ValueError: Duplicated timeseries in CollectorRegistry: {'PartitionEnumerationTask_duration_seconds_created', 'PartitionEnumerationTask_duration_seconds_count', 'PartitionEnumerationTask_duration_seconds_bucket', 'PartitionEnumerationTask_duration_seconds', 'PartitionEnumerationTask_duration_seconds_sum'}
```


Discussion:
@alimez Could you take a look at this one?
this is impacting current e2e tests as well as it makes the FileExtractionTask fail (and all upstream tasks as such).

Resolving patch:
diff --git a/turbinia/client.py b/turbinia/client.py
--- a/turbinia/client.py
+++ b/turbinia/client.py
@@ -78,6 +78,7 @@
     'jenkinsanalysistask': JenkinsAnalysisTask,
     'JupyterAnalysisTask': JupyterAnalysisTask,
     'greptask': GrepTask,
+    'fsstattask': FsstatTask,
     'hadoopanalysistask': HadoopAnalysisTask,
     'hindsighttask': HindsightTask,
     'partitionenumerationtask': PartitionEnumerationTask,
@@ -90,7 +91,7 @@
     'tomcatanalysistask': TomcatAnalysisTask,
     'volatilitytask': VolatilityTask,
     'stattask': StatTask,
-    'binaryextractor': BinaryExtractorTask,
+    'binaryextractortask': BinaryExtractorTask,
     'bulkextractortask': BulkExtractorTask,
     'dockertask': DockerContainersEnumerationTask,
     'photorectask': PhotorecTask
diff --git a/turbinia/workers/__init__.py b/turbinia/workers/__init__.py
--- a/turbinia/workers/__init__.py
+++ b/turbinia/workers/__init__.py
@@ -46,6 +46,8 @@
 from prometheus_client import Gauge
 from prometheus_client import Histogram
 
+METRICS = {}
+
 log = logging.getLogger('turbinia')
 
 turbinia_worker_tasks_started_total = Gauge(
@@ -507,6 +509,15 @@ def evidence_setup(self, evidence):
             'information.'.format(
                 evidence, self.name, state.name, evidence.format_state()))
 
+  def get_metrics(self):
+    """Gets histogram metric for current Task.
+
+    Returns:
+      prometheus_client.Historgram: For the current task, or None if they are not initialized.
+    """
+    global METRICS
+    return METRICS.get(self.name.lower())
+
   def execute(
       self, cmd, result, save_files=None, log_files=None, new_evidence=None,
       close=False, shell=False, stderr_file=None, stdout_file=None,
@@ -665,6 +676,7 @@ def setup(self, evidence):
     Raises:
       TurbiniaException: If the evidence can not be found.
     """
+    self.setup_metrics()
     self.output_manager.setup(self.name, self.id)
     self.tmp_dir, self.output_dir = self.output_manager.get_local_output_dirs()
     if not self.result:
@@ -680,6 +692,40 @@ def setup(self, evidence):
               evidence.source_path))
     return self.result
 
+  def setup_metrics(self, task_map=None):
+    """Sets up the application metrics.
+
+    Returns early with metrics if they are already setup.
+
+    Arguments:
+      task_map(dict): Map of task names to task objects
+
+    Returns:
+      Dict: Mapping of task names to metrics objects.
+    """
+    global METRICS
+
+    if METRICS:
+      return METRICS
+
+    if not task_map:
+      # Late import to avoid circular dependencies
+      from turbinia.client import TASK_MAP
+      task_map = TASK_MAP
+
+    for task_name in task_map:
+      task_name = task_name.lower()
+      if task_name in METRICS:
+        continue
+      metric = Histogram(
+          '{0:s}_duration_seconds'.format(task_name),
+          'Seconds to run {0:s}'.format(task_name))
+      METRICS[task_name] = metric
+
+    log.debug('Registered {0:d} task metrics'.format(len(METRICS)))
+
+    return METRICS
+
   def touch(self):
     """Updates the last_update time of the task."""
     self.last_update = datetime.now()
@@ -817,9 +863,7 @@ def run_wrapper(self, evidence):
       log.info('Starting Task {0:s} {1:s}'.format(self.name, self.id))
       original_result_id = None
       turbinia_worker_tasks_started_total.inc()
-      task_runtime_metrics = Histogram(
-          '{0:s}_duration_seconds'.format(self.name),
-          'Seconds to run {0:s}'.format(self.name))
+      task_runtime_metrics = self.get_metrics()
       with task_runtime_metrics.time():
         try:
           original_result_id = self.result.id


Test patch:
diff --git a/turbinia/workers/workers_test.py b/turbinia/workers/workers_test.py
--- a/turbinia/workers/workers_test.py
+++ b/turbinia/workers/workers_test.py
@@ -61,6 +61,7 @@ def setUp(self, task_class=TurbiniaTask, evidence_class=evidence.RawDisk):
     self.task.job_name = 'PlasoJob'
     self.task.output_manager = mock.MagicMock()
     self.task.output_manager.get_local_output_dirs.return_value = (None, None)
+    self.task.get_metrics = mock.MagicMock()
 
     # Set up RawDisk Evidence
     test_disk_path = tempfile.mkstemp(dir=self.base_output_dir)[1]
@@ -337,6 +338,16 @@ def testTurbiniaTaskExecuteEvidenceExistsButEmpty(self, popen_mock):
         cmd, self.result, new_evidence=[self.evidence], close=True)
     self.assertNotIn(self.evidence, self.result.evidence)
 
+  @mock.patch('turbinia.workers.Histogram')
+  def testTurbiniaSetupMetrics(self, mock_histogram):
+    """Tests that metrics are set up correctly."""
+    mock_task_map = {'TestTask1': None, 'TestTask2': None}
+    mock_histogram.return_value = "test_metrics"
+    metrics = self.task.setup_metrics(task_map=mock_task_map)
+    self.assertEqual(len(metrics), len(mock_task_map))
+    self.assertEqual(metrics['testtask1'], 'test_metrics')
+    self.assertIn('testtask1', metrics)
+
   def testEvidenceSetup(self):
     """Tests basic run of evidence_setup."""
     self.evidence.preprocess = mock.MagicMock()
@@ -357,4 +368,4 @@ def testEvidenceSetupStateNotFulfilled(self):
 
     # Runs fine after setting the state
     self.evidence.state[evidence.EvidenceState.ATTACHED] = True
-    self.task.evidence_setup(self.evidence)
+    self.task.evidence_setup(self.evidence)
\ No newline at end of file