Redis应用之Lua脚本

最近工作中遇到这样一个需求：就是用户每天对某个接口的调用次数不得超过一定的频次。

回想到之前看到的文章中有处理类似问题的场景：就是实现访问频率的限制，这类问题的使用场景特别多。

处理这类问题，使用Redis + Lua是一个非常正确的做法。

Redis + Lua介绍

Lua嵌入 Redis 的优势:

减少网络开销: 不使用  Lua  的代码需要向 Redis 发送多次请求, 而脚本只需一次即可, 减少网络传输
原子操作: Redis 将整个脚本作为一个原子执行, 无需担心并发, 也就无需事务
复用: 脚本会永久保存 Redis 中, 其他客户端可继续使用

基本概念及使用方法

Redis执行Lua脚本的基础知识，官网的文档介绍的很详细，这里不再赘述。通过如下两篇文章基本可以对Redis的概念和使用做到很好的了解。

Redis设计与实现—Lua脚本

Redis EVAL命令

Redis执行Lua脚本的过程

这里重点对Redis服务执行Lua脚本的实现原理进行介绍一下。

创建Lua环境

A:为了在 Redis 服务器中执行 Lua 脚本, Redis 内嵌了一个 Lua 环境, 并对该环境进行了一系列修改, 从而确保满足 Redis 的需要.
B:整个 Redis 服务器只需创建一个 Lua 环境.

Lua环境协作组件

Redis创建两个用于与Lua环境协作的组件: 伪客户端 - 负责执行 Lua 脚本中的 Redis 命令, lua_scripts 字典 - 保存 Lua 脚本。

伪客户端 ：执行Reids命令必须有对应的客户端状态, 因此执行 Lua 脚本内的 Redis 命令 必须为 Lua 环境专门创建一个伪客户端, 由该客户端处理 Lua 内所有命令: redis.call() / redis.pcall() 。

lua_scritps 字典：字典key为脚本 SHA1 校验和, value为 SHA1 对应脚本内容, 所有被 EVAL 和 SCRIPT LOAD 载入过的脚本都被记录到 lua_scripts 中, 便于实现 SCRIPT EXISTS 命令和脚本复制功能。

Eval命令原理

EVAL 命令执行分为以下三个步骤：

1. 定义Lua函数:

   在 Lua 环境内定义 Lua函数 : 名为 f_ 前缀+脚本 SHA1 校验和, 体为 脚本内容本身 . 优势:

   1、执行脚本步骤简单, 调用函数即可;

   2、函数的局部性可保持 Lua 环境清洁, 减少垃圾回收工作量, 且避免使用全局变量;

   3、只要记住 SHA1 校验和, 即可在不知脚本内容的情况下, 直接调用 Lua 函数执行脚本( EVALSHA 命令实现).

2. 将脚本保存到 lua_scripts 字典;

3. 执行脚本函数:

   1、将 EVAL 命令中输入的 KEYS 参数和 ARGV 参数以全局数组的方式传入到 Lua 环境中。

   2、设置伪客户端的目标数据库为调用者客户端的目标数据库： fake_client->db = caller_client->db ，确保脚本中执行的 Redis 命令访问的是正确的数据库。

   3、为 Lua 环境装载超时钩子，保证在脚本执行出现超时时可以杀死脚本，或者停止 Redis 服务器。

   4、执行脚本对应的 Lua 函数。

   5、如果被执行的 Lua 脚本中带有 SELECT 命令，那么在脚本执行完毕之后，伪客户端中的数据库可能已经有所改变，所以需要对调用者客户端的目标数据库进行更新： caller_client->db = fake_client->db 。

   6、执行清理操作：清除钩子；清除指向调用者客户端的指针；等等。

   7、将 Lua 函数执行所得的结果转换成 Redis 回复，然后传给调用者客户端。

   8、对 Lua 环境进行一次单步的渐进式 GC 。

数据流表示图

执行命令 EVAL "return redis.call('DBSIZE')" 0 时，调用者客户端（caller）、伪客户端（fake client）、Redis 服务器和 Lua 环境之间的数据流表示图：

Redis + Lua 应用

这里主要就是结合 Python 开发环境，针对指定时间内IP请求频次限制做一个示例讲解。Python的Redis客户端实现执行Lua脚本的相关功能。客户端执行Lua脚本的源码如下：

classScript(object):
    "An executable Lua script object returned by ``register_script``"

    def__init__(self, registered_client, script):
        self.registered_client = registered_client
        self.script = script
        self.sha = ''

    def__call__(self, keys=[], args=[], client=None):
        "Execute the script, passing any required ``args``"
        if client is None:
            client = self.registered_client
        args = tuple(keys) + tuple(args)
        # make sure the Redis server knows about the script
        if isinstance(client, BasePipeline):
            # make sure this script is good to go on pipeline
            client.script_load_for_pipeline(self)
        try:
            return client.evalsha(self.sha, len(keys), *args)
        except NoScriptError:
            # Maybe the client is pointed to a differnet server than the client
            # that created this instance?
            self.sha = client.script_load(self.script)
            return client.evalsha(self.sha, len(keys), *args)

Lua脚本实现的IP频次限制代码

--
-- Created by IntelliJ IDEA.
-- User: DragonRiver
-- Date: 17/03/11
-- Time: 下午6:11
--

local key = "rate:limit:" .. KEYS[1]
local limit = tonumber(ARGV[1])
local expire_time = ARGV[2]
-- 首先判断键是否存在
local is_exists = redis.call("EXISTS", key)
-- 键存在就加1
if is_exists == 1 then
    if redis.call("INCR", key) > limit then
        return 0
    else
        return 1
    end
-- 不存在的话，就设置键，同时设置有效期
else
    redis.call("SET", key, 1)
    redis.call("EXPIRE", key, expire_time)
    return 1

Python执行乱脚本helper

import redis

classLuaRedisHelper(object):
    def__init__(self, redis_conf, lua_script_path):
        self.redis_client = redis.StrictRedis(
            host=redis_conf.get("host", "localhost"),
            port=redis_conf.get("port", 6379), db=0)
        self.lua_script = open(lua_script_path).read()
        self.multiply = self.redis_client.register_script(self.lua_script)

    defmultiply_excete(self, keys=list(), args=list()):
        return self.multiply(keys=keys, args=args)

实际应用

#!/usr/bin/env python
# -*-coding:utf-8-*-

"""
    实现访问频率限制: 实现访问者 $ip 在一定的时间 $time 内只能访问 $limit 次.
"""
from lua_script_helper import LuaRedisHelper

if __name__ == "__main__":
    redis_conf = {}
    lua_redis_helper = LuaRedisHelper(
        redis_conf=redis_conf, lua_script_path="./ip_limit.lua")
    keys, args = ["127.0.0.1"], ["5", "10"]
    res = lua_redis_helper.multiply(keys, args)
    print res

Redis + Lua的其他应用请访问如下资源：

DraognRiver Github