使用 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