wifitables/iperf.py

"""iperf: run a series of iperf tests over a wireless network."""

import collections
import glob
import json
import os
import re
import sys


def Mbps(mbps):
  """Given a bandwidth in Mbps, return the bandwidth in bit/s."""
  return mbps * 1000000


class Iperf(object):
  """Iperf collects `iperf` measurements across a wireless network."""

  def __init__(self, runner=None, bind=None):
    self.run = runner
    self.bind = bind

  def Ping(self, host):
    args = ['-c', '1', host]
    return (self.run(['ping'] + args).succeeded or
            self.run(['ping6'] + args).succeeded)

  def _Iperf(self, host, udp=False, bandwidth=20):
    """Run iperf against host and return string containing stdout from run."""
    args = ['iperf', '-c', host]
    if udp:
      args += ['-u', '-b', str(Mbps(bandwidth))]
    if self.bind:
      args += ['-B', self.bind]

    return self.run(args)

  def IperfTCP(self, host='127.0.0.1'):
    return self._Iperf(host)

  def IperfUDP(self, host='127.0.0.1', bandwidth=20):
    return self._Iperf(host, udp=True, bandwidth=bandwidth)

  def Iperf3(self, host, udp=False, reverse=False, bandwidth=1e6):
    """Run iperf against host and return string containing stdout from run."""
    args = ['iperf3', '-c', host, '-J']
    if udp:
      args += ['-u', '-b', str(bandwidth)]
    if reverse:
      args += ['-R']

    if self.bind:
      args += ['-B', self.bind]

    return self.run(args)

  def RunTestSeries(self, host):
    """RunTestSeries runs iperf tests and returns their results.

    Args:
      host: A string containing the hostname to run tests against.
    Returns:
      A dict mapping unique names (strings, suitable for use as filenames) to
      strings containing JSON-formatted iperf3 results.
    """
    outputs = {}
    if not self.Ping(host):
      print >> sys.stderr, ('Could not ping destination host {}; '
                            'skipping performance tests').format(host)
      return outputs

    counter = collections.Counter()
    def NextName(counter, basename):
      """Return `logrotate(1)`-style names for a bunch of iperf3 files."""
      count = counter[basename]
      if count:
        name = '{}.{}'.format(basename, count)
      else:
        name = basename

      counter[basename] += 1
      return name

    for reverse in (False, True):
      it = self.Iperf3(host, reverse=reverse)
      outputs[NextName(counter, 'iperf3')] = it

      # Empirically about 1.25x more packets make it through in UDP than TCP.
      # Try to saturate the channel by sending a bit more than that over UDP.
      try:
        itp = ParseIperf3(it)
      except ValueError:
        print >> sys.stderr, ('Could not decode iperf3 TCP result, proceeding '
                              'as if it were empty.')
        itp = {}

      bw = (itp.get('end', {})
            .get('sum_received', {})
            .get('bits_per_second', 1))

      iu = self.Iperf3(host, udp=True, reverse=reverse, bandwidth=bw * 1.5)
      outputs[NextName(counter, 'iperf3')] = iu

    return outputs


def _ParseIperf(text, udp=False):
  """Parse summary line written by an `iperf` run into a Python dict."""
  pattern = (r'\[(.{3})\]\s+(?P<interval>.*?sec)\s+'
             r'(?P<transfer>.*?Bytes|bits)'
             r'\s+(?P<bandwidth>.*?/sec)')
  if udp:
    pattern += r'\s+(?P<jitter>.*?s)\s+(?P<datagrams>.*)'

  iperf_re = re.compile(pattern)

  for line in text.splitlines():
    match = iperf_re.match(line)
    if match:
      iperf = match.groupdict()
      bval, bunit = iperf['bandwidth'].split()
      iperf['bandwidth'] = float(bval)
      iperf['bandwidth_unit'] = bunit

      tval, tunit = iperf['transfer'].split()
      iperf['transfer'] = float(tval)
      iperf['transfer_unit'] = tunit
      return iperf

  return {}


def ParseIperfTCP(text):
  # sample line: [  4]  0.0-10.0 sec   245 MBytes   206 Mbits/sec
  return _ParseIperf(text)


def ParseIperfUDP(text):
  # pylint: disable=line-too-long
  # sample line: [  5]  0.0-10.0 sec  1.25 MBytes  1.05 Mbits/sec   0.593 ms    0/  893 (0%)
  return _ParseIperf(text, udp=True)


def ParseIperf3(text):
  return json.loads(text)


def Restore(report_dir):
  """Restores an `Iperf` cache from data on the filesystem."""
  cache = {}
  for name in glob.glob(os.path.join(report_dir, 'iperf*')):
    try:
      with open(name) as infile:
        cache[os.path.basename(name)] = infile.read()
    except IOError:
      cache[name] = ''

  return cache