使用 Redis 实现类表级锁和类行级锁

使用 Redis 实现类表级锁和类行级锁

摘要

Redis 本身带有加锁的功能，无论是使用 redis.set(key, value, nx=True) 还是 with redis.lock(key) 都可以方便的对指定的代码块加锁。但现在我们面临一个新的需求，有两组进程，A 组进程会访问全部的数据，并对这些数据做批量处理，而 B 组进程则会分别针对某个数据做单独处理。现在我们希望 A, B 两组进程之间互斥，A 组进程内部互斥，但 B 组进程在处理不同的数据时不互斥。概念上类似于数据库的表级锁和行级锁，但可惜并不能直接用数据库锁。

流程

首先我们需要定义一个 namespace, 用以确定 A 和 B 两组进程的作用范围。在此基础上，加锁的逻辑应该是这样的：

A:

• 尝试获取 namespace 下的所有 keys. 如果存在，说明 B 组进程在运行，需要等待
• 如果 B 组进程清空，则继续尝试设置 key, 这里仍然可以用 redis.set(key, value, nx=True).
• 如果设置成功，说明加锁成功，进入代码执行。否则返回第一步

B:

• 尝试获取 namespace 本身，如果存在，说明A组进程在运行，需要等待
• 如果 A 组进程清空，则继续尝试设置 key, 在 namespace 后加上自己的标记，比如数据的行号。如果有 ID 就更好了
• 如果设置成功，说明加锁成功，进入代码执行。否则返回第一步

通过上述的流程，我们保证了 A, B 两组进程间互斥，A 组进程内部互斥，B 组进程内部在处理相同数据时互斥，但 B 组进程在处理不同数据时是可以并发的。

实现

考虑到原子性，我们需要使用 redis 的 lua 脚本来实现上述步骤。同时，上述步骤仅需要改动加锁部分，而释放部分并不需要做改动，因此我们可以就 redis 库本身提供的 Lock 类做扩展。具体实现如下：

class LockExt(Lock):
    lua_acquire = None

    LUA_ACQUIRE_SCRIPT = """
        if redis.call('set', KEYS[1], ARGV[1], 'PX', ARGV[2], 'NX') then
            return 1
        end
        return 0
    """

    def register_scripts(self):
        """
        额外注册我们自己的 ACQUIRE 脚本
        """
        super(LockExt, self).register_scripts