lua-counter example

counter.lua at tip
anonymous

counter.lua at tip

File counter.lua from the latest check-in


-- module counter - counts some junk

local _counter = {}

local redis = require 'resty.redis'
local cjson = require 'cjson'

local closed = false

function _counter.exit()
	if not closed then
		local ok, err = _counter.red:set_keepalive(60000, 10)
		if not ok then
			ngx.log(ngx.ERR, "redis keepalive: " .. err)
		end
		closed = true
	end
end

function _counter.decode_bulk(reply)
	-- decodes values from array into dict
	local data = {}
	for j=1, #reply, 2 do
		data[reply[j] ] = reply[j+1]
	end
	return data
end

function _counter.connect()
	-- connect to redis
	_counter.red = redis:new()
	local ok, err = _counter.red:connect('127.0.0.1', 6379)
	if not ok then
		ngx.log(ngx.ERR, 'redis connection failed: ', err)
		ngx.exit(ngx.HTTP_INTERNAL_SERVER_ERROR)
	end
	closed = false
end

function _counter.check_schema()
	if ngx.var.site_schema == nil then
		ngx.log(ngx.ERR, "No 'site_schema' specified");
		ngx.exit(ngx.HTTP_INTERNAL_SERVER_ERROR)
	end
end

function _counter.update()
	_counter.check_schema()
	local today = os.date('*t')
	local timestamp = os.time{year = today['year'], month = today['month'], day = today['day']}
	local referer = '-';

	if ngx.var.http_referer ~= nil then
		_, _, referer = string.find(ngx.var.http_referer, '^https?://([^?&]+)')
		--ngx.log(ngx.ERR, "Referer: " .. referer)
	end

	local key = ngx.var.site_schema .. '_counter_' .. timestamp .. '_' .. referer
	-- ngx.log(ngx.ERR, "Using key: " .. key)

	local res, err
	_counter.red:init_pipeline(4)
	_counter.red:multi()
	_counter.red:hincrby(key, 'today', 1)
	_counter.red:hgetall(key)
	_counter.red:exec()
	res, err = _counter.red:commit_pipeline()
	if not res then
		ngx.log(ngx.ERR, 'redis pipeline failed: ', err)
		ngx.exit(ngx.HTTP_INTERNAL_SERVER_ERROR)
	end

	local data = _counter.red:array_to_hash(res[4][2])
	-- ngx.log(ngx.ERR, "Got data: " .. cjson.encode(data))

	if tonumber(data.today) == 1 then
		ngx.log(ngx.ERR, "Reading postgres")
		-- postgres fallback
		result = ngx.location.capture('/postgres', {
			method = ngx.HTTP_PUT,
			body = "select * from " .. ngx.var.site_schema .. ".get_stats('" .. referer .. "') as (today int, lastday int, week bigint, whole bigint);"
		})
		if result.status ~= 200 or not result.body then
			ngx.log(ngx.ERR, 'postgres access failed')
			ngx.exit(ngx.HTTP_INTERNAL_SERVER_ERROR)
		else
			local unrds = require "rds.parser"
			local res, err = unrds.parse(result.body)

			if res == nil then
				ngx.log(ngx.ERR, 'failed to obtain data: ' .. err)
				ngx.exit(ngx.HTTP_INTERNAL_SERVER_ERROR)
			else
				data = res.resultset[1]
				data.today = data.today + 1
				ngx.log(ngx.ERR, "Got data: " .. cjson.encode(data))

				_counter.red:init_pipeline(4)
				_counter.red:multi()
				_counter.red:hmset(key, data)
				_counter.red:expire(key, 129600)
				_counter.red:exec()
				res, err = _counter.red:commit_pipeline()
				if not res then
					ngx.log(ngx.ERR, 'redis pipeline failed: ', err)
					ngx.exit(ngx.HTTP_INTERNAL_SERVER_ERROR)
				end
			end
		end
	end

	ngx.say(cjson.encode(data))
	ngx.eof()

	local uid = ''

	if ngx.var.uid_got ~= nil and string.find(ngx.var.uid_got, 'uid=') == 1 then
		uid = string.sub(ngx.var.uid_got, 5)
	elseif ngx.var.uid_set ~= nil and string.find(ngx.var.uid_set, 'uid=') == 1 then
		uid = string.sub(ngx.var.uid_set, 5)
	end

	local hit_key = uid .. '_' .. referer .. '_' .. ngx.var.remote_addr
	local key_pending = ngx.var.site_schema .. '_counter_pending'

	_counter.red:init_pipeline(4)
	_counter.red:multi()
	_counter.red:hincrby(key_pending, hit_key, 1)
	_counter.red:expire(key_pending, 604800)
	_counter.red:exec()
	res, err = _counter.red:commit_pipeline()
	if not res then
		ngx.log(ngx.ERR, 'redis transaction failed: ', err)
		ngx.exit(ngx.HTTP_INTERNAL_SERVER_ERROR)
	end
end

return _counter