import threading import weakref from base64 import b64encode from logging import getLogger from os import urandom from typing import Union from redis import StrictRedis __version__ = '4.0.0' loggers = { k: getLogger("rhodecode." + ".".join((__name__, k))) for k in [ "acquire", "refresh.thread.start", "refresh.thread.stop", "refresh.thread.exit", "refresh.start", "refresh.shutdown", "refresh.exit", "release", ] } text_type = str binary_type = bytes # Check if the id match. If not, return an error code. UNLOCK_SCRIPT = b""" if redis.call("get", KEYS[1]) ~= ARGV[1] then return 1 else redis.call("del", KEYS[2]) redis.call("lpush", KEYS[2], 1) redis.call("pexpire", KEYS[2], ARGV[2]) redis.call("del", KEYS[1]) return 0 end """ # Covers both cases when key doesn't exist and doesn't equal to lock's id EXTEND_SCRIPT = b""" if redis.call("get", KEYS[1]) ~= ARGV[1] then return 1 elseif redis.call("ttl", KEYS[1]) < 0 then return 2 else redis.call("expire", KEYS[1], ARGV[2]) return 0 end """ RESET_SCRIPT = b""" redis.call('del', KEYS[2]) redis.call('lpush', KEYS[2], 1) redis.call('pexpire', KEYS[2], ARGV[2]) return redis.call('del', KEYS[1]) """ RESET_ALL_SCRIPT = b""" local locks = redis.call('keys', 'lock:*') local signal for _, lock in pairs(locks) do signal = 'lock-signal:' .. string.sub(lock, 6) redis.call('del', signal) redis.call('lpush', signal, 1) redis.call('expire', signal, 1) redis.call('del', lock) end return #locks """ class AlreadyAcquired(RuntimeError): pass class NotAcquired(RuntimeError): pass class AlreadyStarted(RuntimeError): pass class TimeoutNotUsable(RuntimeError): pass class InvalidTimeout(RuntimeError): pass class TimeoutTooLarge(RuntimeError): pass class NotExpirable(RuntimeError): pass class Lock(object): """ A Lock context manager implemented via redis SETNX/BLPOP. """ unlock_script = None extend_script = None reset_script = None reset_all_script = None _lock_renewal_interval: float _lock_renewal_thread: Union[threading.Thread, None] def __init__(self, redis_client, name, expire=None, id=None, auto_renewal=False, strict=True, signal_expire=1000): """ :param redis_client: An instance of :class:`~StrictRedis`. :param name: The name (redis key) the lock should have. :param expire: The lock expiry time in seconds. If left at the default (None) the lock will not expire. :param id: The ID (redis value) the lock should have. A random value is generated when left at the default. Note that if you specify this then the lock is marked as "held". Acquires won't be possible. :param auto_renewal: If set to ``True``, Lock will automatically renew the lock so that it doesn't expire for as long as the lock is held (acquire() called or running in a context manager). Implementation note: Renewal will happen using a daemon thread with an interval of ``expire*2/3``. If wishing to use a different renewal time, subclass Lock, call ``super().__init__()`` then set ``self._lock_renewal_interval`` to your desired interval. :param strict: If set ``True`` then the ``redis_client`` needs to be an instance of ``redis.StrictRedis``. :param signal_expire: Advanced option to override signal list expiration in milliseconds. Increase it for very slow clients. Default: ``1000``. """ if strict and not isinstance(redis_client, StrictRedis): raise ValueError("redis_client must be instance of StrictRedis. " "Use strict=False if you know what you're doing.") if auto_renewal and expire is None: raise ValueError("Expire may not be None when auto_renewal is set") self._client = redis_client if expire: expire = int(expire) if expire < 0: raise ValueError("A negative expire is not acceptable.") else: expire = None self._expire = expire self._signal_expire = signal_expire if id is None: self._id = b64encode(urandom(18)).decode('ascii') elif isinstance(id, binary_type): try: self._id = id.decode('ascii') except UnicodeDecodeError: self._id = b64encode(id).decode('ascii') elif isinstance(id, text_type): self._id = id else: raise TypeError(f"Incorrect type for `id`. Must be bytes/str not {type(id)}.") self._name = 'lock:' + name self._signal = 'lock-signal:' + name self._lock_renewal_interval = (float(expire) * 2 / 3 if auto_renewal else None) self._lock_renewal_thread = None self.register_scripts(redis_client) @classmethod def register_scripts(cls, redis_client): global reset_all_script if reset_all_script is None: cls.unlock_script = redis_client.register_script(UNLOCK_SCRIPT) cls.extend_script = redis_client.register_script(EXTEND_SCRIPT) cls.reset_script = redis_client.register_script(RESET_SCRIPT) cls.reset_all_script = redis_client.register_script(RESET_ALL_SCRIPT) reset_all_script = redis_client.register_script(RESET_ALL_SCRIPT) @property def _held(self): return self.id == self.get_owner_id() def reset(self): """ Forcibly deletes the lock. Use this with care. """ self.reset_script(client=self._client, keys=(self._name, self._signal), args=(self.id, self._signal_expire)) @property def id(self): return self._id def get_owner_id(self): owner_id = self._client.get(self._name) if isinstance(owner_id, binary_type): owner_id = owner_id.decode('ascii', 'replace') return owner_id def acquire(self, blocking=True, timeout=None): """ :param blocking: Boolean value specifying whether lock should be blocking or not. :param timeout: An integer value specifying the maximum number of seconds to block. """ logger = loggers["acquire"] logger.debug("Getting blocking: %s acquire on %r ...", blocking, self._name) if self._held: owner_id = self.get_owner_id() raise AlreadyAcquired("Already acquired from this Lock instance. Lock id: {}".format(owner_id)) if not blocking and timeout is not None: raise TimeoutNotUsable("Timeout cannot be used if blocking=False") if timeout: timeout = int(timeout) if timeout < 0: raise InvalidTimeout(f"Timeout ({timeout}) cannot be less than or equal to 0") if self._expire and not self._lock_renewal_interval and timeout > self._expire: raise TimeoutTooLarge(f"Timeout ({timeout}) cannot be greater than expire ({self._expire})") busy = True blpop_timeout = timeout or self._expire or 0 timed_out = False while busy: busy = not self._client.set(self._name, self._id, nx=True, ex=self._expire) if busy: if timed_out: return False elif blocking: timed_out = not self._client.blpop(self._signal, blpop_timeout) and timeout else: logger.warning("Failed to acquire Lock(%r).", self._name) return False logger.debug("Acquired Lock(%r).", self._name) if self._lock_renewal_interval is not None: self._start_lock_renewer() return True def extend(self, expire=None): """ Extends expiration time of the lock. :param expire: New expiration time. If ``None`` - `expire` provided during lock initialization will be taken. """ if expire: expire = int(expire) if expire < 0: raise ValueError("A negative expire is not acceptable.") elif self._expire is not None: expire = self._expire else: raise TypeError( "To extend a lock 'expire' must be provided as an " "argument to extend() method or at initialization time." ) error = self.extend_script(client=self._client, keys=(self._name, self._signal), args=(self._id, expire)) if error == 1: raise NotAcquired(f"Lock {self._name} is not acquired or it already expired.") elif error == 2: raise NotExpirable(f"Lock {self._name} has no assigned expiration time") elif error: raise RuntimeError(f"Unsupported error code {error} from EXTEND script") @staticmethod def _lock_renewer(name, lockref, interval, stop): """ Renew the lock key in redis every `interval` seconds for as long as `self._lock_renewal_thread.should_exit` is False. """ while not stop.wait(timeout=interval): loggers["refresh.thread.start"].debug("Refreshing Lock(%r).", name) lock: "Lock" = lockref() if lock is None: loggers["refresh.thread.stop"].debug( "Stopping loop because Lock(%r) was garbage collected.", name ) break lock.extend(expire=lock._expire) del lock loggers["refresh.thread.exit"].debug("Exiting renewal thread for Lock(%r).", name) def _start_lock_renewer(self): """ Starts the lock refresher thread. """ if self._lock_renewal_thread is not None: raise AlreadyStarted("Lock refresh thread already started") loggers["refresh.start"].debug( "Starting renewal thread for Lock(%r). Refresh interval: %s seconds.", self._name, self._lock_renewal_interval ) self._lock_renewal_stop = threading.Event() self._lock_renewal_thread = threading.Thread( group=None, target=self._lock_renewer, kwargs={ 'name': self._name, 'lockref': weakref.ref(self), 'interval': self._lock_renewal_interval, 'stop': self._lock_renewal_stop, }, ) self._lock_renewal_thread.daemon = True self._lock_renewal_thread.start() def _stop_lock_renewer(self): """ Stop the lock renewer. This signals the renewal thread and waits for its exit. """ if self._lock_renewal_thread is None or not self._lock_renewal_thread.is_alive(): return loggers["refresh.shutdown"].debug("Signaling renewal thread for Lock(%r) to exit.", self._name) self._lock_renewal_stop.set() self._lock_renewal_thread.join() self._lock_renewal_thread = None loggers["refresh.exit"].debug("Renewal thread for Lock(%r) exited.", self._name) def __enter__(self): acquired = self.acquire(blocking=True) if not acquired: raise AssertionError(f"Lock({self._name}) wasn't acquired, but blocking=True was used!") return self def __exit__(self, exc_type=None, exc_value=None, traceback=None): self.release() def release(self): """Releases the lock, that was acquired with the same object. .. note:: If you want to release a lock that you acquired in a different place you have two choices: * Use ``Lock("name", id=id_from_other_place).release()`` * Use ``Lock("name").reset()`` """ if self._lock_renewal_thread is not None: self._stop_lock_renewer() loggers["release"].debug("Releasing Lock(%r).", self._name) error = self.unlock_script(client=self._client, keys=(self._name, self._signal), args=(self._id, self._signal_expire)) if error == 1: raise NotAcquired(f"Lock({self._name}) is not acquired or it already expired.") elif error: raise RuntimeError(f"Unsupported error code {error} from EXTEND script.") def locked(self): """ Return true if the lock is acquired. Checks that lock with same name already exists. This method returns true, even if lock have another id. """ return self._client.exists(self._name) == 1 reset_all_script = None def reset_all(redis_client): """ Forcibly deletes all locks if its remains (like a crash reason). Use this with care. :param redis_client: An instance of :class:`~StrictRedis`. """ Lock.register_scripts(redis_client) reset_all_script(client=redis_client) # noqa