[kombu] add Kombu integration (#515)

Adding routing key to tags

Initial tracing of Consumer.receive

Using the Propagator for kombu publish

Addressing first bit of feedback, more to come.

Rolling out tests for python 2.7-3.6 and Kombu 4.0,4.1,4.2

Fixing flake8 issue, adding a second test, cleaning up utils.

Actually tell CircleCI to run Kombu (oops!)

Use the config settings for the service name

Remove superflous function wrapper, remove extra patch call in example

PR Feedback, documentation improvements, rebase from master

Updating circleci config

Fixing CircleCI

Update ddtrace/contrib/kombu/patch.py

Co-Authored-By: tebriel <[email protected]>

Update ddtrace/contrib/kombu/patch.py

Co-Authored-By: tebriel <[email protected]>

Author: tebriel

.circleci/config.yml

@@ -695,6 +695,21 @@ jobs:
             - vertica.results
       - *save_cache_step
 
+  kombu:
+    docker:
+    - *test_runner
+    - image: rabbitmq:3.7-alpine
+    steps:
+    - checkout
+    - *restore_cache_step
+    - run: tox -e wait rabbitmq
+    - run: tox -e 'kombu_contrib-{py27,py34,py35,py36}-kombu{40,41,42}' --result-json /tmp/kombu.results
+    - persist_to_workspace:
+        root: /tmp
+        paths:
+        - kombu.results
+    - *save_cache_step
+
   sqlite3:
     docker:
       - *test_runner
@@ -860,6 +875,7 @@ workflows:
       - aiopg
       - redis
       - rediscluster
+      - kombu
       - sqlite3
       - msgpack
       - vertica
@@ -906,6 +922,7 @@ workflows:
             - aiopg
             - redis
             - rediscluster
+            - kombu
             - sqlite3
             - msgpack
             - vertica

ddtrace/contrib/kombu/__init__.py

@@ -0,0 +1,43 @@
+"""Instrument kombu to report AMQP messaging.
+
+``patch_all`` will not automatically patch your Kombu client to make it work, as this would conflict with the
+Celery integration. You must specifically request kombu be patched, as in the example below.
+
+Note: To permit distributed tracing for the kombu integration you must enable the tracer with priority
+sampling. Refer to the documentation here:
+http://pypi.datadoghq.com/trace/docs/advanced_usage.html#priority-sampling
+
+Without enabling distributed tracing, spans within a trace generated by the kombu integration might be dropped
+without the whole trace being dropped.
+::
+
+    from ddtrace import Pin, patch
+    import kombu
+
+    # If not patched yet, you can patch kombu specifically
+    patch(kombu=True)
+
+    # This will report a span with the default settings
+    conn = kombu.Connection("amqp://guest:[email protected]:5672//")
+    conn.connect()
+    task_queue = kombu.Queue('tasks', kombu.Exchange('tasks'), routing_key='tasks')
+    to_publish = {'hello': 'world'}
+    producer = conn.Producer()
+    producer.publish(to_publish,
+                     exchange=task_queue.exchange,
+                     routing_key=task_queue.routing_key,
+                     declare=[task_queue])
+
+    # Use a pin to specify metadata related to this client
+    Pin.override(producer, service='kombu-consumer')
+"""
+
+from ...utils.importlib import require_modules
+
+required_modules = ['kombu', 'kombu.messaging']
+
+with require_modules(required_modules) as missing_modules:
+    if not missing_modules:
+        from .patch import patch
+
+        __all__ = ['patch']

ddtrace/contrib/kombu/constants.py

@@ -0,0 +1 @@
+DEFAULT_SERVICE = 'kombu'

ddtrace/contrib/kombu/patch.py

@@ -0,0 +1,110 @@
+# 3p
+import kombu
+import wrapt
+
+# project
+from ddtrace import config
+
+from ...pin import Pin
+from ...utils.formats import get_env
+from .constants import DEFAULT_SERVICE
+from ...ext import kombu as kombux
+from ...ext import AppTypes
+from ...utils.wrappers import unwrap
+from ...propagation.http import HTTPPropagator
+from .utils import (
+    get_exchange_from_args,
+    get_body_length_from_args,
+    get_routing_key_from_args,
+    extract_conn_tags,
+    HEADER_POS
+)
+
+# kombu default settings
+config._add('kombu',{
+    'service_name': get_env('kombu', 'service_name', DEFAULT_SERVICE)
+})
+
+propagator = HTTPPropagator()
+
+
+def patch():
+    """Patch the instrumented methods
+
+    This duplicated doesn't look nice. The nicer alternative is to use an ObjectProxy on top
+    of Kombu. However, it means that any "import kombu.Connection" won't be instrumented.
+    """
+    if getattr(kombu, '_datadog_patch', False):
+        return
+    setattr(kombu, '_datadog_patch', True)
+
+    _w = wrapt.wrap_function_wrapper
+    # We wrap the _publish method because the publish method:
+    # *  defines defaults in its kwargs
+    # *  potentially overrides kwargs with values from self
+    # *  extracts/normalizes things like exchange
+    _w(kombux.TYPE, 'Producer._publish', traced_publish)
+    _w(kombux.TYPE, 'Consumer.receive', traced_receive)
+    Pin(
+        service=config.kombu['service_name'],
+        app='kombu',
+        app_type=AppTypes.worker,
+    ).onto(kombu.messaging.Producer)
+
+    Pin(
+        service=config.kombu['service_name'],
+        app='kombu',
+        app_type=AppTypes.worker,
+    ).onto(kombu.messaging.Consumer)
+
+
+def unpatch():
+    if getattr(kombu, '_datadog_patch', False):
+        setattr(kombu, '_datadog_patch', False)
+        unwrap(kombu.Producer, '_publish')
+        unwrap(kombu.Consumer, 'receive')
+
+#
+# tracing functions
+#
+
+
+def traced_receive(func, instance, args, kwargs):
+    pin = Pin.get_from(instance)
+    if not pin or not pin.enabled():
+        return func(*args, **kwargs)
+
+    # Signature only takes 2 args: (body, message)
+    message = args[1]
+    context = propagator.extract(message.headers)
+    # only need to active the new context if something was propagated
+    if context.trace_id:
+        pin.tracer.context_provider.activate(context)
+    with pin.tracer.trace(kombux.RECEIVE_NAME, service=pin.service, span_type='kombu') as s:
+        # run the command
+        exchange = message.delivery_info['exchange']
+        s.resource = exchange
+        s.set_tag(kombux.EXCHANGE, exchange)
+
+        s.set_tags(extract_conn_tags(message.channel.connection))
+        s.set_tag(kombux.ROUTING_KEY, message.delivery_info['routing_key'])
+        return func(*args, **kwargs)
+
+
+def traced_publish(func, instance, args, kwargs):
+    pin = Pin.get_from(instance)
+    if not pin or not pin.enabled():
+        return func(*args, **kwargs)
+
+    with pin.tracer.trace(kombux.PUBLISH_NAME, service=pin.service, span_type='kombu') as s:
+        exchange_name = get_exchange_from_args(args)
+        s.resource = exchange_name
+        s.set_tag(kombux.EXCHANGE, exchange_name)
+        if pin.tags:
+            s.set_tags(pin.tags)
+        s.set_tag(kombux.ROUTING_KEY, get_routing_key_from_args(args))
+        s.set_tags(extract_conn_tags(instance.channel.connection))
+        s.set_metric(kombux.BODY_LEN, get_body_length_from_args(args))
+        # run the command
+        propagator.inject(s.context, args[HEADER_POS])
+        return func(*args, **kwargs)

ddtrace/contrib/kombu/utils.py

@@ -0,0 +1,47 @@
+"""
+Some utils used by the dogtrace kombu integration
+"""
+from ...ext import kombu as kombux, net
+
+PUBLISH_BODY_IDX = 0
+PUBLISH_ROUTING_KEY = 6
+PUBLISH_EXCHANGE_IDX = 9
+
+HEADER_POS = 4
+
+
+def extract_conn_tags(connection):
+    """ Transform kombu conn info into dogtrace metas """
+    try:
+        host, port = connection.host.split(':')
+        return {
+            net.TARGET_HOST: host,
+            net.TARGET_PORT: port,
+            kombux.VHOST: connection.virtual_host,
+        }
+    except AttributeError:
+        # Unlikely that we don't have .host or .virtual_host but let's not die over it
+        return {}
+
+
+def get_exchange_from_args(args):
+    """Extract the exchange
+
+    The publish method extracts the name and hands that off to _publish (what we patch)
+    """
+
+    return args[PUBLISH_EXCHANGE_IDX]
+
+
+def get_routing_key_from_args(args):
+    """Extract the routing key"""
+
+    name = args[PUBLISH_ROUTING_KEY]
+    return name
+
+
+def get_body_length_from_args(args):
+    """Extract the length of the body"""
+
+    length = len(args[PUBLISH_BODY_IDX])
+    return length

ddtrace/ext/kombu.py

@@ -0,0 +1,13 @@
+# type of the spans
+TYPE = 'kombu'
+
+# net extension
+VHOST = 'out.vhost'
+
+# standard tags
+EXCHANGE = 'kombu.exchange'
+BODY_LEN = 'kombu.body_length'
+ROUTING_KEY = 'kombu.routing_key'
+
+PUBLISH_NAME = 'kombu.publish'
+RECEIVE_NAME = 'kombu.receive'

ddtrace/monkey.py

@@ -46,6 +46,7 @@
     'vertica': True,
     'jinja2': True,
     'flask': True,
+    'kombu': False,
 
     # Ignore some web framework integrations that might be configured explicitly in code
     "django": False,

docker-compose.yml

@@ -62,6 +62,10 @@ services:
             - "127.0.0.1:5003:5003"
             - "127.0.0.1:5004:5004"
             - "127.0.0.1:5005:5005"
+    rabbitmq:
+        image: rabbitmq:3.7-alpine
+        ports:
+            - "127.0.0.1:5672:5672"
     ddagent:
         image: datadog/docker-dd-agent
         environment:

tests/contrib/config.py

@@ -18,19 +18,19 @@
 }
 
 POSTGRES_CONFIG = {
-    'host' : 'localhost',
+    'host': 'localhost',
     'port': int(os.getenv("TEST_POSTGRES_PORT", 5432)),
-    'user' : os.getenv("TEST_POSTGRES_USER", "postgres"),
-    'password' : os.getenv("TEST_POSTGRES_PASSWORD", "postgres"),
-    'dbname' : os.getenv("TEST_POSTGRES_DB", "postgres"),
+    'user': os.getenv("TEST_POSTGRES_USER", "postgres"),
+    'password': os.getenv("TEST_POSTGRES_PASSWORD", "postgres"),
+    'dbname': os.getenv("TEST_POSTGRES_DB", "postgres"),
 }
 
 MYSQL_CONFIG = {
-    'host' : '127.0.0.1',
-    'port' : int(os.getenv("TEST_MYSQL_PORT", 3306)),
-    'user' : os.getenv("TEST_MYSQL_USER", 'test'),
-    'password' : os.getenv("TEST_MYSQL_PASSWORD", 'test'),
-    'database' : os.getenv("TEST_MYSQL_DATABASE", 'test'),
+    'host': '127.0.0.1',
+    'port': int(os.getenv("TEST_MYSQL_PORT", 3306)),
+    'user': os.getenv("TEST_MYSQL_USER", 'test'),
+    'password': os.getenv("TEST_MYSQL_PASSWORD", 'test'),
+    'database': os.getenv("TEST_MYSQL_DATABASE", 'test'),
 }
 
 REDIS_CONFIG = {
@@ -47,7 +47,7 @@
 }
 
 MEMCACHED_CONFIG = {
-    'host' : os.getenv('TEST_MEMCACHED_HOST', '127.0.0.1'),
+    'host': os.getenv('TEST_MEMCACHED_HOST', '127.0.0.1'),
     'port': int(os.getenv("TEST_MEMCACHED_PORT", 11211)),
 }
 
@@ -58,3 +58,10 @@
     'password': os.getenv('TEST_VERTICA_PASSWORD', 'abc123'),
     'database': os.getenv('TEST_VERTICA_DATABASE', 'docker'),
 }
+
+RABBITMQ_CONFIG = {
+    'host': os.getenv('TEST_RABBITMQ_HOST', '127.0.0.1'),
+    'user': os.getenv('TEST_RABBITMQ_USER', 'guest'),
+    'password': os.getenv('TEST_RABBITMQ_PASSWORD', 'guest'),
+    'port': int(os.getenv("TEST_RABBITMQ_PORT", 5672)),
+}