Asynqp consume support (#114)

* Use a callback generator to wrap consumers

* Update tests to follow changes

* Cleanup

Author: pglombardo

instana/instrumentation/asynqp.py

@@ -1,7 +1,5 @@
 from __future__ import absolute_import
 
-import sys
-
 import opentracing
 import wrapt
 
@@ -65,32 +63,36 @@ def get_with_instana(wrapped, instance, argv, kwargs):
 
             return msg
 
-    @wrapt.patch_function_wrapper('asynqp.queue','Consumers.deliver')
-    def deliver_with_instana(wrapped, instance, argv, kwargs):
-
-        ctx = None
-        msg = argv[1]
-        if msg.headers is not None:
-            ctx = async_tracer.extract(opentracing.Format.HTTP_HEADERS, dict(msg.headers))
-
-        with async_tracer.start_active_span("rabbitmq", child_of=ctx) as scope:
-            host, port = argv[1].sender.protocol.transport._sock.getsockname()
-
-            try:
-                scope.span.set_tag("exchange", msg.exchange_name)
-                scope.span.set_tag("sort", "consume")
-                scope.span.set_tag("address", host + ":" + str(port) )
-                scope.span.set_tag("key", msg.routing_key)
-
-                rv = wrapped(*argv, **kwargs)
-            except Exception as e:
-                scope.span.log_kv({'message': e})
-                scope.span.set_tag("error", True)
-                ec = scope.span.tags.get('ec', 0)
-                scope.span.set_tag("ec", ec+1)
-                raise
-            else:
-                return rv
+    @wrapt.patch_function_wrapper('asynqp.queue','Queue.consume')
+    def consume_with_instana(wrapped, instance, argv, kwargs):
+        def callback_generator(original_callback):
+            def callback_with_instana(*argv, **kwargs):
+                ctx = None
+                msg = argv[0]
+                if msg.headers is not None:
+                    ctx = async_tracer.extract(opentracing.Format.HTTP_HEADERS, dict(msg.headers))
+
+                with async_tracer.start_active_span("rabbitmq", child_of=ctx) as scope:
+                    host, port = msg.sender.protocol.transport._sock.getsockname()
+
+                    try:
+                        scope.span.set_tag("exchange", msg.exchange_name)
+                        scope.span.set_tag("sort", "consume")
+                        scope.span.set_tag("address", host + ":" + str(port) )
+                        scope.span.set_tag("key", msg.routing_key)
+
+                        original_callback(*argv, **kwargs)
+                    except Exception as e:
+                        scope.span.log_kv({'message': e})
+                        scope.span.set_tag("error", True)
+                        ec = scope.span.tags.get('ec', 0)
+                        scope.span.set_tag("ec", ec+1)
+                        raise
+            return callback_with_instana
+
+        cb = argv[0]
+        argv = (callback_generator(cb),)
+        return wrapped(*argv, **kwargs)
 
     logger.debug("Instrumenting asynqp")
 except ImportError:

tests/test_asynqp.py

@@ -218,27 +218,24 @@ def test():
         self.loop.run_until_complete(test())
 
         spans = self.recorder.queued_spans()
-        self.assertEqual(5, len(spans))
+        self.assertEqual(4, len(spans))
 
         publish1_span = spans[0]
-        consume1_span = spans[1]
-        publish2_span = spans[2]
-        consume2_span = spans[3]
-        test_span = spans[4]
+        publish2_span = spans[1]
+        consume1_span = spans[2]
+        test_span = spans[3]
 
         self.assertIsNone(async_tracer.active_span)
 
         # Same traceId
         self.assertEqual(test_span.t, publish1_span.t)
         self.assertEqual(test_span.t, publish2_span.t)
         self.assertEqual(test_span.t, consume1_span.t)
-        self.assertEqual(test_span.t, consume2_span.t)
 
         # Parent relationships
         self.assertEqual(publish1_span.p, test_span.s)
-        self.assertEqual(publish2_span.p, test_span.s)
         self.assertEqual(consume1_span.p, publish1_span.s)
-        self.assertEqual(consume2_span.p, publish2_span.s)
+        self.assertEqual(publish2_span.p, consume1_span.s)
 
         # publish
         self.assertEqual('test.exchange', publish1_span.data.rabbitmq.exchange)
@@ -266,21 +263,11 @@ def test():
         self.assertTrue(type(consume1_span.stack) is list)
         self.assertGreater(len(consume1_span.stack), 0)
 
-        self.assertEqual('test.exchange', consume2_span.data.rabbitmq.exchange)
-        self.assertEqual('consume', consume2_span.data.rabbitmq.sort)
-        self.assertIsNotNone(consume2_span.data.rabbitmq.address)
-        self.assertEqual('another.key', consume2_span.data.rabbitmq.key)
-        self.assertIsNotNone(consume2_span.stack)
-        self.assertTrue(type(consume2_span.stack) is list)
-        self.assertGreater(len(consume2_span.stack), 0)
-
         # Error logging
         self.assertFalse(test_span.error)
         self.assertIsNone(test_span.ec)
         self.assertFalse(consume1_span.error)
         self.assertIsNone(consume1_span.ec)
-        self.assertFalse(consume2_span.error)
-        self.assertIsNone(consume2_span.ec)
         self.assertFalse(publish1_span.error)
         self.assertIsNone(publish1_span.ec)
         self.assertFalse(publish2_span.error)