1#!/usr/bin/env python
2
3#
4# Copyright (C) 2010 The Android Open Source Project
5#
6# Licensed under the Apache License, Version 2.0 (the "License");
7# you may not use this file except in compliance with the License.
8# You may obtain a copy of the License at
9#
10#      http://www.apache.org/licenses/LICENSE-2.0
11#
12# Unless required by applicable law or agreed to in writing, software
13# distributed under the License is distributed on an "AS IS" BASIS,
14# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
15# See the License for the specific language governing permissions and
16# limitations under the License.
17#
18
19import math
20import optparse
21import sched
22import subprocess
23import sys
24import time
25
26
27def fxrange(start, finish, increment=1.0):
28    """Like xrange, but with float arguments."""
29    steps = int(math.ceil(float(finish - start) / increment))
30
31    if steps < 0:
32        raise ValueError
33
34    for i in xrange(steps):
35        yield start + i * increment
36
37
38def hms(seconds):
39    hours = int(seconds / (60 * 60))
40    seconds -= hours * 60 * 60
41    minutes = int(seconds / 60)
42    seconds -= minutes * 60
43    return '%d:%02d:%02d' % (hours, minutes, seconds)
44
45
46class PeriodicExperiment(object):
47    """Uses the scheduler to run the specified function repeatedly."""
48    def __init__(self,
49                 scheduler=None,
50                 total_duration=8 * 60 * 60,
51                 test_interval=60,
52                 test_function=None):
53        self._scheduler = scheduler
54        self._total_duration = total_duration
55        self._test_interval = test_interval
56        self._test_function = test_function
57        self._start = self._scheduler.timefunc()
58        self._finish = self._start + self._total_duration
59
60    def Run(self):
61        for start_one in fxrange(self._start,
62                                 self._finish,
63                                 self._test_interval):
64            time_remaining = self._finish - start_one
65            self._scheduler.enterabs(start_one,
66                                     1,     # Priority
67                                     self._test_function,
68                                     [time_remaining])
69        self._scheduler.run()
70
71
72class ManualExperiment(object):
73    """Runs the experiment repeatedly, prompting for input each time."""
74    def __init__(self, test_function):
75        self._test_function = test_function
76
77    def Run(self):
78        try:
79            while True:
80                self._test_function(0)    # Pass in a fake time remaining
81                _ = raw_input('Press return to run the test again.  '
82                              'Control-c to exit.')
83        except KeyboardInterrupt:
84            return
85
86class IperfTest(object):
87    def __init__(self, filename, servername, individual_length):
88        self._file = file(filename, 'a')
89        self._servername = servername
90        self._individual_length = individual_length
91
92    def Run(self, remaining):
93        """Run iperf, log output to file, and print."""
94        iperf = ['iperf',
95                 '--client', self._servername,
96                 # Transfer time in seconds.
97                 '--time', str(self._individual_length),
98                 '--reportstyle', 'c' # CSV output
99                 ]
100        print '%s remaining.  Running %s' % (hms(remaining), ' '.join(iperf))
101        result = subprocess.Popen(iperf,
102                                  stdout=subprocess.PIPE).communicate()[0]
103        print result.rstrip()
104        sys.stdout.flush()
105        self._file.write(result)
106        self._file.flush()
107
108    def teardown(self):
109        self._file.close()
110
111def main():
112    default_output = 'stability-' + time.strftime('%Y-%m-%d-%H-%M-%S')
113
114    parser = optparse.OptionParser()
115    parser.add_option('--server', default=None,
116                      help='Machine running the iperf server')
117    parser.add_option('--test_interval', default=60 * 5, type='int',
118                      help='Interval (in seconds) between tests')
119    parser.add_option('--individual_length', default=10, type='int',
120                      help='length (in seconds) of each individual test')
121    parser.add_option('--total_duration', default=8 * 60 * 60, type='int',
122                      help='length (in seconds) for entire test')
123    parser.add_option('--output', default=default_output,
124                      help='Output file')
125    parser.add_option('--manual', default=False, action='store_true',
126                      help='Manual mode; wait for input between every test')
127
128    (options, _) = parser.parse_args()
129
130    if not options.server:
131        print 'No server specified.  Specify a server with --server=SERVER.'
132        exit(2)
133
134    if options.individual_length > options.test_interval:
135        print ('The length of a given bandwidth test must be lower than the '
136               'interval between tests')
137        exit(2)
138
139    s = sched.scheduler(time.time, time.sleep)
140
141    iperf = IperfTest(filename=options.output,
142                      servername=options.server,
143                      individual_length=options.individual_length)
144
145    if options.manual:
146        e = ManualExperiment(test_function=iperf.Run)
147    else:
148        e = PeriodicExperiment(scheduler=s,
149                               total_duration=options.total_duration,
150                               test_interval=options.test_interval,
151                               test_function=iperf.Run)
152    e.Run()
153    iperf.teardown()
154
155if __name__ == '__main__':
156    main()
157