From 050eb371761a605952745e6eec3821ec5acd8751 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

cherry-picked from https://github.com/open-telemetry/opentelemetry-python-contrib/pull/2637
---
 CHANGELOG.md                                        |  7 +++++++
 .../instrumentation/asyncio/__init__.py             |  3 +++
 .../tests/test_asyncio_wait.py                      | 13 +++++++++++++
 3 files changed, 23 insertions(+)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 5184e4c9b4..0c16073b27 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -5,6 +5,13 @@ All notable changes to this project will be documented in this file.
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
+## Version 1.25.0/0.46b0+asyncio-wait-for-fix (2024-06-26)
+
+### Fixed
+
+- `opentelemetry-instrumentation-asyncio` instrumented `asyncio.wait_for` properly raises `asyncio.TimeoutError` as expected
+  ([#2637](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/2637))
+
 ## Version 1.25.0/0.46b0 (2024-05-30)
 
 ### Breaking changes
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):