tools/telemetry/third_party/webpagereplay/proxyshaper.py

   1 #!/usr/bin/env python
   2 # Copyright 2012 Google Inc. All Rights Reserved.
   3 #
   4 # Licensed under the Apache License, Version 2.0 (the "License");
   5 # you may not use this file except in compliance with the License.
   6 # You may obtain a copy of the License at
   7 #
   8 #      http://www.apache.org/licenses/LICENSE-2.0
   9 #
  10 # Unless required by applicable law or agreed to in writing, software
  11 # distributed under the License is distributed on an "AS IS" BASIS,
  12 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  13 # See the License for the specific language governing permissions and
  14 # limitations under the License.
  15
  16 """Simulate network characteristics directly in Python.
  17
  18 Allows running replay without dummynet.
  19 """
  20
  21 import logging
  22 import platformsettings
  23 import re
  24 import time
  25
  26
  27 TIMER = platformsettings.timer
  28
  29
  30 class ProxyShaperError(Exception):
  31   """Module catch-all error."""
  32   pass
  33
  34 class BandwidthValueError(ProxyShaperError):
  35   """Raised for unexpected dummynet-style bandwidth value."""
  36   pass
  37
  38
  39 class RateLimitedFile(object):
  40   """Wrap a file like object with rate limiting.
  41
  42   TODO(slamm): Simulate slow-start.
  43       Each RateLimitedFile corresponds to one-direction of a
  44       bidirectional socket. Slow-start can be added here (algorithm needed).
  45       Will consider changing this class to take read and write files and
  46       corresponding bit rates for each.
  47   """
  48   BYTES_PER_WRITE = 1460
  49
  50   def __init__(self, request_counter, f, bps):
  51     """Initialize a RateLimiter.
  52
  53     Args:
  54       request_counter: callable to see how many requests share the limit.
  55       f: file-like object to wrap.
  56       bps: an integer of bits per second.
  57     """
  58     self.request_counter = request_counter
  59     self.original_file = f
  60     self.bps = bps
  61
  62   def transfer_seconds(self, num_bytes):
  63     """Seconds to read/write |num_bytes| with |self.bps|."""
  64     return 8.0 * num_bytes / self.bps
  65
  66   def write(self, data):
  67     num_bytes = len(data)
  68     num_sent_bytes = 0
  69     while num_sent_bytes < num_bytes:
  70       num_write_bytes = min(self.BYTES_PER_WRITE, num_bytes - num_sent_bytes)
  71       num_requests = self.request_counter()
  72       wait = self.transfer_seconds(num_write_bytes) * num_requests
  73       logging.debug('write sleep: %0.4fs (%d requests)', wait, num_requests)
  74       time.sleep(wait)
  75
  76       self.original_file.write(
  77           data[num_sent_bytes:num_sent_bytes + num_write_bytes])
  78       num_sent_bytes += num_write_bytes
  79
  80   def _read(self, read_func, size):
  81     start = TIMER()
  82     data = read_func(size)
  83     read_seconds = TIMER() - start
  84     num_bytes = len(data)
  85     num_requests = self.request_counter()
  86     wait = self.transfer_seconds(num_bytes) * num_requests - read_seconds
  87     if wait > 0:
  88       logging.debug('read sleep: %0.4fs %d requests)', wait, num_requests)
  89       time.sleep(wait)
  90     return data
  91
  92   def readline(self, size=-1):
  93     return self._read(self.original_file.readline, size)
  94
  95   def read(self, size=-1):
  96     return self._read(self.original_file.read, size)
  97
  98   def __getattr__(self, name):
  99     """Forward any non-overriden calls."""
 100     return getattr(self.original_file, name)
 101
 102
 103 def GetBitsPerSecond(bandwidth):
 104   """Return bits per second represented by dummynet bandwidth option.
 105
 106   See ipfw/dummynet.c:read_bandwidth for how it is really done.
 107
 108   Args:
 109     bandwidth: a dummynet-style bandwidth specification (e.g. "10Kbit/s")
 110   """
 111   if bandwidth == '0':
 112     return 0
 113   bw_re = r'^(\d+)(?:([KM])?(bit|Byte)/s)?$'
 114   match = re.match(bw_re, str(bandwidth))
 115   if not match:
 116     raise BandwidthValueError('Value, "%s", does not match regex: %s' % (
 117         bandwidth, bw_re))
 118   bw = int(match.group(1))
 119   if match.group(2) == 'K':
 120     bw *= 1000
 121   if match.group(2) == 'M':
 122     bw *= 1000000
 123   if match.group(3) == 'Byte':
 124     bw *= 8
 125   return bw