proxyshaper.py

#!/usr/bin/env python
# Copyright 2012 Google Inc. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Simulate network characteristics directly in Python.

Allows running replay without dummynet.
"""

import logging
import platformsettings
import re
import time


TIMER = platformsettings.timer


class ProxyShaperError(Exception):
  """Module catch-all error."""
  pass

class BandwidthValueError(ProxyShaperError):
  """Raised for unexpected dummynet-style bandwidth value."""
  pass


class RateLimitedFile(object):
  """Wrap a file like object with rate limiting.

  TODO(slamm): Simulate slow-start.
      Each RateLimitedFile corresponds to one-direction of a
      bidirectional socket. Slow-start can be added here (algorithm needed).
      Will consider changing this class to take read and write files and
      corresponding bit rates for each.
  """
  BYTES_PER_WRITE = 1460

  def __init__(self, request_counter, f, bps):
    """Initialize a RateLimiter.

    Args:
      request_counter: callable to see how many requests share the limit.
      f: file-like object to wrap.
      bps: an integer of bits per second.
    """
    self.request_counter = request_counter
    self.original_file = f
    self.bps = bps

  def transfer_seconds(self, num_bytes):
    """Seconds to read/write |num_bytes| with |self.bps|."""
    return 8.0 * num_bytes / self.bps

  def write(self, data):
    num_bytes = len(data)
    num_sent_bytes = 0
    while num_sent_bytes < num_bytes:
      num_write_bytes = min(self.BYTES_PER_WRITE, num_bytes - num_sent_bytes)
      num_requests = self.request_counter()
      wait = self.transfer_seconds(num_write_bytes) * num_requests
      logging.debug('write sleep: %0.4fs (%d requests)', wait, num_requests)
      time.sleep(wait)

      self.original_file.write(
          data[num_sent_bytes:num_sent_bytes + num_write_bytes])
      num_sent_bytes += num_write_bytes

  def _read(self, read_func, size):
    start = TIMER()
    data = read_func(size)
    read_seconds = TIMER() - start
    num_bytes = len(data)
    num_requests = self.request_counter()
    wait = self.transfer_seconds(num_bytes) * num_requests - read_seconds
    if wait > 0:
      logging.debug('read sleep: %0.4fs %d requests)', wait, num_requests)
      time.sleep(wait)
    return data

  def readline(self, size=-1):
    return self._read(self.original_file.readline, size)

  def read(self, size=-1):
    return self._read(self.original_file.read, size)

  def __getattr__(self, name):
    """Forward any non-overriden calls."""
    return getattr(self.original_file, name)


def GetBitsPerSecond(bandwidth):
  """Return bits per second represented by dummynet bandwidth option.

  See ipfw/dummynet.c:read_bandwidth for how it is really done.

  Args:
    bandwidth: a dummynet-style bandwidth specification (e.g. "10Kbit/s")
  """
  if bandwidth == '0':
    return 0
  bw_re = r'^(\d+)(?:([KM])?(bit|Byte)/s)?$'
  match = re.match(bw_re, str(bandwidth))
  if not match:
    raise BandwidthValueError('Value, "%s", does not match regex: %s' % (
        bandwidth, bw_re))
  bw = int(match.group(1))
  if match.group(2) == 'K':
    bw *= 1000
  if match.group(2) == 'M':
    bw *= 1000000
  if match.group(3) == 'Byte':
    bw *= 8
  return bw