Supports the backend port being sent in the HAProxy send-state header

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'):

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),

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']),