Repository: census-instrumentation/opencensus-python
Pull request #12 (2017-08-05 00:54:39)

Add interceptors for grpc tracing in client side and server side
For #3 

duggelz: Approved, and now we just wait for GRPC work.

Patch:
diff --git a/opencensus/trace/attributes_helper.py b/opencensus/trace/attributes_helper.py
index 3cb5fdf..4f74243 100644
--- a/opencensus/trace/attributes_helper.py
+++ b/opencensus/trace/attributes_helper.py
@@ -33,3 +33,9 @@ COMMON_ATTRIBUTES = {
     'STACKTRACE': '/stacktrace',
     'TID': '/tid',
 }
+
+
+GRPC_ATTRIBUTES = {
+    'GRPC_HOST_PORT': '/grpc/host_port',
+    'GRPC_METHOD': '/grpc/method',
+}
diff --git a/opencensus/trace/ext/grpc/__init__.py b/opencensus/trace/ext/grpc/__init__.py
new file mode 100644
index 00000000..0e36c6f
--- /dev/null
+++ b/opencensus/trace/ext/grpc/__init__.py
@@ -0,0 +1,20 @@
+# Copyright 2017, OpenCensus Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+UNARY_UNARY = 'unary_unary'
+UNARY_STREAM = 'unary_stream'
+STREAM_UNARY = 'stream_unary'
+STREAM_STREAM = 'stream_stream'
+
+GRPC_TRACE_KEY = 'grpc-trace-bin'
diff --git a/opencensus/trace/ext/grpc/client_interceptor.py b/opencensus/trace/ext/grpc/client_interceptor.py
new file mode 100644
index 00000000..ed2636f
--- /dev/null
+++ b/opencensus/trace/ext/grpc/client_interceptor.py
@@ -0,0 +1,186 @@
+# Copyright 2017, OpenCensus Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import collections
+import logging
+import grpc
+import six
+
+from opencensus.trace import execution_context
+from opencensus.trace import attributes_helper
+from opencensus.trace.ext import grpc as oc_grpc
+from opencensus.trace.propagation import binary_format
+
+log = logging.getLogger(__name__)
+
+ATTRIBUTE_COMPONENT = 'COMPONENT'
+ATTRIBUTE_ERROR_NAME = 'ERROR_NAME'
+ATTRIBUTE_ERROR_MESSAGE = 'ERROR_MESSAGE'
+GRPC_HOST_PORT = 'GRPC_HOST_PORT'
+GRPC_METHOD = 'GRPC_METHOD'
+
+TIMEOUT = 3
+
+
+class _ClientCallDetails(
+    collections.namedtuple(
+        '_ClientCallDetails',
+        ('method', 'timeout', 'metadata', 'credentials')),
+        grpc.ClientCallDetails):
+    pass
+
+
+class OpenCensusClientInterceptor(grpc.UnaryUnaryClientInterceptor,
+                                  grpc.UnaryStreamClientInterceptor,
+                                  grpc.StreamUnaryClientInterceptor,
+                                  grpc.StreamStreamClientInterceptor):
+
+    def __init__(self, tracer=None, host_port=None):
+        if tracer is None:
+            tracer = execution_context.get_opencensus_tracer()
+
+        self._tracer = tracer
+        self.host_port = host_port
+        self._propagator = binary_format.BinaryFormatPropagator()
+
+    def _start_client_span(self, method, grpc_type):
+        span = self._tracer.start_span(
+            name='[gRPC_client][{}]{}'.format(grpc_type, str(method)))
+
+        # Add the component grpc to span attribute
+        self._tracer.add_attribute_to_current_span(
+            attribute_key=attributes_helper.COMMON_ATTRIBUTES.get(
+                ATTRIBUTE_COMPONENT),
+            attribute_value='grpc')
+
+        # Add the host:port info to span attribute
+        self._tracer.add_attribute_to_current_span(
+            attribute_key=attributes_helper.GRPC_ATTRIBUTES.get(
+                GRPC_HOST_PORT),
+            attribute_value=self.host_port)
+
+        # Add the method to span attribute
+        self._tracer.add_attribute_to_current_span(
+            attribute_key=attributes_helper.GRPC_ATTRIBUTES.get(GRPC_METHOD),
+            attribute_value=str(method))
+
+        return span
+
+    def _end_span_between_context(self, current_span):
+        execution_context.set_current_span(current_span)
+        self._tracer.end_span()
+
+    def _intercept_call(
+            self, client_call_details, request_iterator, grpc_type):
+        metadata = ()
+        if client_call_details.metadata is not None:
+            metadata = client_call_details.metadata
+
+        # Start a span
+        current_span = self._start_client_span(
+            client_call_details.method,
+            grpc_type)
+
+        span_context = self._tracer.span_context
+        header = self._propagator.to_header(span_context)
+        grpc_trace_metadata = {
+            oc_grpc.GRPC_TRACE_KEY: header,
+        }
+        metadata = metadata + tuple(six.iteritems(grpc_trace_metadata))
+
+        client_call_details = _ClientCallDetails(
+            client_call_details.method,
+            client_call_details.timeout,
+            metadata,
+            client_call_details.credentials)
+
+        return client_call_details, request_iterator, current_span
+
+    def _callback(self, current_span):
+        def callback(future_response):
+            execution_context.set_current_span(current_span)
+            self._trace_future_exception(future_response)
+            self._tracer.end_span()
+
+        return callback
+
+    def _trace_future_exception(self, response):
+        # Trace the exception for a grpc.Future if any
+        exception = response.exception(timeout=TIMEOUT)
+
+        if exception is not None:
+            exception = str(exception)
+
+        self._tracer.add_attribute_to_current_span(
+            attribute_key=attributes_helper.COMMON_ATTRIBUTES.get(
+                ATTRIBUTE_ERROR_MESSAGE),
+            attribute_value=exception)
+
+    def intercept_unary_unary(
+            self, continuation, client_call_details, request):
+        new_details, new_request, current_span = self._intercept_call(
+            client_call_details=client_call_details,
+            request_iterator=iter((request,)),
+            grpc_type=oc_grpc.UNARY_UNARY)
+
+        response = continuation(
+            new_details,
+            next(new_request))
+
+        response.add_done_callback(self._callback(current_span))
+
+        return response
+
+    def intercept_unary_stream(self, continuation, client_call_details,
+                               request):
+        new_details, new_request_iterator, current_span = self._intercept_call(
+            client_call_details=client_call_details,
+            request_iterator=iter((request,)),
+            grpc_type=oc_grpc.UNARY_STREAM)
+
+        response_it = continuation(
+            new_details,
+            next(new_request_iterator))
+        self._end_span_between_context(current_span)
+
+        return response_it
+
+    def intercept_stream_unary(self, continuation, client_call_details,
+                               request_iterator):
+        new_details, new_request_iterator, current_span = self._intercept_call(
+            client_call_details=client_call_details,
+            request_iterator=request_iterator,
+            grpc_type=oc_grpc.STREAM_UNARY)
+
+        response = continuation(
+            new_details,
+            new_request_iterator)
+
+        response.add_done_callback(self._callback(current_span))
+
+        return response
+
+    def intercept_stream_stream(self, continuation, client_call_details,
+                                request_iterator):
+        new_details, new_request_iterator, current_span = self._intercept_call(
+            client_call_details=client_call_details,
+            request_iterator=request_iterator,
+            grpc_type=oc_grpc.STREAM_STREAM)
+
+        response_it = continuation(
+            new_details,
+            new_request_iterator)
+        self._end_span_between_context(current_span)
+
+        return response_it
diff --git a/opencensus/trace/ext/grpc/server_interceptor.py b/opencensus/trace/ext/grpc/server_interceptor.py
new file mode 100644
index 00000000..86e633e
--- /dev/null
+++ b/opencensus/trace/ext/grpc/server_interceptor.py
@@ -0,0 +1,75 @@
+# Copyright 2017, OpenCensus Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import grpc
+import logging
+
+from opencensus.trace import attributes_helper
+from opencensus.trace import tracer as tracer_module
+from opencensus.trace.ext import grpc as oc_grpc
+from opencensus.trace.propagation import binary_format
+
+ATTRIBUTE_COMPONENT = 'COMPONENT'
+ATTRIBUTE_ERROR_NAME = 'ERROR_NAME'
+ATTRIBUTE_ERROR_MESSAGE = 'ERROR_MESSAGE'
+
+
+class OpenCensusServerInterceptor(grpc.ServerInterceptor):
+
+    def __init__(self, sampler=None, exporter=None):
+        self.sampler = sampler
+        self.exporter = exporter
+
+    def _start_server_span(self, tracer):
+        span = tracer.start_span(name='grpc_server')
+        tracer.add_attribute_to_current_span(
+            attribute_key=attributes_helper.COMMON_ATTRIBUTES.get(
+                ATTRIBUTE_COMPONENT),
+            attribute_value='grpc')
+
+        return span
+
+    def intercept_handler(self, continuation, handler_call_details):
+        metadata = handler_call_details.invocation_metadata
+        span_context = None
+
+        if metadata is not None:
+            propagator = binary_format.BinaryFormatPropagator()
+            metadata_dict = dict(metadata)
+            trace_header = metadata_dict.get(oc_grpc.GRPC_TRACE_KEY)
+
+            span_context = propagator.from_header(trace_header)
+
+        tracer = tracer_module.Tracer(span_context=span_context,
+                                      sampler=self.sampler,
+                                      exporter=self.exporter)
+
+        with self._start_server_span(tracer):
+            response = None
+
+            try:
+                response = continuation(handler_call_details)
+            except Exception as e:  # pragma: NO COVER
+                logging.error(e)
+                tracer.add_attribute_to_current_span(
+                    attributes_helper.COMMON_ATTRIBUTES.get(
+                        ATTRIBUTE_ERROR_MESSAGE),
+                    str(e))
+                tracer.end_span()
+                raise
+
+        return response
+
+    def intercept_service(self, continuation, handler_call_details):
+        return self.intercept_handler(continuation, handler_call_details)
diff --git a/opencensus/trace/propagation/binary_format.py b/opencensus/trace/propagation/binary_format.py
index cce16f6..b73492b 100644
--- a/opencensus/trace/propagation/binary_format.py
+++ b/opencensus/trace/propagation/binary_format.py
@@ -104,6 +104,10 @@ class BinaryFormatPropagator(object):
         :rtype: :class:`~opencensus.trace.span_context.SpanContext`
         :returns: SpanContext generated from the trace context header.
         """
+        # If no binary provided, generate a new SpanContext
+        if binary is None:
+            return SpanContext(from_header=False)
+
         # If cannot parse, return a new SpanContext and ignore the context
         # from binary.
         try:


Test patch:
diff --git a/requirements-test.txt b/requirements-test.txt
index 8de2c90..654441f 100644
--- a/requirements-test.txt
+++ b/requirements-test.txt
@@ -1,6 +1,7 @@
 Django==1.11.7
 Flask==0.12.2
 google-cloud-trace==0.17.0
+grpcio==1.8.0rc3
 mock==2.0.0
 mysql-connector==2.1.6
 psycopg2==2.7.3.1
diff --git a/tests/unit/trace/ext/grpc/test_client_interceptor.py b/tests/unit/trace/ext/grpc/test_client_interceptor.py
new file mode 100644
index 00000000..cf1c7a0
--- /dev/null
+++ b/tests/unit/trace/ext/grpc/test_client_interceptor.py
@@ -0,0 +1,191 @@
+# Copyright 2017, OpenCensus Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import unittest
+
+import mock
+
+from opencensus.trace.ext.grpc import client_interceptor
+
+
+class TestOpenCensusClientInterceptor(unittest.TestCase):
+    def test_constructor_default(self):
+        from opencensus.trace.tracers import noop_tracer
+        from opencensus.trace.propagation import binary_format
+
+        interceptor = client_interceptor.OpenCensusClientInterceptor()
+
+        self.assertTrue(isinstance(
+            interceptor._tracer, noop_tracer.NoopTracer))
+        self.assertIsNone(interceptor.host_port)
+        self.assertTrue(isinstance(
+            interceptor._propagator, binary_format.BinaryFormatPropagator))
+
+    def test_constructor_explicit(self):
+        from opencensus.trace.propagation import binary_format
+
+        tracer = mock.Mock()
+        host_port = 'localhost:50051'
+        interceptor = client_interceptor.OpenCensusClientInterceptor(
+            tracer=tracer, host_port=host_port)
+
+        self.assertEqual(interceptor._tracer, tracer)
+        self.assertEqual(interceptor.host_port, host_port)
+        self.assertTrue(isinstance(
+            interceptor._propagator, binary_format.BinaryFormatPropagator))
+
+    def test__start_client_span(self):
+        tracer = mock.Mock()
+        interceptor = client_interceptor.OpenCensusClientInterceptor(
+            tracer=tracer, host_port='test')
+        interceptor._start_client_span('test_method', 'unary_unary')
+
+        self.assertTrue(tracer.start_span.called)
+        self.assertTrue(tracer.add_attribute_to_current_span.called)
+
+    def test__end_span_between_context(self):
+        from opencensus.trace import execution_context
+
+        current_span = mock.Mock()
+        tracer = mock.Mock()
+        interceptor = client_interceptor.OpenCensusClientInterceptor(
+            tracer=tracer, host_port='test')
+        interceptor._end_span_between_context(current_span)
+
+        span_in_context = execution_context.get_current_span()
+
+        self.assertEqual(span_in_context, current_span)
+        self.assertTrue(tracer.end_span.called)
+
+    def test__intercept_call_metadata_none(self):
+        tracer = mock.Mock()
+        tracer.span_context = mock.Mock()
+        test_header = 'test header'
+        mock_propagator = mock.Mock()
+        mock_propagator.to_header.return_value = test_header
+
+        interceptor = client_interceptor.OpenCensusClientInterceptor(
+            tracer=tracer, host_port='test')
+        interceptor._propagator = mock_propagator
+        mock_client_call_details = mock.Mock()
+        mock_client_call_details.metadata = None
+
+        client_call_details, request_iterator, current_span = interceptor._intercept_call(
+            mock_client_call_details,
+            mock.Mock(),
+            'unary_unary')
+
+        expected_metadata = (('grpc-trace-bin', test_header),)
+
+        self.assertEqual(expected_metadata, client_call_details.metadata)
+
+    def test__intercept_call(self):
+        tracer = mock.Mock()
+        tracer.span_context = mock.Mock()
+        test_header = 'test header'
+        mock_propagator = mock.Mock()
+        mock_propagator.to_header.return_value = test_header
+
+        interceptor = client_interceptor.OpenCensusClientInterceptor(
+            tracer=tracer, host_port='test')
+        interceptor._propagator = mock_propagator
+        mock_client_call_details = mock.Mock()
+        mock_client_call_details.metadata = (('test_key', 'test_value'),)
+
+        client_call_details, request_iterator, current_span = interceptor._intercept_call(
+            mock_client_call_details,
+            mock.Mock(),
+            'unary_unary')
+
+        expected_metadata = (('test_key', 'test_value'), ('grpc-trace-bin', test_header),)
+
+        self.assertEqual(expected_metadata, client_call_details.metadata)
+
+    def test__callback(self):
+        current_span = mock.Mock()
+        tracer = MockTracer(current_span)
+        interceptor = client_interceptor.OpenCensusClientInterceptor(
+            tracer=tracer, host_port='test')
+        current_span.attributes = {}
+        callback = interceptor._callback(current_span)
+        response = mock.Mock()
+        response.exception.return_value = 'test_exception'
+        callback(response)
+        expected_attributes = {
+            '/error/message': 'test_exception',
+        }
+
+        self.assertEqual(current_span.attributes, expected_attributes)
+
+    def test__callback_no_exception(self):
+        current_span = mock.Mock()
+        tracer = MockTracer(current_span)
+        interceptor = client_interceptor.OpenCensusClientInterceptor(
+            tracer=tracer, host_port='test')
+        current_span.attributes = {}
+        callback = interceptor._callback(current_span)
+        response = mock.Mock()
+        response.exception.return_value = None
+        callback(response)
+        expected_attributes = {'/error/message': None}
+
+        self.assertEqual(current_span.attributes, expected_attributes)
+
+    def _unary_helper(self):
+        continuation = mock.Mock()
+        mock_response = mock.Mock()
+        continuation.return_value = mock_response
+        interceptor = client_interceptor.OpenCensusClientInterceptor()
+        interceptor._intercept_call = mock.Mock(return_value=(None, iter([mock.Mock()]), None))
+        return interceptor, continuation, mock_response
+
+    def _stream_helper(self):
+        continuation = mock.Mock()
+        mock_response = mock.Mock()
+        continuation.return_value = mock_response
+        mock_tracer = mock.Mock()
+        interceptor = client_interceptor.OpenCensusClientInterceptor(tracer=mock_tracer)
+        interceptor._intercept_call = mock.Mock(return_value=(None, iter([mock.Mock()]), None))
+        return interceptor, continuation, mock_tracer
+
+    def test_intercept_unary_unary(self):
+        interceptor, continuation, mock_response = self._unary_helper()
+        interceptor.intercept_unary_unary(continuation, mock.Mock(), [])
+        self.assertTrue(mock_response.add_done_callback.called)
+
+    def test_intercept_unary_stream(self):
+        interceptor, continuation, mock_tracer = self._stream_helper()
+        interceptor.intercept_unary_stream(continuation, mock.Mock(), [])
+        self.assertTrue(mock_tracer.end_span.called)
+
+    def test_intercept_stream_unary(self):
+        interceptor, continuation, mock_response = self._unary_helper()
+        interceptor.intercept_stream_unary(continuation, mock.Mock(), [])
+        self.assertTrue(mock_response.add_done_callback.called)
+
+    def test_intercept_stream_stream(self):
+        interceptor, continuation, mock_tracer = self._stream_helper()
+        interceptor.intercept_stream_stream(continuation, mock.Mock(), [])
+        self.assertTrue(mock_tracer.end_span.called)
+
+
+class MockTracer(object):
+    def __init__(self, current_span):
+        self.current_span = current_span
+
+    def add_attribute_to_current_span(self, attribute_key, attribute_value):
+        self.current_span.attributes[attribute_key] = attribute_value
+
+    def end_span(self):
+        return
diff --git a/tests/unit/trace/ext/grpc/test_server_interceptor.py b/tests/unit/trace/ext/grpc/test_server_interceptor.py
new file mode 100644
index 00000000..dee3aaa
--- /dev/null
+++ b/tests/unit/trace/ext/grpc/test_server_interceptor.py
@@ -0,0 +1,102 @@
+# Copyright 2017, OpenCensus Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import unittest
+
+import mock
+
+from opencensus.trace.ext.grpc import server_interceptor
+from opencensus.trace import execution_context
+
+
+class TestOpenCensusServerInterceptor(unittest.TestCase):
+    def test_constructor(self):
+        sampler = mock.Mock()
+        exporter = mock.Mock()
+        interceptor = server_interceptor.OpenCensusServerInterceptor(
+            sampler=sampler, exporter=exporter)
+        self.assertEqual(interceptor.sampler, sampler)
+        self.assertEqual(interceptor.exporter, exporter)
+
+    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)
+        mock_details = mock.Mock()
+        mock_details.invocation_metadata = None
+        interceptor = server_interceptor.OpenCensusServerInterceptor(
+            None, None)
+
+        with patch:
+            interceptor.intercept_handler(mock.Mock(), mock_details)
+
+        expected_attributes = {
+            '/component': 'grpc',
+        }
+
+        self.assertEqual(
+            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)
+        mock_details = mock.Mock()
+        mock_details.invocation_metadata = (('test_key', b'test_value'),)
+        interceptor = server_interceptor.OpenCensusServerInterceptor(
+            None, None)
+
+        with patch:
+            interceptor.intercept_handler(mock.Mock(), mock_details)
+
+        expected_attributes = {
+            '/component': 'grpc',
+        }
+
+        self.assertEqual(
+            execution_context.get_opencensus_tracer().current_span.attributes,
+            expected_attributes)
+
+    def test_intercept_service(self):
+        interceptor = server_interceptor.OpenCensusServerInterceptor(
+            None, None)
+        mock_handler = mock.Mock()
+        interceptor.intercept_handler = mock_handler
+        interceptor.intercept_service(None, None)
+        self.assertTrue(mock_handler.called)
+
+
+class MockTracer(object):
+    def __init__(self, *args, **kwargs):
+        self.current_span = mock.Mock()
+        self.current_span.attributes = {}
+        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
+
+    def end_span(self):
+        return
+
+    def add_attribute_to_current_span(self, attribute_key, attribute_value):
+        self.current_span.attributes[attribute_key] = attribute_value
diff --git a/tests/unit/trace/propagation/test_binary_format.py b/tests/unit/trace/propagation/test_binary_format.py
index e48ad6e..769639a 100644
--- a/tests/unit/trace/propagation/test_binary_format.py
+++ b/tests/unit/trace/propagation/test_binary_format.py
@@ -30,6 +30,14 @@ class TestBinaryFormat(unittest.TestCase):
 
         self.assertFalse(span_context.from_header)
 
+    def test_from_header_none(self):
+        binary_header = None
+
+        propagator = binary_format.BinaryFormatPropagator()
+        span_context = propagator.from_header(binary_header)
+
+        self.assertFalse(span_context.from_header)
+
     def test_from_header(self):
         binary_header = b'\x00\x00\xa0\xb7,\xa1\\\x1aK\xd1\x89b\xd0' \
                         b'\xacY\xdc\x90\xb9\x01g)U\xf6\xf5\x01\x12' \