3 from dataclasses import dataclass
10 from typing import Optional
12 import decorator_utils
15 logger = logging.getLogger(__name__)
18 class LockFileException(Exception):
23 class LockFileContents:
26 expiration_timestamp: float
29 class LockFile(object):
30 """A file locking mechanism that has context-manager support so you
31 can use it in a with statement. e.g.
33 with LockFile('./foo.lock'):
34 # do a bunch of stuff... if the process dies we have a signal
35 # handler to do cleanup. Other code (in this process or another)
36 # that tries to take the same lockfile will block. There is also
37 # some logic for detecting stale locks.
44 do_signal_cleanup: bool = True,
45 expiration_timestamp: Optional[float] = None,
46 override_command: Optional[str] = None,
48 self.is_locked = False
49 self.lockfile = lockfile_path
50 self.override_command = override_command
52 signal.signal(signal.SIGINT, self._signal)
53 signal.signal(signal.SIGTERM, self._signal)
54 self.expiration_timestamp = expiration_timestamp
60 return not os.path.exists(self.lockfile)
62 def try_acquire_lock_once(self) -> bool:
63 logger.debug(f"Trying to acquire {self.lockfile}.")
65 # Attempt to create the lockfile. These flags cause
66 # os.open to raise an OSError if the file already
68 fd = os.open(self.lockfile, os.O_CREAT | os.O_EXCL | os.O_RDWR)
69 with os.fdopen(fd, "a") as f:
70 contents = self._get_lockfile_contents()
71 logger.debug(contents)
73 logger.debug(f'Success; I own {self.lockfile}.')
78 logger.debug(f'Failed; I could not acquire {self.lockfile}.')
81 def acquire_with_retries(
84 initial_delay: float = 1.0,
85 backoff_factor: float = 2.0,
89 @decorator_utils.retry_if_false(tries = max_attempts,
90 delay_sec = initial_delay,
91 backoff = backoff_factor)
92 def _try_acquire_lock_with_retries() -> bool:
93 success = self.try_acquire_lock_once()
94 if not success and os.path.exists(self.lockfile):
95 self._detect_stale_lockfile()
98 if os.path.exists(self.lockfile):
99 self._detect_stale_lockfile()
100 return _try_acquire_lock_with_retries()
104 os.unlink(self.lockfile)
105 except Exception as e:
107 self.is_locked = False
110 if self.acquire_with_retries():
112 msg = f"Couldn't acquire {self.lockfile}; giving up."
114 raise LockFileException(msg)
116 def __exit__(self, type, value, traceback):
123 def _signal(self, *args):
127 def _get_lockfile_contents(self) -> str:
128 if self.override_command:
129 cmd = self.override_command
131 cmd = ' '.join(sys.argv)
133 contents = LockFileContents(
136 expiration_timestamp = self.expiration_timestamp,
138 return json.dumps(contents.__dict__)
140 def _detect_stale_lockfile(self) -> None:
142 with open(self.lockfile, 'r') as rf:
143 lines = rf.readlines()
146 line_dict = json.loads(line)
147 contents = LockFileContents(**line_dict)
148 logger.debug(f'Blocking lock contents="{contents}"')
150 # Does the PID exist still?
152 os.kill(contents.pid, 0)
154 logger.debug('The pid seems stale; killing the lock.')
157 # Has the lock expiration expired?
158 if contents.expiration_timestamp is not None:
159 now = datetime.datetime.now().timestamp()
160 if now > contents.expiration_datetime:
161 logger.debug('The expiration time has passed; ' +