asyncio.py 12 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342
  1. """An asyncio-based implementation of the file lock.""" # noqa: A005
  2. from __future__ import annotations
  3. import asyncio
  4. import contextlib
  5. import logging
  6. import os
  7. import time
  8. from dataclasses import dataclass
  9. from threading import local
  10. from typing import TYPE_CHECKING, Any, Callable, NoReturn, cast
  11. from ._api import BaseFileLock, FileLockContext, FileLockMeta
  12. from ._error import Timeout
  13. from ._soft import SoftFileLock
  14. from ._unix import UnixFileLock
  15. from ._windows import WindowsFileLock
  16. if TYPE_CHECKING:
  17. import sys
  18. from concurrent import futures
  19. from types import TracebackType
  20. if sys.version_info >= (3, 11): # pragma: no cover (py311+)
  21. from typing import Self
  22. else: # pragma: no cover (<py311)
  23. from typing_extensions import Self
  24. _LOGGER = logging.getLogger("filelock")
  25. @dataclass
  26. class AsyncFileLockContext(FileLockContext):
  27. """A dataclass which holds the context for a ``BaseAsyncFileLock`` object."""
  28. #: Whether run in executor
  29. run_in_executor: bool = True
  30. #: The executor
  31. executor: futures.Executor | None = None
  32. #: The loop
  33. loop: asyncio.AbstractEventLoop | None = None
  34. class AsyncThreadLocalFileContext(AsyncFileLockContext, local):
  35. """A thread local version of the ``FileLockContext`` class."""
  36. class AsyncAcquireReturnProxy:
  37. """A context-aware object that will release the lock file when exiting."""
  38. def __init__(self, lock: BaseAsyncFileLock) -> None: # noqa: D107
  39. self.lock = lock
  40. async def __aenter__(self) -> BaseAsyncFileLock: # noqa: D105
  41. return self.lock
  42. async def __aexit__( # noqa: D105
  43. self,
  44. exc_type: type[BaseException] | None,
  45. exc_value: BaseException | None,
  46. traceback: TracebackType | None,
  47. ) -> None:
  48. await self.lock.release()
  49. class AsyncFileLockMeta(FileLockMeta):
  50. def __call__( # type: ignore[override] # noqa: PLR0913
  51. cls, # noqa: N805
  52. lock_file: str | os.PathLike[str],
  53. timeout: float = -1,
  54. mode: int = 0o644,
  55. thread_local: bool = False, # noqa: FBT001, FBT002
  56. *,
  57. blocking: bool = True,
  58. is_singleton: bool = False,
  59. loop: asyncio.AbstractEventLoop | None = None,
  60. run_in_executor: bool = True,
  61. executor: futures.Executor | None = None,
  62. ) -> BaseAsyncFileLock:
  63. if thread_local and run_in_executor:
  64. msg = "run_in_executor is not supported when thread_local is True"
  65. raise ValueError(msg)
  66. instance = super().__call__(
  67. lock_file=lock_file,
  68. timeout=timeout,
  69. mode=mode,
  70. thread_local=thread_local,
  71. blocking=blocking,
  72. is_singleton=is_singleton,
  73. loop=loop,
  74. run_in_executor=run_in_executor,
  75. executor=executor,
  76. )
  77. return cast(BaseAsyncFileLock, instance)
  78. class BaseAsyncFileLock(BaseFileLock, metaclass=AsyncFileLockMeta):
  79. """Base class for asynchronous file locks."""
  80. def __init__( # noqa: PLR0913
  81. self,
  82. lock_file: str | os.PathLike[str],
  83. timeout: float = -1,
  84. mode: int = 0o644,
  85. thread_local: bool = False, # noqa: FBT001, FBT002
  86. *,
  87. blocking: bool = True,
  88. is_singleton: bool = False,
  89. loop: asyncio.AbstractEventLoop | None = None,
  90. run_in_executor: bool = True,
  91. executor: futures.Executor | None = None,
  92. ) -> None:
  93. """
  94. Create a new lock object.
  95. :param lock_file: path to the file
  96. :param timeout: default timeout when acquiring the lock, in seconds. It will be used as fallback value in \
  97. the acquire method, if no timeout value (``None``) is given. If you want to disable the timeout, set it \
  98. to a negative value. A timeout of 0 means that there is exactly one attempt to acquire the file lock.
  99. :param mode: file permissions for the lockfile
  100. :param thread_local: Whether this object's internal context should be thread local or not. If this is set to \
  101. ``False`` then the lock will be reentrant across threads.
  102. :param blocking: whether the lock should be blocking or not
  103. :param is_singleton: If this is set to ``True`` then only one instance of this class will be created \
  104. per lock file. This is useful if you want to use the lock object for reentrant locking without needing \
  105. to pass the same object around.
  106. :param loop: The event loop to use. If not specified, the running event loop will be used.
  107. :param run_in_executor: If this is set to ``True`` then the lock will be acquired in an executor.
  108. :param executor: The executor to use. If not specified, the default executor will be used.
  109. """
  110. self._is_thread_local = thread_local
  111. self._is_singleton = is_singleton
  112. # Create the context. Note that external code should not work with the context directly and should instead use
  113. # properties of this class.
  114. kwargs: dict[str, Any] = {
  115. "lock_file": os.fspath(lock_file),
  116. "timeout": timeout,
  117. "mode": mode,
  118. "blocking": blocking,
  119. "loop": loop,
  120. "run_in_executor": run_in_executor,
  121. "executor": executor,
  122. }
  123. self._context: AsyncFileLockContext = (AsyncThreadLocalFileContext if thread_local else AsyncFileLockContext)(
  124. **kwargs
  125. )
  126. @property
  127. def run_in_executor(self) -> bool:
  128. """::return: whether run in executor."""
  129. return self._context.run_in_executor
  130. @property
  131. def executor(self) -> futures.Executor | None:
  132. """::return: the executor."""
  133. return self._context.executor
  134. @executor.setter
  135. def executor(self, value: futures.Executor | None) -> None: # pragma: no cover
  136. """
  137. Change the executor.
  138. :param value: the new executor or ``None``
  139. :type value: futures.Executor | None
  140. """
  141. self._context.executor = value
  142. @property
  143. def loop(self) -> asyncio.AbstractEventLoop | None:
  144. """::return: the event loop."""
  145. return self._context.loop
  146. async def acquire( # type: ignore[override]
  147. self,
  148. timeout: float | None = None,
  149. poll_interval: float = 0.05,
  150. *,
  151. blocking: bool | None = None,
  152. ) -> AsyncAcquireReturnProxy:
  153. """
  154. Try to acquire the file lock.
  155. :param timeout: maximum wait time for acquiring the lock, ``None`` means use the default
  156. :attr:`~BaseFileLock.timeout` is and if ``timeout < 0``, there is no timeout and
  157. this method will block until the lock could be acquired
  158. :param poll_interval: interval of trying to acquire the lock file
  159. :param blocking: defaults to True. If False, function will return immediately if it cannot obtain a lock on the
  160. first attempt. Otherwise, this method will block until the timeout expires or the lock is acquired.
  161. :raises Timeout: if fails to acquire lock within the timeout period
  162. :return: a context object that will unlock the file when the context is exited
  163. .. code-block:: python
  164. # You can use this method in the context manager (recommended)
  165. with lock.acquire():
  166. pass
  167. # Or use an equivalent try-finally construct:
  168. lock.acquire()
  169. try:
  170. pass
  171. finally:
  172. lock.release()
  173. """
  174. # Use the default timeout, if no timeout is provided.
  175. if timeout is None:
  176. timeout = self._context.timeout
  177. if blocking is None:
  178. blocking = self._context.blocking
  179. # Increment the number right at the beginning. We can still undo it, if something fails.
  180. self._context.lock_counter += 1
  181. lock_id = id(self)
  182. lock_filename = self.lock_file
  183. start_time = time.perf_counter()
  184. try:
  185. while True:
  186. if not self.is_locked:
  187. _LOGGER.debug("Attempting to acquire lock %s on %s", lock_id, lock_filename)
  188. await self._run_internal_method(self._acquire)
  189. if self.is_locked:
  190. _LOGGER.debug("Lock %s acquired on %s", lock_id, lock_filename)
  191. break
  192. if blocking is False:
  193. _LOGGER.debug("Failed to immediately acquire lock %s on %s", lock_id, lock_filename)
  194. raise Timeout(lock_filename) # noqa: TRY301
  195. if 0 <= timeout < time.perf_counter() - start_time:
  196. _LOGGER.debug("Timeout on acquiring lock %s on %s", lock_id, lock_filename)
  197. raise Timeout(lock_filename) # noqa: TRY301
  198. msg = "Lock %s not acquired on %s, waiting %s seconds ..."
  199. _LOGGER.debug(msg, lock_id, lock_filename, poll_interval)
  200. await asyncio.sleep(poll_interval)
  201. except BaseException: # Something did go wrong, so decrement the counter.
  202. self._context.lock_counter = max(0, self._context.lock_counter - 1)
  203. raise
  204. return AsyncAcquireReturnProxy(lock=self)
  205. async def release(self, force: bool = False) -> None: # type: ignore[override] # noqa: FBT001, FBT002
  206. """
  207. Releases the file lock. Please note, that the lock is only completely released, if the lock counter is 0.
  208. Also note, that the lock file itself is not automatically deleted.
  209. :param force: If true, the lock counter is ignored and the lock is released in every case/
  210. """
  211. if self.is_locked:
  212. self._context.lock_counter -= 1
  213. if self._context.lock_counter == 0 or force:
  214. lock_id, lock_filename = id(self), self.lock_file
  215. _LOGGER.debug("Attempting to release lock %s on %s", lock_id, lock_filename)
  216. await self._run_internal_method(self._release)
  217. self._context.lock_counter = 0
  218. _LOGGER.debug("Lock %s released on %s", lock_id, lock_filename)
  219. async def _run_internal_method(self, method: Callable[[], Any]) -> None:
  220. if asyncio.iscoroutinefunction(method):
  221. await method()
  222. elif self.run_in_executor:
  223. loop = self.loop or asyncio.get_running_loop()
  224. await loop.run_in_executor(self.executor, method)
  225. else:
  226. method()
  227. def __enter__(self) -> NoReturn:
  228. """
  229. Replace old __enter__ method to avoid using it.
  230. NOTE: DO NOT USE `with` FOR ASYNCIO LOCKS, USE `async with` INSTEAD.
  231. :return: none
  232. :rtype: NoReturn
  233. """
  234. msg = "Do not use `with` for asyncio locks, use `async with` instead."
  235. raise NotImplementedError(msg)
  236. async def __aenter__(self) -> Self:
  237. """
  238. Acquire the lock.
  239. :return: the lock object
  240. """
  241. await self.acquire()
  242. return self
  243. async def __aexit__(
  244. self,
  245. exc_type: type[BaseException] | None,
  246. exc_value: BaseException | None,
  247. traceback: TracebackType | None,
  248. ) -> None:
  249. """
  250. Release the lock.
  251. :param exc_type: the exception type if raised
  252. :param exc_value: the exception value if raised
  253. :param traceback: the exception traceback if raised
  254. """
  255. await self.release()
  256. def __del__(self) -> None:
  257. """Called when the lock object is deleted."""
  258. with contextlib.suppress(RuntimeError):
  259. loop = self.loop or asyncio.get_running_loop()
  260. if not loop.is_running(): # pragma: no cover
  261. loop.run_until_complete(self.release(force=True))
  262. else:
  263. loop.create_task(self.release(force=True))
  264. class AsyncSoftFileLock(SoftFileLock, BaseAsyncFileLock):
  265. """Simply watches the existence of the lock file."""
  266. class AsyncUnixFileLock(UnixFileLock, BaseAsyncFileLock):
  267. """Uses the :func:`fcntl.flock` to hard lock the lock file on unix systems."""
  268. class AsyncWindowsFileLock(WindowsFileLock, BaseAsyncFileLock):
  269. """Uses the :func:`msvcrt.locking` to hard lock the lock file on windows systems."""
  270. __all__ = [
  271. "AsyncAcquireReturnProxy",
  272. "AsyncSoftFileLock",
  273. "AsyncUnixFileLock",
  274. "AsyncWindowsFileLock",
  275. "BaseAsyncFileLock",
  276. ]