#!/usr/bin/python

import sys, os, signal, subprocess, errno, time

WARMUP = 3
DURATION = 15

CORES = file("/proc/cpuinfo").read().count("processor\t")

def usage():
    print >> sys.stderr, "Usage: %s count port ip duration" % sys.argv[0]
    sys.exit(2)

def setup():
    # Put ourselves in a new process group so we can broadcast signals
    os.setpgid(0, 0)
    # Ignore USR1 and USR2, which we broadcast
    signal.signal(signal.SIGUSR1, signal.SIG_IGN)
    signal.signal(signal.SIGUSR2, signal.SIG_IGN)
    # Get into the right directory
    os.chdir(os.path.dirname(sys.argv[0]))

def startLoad(count, port, ip):
    # Start smtpbm
    procs = []
    for i in range(count):
        # smtpbm will exit if its parent dies, which makes our
        # exception handling much simpler (and often absent
        # altogether).
        procs.append(
            subprocess.Popen(["numactl", "-C", str(i % CORES),
                              "./smtpbm", ip, str(port),
                              "%d@mosbench.org" % i, "mosbench@mosbench.org", str(DURATION)
                              ],
                             stdout = subprocess.PIPE))
        # Don't hit the system too hard
        time.sleep(0.05)
    return procs

def startCounting(procs):
    print "Starting"
    sys.stdout.flush()
    os.kill(0, signal.SIGUSR1)

def timedWait(timeout):
    expired = [False]
    def onAlarm(signum, frame):
        expired[0] = True
    signal.signal(signal.SIGALRM, onAlarm)
    signal.alarm(timeout)
    try:
        return os.wait()
    except OSError, e:
        if e.errno != errno.EINTR or not expired[0]:
            raise
        return (0, 0)
    finally:
        signal.alarm(0)

def run(procs, duration):
    # Pause for the duration, paying attention to any smtpbm that dies
    # unexpectedly
    (pid, status) = timedWait(duration)
    if pid != 0:
        print >> sys.stderr, "smtpbm PID %d exited unexpectedly with %s" % \
            (pid, prettyWait(status))
        killall(procs, pid)
        sys.exit(1)

def stopCounting(procs):
    os.kill(0, signal.SIGUSR2)
    latency_sum = [0] * DURATION
    cnt_sum = [0] * DURATION
    for p in procs:
        def onAlarm(signum, frame):
            raise IOError("Timed out reading from smtpbm %d" % p.pid)
        signal.signal(signal.SIGALRM, onAlarm)
        # If it takes longer than a second to get the output, we're
        # skewed and screwed.  This is mostly a safeguard to catch
        # misbehaving children.  The IOError will kill this process,
        # which will kill the children.
        for i in range(DURATION):
            signal.alarm(1)
            res = p.stdout.readline().split()
            signal.alarm(0)
            cnt_sum[i] += int(res[0])
            latency_sum[i] += float(res[1])
    print "Stopped"
    sys.stdout.flush()
    for i in range(DURATION):
        if cnt_sum[i] == 0:
            print "0 -1" # indicates no message sent
        else:
            print "%d %.2f" % (cnt_sum[i], latency_sum[i] / cnt_sum[i])
    print "overall throughput: %f msg/sec" % (float(sum(cnt_sum[: DURATION])) / DURATION)
    print "overall latency: %f sec/msg" % (sum(latency_sum[: DURATION]) / sum(cnt_sum[: DURATION]))

def killall(procs, butPid = -1):
    wait = [p for p in procs if p.pid != butPid and p.poll() == None]
    for p in wait:
        try:
            os.kill(p.pid, signal.SIGINT)
        except EnvironmentError, e:
            print >> sys.stderr, e
    for p in wait:
        try:
            p.wait()
        except EnvironmentError, e:
            print >> sys.stderr, e

def prettyWait(status):
    if os.WIFEXITED(status):
        return "status %d" % os.WEXITSTATUS(status)
    if os.WIFSIGNALED(status):
        return "signal %d" % os.WTERMSIG(status)
    return "unknown wait status %d" % status

if len(sys.argv) != 5 or not sys.argv[1].isdigit() or not sys.argv[2].isdigit():
    usage()

count, port, ip, DURATION = int(sys.argv[1]), int(sys.argv[2]), sys.argv[3], int(sys.argv[4])

setup()
procs = startLoad(count, port, ip)
time.sleep(WARMUP)
startCounting(procs)
run(procs, DURATION)
stopCounting(procs)
killall(procs)