Thread pools integration (#363)

* Handle thread pools

* Clean tracing of debug informations

* Improve testing

* Add threadpool executor

* Implement testing

* Fix lint

Author: geobeau

opencensus/trace/execution_context.py

@@ -38,6 +38,10 @@ def set_opencensus_attr(attr_key, attr_value):
     setattr(_thread_local, 'attrs', attrs)
 
 
+def set_opencensus_attrs(attrs):
+    setattr(_thread_local, 'attrs', attrs)
+
+
 def get_opencensus_attr(attr_key):
     attrs = getattr(_thread_local, 'attrs', None)
 
@@ -47,6 +51,10 @@ def get_opencensus_attr(attr_key):
     return None
 
 
+def get_opencensus_attrs():
+    return getattr(_thread_local, 'attrs', None)
+
+
 def get_current_span():
     return getattr(_thread_local, 'current_span', None)
 
@@ -55,6 +63,30 @@ def set_current_span(current_span):
     setattr(_thread_local, 'current_span', current_span)
 
 
+def get_opencensus_full_context():
+    _tracer = get_opencensus_tracer()
+    _span = get_current_span()
+    _attrs = get_opencensus_attrs()
+    return _tracer, _span, _attrs
+
+
+def set_opencensus_full_context(tracer, span, attrs):
+    set_opencensus_tracer(tracer)
+    set_current_span(span)
+    if not attrs:
+        set_opencensus_attrs({})
+    else:
+        set_opencensus_attrs(attrs)
+
+
+def clean():
+    setattr(_thread_local, 'attrs', {})
+    if hasattr(_thread_local, 'current_span'):
+        delattr(_thread_local, 'current_span')
+    if hasattr(_thread_local, 'tracer'):
+        delattr(_thread_local, 'tracer')
+
+
 def clear():
     """Clear the thread local, used in test."""
     _thread_local.__dict__.clear()

opencensus/trace/ext/threading/trace.py

@@ -14,8 +14,12 @@
 
 import logging
 import threading
+from multiprocessing import pool
+from concurrent import futures
 
 from opencensus.trace import execution_context
+from opencensus.trace import tracer
+from opencensus.trace.propagation import binary_format
 
 log = logging.getLogger(__name__)
 
@@ -27,13 +31,29 @@ def trace_integration(tracer=None):
     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))
+    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))
+    setattr(threading.Thread, run_func.__name__, wrap_threading_run(run_func))
+
+    # Wrap the threading run function
+    apply_async_func = getattr(pool.Pool, "apply_async")
+    setattr(
+        pool.Pool,
+        apply_async_func.__name__,
+        wrap_apply_async(apply_async_func),
+    )
+
+    # Wrap the threading run function
+    submit_func = getattr(futures.ThreadPoolExecutor, "submit")
+    setattr(
+        futures.ThreadPoolExecutor,
+        submit_func.__name__,
+        wrap_submit(submit_func),
+    )
 
 
 def wrap_threading_start(start_func):
@@ -42,7 +62,9 @@ def wrap_threading_start(start_func):
     """
 
     def call(self):
-        self.__opencensus_tracer = execution_context.get_opencensus_tracer()
+        self._opencensus_context = (
+            execution_context.get_opencensus_full_context()
+        )
         return start_func(self)
 
     return call
@@ -54,7 +76,80 @@ def wrap_threading_run(run_func):
     """
 
     def call(self):
-        execution_context.set_opencensus_tracer(self.__opencensus_tracer)
+        execution_context.set_opencensus_full_context(
+            *self._opencensus_context
+        )
         return run_func(self)
 
     return call
+
+
+def wrap_apply_async(apply_async_func):
+    """Wrap the apply_async function of multiprocessing.pools. Get the function
+    that will be called and wrap it then add the opencensus context."""
+
+    def call(self, func, args=(), kwds={}, **kwargs):
+        wrapped_func = wrap_task_func(func)
+        _tracer = execution_context.get_opencensus_tracer()
+        propagator = binary_format.BinaryFormatPropagator()
+
+        wrapped_kwargs = {}
+        print(_tracer)
+        wrapped_kwargs["span_context_binary"] = propagator.to_header(
+            _tracer.span_context
+        )
+        wrapped_kwargs["kwds"] = kwds
+        wrapped_kwargs["sampler"] = _tracer.sampler
+        wrapped_kwargs["exporter"] = _tracer.exporter
+        wrapped_kwargs["propagator"] = _tracer.propagator
+
+        return apply_async_func(
+            self, wrapped_func, args=args, kwds=wrapped_kwargs, **kwargs
+        )
+
+    return call
+
+
+def wrap_submit(submit_func):
+    """Wrap the apply_async function of multiprocessing.pools. Get the function
+    that will be called and wrap it then add the opencensus context."""
+
+    def call(self, func, *args, **kwargs):
+        wrapped_func = wrap_task_func(func)
+        _tracer = execution_context.get_opencensus_tracer()
+        propagator = binary_format.BinaryFormatPropagator()
+
+        wrapped_kwargs = {}
+        wrapped_kwargs["span_context_binary"] = propagator.to_header(
+            _tracer.span_context
+        )
+        wrapped_kwargs["kwds"] = kwargs
+        wrapped_kwargs["sampler"] = _tracer.sampler
+        wrapped_kwargs["exporter"] = _tracer.exporter
+        wrapped_kwargs["propagator"] = _tracer.propagator
+
+        return submit_func(self, wrapped_func, *args, **wrapped_kwargs)
+
+    return call
+
+
+class wrap_task_func(object):
+    """Wrap the function given to apply_async to get the tracer from context,
+    execute the function then clear the context."""
+
+    def __init__(self, func):
+        self.func = func
+
+    def __call__(self, *args, **kwargs):
+        kwds = kwargs.pop("kwds")
+
+        span_context_binary = kwargs.pop("span_context_binary")
+        propagator = binary_format.BinaryFormatPropagator()
+        kwargs["span_context"] = propagator.from_header(span_context_binary)
+
+        _tracer = tracer.Tracer(**kwargs)
+        execution_context.set_opencensus_tracer(_tracer)
+        with _tracer.span(name=threading.current_thread().name):
+            result = self.func(*args, **kwds)
+        execution_context.clean()
+        return result

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

@@ -15,10 +15,11 @@
 import unittest
 import threading
 import mock
+from multiprocessing.pool import Pool
+from concurrent.futures import ThreadPoolExecutor
 
-from opencensus.trace import span as span_module
 from opencensus.trace.ext.threading import trace
-from opencensus.trace import execution_context
+from opencensus.trace import execution_context, tracer
 
 class Test_threading_trace(unittest.TestCase):
 
@@ -29,36 +30,51 @@ def tearDown(self):
     def test_trace_integration(self):
         mock_wrap_start = mock.Mock()
         mock_wrap_run = mock.Mock()
+        mock_wrap_apply_async = mock.Mock()
+
         mock_threading = mock.Mock()
+        mock_pool = mock.Mock()
 
         wrap_start_result = 'wrap start result'
         wrap_run_result = 'wrap run result'
+        wrap_apply_async_result = 'wrap apply_async result'
         mock_wrap_start.return_value = wrap_start_result
         mock_wrap_run.return_value = wrap_run_result
+        mock_wrap_apply_async.return_value = wrap_apply_async_result
 
         mock_start_func = mock.Mock()
         mock_run_func = mock.Mock()
+        mock_apply_async_func = mock.Mock()
         mock_start_func.__name__ = 'start'
         mock_run_func.__name__ = 'run'
+        mock_apply_async_func.__name__ = 'apply_async'
         setattr(mock_threading.Thread, 'start', mock_start_func)
         setattr(mock_threading.Thread, 'run', mock_run_func)
+        setattr(mock_pool.Pool, 'apply_async', mock_apply_async_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_wrap_apply_async = mock.patch(
+            'opencensus.trace.ext.threading.trace.wrap_apply_async',
+            mock_wrap_apply_async)
         patch_threading = mock.patch(
             'opencensus.trace.ext.threading.trace.threading', mock_threading)
+        patch_pool = mock.patch(
+            'opencensus.trace.ext.threading.trace.pool', mock_pool)
 
-        with patch_wrap_start, patch_wrap_run, patch_threading:
+        with patch_wrap_start, patch_wrap_run, patch_wrap_apply_async, patch_threading, patch_pool:
             trace.trace_integration()
 
         self.assertEqual(getattr(mock_threading.Thread, 'start'),
                          wrap_start_result)
         self.assertEqual(getattr(mock_threading.Thread, 'run'),
                          wrap_run_result)
+        self.assertEqual(getattr(mock_pool.Pool, 'apply_async'),
+                         wrap_apply_async_result)
 
     def test_wrap_threading(self):
         global global_tracer
@@ -75,11 +91,45 @@ def test_wrap_threading(self):
         t.join()
         assert isinstance(global_tracer, MockTracer)
 
+    def test_wrap_pool(self):
+        _tracer = tracer.Tracer()
+        execution_context.set_opencensus_tracer(tracer)
+
+        trace.trace_integration()
+        context = tracer.Tracer().span_context
+        print(context.trace_id)
+
+        pool = Pool(processes=1)
+        with _tracer.span(name='span1'):
+            result = pool.apply_async(fake_pooled_func, ()).get(timeout=1)
+
+        self.assertEqual(result, context.trace_id)
+
+    def test_wrap_futures(self):
+        _tracer = tracer.Tracer()
+        execution_context.set_opencensus_tracer(tracer)
+
+        trace.trace_integration()
+        context = tracer.Tracer().span_context
+        print(context.trace_id)
+
+        pool = ThreadPoolExecutor(max_workers=1)
+        with _tracer.span(name='span1'):
+            future = pool.submit(fake_pooled_func)
+            result = future.result()
+
+        self.assertEqual(result, context.trace_id)
+
     def fake_threaded_func(self):
         global global_tracer
         global_tracer = execution_context.get_opencensus_tracer()
 
 
+def fake_pooled_func():
+    _tracer = execution_context.get_opencensus_tracer()
+    return _tracer.span_context.trace_id
+
+
 class MockTracer(object):
     def __init__(self, span=None):
         self.span = span

tests/unit/trace/test_execution_context.py

@@ -13,6 +13,8 @@
 # limitations under the License.
 
 import unittest
+import mock
+import threading
 
 from opencensus.trace import execution_context
 
@@ -38,3 +40,60 @@ def test_has_attrs(self):
         result = execution_context.get_opencensus_attr(key)
 
         self.assertEqual(result, value)
+
+    def test_get_and_set_full_context(self):
+        mock_tracer_get = mock.Mock()
+        mock_span_get = mock.Mock()
+        execution_context.set_opencensus_tracer(mock_tracer_get)
+        execution_context.set_current_span(mock_span_get)
+
+        execution_context.set_opencensus_attr("test", "test_value")
+
+        tracer, span, attrs = execution_context.get_opencensus_full_context()
+
+        self.assertEqual(mock_tracer_get, tracer)
+        self.assertEqual(mock_span_get, span)
+        self.assertEqual({"test":"test_value"}, attrs)
+
+        mock_tracer_set = mock.Mock()
+        mock_span_set = mock.Mock()
+
+        execution_context.set_opencensus_full_context(mock_tracer_set, mock_span_set, None)
+        self.assertEqual(mock_tracer_set, execution_context.get_opencensus_tracer())
+        self.assertEqual(mock_span_set, execution_context.get_current_span())
+        self.assertEqual({}, execution_context.get_opencensus_attrs())
+
+        execution_context.set_opencensus_full_context(mock_tracer_set, mock_span_set, {"test": "test_value"})
+        self.assertEqual("test_value", execution_context.get_opencensus_attr("test"))
+
+    def test_clean_tracer(self):
+        mock_tracer = mock.Mock()
+        some_value = mock.Mock()
+        execution_context.set_opencensus_tracer(mock_tracer)
+
+        thread_local = threading.local()
+        setattr(thread_local, 'random_non_oc_attr', some_value)
+
+        execution_context.clean()
+
+        self.assertNotEqual(mock_tracer, execution_context.get_opencensus_tracer())
+        self.assertEqual(some_value, getattr(thread_local, 'random_non_oc_attr'))
+
+    def test_clean_span(self):
+        mock_span = mock.Mock()
+        some_value = mock.Mock()
+        execution_context.set_current_span(mock_span)
+
+        thread_local = threading.local()
+        setattr(thread_local, 'random_non_oc_attr', some_value)
+
+        execution_context.clean()
+
+        self.assertNotEqual(mock_span, execution_context.get_current_span())
+        self.assertEqual(some_value, getattr(thread_local, 'random_non_oc_attr'))
+
+
+
+
+
+