add instrumentation to generator functions

beeline/aiotrace.py

@@ -6,6 +6,7 @@
 import asyncio
 import contextvars  # pylint: disable=import-error
 import functools
+import inspect
 
 from beeline.trace import Tracer
 
@@ -77,7 +78,14 @@ async def async_inner(*args, **kwargs):
                     return await fn(*args, **kwargs)
 
             return async_inner
+        elif inspect.isgeneratorfunction(fn):
+            @functools.wraps(fn)
+            def inner(*args, **kwargs):
+                inner_generator = fn(*args, **kwargs)
+                with tracer_fn(name=name, trace_id=trace_id, parent_id=parent_id):
+                    yield from inner_generator
 
+            return inner
         else:
             @functools.wraps(fn)
             def inner(*args, **kwargs):

beeline/test_beeline.py

@@ -150,6 +150,41 @@ def my_sum(a, b):
             # check that an event was sent, from which we can infer that the function was wrapped
             self.assertTrue(_beeline.tracer_impl._run_hooks_and_send.called)
 
+    def test_generator_wrapper(self):
+        ''' ensure that the trace wrapper decorates a generator function and starts a trace
+            also ensure that child traces get the parent trace correctly set
+        '''
+
+        _beeline = beeline.Beeline()
+
+        with patch('beeline.get_beeline') as m_gbl:
+            m_gbl.return_value = _beeline
+            _beeline.tracer_impl._run_hooks_and_send = Mock()
+
+            @beeline.traced(name="return_integer_n")
+            def return_integer(n):
+                return n
+
+            @beeline.traced(name="output_integers_to")
+            def output_integers_to(n):
+                for i in range(n):
+                    yield return_integer(i)
+
+            # should accept the function's arguments normally and yield the items from the
+            # generator
+            self.assertEqual(list(output_integers_to(3)), [0, 1, 2])
+
+            self.assertTrue(_beeline.tracer_impl._run_hooks_and_send.called)
+
+            spans = [x[0][0] for x in _beeline.tracer_impl._run_hooks_and_send.call_args_list]
+
+            # check the child spans now
+            parent_span = spans[-1]
+            child_spans = spans[:-1]
+
+            for child_span in child_spans:
+                self.assertEqual(child_span.parent_id, parent_span.id)
+
     @staticmethod
     def raising_run_in_thread(target):
         closure_dict = {}
@@ -246,4 +281,4 @@ def my_sum(a, b):
 
         # this should not crash if the beeline isn't initialized
         # it should also accept arguments normally and return the function's value
-        self.assertEqual(my_sum(1, 2), 3)
+        self.assertEqual(my_sum(1, 2), 3)

beeline/trace.py

@@ -8,6 +8,7 @@
 import struct
 import threading
 import uuid
+import inspect
 from collections import defaultdict
 
 from contextlib import contextmanager
@@ -359,11 +360,18 @@ def unmarshal_trace_context(trace_context):
 def traced_impl(tracer_fn, name, trace_id, parent_id):
     """Implementation of the traced decorator without async support."""
     def wrapped(fn):
-        @functools.wraps(fn)
-        def inner(*args, **kwargs):
-            with tracer_fn(name=name, trace_id=trace_id, parent_id=parent_id):
-                return fn(*args, **kwargs)
-
-        return inner
-
+        if inspect.isgeneratorfunction(fn):
+            @functools.wraps(fn)
+            def inner(*args, **kwargs):
+                inner_generator = fn(*args, **kwargs)
+                with tracer_fn(name=name, trace_id=trace_id, parent_id=parent_id):
+                    for value in inner_generator:
+                        yield value
+            return inner
+        else:
+            @functools.wraps(fn)
+            def inner(*args, **kwargs):
+                with tracer_fn(name=name, trace_id=trace_id, parent_id=parent_id):
+                    return fn(*args, **kwargs)
+            return inner
     return wrapped