← Back to branch summary

~bzr-pqm/bzr/bzr.dev

« back to all changes in this revision

Viewing changes to bzrlib/transport/sftp.py

merge bzr.dev

Show diffs side-by-side

added added

removed removed

Lines of Context:
bzrlib/transport/sftp.py
1065
1065
                        x)
1066
1066
 
1067
1067
 
 
1068
class SocketDelay(object):
 
1069
    """A socket decorator to make TCP appear slower.
 
1070
 
 
1071
    This changes recv, send, and sendall to add a fixed latency to each python
 
1072
    call if a new roundtrip is detected. That is, when a recv is called and the
 
1073
    flag new_roundtrip is set, latency is charged. Every send and send_all
 
1074
    sets this flag.
 
1075
 
 
1076
    In addition every send, sendall and recv sleeps a bit per character send to
 
1077
    simulate bandwidth.
 
1078
 
 
1079
    Not all methods are implemented, this is deliberate as this class is not a
 
1080
    replacement for the builtin sockets layer. fileno is not implemented to
 
1081
    prevent the proxy being bypassed. 
 
1082
    """
 
1083
 
 
1084
    simulated_time = 0
 
1085
    _proxied_arguments = dict.fromkeys([
 
1086
        "close", "getpeername", "getsockname", "getsockopt", "gettimeout",
 
1087
        "setblocking", "setsockopt", "settimeout", "shutdown"])
 
1088
 
 
1089
    def __init__(self, sock, latency, bandwidth=1.0, 
 
1090
                 really_sleep=True):
 
1091
        """ 
 
1092
        :param bandwith: simulated bandwith (MegaBit)
 
1093
        :param really_sleep: If set to false, the SocketDelay will just
 
1094
        increase a counter, instead of calling time.sleep. This is useful for
 
1095
        unittesting the SocketDelay.
 
1096
        """
 
1097
        self.sock = sock
 
1098
        self.latency = latency
 
1099
        self.really_sleep = really_sleep
 
1100
        self.time_per_byte = 1 / (bandwidth / 8.0 * 1024 * 1024) 
 
1101
        self.new_roundtrip = False
 
1102
 
 
1103
    def sleep(self, s):
 
1104
        if self.really_sleep:
 
1105
            time.sleep(s)
 
1106
        else:
 
1107
            SocketDelay.simulated_time += s
 
1108
 
 
1109
    def __getattr__(self, attr):
 
1110
        if attr in SocketDelay._proxied_arguments:
 
1111
            return getattr(self.sock, attr)
 
1112
        raise AttributeError("'SocketDelay' object has no attribute %r" %
 
1113
                             attr)
 
1114
 
 
1115
    def dup(self):
 
1116
        return SocketDelay(self.sock.dup(), self.latency, self.time_per_byte,
 
1117
                           self._sleep)
 
1118
 
 
1119
    def recv(self, *args):
 
1120
        data = self.sock.recv(*args)
 
1121
        if data and self.new_roundtrip:
 
1122
            self.new_roundtrip = False
 
1123
            self.sleep(self.latency)
 
1124
        self.sleep(len(data) * self.time_per_byte)
 
1125
        return data
 
1126
 
 
1127
    def sendall(self, data, flags=0):
 
1128
        if not self.new_roundtrip:
 
1129
            self.new_roundtrip = True
 
1130
            self.sleep(self.latency)
 
1131
        self.sleep(len(data) * self.time_per_byte)
 
1132
        return self.sock.sendall(data, flags)
 
1133
 
 
1134
    def send(self, data, flags=0):
 
1135
        if not self.new_roundtrip:
 
1136
            self.new_roundtrip = True
 
1137
            self.sleep(self.latency)
 
1138
        bytes_sent = self.sock.send(data, flags)
 
1139
        self.sleep(bytes_sent * self.time_per_byte)
 
1140
        return bytes_sent
 
1141
 
 
1142
 
1068
1143
class SFTPServer(Server):
1069
1144
    """Common code for SFTP server facilities."""
1070
1145
 
1077
1152
        self._vendor = 'none'
1078
1153
        # sftp server logs
1079
1154
        self.logs = []
 
1155
        self.add_latency = 0
1080
1156
 
1081
1157
    def _get_sftp_url(self, path):
1082
1158
        """Calculate an sftp url to this server for path."""
1086
1162
        """StubServer uses this to log when a new server is created."""
1087
1163
        self.logs.append(message)
1088
1164
 
 
1165
    def _run_server_entry(self, sock):
 
1166
        """Entry point for all implementations of _run_server.
 
1167
        
 
1168
        If self.add_latency is > 0.000001 then sock is given a latency adding
 
1169
        decorator.
 
1170
        """
 
1171
        if self.add_latency > 0.000001:
 
1172
            sock = SocketDelay(sock, self.add_latency)
 
1173
        return self._run_server(sock)
 
1174
 
1089
1175
    def _run_server(self, s):
1090
1176
        ssh_server = paramiko.Transport(s)
1091
1177
        key_file = pathjoin(self._homedir, 'test_rsa.key')
1117
1203
        self._root = '/'
1118
1204
        if sys.platform == 'win32':
1119
1205
            self._root = ''
1120
 
        self._listener = SocketListener(self._run_server)
 
1206
        self._listener = SocketListener(self._run_server_entry)
1121
1207
        self._listener.setDaemon(True)
1122
1208
        self._listener.start()
1123
1209