Sends stats about a gearmand process to StatsD

gearman_statsd.lua

#!/usr/bin/lua
-- Requires lua 5.1+ and lua socket
--
-- Copyright (c) 2015, Brian Moon 
-- All rights reserved. 
-- 
-- Redistribution and use in source and binary forms, with or without 
-- modification, are permitted provided that the following conditions are met: 
-- 
--  * Redistributions of source code must retain the above copyright notice, 
--    this list of conditions and the following disclaimer. 
--  * Redistributions in binary form must reproduce the above copyright 
--    notice, this list of conditions and the following disclaimer in the 
--    documentation and/or other materials provided with the distribution. 
--  * Neither the name of DealNews.com, Inc. nor the names of its contributors 
--    may be used to endorse or promote products derived from this software 
--    without specific prior written permission. 
-- 
-- THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" 
-- AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE 
-- IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE 
-- ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE 
-- LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR 
-- CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF 
-- SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS 
-- INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN 
-- CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) 
-- ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE 
-- POSSIBILITY OF SUCH DAMAGE. 

require "socket"

function show_help()
    print("Gearmand Stats Collector for StatsD. Metrics are sent to")
    print("StatsD as guage type stats.")
    print(" -d          Turns on debug mode. No stats sent to StatsD. All stats")
    print("             are printed to standard out.")
    print(" -H HOSTNAME Host name for the StatsD server. default: "..statsd_host)
    print(" -h          Shows this help")
    print(" -i PID_FILE gearmand PID file")
    print(" -n HOSTNAME Hostname to add to the stat. If not provided, the")
    print("             application will try to determine it from the system")
    print("             Currently it is "..hostname.." for this server")
    print(" -p PORT     Port of the StatsD server. 8125 by default")
    print(" -s STAT     Prefix for the statsd stat. default: "..statsd_path)
    print(" -t TIMEOUT  Timeout for sending stats to StatsD. default: "..timeout)
    print(" -v          Makes the output verbose")
    os.exit(0)
end

function main (...)

    verbose = false
    debug = false
    pid_file = ""
    statsd_host = "localhost"
    statsd_port = 8125
    statsd_path = "application.gearman.daemons"
    timeout = 10
    hostname = exec("hostname"):gsub("^%s*(.-)%s*$", "%1")

    local proc_vars = {VmRSS = "cur_mem", FDSize = "max_fd", VmHWM = "max_mem"}

    for opt, arg in getopt("i:H:p:vs:t:n:dh", ...) do
        if     opt == "h" then
            show_help()
        elseif opt == "v" then
            verbose = true
        elseif opt == "d" then
            debug = true
        elseif opt == "i" then
            pid_file = arg
        elseif opt == "H" then
            statsd_host = arg
        elseif opt == "p" then
            statsd_port = arg:match("%d+") + 0
        elseif opt == "s" then
            statsd_path = arg
        elseif opt == "t" then
            timeout = arg:match("%d+") + 0
        elseif opt == "n" then
            hostname = arg
        end
    end

    if verbose then
        print("Timeout:     "..timeout)
        print("Hostname:    "..hostname)
        print("StatsD Host: "..statsd_host)
        print("StatsD Port: "..statsd_port)
    end

    udp = socket.udp()
    udp:setpeername(statsd_host, statsd_port)

    while 1 do

        local start_clock = 0

        if verbose then
            start_clock = socket.gettime()
        end

        local pid = 0

        if pid_file ~= "" then
            pid = io.input(pid_file, "r"):read("*n")
        else
            local stat = exec("svstat /service/gearmand/")
            pid = stat:match("pid %d+"):match("%d+")
        end

        local proc_file = "/proc/"..pid.."/status"

        if verbose then
            print("PID is "..pid)
            print("Reading from "..proc_file)
        end

        -- read data from the proc status file
        local proc_handle = io.open(proc_file)
        if proc_handle ~= nil then
            for line in proc_handle:lines() do
                local prefix = line:match("%a+")
                local value = line:match("%d+")
                if proc_vars[prefix] ~= nil then
                    statsd_send(proc_vars[prefix], value)
                end
            end
            proc_handle:close()
        elseif verbose then
            print("Failed to read from "..proc_file)
        end

        -- read data from lsof
        local total_fd = -1; -- the first line is a header
        local lsof_handle = io.popen("lsof -p "..pid)
        if lsof_handle ~= nil then
            local tcp_connections = {}
            local total_tcp_conn = 0
            for line in lsof_handle:lines() do
                total_fd = total_fd + 1
                if total_fd ~= 0 then
                    if line:find("TCP") ~= nil then
                        total_tcp_conn = total_tcp_conn + 1
                        local con_status = ""
                        if line:find("%(%a+%)") ~= nil then
                            con_status = line:match("%(%a+%)"):match("%a+"):lower()
                        else
                            con_status = "unknown"
                        end


                        if tcp_connections[con_status] == nil then
                            tcp_connections[con_status] = 0
                        end
                        tcp_connections[con_status] = tcp_connections[con_status] + 1
                    end
                end
            end
            lsof_handle:close()

            statsd_send("cur_tcp", total_tcp_conn)
            statsd_send("cur_fd", total_fd)

            for con_status, value in pairs(tcp_connections) do
                statsd_send("cur_tcp_"..con_status, value)
            end

        elseif verbose then
            print("Failed to read data from lsof")
        end

        if verbose then
            local elapsed = socket.gettime() - start_clock
            print("time for checks: "..elapsed.." secs")
        end

        socket.sleep(timeout)
    end

end

function statsd_send(metric, value)
    local string = statsd_path.."."..hostname.."."..metric..":"..value.."|g"
    if debug then
        print(string)
    else
        udp:send(string)
    end
end

function exec(command)
    local handle = io.popen(command)
    local result = handle:read("*a")
    handle:close()
    return result
end

--
-- getopt(":a:b", ...) -- works just like getopt(3).
--
-- Send bug reports to william@25thandClement.com.
--
function getopt(optstring, ...)
    local opts = { }
    local args = { ... }

    for optc, optv in optstring:gmatch"(%a)(:?)" do
        opts[optc] = { hasarg = optv == ":" }
    end

    return coroutine.wrap(function()
        local yield = coroutine.yield
        local i = 1

        while i <= #args do
            local arg = args[i]

            i = i + 1

            if arg == "--" then
                break
            elseif arg:sub(1, 1) == "-" then
                for j = 2, #arg do
                    local opt = arg:sub(j, j)

                    if opts[opt] then
                        if opts[opt].hasarg then
                            if j == #arg then
                                if args[i] then
                                    yield(opt, args[i])
                                    i = i + 1
                                elseif optstring:sub(1, 1) == ":" then
                                    yield(':', opt)
                                else
                                    yield('?', opt)
                                end
                            else
                                yield(opt, arg:sub(j + 1))
                            end

                            break
                        else
                            yield(opt, false)
                        end
                    else
                        yield('?', opt)
                    end
                end
            else
                yield(false, arg)
            end
        end

        for i = i, #args do
            yield(false, args[i])
        end
    end)
end

local status, err = pcall(main, ...)
if status ~= true and err:find("interrupted!") == nil then
    print(err)
    os.exit(1)
end