Repository: open-telemetry/opentelemetry-python

Issue:
Make `tracer.start_as_current_span()` decorator work with async functions
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()`.


Discussion:
Any news on this subject ?
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 🙂 
Same here, would you like to implement this, @aabmass?
I don't have time right now if someone else wants to take it
I've successfully make it works with a little bit of hacking :

```python
    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 :
```python
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 :

```bash
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
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](https://github.com/search?type=code&q=lang%3Apython+%2F%40.%2B%5C.start_as_current_span%5C%28.*%5C%29%5Cn.*async+def%2F) 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.
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](https://github.com/open-telemetry/opentelemetry-python/blob/975733c71473cddddd0859c6fcbd2b02405f7e12/opentelemetry-api/src/opentelemetry/trace/__init__.py#L269) 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`](https://github.com/python/cpython/blob/d63e4c49d3d1080ff53338fea3b37220161d1891/Lib/contextlib.py#L17)...
> However, I was wondering if [the API Tracer abstract class](https://github.com/open-telemetry/opentelemetry-python/blob/975733c71473cddddd0859c6fcbd2b02405f7e12/opentelemetry-api/src/opentelemetry/trace/__init__.py#L269) 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`](https://github.com/python/cpython/blob/d63e4c49d3d1080ff53338fea3b37220161d1891/Lib/contextlib.py#L17)...

I think it should be OK as long as the returned object matches the existing [public protocol of `@contextmanager`](https://github.com/python/typeshed/blob/5286722ef319b3931aa3e1af6b8948be2f6fa333/stdlib/contextlib.pyi#L59-L74) and we only need to extend the behavior of `__call__()`. Do you think that is feasible?

Resolving patch:
diff --git a/opentelemetry-api/src/opentelemetry/trace/__init__.py b/opentelemetry-api/src/opentelemetry/trace/__init__.py
--- a/opentelemetry-api/src/opentelemetry/trace/__init__.py
+++ b/opentelemetry-api/src/opentelemetry/trace/__init__.py
@@ -77,7 +77,6 @@
 import os
 import typing
 from abc import ABC, abstractmethod
-from contextlib import contextmanager
 from enum import Enum
 from logging import getLogger
 from typing import Iterator, Optional, Sequence, cast
@@ -109,6 +108,7 @@
 )
 from opentelemetry.trace.status import Status, StatusCode
 from opentelemetry.util import types
+from opentelemetry.util._decorator import _agnosticcontextmanager
 from opentelemetry.util._once import Once
 from opentelemetry.util._providers import _load_provider
 
@@ -324,7 +324,7 @@ def start_span(
             The newly-created span.
         """
 
-    @contextmanager
+    @_agnosticcontextmanager
     @abstractmethod
     def start_as_current_span(
         self,
@@ -431,7 +431,7 @@ def _tracer(self) -> Tracer:
     def start_span(self, *args, **kwargs) -> Span:  # type: ignore
         return self._tracer.start_span(*args, **kwargs)  # type: ignore
 
-    @contextmanager  # type: ignore
+    @_agnosticcontextmanager  # type: ignore
     def start_as_current_span(self, *args, **kwargs) -> Iterator[Span]:
         with self._tracer.start_as_current_span(*args, **kwargs) as span:  # type: ignore
             yield span
@@ -457,7 +457,7 @@ def start_span(
         # pylint: disable=unused-argument,no-self-use
         return INVALID_SPAN
 
-    @contextmanager
+    @_agnosticcontextmanager
     def start_as_current_span(
         self,
         name: str,
@@ -543,7 +543,7 @@ def get_tracer_provider() -> TracerProvider:
     return cast("TracerProvider", _TRACER_PROVIDER)
 
 
-@contextmanager
+@_agnosticcontextmanager
 def use_span(
     span: Span,
     end_on_exit: bool = False,
diff --git a/opentelemetry-api/src/opentelemetry/util/_decorator.py b/opentelemetry-api/src/opentelemetry/util/_decorator.py
new file mode 100644
--- /dev/null
+++ b/opentelemetry-api/src/opentelemetry/util/_decorator.py
@@ -0,0 +1,81 @@
+# Copyright The OpenTelemetry Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import asyncio
+import contextlib
+import functools
+import typing
+from typing import Callable, Generic, Iterator, TypeVar
+
+V = TypeVar("V")
+R = TypeVar("R")  # Return type
+Pargs = TypeVar("Pargs")  # Generic type for arguments
+Pkwargs = TypeVar("Pkwargs")  # Generic type for arguments
+
+if hasattr(typing, "ParamSpec"):
+    # only available in python 3.10+
+    # https://peps.python.org/pep-0612/
+    P = typing.ParamSpec("P")  # Generic type for all arguments
+
+
+class _AgnosticContextManager(
+    contextlib._GeneratorContextManager, Generic[R]  # type: ignore  # FIXME use contextlib._GeneratorContextManager[R] when we drop the python 3.8 support
+):  # pylint: disable=protected-access
+    """Context manager that can decorate both async and sync functions.
+
+    This is an overridden version of the contextlib._GeneratorContextManager
+    class that will decorate async functions with an async context manager
+    to end the span AFTER the entire async function coroutine finishes.
+
+    Else it will report near zero spans durations for async functions.
+
+    We are overriding the contextlib._GeneratorContextManager class as
+    reimplementing it is a lot of code to maintain and this class (even if it's
+    marked as protected) doesn't seems like to be evolving a lot.
+
+    For more information, see:
+    https://github.com/open-telemetry/opentelemetry-python/pull/3633
+    """
+
+    def __enter__(self) -> R:
+        """Reimplementing __enter__ to avoid the type error.
+
+        The original __enter__ method returns Any type, but we want to return R.
+        """
+        del self.args, self.kwds, self.func  # type: ignore
+        try:
+            return next(self.gen)  # type: ignore
+        except StopIteration:
+            raise RuntimeError("generator didn't yield") from None
+
+    def __call__(self, func: V) -> V:
+        if asyncio.iscoroutinefunction(func):
+
+            @functools.wraps(func)  # type: ignore
+            async def async_wrapper(*args: Pargs, **kwargs: Pkwargs) -> R:
+                with self._recreate_cm():  # type: ignore
+                    return await func(*args, **kwargs)  # type: ignore
+
+            return async_wrapper  # type: ignore
+        return super().__call__(func)  # type: ignore
+
+
+def _agnosticcontextmanager(
+    func: "Callable[P, Iterator[R]]",
+) -> "Callable[P, _AgnosticContextManager[R]]":
+    @functools.wraps(func)
+    def helper(*args: Pargs, **kwargs: Pkwargs) -> _AgnosticContextManager[R]:
+        return _AgnosticContextManager(func, args, kwargs)
+
+    return helper
diff --git a/opentelemetry-sdk/src/opentelemetry/sdk/trace/__init__.py b/opentelemetry-sdk/src/opentelemetry/sdk/trace/__init__.py
--- a/opentelemetry-sdk/src/opentelemetry/sdk/trace/__init__.py
+++ b/opentelemetry-sdk/src/opentelemetry/sdk/trace/__init__.py
@@ -21,7 +21,6 @@
 import threading
 import traceback
 import typing
-from contextlib import contextmanager
 from os import environ
 from time import time_ns
 from types import MappingProxyType, TracebackType
@@ -66,6 +65,7 @@
 from opentelemetry.trace import SpanContext
 from opentelemetry.trace.status import Status, StatusCode
 from opentelemetry.util import types
+from opentelemetry.util._decorator import _agnosticcontextmanager
 
 logger = logging.getLogger(__name__)
 
@@ -1038,7 +1038,7 @@ def __init__(
         self._span_limits = span_limits
         self._instrumentation_scope = instrumentation_scope
 
-    @contextmanager
+    @_agnosticcontextmanager  # pylint: disable=protected-access
     def start_as_current_span(
         self,
         name: str,


Test patch:
diff --git a/opentelemetry-api/tests/trace/test_proxy.py b/opentelemetry-api/tests/trace/test_proxy.py
--- a/opentelemetry-api/tests/trace/test_proxy.py
+++ b/opentelemetry-api/tests/trace/test_proxy.py
@@ -15,7 +15,6 @@
 # pylint: disable=W0212,W0222,W0221
 import typing
 import unittest
-from contextlib import contextmanager
 
 from opentelemetry import trace
 from opentelemetry.test.globals_test import TraceGlobalsTest
@@ -24,6 +23,7 @@
     NonRecordingSpan,
     Span,
 )
+from opentelemetry.util._decorator import _agnosticcontextmanager
 
 
 class TestProvider(trace.NoOpTracerProvider):
@@ -40,7 +40,7 @@ class TestTracer(trace.NoOpTracer):
     def start_span(self, *args, **kwargs):
         return TestSpan(INVALID_SPAN_CONTEXT)
 
-    @contextmanager
+    @_agnosticcontextmanager  # pylint: disable=protected-access
     def start_as_current_span(self, *args, **kwargs):  # type: ignore
         with trace.use_span(self.start_span(*args, **kwargs)) as span:  # type: ignore
             yield span
diff --git a/opentelemetry-api/tests/trace/test_tracer.py b/opentelemetry-api/tests/trace/test_tracer.py
--- a/opentelemetry-api/tests/trace/test_tracer.py
+++ b/opentelemetry-api/tests/trace/test_tracer.py
@@ -13,15 +13,15 @@
 # limitations under the License.
 
 
-from contextlib import contextmanager
+import asyncio
 from unittest import TestCase
-from unittest.mock import Mock
 
 from opentelemetry.trace import (
     INVALID_SPAN,
     NoOpTracer,
     Span,
     Tracer,
+    _agnosticcontextmanager,
     get_current_span,
 )
 
@@ -39,29 +39,42 @@ def test_start_as_current_span_context_manager(self):
             self.assertIsInstance(span, Span)
 
     def test_start_as_current_span_decorator(self):
-
-        mock_call = Mock()
+        # using a list to track the mock call order
+        calls = []
 
         class MockTracer(Tracer):
             def start_span(self, *args, **kwargs):
                 return INVALID_SPAN
 
-            @contextmanager
+            @_agnosticcontextmanager  # pylint: disable=protected-access
             def start_as_current_span(self, *args, **kwargs):  # type: ignore
-                mock_call()
+                calls.append(1)
                 yield INVALID_SPAN
+                calls.append(9)
 
         mock_tracer = MockTracer()
 
+        # test 1 : sync function
         @mock_tracer.start_as_current_span("name")
-        def function():  # type: ignore
-            pass
+        def function_sync(data: str) -> int:
+            calls.append(5)
+            return len(data)
 
-        function()  # type: ignore
-        function()  # type: ignore
-        function()  # type: ignore
+        calls = []
+        res = function_sync("123")
+        self.assertEqual(res, 3)
+        self.assertEqual(calls, [1, 5, 9])
 
-        self.assertEqual(mock_call.call_count, 3)
+        # test 2 : async function
+        @mock_tracer.start_as_current_span("name")
+        async def function_async(data: str) -> int:
+            calls.append(5)
+            return len(data)
+
+        calls = []
+        res = asyncio.run(function_async("123"))
+        self.assertEqual(res, 3)
+        self.assertEqual(calls, [1, 5, 9])
 
     def test_get_current_span(self):
         with self.tracer.start_as_current_span("test") as span:
diff --git a/opentelemetry-api/tests/util/test_contextmanager.py b/opentelemetry-api/tests/util/test_contextmanager.py
new file mode 100644
--- /dev/null
+++ b/opentelemetry-api/tests/util/test_contextmanager.py
@@ -0,0 +1,68 @@
+# Copyright The OpenTelemetry Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import asyncio
+import unittest
+from typing import Callable, Iterator
+
+from opentelemetry.util._decorator import _agnosticcontextmanager
+
+
+@_agnosticcontextmanager
+def cm() -> Iterator[int]:
+    yield 3
+
+
+@_agnosticcontextmanager
+def cm_call_when_done(f: Callable[[], None]) -> Iterator[int]:
+    yield 3
+    f()
+
+
+class TestContextManager(unittest.TestCase):
+    def test_sync_with(self):
+        with cm() as val:
+            self.assertEqual(val, 3)
+
+    def test_decorate_sync_func(self):
+        @cm()
+        def sync_func(a: str) -> str:
+            return a + a
+
+        res = sync_func("a")
+        self.assertEqual(res, "aa")
+
+    def test_decorate_async_func(self):
+        # Test that a universal context manager decorating an async function runs it's cleanup
+        # code after the entire async function coroutine finishes. This silently fails when
+        # using the normal @contextmanager decorator, which runs it's __exit__() after the
+        # un-started coroutine is returned.
+        #
+        # To see this behavior, change cm_call_when_done() to
+        # be decorated with @contextmanager.
+
+        events = []
+
+        @cm_call_when_done(lambda: events.append("cm_done"))
+        async def async_func(a: str) -> str:
+            events.append("start_async_func")
+            await asyncio.sleep(0)
+            events.append("finish_sleep")
+            return a + a
+
+        res = asyncio.run(async_func("a"))
+        self.assertEqual(res, "aa")
+        self.assertEqual(
+            events, ["start_async_func", "finish_sleep", "cm_done"]
+        )