Attach stack traces and status to the current span on gRPC server failures (#140)

Author: wkiser

opencensus/trace/ext/grpc/server_interceptor.py

@@ -13,12 +13,16 @@
 # limitations under the License.
 import collections
 import logging
+import sys
 
+from google.rpc import code_pb2
 import grpc
 
 from opencensus.trace import attributes_helper
-from opencensus.trace import tracer as tracer_module
 from opencensus.trace import execution_context
+from opencensus.trace import stack_trace as stack_trace
+from opencensus.trace import status
+from opencensus.trace import tracer as tracer_module
 from opencensus.trace.ext import grpc as oc_grpc
 from opencensus.trace.propagation import binary_format
 
@@ -29,8 +33,10 @@
 RpcRequestInfo = collections.namedtuple(
     'RPCRequestInfo', ('request', 'context')
 )
+# exc_info is the three tuple defined at:
+# https://docs.python.org/3/library/sys.html#sys.exc_info
 RpcResponseInfo = collections.namedtuple(
-    'RPCCallbackInfo', ('request', 'context', 'response', 'exc')
+    'RPCCallbackInfo', ('request', 'context', 'response', 'exc_info')
 )
 
 
@@ -60,19 +66,21 @@ def proxy(self, prop_name):
         def _wrapper(request, context, *args, **kwargs):
             for callback in self._pre_handler_callbacks:
                 callback(RpcRequestInfo(request, context))
-            exc = None
+            exc_info = (None, None, None)
             response = None
             try:
                 response = getattr(
                     self.handler, prop_name
                 )(request, context, *args, **kwargs)
             except Exception as e:
-                logging.error(e)
-                exc = e
+                logging.exception(e)
+                exc_info = sys.exc_info()
                 raise
             finally:
                 for callback in self._post_handler_callbacks:
-                    callback(RpcResponseInfo(request, context, response, exc))
+                    callback(RpcResponseInfo(
+                        request, context, response, exc_info)
+                    )
             return response
 
         return _wrapper
@@ -116,11 +124,22 @@ def _start_server_span(self, rpc_request_info):
 
     def _end_server_span(self, rpc_response_info):
         tracer = execution_context.get_opencensus_tracer()
-        if rpc_response_info.exc is not None:
-            tracer.add_attribute_to_current_span(
+        exc_type, exc_value, tb = rpc_response_info.exc_info
+        if exc_type is not None:
+            current_span = tracer.current_span()
+            current_span.add_attribute(
                 attributes_helper.COMMON_ATTRIBUTES.get(
                     ATTRIBUTE_ERROR_MESSAGE),
-                str(rpc_response_info.exc))
+                str(exc_value)
+            )
+            current_span.stack_trace = stack_trace.StackTrace.from_traceback(
+                tb
+            )
+            current_span.status = status.Status(
+                code=code_pb2.UNKNOWN,
+                message=str(exc_value)
+            )
+
         tracer.end_span()
 
     def intercept_handler(self, continuation, handler_call_details):

tests/unit/trace/ext/grpc/test_server_interceptor.py

@@ -15,8 +15,10 @@
 import unittest
 
 import mock
+from google.rpc import code_pb2
 
 from opencensus.trace import execution_context
+from opencensus.trace import span as span_module
 from opencensus.trace.ext.grpc import server_interceptor
 
 
@@ -37,8 +39,6 @@ def test_rpc_handler_wrapper(self):
         self.assertEqual(wrapper.response_streaming, False)
 
     def test_intercept_handler_no_metadata(self):
-        current_span = mock.Mock()
-        mock_tracer = MockTracer(None, None, None)
         patch = mock.patch(
             'opencensus.trace.ext.grpc.server_interceptor.tracer_module.Tracer',
             MockTracer)
@@ -57,12 +57,10 @@ def test_intercept_handler_no_metadata(self):
         }
 
         self.assertEqual(
-            execution_context.get_opencensus_tracer().current_span.attributes,
+            execution_context.get_opencensus_tracer().current_span().attributes,
             expected_attributes)
 
     def test_intercept_handler(self):
-        current_span = mock.Mock()
-        mock_tracer = MockTracer(None, None, None)
         patch = mock.patch(
             'opencensus.trace.ext.grpc.server_interceptor.tracer_module.Tracer',
             MockTracer)
@@ -83,7 +81,7 @@ def test_intercept_handler(self):
         }
 
         self.assertEqual(
-            execution_context.get_opencensus_tracer().current_span.attributes,
+            execution_context.get_opencensus_tracer().current_span().attributes,
             expected_attributes)
 
     def test_intercept_service(self):
@@ -95,8 +93,6 @@ def test_intercept_service(self):
         self.assertTrue(mock_handler.called)
 
     def test_intercept_handler_exception(self):
-        current_span = mock.Mock()
-        mock_tracer = MockTracer(None, None, None)
         patch = mock.patch(
             'opencensus.trace.ext.grpc.server_interceptor.tracer_module.Tracer',
             MockTracer)
@@ -120,26 +116,36 @@ def test_intercept_handler_exception(self):
             '/error/message': 'Test'
         }
 
+        current_span = execution_context.get_opencensus_tracer().current_span()
         self.assertEqual(
-            execution_context.get_opencensus_tracer().current_span.attributes,
+            execution_context.get_opencensus_tracer().current_span().attributes,
             expected_attributes)
 
+        # check that the stack trace is attached to the current span
+        self.assertIsNotNone(current_span.stack_trace)
+        self.assertIsNotNone(current_span.stack_trace.stack_trace_hash_id)
+        self.assertNotEqual(current_span.stack_trace.stack_frames, [])
+
+        # check that the status obj is attached to the current span
+        self.assertIsNotNone(current_span.status)
+        self.assertEqual(current_span.status.code, code_pb2.UNKNOWN)
+        self.assertEqual(current_span.status.message, 'Test')
+
 
 class MockTracer(object):
     def __init__(self, *args, **kwargs):
-        self.current_span = mock.Mock()
-        self.current_span.attributes = {}
+        self._current_span = span_module.Span('mock_span')
         execution_context.set_opencensus_tracer(self)
 
     def start_span(self, name):
-        self.current_span.name = name
-        span = mock.Mock()
-        span.__enter__ = mock.Mock()
-        span.__exit__ = mock.Mock()
-        return span
+        self._current_span.name = name
+        return self._current_span
 
     def end_span(self):
         return
 
     def add_attribute_to_current_span(self, attribute_key, attribute_value):
-        self.current_span.attributes[attribute_key] = attribute_value
+        self._current_span.attributes[attribute_key] = attribute_value
+
+    def current_span(self):
+        return self._current_span