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

1import asyncio 

2import logging 

3import threading 

4import uuid 

5from types import SimpleNamespace 

6from typing import TYPE_CHECKING, Awaitable, Literal, Optional, Union 

7 

8from redis.exceptions import LockError, LockNotOwnedError 

9from redis.typing import Number 

10 

11if TYPE_CHECKING: 

12 from redis.asyncio import Redis, RedisCluster 

13 

14logger = logging.getLogger(__name__) 

15 

16 

17class Lock: 

18 """ 

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

20 to be shared across processes and/or machines. 

21 

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

23 multiple clients play nicely together. 

24 """ 

25 

26 lua_release = None 

27 lua_extend = None 

28 lua_reacquire = None 

29 

30 # KEYS[1] - lock name 

31 # ARGV[1] - token 

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

33 LUA_RELEASE_SCRIPT = """ 

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

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

36 return 0 

37 end 

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

39 return 1 

40 """ 

41 

42 # KEYS[1] - lock name 

43 # ARGV[1] - token 

44 # ARGV[2] - additional milliseconds 

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

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

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

48 LUA_EXTEND_SCRIPT = """ 

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

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

51 return 0 

52 end 

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

54 if not expiration then 

55 expiration = 0 

56 end 

57 if expiration < 0 then 

58 return 0 

59 end 

60 

61 local newttl = ARGV[2] 

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

63 newttl = ARGV[2] + expiration 

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

79 return 1 

80 """ 

81 

82 def __init__( 

83 self, 

84 redis: Union["Redis", "RedisCluster"], 

85 name: Union[str, bytes, memoryview], 

86 timeout: Optional[float] = None, 

87 sleep: float = 0.1, 

88 blocking: bool = True, 

89 blocking_timeout: Optional[Number] = None, 

90 thread_local: bool = True, 

91 raise_on_release_error: bool = True, 

92 ): 

93 """ 

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

95 supplied by ``redis``. 

96 

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

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

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

100 the number of seconds to wait. 

101 

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

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

104 currently holding the lock. 

105 

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

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

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

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

110 argument to ``acquire``. 

111 

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

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

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

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

116 

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

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

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

120 another thread. Consider the following timeline: 

121 

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

123 thread-1 sets the token to "abc" 

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

125 Lock instance. 

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

127 key. 

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

129 thread-2 sets the token to "xyz" 

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

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

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

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

134 

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

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

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

138 will be logged and the exception will be suppressed. 

139 

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

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

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

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

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

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

146 thread local storage. 

147 """ 

148 self.redis = redis 

149 self.name = name 

150 self.timeout = timeout 

151 self.sleep = sleep 

152 self.blocking = blocking 

153 self.blocking_timeout = blocking_timeout 

154 self.thread_local = bool(thread_local) 

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

156 self.raise_on_release_error = raise_on_release_error 

157 self.local.token = None 

158 self.register_scripts() 

159 

160 def register_scripts(self): 

161 cls = self.__class__ 

162 client = self.redis 

163 if cls.lua_release is None: 

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

165 if cls.lua_extend is None: 

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

167 if cls.lua_reacquire is None: 

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

169 

170 async def __aenter__(self): 

171 if await self.acquire(): 

172 return self 

173 raise LockError( 

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

175 lock_name=self.name, 

176 ) 

177 

178 async def __aexit__(self, exc_type, exc_value, traceback): 

179 try: 

180 await self.release() 

181 except LockError: 

182 if self.raise_on_release_error: 

183 raise 

184 logger.warning( 

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

186 ) 

187 

188 async def acquire( 

189 self, 

190 blocking: Optional[bool] = None, 

191 blocking_timeout: Optional[Number] = None, 

192 token: Optional[Union[str, bytes]] = None, 

193 ): 

194 """ 

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

196 Returns True once the lock is acquired. 

197 

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

199 was acquired, return True, otherwise return False. 

200 

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

202 wait trying to acquire the lock. 

203 

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

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

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

207 will be generated. 

208 """ 

209 sleep = self.sleep 

210 if token is None: 

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

212 else: 

213 try: 

214 encoder = self.redis.connection_pool.get_encoder() 

215 except AttributeError: 

216 # Cluster 

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 = asyncio.get_running_loop().time() + blocking_timeout 

226 while True: 

227 if await self.do_acquire(token): 

228 self.local.token = token 

229 return True 

230 if not blocking: 

231 return False 

232 next_try_at = asyncio.get_running_loop().time() + sleep 

233 if stop_trying_at is not None and next_try_at > stop_trying_at: 

234 return False 

235 await asyncio.sleep(sleep) 

236 

237 async def do_acquire(self, token: Union[str, bytes]) -> bool: 

238 if self.timeout: 

239 # convert to milliseconds 

240 timeout = int(self.timeout * 1000) 

241 else: 

242 timeout = None 

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

244 return True 

245 return False 

246 

247 async def locked(self) -> bool: 

248 """ 

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

250 """ 

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

252 

253 async def owned(self) -> bool: 

254 """ 

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

256 """ 

257 stored_token = await 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 try: 

262 encoder = self.redis.connection_pool.get_encoder() 

263 except AttributeError: 

264 # Cluster 

265 encoder = self.redis.get_encoder() 

266 stored_token = encoder.encode(stored_token) 

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

268 

269 async def release(self) -> None: 

270 """Releases the already acquired lock. 

271 

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

273 successfully. This ensures that if the release is cancelled mid-operation, 

274 the lock state remains consistent and can be retried. 

275 """ 

276 expected_token = self.local.token 

277 if expected_token is None: 

278 raise LockError( 

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

280 lock_name=self.name, 

281 ) 

282 try: 

283 await self.do_release(expected_token) 

284 except LockNotOwnedError: 

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

286 # ours. threading.local() is shared by every task on one event loop, so a 

287 # Lock shared between tasks may already hold another acquirer's token; 

288 # overwriting it would leave the new owner unable to release. 

289 if self.local.token == expected_token: 

290 self.local.token = None 

291 raise 

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

293 if self.local.token == expected_token: 

294 self.local.token = None 

295 

296 async def do_release(self, expected_token: bytes) -> None: 

297 if not bool( 

298 await self.lua_release( 

299 keys=[self.name], args=[expected_token], client=self.redis 

300 ) 

301 ): 

302 raise LockNotOwnedError( 

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

304 lock_name=self.name, 

305 ) 

306 

307 def extend( 

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

309 ) -> Awaitable[Literal[True]]: 

310 """ 

311 Adds more time to an already acquired lock. 

312 

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

314 representing the number of seconds to add. 

315 

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

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

318 `additional_time`. 

319 """ 

320 if self.local.token is None: 

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

322 if self.timeout is None: 

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

324 return self.do_extend(additional_time, replace_ttl) 

325 

326 async def do_extend(self, additional_time, replace_ttl) -> Literal[True]: 

327 additional_time = int(additional_time * 1000) 

328 if not bool( 

329 await self.lua_extend( 

330 keys=[self.name], 

331 args=[self.local.token, additional_time, replace_ttl and "1" or "0"], 

332 client=self.redis, 

333 ) 

334 ): 

335 raise LockNotOwnedError( 

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

337 lock_name=self.name, 

338 ) 

339 return True 

340 

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

342 """ 

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

344 """ 

345 if self.local.token is None: 

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

347 if self.timeout is None: 

348 raise LockError( 

349 "Cannot reacquire a lock with no timeout", lock_name=self.name 

350 ) 

351 return self.do_reacquire() 

352 

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

354 timeout = int(self.timeout * 1000) 

355 if not bool( 

356 await self.lua_reacquire( 

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

358 ) 

359 ): 

360 raise LockNotOwnedError( 

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

362 lock_name=self.name, 

363 ) 

364 return True