extras.py

import os
import subprocess
import sys
import threading

import termcolor as T

class MovingRecord:
    def __init__(self, num_samples):
        self.samples = []
        self.num_samples = num_samples

    def add_sample(self, sample):
        if len(self.samples) == self.num_samples:
            self.samples.pop(0)
        self.samples.append(sample)

    def average(self):
        if len(self.samples) == 0:
            return 0
        return sum(self.samples) / len(self.samples)

    def median(self):
        n = len(self.samples)
        if n == 0:
            return 0
        s = sorted(self.samples)
        if n % 2 == 0:
            return (s[n/2] + s[n/2 - 1]) / 2
        else:
            return s[n/2]

def start_tcpprobe(fname):
    "Install tcp_probe module and dump to file"
    os.system("rmmod tcp_probe 2>/dev/null; modprobe tcp_probe;")
    subprocess.Popen("cat /proc/net/tcpprobe > %s" %
          fname, shell=True)

def stop_tcpprobe():
    os.system("killall -9 cat; rmmod tcp_probe &>/dev/null;")

def cprint(s, color, cr=True):
    """Print in color
       s: string to print
       color: color to use"""
    if cr:
        print T.colored(s, color)
    else:
        print T.colored(s, color),

def monitor_stderr(p):
    lines_iterator = iter(p.stderr.readline, b"")
    for line in lines_iterator:
        print(line.strip())

def monitor_stdout(p):
    lines_iterator = iter(p.stdout.readline, b"")
    for line in lines_iterator:
        print(line.strip())

def exec_and_monitor(host, cmd):
    if host is not None:
        p = host.popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
    else:
        p = subprocess.Popen(args=cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)

    t1 = threading.Thread(target=monitor_stdout, args=(p,))
    t1.setDaemon(True)
    t1.start()
    t2 = threading.Thread(target=monitor_stderr, args=(p,))
    t2.setDaemon(True)
    t2.start()
    return p

def verify_rtt_between_endpoints(e1, e2, expected_rtt, tolerance):
    proc = e1.popen('ping -c 2 %s' % e2.IP())
    output = proc.communicate()[0]
    line = output.split('\n')[2]
    # "64 bytes from 10.0.0.2: icmp_seq=2 ttl=64 time=174 ms"
    rtt = float(line.split(' ')[-2][5:])
    if rtt < expected_rtt*(1-tolerance) or rtt > expected_rtt*(1+tolerance):
        print 'Latency verification failed'
        print 'Got %f instead of %f' % (rtt, expected_rtt)
        sys.exit()

def verify_latency(h1, h2, expected_rtt):
    print 'Verifying client-server latency...'
    verify_rtt_between_endpoints(h1, h2, expected_rtt, 0.05)
    verify_rtt_between_endpoints(h2, h2, expected_rtt, 0.05)

def verify_bandwidth(h1, h2, expected_bandwidth):
    print 'Verifying bandwidth of bottleneck link...'
    server_proc = h2.popen('iperf -s -p 5001')
    client_proc = h1.popen('iperf -c %s -t 20' % h2.IP());
    output = client_proc.communicate()[0]
    # "[ 15]  0.0- 5.1 sec  38.0 MBytes  61.9 Mbits/sec"
    line = output.split('\n')[6]
    bandwidth = float(line.split(' ')[-2])
    tolerance = 0.1
    if bandwidth < expected_bandwidth*(1-tolerance) or \
        bandwidth > expected_bandwidth*(1+tolerance):
        print 'Bandwidth verification failed'
        print 'Got %f instead of %f' % (bandwidth, expected_bandwidth)
    server_proc.kill()