Add error data to span when an exception is raised in a span context (#122)

Author: wkiser

opencensus/trace/span.py

@@ -17,6 +17,8 @@
 
 from opencensus.trace import attributes
 from opencensus.trace import link as link_module
+from opencensus.trace import stack_trace
+from opencensus.trace import status
 from opencensus.trace import time_event as time_event_module
 from opencensus.trace.span_context import generate_span_id
 from opencensus.trace.tracers import base
@@ -206,6 +208,10 @@ def __enter__(self):
 
     def __exit__(self, exception_type, exception_value, traceback):
         """Finish a span."""
+        if traceback is not None:
+            self.stack_trace = stack_trace.StackTrace.from_traceback(traceback)
+        if exception_value is not None:
+            self.status = status.Status.from_exception(exception_value)
         if self.context_tracer is not None:
             self.context_tracer.end_span()
             return

opencensus/trace/stack_trace.py

@@ -12,10 +12,18 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import hashlib
+import os
 import random
+import traceback
 
 from opencensus.trace.utils import _get_truncatable_str
 
+MAX_FRAMES = 128
+
+BUILD_ID = os.environ.get('BUILD_ID', 'unknown')
+SOURCE_VERSION = os.environ.get('SOURCE_VERSION', 'unknown')
+
 
 class StackFrame(object):
     """Represents a single stack frame in a stack trace.
@@ -84,7 +92,7 @@ def format_stack_frame_json(self):
             self.original_func_name)
         stack_frame_json['file_name'] = _get_truncatable_str(self.file_name)
         stack_frame_json['line_number'] = self.line_num
-        stack_frame_json['col_number'] = self.col_num
+        stack_frame_json['column_number'] = self.col_num
         stack_frame_json['load_module'] = {
             'module': _get_truncatable_str(self.load_module),
             'build_id': _get_truncatable_str(self.build_id),
@@ -110,23 +118,57 @@ class StackTrace(object):
     def __init__(self, stack_frames=None, stack_trace_hash_id=None):
         if stack_frames is None:
             stack_frames = []
+        if len(stack_frames) > MAX_FRAMES:
+            self.dropped_frames_count = len(stack_frames) - MAX_FRAMES
+            stack_frames = stack_frames[-MAX_FRAMES:]
+        else:
+            self.dropped_frames_count = 0
 
         if stack_trace_hash_id is None:
             stack_trace_hash_id = generate_hash_id()
 
         self.stack_frames = stack_frames
         self.stack_trace_hash_id = stack_trace_hash_id
 
+    @classmethod
+    def from_traceback(cls, tb):
+        """Initializes a StackTrace from a python traceback instance"""
+        stack_trace = cls(
+            stack_trace_hash_id=generate_hash_id_from_traceback(tb)
+        )
+        # use the add_stack_frame so that json formatting is applied
+        for tb_frame_info in traceback.extract_tb(tb):
+            filename, line_num, fn_name, _ = tb_frame_info
+            stack_trace.add_stack_frame(
+                StackFrame(
+                    func_name=fn_name,
+                    original_func_name=fn_name,
+                    file_name=filename,
+                    line_num=line_num,
+                    col_num=0,  # I don't think this is available in python
+                    load_module=filename,
+                    build_id=BUILD_ID,
+                    source_version=SOURCE_VERSION
+                )
+            )
+        return stack_trace
+
     def add_stack_frame(self, stack_frame):
         """Add StackFrame to frames list."""
-        self.stack_frames.append(stack_frame.format_stack_frame_json())
+        if len(self.stack_frames) >= MAX_FRAMES:
+            self.dropped_frames_count += 1
+        else:
+            self.stack_frames.append(stack_frame.format_stack_frame_json())
 
     def format_stack_trace_json(self):
         """Convert a StackTrace object to json format."""
         stack_trace_json = {}
 
         if self.stack_frames:
-            stack_trace_json['stack_frames'] = self.stack_frames
+            stack_trace_json['stack_frames'] = {
+                'frame': self.stack_frames,
+                'dropped_frames_count': self.dropped_frames_count
+            }
 
         stack_trace_json['stack_trace_hash_id'] = self.stack_trace_hash_id
 
@@ -136,3 +178,12 @@ def format_stack_trace_json(self):
 def generate_hash_id():
     """Generate a hash id."""
     return random.getrandbits(64)
+
+
+def generate_hash_id_from_traceback(tb):
+    m = hashlib.md5()
+    for tb_line in traceback.format_tb(tb):
+        m.update(tb_line.encode('utf-8'))
+    # truncate the hash for easier compatibility with StackDriver,
+    # should still be unique enough to avoid collisions
+    return int(m.hexdigest()[:12], 16)

opencensus/trace/status.py

@@ -12,6 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from google.rpc import code_pb2
+
 
 class Status(object):
     """The Status type defines a logical error model that is suitable for
@@ -53,3 +55,10 @@ def format_status_json(self):
             status_json['details'] = self.details
 
         return status_json
+
+    @classmethod
+    def from_exception(cls, exc):
+        return cls(
+            code=code_pb2.UNKNOWN,
+            message=str(exc)
+        )

tests/unit/trace/test_span.py

@@ -16,6 +16,8 @@
 
 import mock
 
+from google.rpc import code_pb2
+
 from opencensus.trace.stack_trace import StackTrace
 from opencensus.trace.status import Status
 from opencensus.trace.time_event import TimeEvent
@@ -219,6 +221,39 @@ def test___iter__(self):
             span_iter_list,
             [child1_child1_span, child1_span, child2_span, root_span])
 
+    def test_exception_in_span(self):
+        """Make sure that an exception within a span context is
+        attached to the span"""
+        root_span = self._make_one('root_span')
+        exception_message = 'error'
+        with self.assertRaises(AssertionError):
+            with root_span:
+                raise AssertionError(exception_message)
+        stack_trace = root_span.stack_trace
+        # make sure a stack trace has been attached and populated
+        self.assertIsNotNone(stack_trace)
+        self.assertIsNotNone(stack_trace.stack_trace_hash_id)
+        self.assertEqual(len(stack_trace.stack_frames), 1)
+
+        stack_frame = stack_trace.stack_frames[0]
+        self.assertEqual(stack_frame['file_name']['value'], __file__)
+        self.assertEqual(
+            stack_frame['function_name']['value'], 'test_exception_in_span'
+        )
+        self.assertEqual(
+            stack_frame['load_module']['module']['value'], __file__
+        )
+        self.assertEqual(
+            stack_frame['original_function_name']['value'],
+            'test_exception_in_span'
+        )
+        self.assertIsNotNone(stack_frame['source_version']['value'])
+        self.assertIsNotNone(stack_frame['load_module']['build_id']['value'])
+
+        self.assertIsNotNone(root_span.status)
+        self.assertEqual(root_span.status.message, exception_message)
+        self.assertEqual(root_span.status.code, code_pb2.UNKNOWN)
+
 
 class Test_format_span_json(unittest.TestCase):
 

tests/unit/trace/test_stack_trace.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import sys
 import unittest
 
 import mock
@@ -81,7 +82,7 @@ def mock_get_truncatable_str(str):
             'original_function_name': original_func_name,
             'file_name': file_name,
             'line_number': line_num,
-            'col_number': col_num,
+            'column_number': col_num,
             'load_module': {
                 'module': load_module,
                 'build_id': build_id
@@ -116,6 +117,15 @@ def test_constructor_explicit(self):
         self.assertEqual(stack_trace.stack_frames, stack_frames)
         self.assertEqual(stack_trace.stack_trace_hash_id, hash_id)
 
+    def test_constructor_max_frames(self):
+        stack_frames = [mock.Mock()] * (stack_trace_module.MAX_FRAMES + 1)
+        stack_trace = stack_trace_module.StackTrace(stack_frames, 100)
+        self.assertEqual(stack_trace.dropped_frames_count, 1)
+        self.assertEqual(
+            len(stack_trace.stack_frames),
+            stack_trace_module.MAX_FRAMES
+        )
+
     def test_add_stack_frame(self):
         stack_trace = stack_trace_module.StackTrace()
         stack_frame = mock.Mock()
@@ -136,7 +146,10 @@ def test_format_stack_trace_json_with_stack_frame(self):
         stack_trace_json = stack_trace.format_stack_trace_json()
 
         expected_stack_trace_json = {
-            'stack_frames': stack_frame,
+            'stack_frames': {
+                'frame': stack_frame,
+                'dropped_frames_count': 0
+            },
             'stack_trace_hash_id': hash_id
         }
 
@@ -155,3 +168,50 @@ def test_format_stack_trace_json_without_stack_frame(self):
         }
 
         self.assertEqual(expected_stack_trace_json, stack_trace_json)
+
+    def test_create_from_traceback(self):
+        try:
+            raise AssertionError('something went wrong')
+        except AssertionError:
+            _, _, tb = sys.exc_info()
+
+        stack_trace = stack_trace_module.StackTrace.from_traceback(tb)
+        self.assertIsNotNone(stack_trace)
+        self.assertIsNotNone(stack_trace.stack_trace_hash_id)
+        self.assertEqual(len(stack_trace.stack_frames), 1)
+
+        stack_frame = stack_trace.stack_frames[0]
+        self.assertEqual(stack_frame['file_name']['value'], __file__)
+        self.assertEqual(
+            stack_frame['function_name']['value'], 'test_create_from_traceback'
+        )
+        self.assertEqual(
+            stack_frame['load_module']['module']['value'], __file__
+        )
+        self.assertEqual(
+            stack_frame['original_function_name']['value'],
+            'test_create_from_traceback'
+        )
+        self.assertIsNotNone(stack_frame['source_version']['value'])
+        self.assertIsNotNone(stack_frame['load_module']['build_id']['value'])
+
+    def test_dropped_frames(self):
+        """Make sure the limit of 128 frames is enforced"""
+        def recur(max_depth):
+            def _recur_helper(depth):
+                if depth >= max_depth:
+                    raise AssertionError('reached max depth')
+                _recur_helper(depth+1)
+            _recur_helper(0)
+        try:
+            recur(stack_trace_module.MAX_FRAMES)
+        except AssertionError:
+            _, _, tb = sys.exc_info()
+
+        stack_trace = stack_trace_module.StackTrace.from_traceback(tb)
+        # total frames should be MAX_FRAMES + 3 (1 for test function,
+        # 1 for recursion start, one for exception, MAX_FRAMES in helper)
+        self.assertEqual(stack_trace.dropped_frames_count, 3)
+
+
+

tests/unit/trace/test_status.py

@@ -14,7 +14,7 @@
 
 import unittest
 
-import mock
+from google.rpc import code_pb2
 
 from opencensus.trace import status as status_module
 
@@ -63,3 +63,10 @@ def test_format_status_json_without_details(self):
         }
 
         self.assertEqual(expected_status_json, status_json)
+
+    def test_create_from_exception(self):
+        message = 'test message'
+        exc = ValueError(message)
+        status = status_module.Status.from_exception(exc)
+        self.assertEqual(status.message, message)
+        self.assertEqual(status.code, code_pb2.UNKNOWN)