Repository: census-instrumentation/opencensus-python
Pull request #133 (2018-03-08 01:29:23)

Improve the integration modules to share tracer between threads


Patch:
diff --git a/opencensus/trace/config_integration.py b/opencensus/trace/config_integration.py
index 33270ed..fcc7d76 100644
--- a/opencensus/trace/config_integration.py
+++ b/opencensus/trace/config_integration.py
@@ -23,7 +23,7 @@ SUPPORTED_INTEGRATIONS = ['httplib', 'mysql', 'postgresql', 'pymysql',
 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(
diff --git a/opencensus/trace/ext/google_cloud_clientlibs/trace.py b/opencensus/trace/ext/google_cloud_clientlibs/trace.py
index 862a9a1..91d8351 100644
--- a/opencensus/trace/ext/google_cloud_clientlibs/trace.py
+++ b/opencensus/trace/ext/google_cloud_clientlibs/trace.py
@@ -35,25 +35,25 @@ INSECURE_CHANNEL = 'insecure_channel'
 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,7 +62,7 @@ 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,
@@ -70,26 +70,29 @@ def trace_grpc():
 
     # 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 wrap_make_secure_channel(make_secure_channel_func):
     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 wrap_insecure_channel(insecure_channel_func):
     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)
diff --git a/opencensus/trace/ext/grpc/client_interceptor.py b/opencensus/trace/ext/grpc/client_interceptor.py
index 09af2b7..82c25c3 100644
--- a/opencensus/trace/ext/grpc/client_interceptor.py
+++ b/opencensus/trace/ext/grpc/client_interceptor.py
@@ -32,6 +32,9 @@ GRPC_METHOD = 'GRPC_METHOD'
 
 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 @@ class OpenCensusClientInterceptor(grpc.UnaryUnaryClientInterceptor,
 
     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 @@ class OpenCensusClientInterceptor(grpc.UnaryUnaryClientInterceptor,
 
     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 @@ class OpenCensusClientInterceptor(grpc.UnaryUnaryClientInterceptor,
 
     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 @@ class OpenCensusClientInterceptor(grpc.UnaryUnaryClientInterceptor,
 
     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,
diff --git a/opencensus/trace/ext/httplib/trace.py b/opencensus/trace/ext/httplib/trace.py
index eb8168d..c077348 100644
--- a/opencensus/trace/ext/httplib/trace.py
+++ b/opencensus/trace/ext/httplib/trace.py
@@ -36,7 +36,7 @@ HTTP_URL = attributes_helper.COMMON_ATTRIBUTES['HTTP_URL']
 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))
 
diff --git a/opencensus/trace/ext/mysql/trace.py b/opencensus/trace/ext/mysql/trace.py
index 58f9783..f7a4e8a 100644
--- a/opencensus/trace/ext/mysql/trace.py
+++ b/opencensus/trace/ext/mysql/trace.py
@@ -23,7 +23,7 @@ MODULE_NAME = 'mysql'
 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)
diff --git a/opencensus/trace/ext/postgresql/trace.py b/opencensus/trace/ext/postgresql/trace.py
index c2a5560..44e4572 100644
--- a/opencensus/trace/ext/postgresql/trace.py
+++ b/opencensus/trace/ext/postgresql/trace.py
@@ -30,7 +30,7 @@ CURSOR_WRAP_METHOD = 'cursor'
 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)
diff --git a/opencensus/trace/ext/pymysql/trace.py b/opencensus/trace/ext/pymysql/trace.py
index 8d839e6..e912bb9 100644
--- a/opencensus/trace/ext/pymysql/trace.py
+++ b/opencensus/trace/ext/pymysql/trace.py
@@ -23,7 +23,7 @@ MODULE_NAME = 'pymysql'
 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)
diff --git a/opencensus/trace/ext/requests/trace.py b/opencensus/trace/ext/requests/trace.py
index 01dd4cc..db688d7 100644
--- a/opencensus/trace/ext/requests/trace.py
+++ b/opencensus/trace/ext/requests/trace.py
@@ -27,10 +27,12 @@ SESSION_WRAP_METHODS = 'request'
 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)
diff --git a/opencensus/trace/ext/sqlalchemy/trace.py b/opencensus/trace/ext/sqlalchemy/trace.py
index aa1cf2c..ce3f913 100644
--- a/opencensus/trace/ext/sqlalchemy/trace.py
+++ b/opencensus/trace/ext/sqlalchemy/trace.py
@@ -25,7 +25,7 @@ log = logging.getLogger(__name__)
 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
diff --git a/opencensus/trace/tracers/noop_tracer.py b/opencensus/trace/tracers/noop_tracer.py
index 2f54be4..6fa65cf 100644
--- a/opencensus/trace/tracers/noop_tracer.py
+++ b/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


Test patch:
diff --git a/tests/unit/trace/ext/grpc/test_client_interceptor.py b/tests/unit/trace/ext/grpc/test_client_interceptor.py
index b4281a4..b469667 100644
--- a/tests/unit/trace/ext/grpc/test_client_interceptor.py
+++ b/tests/unit/trace/ext/grpc/test_client_interceptor.py
@@ -20,6 +20,12 @@ from opencensus.trace.ext.grpc import client_interceptor
 
 
 class TestOpenCensusClientInterceptor(unittest.TestCase):
+
+    def setUp(self):
+        from opencensus.trace import execution_context
+
+        execution_context.clear()
+
     def test_constructor_default(self):
         from opencensus.trace.tracers import noop_tracer
         from opencensus.trace.propagation import binary_format
@@ -181,26 +187,66 @@ class TestOpenCensusClientInterceptor(unittest.TestCase):
         interceptor._intercept_call = mock.Mock(return_value=(None, iter([mock.Mock()]), None))
         return interceptor, continuation, mock_tracer
 
-    def test_intercept_unary_unary(self):
+    def test_intercept_unary_unary_trace(self):
         interceptor, continuation, mock_response = self._unary_helper()
-        interceptor.intercept_unary_unary(continuation, mock.Mock(), [])
+        client_call_details = mock.Mock()
+        client_call_details.method = 'test'
+        interceptor.intercept_unary_unary(continuation, client_call_details, [])
         self.assertTrue(mock_response.add_done_callback.called)
 
-    def test_intercept_unary_stream(self):
+    def test_intercept_unary_unary_not_trace(self):
+        interceptor, continuation, mock_response = self._unary_helper()
+        client_call_details = mock.Mock()
+        client_call_details.method = 'google.devtools.cloudtrace'
+        interceptor.intercept_unary_unary(continuation, client_call_details, [])
+        # Should skip tracing the cloud trace activities
+        self.assertFalse(mock_response.add_done_callback.called)
+
+    def test_intercept_unary_stream_trace(self):
         interceptor, continuation, mock_tracer = self._stream_helper()
-        interceptor.intercept_unary_stream(continuation, mock.Mock(), [])
+        client_call_details = mock.Mock()
+        client_call_details.method = 'test'
+        interceptor.intercept_unary_stream(continuation, client_call_details, [])
         self.assertTrue(mock_tracer.end_span.called)
 
-    def test_intercept_stream_unary(self):
+    def test_intercept_unary_stream_not_trace(self):
+        interceptor, continuation, mock_tracer = self._stream_helper()
+        client_call_details = mock.Mock()
+        client_call_details.method = 'google.devtools.cloudtrace'
+        interceptor.intercept_unary_stream(continuation, client_call_details, [])
+        # Should skip tracing the cloud trace activities
+        self.assertFalse(mock_tracer.end_span.called)
+
+    def test_intercept_stream_unary_trace(self):
         interceptor, continuation, mock_response = self._unary_helper()
-        interceptor.intercept_stream_unary(continuation, mock.Mock(), [])
+        client_call_details = mock.Mock()
+        client_call_details.method = 'test'
+        interceptor.intercept_stream_unary(continuation, client_call_details, [])
         self.assertTrue(mock_response.add_done_callback.called)
 
-    def test_intercept_stream_stream(self):
+    def test_intercept_stream_unary_not_trace(self):
+        interceptor, continuation, mock_tracer = self._stream_helper()
+        client_call_details = mock.Mock()
+        client_call_details.method = 'google.devtools.cloudtrace'
+        interceptor.intercept_stream_unary(continuation, client_call_details, [])
+        # Should skip tracing the cloud trace activities
+        self.assertFalse(mock_tracer.end_span.called)
+
+    def test_intercept_stream_stream_trace(self):
         interceptor, continuation, mock_tracer = self._stream_helper()
-        interceptor.intercept_stream_stream(continuation, mock.Mock(), [])
+        client_call_details = mock.Mock()
+        client_call_details.method = 'test'
+        interceptor.intercept_stream_stream(continuation, client_call_details, [])
         self.assertTrue(mock_tracer.end_span.called)
 
+    def test_intercept_stream_stream_not_trace(self):
+        interceptor, continuation, mock_tracer = self._stream_helper()
+        client_call_details = mock.Mock()
+        client_call_details.method = 'google.devtools.cloudtrace'
+        interceptor.intercept_stream_stream(continuation, client_call_details, [])
+        # Should skip tracing the cloud trace activities
+        self.assertFalse(mock_tracer.end_span.called)
+
 
 class MockTracer(object):
     def __init__(self, current_span):
diff --git a/tests/unit/trace/test_config_integration.py b/tests/unit/trace/test_config_integration.py
index 01a1fda..548c8fd 100644
--- a/tests/unit/trace/test_config_integration.py
+++ b/tests/unit/trace/test_config_integration.py
@@ -22,9 +22,7 @@ from opencensus.trace import config_integration
 class Test_trace_integrations(unittest.TestCase):
 
     def test_trace_integrations_not_exist(self):
-        integrations = {
-            'test_not_exists': True,
-        }
+        integrations = ['test_not_exists']
 
         integrated = config_integration.trace_integrations(integrations)