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 redis.call('pexpire', KEYS[1], newttl) 

63 return 1 

64 """ 

65 

66 # KEYS[1] - lock name 

67 # ARGV[1] - token 

68 # ARGV[2] - milliseconds 

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

70 LUA_REACQUIRE_SCRIPT = """ 

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

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

73 return 0 

74 end 

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

76 return 1 

77 """ 

78 

79 def __init__( 

80 self, 

81 redis, 

82 name: str, 

83 timeout: Optional[Number] = None, 

84 sleep: Number = 0.1, 

85 blocking: bool = True, 

86 blocking_timeout: Optional[Number] = None, 

87 thread_local: bool = True, 

88 raise_on_release_error: bool = True, 

89 ): 

90 """ 

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

92 supplied by ``redis``. 

93 

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

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

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

97 the number of seconds to wait. 

98 

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

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

101 currently holding the lock. 

102 

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

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

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

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

107 argument to ``acquire``. 

108 

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

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

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

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

113 

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

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

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

117 another thread. Consider the following timeline: 

118 

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

120 thread-1 sets the token to "abc" 

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

122 Lock instance. 

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

124 key. 

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

126 thread-2 sets the token to "xyz" 

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

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

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

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

131 

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

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

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

135 will be logged and the exception will be suppressed. 

136 

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

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

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

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

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

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

143 thread local storage. 

144 """ 

145 self.redis = redis 

146 self.name = name 

147 self.timeout = timeout 

148 self.sleep = sleep 

149 self.blocking = blocking 

150 self.blocking_timeout = blocking_timeout 

151 self.thread_local = bool(thread_local) 

152 self.raise_on_release_error = raise_on_release_error 

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

154 self.local.token = None 

155 self.register_scripts() 

156 

157 def register_scripts(self) -> None: 

158 cls = self.__class__ 

159 client = self.redis 

160 if cls.lua_release is None: 

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

162 if cls.lua_extend is None: 

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

164 if cls.lua_reacquire is None: 

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

166 

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

168 if self.acquire(): 

169 return self 

170 raise LockError( 

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

172 lock_name=self.name, 

173 ) 

174 

175 def __exit__( 

176 self, 

177 exc_type: Optional[Type[BaseException]], 

178 exc_value: Optional[BaseException], 

179 traceback: Optional[TracebackType], 

180 ) -> None: 

181 try: 

182 self.release() 

183 except LockError: 

184 if self.raise_on_release_error: 

185 raise 

186 logger.warning( 

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

188 ) 

189 

190 def acquire( 

191 self, 

192 sleep: Optional[Number] = None, 

193 blocking: Optional[bool] = None, 

194 blocking_timeout: Optional[Number] = None, 

195 token: Optional[str] = None, 

196 ): 

197 """ 

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

199 Returns True once the lock is acquired. 

200 

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

202 was acquired, return True, otherwise return False. 

203 

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

205 wait trying to acquire the lock. 

206 

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

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

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

210 will be generated. 

211 """ 

212 if sleep is None: 

213 sleep = self.sleep 

214 if token is None: 

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

216 else: 

217 encoder = self.redis.get_encoder() 

218 token = encoder.encode(token) 

219 if blocking is None: 

220 blocking = self.blocking 

221 if blocking_timeout is None: 

222 blocking_timeout = self.blocking_timeout 

223 stop_trying_at = None 

224 if blocking_timeout is not None: 

225 stop_trying_at = mod_time.monotonic() + blocking_timeout 

226 while True: 

227 if self.do_acquire(token): 

228 self.local.token = token 

229 return True 

230 if not blocking: 

231 return False 

232 next_try_at = mod_time.monotonic() + sleep 

233 if stop_trying_at is not None and next_try_at > stop_trying_at: 

234 return False 

235 mod_time.sleep(sleep) 

236 

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

238 if self.timeout: 

239 # convert to milliseconds 

240 timeout = int(self.timeout * 1000) 

241 else: 

242 timeout = None 

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

244 return True 

245 return False 

246 

247 def locked(self) -> bool: 

248 """ 

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

250 """ 

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

252 

253 def owned(self) -> bool: 

254 """ 

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

256 """ 

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

258 # need to always compare bytes to bytes 

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

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

261 encoder = self.redis.get_encoder() 

262 stored_token = encoder.encode(stored_token) 

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

264 

265 def release(self) -> None: 

266 """ 

267 Releases the already acquired lock 

268 

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

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

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

272 release can be retried. 

273 """ 

274 expected_token = self.local.token 

275 if expected_token is None: 

276 raise LockError( 

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

278 lock_name=self.name, 

279 ) 

280 try: 

281 self.do_release(expected_token) 

282 except LockNotOwnedError: 

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

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

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

286 # owner unable to release. 

287 if self.local.token == expected_token: 

288 self.local.token = None 

289 raise 

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

291 if self.local.token == expected_token: 

292 self.local.token = None 

293 

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

295 if not bool( 

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

297 ): 

298 raise LockNotOwnedError( 

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

300 lock_name=self.name, 

301 ) 

302 

303 def extend( 

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

305 ) -> Literal[True]: 

306 """ 

307 Adds more time to an already acquired lock. 

308 

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

310 representing the number of seconds to add. 

311 

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

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

314 `additional_time`. 

315 """ 

316 if self.local.token is None: 

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

318 if self.timeout is None: 

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

320 return self.do_extend(additional_time, replace_ttl) 

321 

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

323 additional_time = int(additional_time * 1000) 

324 if not bool( 

325 self.lua_extend( 

326 keys=[self.name], 

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

328 client=self.redis, 

329 ) 

330 ): 

331 raise LockNotOwnedError( 

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

333 lock_name=self.name, 

334 ) 

335 return True 

336 

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

338 """ 

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

340 """ 

341 if self.local.token is None: 

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

343 if self.timeout is None: 

344 raise LockError( 

345 "Cannot reacquire a lock with no timeout", 

346 lock_name=self.name, 

347 ) 

348 return self.do_reacquire() 

349 

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

351 timeout = int(self.timeout * 1000) 

352 if not bool( 

353 self.lua_reacquire( 

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

355 ) 

356 ): 

357 raise LockNotOwnedError( 

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

359 lock_name=self.name, 

360 ) 

361 return True