1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
|
#!/usr/bin/env python3
"""File-based locking helper."""
from __future__ import annotations
import contextlib
import datetime
import json
import logging
import os
import signal
import sys
import warnings
from dataclasses import dataclass
from typing import Literal, Optional
import config
import datetime_utils
import decorator_utils
cfg = config.add_commandline_args(f'Lockfile ({__file__})', 'Args related to lockfiles')
cfg.add_argument(
'--lockfile_held_duration_warning_threshold_sec',
type=float,
default=60.0,
metavar='SECONDS',
help='If a lock is held for longer than this threshold we log a warning',
)
logger = logging.getLogger(__name__)
class LockFileException(Exception):
"""An exception related to lock files."""
pass
@dataclass
class LockFileContents:
"""The contents we'll write to each lock file."""
pid: int
commandline: str
expiration_timestamp: Optional[float]
class LockFile(contextlib.AbstractContextManager):
"""A file locking mechanism that has context-manager support so you
can use it in a with statement. e.g.
with LockFile('./foo.lock'):
# do a bunch of stuff... if the process dies we have a signal
# handler to do cleanup. Other code (in this process or another)
# that tries to take the same lockfile will block. There is also
# some logic for detecting stale locks.
"""
def __init__(
self,
lockfile_path: str,
*,
do_signal_cleanup: bool = True,
expiration_timestamp: Optional[float] = None,
override_command: Optional[str] = None,
) -> None:
self.is_locked: bool = False
self.lockfile: str = lockfile_path
self.locktime: Optional[int] = None
self.override_command: Optional[str] = override_command
if do_signal_cleanup:
signal.signal(signal.SIGINT, self._signal)
signal.signal(signal.SIGTERM, self._signal)
self.expiration_timestamp = expiration_timestamp
def locked(self):
return self.is_locked
def available(self):
return not os.path.exists(self.lockfile)
def try_acquire_lock_once(self) -> bool:
logger.debug("Trying to acquire %s.", self.lockfile)
try:
# Attempt to create the lockfile. These flags cause
# os.open to raise an OSError if the file already
# exists.
fd = os.open(self.lockfile, os.O_CREAT | os.O_EXCL | os.O_RDWR)
with os.fdopen(fd, "a") as f:
contents = self._get_lockfile_contents()
logger.debug(contents)
f.write(contents)
logger.debug('Success; I own %s.', self.lockfile)
self.is_locked = True
return True
except OSError:
pass
logger.warning('Couldn\'t acquire %s.', self.lockfile)
return False
def acquire_with_retries(
self,
*,
initial_delay: float = 1.0,
backoff_factor: float = 2.0,
max_attempts=5,
) -> bool:
@decorator_utils.retry_if_false(
tries=max_attempts, delay_sec=initial_delay, backoff=backoff_factor
)
def _try_acquire_lock_with_retries() -> bool:
success = self.try_acquire_lock_once()
if not success and os.path.exists(self.lockfile):
self._detect_stale_lockfile()
return success
if os.path.exists(self.lockfile):
self._detect_stale_lockfile()
return _try_acquire_lock_with_retries()
def release(self):
try:
os.unlink(self.lockfile)
except Exception as e:
logger.exception(e)
self.is_locked = False
def __enter__(self):
if self.acquire_with_retries():
self.locktime = datetime.datetime.now().timestamp()
return self
msg = f"Couldn't acquire {self.lockfile}; giving up."
logger.warning(msg)
raise LockFileException(msg)
def __exit__(self, _, value, traceback) -> Literal[False]:
if self.locktime:
ts = datetime.datetime.now().timestamp()
duration = ts - self.locktime
if duration >= config.config['lockfile_held_duration_warning_threshold_sec']:
# Note: describe duration briefly only does 1s granularity...
str_duration = datetime_utils.describe_duration_briefly(int(duration))
msg = f'Held {self.lockfile} for {str_duration}'
logger.warning(msg)
warnings.warn(msg, stacklevel=2)
self.release()
return False
def __del__(self):
if self.is_locked:
self.release()
def _signal(self, *args):
if self.is_locked:
self.release()
def _get_lockfile_contents(self) -> str:
if self.override_command:
cmd = self.override_command
else:
cmd = ' '.join(sys.argv)
contents = LockFileContents(
pid=os.getpid(),
commandline=cmd,
expiration_timestamp=self.expiration_timestamp,
)
return json.dumps(contents.__dict__)
def _detect_stale_lockfile(self) -> None:
try:
with open(self.lockfile, 'r') as rf:
lines = rf.readlines()
if len(lines) == 1:
line = lines[0]
line_dict = json.loads(line)
contents = LockFileContents(**line_dict)
logger.debug('Blocking lock contents="%s"', contents)
# Does the PID exist still?
try:
os.kill(contents.pid, 0)
except OSError:
logger.warning(
'Lockfile %s\'s pid (%d) is stale; force acquiring...',
self.lockfile,
contents.pid,
)
self.release()
# Has the lock expiration expired?
if contents.expiration_timestamp is not None:
now = datetime.datetime.now().timestamp()
if now > contents.expiration_timestamp:
logger.warning(
'Lockfile %s\'s expiration time has passed; force acquiring',
self.lockfile,
)
self.release()
except Exception:
pass # If the lockfile doesn't exist or disappears, good.
|
