From 094a08559233bed730cea6809ac657c69d3ac645 Mon Sep 17 00:00:00 2001
From: Radoslav Kirilov <rkirilow@gmail.com>
Date: Wed, 26 Jun 2024 15:54:58 +0300
Subject: [PATCH] fix(async-io): instrumented `asyncio.wait_for` properly
 raises `asyncio.TimeoutError`

@see https://docs.python.org/3.13/library/asyncio-task.html#asyncio.wait_for
---
 CHANGELOG.md                                        |  2 ++
 .../instrumentation/asyncio/__init__.py             |  3 +++
 .../tests/test_asyncio_wait.py                      | 13 +++++++++++++
 3 files changed, 18 insertions(+)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 5ddb3f488b..acccfba2e1 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -7,6 +7,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## Unreleased
 
+- `opentelemetry-instrumentation-asyncio` instrumented `asyncio.wait_for` properly raises `asyncio.TimeoutError` as expected
+  ([#2637](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/2637))
 - `opentelemetry-instrumentation-aws-lambda` Bugfix: AWS Lambda event source key incorrect for SNS in instrumentation library. 
   ([#2612](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/2612))
 - `opentelemetry-instrumentation-system-metrics` Permit to use psutil 6.0+.
diff --git a/instrumentation/opentelemetry-instrumentation-asyncio/src/opentelemetry/instrumentation/asyncio/__init__.py b/instrumentation/opentelemetry-instrumentation-asyncio/src/opentelemetry/instrumentation/asyncio/__init__.py
index fc1b535270..f5f0d34c4f 100644
--- a/instrumentation/opentelemetry-instrumentation-asyncio/src/opentelemetry/instrumentation/asyncio/__init__.py
+++ b/instrumentation/opentelemetry-instrumentation-asyncio/src/opentelemetry/instrumentation/asyncio/__init__.py
@@ -280,8 +280,11 @@ async def trace_coroutine(self, coro):
         # CancelledError is raised when a coroutine is cancelled
         # before it has a chance to run. We don't want to record
         # this as an error.
+        # Still it needs to be raised in order for `asyncio.wait_for`
+        # to properly work with timeout and raise accordingly `asyncio.TimeoutError`
         except asyncio.CancelledError:
             attr["state"] = "cancelled"
+            raise
         except Exception as exc:
             exception = exc
             state = determine_state(exception)
diff --git a/instrumentation/opentelemetry-instrumentation-asyncio/tests/test_asyncio_wait.py b/instrumentation/opentelemetry-instrumentation-asyncio/tests/test_asyncio_wait.py
index 77064aeafa..312f035d36 100644
--- a/instrumentation/opentelemetry-instrumentation-asyncio/tests/test_asyncio_wait.py
+++ b/instrumentation/opentelemetry-instrumentation-asyncio/tests/test_asyncio_wait.py
@@ -68,6 +68,19 @@ async def main():
         spans = self.memory_exporter.get_finished_spans()
         self.assertEqual(len(spans), 2)
 
+    def test_asyncio_wait_for_with_timeout(self):
+        expected_timeout_error = None
+
+        async def main():
+            nonlocal expected_timeout_error
+            try:
+                await asyncio.wait_for(async_func(), 0.01)
+            except asyncio.TimeoutError as timeout_error:
+                expected_timeout_error = timeout_error
+
+        asyncio.run(main())
+        self.assertNotEqual(expected_timeout_error, None)
+
     def test_asyncio_as_completed(self):
         async def main():
             if sys.version_info >= (3, 11):