Make `tracer.start_as_current_span()` decorator work with async functions

[aabmass]

Originaly from https://github.com/open-telemetry/opentelemetry-python/issues/62#issuecomment-1268415358

Alternatives:

• Document that it doesn't work. Adding a note doesn't hurt anyone, but I think the type annotations are pretty clear here and a type checker should catch this issue as well.
• Add a separate tracer.start_as_current_span_async() https://github.com/open-telemetry/opentelemetry-python/issues/62#issuecomment-1268910659. IMO this is confusing because it would only be a decorator unlike start_as_current_span().

[QuentinN42]

Any news on this subject ?

[aabmass]

If you check the original #62 you'll find some workarounds.

I think we need some consensus from the project maintainers if we want to go ahead with this feature request @open-telemetry/python-approvers.

I'm in favor personally 🙂

[ocelotl]

Same here, would you like to implement this, @aabmass?

[aabmass]

I don't have time right now if someone else wants to take it

[QuentinN42]

I've successfully make it works with a little bit of hacking :

    def _with_span_sync(self, func: Callable) -> Callable:
        """Decorate sync functions."""

        @wraps(func)
        def func_wrapper(*args, **kwargs):
            with self.start_as_current_span(func.__name__):
                return func(*args, **kwargs)

        return func_wrapper


    def _with_span_async(self, func: Callable) -> Callable:
        """Decorate async functions."""

        @wraps(func)
        async def func_wrapper(*args, **kwargs):
            with self.start_as_current_span(func.__name__):
                return await func(*args, **kwargs)

        return func_wrapper


    def decorate(self, func: Callable) -> Callable:
        # define if a function is async or not
        CO_COROUTINE = 0x0080
        
        if bool(func.__code__.co_flags & CO_COROUTINE):
            return self._with_span_async(func)

        return self._with_span_sync(func)

If you overload the main trace api with this method, you can have it working with both sync and async traces. https://github.com/open-telemetry/opentelemetry-python/blob/3dfe2249cc4a203bf24578483b192fec7266596b/opentelemetry-api/src/opentelemetry/trace/init.py#L269-L275

Here is my code example :

import asyncio
import time

from opentelemetry import trace
from opentelemetry.sdk.trace import TracerProvider
from opentelemetry.sdk.trace.export import BatchSpanProcessor, ConsoleSpanExporter
from opentelemetry.sdk.resources import SERVICE_NAME, Resource

# Setup console tracing :
# https://opentelemetry.io/docs/instrumentation/python/exporters/
resource = Resource(attributes={SERVICE_NAME: "your-service-name"})
provider = TracerProvider(resource=resource)
processor = BatchSpanProcessor(ConsoleSpanExporter())
provider.add_span_processor(processor)
trace.set_tracer_provider(provider)


tracer = trace.get_tracer("diceroller.tracer")


@tracer.start_as_current_span("sync")   # v1
# @tracer.decorate                      # v2
def sync() -> None:
    time.sleep(1)


@tracer.start_as_current_span("async")  # v1
# @tracer.decorate                      # v2
async def async_() -> None:
    await asyncio.sleep(1)


async def main() -> None:
    sync()
    await async_()


if __name__ == "__main__":
    asyncio.run(main())

Install the dependencies with :

python3 -m venv venv    
source ./venv/bin/activate
pip install opentelemetry-distro
python3 main.py

If you just launch this program with the latest version of OTLP it will print you two traces (sync with ~1s and async with ~0s).

If you paste my code inside the ./venv/lib/python3.11/site-packages/opentelemetry/trace/__init__.py (and add from typing import Callable & from functools import wraps), you will see two traces with the right times.

Here are my current changes : https://github.com/open-telemetry/opentelemetry-python/compare/main...QuentinN42:opentelemetry-python:main

[aabmass]

Discuss in SIG. There are a few options:

1. Add a new separate Tracer.decorate() decorator that works for both sync and async functions. This is what #3595 currently does
2. Add a new separate Tracer.start_as_current_span_async(), decorated with @asynccontextmanager
3. Make Tracer.start_as_current_span() behave as expected when decorating async functions

1 adds a new Tracer.decorate() function which is redundant for decorating sync functions which already work fine with Tracer.start_as_current_span(). 2 implies that start_as_current_span is blocking or slow, which is not the case–2 is really only useful as a decorator.

This github search turns up a few cases where start_as_current_span() is being misused with async functions. If we take approach 3, it would fix these cases.

[QuentinN42]

After some investigation, it seems possible to make the Tracer.start_as_current_span() work out of the box.

However, I was wondering if the API Tracer abstract class could be implemented by external libraries ?

If so, I'm afraid my changes will break the API as it's removing the @contextmanager and return a classes that implements the AbstractContextManager...

[aabmass]

However, I was wondering if the API Tracer abstract class could be implemented by external libraries ?

Yes it could be implemented by other SDKs. I'm not aware of any that exist.

If so, I'm afraid my changes will break the API as it's removing the @contextmanager and return a classes that implements the AbstractContextManager...

I think it should be OK as long as the returned object matches the existing public protocol of @contextmanager and we only need to extend the behavior of __call__(). Do you think that is feasible?

[QuentinN42]

import asyncio
import contextlib
import functools
import time


class AgnosticContextManager(contextlib._GeneratorContextManager):
    def __call__(self, func):
        if asyncio.iscoroutinefunction(func):

            @functools.wraps(func)
            async def async_wrapper(*args, **kwargs):
                with self._recreate_cm():
                    return await func(*args, **kwargs)

            return async_wrapper
        else:

            @functools.wraps(func)
            def wrapper(*args, **kwargs):
                with self._recreate_cm():
                    return func(*args, **kwargs)

            return wrapper


def agnosticcontextmanager(func):
    @functools.wraps(func)
    def helper(*args, **kwds):
        return AgnosticContextManager(func, args, kwds)

    return helper


@agnosticcontextmanager
def start_as_current_span():
    start_time = time.time()
    yield
    end_time = time.time()
    print(f"Elapsed time: {end_time - start_time} seconds")


@start_as_current_span()
def sync_example():
    time.sleep(0.01)


@start_as_current_span()
async def async_example():
    await asyncio.sleep(0.01)


# Using examples
sync_example()
asyncio.run(async_example())

Here is a minimal fix to the tracer.start_as_current_span() method, I'm trying to implement it inside the codebase on my fork https://github.com/open-telemetry/opentelemetry-python/compare/main...QuentinN42:opentelemetry-python:feat/make-tracer-start-as-current-span-decorator-work-with-async-functions

[Chise1]

Maybe we don't have to make tracer.start_as_current_span() decorator work with async functions,we can add this task to coroutine loop,code like this:

import asyncio
from asyncio.queues import Queue
async def start(q :Queue):
    await asyncio.sleep(1)
    print("start")
    await q.put(0)
async def end(q :Queue):
    await q.get()
    print("end")
def xx():
    q=Queue()
    asyncio.create_task(start(q))
    print("xx")
    asyncio.create_task(end(q))

async def main():
    xx()
    await asyncio.sleep(2)
    print("main")


if __name__ == '__main__':
    asyncio.run(main())

We can use an environment variable to control whether to use coroutines or multithreads.