timed-os-stats

#!/usr/bin/env python3

# timed-os-stats
# Copyright 2012-2020 Greg Smith gregs1104@gmail.com
#
# Runs a command that produces regularly timed 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

def to_utf8(s):
    return s if isinstance(s, str) else s.decode('utf-8')

class FlushFile:
    """Write-only flushing wrapper for file-type objects."""
    def __init__(self, f):
        self.f = f
        self.flush = f.flush;
    def write(self, x):
        self.f.write(to_utf8(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:
        sys.stderr.write("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

# Popen parameters have to be different for text psql output
def capture_text(cmd,tag_all=True,tag_when=None):
    try:
        os.setsid()
    except OSError:
        sys.stderr.write("Unable to setsid(). Already process group leader?")

    # When executing via shell, use a string for command instead of a list
    cmdstr=' '.join(cmd)

    p = subprocess.Popen(cmdstr, shell = True,
                         universal_newlines=True,  # text mode
                         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=1):
    """
    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:
        sys.stderr.write("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():
    sys.stderr.write("Usage:  %s [vmstat | iostat | meminfo | psql]\n" % sys.argv[0])
    sys.stderr.write("Supported platforms are linux and darwin\n")
    sys.stderr.write("meminfo is only available on linux\n")
    sys.exit(1)

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

    cmd=sys.argv[1]

    if cmd=="psql":
        capture_text(sys.argv[1:])
    elif sys.platform.startswith('linux'):
        if cmd=='vmstat':
            capture("vmstat 1")
        elif cmd=='iostat':
            # Originally collection only added a timestamp per output set
            # using this call:
            #
            #   capture("iostat -mx 1",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 1")
        elif cmd=='meminfo':
            linux_meminfo(1)
        else:
            usage()

    elif sys.platform.startswith('darwin'):
        if cmd=='vmstat':
            capture("vm_stat 1")
        elif cmd=='iostat':
            capture("iostat -d -C -U -K 1")
        elif cmd=='meminfo':
            sys.stderr.write("meminfo unavailable on OS X, skipping\n")
        else:
            usage()
    # TODO Test this section actually works
    elif sys.platform.startswith('freebsd'):
        if cmd=='vmstat':
            capture("vmstat 1")
        elif cmd=='iostat':
            capture("iostat -Kx 1")
        else:
            usage()
    else:
        sys.stderr.write("Unsupported platform %s\n" % sys.platform)