file_baton.py 2.0 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263
  1. # mypy: allow-untyped-defs
  2. import os
  3. import time
  4. import warnings
  5. class FileBaton:
  6. """A primitive, file-based synchronization utility."""
  7. def __init__(self, lock_file_path, wait_seconds=0.1, warn_after_seconds=None):
  8. """
  9. Create a new :class:`FileBaton`.
  10. Args:
  11. lock_file_path: The path to the file used for locking.
  12. wait_seconds: The seconds to periodically sleep (spin) when
  13. calling ``wait()``.
  14. warn_after_seconds: The seconds to wait before showing
  15. lock file path to warn existing lock file.
  16. """
  17. self.lock_file_path = lock_file_path
  18. self.wait_seconds = wait_seconds
  19. self.fd = None
  20. self.warn_after_seconds = warn_after_seconds
  21. def try_acquire(self):
  22. """
  23. Try to atomically create a file under exclusive access.
  24. Returns:
  25. True if the file could be created, else False.
  26. """
  27. try:
  28. self.fd = os.open(self.lock_file_path, os.O_CREAT | os.O_EXCL)
  29. return True
  30. except FileExistsError:
  31. return False
  32. def wait(self):
  33. """
  34. Periodically sleeps for a certain amount until the baton is released.
  35. The amount of time slept depends on the ``wait_seconds`` parameter
  36. passed to the constructor.
  37. """
  38. has_warned = False
  39. start_time = time.time()
  40. while os.path.exists(self.lock_file_path):
  41. time.sleep(self.wait_seconds)
  42. if self.warn_after_seconds is not None:
  43. if time.time() - start_time > self.warn_after_seconds and not has_warned:
  44. warnings.warn(f'Waited on lock file "{self.lock_file_path}" for '
  45. f'{self.warn_after_seconds} seconds.')
  46. has_warned = True
  47. def release(self):
  48. """Release the baton and removes its file."""
  49. if self.fd is not None:
  50. os.close(self.fd)
  51. os.remove(self.lock_file_path)