First pass at adding ability to test using newer socket based connection
This commit is contained in:
Matt Martz
parent
9c2977acfc
commit
0ef2f6b04c
166
speedtest.py
166
speedtest.py
|
|
@ -713,7 +713,7 @@ class HTTPDownloader(threading.Thread):
|
|||
shutdown_event=None):
|
||||
threading.Thread.__init__(self)
|
||||
self.request = request
|
||||
self.result = [0]
|
||||
self.result = 0
|
||||
self.starttime = start
|
||||
self.timeout = timeout
|
||||
self.i = i
|
||||
|
|
@ -734,14 +734,62 @@ class HTTPDownloader(threading.Thread):
|
|||
while (not self._shutdown_event.isSet() and
|
||||
(timeit.default_timer() - self.starttime) <=
|
||||
self.timeout):
|
||||
self.result.append(len(f.read(10240)))
|
||||
if self.result[-1] == 0:
|
||||
data = len(f.read(10240))
|
||||
if data == 0:
|
||||
break
|
||||
self.result += data
|
||||
f.close()
|
||||
except IOError:
|
||||
pass
|
||||
|
||||
|
||||
class SocketDownloader(threading.Thread):
|
||||
def __init__(self, i, address, size, start, timeout, shutdown_event=None,
|
||||
source_address=None):
|
||||
threading.Thread.__init__(self)
|
||||
self.result = 0
|
||||
self.starttime = start
|
||||
self.timeout = timeout
|
||||
self.i = i
|
||||
self.size = size
|
||||
self.remaining = self.size
|
||||
|
||||
if shutdown_event:
|
||||
self._shutdown_event = shutdown_event
|
||||
else:
|
||||
self._shutdown_event = FakeShutdownEvent()
|
||||
|
||||
self.sock = create_connection(address, timeout=timeout,
|
||||
source_address=source_address)
|
||||
|
||||
def run(self):
|
||||
try:
|
||||
if (timeit.default_timer() - self.starttime) <= self.timeout:
|
||||
self.sock.sendall('HI\n'.encode())
|
||||
self.sock.recv(1024)
|
||||
|
||||
while (self.remaining and not self._shutdown_event.isSet() and
|
||||
(timeit.default_timer() - self.starttime) <=
|
||||
self.timeout):
|
||||
|
||||
if self.remaining > 1000000:
|
||||
ask = 1000000
|
||||
else:
|
||||
ask = self.remaining
|
||||
|
||||
down = 0
|
||||
self.sock.sendall(('DOWNLOAD %d\n' % ask).encode())
|
||||
while down < ask:
|
||||
down += len(self.sock.recv(10240))
|
||||
|
||||
self.result += down
|
||||
self.remaining -= down
|
||||
|
||||
self.sock.close()
|
||||
except IOError:
|
||||
pass
|
||||
|
||||
|
||||
class HTTPUploaderData(object):
|
||||
"""File like object to improve cutting off the upload once the timeout
|
||||
has been reached
|
||||
|
|
@ -759,7 +807,7 @@ class HTTPUploaderData(object):
|
|||
|
||||
self._data = None
|
||||
|
||||
self.total = [0]
|
||||
self.total = 0
|
||||
|
||||
def pre_allocate(self):
|
||||
chars = '0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZ'
|
||||
|
|
@ -787,7 +835,7 @@ class HTTPUploaderData(object):
|
|||
if ((timeit.default_timer() - self.start) <= self.timeout and
|
||||
not self._shutdown_event.isSet()):
|
||||
chunk = self.data.read(n)
|
||||
self.total.append(len(chunk))
|
||||
self.total += len(chunk)
|
||||
return chunk
|
||||
else:
|
||||
raise SpeedtestUploadTimeout()
|
||||
|
|
@ -835,11 +883,59 @@ class HTTPUploader(threading.Thread):
|
|||
f = self._opener(request)
|
||||
f.read(11)
|
||||
f.close()
|
||||
self.result = sum(self.request.data.total)
|
||||
self.result = self.request.data.total
|
||||
else:
|
||||
self.result = 0
|
||||
except (IOError, SpeedtestUploadTimeout):
|
||||
self.result = sum(self.request.data.total)
|
||||
self.result = self.request.data.total
|
||||
|
||||
|
||||
class SocketUploader(threading.Thread):
|
||||
def __init__(self, i, address, size, start, timeout, shutdown_event=None,
|
||||
source_address=None):
|
||||
threading.Thread.__init__(self)
|
||||
self.result = 0
|
||||
self.starttime = start
|
||||
self.timeout = timeout
|
||||
self.i = i
|
||||
self.size = size
|
||||
self.remaining = self.size
|
||||
|
||||
if shutdown_event:
|
||||
self._shutdown_event = shutdown_event
|
||||
else:
|
||||
self._shutdown_event = FakeShutdownEvent()
|
||||
|
||||
self.sock = create_connection(address, timeout=timeout,
|
||||
source_address=source_address)
|
||||
|
||||
def run(self):
|
||||
try:
|
||||
if (timeit.default_timer() - self.starttime) <= self.timeout:
|
||||
self.sock.sendall('HI\n'.encode())
|
||||
self.sock.recv(1024)
|
||||
|
||||
while (self.remaining and not self._shutdown_event.isSet() and
|
||||
(timeit.default_timer() - self.starttime) <=
|
||||
self.timeout):
|
||||
|
||||
if self.remaining > 100000:
|
||||
give = 100000
|
||||
else:
|
||||
give = self.remaining
|
||||
|
||||
header = ('UPLOAD %d 0\n' % give).encode()
|
||||
data = '0'.encode() * (give - len(header))
|
||||
|
||||
self.sock.sendall(header)
|
||||
self.sock.sendall(data)
|
||||
self.sock.recv(24)
|
||||
self.result += give
|
||||
self.remaining -= give
|
||||
|
||||
self.sock.close()
|
||||
except IOError:
|
||||
pass
|
||||
|
||||
|
||||
class SpeedtestResults(object):
|
||||
|
|
@ -997,7 +1093,7 @@ class Speedtest(object):
|
|||
"""Class for performing standard speedtest.net testing operations"""
|
||||
|
||||
def __init__(self, config=None, source_address=None, timeout=10,
|
||||
secure=False, shutdown_event=None):
|
||||
secure=False, shutdown_event=None, use_socket=False):
|
||||
self.config = {}
|
||||
|
||||
self._source_address = source_address
|
||||
|
|
@ -1011,6 +1107,8 @@ class Speedtest(object):
|
|||
else:
|
||||
self._shutdown_event = FakeShutdownEvent()
|
||||
|
||||
self._use_socket = use_socket
|
||||
|
||||
self.get_config()
|
||||
if config is not None:
|
||||
self.config.update(config)
|
||||
|
|
@ -1105,9 +1203,14 @@ class Speedtest(object):
|
|||
up_sizes = [32768, 65536, 131072, 262144, 524288, 1048576, 7340032]
|
||||
sizes = {
|
||||
'upload': up_sizes[ratio - 1:],
|
||||
'download': [350, 500, 750, 1000, 1500, 2000, 2500,
|
||||
3000, 3500, 4000]
|
||||
}
|
||||
if self._use_socket:
|
||||
sizes['download'] = [245388, 505544, 1118012, 1986284, 4468241,
|
||||
7907740, 12407926, 17816816, 24262167,
|
||||
31625365]
|
||||
else:
|
||||
sizes['download'] = [350, 500, 750, 1000, 1500, 2000, 2500,
|
||||
3000, 3500, 4000]
|
||||
|
||||
size_count = len(sizes['upload'])
|
||||
|
||||
|
|
@ -1252,6 +1355,9 @@ class Speedtest(object):
|
|||
or int(attrib.get('id')) in exclude):
|
||||
continue
|
||||
|
||||
host, port = attrib['host'].split(':')
|
||||
attrib['host'] = (host, int(port))
|
||||
|
||||
try:
|
||||
d = distance(self.lat_lon,
|
||||
(float(attrib.get('lat')),
|
||||
|
|
@ -1429,6 +1535,14 @@ class Speedtest(object):
|
|||
def download(self, callback=do_nothing):
|
||||
"""Test download speed against speedtest.net"""
|
||||
|
||||
if self._use_socket:
|
||||
requests = []
|
||||
for size in self.config['sizes']['download']:
|
||||
for _ in range(0, self.config['counts']['download']):
|
||||
requests.append(size)
|
||||
|
||||
request_count = len(requests)
|
||||
else:
|
||||
urls = []
|
||||
for size in self.config['sizes']['download']:
|
||||
for _ in range(0, self.config['counts']['download']):
|
||||
|
|
@ -1444,6 +1558,17 @@ class Speedtest(object):
|
|||
|
||||
def producer(q, requests, request_count):
|
||||
for i, request in enumerate(requests):
|
||||
if self._use_socket:
|
||||
thread = SocketDownloader(
|
||||
i,
|
||||
self.best['host'],
|
||||
request,
|
||||
start,
|
||||
self.config['length']['download'],
|
||||
shutdown_event=self._shutdown_event,
|
||||
source_address=self._source_address
|
||||
)
|
||||
else:
|
||||
thread = HTTPDownloader(
|
||||
i,
|
||||
request,
|
||||
|
|
@ -1463,7 +1588,7 @@ class Speedtest(object):
|
|||
thread = q.get(True)
|
||||
while thread.isAlive():
|
||||
thread.join(timeout=0.1)
|
||||
finished.append(sum(thread.result))
|
||||
finished.append(thread.result)
|
||||
callback(thread.i, request_count, end=True)
|
||||
|
||||
q = Queue(self.config['threads']['download'])
|
||||
|
|
@ -1502,6 +1627,9 @@ class Speedtest(object):
|
|||
|
||||
requests = []
|
||||
for i, size in enumerate(sizes):
|
||||
if self._use_socket:
|
||||
requests.append(size)
|
||||
else:
|
||||
# We set ``0`` for ``start`` and handle setting the actual
|
||||
# ``start`` in ``HTTPUploader`` to get better measurements
|
||||
data = HTTPUploaderData(
|
||||
|
|
@ -1521,6 +1649,17 @@ class Speedtest(object):
|
|||
|
||||
def producer(q, requests, request_count):
|
||||
for i, request in enumerate(requests[:request_count]):
|
||||
if self._use_socket:
|
||||
thread = SocketUploader(
|
||||
i,
|
||||
self.best['host'],
|
||||
request,
|
||||
start,
|
||||
self.config['length']['upload'],
|
||||
shutdown_event=self._shutdown_event,
|
||||
source_address=self._source_address
|
||||
)
|
||||
else:
|
||||
thread = HTTPUploader(
|
||||
i,
|
||||
request[0],
|
||||
|
|
@ -1652,6 +1791,8 @@ def parse_args():
|
|||
parser.add_argument('--secure', action='store_true',
|
||||
help='Use HTTPS instead of HTTP when communicating '
|
||||
'with speedtest.net operated servers')
|
||||
parser.add_argument('--socket', action='store_true',
|
||||
help='Use socket test instead of HTTP based tests')
|
||||
parser.add_argument('--no-pre-allocate', dest='pre_allocate',
|
||||
action='store_const', default=True, const=False,
|
||||
help='Do not pre allocate upload data. Pre allocation '
|
||||
|
|
@ -1764,7 +1905,8 @@ def shell():
|
|||
speedtest = Speedtest(
|
||||
source_address=args.source,
|
||||
timeout=args.timeout,
|
||||
secure=args.secure
|
||||
secure=args.secure,
|
||||
use_socket=args.socket
|
||||
)
|
||||
except (ConfigRetrievalError,) + HTTP_ERRORS:
|
||||
printer('Cannot retrieve speedtest configuration', error=True)
|
||||
|
|
|
|||
Loading…
Reference in New Issue