Home | History | Annotate | Download | only in shard 
      1 #! /usr/bin/python
      2 
      3 """A simple heartbeat client.
      4 
      5 Executes heartbeats against a simple_heartbeat_server running on the give
      6 --server address and deserializes records into an in memory sqlite database.
      7 
      8 Usage:
      9 1. heartbeat_client.py
     10     --server http://localhost:8080
     11     --board lumpy
     12 
     13     Perform a heartbeat against the given server for the given board,
     14     and deserialize records into a sqlite database.
     15 
     16 2. heartbeat_client.py
     17     --server http://localhost:8080
     18     --board lumpy
     19     --host_limit 1 --job_limit 100
     20 
     21     Do the same as 1, but instruct the server to limit the hosts to 1
     22     and jobs to 100. This is useful for debugging issues with only jobs/hosts.
     23 """
     24 
     25 
     26 from json import decoder
     27 import argparse
     28 import sys
     29 import urllib2
     30 
     31 import common
     32 
     33 from autotest_lib.scheduler.shard import simple_heartbeat_server
     34 from autotest_lib.frontend import setup_django_environment
     35 from autotest_lib.frontend.afe import frontend_test_utils
     36 from autotest_lib.frontend.afe import models
     37 
     38 json_decoder = decoder.JSONDecoder()
     39 
     40 
     41 class HeartbeatHandler(frontend_test_utils.FrontendTestMixin):
     42     """Performs heartbeats and deserializes into an in memory database."""
     43 
     44     _config_section = 'AUTOTEST_WEB'
     45 
     46 
     47     def __init__(self, server):
     48         """Initialize a heartbeat server.
     49 
     50         @param server: The address of a simple_heartbeat_server.
     51         """
     52         self.server = server
     53         self._frontend_common_setup(setup_tables=True, fill_data=False)
     54 
     55 
     56     @staticmethod
     57     @simple_heartbeat_server.time_call
     58     def get_heartbeat_packet(server, board, host_limit, job_limit):
     59         """Perform the heartbeat.
     60 
     61         Constructs a url like: http://localhost:8080/lumpy?raw&host_limit=3
     62         and does a urlopen.
     63 
     64         @param server: The address of a simple_heartbeat_server.
     65         @param host_limit: The number of hosts to include in the heartbeat.
     66         @param job_limit: The number of jobs to include in the heartbeat.
     67 
     68         @return: A string containing the heartbeat packet.
     69         """
     70         url = '%s/%s?raw' % (server, board)
     71         if job_limit:
     72             url = '%s&job_limit=%s' % (url, job_limit)
     73         if host_limit:
     74             url = '%s&host_limit=%s' % (url, host_limit)
     75         print 'Performing heartbeat against %s' % url
     76         return urllib2.urlopen(url).read()
     77 
     78 
     79     @staticmethod
     80     @simple_heartbeat_server.time_call
     81     def deserialize_heartbeat(packet):
     82         """Deserialize the given heartbeat packet into an in memory database.
     83 
     84         @param packet: A string representing the heartbeat packet containing
     85                 jobs and hosts.
     86 
     87         @return: The json decoded heartbeat.
     88         """
     89         response = json_decoder.decode(packet)
     90         [models.Host.deserialize(h) for h in response['hosts']]
     91         [models.Job.deserialize(j) for j in response['jobs']]
     92         return response
     93 
     94 
     95     def perform_heartbeat(self, board, host_limit, job_limit):
     96         """Perform a heartbeat against the given server, for the given board.
     97 
     98         @param board: Boardname, eg: lumpy.
     99         @param host_limit: Limit number of hosts retrieved.
    100         @param job_limit: Limit number of jobs retrieved.
    101         """
    102         timing, packet = self.get_heartbeat_packet(
    103                 self.server, board, host_limit, job_limit)
    104         print 'Time to perform heartbeat %s' % timing
    105         timing, response = self.deserialize_heartbeat(packet)
    106         print 'Time to deserialize hearbeat %s' % timing
    107         print ('Jobs: %s, Hosts: %s' %
    108                 (len(response['jobs']), len(response['hosts'])))
    109 
    110 
    111 def _parse_args(args):
    112     parser = argparse.ArgumentParser(
    113             description='Start up a simple heartbeat client.')
    114     parser.add_argument(
    115             '--server', default='http://localhost:8080',
    116             help='Address of a simple_heartbeat_server to heartbeat against.')
    117     parser.add_argument(
    118             '--board', default='lumpy',
    119             help='Heartbeats can only be performed '
    120                  'against a specific board, eg: lumpy.')
    121     parser.add_argument(
    122             '--host_limit', default='',
    123             help='Limit hosts in the heartbeat.')
    124     parser.add_argument(
    125             '--job_limit', default='',
    126             help='Limit jobs in the heartbeat.')
    127     args = parser.parse_args(args)
    128     args.board = args.board.lstrip(
    129             simple_heartbeat_server.BoardHandler.board_prefix)
    130     return args
    131 
    132 
    133 if __name__ == '__main__':
    134     args = _parse_args(sys.argv[1:])
    135     HeartbeatHandler(args.server).perform_heartbeat(
    136             args.board, args.host_limit, args.job_limit)
    137