investigate tags implementation for python coroutines

[korniltsev]

currently our tag_wrapper only supports os threads

[Bruno-DaSilva]

# Support async functions with decorator
# https://stackoverflow.com/a/63156433
def profile_tag_wrapper_async(*outer_args, **outer_kwargs):

    def inner_decorator(func):

        async def helper(func, *args, **kwargs):
            if asyncio.iscoroutinefunction(func):
                return await func(*args, **kwargs)

            return func(*args, **kwargs)

        @functools.wraps(func)
        async def wrapper(*args, **kwargs):
            with pyroscope.tag_wrapper(*outer_args, **outer_kwargs):
                result = await helper(func, *args, **kwargs)

            return result

        return wrapper

    return inner_decorator

Here's an implementation i wrote of a decorator that works for async functions. I'm sure it can be improved, but here it is in case its useful for implementing an async_tag_wrapper (or making the existing one sync compatible)

[korniltsev]

@Bruno-DaSilva Thanks for feedback, appreciate it

Well the problem is that pyroscope.tag_wrapper attaches tags to an os thread
And if I understand correctly, there maybe multiple corouines running on the same os threads at a time.
example:
Lets say we have two coroutines A, B running on OS Thread T

1. A is executing on thread T and enters tag_wrapper and sets tags controller="A"
2. A suspends
3. B is executing on thread T and enters tag_wrapper and sets tag controller="B"
4. B suspends
5. A resumes and burn some cpu

Now at step 5 we still have controller=B while the coroutine A is executing

Does it makes sense or am I missing something?
If my understanding is correct, then your decorator is likely incorrect for coroutines because of pyroscope.tag_wrapper implementaion

[Bruno-DaSilva]

ah, you're right! Interleaving on the same thread will lead to this behaviour, makes sense.

[caspervdw]

Would it be complex to use contextvars instead of threadlocals for tag_wrapper?

Context managers that have state should use Context Variables instead of threading.local() to prevent their state from bleeding to other code unexpectedly, when used in concurrent code.

Ref. https://docs.python.org/3/library/contextvars.html