mozilla-services · grahamalama · Aug 5, 2023 · Aug 5, 2023 · Aug 5, 2023 · Aug 5, 2023
diff --git a/src/dockerflow/checks/__init__.py b/src/dockerflow/checks/__init__.py
@@ -1,6 +1,12 @@
 # This Source Code Form is subject to the terms of the Mozilla Public
 # License, v. 2.0. If a copy of the MPL was not distributed with this
 # file, you can obtain one at http://mozilla.org/MPL/2.0/.
+import asyncio
+import functools
+import inspect
+from dataclasses import dataclass
+from typing import Any, Callable, Dict, Iterable, List, Tuple
+
 from .messages import (  # noqa
     CRITICAL,
     DEBUG,
@@ -16,3 +22,147 @@
     Warning,
     level_to_text,
 )
+
+CheckFn = Callable[..., List[CheckMessage]]
+
+
+@dataclass
+class ChecksResults:
+    """
+    Represents the results of running checks.
+
+    This data class holds the results of running a collection of checks. It includes
+    details about each check's outcome, their statuses, and the overall result level.
+
+    :param details: A dictionary containing detailed information about each check's
+        outcome, with check names as keys and dictionaries of details as values.
+    :type details: Dict[str, Dict[str, Any]]
+
+    :param statuses: A dictionary containing the status of each check, with check names
+        as keys and statuses as values (e.g., 'pass', 'fail', 'warning').
+    :type statuses: Dict[str, str]
+
+    :param level: An integer representing the overall result level of the checks
+    :type level: int
+    """
+
+    details: Dict[str, Dict[str, Any]]
+    statuses: Dict[str, str]
+    level: int
+
+
+def iscoroutinefunction_or_partial(obj):
+    """
+    Determine if the provided object is a coroutine function or a partial function
+    that wraps a coroutine function.
+
+    This function checks whether the given object is a coroutine function or a partial
+    function that wraps a coroutine function. This function should be removed when we
+    drop support for Python 3.7, as this is handled directly by `inspect.iscoroutinefunction`
+    in Python 3.8.
+
+    :param obj: The object to be checked for being a coroutine function or partial.
+    :type obj: object
+
+    :return: True if the object is a coroutine function or a partial function wrapping
+             a coroutine function, False otherwise.
+    :rtype: bool
+    """
+    while isinstance(obj, functools.partial):
+        obj = obj.func
+    return inspect.iscoroutinefunction(obj)
+
+
+async def _run_check_async(check):
+    name, check_fn = check
+    if iscoroutinefunction_or_partial(check_fn):
+        errors = await check_fn()
+    else:
+        loop = asyncio.get_event_loop()
+        errors = await loop.run_in_executor(None, check_fn)
+
+    return (name, errors)
+
+
+async def run_checks_async(
+    checks: Iterable[Tuple[str, CheckFn]],
+    silenced_check_ids=None,
+) -> ChecksResults:
+    """
+    Run checks concurrently and return the results.
+
+    Executes a collection of checks concurrently, supporting both synchronous and
+    asynchronous checks. The results include the outcome of each check and can be
+    further processed.
+
+    :param checks: An iterable of tuples where each tuple contains a check name and a
+        check function.
+    :type checks: Iterable[Tuple[str, CheckFn]]
+
+    :param silenced_check_ids: A list of check IDs that should be omitted from the
+        results.
+    :type silenced_check_ids: List[str]
+
+    :return: An instance of ChecksResults containing detailed information about each
+        check's outcome, their statuses, and the overall result level.
+    :rtype: ChecksResults
+    """
+    if not silenced_check_ids:
+        silenced_check_ids = []
+
+    tasks = (_run_check_async(check) for check in checks)
+    results = await asyncio.gather(*tasks)
+    return _build_results_payload(results, silenced_check_ids)
+
+
+def run_checks(
+    checks: Iterable[Tuple[str, CheckFn]],
+    silenced_check_ids=None,
+) -> ChecksResults:
+    """
+    Run checks synchronously and return the results.
+
+    Executes a collection of checks and returns the results. The results include the
+    outcome of each check and can be further processed.
+
+    :param checks: An iterable of tuples where each tuple contains a check name and a
+        check function.
+    :type checks: Iterable[Tuple[str, CheckFn]]
+
+    :param silenced_check_ids: A list of check IDs that should be omitted from the
+        results.
+    :type silenced_check_ids: List[str]
+
+    :return: An instance of ChecksResults containing detailed information about each
+        check's outcome, their statuses, and the overall result level.
+    :rtype: ChecksResults
+    """
+    if not silenced_check_ids:
+        silenced_check_ids = []
+    results = [(name, check()) for name, check in checks]
+    return _build_results_payload(results, silenced_check_ids)
+
+
+def _build_results_payload(
+    checks_results: Iterable[Tuple[str, Iterable[CheckMessage]]],
+    silenced_check_ids,
+):
+    details = {}
+    statuses = {}
+    max_level = 0
+
+    for name, errors in checks_results:
+        errors = [e for e in errors if e.id not in silenced_check_ids]
+        level = max([0] + [e.level for e in errors])
+
+        detail = {
+            "status": level_to_text(level),
+            "level": level,
+            "messages": {e.id: e.msg for e in errors},
+        }
+        statuses[name] = level_to_text(level)
+        max_level = max(max_level, level)
+        if level > 0:
+            details[name] = detail
+
+    return ChecksResults(statuses=statuses, details=details, level=max_level)
diff --git a/src/dockerflow/django/checks.py b/src/dockerflow/django/checks.py
@@ -11,18 +11,6 @@
 from .. import health
 
 
-def level_to_text(level):
-    statuses = {
-        0: "ok",
-        checks.messages.DEBUG: "debug",
-        checks.messages.INFO: "info",
-        checks.messages.WARNING: "warning",
-        checks.messages.ERROR: "error",
-        checks.messages.CRITICAL: "critical",
-    }
-    return statuses.get(level, "unknown")
-
-
 def check_database_connected(app_configs, **kwargs):
     """
     A Django check to see if connecting to the configured default

diff --git a/src/dockerflow/django/views.py b/src/dockerflow/django/views.py
@@ -1,12 +1,14 @@
 # This Source Code Form is subject to the terms of the Mozilla Public
 # License, v. 2.0. If a copy of the MPL was not distributed with this
 # file, you can obtain one at http://mozilla.org/MPL/2.0/.
+import functools
+
 from django.conf import settings
 from django.core import checks
 from django.http import HttpResponse, HttpResponseNotFound, JsonResponse
 from django.utils.module_loading import import_string
 
-from .checks import level_to_text
+from ..checks import level_to_text, run_checks
 from .signals import heartbeat_failed, heartbeat_passed
 
 version_callback = getattr(
@@ -42,42 +44,27 @@ def heartbeat(request):
     Any check that returns a warning or worse (error, critical) will
     return a 500 response.
     """
-    all_checks = checks.registry.registry.get_checks(
+    registered_checks = checks.registry.registry.get_checks(
         include_deployment_checks=not settings.DEBUG
     )
-
-    details = {}
-    statuses = {}
-    level = 0
-
-    for check in all_checks:
-        detail = heartbeat_check_detail(check)
-        statuses[check.__name__] = detail["status"]
-        level = max(level, detail["level"])
-        if detail["level"] > 0:
-            details[check.__name__] = detail
-
-    if level < checks.messages.ERROR:
+    # partially apply app_configs=None to our check functions so we can run them with our check
+    # running function, since `app_configs` is a Django-specific kwarg
+    checks_to_run = (
+        (check.__name__, functools.partial(check, app_configs=None))
+        for check in registered_checks
+    )
+    check_results = run_checks(
+        checks_to_run,
+        silenced_check_ids=settings.SILENCED_SYSTEM_CHECKS,
+    )
+    if check_results.level < checks.messages.ERROR:
         status_code = 200
-        heartbeat_passed.send(sender=heartbeat, level=level)
+        heartbeat_passed.send(sender=heartbeat, level=check_results.level)
     else:
         status_code = 500
-        heartbeat_failed.send(sender=heartbeat, level=level)
-
-    payload = {"status": level_to_text(level)}
+        heartbeat_failed.send(sender=heartbeat, level=check_results.level)
+    payload = {"status": level_to_text(check_results.level)}
     if settings.DEBUG:
-        payload["checks"] = statuses
-        payload["details"] = details
+        payload["checks"] = check_results.statuses
+        payload["details"] = check_results.details
     return JsonResponse(payload, status=status_code)
-
-
-def heartbeat_check_detail(check):
-    errors = check(app_configs=None)
-    errors = list(filter(lambda e: e.id not in settings.SILENCED_SYSTEM_CHECKS, errors))
-    level = max([0] + [e.level for e in errors])
-
-    return {
-        "status": level_to_text(level),
-        "level": level,
-        "messages": {e.id: e.msg for e in errors},
-    }
diff --git a/src/dockerflow/flask/app.py b/src/dockerflow/flask/app.py
@@ -11,6 +11,8 @@
 import flask
 from werkzeug.exceptions import InternalServerError
 
+from dockerflow.checks import run_checks
+
 from .. import version
 from . import checks
 from .signals import heartbeat_failed, heartbeat_passed
@@ -303,16 +305,6 @@ def _lbheartbeat_view(self):
         """
         return "", 200
 
-    def _heartbeat_check_detail(self, check):
-        errors = list(filter(lambda e: e.id not in self.silenced_checks, check()))
-        level = max([0] + [e.level for e in errors])
-
-        return {
-            "status": checks.level_to_text(level),
-            "level": level,
-            "messages": {e.id: e.msg for e in errors},
-        }
-
     def _heartbeat_view(self):
         """
         Runs all the registered checks and returns a JSON response with either
@@ -321,33 +313,28 @@ def _heartbeat_view(self):
         Any check that returns a warning or worse (error, critical) will
         return a 500 response.
         """
-        details = {}
-        statuses = {}
-        level = 0
 
-        for name, check in self.checks.items():
-            detail = self._heartbeat_check_detail(check)
-            statuses[name] = detail["status"]
-            level = max(level, detail["level"])
-            if detail["level"] > 0:
-                details[name] = detail
+        check_results = run_checks(
+            self.checks.items(),
+            silenced_check_ids=self.silenced_checks,
+        )
 
         payload = {
-            "status": checks.level_to_text(level),
-            "checks": statuses,
-            "details": details,
+            "status": checks.level_to_text(check_results.level),
+            "checks": check_results.statuses,
+            "details": check_results.details,
         }
 
         def render(status_code):
             return flask.make_response(flask.jsonify(payload), status_code)
 
-        if level < checks.ERROR:
+        if check_results.level < checks.ERROR:
             status_code = 200
-            heartbeat_passed.send(self, level=level)
+            heartbeat_passed.send(self, level=check_results.level)
             return render(status_code)
         else:
             status_code = 500
-            heartbeat_failed.send(self, level=level)
+            heartbeat_failed.send(self, level=check_results.level)
             raise HeartbeatFailure(response=render(status_code))
 
     def version_callback(self, func):