Add Threading integration (#329)

* Add Threading integration

* Fix bad descriptions and documentation to remove ref to httplib/request

* Fix date of copyright

Author: geobeau

README.rst

@@ -430,6 +430,12 @@ You can enable Google Cloud client libraries integration by specifying ``'google
 
 .. _Cloud client libraries: https://github.com/GoogleCloudPlatform/google-cloud-python#google-cloud-python-client
 
+Threading
+~~~~~~~~~
+
+Census can propagate trace across threads when using the Threading package.
+
+You can enable Threading integration by specifying ``'threading'`` to ``trace_integrations``.
 
 ------
  Stats

opencensus/trace/config_integration.py

@@ -18,7 +18,8 @@
 log = logging.getLogger(__name__)
 
 SUPPORTED_INTEGRATIONS = ['httplib', 'mysql', 'postgresql', 'pymysql',
-                          'requests', 'sqlalchemy', 'google_cloud_clientlibs']
+                          'requests', 'sqlalchemy', 'google_cloud_clientlibs',
+                          'threading']
 
 PATH_PREFIX = 'opencensus.trace.ext'
 

opencensus/trace/ext/threading/__init__.py

@@ -0,0 +1,13 @@
+# Copyright 2018, 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.

opencensus/trace/ext/threading/trace.py

@@ -0,0 +1,60 @@
+# Copyright 2018, 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 logging
+import threading
+
+from opencensus.trace import execution_context
+
+log = logging.getLogger(__name__)
+
+MODULE_NAME = "threading"
+
+
+def trace_integration(tracer=None):
+    """Wrap threading functions to trace."""
+    log.info("Integrated module: {}".format(MODULE_NAME))
+    # Wrap the threading start function
+    start_func = getattr(threading.Thread, "start")
+    setattr(threading.Thread, start_func.__name__,
+            wrap_threading_start(start_func))
+
+    # Wrap the threading run function
+    run_func = getattr(threading.Thread, "run")
+    setattr(threading.Thread, run_func.__name__,
+            wrap_threading_run(run_func))
+
+
+def wrap_threading_start(start_func):
+    """Wrap the start function from thread. Put the tracer informations in the
+    threading object.
+    """
+
+    def call(self):
+        self.__opencensus_tracer = execution_context.get_opencensus_tracer()
+        return start_func(self)
+
+    return call
+
+
+def wrap_threading_run(run_func):
+    """Wrap the run function from thread. Get the tracer informations from the
+    threading object and set it as current tracer.
+    """
+
+    def call(self):
+        execution_context.set_opencensus_tracer(self.__opencensus_tracer)
+        return run_func(self)
+
+    return call

tests/unit/trace/ext/threading/test_threading_trace.py

@@ -0,0 +1,105 @@
+# Copyright 2018, 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 threading
+import mock
+
+from opencensus.trace import span as span_module
+from opencensus.trace.ext.threading import trace
+from opencensus.trace import execution_context
+
+class Test_threading_trace(unittest.TestCase):
+
+    def tearDown(self):
+        execution_context.clear()
+
+
+    def test_trace_integration(self):
+        mock_wrap_start = mock.Mock()
+        mock_wrap_run = mock.Mock()
+        mock_threading = mock.Mock()
+
+        wrap_start_result = 'wrap start result'
+        wrap_run_result = 'wrap run result'
+        mock_wrap_start.return_value = wrap_start_result
+        mock_wrap_run.return_value = wrap_run_result
+
+        mock_start_func = mock.Mock()
+        mock_run_func = mock.Mock()
+        mock_start_func.__name__ = 'start'
+        mock_run_func.__name__ = 'run'
+        setattr(mock_threading.Thread, 'start', mock_start_func)
+        setattr(mock_threading.Thread, 'run', mock_run_func)
+
+        patch_wrap_start = mock.patch(
+            'opencensus.trace.ext.threading.trace.wrap_threading_start',
+            mock_wrap_start)
+        patch_wrap_run = mock.patch(
+            'opencensus.trace.ext.threading.trace.wrap_threading_run',
+            mock_wrap_run)
+        patch_threading = mock.patch(
+            'opencensus.trace.ext.threading.trace.threading', mock_threading)
+
+        with patch_wrap_start, patch_wrap_run, patch_threading:
+            trace.trace_integration()
+
+        self.assertEqual(getattr(mock_threading.Thread, 'start'),
+                         wrap_start_result)
+        self.assertEqual(getattr(mock_threading.Thread, 'run'),
+                         wrap_run_result)
+
+    def test_wrap_threading(self):
+        global global_tracer
+        mock_span = mock.Mock()
+        span_id = '1234'
+        mock_span.span_id = span_id
+        mock_tracer = MockTracer(mock_span)
+        execution_context.set_opencensus_tracer(mock_tracer)
+
+        trace.trace_integration()
+
+        t = threading.Thread(target=self.fake_threaded_func)
+        t.start()
+        t.join()
+        assert isinstance(global_tracer, MockTracer)
+
+    def fake_threaded_func(self):
+        global global_tracer
+        global_tracer = execution_context.get_opencensus_tracer()
+
+
+class MockTracer(object):
+    def __init__(self, span=None):
+        self.span = span
+
+    def current_span(self):
+        return self.span
+
+    def start_span(self):
+        span = mock.Mock()
+        span.attributes = {}
+        span.context_tracer = mock.Mock()
+        span.context_tracer.span_context = mock.Mock()
+        span.context_tracer.span_context.trace_id = '123'
+        span.context_tracer.span_context.span_id = '456'
+        span.context_tracer.span_context.tracestate = None
+        self.span = span
+        return span
+
+    def end_span(self):
+        pass
+
+    def add_attribute_to_current_span(self, key, value):
+        self.span.attributes[key] = value