concurrency.py 1.5 KB

1234567891011121314151617181920212223242526272829303132333435363738394041
  1. from collections.abc import AsyncGenerator
  2. from contextlib import AbstractContextManager
  3. from contextlib import asynccontextmanager as asynccontextmanager
  4. from typing import TypeVar
  5. import anyio.to_thread
  6. from anyio import CapacityLimiter
  7. from starlette.concurrency import iterate_in_threadpool as iterate_in_threadpool # noqa
  8. from starlette.concurrency import run_in_threadpool as run_in_threadpool # noqa
  9. from starlette.concurrency import ( # noqa
  10. run_until_first_complete as run_until_first_complete,
  11. )
  12. _T = TypeVar("_T")
  13. @asynccontextmanager
  14. async def contextmanager_in_threadpool(
  15. cm: AbstractContextManager[_T],
  16. ) -> AsyncGenerator[_T, None]:
  17. # blocking __exit__ from running waiting on a free thread
  18. # can create race conditions/deadlocks if the context manager itself
  19. # has its own internal pool (e.g. a database connection pool)
  20. # to avoid this we let __exit__ run without a capacity limit
  21. # since we're creating a new limiter for each call, any non-zero limit
  22. # works (1 is arbitrary)
  23. exit_limiter = CapacityLimiter(1)
  24. try:
  25. yield await run_in_threadpool(cm.__enter__)
  26. except Exception as e:
  27. ok = bool(
  28. await anyio.to_thread.run_sync(
  29. cm.__exit__, type(e), e, e.__traceback__, limiter=exit_limiter
  30. )
  31. )
  32. if not ok:
  33. raise e
  34. else:
  35. await anyio.to_thread.run_sync(
  36. cm.__exit__, None, None, None, limiter=exit_limiter
  37. )