3 """File-based locking helper."""
5 from __future__ import annotations
15 from dataclasses import dataclass
16 from typing import Literal, Optional
20 import decorator_utils
22 cfg = config.add_commandline_args(f'Lockfile ({__file__})', 'Args related to lockfiles')
24 '--lockfile_held_duration_warning_threshold_sec',
28 help='If a lock is held for longer than this threshold we log a warning',
30 logger = logging.getLogger(__name__)
33 class LockFileException(Exception):
34 """An exception related to lock files."""
40 class LockFileContents:
41 """The contents we'll write to each lock file."""
45 expiration_timestamp: Optional[float]
48 class LockFile(contextlib.AbstractContextManager):
49 """A file locking mechanism that has context-manager support so you
50 can use it in a with statement. e.g.
52 with LockFile('./foo.lock'):
53 # do a bunch of stuff... if the process dies we have a signal
54 # handler to do cleanup. Other code (in this process or another)
55 # that tries to take the same lockfile will block. There is also
56 # some logic for detecting stale locks.
64 do_signal_cleanup: bool = True,
65 expiration_timestamp: Optional[float] = None,
66 override_command: Optional[str] = None,
68 self.is_locked: bool = False
69 self.lockfile: str = lockfile_path
70 self.locktime: Optional[int] = None
71 self.override_command: Optional[str] = override_command
73 signal.signal(signal.SIGINT, self._signal)
74 signal.signal(signal.SIGTERM, self._signal)
75 self.expiration_timestamp = expiration_timestamp
81 return not os.path.exists(self.lockfile)
83 def try_acquire_lock_once(self) -> bool:
84 logger.debug("Trying to acquire %s.", self.lockfile)
86 # Attempt to create the lockfile. These flags cause
87 # os.open to raise an OSError if the file already
89 fd = os.open(self.lockfile, os.O_CREAT | os.O_EXCL | os.O_RDWR)
90 with os.fdopen(fd, "a") as f:
91 contents = self._get_lockfile_contents()
92 logger.debug(contents)
94 logger.debug('Success; I own %s.', self.lockfile)
99 logger.warning('Couldn\'t acquire %s.', self.lockfile)
102 def acquire_with_retries(
105 initial_delay: float = 1.0,
106 backoff_factor: float = 2.0,
109 @decorator_utils.retry_if_false(
110 tries=max_attempts, delay_sec=initial_delay, backoff=backoff_factor
112 def _try_acquire_lock_with_retries() -> bool:
113 success = self.try_acquire_lock_once()
114 if not success and os.path.exists(self.lockfile):
115 self._detect_stale_lockfile()
118 if os.path.exists(self.lockfile):
119 self._detect_stale_lockfile()
120 return _try_acquire_lock_with_retries()
124 os.unlink(self.lockfile)
125 except Exception as e:
127 self.is_locked = False
130 if self.acquire_with_retries():
131 self.locktime = datetime.datetime.now().timestamp()
133 msg = f"Couldn't acquire {self.lockfile}; giving up."
135 raise LockFileException(msg)
137 def __exit__(self, _, value, traceback) -> Literal[False]:
139 ts = datetime.datetime.now().timestamp()
140 duration = ts - self.locktime
141 if duration >= config.config['lockfile_held_duration_warning_threshold_sec']:
142 # Note: describe duration briefly only does 1s granularity...
143 str_duration = datetime_utils.describe_duration_briefly(int(duration))
144 msg = f'Held {self.lockfile} for {str_duration}'
146 warnings.warn(msg, stacklevel=2)
154 def _signal(self, *args):
158 def _get_lockfile_contents(self) -> str:
159 if self.override_command:
160 cmd = self.override_command
162 cmd = ' '.join(sys.argv)
163 contents = LockFileContents(
166 expiration_timestamp=self.expiration_timestamp,
168 return json.dumps(contents.__dict__)
170 def _detect_stale_lockfile(self) -> None:
172 with open(self.lockfile, 'r') as rf:
173 lines = rf.readlines()
176 line_dict = json.loads(line)
177 contents = LockFileContents(**line_dict)
178 logger.debug('Blocking lock contents="%s"', contents)
180 # Does the PID exist still?
182 os.kill(contents.pid, 0)
185 'Lockfile %s\'s pid (%d) is stale; force acquiring...',
191 # Has the lock expiration expired?
192 if contents.expiration_timestamp is not None:
193 now = datetime.datetime.now().timestamp()
194 if now > contents.expiration_timestamp:
196 'Lockfile %s\'s expiration time has passed; force acquiring',
201 pass # If the lockfile doesn't exist or disappears, good.