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