feat: run user-supplied health probes on every health check route (#294) (#297)

Author: lesnik512

CONTEXT.md

@@ -46,6 +46,13 @@ signals: not configured is silent and lands in `skipped_instruments`, so it show
 raises `InstrumentDependencyMissingWarning` plus a `logger.warning` and lands in neither list. The
 quiet skip is the one you have to go looking for; the loud one finds you.
 
+**Health probe**:
+One check of something the service depends on, run each time the health check route is requested. A
+single failing probe makes the route report the service unhealthy. The health check is the route;
+health probes are what it runs.
+_Avoid_: health check, for this sense; ping, which is FastStream's broker method and only one
+possible probe.
+
 **Free**:
 Without a web framework: `FreeBootstrapper`, `FreeConfig`, the `free-all` extra. It has no
 application, so nothing to attach teardown to and no double-bootstrap guard.

docs/introduction/configuration.md

@@ -391,6 +391,35 @@ Additional params:
 
 - `health_checks_path`
 - `health_checks_include_in_schema`
+- `health_checks_probes`: named async callables, each returning `bool`, that the route runs on every request
+
+The health check route on Litestar, FastAPI, FastMCP and FastStream runs every probe concurrently. On FastStream the probes run alongside the broker pings. The route returns 200 with `health_status: true` when everything passes. If any probe returns `False` or raises, the route returns 500 with the same JSON body and `health_status: false`. Each failed probe logs one warning with its name; the response body does not say which probe failed.
+
+```python
+import functools
+
+import sqlalchemy
+
+from lite_bootstrap import LitestarConfig
+
+
+async def database_probe() -> bool:
+    async with engine.connect() as connection:
+        await connection.execute(sqlalchemy.text("SELECT 1"))
+    return True
+
+
+config = LitestarConfig(
+    health_checks_probes={
+        "database": database_probe,
+        "outbox-broker": functools.partial(outbox_broker.ping, timeout=5),
+    },
+)
+```
+
+lite-bootstrap puts no timeout on probes, so a probe that hangs makes the route hang too. Give each probe its own timeout, as `functools.partial(broker.ping, timeout=5)` does.
+
+Be careful with probes on dependencies outside the service, such as a database, when the route also serves a Kubernetes liveness probe. During the database outage every pod fails its liveness check and gets restarted, and the restarts add load once the database comes back. Readiness probes are the usual place for dependency checks.
 
 ## Skipped instruments
 

lite_bootstrap/bootstrappers/fastapi_bootstrapper.py

@@ -178,8 +178,11 @@ def build_fastapi_health_check_router(self) -> "fastapi.APIRouter":
         )
 
         @fastapi_router.get(self.bootstrap_config.health_checks_path)
-        async def health_check_handler() -> HealthCheckTypedDict:
-            return self.render_health_check_data()
+        async def health_check_handler(response: fastapi.Response) -> HealthCheckTypedDict:
+            health_status = await self.define_health_status()
+            if not health_status:
+                response.status_code = fastapi.status.HTTP_500_INTERNAL_SERVER_ERROR
+            return self.render_health_check_data(health_status=health_status)
 
         return fastapi_router
 

lite_bootstrap/bootstrappers/fastmcp_bootstrapper.py

@@ -22,6 +22,7 @@
     from fastmcp.server.http import StarletteWithLifespan
     from fastmcp.server.middleware import Middleware, MiddlewareContext
     from fastmcp.server.providers import Provider
+    from starlette import status
     from starlette.requests import Request
     from starlette.responses import JSONResponse, Response
 
@@ -210,7 +211,11 @@ def bootstrap(self) -> None:
             include_in_schema=config.health_checks_include_in_schema,
         )
         async def health_check_handler(_: "Request") -> "JSONResponse":
-            return JSONResponse(dict(self.render_health_check_data()))
+            health_status = await self.define_health_status()
+            return JSONResponse(
+                dict(self.render_health_check_data(health_status=health_status)),
+                status_code=status.HTTP_200_OK if health_status else status.HTTP_500_INTERNAL_SERVER_ERROR,
+            )
 
 
 @dataclasses.dataclass(kw_only=True)

lite_bootstrap/bootstrappers/faststream_bootstrapper.py

@@ -88,22 +88,19 @@ def bootstrap(self) -> None:
 
         @handle_get
         async def check_health(_: object) -> "AsgiResponse":
-            return (
-                AsgiResponse(
-                    json.dumps(self.render_health_check_data()).encode(),
-                    200,
-                    headers={"content-type": "application/json"},
-                )
-                if await self._define_health_status()
-                else AsgiResponse(b"Service is unhealthy", 500, headers={"content-type": "text/plain"})
+            health_status = await self.define_health_status()
+            return AsgiResponse(
+                json.dumps(self.render_health_check_data(health_status=health_status)).encode(),
+                200 if health_status else 500,
+                headers={"content-type": "application/json"},
             )
 
         if config.opentelemetry_generate_health_check_spans and import_checker.is_opentelemetry_installed:
             check_health = tracer.start_as_current_span(f"GET {config.health_checks_path}")(check_health)
 
         config.application.mount(config.health_checks_path, check_health)
 
-    async def _define_health_status(self) -> bool:
+    async def _define_framework_health_status(self) -> bool:
         config = self.bootstrap_config
         if not config.application or not config.application.brokers:
             return False

lite_bootstrap/bootstrappers/litestar_bootstrapper.py

@@ -28,6 +28,7 @@
 
 if import_checker.is_litestar_installed:
     import litestar
+    from litestar import status_codes
     from litestar.config.app import AppConfig
     from litestar.config.cors import CORSConfig
     from litestar.logging.config import StructLoggingConfig
@@ -149,8 +150,12 @@ def bootstrap(self) -> None:
 class LitestarHealthChecksInstrument(HealthChecksInstrument[LitestarConfig]):
     def build_litestar_health_check_router(self) -> "litestar.Router":
         @litestar.get(media_type=litestar.MediaType.JSON)
-        async def health_check_handler() -> HealthCheckTypedDict:
-            return self.render_health_check_data()
+        async def health_check_handler() -> litestar.Response[HealthCheckTypedDict]:
+            health_status = await self.define_health_status()
+            return litestar.Response(
+                self.render_health_check_data(health_status=health_status),
+                status_code=status_codes.HTTP_200_OK if health_status else status_codes.HTTP_500_INTERNAL_SERVER_ERROR,
+            )
 
         return litestar.Router(
             path=self.bootstrap_config.health_checks_path,

lite_bootstrap/instruments/healthchecks_instrument.py

@@ -1,11 +1,19 @@
+import asyncio
 import dataclasses
+import logging
 import typing
+from collections.abc import Awaitable, Callable, Mapping
 
 import typing_extensions
 
 from lite_bootstrap.instruments.base import BaseConfig, BaseInstrument
 
 
+logger = logging.getLogger(__name__)
+
+HealthProbe = Callable[[], Awaitable[bool]]
+
+
 class HealthCheckTypedDict(typing_extensions.TypedDict, total=False):
     service_version: str | None
     service_name: str | None
@@ -17,6 +25,7 @@ class HealthChecksConfig(BaseConfig):
     health_checks_enabled: bool = True
     health_checks_path: str = "/health/"
     health_checks_include_in_schema: bool = False
+    health_checks_probes: Mapping[str, HealthProbe] = dataclasses.field(default_factory=dict)
 
     @property
     def health_check_data(self) -> HealthCheckTypedDict:
@@ -38,5 +47,28 @@ class HealthChecksInstrument(BaseInstrument[HealthChecksConfigT]):
     def is_configured(cls, bootstrap_config: "HealthChecksConfig") -> bool:
         return bootstrap_config.health_checks_enabled
 
-    def render_health_check_data(self) -> HealthCheckTypedDict:
-        return self.bootstrap_config.health_check_data
+    def render_health_check_data(self, *, health_status: bool = True) -> HealthCheckTypedDict:
+        health_check_data = self.bootstrap_config.health_check_data
+        health_check_data["health_status"] = health_status
+        return health_check_data
+
+    async def define_health_status(self) -> bool:
+        results = await asyncio.gather(
+            self._define_framework_health_status(),
+            *(self._run_probe(name, probe) for name, probe in self.bootstrap_config.health_checks_probes.items()),
+        )
+        return all(results)
+
+    async def _define_framework_health_status(self) -> bool:
+        return True
+
+    @staticmethod
+    async def _run_probe(name: str, probe: HealthProbe) -> bool:
+        try:
+            passed = await probe()
+        except Exception:
+            logger.warning("Health probe %s raised", name, exc_info=True)
+            return False
+        if not passed:
+            logger.warning("Health probe %s failed", name)
+        return bool(passed)

tests/conftest.py

@@ -4,6 +4,7 @@
 import warnings
 from collections.abc import Generator
 from importlib import reload
+from unittest.mock import MagicMock, patch
 
 import opentelemetry.metrics._internal
 import opentelemetry.trace
@@ -16,6 +17,7 @@
 
 from lite_bootstrap import import_checker
 from lite_bootstrap.exceptions import InstrumentSkippedWarning
+from lite_bootstrap.instruments import healthchecks_instrument
 
 
 def pytest_configure() -> None:
@@ -36,6 +38,40 @@ def _uninstrument(self, **kwargs: typing.Mapping[str, typing.Any]) -> None:
         pass
 
 
+class ProbeError(RuntimeError):
+    pass
+
+
+async def passing_probe() -> bool:
+    return True
+
+
+async def failing_probe() -> bool:
+    return False
+
+
+async def raising_probe() -> bool:
+    message = "dependency unreachable"
+    raise ProbeError(message)
+
+
+HEALTH_PROBE_CASES: typing.Final = [
+    pytest.param({"db": passing_probe, "cache": passing_probe}, 200, True, [], id="all-pass"),
+    pytest.param({"db": passing_probe, "cache": failing_probe}, 500, False, ["cache"], id="one-returns-false"),
+    pytest.param({"db": passing_probe, "cache": raising_probe}, 500, False, ["cache"], id="one-raises"),
+]
+
+
+@contextlib.contextmanager
+def capture_probe_warnings() -> Generator[MagicMock]:
+    with patch.object(healthchecks_instrument, "logger") as mock_logger:
+        yield mock_logger.warning
+
+
+def failed_probe_names(warning_mock: MagicMock) -> list[str]:
+    return [call.args[1] for call in warning_mock.call_args_list]
+
+
 class SentryTestTransport(sentry_sdk.Transport):
     def __init__(self, *args: typing.Any, **kwargs: typing.Any) -> None:  # noqa: ANN401
         super().__init__(*args, **kwargs)

tests/instruments/test_healthchecks_instrument.py

@@ -1,3 +1,6 @@
+import asyncio
+import time
+
 from lite_bootstrap.instruments.healthchecks_instrument import HealthChecksConfig, HealthChecksInstrument
 
 
@@ -39,7 +42,33 @@ def test_healthchecks_config_defaults() -> None:
     assert config.health_checks_enabled is True
     assert config.health_checks_path == "/health/"
     assert config.health_checks_include_in_schema is False
+    assert config.health_checks_probes == {}
 
 
 def test_healthchecks_dependencies_installed() -> None:
     assert HealthChecksInstrument.dependencies_installed() is True
+
+
+async def test_healthchecks_probes_run_concurrently() -> None:
+    interval = 0.2
+
+    async def slow_probe() -> bool:
+        await asyncio.sleep(interval)
+        return True
+
+    instrument = HealthChecksInstrument(
+        bootstrap_config=HealthChecksConfig(health_checks_probes={"first": slow_probe, "second": slow_probe}),
+    )
+    started = time.perf_counter()
+    assert await instrument.define_health_status()
+    assert time.perf_counter() - started < interval * 1.5
+
+
+async def test_healthchecks_status_is_healthy_without_probes() -> None:
+    instrument = HealthChecksInstrument(bootstrap_config=HealthChecksConfig())
+    assert await instrument.define_health_status()
+
+
+def test_healthchecks_render_data_unhealthy() -> None:
+    instrument = HealthChecksInstrument(bootstrap_config=HealthChecksConfig())
+    assert instrument.render_health_check_data(health_status=False)["health_status"] is False

tests/test_fastapi_bootstrap.py

@@ -23,8 +23,17 @@
 from lite_bootstrap import FastAPIBootstrapper, FastAPIConfig, import_checker
 from lite_bootstrap.bootstrappers import fastapi_bootstrapper
 from lite_bootstrap.exceptions import ConfigurationError, InstrumentDependencyMissingWarning
+from lite_bootstrap.instruments.healthchecks_instrument import HealthProbe
 from lite_bootstrap.types import UNSET
-from tests.conftest import CustomInstrumentor, SentryTestTransport, emulate_package_missing, warning_source_files
+from tests.conftest import (
+    HEALTH_PROBE_CASES,
+    CustomInstrumentor,
+    SentryTestTransport,
+    capture_probe_warnings,
+    emulate_package_missing,
+    failed_probe_names,
+    warning_source_files,
+)
 
 
 logger = structlog.getLogger(__name__)
@@ -501,3 +510,26 @@ def test_fastapi_otel_honors_excluded_urls_from_environment(
     span_names = _server_span_names_for(fastapi_config, ["/from-env", "/custom-healthy"])
 
     assert span_names == ["GET /custom-healthy"]
+
+
+@pytest.mark.parametrize(("probes", "expected_status", "expected_health_status", "expected_failed"), HEALTH_PROBE_CASES)
+def test_fastapi_health_check_runs_probes(
+    fastapi_config: FastAPIConfig,
+    probes: dict[str, HealthProbe],
+    expected_status: int,
+    expected_health_status: bool,
+    expected_failed: list[str],
+) -> None:
+    config = dataclasses.replace(fastapi_config, health_checks_probes=probes)
+    bootstrapper = FastAPIBootstrapper(bootstrap_config=config)
+    application = bootstrapper.bootstrap()
+    with capture_probe_warnings() as probe_warnings, TestClient(application) as test_client:
+        response = test_client.get(config.health_checks_path)
+
+    assert response.status_code == expected_status
+    assert response.json() == {
+        "health_status": expected_health_status,
+        "service_name": "microservice",
+        "service_version": "2.0.0",
+    }
+    assert failed_probe_names(probe_warnings) == expected_failed