Asynchronous request cloning using Lua + openresty

Question

I'm working on a project to clone a subset of requests going to our production servers to one or more non-production endpoints. I opted to do this in Nginx and Lua because at the time, I could not find a tool that fit the requirements I needed, namely, only forwarding all requests for a subset of session IDs (as opposed to a percentage of all requests).

Initially, I was simply using ngx.timer.at to dispatch the subrequest, but I was informed that the server could run out of threads if the destination server was unavailable. So I looked at lua-resty-logger-socket and used the code that they use to buffer requests and modified it for my usage.

This is my first Lua project in Nginx, and actually my first time touching Lua years, and nobody at work knows Lua well enough to give any advice, so I'd appreciate if anyone could point out any flaws in my methods or things that can be done better.

example nginx.conf:

location = /GenerateToken {
    set $session_token "";
    access_by_lua_file store_token_from_uri.lua;

    body_filter_by_lua_file store_token_from_body.lua;

    log_by_lua_block {
        local dispatcher = require "dispatcher"
        if not dispatcher.initialized() then
            local ok, err = dispatcher.init {
                servers = {
                    ["127.0.0.1:" .. ngx.var.server_port] = .99
                }
            }
            if not ok then
                ngx.log(ngx.ERR, "failed to initialize the logger: ",
                        err)
                return
            end
        end
        dispatcher.queue_request(ngx.var.session_token)
    }
}

store_token_from_uri.lua

ngx.req.read_body()
-- Extract the token from the URI if it is present
-- add "set $session_token "";" before to initialize variable
ngx.ctx.args = ngx.req.get_uri_args()
-- If the id argument is present, in the query string, that is our token.
if ngx.ctx.args.token then
ngx.var.session_token = ngx.ctx.args.token
elseif ngx.ctx.args.id then
    ngx.var.session_token = ngx.ctx.args.id
end

store_token_from_body.lua

-- Stores token from the response body in an nginx variable
-- Runs in body_filter_by_lua_file
-- For GenerateToken requests, the token must be extracted from the response.
if string.find(ngx.var.uri, '/GenerateToken') then
    -- Build the response body from the buffer
    local resp_body = string.sub(ngx.arg[1], 1, 8000)
    ngx.ctx.buffered = (ngx.ctx.buffered or "") .. resp_body
    if ngx.arg[2] then
        -- For GET requests with a callback, extract the token from the callback response
        if string.find(ngx.var.uri, '/GenerateToken.js') and ngx.ctx.args.callback then
            _, _, ngx.var.session_token = string.find(ngx.ctx.buffered, "(%x+-%x+-%x+-%x+%-%x+)")
            -- Otherwise, the response should just contain the token.
        else
            -- Trim the response of any trailing newlines
            ngx.var.session_token = ngx.ctx.buffered:gsub("^%s*(.-)%s*$", "%1")
        end
        -- Add the token to the query string so the remote server will not re-create a new token
        ngx.ctx.args.token = ngx.var.session_token
        ngx.req.set_uri_args(ngx.ctx.args)
    end
end

dispatcher.lua

-- Add to log_by_lua_file
local http = require "resty.http"

local insert = table.insert
local timer_at = ngx.timer.at
local ngx_log = ngx.log
local type = type
local pairs = pairs
local ipairs = ipairs
local debug = ngx.config.debug
local DEBUG = ngx.DEBUG

-- table.new(narr, nrec)
local succ, new_tab = pcall(require, "table.new")
if not succ then
    new_tab = function () return {} end
end

local _M = new_tab(0, 4)

-- user config
local servers = {}
-- 1MB
local flush_limit = 1048576
-- drop at 50MB
local drop_limit = 52428800
-- reuse buffer for at most 10000 times
local max_buffer_reuse = 10000

local periodic_flush
local need_periodic_flush

local flushing
local dispatcher_initialized
local counter  = 0


-- internal variables
local buffer_size = 0
-- 1st level buffer, it stores incoming requests
local incoming_buffer = new_tab(20000, 0)
-- number of requests in current 1st level buffer, starts from 0
local incoming_buffer_index = 0
-- 2nd level buffer, it stores requests ready to be sent out
local send_buffer = new_tab(1000, 0)
local send_buffer_index = 0
local send_buffer_size = 0

---
-- Populates the second-level buffer
local function _prepare_send_buffer()
    for i=1, incoming_buffer_index do
        send_buffer_index = send_buffer_index + 1
        send_buffer[i] = incoming_buffer[i]
    end

    send_buffer_size = buffer_size
    incoming_buffer_index = 0

    counter = counter + 1
    if counter > max_buffer_reuse then
        incoming_buffer = new_tab(20000, 0)
        send_buffer = new_tab(1000, 0)
        counter = 0
        if debug then
            ngx_log(DEBUG, "request buffer reuse limit (" .. max_buffer_reuse
                    .. ") reached, create a new \"request_buffer_data\"")
        end
    end
end


-- this is expensive and should only be used to tidy up in case of an error
local function _pop_send_buffer(count)
    for i=1, count do
        local packet = send_buffer.remove(i)
        send_buffer_index = send_buffer_index - 1
        send_buffer_size = send_buffer_size - #packet
    end
end

---
-- Performs the actual buffered subrequests
local function _do_flush()
    local ok, err
    local errors = ""
    local all_ok = true

    for i ,request in ipairs(send_buffer) do
        local httpc = http.new()
        httpc:set_timeout(0)
        httpc:connect(request[1], request[2] or 80)
        request[5].host = request[1]

        -- Use send_request instead of request to prevent needing to read the response back.
        ok, err = httpc:send_request{
            method = request[4],
            path = request[3],
            headers = request[5],
            body = request[6]
        }
        if not ok then
            -- ensure we don't resend packets later that we've already sent
            _pop_send_buffer(i - 1)
            all_ok = false
        end
        if err then
            errors = errors .. " ".. err
        end
    end
    return all_ok, errors
end

---
-- Determines if a flush is required
local function _need_flush()
    if incoming_buffer_index > 0 or send_buffer_index > 0 then
        return true
    end
    return false
end

---
-- Locks the buffer flush so only one flush can be happening at once
local function _flush_lock()
    if not flushing then
        if debug then
            ngx_log(DEBUG, "flush lock acquired")
        end
        flushing = true
        return true
    end
    return false
end

---
-- Unlocks the buffer flush once finished
local function _flush_unlock()
    if debug then
        ngx_log(DEBUG, "flush lock released")
    end
    flushing = false
end

---
-- Performs the buffer flush if we can
local function _flush()
    local ok, err

    -- pre check
    if not _flush_lock() then
        if debug then
            ngx_log(DEBUG, "previous flush not finished")
        end
        -- do this later
        return true
    end

    if not _need_flush() then
        if debug then
            ngx_log(DEBUG, "no need to flush:", incoming_buffer_index)
        end
        _flush_unlock()
        return true
    end

    -- start flushing
    if debug then
        ngx_log(DEBUG, "start flushing")
    end

    local bytes
    if incoming_buffer_index > 0 then
        _prepare_send_buffer()
    end

    ok, err = _do_flush()

    _flush_unlock()

    if not ok then
        local err_msg = "buffer flush failed: " .. err
        return nil, err_msg
    else
        if debug then
            ngx_log(DEBUG, "request sent")
        end
    end

    buffer_size = buffer_size - send_buffer_size

    return ok
end

---
-- If periodic flush is enabled, determines if a flush is needed. A periodic flush
-- will only happen if no flush had occured during the waiting period between periodic
-- flushes
local function _periodic_flush()
    if need_periodic_flush then
        -- no regular flush happened after periodic flush timer had been set
        if debug then
            ngx_log(DEBUG, "performing periodic flush")
        end
        _flush()
    else
        if debug then
            ngx_log(DEBUG, "no need to perform periodic flush: regular flush "
                    .. "happened before")
        end
        need_periodic_flush = true
    end

    timer_at(periodic_flush, _periodic_flush)
end

---
-- Performs the flush of the request buffer. The flush is done via
-- a 0-second timer so it can be performed in the background
local function _flush_buffer()
    local ok, err = timer_at(0, _flush)

    need_periodic_flush = false

    if not ok then
        return nil, err
    end
end

---
-- Writes the subrequest to be made to the buffer
-- @param hostname
-- @param port
-- @param uri
-- @param method
-- @param headers
-- @param body
-- @param req_length
local function _write_buffer(hostname, port, uri, method, headers, body, req_length)
    incoming_buffer_index = incoming_buffer_index + 1
    incoming_buffer[incoming_buffer_index] = {
        hostname,
        port,
        uri,
        method,
        headers,
        body,
        req_length
    }

    buffer_size = buffer_size + req_length

    return buffer_size
end

---
-- Queue up this request to be replayed if the token passes the deterministic selector
-- @param token
function _M.queue_request(token)
    if not dispatcher_initialized then
        return false
    end

    local bytes
    -- If we have a token stored, clone this request to other servers
    if token and token ~= "" then
        for hostname, percentage in ipairs(servers) do
            -- Deterministic way of selecting a percentage of tokens. It takes the last 4 bytes of the
            -- uuid, converts it to an integer, and modulos it by 10K, (which should generate a number
            -- between 0 and 9999). If it's less than the configured percentage for this host multiplied
            -- by 10K, then this request is a candidate for cloning. We use 10K as this would give us a
            -- precision of .01%
            if (tonumber(token:sub(-8), 16) % 10000) < (percentage * 10000) then
                -- Slit the IP and port up
                local remote_addr = {}
                for part in string.gmatch(hostname, '([^:]+)') do
                    table.insert(remote_addr, part)
                end

                -- Append the query string if we have it, othwerise just use the URI
                local uri = ngx.var.args and ngx.var.uri .. "?" .. ngx.var.args or ngx.var.uri
                local method = ngx.req.get_method()
                local headers = ngx.req.get_headers()
                local body = ngx.req.get_body_data()
                local req_length = ngx.var.request_length

                if (debug) then
                    ngx.update_time()
                    ngx_log(DEBUG, ngx.now(), ":request length: " .. req_length)
                end


                -- Flush if we're beyond the flush limit or drop limit but if this request will put us beyond thee
                -- drop limit, then don't queue up this request.
                if (req_length + buffer_size < flush_limit) then
                    _write_buffer(remote_addr[1], remote_addr[2], uri, method, headers, body, req_length)
                    return req_length
                elseif (req_length + buffer_size <= drop_limit) then
                    _write_buffer(remote_addr[1], remote_addr[2], uri, method, headers, body, req_length)
                    _flush_buffer()
                    return req_length
                else
                    -- this log message doesn't fit in buffer, drop it
                    _flush_buffer()
                    if (debug) then
                        ngx_log(DEBUG, "request buffer is full, this request will be dropped")
                    end
                    return false
                end
            end
        end
    end
end

---
-- Initializes the dispatcher module, validates and applies the configuration settings passed
-- @param config
function _M.init(config)
    if (type(config) ~= "table") then
        return nil, "user_config must be a table"
    end

    for k, v in pairs(config) do
        if k == "servers" then
            if type(v) ~= "table" then
                return nil, '"servers" must be a table'
            end
            servers = v
        elseif k == "flush_limit" then
            if type(v) ~= "number" or v < 0 then
                return nil, 'invalid "flush_limit"'
            end
            flush_limit = v
        elseif k == "drop_limit" then
            if type(v) ~= "number" or v < 0 then
                return nil, 'invalid "drop_limit"'
            end
            drop_limit = v
        elseif k == "max_buffer_reuse" then
            if type(v) ~= "number" or v < 0 then
                return nil, 'invalid "max_buffer_reuse"'
            end
            max_buffer_reuse = v
        elseif k == "periodic_flush" then
            if type(v) ~= "number" or v < 0 then
                return nil, 'invalid "periodic_flush"'
            end
            periodic_flush = v
        end
    end

    if (flush_limit >= drop_limit) then
        return nil, "\"flush_limit\" should be < \"drop_limit\""
    end

    flushing = false

    dispatcher_initialized = true

    if periodic_flush then
        if debug then
            ngx_log(DEBUG, "periodic flush enabled for every " .. periodic_flush .. " seconds")
        end
        need_periodic_flush = true
        timer_at(periodic_flush, _periodic_flush)
    end

    return dispatcher_initialized
end

function _M.initialized()
    return dispatcher_initialized
end

_M.flush = _flush

return _M

Answer 1

The ngx.var.session_token can be easily setup in store_token_from_uri.lua as follows:

ngx.var.session_token = ngx.ctx.args.token or ngx.ctx.args.id or ngx.var.session_token

where, the value of ngx.var.session_token will be restored in case both ngx.ctx.args.token and ngx.ctx.args.id are nil. The if-elseif blocks are not required anymore.

---

If you do not want the start and end positions when using string.find, use the string.match function instead. The same can be used instead of your gsub call later in store_token_from_body.lua file. The characters - and . in a lua-pattern have special meaning. Escape them using a % character for a literal matching. Therefore:

if string.find(ngx.var.uri, '/GenerateToken') then
    local resp_body = ngx.arg[1]:sub(1, 8000)
    ngx.ctx.buffered = (ngx.ctx.buffered or "") .. resp_body
    if ngx.arg[2] then
        -- For GET requests with a callback, extract the token from the callback response
        if string.find(ngx.var.uri, '/GenerateToken%.js') and ngx.ctx.args.callback then
            ngx.var.session_token = ngx.ctx.buffered:match "(%x+%-%x+%-%x+%-%x+%-%x+)"
            -- Otherwise, the response should just contain the token.
        else
            -- Trim the response of any trailing newlines
            ngx.var.session_token = ngx.ctx.buffered:match "^%s*(.-)%s*$"
        end
        -- Add the token to the query string so the remote server will not re-create a new token
        ngx.ctx.args.token = ngx.var.session_token
        ngx.req.set_uri_args(ngx.ctx.args)
    end
end

---

In the dispatcher.lua file, the code snippet for debug logging, can be put inside a separate function:

function debug_log( msg, ... )
    if debug then
        ngx_log(DEBUG, msg, ...)
    end
end

In the _do_flush function, you are concatenating a string buffer repeatedly. Instead, store all errors inside a table, and return table.concat(errors, " "):

local errors = {}
...
    if err then
        table.insert(errors, err)
    end
...
return all_ok, table.concat(errors, " ")