今天在 openresty 的中文邮件列表看到有同学问了一个这样的问题, openresty 的上传限速方案. 他的问题描述是这样的:

由于业务需要，现阶段需开发一个限制客户端上传速度的模块，在网上看了很多资料，有基于Nginx 第三方模块实现的，有使用内置的Nginx模块的，自己认为是可以通过OpenResty内置的API去实现的，但是不确定会不会有问题，就想问一下各位，有没有什么好的解决方案？多谢！

个人感觉挺疑惑的, 为什么上传还需要限速呢? 不过不考虑目的, 也可以探讨一下功能方面的实现到底应该怎么办. 于是我找了一些 nginx 限制上传速度的实现, 不过结果主要是 limit_rate 这个指令, 这个指令是用来限制回复的速度的. 另一个介绍的比较多的是 client_max_body_size, 这个指令是用来限制客户端上传的文件大小. 和功能描述距离比较远. 最后找到了一个已经不活跃的模块, limit_upload_rate, 这个 nginx 的 c 模块提供了一些指令达到限速的目的, 瞄了一眼源码, 了解到应该是在 input filter 阶段做上传限速这个功能.

static ngx_http_input_body_filter_pt  ngx_http_next_input_body_filter;


static ngx_int_t
ngx_http_limit_upload_input_body_filter(ngx_http_request_t *r, ngx_buf_t *buf)
{
    off_t                          excess;
    ngx_int_t                      rc;
    ngx_msec_t                     delay;
    ngx_http_core_loc_conf_t      *clcf;
    ngx_http_limit_upload_ctx_t   *ctx;
    ngx_http_limit_upload_conf_t  *llcf;

    rc = ngx_http_next_input_body_filter(r, buf);
    if (rc != NGX_OK) {
        return rc;
    }
    ......
}

知道应该在什么阶段做这个事情之后, 应该就去找 openresty 提供的 api 了, 在 openresty 邮件列表找到了这篇帖子, 春哥在这篇帖子里面解释了由于 ngx_lua 的内在机制, 并没有实现跟 nginx 完全相同的 input filter 阶段, 但是也提供了另外的思路, 就是利用 ngx.req.socket 与 ngx.req.init_body(), ngx.req.append_body(), 和 ngx.req.finish_body() 这几个 api 结合, 可以模拟出 input filter 相同的功能, 知道了相关的 api 之后, 实现的话, 就没有太大的难度了. 具体的 api 介绍在 lua-nginx-module , 我不再介绍重复的内容.

附上我的实现代码, 如果你想要试一下的话可以看我的 github仓库, 我在这里阐述一下 api 文档里面找不到的东西.

  • 关于 sock:receive 返回的 err == “closed” 判断

这个示例是我通过在 lua-nginx-module 仓库搜索 ngx.req.init_body 得到的信息,  044-req-body.t 这是一个测试文件, 你可以在 1216 行找到这个判断, 然后我通过 postman, wireshark 抓包, 验证了这个 err 应该就是 openresty 提示请求体读取完毕, 因为我在测试过程中, 并没有发现 postman 有关闭写端, 发送 FIN 包, 这个连接也是 keepalive 的, 所以这个用法应该是正确的.

  • 关于限流算法

这个限流算法比较简单, 累加接收到的数据大小, 然后计算, 这个数据量, 在当前的限速条件下, 应该需要传输多少秒, 这是期待值, 用期待值与真正流逝的时间作差, 这样的话, 我们判断两者的差值, 当期待值大于当前当前花费的时间时, 就应该直接挂起差值这么久, 而期待值小于流逝的时间时, 代表客户端并没有跑满我们的限速, 可以继续读取.

  • 关于参数的调整

sock:receive 的参数是 1 * 1024, 也就是每次读取 1kb, 可能会导致 ngx.req.append_body 的调用次数比较多, 应该可以根据精度要求调整这个值.

ngx.sleep 事实上在需要 sleep 时, 挂起时间可以比期待值与实际时间的差值, 也就是代码中的 interval 稍大, 因为客户端的速度已经快于我们的限速了, 我们可以适当的加上一点惩罚值, 来平衡它之后的速度, 如果我们只是挂起一个恰当好的时间, 它之后还是会因为超速被挂起. 加上惩罚值之后, 从性能方面考虑的话, 对于一个大文件, 应该可以节省很多次 ngx.sleep 的调用.

  • 关于时间精度

ngx.now 返回的是 nginx 缓存的时间, 可能也会对精度有一定的影响. 如果想要更精确, 可以在 ngx.now 之前, 强制 nginx 更新时间. 不过个人感觉没有必要.

# nginx.conf
worker_processes 1;

events {
    worker_connections 1024;
}

http {
    server {
        listen 8888;
        client_max_body_size 100M;
        client_body_buffer_size 10K;
        access_by_lua_block {
            local sock = ngx.req.socket()
            -- limit upload rate in ? kb/second
            local function limit_recv_body(rate)
                local rate_in_byte = rate * 1024
                local function next_data_chunk()
                    local start = ngx.now()
                    local size = 0
                    while true do
                        local chunk, err = sock:receive(1 * 1024)
                        if not chunk then
                            if err == "closed" then
                                break
                            else
                                ngx.say( "failed to read body: ", err)
                                break
                            end
                        else
                            size = size + #chunk
                            coroutine.yield(chunk)
                            -- rate limit here
                            -- real time goes by
                            local delta = ngx.now() - start
                            -- how long we should take to receive `size` of bytes
                            local expected = size / rate_in_byte
                            -- if expected larger than delta, we should sleep for a while
                            local interval = expected - delta
                            if interval > 0 then
                                ngx.sleep(interval)
                            end
                        end
                    end
                end
                return coroutine.wrap(
                    function() next_data_chunk() end
                )
            end
            ngx.req.init_body(128 * 1024)
            -- iter req body, limit upload speed in 200kb/s
            for chunk in limit_recv_body(200) do
                ngx.req.append_body(chunk)
            end
            ngx.req.finish_body()
        }
        location / {
            content_by_lua_block {
                ngx.say("hello world");
            }
        }
    }
}