1# Copyright (c) 2012 The Chromium OS Authors. All rights reserved.
2# Use of this source code is governed by a BSD-style license that can be
3# found in the LICENSE file.
4
5"""A tool to measure single-stream link bandwidth using HTTP connections."""
6
7import logging, random, time, urllib2
8
9import numpy.random
10
11TIMEOUT = 90
12
13
14class Error(Exception):
15  pass
16
17
18def TimeTransfer(url, data):
19    """Transfers data to/from url.  Returns (time, url contents)."""
20    start_time = time.time()
21    result = urllib2.urlopen(url, data=data, timeout=TIMEOUT)
22    got = result.read()
23    transfer_time = time.time() - start_time
24    if transfer_time <= 0:
25        raise Error("Transfer of %s bytes took nonsensical time %s"
26                    % (url, transfer_time))
27    return (transfer_time, got)
28
29
30def TimeTransferDown(url_pattern, size):
31    url = url_pattern % {'size': size}
32    (transfer_time, got) = TimeTransfer(url, data=None)
33    if len(got) != size:
34      raise Error('Got %d bytes, expected %d' % (len(got), size))
35    return transfer_time
36
37
38def TimeTransferUp(url, size):
39    """If size > 0, POST size bytes to URL, else GET url.  Return time taken."""
40    data = numpy.random.bytes(size)
41    (transfer_time, _) = TimeTransfer(url, data)
42    return transfer_time
43
44
45def BenchmarkOneDirection(latency, label, url, benchmark_function):
46    """Transfer a reasonable amount of data and record the speed.
47
48    Args:
49        latency:  Time for a 1-byte transfer
50        label:  Label to add to perf keyvals
51        url:  URL (or pattern) to transfer at
52        benchmark_function:  Function to perform actual transfer
53    Returns:
54        Key-value dictionary, suitable for reporting to write_perf_keyval.
55        """
56
57    size = 1 << 15              # Start with a small download
58    maximum_size = 1 << 24      # Go large, if necessary
59    multiple = 1
60
61    remaining = 2
62    transfer_time = 0
63
64    # Long enough that startup latency shouldn't dominate.
65    target = max(20 * latency, 10)
66    logging.info('Target time: %s' % target)
67
68    while remaining > 0:
69        size = min(int(size * multiple), maximum_size)
70        transfer_time = benchmark_function(url, size)
71        logging.info('Transfer of %s took %s (%s b/s)'
72                     % (size, transfer_time, 8 * size / transfer_time))
73        if transfer_time >= target:
74            break
75        remaining -= 1
76
77        # Take the latency into account when guessing a size for a
78        # larger transfer.  This is a pretty simple model, but it
79        # appears to work.
80        adjusted_transfer_time = max(transfer_time - latency, 0.01)
81        multiple = target / adjusted_transfer_time
82
83    if remaining == 0:
84        logging.warning(
85            'Max size transfer still took less than minimum desired time %s'
86            % target)
87
88    return {'seconds_%s_fetch_time' % label: transfer_time,
89            'bytes_%s_bytes_transferred' % label: size,
90            'bits_second_%s_speed' % label: 8 * size / transfer_time,
91            }
92
93
94def HttpSpeed(download_url_format_string,
95              upload_url):
96    """Measures upload and download performance to the supplied URLs.
97
98    Args:
99        download_url_format_string:  URL pattern with %(size) for payload bytes
100        upload_url:  URL that accepts large POSTs
101    Returns:
102        A dict of perf_keyval
103    """
104    # We want the download to be substantially longer than the
105    # one-byte fetch time that we can isolate bandwidth instead of
106    # latency.
107    latency = TimeTransferDown(download_url_format_string, 1)
108
109    logging.info('Latency is %s'  % latency)
110
111    down = BenchmarkOneDirection(
112        latency,
113        'downlink',
114        download_url_format_string,
115        TimeTransferDown)
116
117    up = BenchmarkOneDirection(
118        latency,
119        'uplink',
120        upload_url,
121        TimeTransferUp)
122
123    up.update(down)
124    return up
125