Improve the integration modules to share tracer between threads (#133)

Author: liyanhui1228

opencensus/trace/config_integration.py

@@ -23,7 +23,7 @@
 PATH_PREFIX = 'opencensus.trace.ext'
 
 
-def trace_integrations(integrations):
+def trace_integrations(integrations, tracer=None):
     """Enable tracing on the selected integrations.
 
     :type integrations: list
@@ -35,7 +35,7 @@ def trace_integrations(integrations):
         try:
             path_to_module = '{}.{}.trace'.format(PATH_PREFIX, item)
             module = importlib.import_module(path_to_module)
-            module.trace_integration()
+            module.trace_integration(tracer=tracer)
             integrated.append(item)
         except Exception:
             log.warning(

opencensus/trace/ext/google_cloud_clientlibs/trace.py

@@ -35,25 +35,25 @@
 CREATE_CHANNEL = 'create_channel'
 
 
-def trace_integration():
+def trace_integration(tracer=None):
     """Trace the Google Cloud Client libraries by integrating with
     the transport level including HTTP and gRPC.
     """
     log.info('Integrated module: {}'.format(MODULE_NAME))
 
     # Integrate with gRPC
-    trace_grpc()
+    trace_grpc(tracer)
 
     # Integrate with HTTP
-    trace_http()
+    trace_http(tracer)
 
 
-def trace_grpc():
+def trace_grpc(tracer=None):
     """Integrate with gRPC."""
     # Wrap google.cloud._helpers.make_secure_channel
     make_secure_channel_func = getattr(_helpers, MAKE_SECURE_CHANNEL)
     make_secure_channel_wrapped = wrap_make_secure_channel(
-        make_secure_channel_func)
+        make_secure_channel_func, tracer)
     setattr(
         _helpers,
         MAKE_SECURE_CHANNEL,
@@ -62,34 +62,37 @@ def trace_grpc():
     # Wrap the grpc.insecure_channel.
     insecure_channel_func = getattr(grpc, INSECURE_CHANNEL)
     insecure_channel_wrapped = wrap_insecure_channel(
-        insecure_channel_func)
+        insecure_channel_func, tracer)
     setattr(
         grpc,
         INSECURE_CHANNEL,
         insecure_channel_wrapped)
 
     # Wrap google.api_core.grpc_helpers.create_channel
     create_channel_func = getattr(grpc_helpers, CREATE_CHANNEL)
-    create_channel_wrapped = wrap_create_channel(create_channel_func)
+    create_channel_wrapped = wrap_create_channel(create_channel_func, tracer)
     setattr(
         grpc_helpers,
         CREATE_CHANNEL,
         create_channel_wrapped)
 
 
-def trace_http():
+def trace_http(tracer=None):
     """Integrate with HTTP (requests library)."""
-    trace_requests()
+    trace_requests(tracer)
 
 
-def wrap_make_secure_channel(make_secure_channel_func):
+def wrap_make_secure_channel(make_secure_channel_func, tracer=None):
     """Wrap the google.cloud._helpers.make_secure_channel."""
     def call(*args, **kwargs):
         channel = make_secure_channel_func(*args, **kwargs)
 
         try:
             host = kwargs.get('host')
-            _tracer = execution_context.get_opencensus_tracer()
+            if tracer is None:
+                _tracer = execution_context.get_opencensus_tracer()
+            else:  # pragma: NO COVER
+                _tracer = tracer
             tracer_interceptor = OpenCensusClientInterceptor(_tracer, host)
             intercepted_channel = grpc.intercept_channel(
                 channel, tracer_interceptor)
@@ -102,14 +105,17 @@ def call(*args, **kwargs):
     return call
 
 
-def wrap_insecure_channel(insecure_channel_func):
+def wrap_insecure_channel(insecure_channel_func, tracer=None):
     """Wrap the grpc.insecure_channel."""
     def call(*args, **kwargs):
         channel = insecure_channel_func(*args, **kwargs)
 
         try:
             target = kwargs.get('target')
-            _tracer = execution_context.get_opencensus_tracer()
+            if tracer is None:
+                _tracer = execution_context.get_opencensus_tracer()
+            else:  # pragma: NO COVER
+                _tracer = tracer
             tracer_interceptor = OpenCensusClientInterceptor(_tracer, target)
             intercepted_channel = grpc.intercept_channel(
                 channel, tracer_interceptor)
@@ -122,14 +128,17 @@ def call(*args, **kwargs):
     return call
 
 
-def wrap_create_channel(create_channel_func):
+def wrap_create_channel(create_channel_func, tracer=None):
     """Wrap the google.api_core.grpc_helpers.create_channel."""
     def call(*args, **kwargs):
         channel = create_channel_func(*args, **kwargs)
 
         try:
             target = kwargs.get('target')
-            _tracer = execution_context.get_opencensus_tracer()
+            if tracer is None:
+                _tracer = execution_context.get_opencensus_tracer()
+            else:  # pragma: NO COVER
+                _tracer = tracer
             tracer_interceptor = OpenCensusClientInterceptor(_tracer, target)
             intercepted_channel = grpc.intercept_channel(
                 channel, tracer_interceptor)

opencensus/trace/ext/grpc/client_interceptor.py

@@ -32,6 +32,9 @@
 
 TIMEOUT = 3
 
+# Do not trace StackDriver Trace exporter activities to avoid deadlock.
+CLOUD_TRACE = 'google.devtools.cloudtrace'
+
 
 class _ClientCallDetails(
     collections.namedtuple(
@@ -137,6 +140,10 @@ def _trace_future_exception(self, response):
 
     def intercept_unary_unary(
             self, continuation, client_call_details, request):
+        if CLOUD_TRACE in client_call_details.method:
+            response = continuation(client_call_details, request)
+            return response
+
         new_details, new_request, current_span = self._intercept_call(
             client_call_details=client_call_details,
             request_iterator=iter((request,)),
@@ -152,6 +159,10 @@ def intercept_unary_unary(
 
     def intercept_unary_stream(self, continuation, client_call_details,
                                request):
+        if CLOUD_TRACE in client_call_details.method:
+            response = continuation(client_call_details, request)
+            return response
+
         new_details, new_request_iterator, current_span = self._intercept_call(
             client_call_details=client_call_details,
             request_iterator=iter((request,)),
@@ -166,6 +177,10 @@ def intercept_unary_stream(self, continuation, client_call_details,
 
     def intercept_stream_unary(self, continuation, client_call_details,
                                request_iterator):
+        if CLOUD_TRACE in client_call_details.method:
+            response = continuation(client_call_details, request_iterator)
+            return response
+
         new_details, new_request_iterator, current_span = self._intercept_call(
             client_call_details=client_call_details,
             request_iterator=request_iterator,
@@ -181,6 +196,10 @@ def intercept_stream_unary(self, continuation, client_call_details,
 
     def intercept_stream_stream(self, continuation, client_call_details,
                                 request_iterator):
+        if CLOUD_TRACE in client_call_details.method:
+            response = continuation(client_call_details, request_iterator)
+            return response
+
         new_details, new_request_iterator, current_span = self._intercept_call(
             client_call_details=client_call_details,
             request_iterator=request_iterator,

opencensus/trace/ext/httplib/trace.py

@@ -36,7 +36,7 @@
 HTTP_STATUS_CODE = attributes_helper.COMMON_ATTRIBUTES['HTTP_STATUS_CODE']
 
 
-def trace_integration():
+def trace_integration(tracer=None):
     """Wrap the httplib to trace."""
     log.info('Integrated module: {}'.format(MODULE_NAME))
 

opencensus/trace/ext/mysql/trace.py

@@ -23,7 +23,7 @@
 CONN_WRAP_METHOD = 'connect'
 
 
-def trace_integration():
+def trace_integration(tracer=None):
     """Wrap the mysql connector to trace it."""
     logging.info('Integrated module: {}'.format(MODULE_NAME))
     conn_func = getattr(mysql.connector, CONN_WRAP_METHOD)

opencensus/trace/ext/postgresql/trace.py

@@ -30,7 +30,7 @@
 QUERY_WRAP_METHODS = ['execute', 'executemany']
 
 
-def trace_integration():
+def trace_integration(tracer=None):
     """Wrap the mysql connector to trace it."""
     log.info('Integrated module: {}'.format(MODULE_NAME))
     conn_func = getattr(psycopg2, CONN_WRAP_METHOD)

opencensus/trace/ext/pymysql/trace.py

@@ -23,7 +23,7 @@
 CONN_WRAP_METHOD = 'connect'
 
 
-def trace_integration():
+def trace_integration(tracer=None):
     """Wrap the mysql connector to trace it."""
     logging.info('Integrated module: {}'.format(MODULE_NAME))
     conn_func = getattr(pymysql, CONN_WRAP_METHOD)

opencensus/trace/ext/requests/trace.py

@@ -27,10 +27,12 @@
 SESSION_CLASS_NAME = 'Session'
 
 
-def trace_integration():
+def trace_integration(tracer=None):
     """Wrap the requests library to trace it."""
     log.info('Integrated module: {}'.format(MODULE_NAME))
 
+    execution_context.set_opencensus_tracer(tracer)
+
     # Wrap the requests functions
     for func in REQUESTS_WRAP_METHODS:
         requests_func = getattr(requests, func)

opencensus/trace/ext/sqlalchemy/trace.py

@@ -25,7 +25,7 @@
 MODULE_NAME = 'sqlalchemy'
 
 
-def trace_integration():
+def trace_integration(tracer=None):
     """Integrate with SQLAlchemy to trace it using event listener.
 
     See: http://docs.sqlalchemy.org/en/latest/core/events.html

opencensus/trace/tracers/noop_tracer.py

@@ -19,6 +19,8 @@ class NoopTracer(base.Tracer):
     """No-op implementation of the :class:`Tracer` interface, all methods are
     no-ops. Should be used when tracing is not enabled or not sampled.
     """
+    span_context = None
+
     def finish(self):
         """End spans and send to reporter."""
         return None