Add liveness probe to services (#150)

Add a --liveness option to zocalo.service which starts
a liveness endpoint on the port defined by --liveness-port
This sends a message to the service, which, if alive, will
return a message to the frontend confirming it is alive.
If confirmation of aliveness is received the endpoint will
return a 200 status code.

This is intended for use as a livenessProbe when running services
on Kubernetes. It should force a container restart in the event
that a pod is blocked indefinitely.

Author: rjgildea

src/workflows/contrib/start_service.py

@@ -82,6 +82,31 @@ def run(
             help="Name of the service to start. Known services: "
             + ", ".join(known_services),
         )
+        parser.add_option(
+            "--liveness",
+            dest="liveness",
+            action="store_true",
+            default=False,
+            help=(
+                "Expose a liveness check endpoint for this service."
+                "A return code of 200 indicates success."
+                "A return code of 408 indicates failure."
+            ),
+        )
+        parser.add_option(
+            "--liveness-port",
+            dest="liveness_port",
+            default=8000,
+            type="int",
+            help="Expose liveness check endpoint on this port.",
+        )
+        parser.add_option(
+            "--liveness-timeout",
+            dest="liveness_timeout",
+            default=30,
+            type="float",
+            help="Timeout for the liveness check (in seconds).",
+        )
         if add_metrics_option:
             parser.add_option(
                 "-m",
@@ -152,6 +177,12 @@ def on_transport_preparation_hook():
         if add_metrics_option and options.metrics:
             kwargs["environment"]["metrics"] = {"port": options.metrics_port}
 
+        if options.liveness:
+            kwargs["environment"]["liveness"] = {
+                "port": options.liveness_port,
+                "timeout": options.liveness_timeout,
+            }
+
         # Call before_frontend_construction hook
         kwargs = self.before_frontend_construction(kwargs) or kwargs
 

src/workflows/frontend/__init__.py

@@ -128,6 +128,9 @@ def __getitem__(self, key):
         else:
             self.update_status()
 
+        if environment and "liveness" in environment:
+            self._start_liveness_endpoint(environment["liveness"]["port"])
+
     def update_status(self, status_code=None):
         """Update the service status kept inside the frontend (_service_status).
         The status is broadcast over the network immediately. If the status
@@ -339,6 +342,11 @@ def parse_band_status_update(self, message):
         self.log.debug("Status update: " + str(message))
         self.update_status(status_code=message["statuscode"])
 
+    def parse_band_liveness_check(self, message):
+        """Respond by sending message to backend to let it know we are alive."""
+        self.log.debug("Service liveness check: alive!")
+        self.__alive = True
+
     def get_host_id(self):
         """Get a cached copy of the host id."""
         return self.__hostid
@@ -451,3 +459,27 @@ def _terminate_service(self):
             if self._service:
                 self._service.join()  # must wait for process to be actually destroyed
             self._service = None
+
+    def _start_liveness_endpoint(self, port: int):
+        from wsgiref.simple_server import make_server
+
+        def alive(environ, start_response):
+            self.__alive = False
+            self.send_command({"band": "command", "payload": "liveness_check"})
+
+            while True:
+                if self.__alive:
+                    status = "200 OK"
+                    headers = [
+                        ("Content-type", "text/plain; charset=utf-8")
+                    ]  # HTTP Headers
+                    start_response(status, headers)
+                    return [b"ok"]
+
+        httpd = make_server("", port, alive)
+        self.log.debug(f"Serving liveness check on port {port}...")
+
+        # Serve until process is killed
+        t = threading.Thread(target=httpd.serve_forever)
+        t.daemon = True
+        t.start()

src/workflows/services/common_service.py

@@ -200,7 +200,6 @@ def start_transport(self):
                 port = metrics["port"]
                 self.log.debug(f"Starting metrics endpoint on port {port}")
                 prometheus_client.start_http_server(port=port)
-
         else:
             self.log.debug("No transport layer defined for service. Skipping.")
 
@@ -500,9 +499,12 @@ def __process_command(self, command):
         """Process an incoming command message from the frontend."""
         if command == Commands.SHUTDOWN:
             self.__shutdown = True
+        elif command == Commands.LIVENESS_CHECK:
+            self.__send_to_frontend({"band": "liveness_check", "payload": "alive"})
 
 
 class Commands:
     """A list of command strings used for communicating with the frontend."""
 
     SHUTDOWN = "shutdown"
+    LIVENESS_CHECK = "liveness_check"

tests/contrib/test_start_service.py

@@ -29,6 +29,7 @@ def test_script_initialises_transport_and_starts_frontend(
     mock_options = mock.Mock()
     mock_options.service = "someservice"
     mock_options.transport = mock.sentinel.transport
+    mock_options.liveness = False
     mock_parser.return_value.parse_args.return_value = (mock_options, mock.Mock())
     mock_services.get_known_services.return_value = {"SomeService": None}
 
@@ -58,6 +59,7 @@ def test_add_metrics_option(mock_services, mock_frontend, mock_tlookup, mock_par
     mock_options = mock.Mock()
     mock_options.service = "someservice"
     mock_options.transport = mock.sentinel.transport
+    mock_options.liveness = False
     mock_options.metrics = False
     mock_options.metrics_port = 4242
     mock_parser.return_value.parse_args.return_value = (mock_options, mock.Mock())

tests/frontend/test_frontend.py

@@ -141,13 +141,14 @@ def test_start_service_in_frontend(mock_transport, mock_mp):
     mock_mp.Pipe.side_effect = [(pipes[0], pipes[1]), (pipes[2], pipes[3]), None]
 
     # initialize frontend
-    fe = workflows.frontend.Frontend(environment=mock.sentinel.environment)
+    environment = {mock.sentinel.foo: mock.sentinel.bar}
+    fe = workflows.frontend.Frontend(environment=environment)
 
     # start service
     fe.switch_service(mock_service)
 
     # check service was started properly
-    mock_service.assert_called_once_with(environment=mock.sentinel.environment)
+    mock_service.assert_called_once_with(environment=environment)
     mock_service.return_value.connect.assert_called_once_with(
         commands=pipes[0], frontend=pipes[3]
     )