python redis自帶門神 lock 方法

金色旭光發表於2021-10-28

redis 支援的資料結構比較豐富,自制一個鎖也很方便,所以極少提到其原生鎖的方法。但是在單機版redis的使用時,自帶鎖的使用還是非常方便的。自己有車還打啥滴滴順風車是吧,本篇主要介紹redis-py模組中原生鎖lock的相關方法。

使用場景

  • 多執行緒資源搶佔
  • 關鍵變數鎖定
  • 防止重複執行程式碼

基本使用

lock使用

ubuntu 安裝redis

apt install redis-server

安裝python redis-py模組

pip install redis

普通使用

import redis

# redis 執行緒池
pool = redis.ConnectionPool(host='localhost', port=6379, decode_responses=True)
r = redis.Redis(connection_pool=pool)


# 建立一個鎖
lock = r.lock('mylock')
try:
    # 獲取鎖
    lock.acquire()
    print('get lock')
except:
    pass
finally:
    # 釋放鎖
    lock.release()

因為獲取了鎖之後一定要釋放鎖,所以用try except finally的錯誤捕獲方法保證不管在獲取鎖之後是否發生錯誤,最後都會釋放鎖,這是安全使用鎖的一種姿勢。
例如,在如下例子中,當獲取鎖之後主動丟擲異常,此時也能保證鎖的正常釋放。

lock = r.lock('mylock')
try:
    lock.acquire()
    print('get lock')
    raise
except:
    pass
finally:
    print('release lock')
    lock.release()

推薦使用 with 方法

推薦的使用方法是with,在檔案操作,執行緒鎖操作時經常使用with方法。在with的語法中,獲取鎖和釋放鎖都已經自動完成,所以是一種更加簡潔和高效的使用方法。

with r.lock('mylock'):
    print('get lock')

with語句在執行程式碼之前加鎖,在退出之前釋放鎖。具體來說就是實現了__enter____exit__方法,這些都可以在最後的原始碼中找答案,不僅能學到lock的方法,也能加深對with語法的理解。

lock支援的引數

lock 函式的定義:

def lock(self, name, timeout=None, sleep=0.1, blocking_timeout=None,
             lock_class=None, thread_local=True):
  • name: 鎖的名字
  • timeout: 鎖的生命週期。如果不設定鎖就不會過期,直到被釋放。預設不設定
  • sleep: 當獲取鎖阻塞時,嘗試獲取鎖迴圈的間隔時間,預設是睡眠間隔0.1s嘗試
  • blocking_timeout:當獲取鎖阻塞時,最長的等待時間,預設一直等待
  • lock_class: 強制執行指定的鎖實現
  • thread_local:用來表示是否將token儲存線上程本地。預設是儲存在本地執行緒的,所以一個執行緒只能看到自己的token,而不能被另一個執行緒使用。比如有如下例子:
    0s: 執行緒1獲取到鎖my-lock,設定過期時間是5s,token是abc。
    1s:執行緒2嘗試獲取鎖。
    5s:執行緒1還沒有完成,redis釋放了鎖。同時執行緒2獲取了鎖,並設定token是xyz
    6s: 執行緒1執行完成,然後呼叫release()釋放鎖。如果token不是儲存在本地,那麼執行緒1將拿到token xyz,然後釋放了執行緒2的鎖
    在某些用例中,有必要禁用執行緒本地儲存。
    例如,如果您有程式碼,其中一個執行緒獲取一個鎖,並將該鎖例項傳遞給工作執行緒,以便稍後釋放。
    如果在這種情況下未禁用執行緒本地儲存,那麼工作執行緒將看不到獲取鎖的執行緒設定的令牌。
    我們的假設是,這些情況並不常見,因此預設使用執行緒本地儲存。

通過建立lock時傳入引數來控制lock的一些屬性,比如獲取鎖的最長等待時間,持有鎖的最長時間等。

設定鎖的生命週期

設定鎖5s,拿到鎖5s之內還能釋放鎖

>>> lock = r.lock('mylock_one', timeout=5)
>>> 
>>> lock.acquire()
True
>>> 
>>> lock.release()
>>> 
>>> 

設定鎖2s,拿到鎖2s之後,鎖自動釋放掉,再次釋放就會報錯。

>>> 
>>> 
>>> 
>>> lock = r.lock('mylock_one', timeout=2)
>>> 
>>> lock.acquire()
True
>>> lock.release()
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/home/ljk/.virtualenvs/work/lib/python3.7/site-packages/redis/lock.py", line 232, in release
    self.do_release(expected_token)
  File "/home/ljk/.virtualenvs/work/lib/python3.7/site-packages/redis/lock.py", line 238, in do_release
    raise LockNotOwnedError("Cannot release a lock"
redis.exceptions.LockNotOwnedError: Cannot release a lock that's no longer owned
>>> 

設定阻塞等待的最長時間

當不設定最長等待時間時,會一直等待key的釋放,當設定了最長等待時間,如果在time內key沒有釋放,那麼就直接返回False,表示獲取不到。

>>> lock = r.lock('mylock_one')
>>> lock = r.lock('mylock_one')
>>> 
>>> 
>>> lock.acquire()
True
>>> 
>>> lock = r.lock('mylock_one', blocking_timeout=3)
>>> lock.acquire()
False
>>> 

注意:需要說明的是 lock 方法針對的某一個key的獲取,即獲取某一個key作為鎖的關鍵字,而不是獲取某一個瑣。這和下面的獲取鎖的方法acquire是有根本的區別的。

lock 的方法

lock擁有的方法並不是很多,所以用法不會花裡胡哨。lock主要的方法如下:

  • acquire:獲取鎖
  • release:釋放鎖
  • owned:key是否被該鎖擁有,擁有返回True
  • locked:鎖是不是否被任何一個執行緒鎖住,鎖住返回True

acquire

acquire 就是獲取鎖的方法,原型如下:

def acquire(self, blocking=None, blocking_timeout=None, token=None):

最簡單的使用

lock.acquire()

當鎖已經被佔用時再次請求,acquire預設會阻塞。

非阻塞使用

當設定了blocking=False時,表示拿不到鎖時不阻塞,直接返回False

lock.acquire(blocking=False)

設定阻塞時長

當拿不到鎖時可以設定阻塞的時長

lock.acquire(blocking_timeout=5)

5s之內拿不到鎖的話,就會放棄嘗試,返回False

owned

owned表示key :mylock_one是不是被該鎖 lock 作為關鍵字。被鎖定返回True,沒有鎖定返回False

>>> lock = r.lock('mylock_one')
>>> 
>>> lock.owned()
False
>>> 
>>> lock.acquire()
True
>>> 
>>> lock.owned()
True
>>> 

locked

是用來看鎖是不是被佔用,佔用返回True,沒有被佔用返回False

>>> lock = r.lock('mylock_two')
>>> 
>>> lock.locked()
False
>>> 
>>> lock.acquire()
True
>>> 
>>> lock.locked()
True
>>> 

附錄 lock 的原始碼

其實使用redis的字串以及過期時間也是可以自己實現一個鎖的,事實上lock的實現也確實是基於字串來實現的。對於想要閱讀優秀原始碼或更深入理解lock特性的同學來說,該原始碼是一個不錯的學習資料。

import threading
import time as mod_time
import uuid
from redis.exceptions import LockError, LockNotOwnedError
from redis.utils import dummy


class Lock(object):
    """
    A shared, distributed Lock. Using Redis for locking allows the Lock
    to be shared across processes and/or machines.

    It's left to the user to resolve deadlock issues and make sure
    multiple clients play nicely together.
    """

    lua_release = None
    lua_extend = None
    lua_reacquire = None

    # KEYS[1] - lock name
    # ARGV[1] - token
    # return 1 if the lock was released, otherwise 0
    LUA_RELEASE_SCRIPT = """
        local token = redis.call('get', KEYS[1])
        if not token or token ~= ARGV[1] then
            return 0
        end
        redis.call('del', KEYS[1])
        return 1
    """

    # KEYS[1] - lock name
    # ARGV[1] - token
    # ARGV[2] - additional milliseconds
    # ARGV[3] - "0" if the additional time should be added to the lock's
    #           existing ttl or "1" if the existing ttl should be replaced
    # return 1 if the locks time was extended, otherwise 0
    LUA_EXTEND_SCRIPT = """
        local token = redis.call('get', KEYS[1])
        if not token or token ~= ARGV[1] then
            return 0
        end
        local expiration = redis.call('pttl', KEYS[1])
        if not expiration then
            expiration = 0
        end
        if expiration < 0 then
            return 0
        end

        local newttl = ARGV[2]
        if ARGV[3] == "0" then
            newttl = ARGV[2] + expiration
        end
        redis.call('pexpire', KEYS[1], newttl)
        return 1
    """

    # KEYS[1] - lock name
    # ARGV[1] - token
    # ARGV[2] - milliseconds
    # return 1 if the locks time was reacquired, otherwise 0
    LUA_REACQUIRE_SCRIPT = """
        local token = redis.call('get', KEYS[1])
        if not token or token ~= ARGV[1] then
            return 0
        end
        redis.call('pexpire', KEYS[1], ARGV[2])
        return 1
    """

    def __init__(self, redis, name, timeout=None, sleep=0.1,
                 blocking=True, blocking_timeout=None, thread_local=True):
        """
        Create a new Lock instance named ``name`` using the Redis client
        supplied by ``redis``.

        ``timeout`` indicates a maximum life for the lock.
        By default, it will remain locked until release() is called.
        ``timeout`` can be specified as a float or integer, both representing
        the number of seconds to wait.

        ``sleep`` indicates the amount of time to sleep per loop iteration
        when the lock is in blocking mode and another client is currently
        holding the lock.

        ``blocking`` indicates whether calling ``acquire`` should block until
        the lock has been acquired or to fail immediately, causing ``acquire``
        to return False and the lock not being acquired. Defaults to True.
        Note this value can be overridden by passing a ``blocking``
        argument to ``acquire``.

        ``blocking_timeout`` indicates the maximum amount of time in seconds to
        spend trying to acquire the lock. A value of ``None`` indicates
        continue trying forever. ``blocking_timeout`` can be specified as a
        float or integer, both representing the number of seconds to wait.

        ``thread_local`` indicates whether the lock token is placed in
        thread-local storage. By default, the token is placed in thread local
        storage so that a thread only sees its token, not a token set by
        another thread. Consider the following timeline:

            time: 0, thread-1 acquires `my-lock`, with a timeout of 5 seconds.
                     thread-1 sets the token to "abc"
            time: 1, thread-2 blocks trying to acquire `my-lock` using the
                     Lock instance.
            time: 5, thread-1 has not yet completed. redis expires the lock
                     key.
            time: 5, thread-2 acquired `my-lock` now that it's available.
                     thread-2 sets the token to "xyz"
            time: 6, thread-1 finishes its work and calls release(). if the
                     token is *not* stored in thread local storage, then
                     thread-1 would see the token value as "xyz" and would be
                     able to successfully release the thread-2's lock.

        In some use cases it's necessary to disable thread local storage. For
        example, if you have code where one thread acquires a lock and passes
        that lock instance to a worker thread to release later. If thread
        local storage isn't disabled in this case, the worker thread won't see
        the token set by the thread that acquired the lock. Our assumption
        is that these cases aren't common and as such default to using
        thread local storage.
        """
        self.redis = redis
        self.name = name
        self.timeout = timeout
        self.sleep = sleep
        self.blocking = blocking
        self.blocking_timeout = blocking_timeout
        self.thread_local = bool(thread_local)
        self.local = threading.local() if self.thread_local else dummy()
        self.local.token = None
        self.register_scripts()

    def register_scripts(self):
        cls = self.__class__
        client = self.redis
        if cls.lua_release is None:
            cls.lua_release = client.register_script(cls.LUA_RELEASE_SCRIPT)
        if cls.lua_extend is None:
            cls.lua_extend = client.register_script(cls.LUA_EXTEND_SCRIPT)
        if cls.lua_reacquire is None:
            cls.lua_reacquire = \
                client.register_script(cls.LUA_REACQUIRE_SCRIPT)

    def __enter__(self):
        # force blocking, as otherwise the user would have to check whether
        # the lock was actually acquired or not.
        if self.acquire(blocking=True):
            return self
        raise LockError("Unable to acquire lock within the time specified")

    def __exit__(self, exc_type, exc_value, traceback):
        self.release()

    def acquire(self, blocking=None, blocking_timeout=None, token=None):
        """
        Use Redis to hold a shared, distributed lock named ``name``.
        Returns True once the lock is acquired.

        If ``blocking`` is False, always return immediately. If the lock
        was acquired, return True, otherwise return False.

        ``blocking_timeout`` specifies the maximum number of seconds to
        wait trying to acquire the lock.

        ``token`` specifies the token value to be used. If provided, token
        must be a bytes object or a string that can be encoded to a bytes
        object with the default encoding. If a token isn't specified, a UUID
        will be generated.
        """
        sleep = self.sleep
        if token is None:
            token = uuid.uuid1().hex.encode()
        else:
            encoder = self.redis.connection_pool.get_encoder()
            token = encoder.encode(token)
        if blocking is None:
            blocking = self.blocking
        if blocking_timeout is None:
            blocking_timeout = self.blocking_timeout
        stop_trying_at = None
        if blocking_timeout is not None:
            stop_trying_at = mod_time.time() + blocking_timeout
        while True:
            if self.do_acquire(token):
                self.local.token = token
                return True
            if not blocking:
                return False
            next_try_at = mod_time.time() + sleep
            if stop_trying_at is not None and next_try_at > stop_trying_at:
                return False
            mod_time.sleep(sleep)

    def do_acquire(self, token):
        if self.timeout:
            # convert to milliseconds
            timeout = int(self.timeout * 1000)
        else:
            timeout = None
        if self.redis.set(self.name, token, nx=True, px=timeout):
            return True
        return False

    def locked(self):
        """
        Returns True if this key is locked by any process, otherwise False.
        """
        return self.redis.get(self.name) is not None

    def owned(self):
        """
        Returns True if this key is locked by this lock, otherwise False.
        """
        stored_token = self.redis.get(self.name)
        # need to always compare bytes to bytes
        # TODO: this can be simplified when the context manager is finished
        if stored_token and not isinstance(stored_token, bytes):
            encoder = self.redis.connection_pool.get_encoder()
            stored_token = encoder.encode(stored_token)
        return self.local.token is not None and \
            stored_token == self.local.token

    def release(self):
        "Releases the already acquired lock"
        expected_token = self.local.token
        if expected_token is None:
            raise LockError("Cannot release an unlocked lock")
        self.local.token = None
        self.do_release(expected_token)

    def do_release(self, expected_token):
        if not bool(self.lua_release(keys=[self.name],
                                     args=[expected_token],
                                     client=self.redis)):
            raise LockNotOwnedError("Cannot release a lock"
                                    " that's no longer owned")

    def extend(self, additional_time, replace_ttl=False):
        """
        Adds more time to an already acquired lock.

        ``additional_time`` can be specified as an integer or a float, both
        representing the number of seconds to add.

        ``replace_ttl`` if False (the default), add `additional_time` to
        the lock's existing ttl. If True, replace the lock's ttl with
        `additional_time`.
        """
        if self.local.token is None:
            raise LockError("Cannot extend an unlocked lock")
        if self.timeout is None:
            raise LockError("Cannot extend a lock with no timeout")
        return self.do_extend(additional_time, replace_ttl)

    def do_extend(self, additional_time, replace_ttl):
        additional_time = int(additional_time * 1000)
        if not bool(
            self.lua_extend(
                keys=[self.name],
                args=[
                    self.local.token,
                    additional_time,
                    replace_ttl and "1" or "0"
                ],
                client=self.redis,
            )
        ):
            raise LockNotOwnedError(
                "Cannot extend a lock that's" " no longer owned"
            )
        return True

    def reacquire(self):
        """
        Resets a TTL of an already acquired lock back to a timeout value.
        """
        if self.local.token is None:
            raise LockError("Cannot reacquire an unlocked lock")
        if self.timeout is None:
            raise LockError("Cannot reacquire a lock with no timeout")
        return self.do_reacquire()

    def do_reacquire(self):
        timeout = int(self.timeout * 1000)
        if not bool(self.lua_reacquire(keys=[self.name],
                                       args=[self.local.token, timeout],
                                       client=self.redis)):
            raise LockNotOwnedError("Cannot reacquire a lock that's"
                                    " no longer owned")
        return True

相關文章