1#! /usr/bin/python 2 3"""A simple heartbeat server. 4 5Executes *readonly* heartbeats against the given database. 6 7Usage: 81. heartbeat_server.py 9 --port 8080 10 11 Start to serve heartbeats on port 8080 using the database credentials 12 found in the shadow_config. One would perform heartbeats for board:lumpy 13 against this server with: 14 curl http://localhost:8080/lumpy. 15 Or just visiting the url through the browser. 16 17 Such a server is capable of handling the following urls: 18 /lumpy: Return formatted heartbeat packets with timing information for 19 each stage, to be viewed in the browser. 20 /lumpy?raw: Return raw json heartbeat packets for lumpy 21 /lumpy?raw&host_limit=1&job_limit=0: Return a 'raw' heartbeat with the 22 first host and not jobs. 23 242. heartbeat_server.py 25 --db_host <ip, eg: production db server> 26 --db_user <user, eg: chromeosqa-admin> 27 --db_password <password, eg: production db password> 28 29 The same as 1. but use the remote db server specified via 30 db_(host,user,password). 31""" 32 33 34import argparse 35import sys 36import time 37import urlparse 38from BaseHTTPServer import BaseHTTPRequestHandler 39from BaseHTTPServer import HTTPServer 40 41import common 42from autotest_lib.client.common_lib.global_config import global_config as config 43from autotest_lib.frontend import setup_django_environment 44 45 46# Populated with command line database credentials. 47DB_SETTINGS = { 48 'ENGINE': 'autotest_lib.frontend.db.backends.afe', 49} 50 51# Indent level used when formatting json for the browser. 52JSON_FORMATTING_INDENT = 4 53 54 55def time_call(func): 56 """A simple timer wrapper. 57 58 @param func: The function to wrap. 59 """ 60 def wrapper(*args, **kwargs): 61 """Wrapper returned by time_call decorator.""" 62 start = time.time() 63 res = func(*args, **kwargs) 64 return time.time()-start, res 65 return wrapper 66 67 68class BoardHandler(BaseHTTPRequestHandler): 69 """Handles heartbeat urls.""" 70 71 # Prefix for all board labels. 72 board_prefix = 'board:' 73 74 75 @staticmethod 76 @time_call 77 def _get_jobs(board, job_limit=None): 78 jobs = models.Job.objects.filter( 79 dependency_labels__name=board).exclude( 80 hostqueueentry__complete=True).exclude( 81 hostqueueentry__active=True) 82 return jobs[:job_limit] if job_limit is not None else jobs 83 84 85 @staticmethod 86 @time_call 87 def _get_hosts(board, host_limit=None): 88 hosts = models.Host.objects.filter( 89 labels__name__in=[board], leased=False) 90 return hosts[:host_limit] if host_limit is not None else hosts 91 92 93 @staticmethod 94 @time_call 95 def _create_packet(hosts, jobs): 96 return { 97 'hosts': [h.serialize() for h in hosts], 98 'jobs': [j.serialize() for j in jobs] 99 } 100 101 102 def do_GET(self): 103 """GET handler. 104 105 Handles urls like: http://localhost:8080/lumpy?raw&host_limit=5 106 and writes the appropriate http response containing the heartbeat. 107 """ 108 parsed_path = urlparse.urlparse(self.path, allow_fragments=True) 109 board = '%s%s' % (self.board_prefix, parsed_path.path.rsplit('/')[-1]) 110 111 raw = False 112 job_limit = None 113 host_limit = None 114 for query in parsed_path.query.split('&'): 115 split_query = query.split('=') 116 if split_query[0] == 'job_limit': 117 job_limit = int(split_query[1]) 118 elif split_query[0] == 'host_limit': 119 host_limit = int(split_query[1]) 120 elif split_query[0] == 'raw': 121 raw = True 122 123 host_time, hosts = self._get_hosts(board, host_limit) 124 job_time, jobs = self._get_jobs(board, job_limit) 125 126 serialize_time, heartbeat_packet = self._create_packet(hosts, jobs) 127 self.send_response(200) 128 self.end_headers() 129 130 # Format browser requests, the heartbeat client will request using ?raw 131 # while the browser will perform a plain request like 132 # http://localhost:8080/lumpy. The latter needs to be human readable and 133 # include more details timing information. 134 json_encoder = django_encoder.DjangoJSONEncoder() 135 if not raw: 136 json_encoder.indent = JSON_FORMATTING_INDENT 137 self.wfile.write('Serialize: %s,\nJob query: %s\nHost query: %s\n' 138 'Hosts: %s\nJobs: %s\n' % 139 (serialize_time, job_time, host_time, 140 len(heartbeat_packet['hosts']), 141 len(heartbeat_packet['jobs']))) 142 self.wfile.write(json_encoder.encode(heartbeat_packet)) 143 return 144 145 146def _parse_args(args): 147 parser = argparse.ArgumentParser( 148 description='Start up a simple heartbeat server on localhost.') 149 parser.add_argument( 150 '--port', default=8080, 151 help='The port to start the heartbeat server.') 152 parser.add_argument( 153 '--db_host', 154 default=config.get_config_value('AUTOTEST_WEB', 'host'), 155 help='Db server ip address.') 156 parser.add_argument( 157 '--db_name', 158 default=config.get_config_value('AUTOTEST_WEB', 'database'), 159 help='Name of the db table.') 160 parser.add_argument( 161 '--db_user', 162 default=config.get_config_value('AUTOTEST_WEB', 'user'), 163 help='User for the db server.') 164 parser.add_argument( 165 '--db_password', 166 default=config.get_config_value('AUTOTEST_WEB', 'password'), 167 help='Password for the db server.') 168 parser.add_argument( 169 '--db_port', 170 default=config.get_config_value('AUTOTEST_WEB', 'port', default=''), 171 help='Port of the db server.') 172 173 return parser.parse_args(args) 174 175 176if __name__ == '__main__': 177 args = _parse_args(sys.argv[1:]) 178 server = HTTPServer(('localhost', args.port), BoardHandler) 179 print ('Starting heartbeat server, query eg: http://localhost:%s/lumpy' % 180 args.port) 181 182 # We need these lazy imports to allow command line specification of 183 # database credentials. 184 from autotest_lib.frontend import settings 185 DB_SETTINGS['HOST'] = args.db_host 186 DB_SETTINGS['NAME'] = args.db_name 187 DB_SETTINGS['USER'] = args.db_user 188 DB_SETTINGS['PASSWORD'] = args.db_password 189 DB_SETTINGS['PORT'] = args.db_port 190 settings.DATABASES['default'] = DB_SETTINGS 191 from autotest_lib.frontend.afe import models 192 from django.core.serializers import json as django_encoder 193 194 server.serve_forever() 195 196