src/tools/swarming_client/tools/swarming_load_test_bot.py

   1 #!/usr/bin/env python
   2 # Copyright 2013 The Swarming Authors. All rights reserved.
   3 # Use of this source code is governed under the Apache License, Version 2.0 that
   4 # can be found in the LICENSE file.
   5
   6 """Triggers a ton of fake jobs to test its handling under high load.
   7
   8 Generates an histogram with the latencies to process the tasks and number of
   9 retries.
  10 """
  11
  12 import hashlib
  13 import json
  14 import logging
  15 import optparse
  16 import os
  17 import Queue
  18 import socket
  19 import StringIO
  20 import sys
  21 import threading
  22 import time
  23 import zipfile
  24
  25 ROOT_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
  26
  27 sys.path.insert(0, ROOT_DIR)
  28
  29 from third_party import colorama
  30
  31 import swarming
  32
  33 from utils import graph
  34 from utils import net
  35 from utils import threading_utils
  36
  37 # Line too long (NN/80)
  38 # pylint: disable=C0301
  39
  40 OS_NAME = 'Comodore64'
  41 TASK_OUTPUT = 'This task ran with great success'
  42
  43
  44 def print_results(results, columns, buckets):
  45   delays = [i for i in results if isinstance(i, float)]
  46   failures = [i for i in results if not isinstance(i, float)]
  47
  48   print('%sDELAYS%s:' % (colorama.Fore.RED, colorama.Fore.RESET))
  49   graph.print_histogram(
  50       graph.generate_histogram(delays, buckets), columns, ' %.3f')
  51   print('')
  52   print('Total items  : %d' % len(results))
  53   average = 0
  54   if delays:
  55     average = sum(delays)/ len(delays)
  56   print('Average delay: %s' % graph.to_units(average))
  57   print('')
  58
  59   if failures:
  60     print('%sEVENTS%s:' % (colorama.Fore.RED, colorama.Fore.RESET))
  61     values = {}
  62     for f in failures:
  63       values.setdefault(f, 0)
  64       values[f] += 1
  65     graph.print_histogram(values, columns, ' %s')
  66     print('')
  67
  68
  69 def calculate_version(url):
  70   """Retrieves the swarm_bot code and returns the SHA-1 for it."""
  71   # Cannot use url_open() since zipfile requires .seek().
  72   archive = zipfile.ZipFile(StringIO.StringIO(net.url_read(url)))
  73   # See
  74   # https://code.google.com/p/swarming/source/browse/services/swarming/common/bot_archive.py
  75   d = hashlib.sha1()
  76   for f in archive.namelist():
  77     d.update(archive.read(f))
  78   return d.hexdigest()
  79
  80
  81 class FakeSwarmBot(object):
  82   """This is a Fake swarm_bot implementation simulating it is running
  83   Comodore64.
  84
  85   It polls for job, acts as if it was processing them and return the fake
  86   result.
  87   """
  88   def __init__(
  89       self, swarming_url, dimensions, swarm_bot_version_hash, index, progress,
  90       duration, events, kill_event):
  91     self._lock = threading.Lock()
  92     self._swarming = swarming_url
  93     self._index = index
  94     self._progress = progress
  95     self._duration = duration
  96     self._events = events
  97     self._kill_event = kill_event
  98     # Use an impossible hostname.
  99     self._machine_id = '%s-%d' % (socket.getfqdn().lower(), index)
 100
 101     # See
 102     # https://code.google.com/p/swarming/source/browse/src/swarm_bot/slave_machine.py?repo=swarming-server
 103     # and
 104     # https://chromium.googlesource.com/chromium/tools/build.git/ \
 105     #    +/master/scripts/tools/swarm_bootstrap/swarm_bootstrap.py
 106     # for more details.
 107     self._attributes = {
 108       'dimensions': dimensions,
 109       'id': self._machine_id,
 110       'try_count': 0,
 111       'tag': self._machine_id,
 112       'version': swarm_bot_version_hash,
 113     }
 114
 115     self._thread = threading.Thread(target=self._run, name='bot%d' % index)
 116     self._thread.daemon = True
 117     self._thread.start()
 118
 119   def join(self):
 120     self._thread.join()
 121
 122   def is_alive(self):
 123     return self._thread.is_alive()
 124
 125   def _run(self):
 126     try:
 127       self._progress.update_item('%d alive' % self._index, bots=1)
 128       while True:
 129         if self._kill_event.is_set():
 130           return
 131         data = {'attributes': json.dumps(self._attributes)}
 132         request = net.url_open(self._swarming + '/poll_for_test', data=data)
 133         if request is None:
 134           self._events.put('poll_for_test_empty')
 135           continue
 136         start = time.time()
 137         try:
 138           manifest = json.load(request)
 139         except ValueError:
 140           self._progress.update_item('Failed to poll')
 141           self._events.put('poll_for_test_invalid')
 142           continue
 143
 144         commands = [c['function'] for c in manifest.get('commands', [])]
 145         if not commands:
 146           # Nothing to run.
 147           self._events.put('sleep')
 148           time.sleep(manifest['come_back'])
 149           continue
 150
 151         if commands == ['UpdateSlave']:
 152           # Calculate the proper SHA-1 and loop again.
 153           # This could happen if the Swarming server is upgraded while this
 154           # script runs.
 155           self._attributes['version'] = calculate_version(
 156               manifest['commands'][0]['args'])
 157           self._events.put('update_slave')
 158           continue
 159
 160         if commands != ['StoreFiles', 'RunCommands']:
 161           self._progress.update_item(
 162               'Unexpected RPC call %s\n%s' % (commands, manifest))
 163           self._events.put('unknown_rpc')
 164           break
 165
 166         # The normal way Swarming works is that it 'stores' a test_run.swarm
 167         # file and then defer control to swarm_bot/local_test_runner.py.
 168         store_cmd = manifest['commands'][0]
 169         assert len(store_cmd['args']) == 1, store_cmd['args']
 170         filepath, filename, test_run_content = store_cmd['args'][0]
 171         assert filepath == ''
 172         assert filename == 'test_run.swarm'
 173         assert 'local_test_runner.py' in manifest['commands'][1]['args'][0], (
 174             manifest['commands'][1])
 175         result_url = manifest['result_url']
 176         test_run = json.loads(test_run_content)
 177         assert result_url == test_run['result_url']
 178         ping_url = test_run['ping_url']
 179         ping_delay = test_run['ping_delay']
 180         self._progress.update_item('%d processing' % self._index, processing=1)
 181
 182         # Fake activity and send pings as requested.
 183         while True:
 184           remaining = max(0, (start + self._duration) - time.time())
 185           if remaining > ping_delay:
 186             # Include empty data to ensure the request is a POST request.
 187             result = net.url_read(ping_url, data={})
 188             assert result == 'Success.', result
 189             remaining = max(0, (start + self._duration) - time.time())
 190           if not remaining:
 191             break
 192           time.sleep(remaining)
 193
 194         data = {
 195           'c': test_run['configuration']['config_name'],
 196           'n': test_run['test_run_name'],
 197           'o': False,
 198           'result_output': TASK_OUTPUT,
 199           's': True,
 200           'x': '0',
 201         }
 202         result = net.url_read(manifest['result_url'], data=data)
 203         self._progress.update_item(
 204             '%d processed' % self._index, processing=-1, processed=1)
 205         if not result:
 206           self._events.put('result_url_fail')
 207         else:
 208           assert result == 'Successfully update the runner results.', result
 209           self._events.put(time.time() - start)
 210     finally:
 211       try:
 212         # Unregister itself. Otherwise the server will have tons of fake slaves
 213         # that the admin will have to remove manually.
 214         response = net.url_open(
 215             self._swarming + '/delete_machine_stats',
 216             data=[('r', self._machine_id)])
 217         if not response:
 218           self._events.put('failed_unregister')
 219         else:
 220           response.read()
 221       finally:
 222         self._progress.update_item('%d quit' % self._index, bots=-1)
 223
 224
 225 def main():
 226   colorama.init()
 227   parser = optparse.OptionParser(description=sys.modules[__name__].__doc__)
 228   parser.add_option(
 229       '-S', '--swarming',
 230       metavar='URL', default='',
 231       help='Swarming server to use')
 232   swarming.add_filter_options(parser)
 233   # Use improbable values to reduce the chance of interferring with real slaves.
 234   parser.set_defaults(
 235       dimensions=[
 236         ('bits', '36'),
 237         ('machine', os.uname()[4] + '-experimental'),
 238         ('os', OS_NAME),
 239       ])
 240
 241   group = optparse.OptionGroup(parser, 'Load generated')
 242   group.add_option(
 243       '--slaves', type='int', default=300, metavar='N',
 244       help='Number of swarm bot slaves, default: %default')
 245   group.add_option(
 246       '-c', '--consume', type='float', default=60., metavar='N',
 247       help='Duration (s) for consuming a request, default: %default')
 248   parser.add_option_group(group)
 249
 250   group = optparse.OptionGroup(parser, 'Display options')
 251   group.add_option(
 252       '--columns', type='int', default=graph.get_console_width(), metavar='N',
 253       help='For histogram display, default:%default')
 254   group.add_option(
 255       '--buckets', type='int', default=20, metavar='N',
 256       help='Number of buckets for histogram display, default:%default')
 257   parser.add_option_group(group)
 258
 259   parser.add_option(
 260       '--dump', metavar='FOO.JSON', help='Dumps to json file')
 261   parser.add_option(
 262       '-v', '--verbose', action='store_true', help='Enables logging')
 263
 264   options, args = parser.parse_args()
 265   logging.basicConfig(level=logging.INFO if options.verbose else logging.FATAL)
 266   if args:
 267     parser.error('Unsupported args: %s' % args)
 268   options.swarming = options.swarming.rstrip('/')
 269   if not options.swarming:
 270     parser.error('--swarming is required.')
 271   if options.consume <= 0:
 272     parser.error('Needs --consume > 0. 0.01 is a valid value.')
 273   swarming.process_filter_options(parser, options)
 274
 275   print(
 276       'Running %d slaves, each task lasting %.1fs' % (
 277         options.slaves, options.consume))
 278   print('Ctrl-C to exit.')
 279   print('[processing/processed/bots]')
 280   columns = [('processing', 0), ('processed', 0), ('bots', 0)]
 281   progress = threading_utils.Progress(columns)
 282   events = Queue.Queue()
 283   start = time.time()
 284   kill_event = threading.Event()
 285   swarm_bot_version_hash = calculate_version(
 286       options.swarming + '/get_slave_code')
 287   slaves = [
 288     FakeSwarmBot(
 289       options.swarming, options.dimensions, swarm_bot_version_hash, i, progress,
 290       options.consume, events, kill_event)
 291     for i in range(options.slaves)
 292   ]
 293   try:
 294     # Wait for all the slaves to come alive.
 295     while not all(s.is_alive() for s in slaves):
 296       time.sleep(0.01)
 297     progress.update_item('Ready to run')
 298     while slaves:
 299       progress.print_update()
 300       time.sleep(0.01)
 301       # The slaves could be told to die.
 302       slaves = [s for s in slaves if s.is_alive()]
 303   except KeyboardInterrupt:
 304     kill_event.set()
 305
 306   progress.update_item('Waiting for slaves to quit.', raw=True)
 307   progress.update_item('')
 308   while slaves:
 309     progress.print_update()
 310     slaves = [s for s in slaves if s.is_alive()]
 311   # At this point, progress is not used anymore.
 312   print('')
 313   print('Ran for %.1fs.' % (time.time() - start))
 314   print('')
 315   results = events.queue
 316   print_results(results, options.columns, options.buckets)
 317   if options.dump:
 318     with open(options.dump, 'w') as f:
 319       json.dump(results, f, separators=(',',':'))
 320   return 0
 321
 322
 323 if __name__ == '__main__':
 324   sys.exit(main())