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 126
