From 977fc98ed2ae349faf33ddbfee419cf19263e4bd Mon Sep 17 00:00:00 2001
From: marcsello <punkosdmarcell@rocketmail.com>
Date: Wed, 4 Aug 2021 15:00:41 +0200
Subject: [PATCH] Fixed tracing

---
 cnn_classification_service/cnn_classifier.py |  8 +++++---
 cnn_classification_service/magic_doer.py     | 12 +++++++-----
 cnn_classification_service/main.py           |  6 +++---
 3 files changed, 15 insertions(+), 11 deletions(-)

diff --git a/cnn_classification_service/cnn_classifier.py b/cnn_classification_service/cnn_classifier.py
index 7d35684..4d368a1 100644
--- a/cnn_classification_service/cnn_classifier.py
+++ b/cnn_classification_service/cnn_classifier.py
@@ -84,11 +84,13 @@ class Classifier(object):
 
         return predicted_class_name, labeled_predictions
 
-    def predict(self, wav_filename: str, span: opentracing.span.Span) -> Tuple[str, dict]:
-        with opentracing.tracer.start_span('createSpectrogram', child_of=span):
+    def predict(self, wav_filename: str) -> Tuple[str, dict]:
+        span = opentracing.tracer.scope_manager.active().span
+
+        with opentracing.tracer.start_active_span('createSpectrogram', child_of=span):
             directory, _ = self.create_spectrogram(wav_filename)
 
-        with opentracing.tracer.start_span('runPredictor', child_of=span):
+        with opentracing.tracer.start_active_span('runPredictor', child_of=span):
             result = self._run_predictor(directory)
 
         shutil.rmtree(directory)  # The image is no longer needed
diff --git a/cnn_classification_service/magic_doer.py b/cnn_classification_service/magic_doer.py
index a858050..337dec2 100644
--- a/cnn_classification_service/magic_doer.py
+++ b/cnn_classification_service/magic_doer.py
@@ -18,14 +18,16 @@ class MagicDoer:
     requests_session = SessionTracing(propagate=True)
 
     @classmethod
-    def run_everything(cls, parameters: dict, span: opentracing.span.Span) -> dict:
+    def run_everything(cls, parameters: dict) -> dict:
+        span = opentracing.tracer.scope_manager.active().span
+
         tag = parameters['tag']
         sample_file_handle, sample_file_path = tempfile.mkstemp(prefix=f"{tag}_", suffix=".wav", dir="/dev/shm")
         span.log_kv({'event': 'sampleFileOpened', 'sampleTag': tag})
         response = None
         try:
 
-            with opentracing.tracer.start_span('downloadSample', child_of=span):
+            with opentracing.tracer.start_active_span('downloadSample', child_of=span):
                 # Download Sample
                 object_path = urljoin(Config.STORAGE_SERVICE_URL, f"object/{tag}")
 
@@ -36,14 +38,14 @@ class MagicDoer:
 
                 logging.debug(f"Downloaded sample to {sample_file_path}")
 
-            with opentracing.tracer.start_span('loadClassifier', child_of=span):
+            with opentracing.tracer.start_active_span('loadClassifier', child_of=span):
                 # Get a classifier that uses the default model
                 model_details, classifier = cls.classifier_cache.get_default_classifier()
 
-            with opentracing.tracer.start_span('runClassifier', child_of=span) as child_span:
+            with opentracing.tracer.start_active_span('runClassifier', child_of=span) as child_span:
                 # do the majic
                 classification_start_time = time.time()
-                predicted_class_name, labeled_predictions = classifier.predict(sample_file_path, child_span)
+                predicted_class_name, labeled_predictions = classifier.predict(sample_file_path)
                 classification_duration = time.time() - classification_start_time
 
             response = {
diff --git a/cnn_classification_service/main.py b/cnn_classification_service/main.py
index 9fe8794..c28b37e 100644
--- a/cnn_classification_service/main.py
+++ b/cnn_classification_service/main.py
@@ -19,7 +19,7 @@ from magic_doer import MagicDoer
 
 
 def message_callback(channel, method, properties, body):
-    with opentracing.tracer.start_span('messageHandling') as span:
+    with opentracing.tracer.start_active_span('messageHandling', finish_on_close=True) as span:
         try:
             msg = json.loads(body.decode('utf-8'))
         except (UnicodeDecodeError, json.JSONDecodeError) as e:
@@ -28,8 +28,8 @@ def message_callback(channel, method, properties, body):
 
         span.log_kv({'event': 'messageParsed', 'sampleTag': msg['tag']})
 
-        with opentracing.tracer.start_span('runAlgorithm', child_of=span) as child_span:
-            results = MagicDoer.run_everything(msg, child_span)  # <- This is where the magic happens
+        with opentracing.tracer.start_active_span('runAlgorithm', child_of=span) as child_span:
+            results = MagicDoer.run_everything(msg)  # <- This is where the magic happens
 
         if results:
             channel.basic_publish(