url shortener base on nginx lua and resty redis

urlshortener.lua

local _M = {
    _VERSION = '0.1'
}

local connect = function()
    local redis = require "resty.redis"
    local red = redis:new()
    red:set_timeout(1000) -- 1 sec
    local ok, err = red:connect("127.0.0.1", 6379)
    if not ok then
        ngx.say(err)
        return ngx.exit(ngx.ERROR)
    end

    return red
end

local basen = function(n)
    local digits = "0123456789abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ"
    local t = {}
    repeat
        local d = (n % 62) + 1
        n = math.floor(n / 62)
        table.insert(t, 1, digits:sub(d, d))
    until n == 0
    return table.concat(t,"");
end

-- ref https://github.com/AlexChittock/SimpleShortener
-- ref https://gist.github.com/MendelGusmao/2356310
-- "database scheme"
-- database 0: id ~> url
-- database 1: id ~> hits
-- database 2: id ~> [{referer|user_agent}]
-- database 3: id ~> hits (when id is not found)
-- database 4: id ~> [{referer|user_agent}] (when id is not found)
-- database 5: key "shorten.count" storing the number of shortened urls;
--             the id is generated by (this number + 1) converted to base 62
-- database 6: url md5 ~> id

function _M.pack(urlPrefix, checkAlready)
    local short_pre = urlPrefix or ("http://" .. ngx.var.host ..
                         ":" .. ngx.var.server_port .. "/")

    -- See if we have the url already
    local url = ngx.var.arg_url
    local hash = ngx.md5(url)
    local red = connect()

    if checkAlready or true then
        red:select(6)
        local id, err = red:get(hash)
        if id ~= ngx.null then
            ngx.say(short_pre .. id)
            return ngx.exit(ngx.OK)
        end
    end

    red:select(5)
    local count, err = red:incr("shorten.count")
    local id, err = basen(count)

    if checkAlready or true then
        red:select(6)
        red:set(hash, id)
    end

    red:select(0)
    red:set(id, url)

    ngx.say(short_pre .. id)
    return ngx.exit(ngx.OK)
end

function _M.unpack(notFoundRedirect, recordHits, recordReferAndUserAgent)
    local id = string.sub(ngx.var.request_uri, 2) -- remove prefix /

    local red = connect()
    red:select(0)
    local url, err = red:get(id)
    if not url then
        ngx.say(err)
        return ngx.exit(ngx.ERROR)
    end

    local referer = ngx.var.http_referer or ""
    if url == ngx.null then -- not found
        if recordHits then
            red:select(3)
            red:incr(id)
        end

        if recordReferAndUserAgent then
            red:select(4)
            red:rpush(id, referer .. "|" .. ngx.var.http_user_agent)
        end

        if notFoundRedirect then
            return ngx.redirect(notFoundRedirect)
        end

        return ngx.exit(ngx.HTTP_NOT_FOUND)
    end

    -- found
    if recordHits then
        red:select(1)
        red:incr(id)
    end

    if recordReferAndUserAgent then
        red:select(2)
        red:rpush(id, referer .. "|" .. ngx.var.http_user_agent)
    end

    return ngx.redirect(url)
end

return _M

referred in nginx configuration:

lua_package_path ";;$prefix/conf/?.lua;";

server {
    listen       17001;
    server_name  localhost;
    default_type text/plain;

    location = /shorten {
        add_header Content-Type text/plain;
        content_by_lua '
            local shorten = require "n3r.urlshortener"
            shorten.pack()
        ';
    }

    location ~ "^/[0-9a-zA-Z]{1,5}$" {
        content_by_lua '
            local shorten = require "n3r.urlshortener"
            local notFoundRedirect = "www.10010.com"
            shorten.unpack(notFoundRedirect)
        ';
    }
}

the pack function can be implemented with redis eval function and got two goodies: the first is the perfermance is promoted by reducing the interaction with redis server, the second is the url existance check will always be valid in favored by the whole eval script execution within a single transaction.