prometheus_taskstat_collector

#!/usr/bin/env python
import os

from pyroute2 import TaskStats
from pyroute2.netlink import NLM_F_REQUEST
from pyroute2.netlink.taskstats import tcmd
from pyroute2.netlink.taskstats import TASKSTATS_CMD_GET

clk_tck = os.sysconf('SC_CLK_TCK')
clk_tck_factor = int(1e9 / clk_tck)

pagesize = os.sysconf('SC_PAGESIZE')
pagesize_factor = pagesize / 1024

interesting_taskstats_keys = (
    'cpu_count',
    'cpu_delay_total',
    'blkio_count',
    'blkio_delay_total',
    'swapin_count',
    'swapin_delay_total',
    'nvcsw',
    'nivcsw',
)

interesting_procstat_fields = (
    ('utime', 14, clk_tck_factor),
    ('stime', 15, clk_tck_factor),
    ('cutime', 16, clk_tck_factor),
    ('cstime', 17, clk_tck_factor),
    ('num_threads', 20, 1),
    ('rss', 24, pagesize_factor),
)

smaps_searchstrings = ('Pss:')

def issue_request(ts, msg):
    return ts.nlm_request(msg, ts.prid, msg_flags=NLM_F_REQUEST)

def unpack_msg(msg):
    return msg[0]['attrs'][0][1]['attrs'][1][1]

def do_query(ts, pid):
    msg = tcmd()
    msg['cmd'] = TASKSTATS_CMD_GET
    msg['version'] = 1
    msg['attrs'].append(['TASKSTATS_CMD_ATTR_TGID', pid])
    return unpack_msg(issue_request(ts, msg))

def print_taskstats(processes):
    ts = TaskStats()
    ts.bind()
    try:
        for name, pid in processes:
            results = do_query(ts, pid)
            for key in interesting_taskstats_keys:
                value = results[key]
                print 'taskstat_{}{{name="{}", pid="{}"}} {}'.format(key, name, pid, value)
    finally:
        ts.close()

def print_procstats(processes):
    for name, pid in processes:
        stat_fields = open('/proc/{}/stat'.format(pid)).readline().rstrip().split(' ')
        for field, pos, factor in interesting_procstat_fields:
            value = int(stat_fields[pos - 1]) * factor
            print 'procstat_{}{{name="{}", pid="{}"}} {}'.format(field, name, pid, value)

def print_mapinfo(processes):
    for name, pid in processes:
        stats = dict((i, 0) for i in searchstrings)
        for line in open('/proc/{}/smaps'.format(pid)):
            for i in searchstrings:
                if line.startswith(i):
                    value = int(line.split()[1])
                    stats[i] += value

        for i in searchstrings:
            print 'procmaps_{}{{name="{}", pid="{}"}} {}'.format(i.lower()[:-1], name, pid, stats[i])


def get_pids_from_file(pidfile):
    return map(int, open(pidfile, 'r'))

def get_pids(directory):
    for basename in os.listdir(directory):
        name = basename.split('.')[0]
        pidfile = os.path.join(directory, basename)
        pids = ()
        try:
            pids = get_pids_from_file(pidfile)
        except (ValueError, IOError):
            pass

        for i in pids:
            yield name, i

def main():
    directory = os.getenv('TASK_PIDFILE_DIR', '/etc/prometheus/task_collector_pidfiles')
    processes = tuple(get_pids(directory))
    print_taskstats(processes)
    print_procstats(processes)
    print_mapinfo(processes)

if __name__ == '__main__':
    main()