[cassandra] trace execute_async operations (#333)

* [cassandra] patch execute_async instead of execute

Author: bmermet

ddtrace/contrib/cassandra/session.py

@@ -1,6 +1,8 @@
 """
 Trace queries along a session to a cassandra cluster
 """
+import sys
+import logging
 # 3p
 import cassandra.cluster
 import wrapt
@@ -9,11 +11,14 @@
 from ddtrace import Pin
 from ddtrace.compat import stringify
 from ...util import deep_getattr, deprecated
-from ...ext import net, cassandra as cassx
+from ...ext import net, cassandra as cassx, errors
 
+log = logging.getLogger(__name__)
 
 RESOURCE_MAX_LENGTH = 5000
-SERVICE = "cassandra"
+SERVICE = 'cassandra'
+CURRENT_SPAN = '_ddtrace_current_span'
+PAGE_NUMBER = '_ddtrace_page_number'
 
 # Original connect connect function
 _connect = cassandra.cluster.Cluster.connect
@@ -31,32 +36,142 @@ def traced_connect(func, instance, args, kwargs):
     session = func(*args, **kwargs)
     if not isinstance(session.execute, wrapt.FunctionWrapper):
         # FIXME[matt] this should probably be private.
-        setattr(session, 'execute', wrapt.FunctionWrapper(session.execute, traced_execute))
+        setattr(session, 'execute_async', wrapt.FunctionWrapper(session.execute_async, traced_execute_async))
     return session
 
-def traced_execute(func, instance, args, kwargs):
-    cluster = getattr(instance, 'cluster', None)
+def _close_span_on_success(result, future):
+    span = getattr(future, CURRENT_SPAN, None)
+    if not span:
+        log.debug('traced_set_final_result was not able to get the current span from the ResponseFuture')
+        return
+    try:
+        span.set_tags(_extract_result_metas(cassandra.cluster.ResultSet(future, result)))
+    except Exception as e:
+        log.debug('an exception occured while setting tags: %s', e)
+    finally:
+        span.finish()
+        delattr(future, CURRENT_SPAN)
+
+def traced_set_final_result(func, instance, args, kwargs):
+    result = args[0]
+    _close_span_on_success(result, instance)
+    return func(*args, **kwargs)
+
+def _close_span_on_error(exc, future):
+    span = getattr(future, CURRENT_SPAN, None)
+    if not span:
+        log.debug('traced_set_final_exception was not able to get the current span from the ResponseFuture')
+        return
+    try:
+        # handling the exception manually because we
+        # don't have an ongoing exception here
+        span.error = 1
+        span.set_tag(errors.ERROR_MSG, exc.args[0])
+        span.set_tag(errors.ERROR_TYPE, exc.__class__.__name__)
+    except Exception as e:
+        log.debug('traced_set_final_exception was not able to set the error, failed with error: %s', e)
+    finally:
+        span.finish()
+        delattr(future, CURRENT_SPAN)
+
+def traced_set_final_exception(func, instance, args, kwargs):
+    exc = args[0]
+    _close_span_on_error(exc, instance)
+    return func(*args, **kwargs)
+
+def traced_start_fetching_next_page(func, instance, args, kwargs):
+    has_more_pages = getattr(instance, 'has_more_pages', True)
+    if not has_more_pages:
+        return func(*args, **kwargs)
+    session = getattr(instance, 'session', None)
+    cluster = getattr(session, 'cluster', None)
     pin = Pin.get_from(cluster)
     if not pin or not pin.enabled():
         return func(*args, **kwargs)
 
-    service = pin.service
-    tracer = pin.tracer
+    # In case the current span is not finished we make sure to finish it
+    old_span = getattr(instance, CURRENT_SPAN, None)
+    if old_span:
+        log.debug('previous span was not finished before fetching next page')
+        old_span.finish()
 
-    query = kwargs.get("kwargs") or args[0]
+    query = getattr(instance, 'query', None)
 
-    with tracer.trace("cassandra.query", service=service, span_type=cassx.TYPE) as span:
-        _sanitize_query(span, query)
-        span.set_tags(_extract_session_metas(instance))     # FIXME[matt] do once?
-        span.set_tags(_extract_cluster_metas(cluster))
-        result = None
-        try:
-            result = func(*args, **kwargs)
-            return result
-        finally:
-            if result:
-                span.set_tags(_extract_result_metas(result))
+    span = _start_span_and_set_tags(pin, query, session, cluster)
 
+    page_number = getattr(instance, PAGE_NUMBER, 1) + 1
+    setattr(instance, PAGE_NUMBER, page_number)
+    setattr(instance, CURRENT_SPAN, span)
+    try:
+        return func(*args, **kwargs)
+    except:
+        with span:
+            span.set_exc_info(*sys.exc_info())
+        raise
+
+def traced_execute_async(func, instance, args, kwargs):
+    cluster = getattr(instance, 'cluster', None)
+    pin = Pin.get_from(cluster)
+    if not pin or not pin.enabled():
+        return func(*args, **kwargs)
+
+    query = kwargs.get("query") or args[0]
+
+    span = _start_span_and_set_tags(pin, query, instance, cluster)
+
+    try:
+        result = func(*args, **kwargs)
+        setattr(result, CURRENT_SPAN, span)
+        setattr(result, PAGE_NUMBER, 1)
+        setattr(
+            result,
+            '_set_final_result',
+            wrapt.FunctionWrapper(
+                result._set_final_result,
+                traced_set_final_result
+            )
+        )
+        setattr(
+            result,
+            '_set_final_exception',
+            wrapt.FunctionWrapper(
+                result._set_final_exception,
+                traced_set_final_exception
+            )
+        )
+        setattr(
+            result,
+            'start_fetching_next_page',
+            wrapt.FunctionWrapper(
+                result.start_fetching_next_page,
+                traced_start_fetching_next_page
+            )
+        )
+        # Since we cannot be sure that the previous methods were overwritten
+        # before the call ended, we add callbacks that will be run
+        # synchronously if the call already returned and we remove them right
+        # after.
+        result.add_callbacks(
+            _close_span_on_success,
+            _close_span_on_error,
+            callback_args=(result,),
+            errback_args=(result,)
+        )
+        result.clear_callbacks()
+        return result
+    except:
+        with span:
+            span.set_exc_info(*sys.exc_info())
+        raise
+
+def _start_span_and_set_tags(pin, query, session, cluster):
+    service = pin.service
+    tracer = pin.tracer
+    span = tracer.trace("cassandra.query", service=service, span_type=cassx.TYPE)
+    _sanitize_query(span, query)
+    span.set_tags(_extract_session_metas(session))     # FIXME[matt] do once?
+    span.set_tags(_extract_cluster_metas(cluster))
+    return span
 
 def _extract_session_metas(session):
     metas = {}
@@ -79,7 +194,7 @@ def _extract_cluster_metas(cluster):
 
 def _extract_result_metas(result):
     metas = {}
-    if not result:
+    if result is None:
         return metas
 
     future = getattr(result, "response_future", None)
@@ -100,12 +215,13 @@ def _extract_result_metas(result):
         if getattr(query, "keyspace", None):
             metas[cassx.KEYSPACE] = query.keyspace.lower()
 
-    if hasattr(result, "has_more_pages"):
-        metas[cassx.PAGINATED] = bool(result.has_more_pages)
+        page_number = getattr(future, PAGE_NUMBER, 1)
+        has_more_pages = getattr(future, "has_more_pages")
+        is_paginated = has_more_pages or page_number > 1
+        metas[cassx.PAGINATED] = is_paginated
+        if is_paginated:
+            metas[cassx.PAGE_NUMBER] = page_number
 
-    # NOTE(aaditya): this number only reflects the first page of results
-    # which could be misleading. But a true count would require iterating through
-    # all pages which is expensive
     if hasattr(result, "current_rows"):
         result_rows = result.current_rows or []
         metas[cassx.ROW_COUNT] = len(result_rows)

ddtrace/ext/cassandra.py

@@ -8,3 +8,4 @@
 CONSISTENCY_LEVEL = "cassandra.consistency_level"
 PAGINATED = "cassandra.paginated"
 ROW_COUNT = "cassandra.row_count"
+PAGE_NUMBER = "cassandra.page_number"