Merge pull request #1598 from manics/health-customisable

Add `health_handler_class` to make `/health` handler is customisable

Author: consideRatio

binderhub/app.py

@@ -41,11 +41,11 @@
 from traitlets.config import Application
 
 from .base import AboutHandler, Custom404, VersionHandler
-from .build import Build, BuildExecutor
+from .build import Build, BuildExecutor, KubernetesBuildExecutor
 from .builder import BuildHandler
 from .config import ConfigHandler
 from .events import EventLog
-from .health import HealthHandler
+from .health import HealthHandler, KubernetesHealthHandler
 from .launcher import Launcher
 from .log import log_request
 from .main import LegacyRedirectHandler, MainHandler, ParameterizedMainHandler
@@ -290,6 +290,18 @@ def _valid_badge_base_url(self, proposal):
         config=True,
     )
 
+    health_handler_class = Type(
+        HealthHandler,
+        help="The Tornado /health handler class",
+        config=True,
+    )
+
+    @default("health_handler_class")
+    def _default_health_handler_class(self):
+        if issubclass(self.build_class, KubernetesBuildExecutor):
+            return KubernetesHealthHandler
+        return HealthHandler
+
     per_repo_quota = Integer(
         0,
         help="""
@@ -924,7 +936,7 @@ def initialize(self, *args, **kwargs):
                 {"path": os.path.join(self.tornado_settings["static_path"], "images")},
             ),
             (r"/about", AboutHandler),
-            (r"/health", HealthHandler, {"hub_url": self.hub_url_local}),
+            (r"/health", self.health_handler_class, {"hub_url": self.hub_url_local}),
             (r"/_config", ConfigHandler),
             (r"/", MainHandler),
             (r".*", Custom404),

binderhub/health.py

@@ -105,34 +105,6 @@ class HealthHandler(BaseHandler):
     def initialize(self, hub_url=None):
         self.hub_url = hub_url
 
-    @at_most_every
-    async def _get_pods(self):
-        """Get information about build and user pods"""
-        namespace = self.settings["build_namespace"]
-        k8s = self.settings["kubernetes_client"]
-        pool = self.settings["executor"]
-
-        app_log.info(f"Getting pod statistics for {namespace}")
-
-        label_selectors = [
-            "app=jupyterhub,component=singleuser-server",
-            "component=binderhub-build",
-        ]
-        requests = [
-            asyncio.wrap_future(
-                pool.submit(
-                    k8s.list_namespaced_pod,
-                    namespace,
-                    label_selector=label_selector,
-                    _preload_content=False,
-                    _request_timeout=KUBE_REQUEST_TIMEOUT,
-                )
-            )
-            for label_selector in label_selectors
-        ]
-        responses = await asyncio.gather(*requests)
-        return [json.loads(resp.read())["items"] for resp in responses]
-
     @false_if_raises
     @retry
     async def check_jupyterhub_api(self, hub_url):
@@ -155,58 +127,43 @@ async def check_docker_registry(self):
         )
         return True
 
-    async def check_pod_quota(self):
-        """Compare number of active pods to available quota"""
-        user_pods, build_pods = await self._get_pods()
-
-        n_user_pods = len(user_pods)
-        n_build_pods = len(build_pods)
+    def get_checks(self, checks):
+        """Add health checks to the `checks` dict
 
-        quota = self.settings["pod_quota"]
-        total_pods = n_user_pods + n_build_pods
-        usage = {
-            "total_pods": total_pods,
-            "build_pods": n_build_pods,
-            "user_pods": n_user_pods,
-            "quota": quota,
-            "ok": total_pods <= quota if quota is not None else True,
-        }
-        return usage
+        checks: Dictionary, updated in-place:
+          key: service name
+          value: a future that resolves to either:
+            - a bool (success/fail)
+            - a dict with the field `"ok": bool` plus other information
+        """
+        if self.settings["use_registry"]:
+            checks["Docker registry"] = self.check_docker_registry()
+        checks["JupyterHub API"] = self.check_jupyterhub_api(self.hub_url)
 
     async def check_all(self):
-        """Runs all health checks and returns a tuple (overall, checks).
+        """Runs all health checks and returns a tuple (overall, results).
 
         `overall` is a bool representing the overall status of the service
-        `checks` contains detailed information on each check's result
+        `results` contains detailed information on each check's result
         """
-        checks = []
-        check_futures = []
-
-        if self.settings["use_registry"]:
-            check_futures.append(self.check_docker_registry())
-            checks.append({"service": "Docker registry", "ok": False})
-
-        check_futures.append(self.check_jupyterhub_api(self.hub_url))
-        checks.append({"service": "JupyterHub API", "ok": False})
-
-        check_futures.append(self.check_pod_quota())
-        checks.append({"service": "Pod quota", "ok": False})
+        checks = {}
+        results = []
+        self.get_checks(checks)
 
-        for result, check in zip(await asyncio.gather(*check_futures), checks):
+        for result, service in zip(
+            await asyncio.gather(*checks.values()), checks.keys()
+        ):
             if isinstance(result, bool):
-                check["ok"] = result
+                results.append({"service": service, "ok": result})
             else:
-                check.update(result)
+                results.append(dict({"service": service}, **result))
 
-        # The pod quota is treated as a soft quota this means being above
-        # quota doesn't mean the service is unhealthy
-        overall = all(
-            check["ok"] for check in checks if check["service"] != "Pod quota"
-        )
+        # Some checks are for information but do not count as a health failure
+        overall = all(r["ok"] for r in results if not r.get("_ignore_failure", False))
         if not overall:
-            unhealthy = [check for check in checks if not check["ok"]]
+            unhealthy = [r for r in results if not r["ok"]]
             app_log.warning(f"Unhealthy services: {unhealthy}")
-        return overall, checks
+        return overall, results
 
     async def get(self):
         overall, checks = await self.check_all()
@@ -218,3 +175,60 @@ async def head(self):
         overall, checks = await self.check_all()
         if not overall:
             self.set_status(503)
+
+
+class KubernetesHealthHandler(HealthHandler):
+    """Serve health status on Kubernetes"""
+
+    @at_most_every
+    async def _get_pods(self):
+        """Get information about build and user pods"""
+        namespace = self.settings["build_namespace"]
+        k8s = self.settings["kubernetes_client"]
+        pool = self.settings["executor"]
+
+        app_log.info(f"Getting pod statistics for {namespace}")
+
+        label_selectors = [
+            "app=jupyterhub,component=singleuser-server",
+            "component=binderhub-build",
+        ]
+        requests = [
+            asyncio.wrap_future(
+                pool.submit(
+                    k8s.list_namespaced_pod,
+                    namespace,
+                    label_selector=label_selector,
+                    _preload_content=False,
+                    _request_timeout=KUBE_REQUEST_TIMEOUT,
+                )
+            )
+            for label_selector in label_selectors
+        ]
+        responses = await asyncio.gather(*requests)
+        return [json.loads(resp.read())["items"] for resp in responses]
+
+    def get_checks(self, checks):
+        super().get_checks(checks)
+        checks["Pod quota"] = self._check_pod_quotas()
+
+    async def _check_pod_quotas(self):
+        """Compare number of active pods to available quota"""
+        user_pods, build_pods = await self._get_pods()
+
+        n_user_pods = len(user_pods)
+        n_build_pods = len(build_pods)
+
+        quota = self.settings["pod_quota"]
+        total_pods = n_user_pods + n_build_pods
+        usage = {
+            "total_pods": total_pods,
+            "build_pods": n_build_pods,
+            "user_pods": n_user_pods,
+            "quota": quota,
+            "ok": total_pods <= quota if quota is not None else True,
+            # The pod quota is treated as a soft quota
+            # Being above quota doesn't mean the service is unhealthy
+            "_ignore_failure": True,
+        }
+        return usage