redis: added better distributed lock implementation.

rhodecode/lib/_vendor/redis_lock/__init__.py

@@ -0,0 +1,389 b''
+import sys
+import threading
+import weakref
+from base64 import b64encode
+from logging import getLogger
+from os import urandom
+
+from redis import StrictRedis
+
+__version__ = '3.7.0'
+
+loggers = {
+    k: getLogger(".".join((__name__, k)))
+    for k in [
+        "acquire",
+        "refresh.thread.start",
+        "refresh.thread.stop",
+        "refresh.thread.exit",
+        "refresh.start",
+        "refresh.shutdown",
+        "refresh.exit",
+        "release",
+    ]
+}
+
+PY3 = sys.version_info[0] == 3
+
+if PY3:
+    text_type = str
+    binary_type = bytes
+else:
+    text_type = unicode  # noqa
+    binary_type = str
+
+
+# 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
+
+    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("Incorrect type for `id`. Must be bytes/str not %s." % 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:
+            reset_all_script = redis_client.register_script(RESET_ALL_SCRIPT)
+            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)
+
+    @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 %r ...", self._name)
+
+        if self._held:
+            raise AlreadyAcquired("Already acquired from this Lock instance.")
+
+        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("Timeout (%d) cannot be less than or equal to 0" % timeout)
+
+            if self._expire and not self._lock_renewal_interval and timeout > self._expire:
+                raise TimeoutTooLarge("Timeout (%d) cannot be greater than expire (%d)" % (timeout, 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 get %r.", self._name)
+                    return False
+
+        logger.info("Got lock for %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("Lock %s is not acquired or it already expired." % self._name)
+        elif error == 2:
+            raise NotExpirable("Lock %s has no assigned expiration time" % self._name)
+        elif error:
+            raise RuntimeError("Unsupported error code %s from EXTEND script" % error)
+
+    @staticmethod
+    def _lock_renewer(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")
+            lock = lockref()
+            if lock is None:
+                loggers["refresh.thread.stop"].debug(
+                    "The lock no longer exists, stopping lock refreshing"
+                )
+                break
+            lock.extend(expire=lock._expire)
+            del lock
+        loggers["refresh.thread.exit"].debug("Exit requested, stopping lock refreshing")
+
+    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 thread to refresh lock every %s seconds",
+            self._lock_renewal_interval
+        )
+        self._lock_renewal_stop = threading.Event()
+        self._lock_renewal_thread = threading.Thread(
+            group=None,
+            target=self._lock_renewer,
+            kwargs={'lockref': weakref.ref(self),
+                    'interval': self._lock_renewal_interval,
+                    'stop': self._lock_renewal_stop}
+        )
+        self._lock_renewal_thread.setDaemon(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("Signalling the lock refresher to stop")
+        self._lock_renewal_stop.set()
+        self._lock_renewal_thread.join()
+        self._lock_renewal_thread = None
+        loggers["refresh.exit"].debug("Lock refresher has stopped")
+
+    def __enter__(self):
+        acquired = self.acquire(blocking=True)
+        assert acquired, "Lock wasn't acquired, but blocking=True"
+        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 %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("Lock %s is not acquired or it already expired." % self._name)
+        elif error:
+            raise RuntimeError("Unsupported error code %s from EXTEND script." % error)
+
+    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

rhodecode/lib/rc_cache/backends.py

@@ -1,305 +1,311 b''
 # -*- coding: utf-8 -*-
 
 # Copyright (C) 2015-2020 RhodeCode GmbH
 #
 # This program is free software: you can redistribute it and/or modify
 # it under the terms of the GNU Affero General Public License, version 3
 # (only), as published by the Free Software Foundation.
 #
 # This program is distributed in the hope that it will be useful,
 # but WITHOUT ANY WARRANTY; without even the implied warranty of
 # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 # GNU General Public License for more details.
 #
 # You should have received a copy of the GNU Affero General Public License
 # along with this program.  If not, see <http://www.gnu.org/licenses/>.
 #
 # This program is dual-licensed. If you wish to learn more about the
 # RhodeCode Enterprise Edition, including its added features, Support services,
 # and proprietary license terms, please see https://rhodecode.com/licenses/
 
 import time
 import errno
 import logging
 
 import msgpack
 import gevent
 import redis
 
 from dogpile.cache.api import CachedValue
 from dogpile.cache.backends import memory as memory_backend
 from dogpile.cache.backends import file as file_backend
 from dogpile.cache.backends import redis as redis_backend
 from dogpile.cache.backends.file import NO_VALUE, compat, FileLock
 from dogpile.cache.util import memoized_property
 
 from rhodecode.lib.memory_lru_dict import LRUDict, LRUDictDebug
 
 
 _default_max_size = 1024
 
 log = logging.getLogger(__name__)
 
 
 class LRUMemoryBackend(memory_backend.MemoryBackend):
     key_prefix = 'lru_mem_backend'
     pickle_values = False
 
     def __init__(self, arguments):
         max_size = arguments.pop('max_size', _default_max_size)
 
         LRUDictClass = LRUDict
         if arguments.pop('log_key_count', None):
             LRUDictClass = LRUDictDebug
 
         arguments['cache_dict'] = LRUDictClass(max_size)
         super(LRUMemoryBackend, self).__init__(arguments)
 
     def delete(self, key):
         try:
             del self._cache[key]
         except KeyError:
             # we don't care if key isn't there at deletion
             pass
 
     def delete_multi(self, keys):
         for key in keys:
             self.delete(key)
 
 
 class PickleSerializer(object):
 
     def _dumps(self, value, safe=False):
         try:
             return compat.pickle.dumps(value)
         except Exception:
             if safe:
                 return NO_VALUE
             else:
                 raise
 
     def _loads(self, value, safe=True):
         try:
             return compat.pickle.loads(value)
         except Exception:
             if safe:
                 return NO_VALUE
             else:
                 raise
 
 
 class MsgPackSerializer(object):
 
     def _dumps(self, value, safe=False):
         try:
             return msgpack.packb(value)
         except Exception:
             if safe:
                 return NO_VALUE
             else:
                 raise
 
     def _loads(self, value, safe=True):
         """
         pickle maintained the `CachedValue` wrapper of the tuple
         msgpack does not, so it must be added back in.
        """
         try:
             value = msgpack.unpackb(value, use_list=False)
             return CachedValue(*value)
         except Exception:
             if safe:
                 return NO_VALUE
             else:
                 raise
 
 
 import fcntl
 flock_org = fcntl.flock
 
 
 class CustomLockFactory(FileLock):
 
     @memoized_property
     def _module(self):
 
         def gevent_flock(fd, operation):
             """
             Gevent compatible flock
             """
             # set non-blocking, this will cause an exception if we cannot acquire a lock
             operation |= fcntl.LOCK_NB
             start_lock_time = time.time()
             timeout = 60 * 15  # 15min
             while True:
                 try:
                     flock_org(fd, operation)
                     # lock has been acquired
                     break
                 except (OSError, IOError) as e:
                     # raise on other errors than Resource temporarily unavailable
                     if e.errno != errno.EAGAIN:
                         raise
                     elif (time.time() - start_lock_time) > timeout:
                         # waited to much time on a lock, better fail than loop for ever
                         log.error('Failed to acquire lock on `%s` after waiting %ss',
                                   self.filename, timeout)
                         raise
                     wait_timeout = 0.03
                     log.debug('Failed to acquire lock on `%s`, retry in %ss',
                               self.filename, wait_timeout)
                     gevent.sleep(wait_timeout)
 
         fcntl.flock = gevent_flock
         return fcntl
 
 
 class FileNamespaceBackend(PickleSerializer, file_backend.DBMBackend):
     key_prefix = 'file_backend'
 
     def __init__(self, arguments):
         arguments['lock_factory'] = CustomLockFactory
         db_file = arguments.get('filename')
 
         log.debug('initialing %s DB in %s', self.__class__.__name__, db_file)
         try:
             super(FileNamespaceBackend, self).__init__(arguments)
         except Exception:
             log.error('Failed to initialize db at: %s', db_file)
             raise
 
     def __repr__(self):
         return '{} `{}`'.format(self.__class__, self.filename)
 
     def list_keys(self, prefix=''):
         prefix = '{}:{}'.format(self.key_prefix, prefix)
 
         def cond(v):
             if not prefix:
                 return True
 
             if v.startswith(prefix):
                 return True
             return False
 
         with self._dbm_file(True) as dbm:
             try:
                 return filter(cond, dbm.keys())
             except Exception:
                 log.error('Failed to fetch DBM keys from DB: %s', self.get_store())
                 raise
 
     def get_store(self):
         return self.filename
 
     def _dbm_get(self, key):
         with self._dbm_file(False) as dbm:
             if hasattr(dbm, 'get'):
                 value = dbm.get(key, NO_VALUE)
             else:
                 # gdbm objects lack a .get method
                 try:
                     value = dbm[key]
                 except KeyError:
                     value = NO_VALUE
             if value is not NO_VALUE:
                 value = self._loads(value)
             return value
 
     def get(self, key):
         try:
             return self._dbm_get(key)
         except Exception:
             log.error('Failed to fetch DBM key %s from DB: %s', key, self.get_store())
             raise
 
     def set(self, key, value):
         with self._dbm_file(True) as dbm:
             dbm[key] = self._dumps(value)
 
     def set_multi(self, mapping):
         with self._dbm_file(True) as dbm:
             for key, value in mapping.items():
                 dbm[key] = self._dumps(value)
 
 
 class BaseRedisBackend(redis_backend.RedisBackend):
 
     def _create_client(self):
         args = {}
 
         if self.url is not None:
             args.update(url=self.url)
 
         else:
             args.update(
                 host=self.host, password=self.password,
                 port=self.port, db=self.db
             )
 
         connection_pool = redis.ConnectionPool(**args)
 
         return redis.StrictRedis(connection_pool=connection_pool)
 
     def list_keys(self, prefix=''):
         prefix = '{}:{}*'.format(self.key_prefix, prefix)
         return self.client.keys(prefix)
 
     def get_store(self):
         return self.client.connection_pool
 
     def get(self, key):
         value = self.client.get(key)
         if value is None:
             return NO_VALUE
         return self._loads(value)
 
     def get_multi(self, keys):
         if not keys:
             return []
         values = self.client.mget(keys)
         loads = self._loads
         return [
             loads(v) if v is not None else NO_VALUE
             for v in values]
 
     def set(self, key, value):
         if self.redis_expiration_time:
             self.client.setex(key, self.redis_expiration_time,
                               self._dumps(value))
         else:
             self.client.set(key, self._dumps(value))
 
     def set_multi(self, mapping):
         dumps = self._dumps
         mapping = dict(
             (k, dumps(v))
             for k, v in mapping.items()
         )
 
         if not self.redis_expiration_time:
             self.client.mset(mapping)
         else:
             pipe = self.client.pipeline()
             for key, value in mapping.items():
                 pipe.setex(key, self.redis_expiration_time, value)
             pipe.execute()
 
     def get_mutex(self, key):
-        u = redis_backend.u
         if self.distributed_lock:
-            lock_key = u('_lock_{0}').format(key)
+            import redis_lock
+            lock_key = redis_backend.u('_lock_{0}').format(key)
             log.debug('Trying to acquire Redis lock for key %s', lock_key)
-            return self.client.lock(lock_key, self.lock_timeout, self.lock_sleep)
+            lock = redis_lock.Lock(
+                redis_client=self.client,
+                name=lock_key,
+                expire=self.lock_timeout,
+                auto_renewal=False,
+            )
+            return lock
         else:
             return None
 
 
 class RedisPickleBackend(PickleSerializer, BaseRedisBackend):
     key_prefix = 'redis_pickle_backend'
     pass
 
 
 class RedisMsgPackBackend(MsgPackSerializer, BaseRedisBackend):
     key_prefix = 'redis_msgpack_backend'
     pass