timed-os-stats

#!/usr/bin/env python

# timed-os-stats
# Copyright 2012 Greg Smith gsmith@gregsmith.com
#
# Runs a command that produces regularly time output and adds time
# stamps for each line when appropriate.  Typical commands it might
# be used for are vmstat and iostat.
#
# On Linux the iostat command produces complicated output for each
# sample, such that a more complicated parser is needed to decode
# everything.  The program presumes that will be handled by a
# downstream tool.  It only tags before the first line in each
# each sample (the one starting with "avg-cpu:") in that
# case.
#

import os
import subprocess
import sys
import datetime
import time

class FlushFile(object):
    """Write-only flushing wrapper for file-type objects."""
    def __init__(self, f):
        self.f = f
    def write(self, x):
        self.f.write(x)
        self.f.flush()

# Replace stdout with an automatically flushing version
sys.stdout = FlushFile(sys.__stdout__)

def capture(cmd,tag_all=True,tag_when=None):
    try:
        # Attempt to call setsid() so that this Python process and its
        # children (shell and cmd) are in a new process group.
        # (On some platforms/shells, this Python process may already
        # be a process group leader, on some it will not be.)
        # We need to be in a new process group
        # because when it is time to kill this script we will want to kill
        # by process group ID from benchwarmer, and we don't wan't
        # benchwarmer to kill itself as well.
        os.setsid()
    except OSError:
        print >> sys.stderr, \
                 "Unable to setsid(). Already process group leader?"

    p = subprocess.Popen(cmd, shell = True,
                         stdout = subprocess.PIPE,
                         stderr = subprocess.STDOUT,
                         )                         

    pid = p.pid

    while True:
        line = p.stdout.readline()
        if line == '' and p.poll() != None:
            break

        if tag_all:
            sys.stdout.write("%s\t" % datetime.datetime.now())
        elif tag_when != None and line.find(tag_when)>=0:
            sys.stdout.write("%s\n" % datetime.datetime.now())

        sys.stdout.write(line)

    retcode = p.returncode

def linux_meminfo(interval=5):
    """
    Read /proc/meminfo on a Linux system and output its
    values with a timestamp.  That file has a mix of lines
    that end in "kB" and numbers that are a count.  The
    "kB" suffixes are stripped out, which means parsers
    of this data need to know whether keys are in kB or not.
    """
    try:
        os.setsid()
    except OSError:
        print >> sys.stderr, \
                 "Unable to setsid(). Already process group leader?"

    while True:
        meminfo="/proc/meminfo"
        mem=open(meminfo,"r")
        lines=mem.readlines()
        ts=datetime.datetime.now()
        for line in lines:
            if line.find(":")<0:  continue
            if line.find("kB")<0:
                (key,val)=line.split()
            else:
                (key,val,kb)=line.split()
            key=key.rstrip(":")
            sys.stdout.write("%s\t%s\t%s\n" % (ts,key,val))
        time.sleep(interval)

def usage():
    print >> sys.stderr, "Usage:  %s [vmstat | iostat | meminfo]" % sys.argv[0]
    print >> sys.stderr, "Supported platforms are linux and darwin"
    print >> sys.stderr, "meminfo is only available on linux"
    sys.exit(1)

if __name__=='__main__':
    if len(sys.argv)<2:
        usage()

    cmd=sys.argv[1]

    if sys.platform.startswith('linux'):
        if cmd=='vmstat':
            capture("vmstat 5")
        elif cmd=='iostat':
            # Originally collection only added a timestamp per output set
            # using this call:
            #
            #   capture("iostat -mx 5",False,"avg-cpu:")
            #
            # While the most accurate approach, that result requires a
            # non-trivial parser to insert the time stamps later.
            # Instead, now every line gets a timestamp.  The main
            # downside is that times won't match up exactly
            # for multiple disks worth of data.  That's annoying,
            # but it doesn't make a real difference to graphers of
            # the resulting data.  The other piece that a real
            # parser would help with is eliminating the first output
            # set, which is a set of averages we'd prefer to throw
            # away.
            # Also:  current iostat on RHEL6 at least has a "-t"
            # option that adds timestmaps before each new "avg-cpu:",
            # exactly where we'd want them to be.  No need for this Python
            # program to get them.  vmstat also has a timestamp option with
            # "-t", which puts them at the end of each line.
            capture("iostat -mx 5")
        elif cmd=='meminfo':
            linux_meminfo(5)
        else:
            usage()

    elif sys.platform.startswith('darwin'):
        if cmd=='vmstat':
            capture("vm_stat 5")
        elif cmd=='iostat':
            capture("iostat -d -C -K 5")
        else:
            usage()
    # TODO Test this section actually works
    elif sys.platform.startswith('freebsd'):
        if cmd=='vmstat':
            capture("vmstat 5")
        elif cmd=='iostat':
            capture("iostat -Kx 5")
        else:
            usage()
    else:
        print >> sys.stderr, 'Unsupported platform', sys.platform