您最多选择25个主题 主题必须以字母或数字开头,可以包含连字符 (-),并且长度不得超过35个字符
 
 
 
 

275 行
11 KiB

  1. import threading
  2. import time as mod_time
  3. import uuid
  4. from redis.exceptions import LockError, LockNotOwnedError
  5. from redis.utils import dummy
  6. class Lock(object):
  7. """
  8. A shared, distributed Lock. Using Redis for locking allows the Lock
  9. to be shared across processes and/or machines.
  10. It's left to the user to resolve deadlock issues and make sure
  11. multiple clients play nicely together.
  12. """
  13. lua_release = None
  14. lua_extend = None
  15. lua_reacquire = None
  16. # KEYS[1] - lock name
  17. # ARGS[1] - token
  18. # return 1 if the lock was released, otherwise 0
  19. LUA_RELEASE_SCRIPT = """
  20. local token = redis.call('get', KEYS[1])
  21. if not token or token ~= ARGV[1] then
  22. return 0
  23. end
  24. redis.call('del', KEYS[1])
  25. return 1
  26. """
  27. # KEYS[1] - lock name
  28. # ARGS[1] - token
  29. # ARGS[2] - additional milliseconds
  30. # return 1 if the locks time was extended, otherwise 0
  31. LUA_EXTEND_SCRIPT = """
  32. local token = redis.call('get', KEYS[1])
  33. if not token or token ~= ARGV[1] then
  34. return 0
  35. end
  36. local expiration = redis.call('pttl', KEYS[1])
  37. if not expiration then
  38. expiration = 0
  39. end
  40. if expiration < 0 then
  41. return 0
  42. end
  43. redis.call('pexpire', KEYS[1], expiration + ARGV[2])
  44. return 1
  45. """
  46. # KEYS[1] - lock name
  47. # ARGS[1] - token
  48. # ARGS[2] - milliseconds
  49. # return 1 if the locks time was reacquired, otherwise 0
  50. LUA_REACQUIRE_SCRIPT = """
  51. local token = redis.call('get', KEYS[1])
  52. if not token or token ~= ARGV[1] then
  53. return 0
  54. end
  55. redis.call('pexpire', KEYS[1], ARGV[2])
  56. return 1
  57. """
  58. def __init__(self, redis, name, timeout=None, sleep=0.1,
  59. blocking=True, blocking_timeout=None, thread_local=True):
  60. """
  61. Create a new Lock instance named ``name`` using the Redis client
  62. supplied by ``redis``.
  63. ``timeout`` indicates a maximum life for the lock.
  64. By default, it will remain locked until release() is called.
  65. ``timeout`` can be specified as a float or integer, both representing
  66. the number of seconds to wait.
  67. ``sleep`` indicates the amount of time to sleep per loop iteration
  68. when the lock is in blocking mode and another client is currently
  69. holding the lock.
  70. ``blocking`` indicates whether calling ``acquire`` should block until
  71. the lock has been acquired or to fail immediately, causing ``acquire``
  72. to return False and the lock not being acquired. Defaults to True.
  73. Note this value can be overridden by passing a ``blocking``
  74. argument to ``acquire``.
  75. ``blocking_timeout`` indicates the maximum amount of time in seconds to
  76. spend trying to acquire the lock. A value of ``None`` indicates
  77. continue trying forever. ``blocking_timeout`` can be specified as a
  78. float or integer, both representing the number of seconds to wait.
  79. ``thread_local`` indicates whether the lock token is placed in
  80. thread-local storage. By default, the token is placed in thread local
  81. storage so that a thread only sees its token, not a token set by
  82. another thread. Consider the following timeline:
  83. time: 0, thread-1 acquires `my-lock`, with a timeout of 5 seconds.
  84. thread-1 sets the token to "abc"
  85. time: 1, thread-2 blocks trying to acquire `my-lock` using the
  86. Lock instance.
  87. time: 5, thread-1 has not yet completed. redis expires the lock
  88. key.
  89. time: 5, thread-2 acquired `my-lock` now that it's available.
  90. thread-2 sets the token to "xyz"
  91. time: 6, thread-1 finishes its work and calls release(). if the
  92. token is *not* stored in thread local storage, then
  93. thread-1 would see the token value as "xyz" and would be
  94. able to successfully release the thread-2's lock.
  95. In some use cases it's necessary to disable thread local storage. For
  96. example, if you have code where one thread acquires a lock and passes
  97. that lock instance to a worker thread to release later. If thread
  98. local storage isn't disabled in this case, the worker thread won't see
  99. the token set by the thread that acquired the lock. Our assumption
  100. is that these cases aren't common and as such default to using
  101. thread local storage.
  102. """
  103. self.redis = redis
  104. self.name = name
  105. self.timeout = timeout
  106. self.sleep = sleep
  107. self.blocking = blocking
  108. self.blocking_timeout = blocking_timeout
  109. self.thread_local = bool(thread_local)
  110. self.local = threading.local() if self.thread_local else dummy()
  111. self.local.token = None
  112. if self.timeout and self.sleep > self.timeout:
  113. raise LockError("'sleep' must be less than 'timeout'")
  114. self.register_scripts()
  115. def register_scripts(self):
  116. cls = self.__class__
  117. client = self.redis
  118. if cls.lua_release is None:
  119. cls.lua_release = client.register_script(cls.LUA_RELEASE_SCRIPT)
  120. if cls.lua_extend is None:
  121. cls.lua_extend = client.register_script(cls.LUA_EXTEND_SCRIPT)
  122. if cls.lua_reacquire is None:
  123. cls.lua_reacquire = \
  124. client.register_script(cls.LUA_REACQUIRE_SCRIPT)
  125. def __enter__(self):
  126. # force blocking, as otherwise the user would have to check whether
  127. # the lock was actually acquired or not.
  128. if self.acquire(blocking=True):
  129. return self
  130. raise LockError("Unable to acquire lock within the time specified")
  131. def __exit__(self, exc_type, exc_value, traceback):
  132. self.release()
  133. def acquire(self, blocking=None, blocking_timeout=None, token=None):
  134. """
  135. Use Redis to hold a shared, distributed lock named ``name``.
  136. Returns True once the lock is acquired.
  137. If ``blocking`` is False, always return immediately. If the lock
  138. was acquired, return True, otherwise return False.
  139. ``blocking_timeout`` specifies the maximum number of seconds to
  140. wait trying to acquire the lock.
  141. ``token`` specifies the token value to be used. If provided, token
  142. must be a bytes object or a string that can be encoded to a bytes
  143. object with the default encoding. If a token isn't specified, a UUID
  144. will be generated.
  145. """
  146. sleep = self.sleep
  147. if token is None:
  148. token = uuid.uuid1().hex.encode()
  149. else:
  150. encoder = self.redis.connection_pool.get_encoder()
  151. token = encoder.encode(token)
  152. if blocking is None:
  153. blocking = self.blocking
  154. if blocking_timeout is None:
  155. blocking_timeout = self.blocking_timeout
  156. stop_trying_at = None
  157. if blocking_timeout is not None:
  158. stop_trying_at = mod_time.time() + blocking_timeout
  159. while True:
  160. if self.do_acquire(token):
  161. self.local.token = token
  162. return True
  163. if not blocking:
  164. return False
  165. if stop_trying_at is not None and mod_time.time() > stop_trying_at:
  166. return False
  167. mod_time.sleep(sleep)
  168. def do_acquire(self, token):
  169. if self.timeout:
  170. # convert to milliseconds
  171. timeout = int(self.timeout * 1000)
  172. else:
  173. timeout = None
  174. if self.redis.set(self.name, token, nx=True, px=timeout):
  175. return True
  176. return False
  177. def locked(self):
  178. """
  179. Returns True if this key is locked by any process, otherwise False.
  180. """
  181. return self.redis.get(self.name) is not None
  182. def owned(self):
  183. """
  184. Returns True if this key is locked by this lock, otherwise False.
  185. """
  186. stored_token = self.redis.get(self.name)
  187. # need to always compare bytes to bytes
  188. # TODO: this can be simplified when the context manager is finished
  189. if stored_token and not isinstance(stored_token, bytes):
  190. encoder = self.redis.connection_pool.get_encoder()
  191. stored_token = encoder.encode(stored_token)
  192. return self.local.token is not None and \
  193. stored_token == self.local.token
  194. def release(self):
  195. "Releases the already acquired lock"
  196. expected_token = self.local.token
  197. if expected_token is None:
  198. raise LockError("Cannot release an unlocked lock")
  199. self.local.token = None
  200. self.do_release(expected_token)
  201. def do_release(self, expected_token):
  202. if not bool(self.lua_release(keys=[self.name],
  203. args=[expected_token],
  204. client=self.redis)):
  205. raise LockNotOwnedError("Cannot release a lock"
  206. " that's no longer owned")
  207. def extend(self, additional_time):
  208. """
  209. Adds more time to an already acquired lock.
  210. ``additional_time`` can be specified as an integer or a float, both
  211. representing the number of seconds to add.
  212. """
  213. if self.local.token is None:
  214. raise LockError("Cannot extend an unlocked lock")
  215. if self.timeout is None:
  216. raise LockError("Cannot extend a lock with no timeout")
  217. return self.do_extend(additional_time)
  218. def do_extend(self, additional_time):
  219. additional_time = int(additional_time * 1000)
  220. if not bool(self.lua_extend(keys=[self.name],
  221. args=[self.local.token, additional_time],
  222. client=self.redis)):
  223. raise LockNotOwnedError("Cannot extend a lock that's"
  224. " no longer owned")
  225. return True
  226. def reacquire(self):
  227. """
  228. Resets a TTL of an already acquired lock back to a timeout value.
  229. """
  230. if self.local.token is None:
  231. raise LockError("Cannot reacquire an unlocked lock")
  232. if self.timeout is None:
  233. raise LockError("Cannot reacquire a lock with no timeout")
  234. return self.do_reacquire()
  235. def do_reacquire(self):
  236. timeout = int(self.timeout * 1000)
  237. if not bool(self.lua_reacquire(keys=[self.name],
  238. args=[self.local.token, timeout],
  239. client=self.redis)):
  240. raise LockNotOwnedError("Cannot reacquire a lock that's"
  241. " no longer owned")
  242. return True