Update gRPC client interceptor to use tracer from execution context (#157)

Author: wkiser

examples/trace/helloworld/flask/custom.py

@@ -13,17 +13,20 @@
 # limitations under the License.
 
 import os
-import requests
-import sys
 
 import flask
+import grpc
 import mysql.connector
 import psycopg2
+import requests
 import sqlalchemy
 
-from opencensus.trace.ext.flask.flask_middleware import FlaskMiddleware
+import hello_world_pb2
+import hello_world_pb2_grpc
 from opencensus.trace import config_integration
 from opencensus.trace.exporters import stackdriver_exporter
+from opencensus.trace.ext.flask.flask_middleware import FlaskMiddleware
+from opencensus.trace.ext.grpc import client_interceptor
 from opencensus.trace.samplers import probability
 
 INTEGRATIONS = ['mysql', 'postgresql', 'sqlalchemy', 'requests']
@@ -38,14 +41,20 @@
 # PostgreSQL settings
 POSTGRES_PASSWORD = os.environ.get('SYSTEST_POSTGRES_PASSWORD')
 
+# hello_world_server location
+HELLO_WORLD_HOST_PORT = 'localhost:50051'
+
 app = flask.Flask(__name__)
 
 # Enable tracing, configure the trace params, send traces to Stackdriver Trace
-exporter = stackdriver_exporter.StackdriverExporter(project_id='yanhuili-sandbox')
+exporter = stackdriver_exporter.StackdriverExporter()
 sampler = probability.ProbabilitySampler(rate=1)
 middleware = FlaskMiddleware(app, exporter=exporter, sampler=sampler)
 config_integration.trace_integrations(INTEGRATIONS)
 
+# cache of (stub_cls, hostport) -> grpc stub
+_stub_cache = {}
+
 
 @app.route('/')
 def hello():
@@ -162,5 +171,29 @@ def sqlalchemy_postgresql_query():
         return msg, 500
 
 
+@app.route('/greet/<name>')
+def greet(name):
+    stub = _get_grpc_stub(
+        hello_world_pb2_grpc.GreeterStub, HELLO_WORLD_HOST_PORT
+    )
+    response = stub.SayHello(hello_world_pb2.HelloRequest(name=name))
+    return str(response)
+
+
+def _get_grpc_stub(stub_cls, host_port):
+    stub = _stub_cache.get((stub_cls, host_port))
+    if stub is not None:
+        return stub
+
+    channel = grpc.insecure_channel(host_port)
+    tracer_interceptor = client_interceptor.OpenCensusClientInterceptor(
+        host_port=HELLO_WORLD_HOST_PORT
+    )
+    channel = grpc.intercept_channel(channel, tracer_interceptor)
+    stub = stub_cls(channel)
+    _stub_cache[(stub_cls, host_port)] = stub
+    return stub
+
+
 if __name__ == '__main__':
     app.run(host='localhost', port=8080)

opencensus/trace/ext/grpc/client_interceptor.py

@@ -51,40 +51,44 @@ class OpenCensusClientInterceptor(grpc.UnaryUnaryClientInterceptor,
                                   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()
 
+    @property
+    def tracer(self):
+        return self._tracer or execution_context.get_opencensus_tracer()
+
     def _start_client_span(self, client_call_details):
-        span = self._tracer.start_span(
+        span = self.tracer.start_span(
             name=_get_span_name(client_call_details)
         )
 
         # Add the component grpc to span attribute
-        self._tracer.add_attribute_to_current_span(
+        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(
+        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(
+        self.tracer.add_attribute_to_current_span(
             attribute_key=attributes_helper.GRPC_ATTRIBUTES.get(GRPC_METHOD),
             attribute_value=str(client_call_details.method))
 
+        execution_context.set_opencensus_tracer(self.tracer)
+        execution_context.set_current_span(span)
+
         return span
 
     def _end_span_between_context(self, current_span):
         execution_context.set_current_span(current_span)
-        self._tracer.end_span()
+        self.tracer.end_span()
 
     def _intercept_call(
             self, client_call_details, request_iterator, grpc_type):
@@ -95,14 +99,12 @@ def _intercept_call(
         # Start a span
         current_span = self._start_client_span(client_call_details)
 
-        span_context = self._tracer.span_context
+        span_context = current_span.context_tracer.span_context
         header = self._propagator.to_header(span_context)
         grpc_trace_metadata = {
             oc_grpc.GRPC_TRACE_KEY: header,
         }
 
-        metadata_to_append = None
-
         if isinstance(metadata, list):
             metadata_to_append = list(six.iteritems(grpc_trace_metadata))
         else:
@@ -122,7 +124,7 @@ 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()
+            self.tracer.end_span()
 
         return callback
 
@@ -133,7 +135,7 @@ def _trace_future_exception(self, response):
         if exception is not None:
             exception = str(exception)
 
-        self._tracer.add_attribute_to_current_span(
+        self.tracer.add_attribute_to_current_span(
             attribute_key=attributes_helper.COMMON_ATTRIBUTES.get(
                 ATTRIBUTE_ERROR_MESSAGE),
             attribute_value=exception)

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

@@ -27,13 +27,11 @@ def setUp(self):
         execution_context.clear()
 
     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._tracer)
         self.assertIsNone(interceptor.host_port)
         self.assertTrue(isinstance(
             interceptor._propagator, binary_format.BinaryFormatPropagator))