Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.11/site-packages/redis/lock.py: 24%

1import logging 

2import threading 

3import time as mod_time 

4import uuid 

5from types import SimpleNamespace, TracebackType 

6from typing import Literal, Optional, Type 

7 

8from redis.exceptions import LockError, LockNotOwnedError 

9from redis.typing import Number 

10 

11logger = logging.getLogger(__name__) 

12 

13 

14class Lock: 

15 """ 

16 A shared, distributed Lock. Using Redis for locking allows the Lock 

17 to be shared across processes and/or machines. 

18 

19 It's left to the user to resolve deadlock issues and make sure 

20 multiple clients play nicely together. 

21 """ 

22 

23 lua_release = None 

24 lua_extend = None 

25 lua_reacquire = None 

26 

27 # KEYS[1] - lock name 

28 # ARGV[1] - token 

29 # return 1 if the lock was released, otherwise 0 

30 LUA_RELEASE_SCRIPT = """ 

31 local token = redis.call('get', KEYS[1]) 

32 if not token or token ~= ARGV[1] then 

33 return 0 

34 end 

35 redis.call('del', KEYS[1]) 

36 return 1 

37 """ 

38 

39 # KEYS[1] - lock name 

40 # ARGV[1] - token 

41 # ARGV[2] - additional milliseconds 

42 # ARGV[3] - "0" if the additional time should be added to the lock's 

43 # existing ttl or "1" if the existing ttl should be replaced 

44 # return 1 if the locks time was extended, otherwise 0 

45 LUA_EXTEND_SCRIPT = """ 

46 local token = redis.call('get', KEYS[1]) 

47 if not token or token ~= ARGV[1] then 

48 return 0 

49 end 

50 local expiration = redis.call('pttl', KEYS[1]) 

51 if not expiration then 

52 expiration = 0 

53 end 

54 if expiration < 0 then 

55 return 0 

56 end 

57 

58 local newttl = ARGV[2] 

59 if ARGV[3] == "0" then 

60 newttl = ARGV[2] + expiration 

61 end 

62 if tonumber(newttl) <= 0 then 

63 return 1 

64 end 

65 redis.call('pexpire', KEYS[1], newttl) 

66 return 1 

67 """ 

68 

69 # KEYS[1] - lock name 

70 # ARGV[1] - token 

71 # ARGV[2] - milliseconds 

72 # return 1 if the locks time was reacquired, otherwise 0 

73 LUA_REACQUIRE_SCRIPT = """ 

74 local token = redis.call('get', KEYS[1]) 

75 if not token or token ~= ARGV[1] then 

76 return 0 

77 end 

78 if tonumber(ARGV[2]) > 0 then 

79 redis.call('pexpire', KEYS[1], ARGV[2]) 

80 end 

81 return 1 

82 """ 

83 

84 def __init__( 

85 self, 

86 redis, 

87 name: str, 

88 timeout: Optional[Number] = None, 

89 sleep: Number = 0.1, 

90 blocking: bool = True, 

91 blocking_timeout: Optional[Number] = None, 

92 thread_local: bool = True, 

93 raise_on_release_error: bool = True, 

94 ): 

95 """ 

96 Create a new Lock instance named ``name`` using the Redis client 

97 supplied by ``redis``. 

98 

99 ``timeout`` indicates a maximum life for the lock in seconds. 

100 By default, it will remain locked until release() is called. 

101 ``timeout`` can be specified as a float or integer, both representing 

102 the number of seconds to wait. 

103 

104 ``sleep`` indicates the amount of time to sleep in seconds per loop 

105 iteration when the lock is in blocking mode and another client is 

106 currently holding the lock. 

107 

108 ``blocking`` indicates whether calling ``acquire`` should block until 

109 the lock has been acquired or to fail immediately, causing ``acquire`` 

110 to return False and the lock not being acquired. Defaults to True. 

111 Note this value can be overridden by passing a ``blocking`` 

112 argument to ``acquire``. 

113 

114 ``blocking_timeout`` indicates the maximum amount of time in seconds to 

115 spend trying to acquire the lock. A value of ``None`` indicates 

116 continue trying forever. ``blocking_timeout`` can be specified as a 

117 float or integer, both representing the number of seconds to wait. 

118 

119 ``thread_local`` indicates whether the lock token is placed in 

120 thread-local storage. By default, the token is placed in thread local 

121 storage so that a thread only sees its token, not a token set by 

122 another thread. Consider the following timeline: 

123 

124 time: 0, thread-1 acquires `my-lock`, with a timeout of 5 seconds. 

125 thread-1 sets the token to "abc" 

126 time: 1, thread-2 blocks trying to acquire `my-lock` using the 

127 Lock instance. 

128 time: 5, thread-1 has not yet completed. redis expires the lock 

129 key. 

130 time: 5, thread-2 acquired `my-lock` now that it's available. 

131 thread-2 sets the token to "xyz" 

132 time: 6, thread-1 finishes its work and calls release(). if the 

133 token is *not* stored in thread local storage, then 

134 thread-1 would see the token value as "xyz" and would be 

135 able to successfully release the thread-2's lock. 

136 

137 ``raise_on_release_error`` indicates whether to raise an exception when 

138 the lock is no longer owned when exiting the context manager. By default, 

139 this is True, meaning an exception will be raised. If False, the warning 

140 will be logged and the exception will be suppressed. 

141 

142 In some use cases it's necessary to disable thread local storage. For 

143 example, if you have code where one thread acquires a lock and passes 

144 that lock instance to a worker thread to release later. If thread 

145 local storage isn't disabled in this case, the worker thread won't see 

146 the token set by the thread that acquired the lock. Our assumption 

147 is that these cases aren't common and as such default to using 

148 thread local storage. 

149 """ 

150 self.redis = redis 

151 self.name = name 

152 self.timeout = timeout 

153 self.sleep = sleep 

154 self.blocking = blocking 

155 self.blocking_timeout = blocking_timeout 

156 self.thread_local = bool(thread_local) 

157 self.raise_on_release_error = raise_on_release_error 

158 self.local = threading.local() if self.thread_local else SimpleNamespace() 

159 self.local.token = None 

160 self.register_scripts() 

161 

162 def register_scripts(self) -> None: 

163 cls = self.__class__ 

164 client = self.redis 

165 if cls.lua_release is None: 

166 cls.lua_release = client.register_script(cls.LUA_RELEASE_SCRIPT) 

167 if cls.lua_extend is None: 

168 cls.lua_extend = client.register_script(cls.LUA_EXTEND_SCRIPT) 

169 if cls.lua_reacquire is None: 

170 cls.lua_reacquire = client.register_script(cls.LUA_REACQUIRE_SCRIPT) 

171 

172 def __enter__(self) -> "Lock": 

173 if self.acquire(): 

174 return self 

175 raise LockError( 

176 "Unable to acquire lock within the time specified", 

177 lock_name=self.name, 

178 ) 

179 

180 def __exit__( 

181 self, 

182 exc_type: Optional[Type[BaseException]], 

183 exc_value: Optional[BaseException], 

184 traceback: Optional[TracebackType], 

185 ) -> None: 

186 try: 

187 self.release() 

188 except LockError: 

189 if self.raise_on_release_error: 

190 raise 

191 logger.warning( 

192 "Lock was unlocked or no longer owned when exiting context manager." 

193 ) 

194 

195 def acquire( 

196 self, 

197 sleep: Optional[Number] = None, 

198 blocking: Optional[bool] = None, 

199 blocking_timeout: Optional[Number] = None, 

200 token: Optional[str] = None, 

201 ): 

202 """ 

203 Use Redis to hold a shared, distributed lock named ``name``. 

204 Returns True once the lock is acquired. 

205 

206 If ``blocking`` is False, always return immediately. If the lock 

207 was acquired, return True, otherwise return False. 

208 

209 ``blocking_timeout`` specifies the maximum number of seconds to 

210 wait trying to acquire the lock. 

211 

212 ``token`` specifies the token value to be used. If provided, token 

213 must be a bytes object or a string that can be encoded to a bytes 

214 object with the default encoding. If a token isn't specified, a UUID 

215 will be generated. 

216 """ 

217 if sleep is None: 

218 sleep = self.sleep 

219 if token is None: 

220 token = uuid.uuid1().hex.encode() 

221 else: 

222 encoder = self.redis.get_encoder() 

223 token = encoder.encode(token) 

224 if blocking is None: 

225 blocking = self.blocking 

226 if blocking_timeout is None: 

227 blocking_timeout = self.blocking_timeout 

228 stop_trying_at = None 

229 if blocking_timeout is not None: 

230 stop_trying_at = mod_time.monotonic() + blocking_timeout 

231 while True: 

232 if self.do_acquire(token): 

233 self.local.token = token 

234 return True 

235 if not blocking: 

236 return False 

237 next_try_at = mod_time.monotonic() + sleep 

238 if stop_trying_at is not None and next_try_at > stop_trying_at: 

239 return False 

240 mod_time.sleep(sleep) 

241 

242 def do_acquire(self, token: str) -> bool: 

243 if self.timeout: 

244 # convert to milliseconds 

245 timeout = int(self.timeout * 1000) 

246 else: 

247 timeout = None 

248 if self.redis.set(self.name, token, nx=True, px=timeout): 

249 return True 

250 return False 

251 

252 def locked(self) -> bool: 

253 """ 

254 Returns True if this key is locked by any process, otherwise False. 

255 """ 

256 return self.redis.get(self.name) is not None 

257 

258 def owned(self) -> bool: 

259 """ 

260 Returns True if this key is locked by this lock, otherwise False. 

261 """ 

262 stored_token = self.redis.get(self.name) 

263 # need to always compare bytes to bytes 

264 # TODO: this can be simplified when the context manager is finished 

265 if stored_token and not isinstance(stored_token, bytes): 

266 encoder = self.redis.get_encoder() 

267 stored_token = encoder.encode(stored_token) 

268 return self.local.token is not None and stored_token == self.local.token 

269 

270 def release(self) -> None: 

271 """ 

272 Releases the already acquired lock 

273 

274 The token is only cleared after the Redis release operation completes 

275 successfully. If the release fails with a connection or timeout error 

276 the lock may still exist in Redis, so the token is preserved and the 

277 release can be retried. 

278 """ 

279 expected_token = self.local.token 

280 if expected_token is None: 

281 raise LockError( 

282 "Cannot release a lock that's not owned or is already unlocked.", 

283 lock_name=self.name, 

284 ) 

285 try: 

286 self.do_release(expected_token) 

287 except LockNotOwnedError: 

288 # Lock doesn't exist in Redis, so clear the token, but only if it is still 

289 # ours. With thread_local=False the namespace is shared, so another acquirer 

290 # may already have stored its own token; overwriting it would leave the new 

291 # owner unable to release. 

292 if self.local.token == expected_token: 

293 self.local.token = None 

294 raise 

295 # Only clear token after successful release, and only if it is still ours. 

296 if self.local.token == expected_token: 

297 self.local.token = None 

298 

299 def do_release(self, expected_token: str) -> None: 

300 if not bool( 

301 self.lua_release(keys=[self.name], args=[expected_token], client=self.redis) 

302 ): 

303 raise LockNotOwnedError( 

304 "Cannot release a lock that's no longer owned", 

305 lock_name=self.name, 

306 ) 

307 

308 def extend( 

309 self, additional_time: Number, replace_ttl: bool = False 

310 ) -> Literal[True]: 

311 """ 

312 Adds more time to an already acquired lock. 

313 

314 ``additional_time`` can be specified as an integer or a float, both 

315 representing the number of seconds to add. 

316 

317 ``replace_ttl`` if False (the default), add `additional_time` to 

318 the lock's existing ttl. If True, replace the lock's ttl with 

319 `additional_time`. 

320 

321 When the resulting TTL is non-positive for a lock with an expiry, the 

322 current expiry is left unchanged and ``True`` is returned. 

323 """ 

324 if self.local.token is None: 

325 raise LockError("Cannot extend an unlocked lock", lock_name=self.name) 

326 if self.timeout is None: 

327 raise LockError("Cannot extend a lock with no timeout", lock_name=self.name) 

328 return self.do_extend(additional_time, replace_ttl) 

329 

330 def do_extend(self, additional_time: Number, replace_ttl: bool) -> Literal[True]: 

331 additional_time = int(additional_time * 1000) 

332 if not bool( 

333 self.lua_extend( 

334 keys=[self.name], 

335 args=[self.local.token, additional_time, "1" if replace_ttl else "0"], 

336 client=self.redis, 

337 ) 

338 ): 

339 raise LockNotOwnedError( 

340 "Cannot extend a lock that's no longer owned", 

341 lock_name=self.name, 

342 ) 

343 return True 

344 

345 def reacquire(self) -> Literal[True]: 

346 """ 

347 Resets a TTL of an already acquired lock back to a timeout value. 

348 

349 When the resulting TTL is non-positive, the current expiry is left 

350 unchanged and ``True`` is returned. 

351 """ 

352 if self.local.token is None: 

353 raise LockError("Cannot reacquire an unlocked lock", lock_name=self.name) 

354 if self.timeout is None: 

355 raise LockError( 

356 "Cannot reacquire a lock with no timeout", 

357 lock_name=self.name, 

358 ) 

359 return self.do_reacquire() 

360 

361 def do_reacquire(self) -> Literal[True]: 

362 timeout = int(self.timeout * 1000) 

363 if not bool( 

364 self.lua_reacquire( 

365 keys=[self.name], args=[self.local.token, timeout], client=self.redis 

366 ) 

367 ): 

368 raise LockNotOwnedError( 

369 "Cannot reacquire a lock that's no longer owned", 

370 lock_name=self.name, 

371 ) 

372 return True