-
Notifications
You must be signed in to change notification settings - Fork 805
asyncio: fix duplicate instrumentation #3408
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Merged
Merged
Changes from 6 commits
Commits
Show all changes
16 commits
Select commit
Hold shift + click to select a range
53f05e0
#3383 fix duplicate instrument
allen-k1m 6acef29
Merge branch 'main' into bugfig/memory-leak
bourbonkk 964c861
#3383 fix duplicate instrument
bourbonkk a102e47
Merge remote-tracking branch 'origin/bugfig/memory-leak' into bugfig/…
allen-k1m da3d617
feedback
allen-k1m dfd7713
feat(asyncio): add weakref-based tracking for instrumented objects
allen-k1m aba1f57
Use WeakKeyDictionary to safely track instrumented objects
allen-k1m 2678389
Merge branch 'main' into bugfig/memory-leak
bourbonkk 5413f60
feedback
allen-k1m c13e3c4
Merge remote-tracking branch 'origin/bugfig/memory-leak' into bugfig/…
allen-k1m 787e31e
feedback
allen-k1m e48d655
feedback
allen-k1m deee6f9
Merge branch 'main' into bugfig/memory-leak
bourbonkk e620541
Merge branch 'main' into bugfig/memory-leak
bourbonkk 5099ba4
Merge branch 'main' into bugfig/memory-leak
bourbonkk 893d5fd
Update instrumentation/opentelemetry-instrumentation-asyncio/src/open…
aabmass File filter
Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
74 changes: 74 additions & 0 deletions
74
...nstrumentation-asyncio/src/opentelemetry/instrumentation/asyncio/instrumentation_state.py
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,74 @@ | ||
| """ | ||
bourbonkk marked this conversation as resolved.
Show resolved
Hide resolved
|
||
| Instrumentation State Tracker | ||
| This module provides helper functions to safely track whether a coroutine, | ||
| Future, or function has already been instrumented by the OpenTelemetry | ||
| asyncio instrumentation layer. | ||
| Because some Python objects (like coroutines or functions) may not support | ||
| adding custom attributes or may not be weak-referenceable, we use a | ||
| weak-reference-based dictionary to track instrumented objects safely and | ||
| efficiently without causing memory leaks. | ||
| If an object cannot be weak-referenced, we skip tracking it to avoid | ||
| runtime errors. | ||
| Usage: | ||
| if not _is_instrumented(obj): | ||
| _mark_instrumented(obj) | ||
| # instrument the object... | ||
| """ | ||
|
|
||
| import weakref | ||
|
|
||
| # A global dictionary to track whether an object has been instrumented. | ||
| # Keys are weak references to avoid preventing garbage collection. | ||
| _instrumented_tasks = {} | ||
bourbonkk marked this conversation as resolved.
Outdated
Show resolved
Hide resolved
|
||
|
|
||
|
|
||
| def _get_weak_key(obj): | ||
bourbonkk marked this conversation as resolved.
Outdated
Show resolved
Hide resolved
|
||
| """ | ||
| Attempt to create a weak reference key for the given object. | ||
| Some object types (e.g., built-in functions or async_generator_asend) | ||
| do not support weak references. In those cases, return None. | ||
| Args: | ||
| obj: The object to generate a weak reference for. | ||
| Returns: | ||
| A weakref.ref to the object if supported, otherwise None. | ||
| """ | ||
| try: | ||
| return weakref.ref(obj) | ||
| except TypeError: | ||
aabmass marked this conversation as resolved.
Outdated
Show resolved
Hide resolved
|
||
| return None | ||
|
|
||
|
|
||
| def _is_instrumented(obj) -> bool: | ||
| """ | ||
| Check if the object has already been instrumented. | ||
| Args: | ||
| obj: The coroutine, function, or Future to check. | ||
| Returns: | ||
| True if the object is already marked as instrumented, False otherwise. | ||
| """ | ||
| key = _get_weak_key(obj) | ||
| return key in _instrumented_tasks if key else False | ||
|
|
||
|
|
||
| def _mark_instrumented(obj): | ||
| """ | ||
| Mark the object as instrumented to avoid double-instrumentation. | ||
| Only objects that support weak references are tracked. Unsupported | ||
| objects are silently skipped. | ||
| Args: | ||
| obj: The coroutine, function, or Future to mark. | ||
| """ | ||
| key = _get_weak_key(obj) | ||
| if key: | ||
| _instrumented_tasks[key] = True | ||
115 changes: 115 additions & 0 deletions
115
...entation/opentelemetry-instrumentation-asyncio/tests/test_asyncio_duplicate_instrument.py
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,115 @@ | ||
| # 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. | ||
|
|
||
| """ | ||
| A general test verifying that when the same Future objects (or coroutines) are | ||
| repeatedly instrumented (for example, via `trace_future`), callback references | ||
| do not leak. In this example, we mimic a typical scenario where a small set of | ||
| Futures might be reused throughout an application's lifecycle. | ||
| """ | ||
|
|
||
| import asyncio | ||
|
|
||
| from opentelemetry.instrumentation.asyncio import AsyncioInstrumentor | ||
| from opentelemetry.test.test_base import TestBase | ||
|
|
||
|
|
||
| class MockSubscription: | ||
| """ | ||
| Example class holding an unsubscribe_future, similar to something like | ||
| aiokafka's subscription. | ||
| """ | ||
|
|
||
| def __init__(self): | ||
| self.unsubscribe_future = asyncio.Future() | ||
|
|
||
|
|
||
| class MockGroupCoordinator: | ||
| """ | ||
| Example class modeling repeated instrumentation of the same Future objects. | ||
| """ | ||
|
|
||
| def __init__(self): | ||
| self._closing = asyncio.Future() | ||
| self.subscription = MockSubscription() | ||
| self._rejoin_needed_fut = asyncio.Future() | ||
|
|
||
| async def run_routine(self, instrumentor): | ||
| """ | ||
| Each time this routine is called, the same 3 Futures are 'traced' again. | ||
| In a real-life scenario, there's often a loop reusing these objects. | ||
| """ | ||
| instrumentor.trace_future(self._closing) | ||
| instrumentor.trace_future(self.subscription.unsubscribe_future) | ||
| instrumentor.trace_future(self._rejoin_needed_fut) | ||
|
|
||
|
|
||
| class TestAsyncioDuplicateInstrument(TestBase): | ||
| """ | ||
| Tests whether repeated instrumentation of the same Futures leads to | ||
| exponential callback growth (potential memory leak). | ||
| """ | ||
|
|
||
| def setUp(self): | ||
| super().setUp() | ||
| self.loop = asyncio.new_event_loop() | ||
| asyncio.set_event_loop(self.loop) | ||
|
|
||
| self.instrumentor = AsyncioInstrumentor() | ||
| self.instrumentor.instrument() | ||
|
|
||
| def tearDown(self): | ||
| self.instrumentor.uninstrument() | ||
| self.loop.close() | ||
| asyncio.set_event_loop(None) | ||
| super().tearDown() | ||
|
|
||
| def test_duplicate_instrumentation_of_futures(self): | ||
| """ | ||
| If instrumentor.trace_future is called multiple times on the same Future, | ||
| we should NOT see an unbounded accumulation of callbacks. | ||
| """ | ||
| coordinator = MockGroupCoordinator() | ||
bourbonkk marked this conversation as resolved.
Outdated
Show resolved
Hide resolved
|
||
|
|
||
| # Simulate calling the routine multiple times | ||
| num_iterations = 10 | ||
| for _ in range(num_iterations): | ||
| self.loop.run_until_complete( | ||
| coordinator.run_routine(self.instrumentor) | ||
| ) | ||
|
|
||
| # Check for callback accumulation | ||
| closing_cb_count = len(coordinator._closing._callbacks) | ||
| unsub_cb_count = len( | ||
| coordinator.subscription.unsubscribe_future._callbacks | ||
| ) | ||
| rejoin_cb_count = len(coordinator._rejoin_needed_fut._callbacks) | ||
|
|
||
| # If instrumentation is properly deduplicated, each Future might have ~1-2 callbacks. | ||
| max_expected_callbacks = 2 | ||
| self.assertLessEqual( | ||
bourbonkk marked this conversation as resolved.
Outdated
Show resolved
Hide resolved
|
||
| closing_cb_count, | ||
| max_expected_callbacks, | ||
| f"_closing Future has {closing_cb_count} callbacks. Potential leak!", | ||
| ) | ||
| self.assertLessEqual( | ||
| unsub_cb_count, | ||
| max_expected_callbacks, | ||
| f"unsubscribe_future has {unsub_cb_count} callbacks. Potential leak!", | ||
| ) | ||
| self.assertLessEqual( | ||
| rejoin_cb_count, | ||
| max_expected_callbacks, | ||
| f"_rejoin_needed_fut has {rejoin_cb_count} callbacks. Potential leak!", | ||
| ) | ||
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Uh oh!
There was an error while loading. Please reload this page.