diff --git a/beeline/__init__.py b/beeline/__init__.py index 6c7d059..e84a3a1 100644 --- a/beeline/__init__.py +++ b/beeline/__init__.py @@ -207,7 +207,7 @@ def traced(self, name, trace_id=None, parent_id=None): return traced_impl(tracer_fn=self.tracer, name=name, trace_id=trace_id, parent_id=parent_id) def traced_thread(self, fn): - trace_copy = self.tracer_impl._state.trace.copy() + trace_copy = self.tracer_impl._trace.copy() @functools.wraps(fn) def wrapped(*args, **kwargs): @@ -660,7 +660,7 @@ def noop(*args, **kwargs): return fn(*args, **kwargs) return noop - trace_copy = bl.tracer_impl._state.trace.copy() + trace_copy = bl.tracer_impl._trace.copy() @functools.wraps(fn) def wrapped(*args, **kwargs): diff --git a/beeline/test_async.py b/beeline/test_async.py index aee7fb9..0cdf3e5 100644 --- a/beeline/test_async.py +++ b/beeline/test_async.py @@ -1,5 +1,7 @@ import asyncio +import concurrent.futures import datetime +import time import unittest import beeline @@ -354,3 +356,58 @@ async def task1(): # Verify that the untraced functions were actually called self.assertTrue("fn0" in calls) self.assertTrue("fn1" in calls) + + @async_test + async def test_traced_thread_should_work_with_async_tracer(self): + """Run traced code in threads from within async code. + + A trace is started and two functions are run in threads via a + ThreadPoolExecutor. Both functions start spans, but only one + of them is decorated with the traced_thread decorator. Only + the span from the decorated function is expected to show up in + the trace. + + """ + loop = asyncio.get_running_loop() # pylint: disable=no-member + executor = concurrent.futures.ThreadPoolExecutor(max_workers=2) + calls = set() + + trace = self.tracer.start_trace(context={"name": "root"}) + + @self.beeline.traced_thread + def traced_worker(): + span = self.tracer.start_span(context={"name": "traced_worker"}) + time.sleep(0.2) + self.tracer.finish_span(span) + + def untraced_worker(): + span = self.tracer.start_span(context={"name": "untraced_worker"}) + time.sleep(0.2) + self.tracer.finish_span(span) + calls.add("untraced_worker") + + future0 = loop.run_in_executor(executor, traced_worker) + future1 = loop.run_in_executor(executor, untraced_worker) + + await asyncio.gather(future0, future1) + + self.tracer.finish_trace(trace) + + self.assertEqual(len(self.finished_spans), 2) + + worker_span, root_span = self.finished_spans # pylint: disable=unbalanced-tuple-unpacking + + # Check that the spans finished in the expected order, with + # the root span last. + self.assertEqual(worker_span["name"], "traced_worker") + self.assertEqual(root_span["name"], "root") + self.assertLessEqual(worker_span["end"], root_span["end"]) + + # Check that the root span was started before the worker span. + self.assertLess(root_span["start"], worker_span["start"]) + + # Check that the worker span is a child of the root span. + self.assertEqual(root_span["span"].id, worker_span["span"].parent_id) + + # Verify that the untraced function was actually called + self.assertTrue("untraced_worker" in calls)