Asyncio: Carry context along with Tasks (#139)

* Carry context across ensure_future calls

* Add test without context

* Add support for create_task (with tests)

* Add package configurator module

* Add section describing configuratior

* Use configurator in tests

Author: pglombardo

Configuration.md

@@ -25,6 +25,22 @@ or
 instana.service_name = "myservice"
 ```
 
+## Package Configuration
+
+The Instana package includes a runtime configuration module that manages the configuration of various components.
+
+_Note: as the package evolves, more options will be added here_
+
+```python
+from instana.configurator import config
+
+# To enable tracing context propagation across Asyncio ensure_future and create_task calls
+# Default is false
+config['asyncio_task_context_propagation']['enabled'] = True
+
+```
+
+
 ## Debugging & More Verbosity
 
 Setting `INSTANA_DEV` to a non nil value will enable extra logging output generally useful

instana/__init__.py

@@ -62,6 +62,7 @@ def boot_agent():
     if "INSTANA_DISABLE_AUTO_INSTR" not in os.environ:
         # Import & initialize instrumentation
         if sys.version_info >= (3, 5, 3):
+            from .instrumentation import asyncio  # noqa
             from .instrumentation.aiohttp import client  # noqa
             from .instrumentation.aiohttp import server  # noqa
             from .instrumentation import asynqp  # noqa

instana/configurator.py

@@ -0,0 +1,28 @@
+from __future__ import absolute_import
+from collections import defaultdict
+
+# This file contains a config object that will hold configuration options for the package.
+# Defaults are set and can be overridden after package load.
+
+
+# Simple implementation of a nested dictionary.
+#
+# Same as:
+#   stan_dictionary = lambda: defaultdict(stan_dictionary)
+# but we use the function form because of PEP 8
+#
+def stan_dictionary():
+    return defaultdict(stan_dictionary)
+
+
+# La Protagonista
+config = stan_dictionary()
+
+
+# This option determines if tasks created via asyncio (with ensure_future or create_task) will
+# automatically carry existing context into the created task.
+config['asyncio_task_context_propagation']['enabled'] = False
+
+
+
+

instana/instrumentation/asyncio.py

@@ -0,0 +1,41 @@
+from __future__ import absolute_import
+
+import wrapt
+
+from ..log import logger
+from ..singletons import async_tracer
+from ..configurator import config
+
+try:
+    import asyncio
+
+    @wrapt.patch_function_wrapper('asyncio','ensure_future')
+    def ensure_future_with_instana(wrapped, instance, argv, kwargs):
+        if config['asyncio_task_context_propagation']['enabled'] is False:
+            return wrapped(*argv, **kwargs)
+
+        scope = async_tracer.scope_manager.active
+        task = wrapped(*argv, **kwargs)
+
+        if scope is not None:
+            async_tracer.scope_manager._set_task_scope(scope, task=task)
+
+        return task
+
+    if hasattr(asyncio, "create_task"):
+        @wrapt.patch_function_wrapper('asyncio','create_task')
+        def create_task_with_instana(wrapped, instance, argv, kwargs):
+            if config['asyncio_task_context_propagation']['enabled'] is False:
+                return wrapped(*argv, **kwargs)
+
+            scope = async_tracer.scope_manager.active
+            task = wrapped(*argv, **kwargs)
+
+            if scope is not None:
+                async_tracer.scope_manager._set_task_scope(scope, task=task)
+
+            return task
+
+    logger.debug("Instrumenting asyncio")
+except ImportError:
+    pass

runtests.py

@@ -5,7 +5,7 @@
 command_line = [__file__, '--verbose']
 
 if (LooseVersion(sys.version) < LooseVersion('3.5.3')):
-    command_line.extend(['-e', 'asynqp', '-e', 'aiohttp'])
+    command_line.extend(['-e', 'asynqp', '-e', 'aiohttp', '-e', 'async'])
 
 if (LooseVersion(sys.version) >= LooseVersion('3.7.0')):
     command_line.extend(['-e', 'sudsjurko'])

tests/test_asyncio.py

@@ -0,0 +1,141 @@
+from __future__ import absolute_import
+
+import asyncio
+import unittest
+
+import aiohttp
+
+from instana.singletons import async_tracer
+from instana.configurator import config
+
+from .helpers import testenv
+
+
+class TestAsyncio(unittest.TestCase):
+    def setUp(self):
+        """ Clear all spans before a test run """
+        self.recorder = async_tracer.recorder
+        self.recorder.clear_spans()
+
+        # New event loop for every test
+        self.loop = asyncio.new_event_loop()
+        asyncio.set_event_loop(None)
+
+        # Restore default
+        config['asyncio_task_context_propagation']['enabled'] = False
+
+    def tearDown(self):
+        """ Purge the queue """
+        pass
+
+    async def fetch(self, session, url, headers=None):
+        try:
+            async with session.get(url, headers=headers) as response:
+                return response
+        except aiohttp.web_exceptions.HTTPException:
+            pass
+
+    def test_ensure_future_with_context(self):
+        async def run_later(msg="Hello"):
+            # print("run_later: %s" % async_tracer.active_span.operation_name)
+            async with aiohttp.ClientSession() as session:
+                return await self.fetch(session, testenv["wsgi_server"] + "/")
+
+        async def test():
+            with async_tracer.start_active_span('test'):
+                asyncio.ensure_future(run_later("Hello"))
+            await asyncio.sleep(0.5)
+
+        # Override default task context propagation
+        config['asyncio_task_context_propagation']['enabled'] = True
+
+        self.loop.run_until_complete(test())
+
+        spans = self.recorder.queued_spans()
+        self.assertEqual(3, len(spans))
+
+        test_span = spans[0]
+        wsgi_span = spans[1]
+        aioclient_span = spans[2]
+
+        self.assertEqual(test_span.t, wsgi_span.t)
+        self.assertEqual(test_span.t, aioclient_span.t)
+
+        self.assertEqual(test_span.p, None)
+        self.assertEqual(wsgi_span.p, aioclient_span.s)
+        self.assertEqual(aioclient_span.p, test_span.s)
+
+    def test_ensure_future_without_context(self):
+        async def run_later(msg="Hello"):
+            # print("run_later: %s" % async_tracer.active_span.operation_name)
+            async with aiohttp.ClientSession() as session:
+                return await self.fetch(session, testenv["wsgi_server"] + "/")
+
+        async def test():
+            with async_tracer.start_active_span('test'):
+                asyncio.ensure_future(run_later("Hello"))
+            await asyncio.sleep(0.5)
+
+        self.loop.run_until_complete(test())
+
+        spans = self.recorder.queued_spans()
+
+        self.assertEqual(2, len(spans))
+        self.assertEqual("sdk", spans[0].n)
+        self.assertEqual("wsgi", spans[1].n)
+
+        # Without the context propagated, we should get two separate traces
+        self.assertNotEqual(spans[0].t, spans[1].t)
+
+    if hasattr(asyncio, "create_task"):
+        def test_create_task_with_context(self):
+            async def run_later(msg="Hello"):
+                # print("run_later: %s" % async_tracer.active_span.operation_name)
+                async with aiohttp.ClientSession() as session:
+                    return await self.fetch(session, testenv["wsgi_server"] + "/")
+
+            async def test():
+                with async_tracer.start_active_span('test'):
+                    asyncio.create_task(run_later("Hello"))
+                await asyncio.sleep(0.5)
+
+            # Override default task context propagation
+            config['asyncio_task_context_propagation']['enabled'] = True
+
+            self.loop.run_until_complete(test())
+
+            spans = self.recorder.queued_spans()
+            self.assertEqual(3, len(spans))
+
+            test_span = spans[0]
+            wsgi_span = spans[1]
+            aioclient_span = spans[2]
+
+            self.assertEqual(test_span.t, wsgi_span.t)
+            self.assertEqual(test_span.t, aioclient_span.t)
+
+            self.assertEqual(test_span.p, None)
+            self.assertEqual(wsgi_span.p, aioclient_span.s)
+            self.assertEqual(aioclient_span.p, test_span.s)
+
+        def test_create_task_without_context(self):
+            async def run_later(msg="Hello"):
+                # print("run_later: %s" % async_tracer.active_span.operation_name)
+                async with aiohttp.ClientSession() as session:
+                    return await self.fetch(session, testenv["wsgi_server"] + "/")
+
+            async def test():
+                with async_tracer.start_active_span('test'):
+                    asyncio.create_task(run_later("Hello"))
+                await asyncio.sleep(0.5)
+
+            self.loop.run_until_complete(test())
+
+            spans = self.recorder.queued_spans()
+
+            self.assertEqual(2, len(spans))
+            self.assertEqual("sdk", spans[0].n)
+            self.assertEqual("wsgi", spans[1].n)
+
+            # Without the context propagated, we should get two separate traces
+            self.assertNotEqual(spans[0].t, spans[1].t)

tests/test_configurator.py

@@ -0,0 +1,16 @@
+from __future__ import absolute_import
+
+import unittest
+
+from instana.configurator import config
+
+
+class TestRedis(unittest.TestCase):
+    def setUp(self):
+        pass
+
+    def tearDown(self):
+        pass
+
+    def test_has_default_config(self):
+        self.assertEqual(config['asyncio_task_context_propagation']['enabled'], False)