fix(async-io): instrumented asyncio.wait_for properly raises asyncio.TimeoutError (open-telemetry#2637)

Author: smoke

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+.

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)

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):