_synchronization.py 21 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772
  1. from __future__ import annotations
  2. import math
  3. from collections import deque
  4. from collections.abc import Callable
  5. from dataclasses import dataclass
  6. from types import TracebackType
  7. from typing import TypeVar
  8. from ..lowlevel import checkpoint_if_cancelled
  9. from ._eventloop import get_async_backend
  10. from ._exceptions import BusyResourceError, NoEventLoopError
  11. from ._tasks import CancelScope
  12. from ._testing import TaskInfo, get_current_task
  13. T = TypeVar("T")
  14. @dataclass(frozen=True)
  15. class EventStatistics:
  16. """
  17. :ivar int tasks_waiting: number of tasks waiting on :meth:`~.Event.wait`
  18. """
  19. tasks_waiting: int
  20. @dataclass(frozen=True)
  21. class CapacityLimiterStatistics:
  22. """
  23. :ivar int borrowed_tokens: number of tokens currently borrowed by tasks
  24. :ivar float total_tokens: total number of available tokens
  25. :ivar tuple borrowers: tasks or other objects currently holding tokens borrowed from
  26. this limiter
  27. :ivar int tasks_waiting: number of tasks waiting on
  28. :meth:`~.CapacityLimiter.acquire` or
  29. :meth:`~.CapacityLimiter.acquire_on_behalf_of`
  30. """
  31. borrowed_tokens: int
  32. total_tokens: float
  33. borrowers: tuple[object, ...]
  34. tasks_waiting: int
  35. @dataclass(frozen=True)
  36. class LockStatistics:
  37. """
  38. :ivar bool locked: flag indicating if this lock is locked or not
  39. :ivar ~anyio.TaskInfo owner: task currently holding the lock (or ``None`` if the
  40. lock is not held by any task)
  41. :ivar int tasks_waiting: number of tasks waiting on :meth:`~.Lock.acquire`
  42. """
  43. locked: bool
  44. owner: TaskInfo | None
  45. tasks_waiting: int
  46. @dataclass(frozen=True)
  47. class ConditionStatistics:
  48. """
  49. :ivar int tasks_waiting: number of tasks blocked on :meth:`~.Condition.wait`
  50. :ivar ~anyio.LockStatistics lock_statistics: statistics of the underlying
  51. :class:`~.Lock`
  52. """
  53. tasks_waiting: int
  54. lock_statistics: LockStatistics
  55. @dataclass(frozen=True)
  56. class SemaphoreStatistics:
  57. """
  58. :ivar int tasks_waiting: number of tasks waiting on :meth:`~.Semaphore.acquire`
  59. """
  60. tasks_waiting: int
  61. class Event:
  62. __slots__ = ("__weakref__",)
  63. def __new__(cls) -> Event:
  64. try:
  65. return get_async_backend().create_event()
  66. except NoEventLoopError:
  67. return EventAdapter()
  68. def set(self) -> None:
  69. """Set the flag, notifying all listeners."""
  70. raise NotImplementedError
  71. def is_set(self) -> bool:
  72. """Return ``True`` if the flag is set, ``False`` if not."""
  73. raise NotImplementedError
  74. async def wait(self) -> None:
  75. """
  76. Wait until the flag has been set.
  77. If the flag has already been set when this method is called, it returns
  78. immediately.
  79. """
  80. raise NotImplementedError
  81. def statistics(self) -> EventStatistics:
  82. """Return statistics about the current state of this event."""
  83. raise NotImplementedError
  84. class EventAdapter(Event):
  85. __slots__ = "_internal_event", "_is_set"
  86. def __new__(cls) -> EventAdapter:
  87. return object.__new__(cls)
  88. def __init__(self) -> None:
  89. self._internal_event: Event | None = None
  90. self._is_set = False
  91. @property
  92. def _event(self) -> Event:
  93. if self._internal_event is None:
  94. self._internal_event = get_async_backend().create_event()
  95. if self._is_set:
  96. self._internal_event.set()
  97. return self._internal_event
  98. def set(self) -> None:
  99. if self._internal_event is None:
  100. self._is_set = True
  101. else:
  102. self._event.set()
  103. def is_set(self) -> bool:
  104. if self._internal_event is None:
  105. return self._is_set
  106. return self._internal_event.is_set()
  107. async def wait(self) -> None:
  108. await self._event.wait()
  109. def statistics(self) -> EventStatistics:
  110. if self._internal_event is None:
  111. return EventStatistics(tasks_waiting=0)
  112. return self._internal_event.statistics()
  113. class Lock:
  114. __slots__ = ("__weakref__",)
  115. def __new__(cls, *, fast_acquire: bool = False) -> Lock:
  116. try:
  117. return get_async_backend().create_lock(fast_acquire=fast_acquire)
  118. except NoEventLoopError:
  119. return LockAdapter(fast_acquire=fast_acquire)
  120. async def __aenter__(self) -> None:
  121. await self.acquire()
  122. async def __aexit__(
  123. self,
  124. exc_type: type[BaseException] | None,
  125. exc_val: BaseException | None,
  126. exc_tb: TracebackType | None,
  127. ) -> None:
  128. self.release()
  129. async def acquire(self) -> None:
  130. """Acquire the lock."""
  131. raise NotImplementedError
  132. def acquire_nowait(self) -> None:
  133. """
  134. Acquire the lock, without blocking.
  135. :raises ~anyio.WouldBlock: if the operation would block
  136. """
  137. raise NotImplementedError
  138. def release(self) -> None:
  139. """Release the lock."""
  140. raise NotImplementedError
  141. def locked(self) -> bool:
  142. """Return True if the lock is currently held."""
  143. raise NotImplementedError
  144. def statistics(self) -> LockStatistics:
  145. """
  146. Return statistics about the current state of this lock.
  147. .. versionadded:: 3.0
  148. """
  149. raise NotImplementedError
  150. class LockAdapter(Lock):
  151. __slots__ = "_internal_lock", "_fast_acquire"
  152. def __new__(cls, *, fast_acquire: bool = False) -> LockAdapter:
  153. return object.__new__(cls)
  154. def __init__(self, *, fast_acquire: bool = False):
  155. self._internal_lock: Lock | None = None
  156. self._fast_acquire = fast_acquire
  157. @property
  158. def _lock(self) -> Lock:
  159. if self._internal_lock is None:
  160. self._internal_lock = get_async_backend().create_lock(
  161. fast_acquire=self._fast_acquire
  162. )
  163. return self._internal_lock
  164. async def __aenter__(self) -> None:
  165. await self._lock.acquire()
  166. async def __aexit__(
  167. self,
  168. exc_type: type[BaseException] | None,
  169. exc_val: BaseException | None,
  170. exc_tb: TracebackType | None,
  171. ) -> None:
  172. if self._internal_lock is not None:
  173. self._internal_lock.release()
  174. async def acquire(self) -> None:
  175. """Acquire the lock."""
  176. await self._lock.acquire()
  177. def acquire_nowait(self) -> None:
  178. """
  179. Acquire the lock, without blocking.
  180. :raises ~anyio.WouldBlock: if the operation would block
  181. """
  182. self._lock.acquire_nowait()
  183. def release(self) -> None:
  184. """Release the lock."""
  185. self._lock.release()
  186. def locked(self) -> bool:
  187. """Return True if the lock is currently held."""
  188. return self._lock.locked()
  189. def statistics(self) -> LockStatistics:
  190. """
  191. Return statistics about the current state of this lock.
  192. .. versionadded:: 3.0
  193. """
  194. if self._internal_lock is None:
  195. return LockStatistics(False, None, 0)
  196. return self._internal_lock.statistics()
  197. class Condition:
  198. __slots__ = "__weakref__", "_owner_task", "_lock", "_waiters"
  199. def __init__(self, lock: Lock | None = None):
  200. self._owner_task: TaskInfo | None = None
  201. self._lock = lock or Lock()
  202. self._waiters: deque[Event] = deque()
  203. async def __aenter__(self) -> None:
  204. await self.acquire()
  205. async def __aexit__(
  206. self,
  207. exc_type: type[BaseException] | None,
  208. exc_val: BaseException | None,
  209. exc_tb: TracebackType | None,
  210. ) -> None:
  211. self.release()
  212. def _check_acquired(self) -> None:
  213. if self._owner_task != get_current_task():
  214. raise RuntimeError("The current task is not holding the underlying lock")
  215. async def acquire(self) -> None:
  216. """Acquire the underlying lock."""
  217. await self._lock.acquire()
  218. self._owner_task = get_current_task()
  219. def acquire_nowait(self) -> None:
  220. """
  221. Acquire the underlying lock, without blocking.
  222. :raises ~anyio.WouldBlock: if the operation would block
  223. """
  224. self._lock.acquire_nowait()
  225. self._owner_task = get_current_task()
  226. def release(self) -> None:
  227. """Release the underlying lock."""
  228. self._lock.release()
  229. def locked(self) -> bool:
  230. """Return True if the lock is set."""
  231. return self._lock.locked()
  232. def notify(self, n: int = 1) -> None:
  233. """Notify exactly n listeners."""
  234. self._check_acquired()
  235. for _ in range(n):
  236. try:
  237. event = self._waiters.popleft()
  238. except IndexError:
  239. break
  240. event.set()
  241. def notify_all(self) -> None:
  242. """Notify all the listeners."""
  243. self._check_acquired()
  244. for event in self._waiters:
  245. event.set()
  246. self._waiters.clear()
  247. async def wait(self) -> None:
  248. """Wait for a notification."""
  249. await checkpoint_if_cancelled()
  250. self._check_acquired()
  251. event = Event()
  252. self._waiters.append(event)
  253. self.release()
  254. try:
  255. await event.wait()
  256. except BaseException:
  257. if not event.is_set():
  258. self._waiters.remove(event)
  259. elif self._waiters:
  260. # This task was notified by could not act on it, so pass
  261. # it on to the next task
  262. self._waiters.popleft().set()
  263. raise
  264. finally:
  265. with CancelScope(shield=True):
  266. await self.acquire()
  267. async def wait_for(self, predicate: Callable[[], T]) -> T:
  268. """
  269. Wait until a predicate becomes true.
  270. :param predicate: a callable that returns a truthy value when the condition is
  271. met
  272. :return: the result of the predicate
  273. .. versionadded:: 4.11.0
  274. """
  275. while not (result := predicate()):
  276. await self.wait()
  277. return result
  278. def statistics(self) -> ConditionStatistics:
  279. """
  280. Return statistics about the current state of this condition.
  281. .. versionadded:: 3.0
  282. """
  283. return ConditionStatistics(len(self._waiters), self._lock.statistics())
  284. class Semaphore:
  285. __slots__ = "__weakref__", "_fast_acquire"
  286. def __new__(
  287. cls,
  288. initial_value: int,
  289. *,
  290. max_value: int | None = None,
  291. fast_acquire: bool = False,
  292. ) -> Semaphore:
  293. try:
  294. return get_async_backend().create_semaphore(
  295. initial_value, max_value=max_value, fast_acquire=fast_acquire
  296. )
  297. except NoEventLoopError:
  298. return SemaphoreAdapter(initial_value, max_value=max_value)
  299. def __init__(
  300. self,
  301. initial_value: int,
  302. *,
  303. max_value: int | None = None,
  304. fast_acquire: bool = False,
  305. ):
  306. if not isinstance(initial_value, int):
  307. raise TypeError("initial_value must be an integer")
  308. if initial_value < 0:
  309. raise ValueError("initial_value must be >= 0")
  310. if max_value is not None:
  311. if not isinstance(max_value, int):
  312. raise TypeError("max_value must be an integer or None")
  313. if max_value < initial_value:
  314. raise ValueError(
  315. "max_value must be equal to or higher than initial_value"
  316. )
  317. self._fast_acquire = fast_acquire
  318. async def __aenter__(self) -> Semaphore:
  319. await self.acquire()
  320. return self
  321. async def __aexit__(
  322. self,
  323. exc_type: type[BaseException] | None,
  324. exc_val: BaseException | None,
  325. exc_tb: TracebackType | None,
  326. ) -> None:
  327. self.release()
  328. async def acquire(self) -> None:
  329. """Decrement the semaphore value, blocking if necessary."""
  330. raise NotImplementedError
  331. def acquire_nowait(self) -> None:
  332. """
  333. Acquire the underlying lock, without blocking.
  334. :raises ~anyio.WouldBlock: if the operation would block
  335. """
  336. raise NotImplementedError
  337. def release(self) -> None:
  338. """Increment the semaphore value."""
  339. raise NotImplementedError
  340. @property
  341. def value(self) -> int:
  342. """The current value of the semaphore."""
  343. raise NotImplementedError
  344. @property
  345. def max_value(self) -> int | None:
  346. """The maximum value of the semaphore."""
  347. raise NotImplementedError
  348. def statistics(self) -> SemaphoreStatistics:
  349. """
  350. Return statistics about the current state of this semaphore.
  351. .. versionadded:: 3.0
  352. """
  353. raise NotImplementedError
  354. class SemaphoreAdapter(Semaphore):
  355. __slots__ = "_internal_semaphore", "_initial_value", "_max_value"
  356. def __new__(
  357. cls,
  358. initial_value: int,
  359. *,
  360. max_value: int | None = None,
  361. fast_acquire: bool = False,
  362. ) -> SemaphoreAdapter:
  363. return object.__new__(cls)
  364. def __init__(
  365. self,
  366. initial_value: int,
  367. *,
  368. max_value: int | None = None,
  369. fast_acquire: bool = False,
  370. ) -> None:
  371. super().__init__(initial_value, max_value=max_value, fast_acquire=fast_acquire)
  372. self._internal_semaphore: Semaphore | None = None
  373. self._initial_value = initial_value
  374. self._max_value = max_value
  375. @property
  376. def _semaphore(self) -> Semaphore:
  377. if self._internal_semaphore is None:
  378. self._internal_semaphore = get_async_backend().create_semaphore(
  379. self._initial_value, max_value=self._max_value
  380. )
  381. return self._internal_semaphore
  382. async def acquire(self) -> None:
  383. await self._semaphore.acquire()
  384. def acquire_nowait(self) -> None:
  385. self._semaphore.acquire_nowait()
  386. def release(self) -> None:
  387. self._semaphore.release()
  388. @property
  389. def value(self) -> int:
  390. if self._internal_semaphore is None:
  391. return self._initial_value
  392. return self._semaphore.value
  393. @property
  394. def max_value(self) -> int | None:
  395. return self._max_value
  396. def statistics(self) -> SemaphoreStatistics:
  397. if self._internal_semaphore is None:
  398. return SemaphoreStatistics(tasks_waiting=0)
  399. return self._semaphore.statistics()
  400. class CapacityLimiter:
  401. __slots__ = ("__weakref__",)
  402. def __new__(cls, total_tokens: float) -> CapacityLimiter:
  403. try:
  404. return get_async_backend().create_capacity_limiter(total_tokens)
  405. except NoEventLoopError:
  406. return CapacityLimiterAdapter(total_tokens)
  407. async def __aenter__(self) -> None:
  408. raise NotImplementedError
  409. async def __aexit__(
  410. self,
  411. exc_type: type[BaseException] | None,
  412. exc_val: BaseException | None,
  413. exc_tb: TracebackType | None,
  414. ) -> None:
  415. raise NotImplementedError
  416. @property
  417. def total_tokens(self) -> float:
  418. """
  419. The total number of tokens available for borrowing.
  420. This is a read-write property. If the total number of tokens is increased, the
  421. proportionate number of tasks waiting on this limiter will be granted their
  422. tokens.
  423. .. versionchanged:: 3.0
  424. The property is now writable.
  425. .. versionchanged:: 4.12
  426. The value can now be set to 0.
  427. """
  428. raise NotImplementedError
  429. @total_tokens.setter
  430. def total_tokens(self, value: float) -> None:
  431. raise NotImplementedError
  432. @property
  433. def borrowed_tokens(self) -> int:
  434. """The number of tokens that have currently been borrowed."""
  435. raise NotImplementedError
  436. @property
  437. def available_tokens(self) -> float:
  438. """The number of tokens currently available to be borrowed"""
  439. raise NotImplementedError
  440. def acquire_nowait(self) -> None:
  441. """
  442. Acquire a token for the current task without waiting for one to become
  443. available.
  444. :raises ~anyio.WouldBlock: if there are no tokens available for borrowing
  445. """
  446. raise NotImplementedError
  447. def acquire_on_behalf_of_nowait(self, borrower: object) -> None:
  448. """
  449. Acquire a token without waiting for one to become available.
  450. :param borrower: the entity borrowing a token
  451. :raises ~anyio.WouldBlock: if there are no tokens available for borrowing
  452. """
  453. raise NotImplementedError
  454. async def acquire(self) -> None:
  455. """
  456. Acquire a token for the current task, waiting if necessary for one to become
  457. available.
  458. """
  459. raise NotImplementedError
  460. async def acquire_on_behalf_of(self, borrower: object) -> None:
  461. """
  462. Acquire a token, waiting if necessary for one to become available.
  463. :param borrower: the entity borrowing a token
  464. """
  465. raise NotImplementedError
  466. def release(self) -> None:
  467. """
  468. Release the token held by the current task.
  469. :raises RuntimeError: if the current task has not borrowed a token from this
  470. limiter.
  471. """
  472. raise NotImplementedError
  473. def release_on_behalf_of(self, borrower: object) -> None:
  474. """
  475. Release the token held by the given borrower.
  476. :raises RuntimeError: if the borrower has not borrowed a token from this
  477. limiter.
  478. """
  479. raise NotImplementedError
  480. def statistics(self) -> CapacityLimiterStatistics:
  481. """
  482. Return statistics about the current state of this limiter.
  483. .. versionadded:: 3.0
  484. """
  485. raise NotImplementedError
  486. class CapacityLimiterAdapter(CapacityLimiter):
  487. __slots__ = "_internal_limiter", "_total_tokens"
  488. def __new__(cls, total_tokens: float) -> CapacityLimiterAdapter:
  489. return object.__new__(cls)
  490. def __init__(self, total_tokens: float) -> None:
  491. self._internal_limiter: CapacityLimiter | None = None
  492. self.total_tokens = total_tokens
  493. @property
  494. def _limiter(self) -> CapacityLimiter:
  495. if self._internal_limiter is None:
  496. self._internal_limiter = get_async_backend().create_capacity_limiter(
  497. self._total_tokens
  498. )
  499. return self._internal_limiter
  500. async def __aenter__(self) -> None:
  501. await self._limiter.__aenter__()
  502. async def __aexit__(
  503. self,
  504. exc_type: type[BaseException] | None,
  505. exc_val: BaseException | None,
  506. exc_tb: TracebackType | None,
  507. ) -> None:
  508. return await self._limiter.__aexit__(exc_type, exc_val, exc_tb)
  509. @property
  510. def total_tokens(self) -> float:
  511. if self._internal_limiter is None:
  512. return self._total_tokens
  513. return self._internal_limiter.total_tokens
  514. @total_tokens.setter
  515. def total_tokens(self, value: float) -> None:
  516. if not isinstance(value, int) and not math.isinf(value):
  517. raise TypeError("total_tokens must be an int or math.inf")
  518. elif value < 0:
  519. raise ValueError("total_tokens must be >= 0")
  520. if self._internal_limiter is None:
  521. self._total_tokens = value
  522. return
  523. self._limiter.total_tokens = value
  524. @property
  525. def borrowed_tokens(self) -> int:
  526. if self._internal_limiter is None:
  527. return 0
  528. return self._internal_limiter.borrowed_tokens
  529. @property
  530. def available_tokens(self) -> float:
  531. if self._internal_limiter is None:
  532. return self._total_tokens
  533. return self._internal_limiter.available_tokens
  534. def acquire_nowait(self) -> None:
  535. self._limiter.acquire_nowait()
  536. def acquire_on_behalf_of_nowait(self, borrower: object) -> None:
  537. self._limiter.acquire_on_behalf_of_nowait(borrower)
  538. async def acquire(self) -> None:
  539. await self._limiter.acquire()
  540. async def acquire_on_behalf_of(self, borrower: object) -> None:
  541. await self._limiter.acquire_on_behalf_of(borrower)
  542. def release(self) -> None:
  543. self._limiter.release()
  544. def release_on_behalf_of(self, borrower: object) -> None:
  545. self._limiter.release_on_behalf_of(borrower)
  546. def statistics(self) -> CapacityLimiterStatistics:
  547. if self._internal_limiter is None:
  548. return CapacityLimiterStatistics(
  549. borrowed_tokens=0,
  550. total_tokens=self.total_tokens,
  551. borrowers=(),
  552. tasks_waiting=0,
  553. )
  554. return self._internal_limiter.statistics()
  555. class ResourceGuard:
  556. """
  557. A context manager for ensuring that a resource is only used by a single task at a
  558. time.
  559. Entering this context manager while the previous has not exited it yet will trigger
  560. :exc:`BusyResourceError`.
  561. :param action: the action to guard against (visible in the :exc:`BusyResourceError`
  562. when triggered, e.g. "Another task is already {action} this resource")
  563. .. versionadded:: 4.1
  564. """
  565. __slots__ = "__weakref__", "action", "_guarded"
  566. def __init__(self, action: str = "using"):
  567. self.action: str = action
  568. self._guarded = False
  569. def __enter__(self) -> None:
  570. if self._guarded:
  571. raise BusyResourceError(self.action)
  572. self._guarded = True
  573. def __exit__(
  574. self,
  575. exc_type: type[BaseException] | None,
  576. exc_val: BaseException | None,
  577. exc_tb: TracebackType | None,
  578. ) -> None:
  579. self._guarded = False