diff --git a/hacheck/handlers.py b/hacheck/handlers.py index 97da557..1ede064 100644 --- a/hacheck/handlers.py +++ b/hacheck/handlers.py @@ -67,9 +67,37 @@ def get(self): class BaseServiceHandler(tornado.web.RequestHandler): CHECKERS = [] + def maybe_get_port_from_haproxy_server_state(self): + """ + Look for the 'X-Haproxy-Server-State' header and try to parse out the + 'port' value. + + Note that only very recent versions of HAProxy support sending + the port in the send-state header. In particular you need + commit 514061c414080701cb046171041a2d00859660e8 from haproxy dev. + + Example server state header: + + X-Haproxy-Server-State: UP 2/3; address=srv2; port=1234; + name=bck/srv2; node=lb1; weight=1/2; scur=13/22; qcur=0 + + returns: the string-typed port if found, else None. + """ + + server_state = self.request.headers.get('X-Haproxy-Server-State', '') + + parts = server_state.split(';') + parts = [part.strip() for part in parts] + parts = [part.split('=') for part in parts] + parts = [part for part in parts if len(part) == 2] + + return dict(parts).get('port') + @tornado.web.asynchronous @tornado.gen.coroutine def get(self, service_name, port, query): + port = self.maybe_get_port_from_haproxy_server_state() or port + seen_services[service_name] = time.time() service_count[service_name][self.request.remote_ip] += 1 with cache.maybe_bust(self.request.headers.get('Pragma', '') == 'no-cache'): diff --git a/hacheck/main.py b/hacheck/main.py index c934363..22a2ca2 100644 --- a/hacheck/main.py +++ b/hacheck/main.py @@ -30,10 +30,10 @@ def log_request(handler): def get_app(): return tornado.web.Application([ - (r'/http/([a-zA-Z0-9_-]+)/([0-9]+)/(.*)', handlers.HTTPServiceHandler), - (r'/tcp/([a-zA-Z0-9_-]+)/([0-9]+)/?(.*)', handlers.TCPServiceHandler), - (r'/mysql/([a-zA-Z0-9_-]+)/([0-9]+)/?(.*)', handlers.MySQLServiceHandler), - (r'/spool/([a-zA-Z0-9_-]+)/([0-9]+)/?(.*)', handlers.SpoolServiceHandler), + (r'/http/([.a-zA-Z0-9_-]+)/([0-9]+)/(.*)', handlers.HTTPServiceHandler), + (r'/tcp/([.a-zA-Z0-9_-]+)/([0-9]+)/?(.*)', handlers.TCPServiceHandler), + (r'/mysql/([.a-zA-Z0-9_-]+)/([0-9]+)/?(.*)', handlers.MySQLServiceHandler), + (r'/spool/([.a-zA-Z0-9_-]+)/([0-9]+)/?(.*)', handlers.SpoolServiceHandler), (r'/recent', handlers.ListRecentHandler), (r'/status/count', handlers.ServiceCountHandler), (r'/status', handlers.StatusHandler), diff --git a/tests/test_application.py b/tests/test_application.py index ef9e646..8cd28a4 100644 --- a/tests/test_application.py +++ b/tests/test_application.py @@ -144,6 +144,28 @@ def test_weird_code(self): response = self.fetch('/http/uncached-weird-code/80/status') self.assertEqual(503, response.code) + def test_haproxy_server_state(self): + rv = tornado.concurrent.Future() + rv.set_result((200, b'OK')) + checker = mock.Mock(return_value=rv) + server_state = 'UP 2/3; addr=srv1; port=1234; name=bck/srv2; node=lb1; weight=1/2; scur=13/22; qcur=0' + with mock.patch.object(handlers.HTTPServiceHandler, 'CHECKERS', [checker]): + response = self.fetch('/http/foo/1/status', headers={'X-Haproxy-Server-State': server_state}) + self.assertEqual(200, response.code) + args, _ = checker.call_args + assert args[1] == 1234 + + def test_old_haproxy_server_state_ignored(self): + rv = tornado.concurrent.Future() + rv.set_result((200, b'OK')) + checker = mock.Mock(return_value=rv) + server_state = 'UP 2/3; name=bck/srv2; node=lb1; weight=1/2; scur=13/22; qcur=0' + with mock.patch.object(handlers.HTTPServiceHandler, 'CHECKERS', [checker]): + response = self.fetch('/http/foo/1/status', headers={'X-Haproxy-Server-State': server_state}) + self.assertEqual(200, response.code) + args, _ = checker.call_args + assert args[1] == 1 + def test_option_parsing(self): with nested( mock.patch('sys.argv', ['ignorethis', '-c', self.config_file.name, '--spool-root', 'foo']),