#!/usr/bin/env python
import re
import sys
import time
from optparse import OptionParser

VERSION = "1.0"

BEAN_COUNTERS = '/proc/user_beancounters'
SPLITTER = re.compile("\s*(?:(?P<uid>\d+):)?\s"
                      "\s*(?P<resource>\w+)\s"
                      "\s*(?P<held>\d+)\s"
                      "\s*(?P<maxheld>\d+)\s"
                      "\s*(?P<barrier>\d+)\s"
                      "\s*(?P<limit>\d+)\s"
                      "\s*(?P<failcnt>\d+)")
SPLITTER_NUMS = ('held', 'maxheld', 'barrier', 'limit', 'failcnt')
PAGE_SIZE = 4096

MEM = ('kmemsize', 'oomguarpages',
        'tcpsndbuf', 'tcprcvbuf', 'othersockbuf', 'dgramrcvbuf')
LOW = ('kmemsize', 'tcpsndbuf', 'tcprcvbuf', 'othersockbuf', 'dgramrcvbuf')
HIGH = ('oomguarpages',)
SHARED = ('shmpages',)
BUFFERS = ('tcpsndbuf', 'tcprcvbuf', 'othersockbuf', 'dgramrcvbuf')
CACHED = ('dcachesize',)

def match_and_normalize(line):
    m = SPLITTER.match(line)
    if m is None:
        return None

    d = m.groupdict()
    for k, v in d.iteritems():
        if k in SPLITTER_NUMS:
            if 'page' in d['resource']:
                d[k] = PAGE_SIZE * int(v)
            else:
                d[k] = int(v)
    return d

def get_counters():
    bfile = open(BEAN_COUNTERS);

    uid = None
    counters = {}
    for line in bfile:
        d = match_and_normalize(line)
        if d is None:
            continue

        if d['uid'] is not None:
            uid = int(d['uid'])
        del d['uid']

        if d['resource'] == 'dummy':
            continue

        counters[d['resource']] = d

    return (uid, counters)

def get_group(counters, group, field):
    return sum([v[field] for v in counters.values() if v['resource'] in group])

def cb_version(option, opt_str, value, parser):
    print """freevz %s
Copyright (C) 2006 Scott Dial
This is free software.  You may redistribute copies of it under the terms of
the GNU General Public License <http://www.gnu.org/licenses/gpl.html>.
There is NO WARRANTY, to the extent permitted by law.""" % (VERSION,)
    sys.exit(0)

def parse_options():
    usage="freevz [-b|-k|-m|-g] [-l] [-o] [-t] [-s DELAY] [-c COUNT] [-V]"

    parser = OptionParser(usage = usage)
    parser.add_option("-b", action="store_const", dest="scale",
                        const=1, help="show output in bytes")
    parser.add_option("-k", action="store_const", dest="scale",
                        const=1024, help="show output in kilobytes")
    parser.add_option("-m", action="store_const", dest="scale",
                        const=1048576, help="show output in megabytes")
    parser.add_option("-g", action="store_const", dest="scale",
                        const=1073741824, help="show output in gigabytes")
    parser.add_option("-l", action="store_true", dest="lowhigh",
                        help="show detailed low and high statistics")
    parser.add_option("-o", action="store_false", dest="buffers",
                        help="use old format (no -/+ buffers/cache line)")
    parser.add_option("-t", action="store_true", dest="total",
                        help="display total for RAM + swap")
    parser.add_option("-s", action="store", dest="delay", type="int",
                        help="update every [DELAY] seconds")
    parser.add_option("-c", action="store", dest="count", type="int",
                        help="update [COUNT] times")
    parser.add_option("-V", action="callback", callback=cb_version,
                        help="show version information and exit")

    parser.set_defaults(scale = 1024, lowhigh = False, buffers = True,
                        total = False, delay = None, count = None)

    return parser.parse_args()

def main():
    (options, _) = parse_options()

    def s(v):
        return v // options.scale

    while True:
        uid, counters =  get_counters()

        print "             total       used       free     shared    buffers     cached"

        print "Mem:    %10d %10d %10d %10d %10d %10d" % (
                s(get_group(counters, MEM, 'limit')),
                s(get_group(counters, MEM, 'held')),
                s(get_group(counters, MEM, 'limit') - get_group(counters, MEM, 'held')),
                s(get_group(counters, SHARED, 'held')),
                s(get_group(counters, BUFFERS, 'held')),
                s(get_group(counters, CACHED, 'held')))

        if options.buffers:
            print "-/+ buffers/cache: %10d %10d" % (
                    s(get_group(counters, MEM, 'held') -
                        get_group(counters, SHARED + BUFFERS + CACHED, 'held')),
                    s(get_group(counters, MEM, 'limit') -
                        get_group(counters, MEM + SHARED + BUFFERS + CACHED, 'held')))

        if options.lowhigh:
            print "Low:    %10d %10d %10d" % (
                s(get_group(counters, LOW, 'limit')),
                s(get_group(counters, LOW, 'held')),
                s(get_group(counters, LOW, 'limit') -
                    get_group(counters, LOW, 'held')))
            print "High:   %10d %10d %10d" % (
                s(get_group(counters, HIGH, 'limit')),
                s(get_group(counters, HIGH, 'held')),
                s(get_group(counters, HIGH, 'limit') -
                    get_group(counters, HIGH, 'held')))

        print "Swap:   %10d %10d %10d" % (
                s(counters['oomguarpages']['limit'] -
                    get_group(counters, LOW, 'limit') - counters['physpages']['held']),
                s(counters['oomguarpages']['held'] - counters['physpages']['held']),
                s(counters['oomguarpages']['limit'] -
                    get_group(counters, LOW, 'limit') - counters['physpages']['held'] -
                    (counters['oomguarpages']['held'] - counters['physpages']['held'])))

        if options.total:
            print "Total:  %10d %10d %10d" % (
                s(get_group(counters, MEM, 'limit')),
                s(get_group(counters, MEM, 'held')),
                s(get_group(counters, MEM, 'limit') - get_group(counters, MEM, 'held')))

        if options.delay is None:
            break
        else:
            print

            if options.count is not None:
                options.count -= 1
                if options.count == 0:
                    break

            time.sleep(options.delay)


if __name__ == '__main__':
    try:
        main()
    except KeyboardInterrupt:
        pass
