Skip to content

Commit 5a6da15

Browse files
QuentinN42aabmassocelotl
authored
Make tracer.start_as_current_span() decorator work with async functions (open-telemetry#3633)
* test: add a minimal test that reproduce the bug Signed-off-by: QuentinN42 <quentin@lieumont.fr> * fix: replaced contextmanager by my own agnosticcontextmanager * docs: changelog Signed-off-by: QuentinN42 <quentin@lieumont.fr> * docs: typo Signed-off-by: QuentinN42 <quentin@lieumont.fr> * fix: linting Signed-off-by: QuentinN42 <quentin@lieumont.fr> * feat: reimplement the contextlib._GeneratorContextManager inside the trace api Signed-off-by: QuentinN42 <quentin@lieumont.fr> * refactor: merge tests Signed-off-by: QuentinN42 <quentin@lieumont.fr> * fix: make agnosticcontextmanager as protected Signed-off-by: QuentinN42 <quentin@lieumont.fr> * fix: changelog Signed-off-by: QuentinN42 <quentin@lieumont.fr> * feat: start typing Signed-off-by: QuentinN42 <quentin@lieumont.fr> * fix: revert to acdfa03 implementation and fix mypy Signed-off-by: QuentinN42 <quentin@lieumont.fr> * fix: use super call on the synchronous branch Co-authored-by: Aaron Abbott <aaronabbott@google.com> * fix: typing compliant with 3.8 Signed-off-by: QuentinN42 <quentin@lieumont.fr> * fix: black Signed-off-by: QuentinN42 <quentin@lieumont.fr> * feat: added tests and fixed lint in python 3.10 Signed-off-by: QuentinN42 <quentin@lieumont.fr> * feat: use_span use the _agnosticcontextmanager Signed-off-by: QuentinN42 <quentin@lieumont.fr> * docs: explain why we have an overriden class Signed-off-by: QuentinN42 <quentin@lieumont.fr> * fix: use typing.Generic for pre 3.9 compatibility Signed-off-by: QuentinN42 <quentin@lieumont.fr> * fix: typo Signed-off-by: QuentinN42 <quentin@lieumont.fr> * fix: mypy api/src Signed-off-by: QuentinN42 <quentin@lieumont.fr> * fix: ignore reference to privat attributes Signed-off-by: QuentinN42 <quentin@lieumont.fr> * fix: mv cm inside test Signed-off-by: QuentinN42 <quentin@lieumont.fr> * fix: define __call__ as Coroutine and not awaitable Signed-off-by: QuentinN42 <quentin@lieumont.fr> * fix: mypy green Signed-off-by: QuentinN42 <quentin@lieumont.fr> * fix: py38 tests ok Signed-off-by: QuentinN42 <quentin@lieumont.fr> * fix: reimplementing __enter__ to avoid the type error. Signed-off-by: QuentinN42 <quentin@lieumont.fr> * fix: lint Signed-off-by: QuentinN42 <quentin@lieumont.fr> * fix: mypy Signed-off-by: QuentinN42 <quentin@lieumont.fr> * test: rm test_wraps_contextlib Signed-off-by: QuentinN42 <quentin@lieumont.fr> * docs: document why we are overriding the contextlib._GeneratorContextManager class Signed-off-by: QuentinN42 <quentin@lieumont.fr> * docs: mv feat to unreleased section Signed-off-by: QuentinN42 <quentin@lieumont.fr> * test: rename lst with a more readable name Signed-off-by: QuentinN42 <quentin@lieumont.fr> * Remove unused type ignore comment * Fix missing symbol The missing symbol error was caused by a rebase on main and subsequent force push by me, sorry. --------- Signed-off-by: QuentinN42 <quentin@lieumont.fr> Co-authored-by: Aaron Abbott <aaronabbott@google.com> Co-authored-by: Diego Hurtado <ocelotl@users.noreply.github.com>
1 parent d6321d6 commit 5a6da15

File tree

8 files changed

+186
-22
lines changed

8 files changed

+186
-22
lines changed

‎.pylintrc

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -166,7 +166,7 @@ notes=FIXME,
166166
# List of decorators that produce context managers, such as
167167
# contextlib.contextmanager. Add to this list to register other decorators that
168168
# produce valid context managers.
169-
contextmanager-decorators=contextlib.contextmanager
169+
contextmanager-decorators=contextlib.contextmanager, _agnosticcontextmanager
170170

171171
# List of members which are set dynamically and missed by pylint inference
172172
# system, and so shouldn't trigger E1101 when accessed. Python regular

‎CHANGELOG.md

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
77

88
## Unreleased
99

10+
- Make `tracer.start_as_current_span()` decorator work with async functions
11+
([#3633](https://github.com/open-telemetry/opentelemetry-python/pull/3633))
1012
- Fix python 3.12 deprecation warning
1113
([#3751](https://github.com/open-telemetry/opentelemetry-python/pull/3751))
1214
- bump mypy to 0.982

‎opentelemetry-api/src/opentelemetry/trace/__init__.py

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -77,7 +77,6 @@
7777
import os
7878
import typing
7979
from abc import ABC, abstractmethod
80-
from contextlib import contextmanager
8180
from enum import Enum
8281
from logging import getLogger
8382
from typing import Iterator, Optional, Sequence, cast
@@ -109,6 +108,7 @@
109108
)
110109
from opentelemetry.trace.status import Status, StatusCode
111110
from opentelemetry.util import types
111+
from opentelemetry.util._decorator import _agnosticcontextmanager
112112
from opentelemetry.util._once import Once
113113
from opentelemetry.util._providers import _load_provider
114114

@@ -324,7 +324,7 @@ def start_span(
324324
The newly-created span.
325325
"""
326326

327-
@contextmanager
327+
@_agnosticcontextmanager
328328
@abstractmethod
329329
def start_as_current_span(
330330
self,
@@ -431,7 +431,7 @@ def _tracer(self) -> Tracer:
431431
def start_span(self, *args, **kwargs) -> Span: # type: ignore
432432
return self._tracer.start_span(*args, **kwargs) # type: ignore
433433

434-
@contextmanager # type: ignore
434+
@_agnosticcontextmanager # type: ignore
435435
def start_as_current_span(self, *args, **kwargs) -> Iterator[Span]:
436436
with self._tracer.start_as_current_span(*args, **kwargs) as span: # type: ignore
437437
yield span
@@ -457,7 +457,7 @@ def start_span(
457457
# pylint: disable=unused-argument,no-self-use
458458
return INVALID_SPAN
459459

460-
@contextmanager
460+
@_agnosticcontextmanager
461461
def start_as_current_span(
462462
self,
463463
name: str,
@@ -543,7 +543,7 @@ def get_tracer_provider() -> TracerProvider:
543543
return cast("TracerProvider", _TRACER_PROVIDER)
544544

545545

546-
@contextmanager
546+
@_agnosticcontextmanager
547547
def use_span(
548548
span: Span,
549549
end_on_exit: bool = False,

‎opentelemetry-api/src/opentelemetry/util/_decorator.py

Lines changed: 81 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,81 @@
1+
# Copyright The OpenTelemetry Authors
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
import asyncio
16+
import contextlib
17+
import functools
18+
import typing
19+
from typing import Callable, Generic, Iterator, TypeVar
20+
21+
V = TypeVar("V")
22+
R = TypeVar("R") # Return type
23+
Pargs = TypeVar("Pargs") # Generic type for arguments
24+
Pkwargs = TypeVar("Pkwargs") # Generic type for arguments
25+
26+
if hasattr(typing, "ParamSpec"):
27+
# only available in python 3.10+
28+
# https://peps.python.org/pep-0612/
29+
P = typing.ParamSpec("P") # Generic type for all arguments
30+
31+
32+
class _AgnosticContextManager(
33+
contextlib._GeneratorContextManager, Generic[R] # type: ignore # FIXME use contextlib._GeneratorContextManager[R] when we drop the python 3.8 support
34+
): # pylint: disable=protected-access
35+
"""Context manager that can decorate both async and sync functions.
36+
37+
This is an overridden version of the contextlib._GeneratorContextManager
38+
class that will decorate async functions with an async context manager
39+
to end the span AFTER the entire async function coroutine finishes.
40+
41+
Else it will report near zero spans durations for async functions.
42+
43+
We are overriding the contextlib._GeneratorContextManager class as
44+
reimplementing it is a lot of code to maintain and this class (even if it's
45+
marked as protected) doesn't seems like to be evolving a lot.
46+
47+
For more information, see:
48+
https://github.com/open-telemetry/opentelemetry-python/pull/3633
49+
"""
50+
51+
def __enter__(self) -> R:
52+
"""Reimplementing __enter__ to avoid the type error.
53+
54+
The original __enter__ method returns Any type, but we want to return R.
55+
"""
56+
del self.args, self.kwds, self.func # type: ignore
57+
try:
58+
return next(self.gen) # type: ignore
59+
except StopIteration:
60+
raise RuntimeError("generator didn't yield") from None
61+
62+
def __call__(self, func: V) -> V:
63+
if asyncio.iscoroutinefunction(func):
64+
65+
@functools.wraps(func) # type: ignore
66+
async def async_wrapper(*args: Pargs, **kwargs: Pkwargs) -> R:
67+
with self._recreate_cm(): # type: ignore
68+
return await func(*args, **kwargs) # type: ignore
69+
70+
return async_wrapper # type: ignore
71+
return super().__call__(func) # type: ignore
72+
73+
74+
def _agnosticcontextmanager(
75+
func: "Callable[P, Iterator[R]]",
76+
) -> "Callable[P, _AgnosticContextManager[R]]":
77+
@functools.wraps(func)
78+
def helper(*args: Pargs, **kwargs: Pkwargs) -> _AgnosticContextManager[R]:
79+
return _AgnosticContextManager(func, args, kwargs)
80+
81+
return helper

‎opentelemetry-api/tests/trace/test_proxy.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,6 @@
1515
# pylint: disable=W0212,W0222,W0221
1616
import typing
1717
import unittest
18-
from contextlib import contextmanager
1918

2019
from opentelemetry import trace
2120
from opentelemetry.test.globals_test import TraceGlobalsTest
@@ -24,6 +23,7 @@
2423
NonRecordingSpan,
2524
Span,
2625
)
26+
from opentelemetry.util._decorator import _agnosticcontextmanager
2727

2828

2929
class TestProvider(trace.NoOpTracerProvider):
@@ -40,7 +40,7 @@ class TestTracer(trace.NoOpTracer):
4040
def start_span(self, *args, **kwargs):
4141
return TestSpan(INVALID_SPAN_CONTEXT)
4242

43-
@contextmanager
43+
@_agnosticcontextmanager # pylint: disable=protected-access
4444
def start_as_current_span(self, *args, **kwargs): # type: ignore
4545
with trace.use_span(self.start_span(*args, **kwargs)) as span: # type: ignore
4646
yield span

‎opentelemetry-api/tests/trace/test_tracer.py

Lines changed: 25 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -13,15 +13,15 @@
1313
# limitations under the License.
1414

1515

16-
from contextlib import contextmanager
16+
import asyncio
1717
from unittest import TestCase
18-
from unittest.mock import Mock
1918

2019
from opentelemetry.trace import (
2120
INVALID_SPAN,
2221
NoOpTracer,
2322
Span,
2423
Tracer,
24+
_agnosticcontextmanager,
2525
get_current_span,
2626
)
2727

@@ -39,29 +39,42 @@ def test_start_as_current_span_context_manager(self):
3939
self.assertIsInstance(span, Span)
4040

4141
def test_start_as_current_span_decorator(self):
42-
43-
mock_call = Mock()
42+
# using a list to track the mock call order
43+
calls = []
4444

4545
class MockTracer(Tracer):
4646
def start_span(self, *args, **kwargs):
4747
return INVALID_SPAN
4848

49-
@contextmanager
49+
@_agnosticcontextmanager # pylint: disable=protected-access
5050
def start_as_current_span(self, *args, **kwargs): # type: ignore
51-
mock_call()
51+
calls.append(1)
5252
yield INVALID_SPAN
53+
calls.append(9)
5354

5455
mock_tracer = MockTracer()
5556

57+
# test 1 : sync function
5658
@mock_tracer.start_as_current_span("name")
57-
def function(): # type: ignore
58-
pass
59+
def function_sync(data: str) -> int:
60+
calls.append(5)
61+
return len(data)
5962

60-
function() # type: ignore
61-
function() # type: ignore
62-
function() # type: ignore
63+
calls = []
64+
res = function_sync("123")
65+
self.assertEqual(res, 3)
66+
self.assertEqual(calls, [1, 5, 9])
6367

64-
self.assertEqual(mock_call.call_count, 3)
68+
# test 2 : async function
69+
@mock_tracer.start_as_current_span("name")
70+
async def function_async(data: str) -> int:
71+
calls.append(5)
72+
return len(data)
73+
74+
calls = []
75+
res = asyncio.run(function_async("123"))
76+
self.assertEqual(res, 3)
77+
self.assertEqual(calls, [1, 5, 9])
6578

6679
def test_get_current_span(self):
6780
with self.tracer.start_as_current_span("test") as span:

‎opentelemetry-api/tests/util/test_contextmanager.py

Lines changed: 68 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,68 @@
1+
# Copyright The OpenTelemetry Authors
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
import asyncio
16+
import unittest
17+
from typing import Callable, Iterator
18+
19+
from opentelemetry.util._decorator import _agnosticcontextmanager
20+
21+
22+
@_agnosticcontextmanager
23+
def cm() -> Iterator[int]:
24+
yield 3
25+
26+
27+
@_agnosticcontextmanager
28+
def cm_call_when_done(f: Callable[[], None]) -> Iterator[int]:
29+
yield 3
30+
f()
31+
32+
33+
class TestContextManager(unittest.TestCase):
34+
def test_sync_with(self):
35+
with cm() as val:
36+
self.assertEqual(val, 3)
37+
38+
def test_decorate_sync_func(self):
39+
@cm()
40+
def sync_func(a: str) -> str:
41+
return a + a
42+
43+
res = sync_func("a")
44+
self.assertEqual(res, "aa")
45+
46+
def test_decorate_async_func(self):
47+
# Test that a universal context manager decorating an async function runs it's cleanup
48+
# code after the entire async function coroutine finishes. This silently fails when
49+
# using the normal @contextmanager decorator, which runs it's __exit__() after the
50+
# un-started coroutine is returned.
51+
#
52+
# To see this behavior, change cm_call_when_done() to
53+
# be decorated with @contextmanager.
54+
55+
events = []
56+
57+
@cm_call_when_done(lambda: events.append("cm_done"))
58+
async def async_func(a: str) -> str:
59+
events.append("start_async_func")
60+
await asyncio.sleep(0)
61+
events.append("finish_sleep")
62+
return a + a
63+
64+
res = asyncio.run(async_func("a"))
65+
self.assertEqual(res, "aa")
66+
self.assertEqual(
67+
events, ["start_async_func", "finish_sleep", "cm_done"]
68+
)

‎opentelemetry-sdk/src/opentelemetry/sdk/trace/__init__.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,6 @@
2121
import threading
2222
import traceback
2323
import typing
24-
from contextlib import contextmanager
2524
from os import environ
2625
from time import time_ns
2726
from types import MappingProxyType, TracebackType
@@ -66,6 +65,7 @@
6665
from opentelemetry.trace import SpanContext
6766
from opentelemetry.trace.status import Status, StatusCode
6867
from opentelemetry.util import types
68+
from opentelemetry.util._decorator import _agnosticcontextmanager
6969

7070
logger = logging.getLogger(__name__)
7171

@@ -1038,7 +1038,7 @@ def __init__(
10381038
self._span_limits = span_limits
10391039
self._instrumentation_scope = instrumentation_scope
10401040

1041-
@contextmanager
1041+
@_agnosticcontextmanager # pylint: disable=protected-access
10421042
def start_as_current_span(
10431043
self,
10441044
name: str,

0 commit comments

Comments
 (0)