tracer.wrap does not work on generators correctly

[MichaelAquilina]

Summary of problem

Generator functions wrapped by tracer.wrap() do not seem to work correctly.

In particular, trying to get the current span inside the wrapped function does not seem to work

Which version of dd-trace-py are you using?

1.6.1 (also tried on latest 1.10.1)

Which version of pip are you using?

poetry

Which libraries and their versions are you using?

`pip freeze`

How can we reproduce your problem?

Here is a minimal code example:

from ddtrace import tracer

@tracer.wrap()
def foobar():
     current_span = tracer.current_span()
     current_span.set_tag("hello", "world")
     yield 1

What is the result that you get?

Calling list(foobar()) results in the following error:

AttributeError: 'NoneType' object has no attribute 'set_tag'

What is the result that you expected?

current_span is found and able to have tags set.

[P403n1x87]

I suspect this is because tracer.wrap is not unwinding the generator, just wrapping the call to foobar that returns the generator object.

[MichaelAquilina]

@P403n1x87 I think that is correct. The main issue here is that its hard for tracer.wrap to work generically for both functions and generators because as soon as you use a yield statement within wrap_decorator, it will start returning generators instead.

I think the best solution is to create a separate wrapped specifically for generators. I am writing what that would look like now but thought it would be good to get early feedback in the meanwhile?

[majorgreys]

@MichaelAquilina Thank you for starting this discussion and starting a PR!

It's been a while since I looked closely at Tracer.wrap. If it's possible, we should keep that interface and have the function also handle generators.

Internally, we use wrapt for wrapping functions and objects. I wonder if we should be extending it's decorator.

For example, here's one approach that doesn't hit the issue you are finding:

import inspect

from ddtrace import tracer
import wrapt


@wrapt.decorator
def wrapper(wrapped, instance, args, kwargs):
  with tracer.trace(wrapped.__name__):
    if not inspect.isgeneratorfunction(wrapped):
      return wrapped(*args, **kwargs)
      
    yield wrapped(*args, **kwargs)


@wrapper
def foobar():
  current_span = tracer.current_span()
  current_span.set_tag("hello", "world")
  yield 1
  
  
if __name__ == "__main__":
  assert inspect.isgeneratorfunction(foobar)
  print("main")
  list(foobar())

Would this resolve your issue?

[MichaelAquilina]

@majorgreys that is actually similar to the approach I started off with, but the problem is that once you include a yield inside a function, python turns it into a generator.

Unless wrapt somehow deals with this (Which I couldnt find in their docs), then the returned value will be incorrect for standard functions.

I used this code to test it out:

import inspect

from ddtrace import tracer
import wrapt


@wrapt.decorator
def wrapper(wrapped, instance, args, kwargs):
  with tracer.trace(wrapped.__name__):
    if not inspect.isgeneratorfunction(wrapped):
      return wrapped(*args, **kwargs)

    yield wrapped(*args, **kwargs)


@wrapper
def foobar():
    current_span = tracer.current_span()
    current_span.set_tag("hello", "world")
    return [1]

@wrapper
def bazbang():
    current_span = tracer.current_span()
    current_span.set_tag("hello", "world")
    yield 1


def main():
    print(foobar())
    print(list(bazbang()))


if __name__ == "__main__":
    main()

The output is:

<generator object wrapper at 0x7fe3f54ce1f0>
[<generator object bazbang at 0x7fe3f54ce260>]

In this case, the code doesnt even seem to return the expected type for either so I guess there is a bug in the implementation, bt notice how both cases are generators.

I implemented something similar in tracer.wrap() with inspect previously :

This is the diff:

diff --git a/ddtrace/tracer.py b/ddtrace/tracer.py
index 4122aae3..038cfee3 100644
--- a/ddtrace/tracer.py
+++ b/ddtrace/tracer.py
@@ -1,3 +1,4 @@
+import inspect
 import functools
 import json
 import logging
@@ -980,6 +981,11 @@ class Tracer(object):
                             span_type=span_type,
                         )
 
+                    if inspect.isgeneratorfunction(f):
+                        # otherwise fallback to a default tracing
+                        with self.trace(span_name, service=service, resource=resource, span_type=span_type):
+                            yield from f(*args, **kwargs)
+
                     # otherwise fallback to a default tracing
                     with self.trace(span_name, service=service, resource=resource, span_type=span_type):
                         return f(*args, **kwargs)

and this is what was returned

<generator object foobar at 0x7f4937ee9070>
[1]

[emmettbutler]

Manually doing what the Github Action was supposed to do (fixed in #7835):

This issue has been automatically closed after six months of inactivity. If it's a feature request, it has been added to the maintainers' internal backlog and will be included in an upcoming round of feature prioritization. Please comment or reopen if you think this issue was closed in error.

[MatejBalantic]

We see a similar issue. Anything that happens inside generator functions fails to get traced. Specifically the SQL queries done with the psycopg do not show up when executed from within the generator functions. This is a big issue, since almost all of the SQL queries get executed from within the generator functions in our case.

Is there no workaround for this issue?

This is our SQL store code:

    @contextmanager
    def with_transaction(self, transaction: Transaction | None = None) -> Generator[Transaction, None, None]:
        with tracer.trace("SQLStore.with_transaction"):
            try:
                if transaction is not None:
                    yield transaction

                else:
                    with self.db_connection.get_connection() as conn, conn.transaction() as tx:
                        with tracer.trace("SQLStore.with_transaction"):
                            yield tx

I can see the manual traces (SQLStore.with_transaction and SQLStore.with_transaction), but nothing happening inside there is traced, not even the SQL queries from the patched psycopg

[tommyhe6]

Wanted to follow up on this! We use generators often in the context of streaming LLM token results back to users, and I think that is a very common use case.