added async_tracer support for pymongo instrumentation (#321)

Added async_tracer support for pymongo instrumentation

* refactor the tracer discovery functionality into a tracer utility method

Author: pdimitra

instana/instrumentation/pymongo.py

@@ -4,25 +4,27 @@
 from __future__ import absolute_import
 
 from ..log import logger
-from ..singletons import tracer
+from ..util.traceutils import get_active_tracer
 
 try:
     import pymongo
     from pymongo import monitoring
     from bson import json_util
 
+
     class MongoCommandTracer(monitoring.CommandListener):
         def __init__(self):
             self.__active_commands = {}
 
         def started(self, event):
-            parent_span = tracer.active_span
-
+            active_tracer = get_active_tracer()
             # return early if we're not tracing
-            if parent_span is None:
+            if active_tracer is None:
                 return
 
-            with tracer.start_active_span("mongo", child_of=parent_span) as scope:
+            parent_span = active_tracer.active_span
+
+            with active_tracer.start_active_span("mongo", child_of=parent_span) as scope:
                 self._collect_connection_tags(scope.span, event)
                 self._collect_command_tags(scope.span, event)
 
@@ -79,7 +81,7 @@ def _collect_command_tags(self, span, event):
             cmd_doc = None
             if cmd in cmd_doc_locations:
                 cmd_doc = event.command.get(cmd_doc_locations[cmd])
-            elif cmd.lower() == "mapreduce": # mapreduce command was renamed to mapReduce in pymongo 3.9.0
+            elif cmd.lower() == "mapreduce":  # mapreduce command was renamed to mapReduce in pymongo 3.9.0
                 # mapreduce command consists of two mandatory parts: map and reduce
                 cmd_doc = {
                     "map": event.command.get("map"),
@@ -89,6 +91,7 @@ def _collect_command_tags(self, span, event):
             if cmd_doc is not None:
                 span.set_tag("json", json_util.dumps(cmd_doc))
 
+
     monitoring.register(MongoCommandTracer())
 
     logger.debug("Instrumenting pymongo")

instana/singletons.py

@@ -11,6 +11,7 @@
 
 agent = None
 tracer = None
+async_tracer = None
 profiler = None
 span_recorder = None
 
@@ -79,11 +80,11 @@ def set_agent(new_agent):
 if sys.version_info >= (3, 4):
     try:
         from opentracing.scope_managers.asyncio import AsyncioScopeManager
+
         async_tracer = InstanaTracer(scope_manager=AsyncioScopeManager(), recorder=span_recorder)
     except Exception:
         logger.debug("Error setting up async_tracer:", exc_info=True)
 
-
 # Mock the tornado tracer until tornado is detected and instrumented first
 tornado_tracer = tracer
 
@@ -117,6 +118,7 @@ def set_tracer(new_tracer):
     global tracer
     tracer = new_tracer
 
+
 def get_profiler():
     """
     Retrieve the globally configured profiler

instana/util/traceutils.py

@@ -1,7 +1,7 @@
 # (c) Copyright IBM Corp. 2021
 # (c) Copyright Instana Inc. 2021
 
-from ..singletons import agent
+from ..singletons import agent, tracer, async_tracer
 from ..log import logger
 
 
@@ -14,3 +14,16 @@ def extract_custom_headers(tracing_scope, headers):
                     tracing_scope.span.set_tag("http.header.%s" % custom_header, value)
     except Exception as e:
         logger.debug("extract_custom_headers: ", exc_info=True)
+
+
+def get_active_tracer():
+    try:
+        if tracer.active_span:
+            return tracer
+        elif async_tracer.active_span:
+            return async_tracer
+        else:
+            return None
+    except Exception as e:
+        logger.debug("error while getting active tracer: ", exc_info=True)
+        return None

instana/version.py

@@ -3,4 +3,4 @@
 
 # Module version file.  Used by setup.py and snapshot reporting.
 
-VERSION = '1.33.1'
+VERSION = '1.33.2'

tests/clients/test_pymongo.py

@@ -19,7 +19,7 @@
 logger = logging.getLogger(__name__)
 
 
-class TestPyMongo(unittest.TestCase):
+class TestPyMongoTracer(unittest.TestCase):
     def setUp(self):
         self.conn = pymongo.MongoClient(host=testenv['mongodb_host'], port=int(testenv['mongodb_port']),
                                         username=testenv['mongodb_user'], password=testenv['mongodb_pw'])
@@ -107,11 +107,11 @@ def test_successful_update_query(self):
 
         payload = json.loads(db_span.data["mongo"]["json"])
         assert_true({
-            "q": {"type": "string"},
-            "u": {"$set": {"type": "int"}},
-            "multi": False,
-            "upsert": False
-        } in payload, db_span.data["mongo"]["json"])
+                        "q": {"type": "string"},
+                        "u": {"$set": {"type": "int"}},
+                        "multi": False,
+                        "upsert": False
+                    } in payload, db_span.data["mongo"]["json"])
 
     def test_successful_delete_query(self):
         with tracer.start_active_span("test"):
@@ -174,7 +174,8 @@ def test_successful_map_reduce_query(self):
         reducer = "function (key, values) { return len(values); }"
 
         with tracer.start_active_span("test"):
-            self.conn.test.records.map_reduce(bson.code.Code(mapper), bson.code.Code(reducer), "results", query={"x": {"$lt": 2}})
+            self.conn.test.records.map_reduce(bson.code.Code(mapper), bson.code.Code(reducer), "results",
+                                              query={"x": {"$lt": 2}})
 
         assert_is_none(tracer.active_span)
 
@@ -192,7 +193,8 @@ def test_successful_map_reduce_query(self):
         self.assertEqual(db_span.n, "mongo")
         self.assertEqual(db_span.data["mongo"]["service"], "%s:%s" % (testenv['mongodb_host'], testenv['mongodb_port']))
         self.assertEqual(db_span.data["mongo"]["namespace"], "test.records")
-        self.assertEqual(db_span.data["mongo"]["command"].lower(), "mapreduce") # mapreduce command was renamed to mapReduce in pymongo 3.9.0
+        self.assertEqual(db_span.data["mongo"]["command"].lower(),
+                         "mapreduce")  # mapreduce command was renamed to mapReduce in pymongo 3.9.0
 
         self.assertEqual(db_span.data["mongo"]["filter"], '{"x": {"$lt": 2}}')
         assert_is_not_none(db_span.data["mongo"]["json"])
@@ -228,3 +230,4 @@ def test_successful_mutiple_queries(self):
 
         # ensure spans are ordered the same way as commands
         assert_list_equal(commands, ["insert", "update", "delete"])
+