feature: implement task manager health check

icij-worker/icij_worker/constants.py

@@ -1,6 +1,11 @@
 # AMQP
 from icij_common.pydantic_utils import to_lower_camel
 
+
+AMQP_HEALTH_X = "exchangeMonitoring"
+AMQP_HEALTH_ROUTING_KEY = "routingKeyMonitoring"
+AMQP_HEALTH_QUEUE = "MONITORING"
+
 AMQP_MANAGER_EVENTS_X = "exchangeManagerEvents"
 AMQP_MANAGER_EVENTS_QUEUE = "MANAGER_EVENT"
 AMQP_MANAGER_EVENTS_ROUTING_KEY = "routingKeyManagerEvents"
@@ -22,6 +27,7 @@
 AMQP_WORKER_EVENTS_ROUTING_KEY = "routingKeyWorkerEvents"
 
 AMQP_TASK_QUEUE_PRIORITY = 1000
+AMQP_HEALTH_POLICY_PRIORITY = 10000
 
 _CREATED_AT = "created_at"
 _TASK_ID = "task_id"

icij-worker/icij_worker/task_manager/__init__.py

@@ -3,7 +3,7 @@
 from abc import ABC, abstractmethod
 from asyncio import Task as AsyncIOTask
 from functools import cached_property
-from typing import ClassVar, Optional, final
+from typing import ClassVar, Dict, Optional, final
 
 from pydantic import Field
 
@@ -177,3 +177,6 @@ async def _enqueue(self, task: Task) -> Task: ...
 
     @abstractmethod
     async def _requeue(self, task: Task): ...
+
+    @abstractmethod
+    async def get_health(self) -> Dict[str, bool]: ...

icij-worker/icij_worker/task_manager/amqp.py

@@ -34,6 +34,7 @@
     AMQPMixin,
     RobustConnection,
     amqp_task_group_policy,
+    health_policy,
 )
 
 logger = logging.getLogger(__name__)
@@ -73,6 +74,7 @@ def __init__(
         self._storage = task_store
         self._task_x: Optional[AbstractExchange] = None
         self._worker_evt_x: Optional[AbstractExchange] = None
+        self._health_x: Optional[AbstractExchange] = None
         self._manager_messages_it: Optional[AbstractQueueIterator] = None
 
         self._task_groups: Dict[str, Optional[str]] = dict()
@@ -240,6 +242,9 @@ async def _connection_workflow(self):
         await self._create_routing(
             self.manager_evt_routing(), declare_exchanges=True, declare_queues=True
         )
+        await self._create_routing(
+            self.health_routing(), declare_exchanges=True, declare_queues=True
+        )
         self._task_x = await self._channel.get_exchange(
             self.default_task_routing().exchange.name, ensure=True
         )
@@ -249,6 +254,11 @@ async def _connection_workflow(self):
         self._worker_evt_x = await self._channel.get_exchange(
             self.worker_evt_routing().exchange.name, ensure=True
         )
+        self._health_x = await self._channel.get_exchange(
+            self.health_routing().exchange.name, ensure=True
+        )
+        healthz_policy = health_policy(self.health_routing())
+        await self._management_client.set_policy(healthz_policy)
         logger.info("connection workflow complete")
 
     @cached_property
@@ -275,3 +285,8 @@ async def _ensure_task_queues(self):
             await self._create_routing(
                 routing, declare_exchanges=True, declare_queues=True
             )
+
+    async def get_health(self) -> Dict[str, bool]:
+        storage_health = await self._storage.get_health()
+        amqp_health = await self._get_amqp_health()
+        return {"storage": storage_health, "amqp": amqp_health}

icij-worker/icij_worker/task_manager/neo4j_.py

@@ -125,6 +125,9 @@ async def shutdown_workers(self):
             async with self._db_session(db.name) as sess:
                 await sess.execute_write(_shutdown_workers_tx)
 
+    async def get_health(self) -> bool:
+        return await super(Neo4jConsumerMixin, self).get_health()
+
 
 async def _enqueue_task_tx(
     tx: neo4j.AsyncTransaction, *, task_id: str, max_queue_size: int

icij-worker/icij_worker/task_storage/__init__.py

@@ -55,3 +55,6 @@ async def save_result(self, result: ResultEvent): ...
 
     @abstractmethod
     async def save_error(self, error: ErrorEvent): ...
+
+    @abstractmethod
+    async def get_health(self) -> bool: ...

icij-worker/icij_worker/task_storage/fs.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import functools
+import logging
 from contextlib import AsyncExitStack
 from copy import deepcopy
 from pathlib import Path
@@ -15,6 +16,8 @@
 from icij_worker.task_storage import TaskStorageConfig
 from icij_worker.task_storage.key_value import DBItem, KeyValueStorage
 
+logger = logging.getLogger(__name__)
+
 
 class FSKeyValueStorageConfig(ICIJModel, TaskStorageConfig):
     db_path: Path
@@ -129,3 +132,12 @@ def _make_group_dbs(self) -> Dict[str, SqliteDict]:
                 self._db_path, name=self._errors_db_name
             ),
         }
+
+    async def get_health(self) -> bool:
+        try:
+            for db in self._dbs.values():
+                len(db)
+        except Exception as e:
+            logger.exception("fs storage health failed: %s", e)
+            return False
+        return True

icij-worker/icij_worker/task_storage/neo4j_.py

@@ -1,5 +1,6 @@
 import itertools
 import json
+import logging
 from contextlib import asynccontextmanager
 from copy import deepcopy
 from datetime import datetime
@@ -9,7 +10,7 @@
 from neo4j import AsyncTransaction
 from neo4j.exceptions import ResultNotSingleError
 
-from icij_common.neo4j.db import db_specific_session
+from icij_common.neo4j.db import db_specific_session, registry_db_session
 from icij_common.neo4j.migrate import retrieve_dbs
 from icij_common.pydantic_utils import jsonable_encoder
 from icij_worker.constants import (
@@ -60,6 +61,8 @@
 from icij_worker.objects import ErrorEvent, ResultEvent, Task, TaskState, TaskUpdate
 from icij_worker.task_storage import TaskStorage
 
+logger = logging.getLogger(__name__)
+
 
 class Neo4jStorage(TaskStorage):
     def __init__(self, driver: neo4j.AsyncDriver):
@@ -174,6 +177,16 @@ async def _refresh_task_meta(self):
                 }
                 self._task_meta.update(task_meta)
 
+    async def get_health(self) -> bool:
+        try:
+            async with registry_db_session(self._driver) as sess:
+                res = await sess.run("RETURN 1 AS health_check")
+                await res.single()
+        except Exception as e:
+            logger.error("neo4j health failed: %s", e)
+            return False
+        return True
+
 
 async def _get_tasks_meta_tx(tx: neo4j.AsyncTransaction) -> List[neo4j.Record]:
     query = f"""MATCH (task:{NEO4J_TASK_NODE})

icij-worker/icij_worker/task_storage/postgres/postgres.py

@@ -124,6 +124,7 @@ async def __aexit__(self, exc_type, exc_value, traceback):
 
 
 class PostgresStorage(TaskStorage):
+
     def __init__(
         self,
         connection_info: PostgresConnectionInfo,
@@ -276,6 +277,18 @@ async def init_database(self, db_name: str):
             migration_throttle_s=self._migration_throttle_s,
         )
 
+    async def get_health(self) -> bool:
+        try:
+            registry_pool = await self._pool_manager.get_pool(self._registry_db_name)
+            async with registry_pool.connection() as conn:
+                health_query = sql.SQL("SELECT 1;")
+                async with conn.cursor() as cur:
+                    await cur.execute(health_query)
+        except Exception as e:
+            logger.error("postgres health check failed: %s", e)
+            return False
+        return True
+
     async def _get_task_db(self, task_id: str) -> str:
         if task_id not in self._task_meta:
             await self._refresh_task_meta()
@@ -522,7 +535,7 @@ async def init_database(
 async def _insert_db_into_registry(registry_con: AsyncConnection, db_name: str):
     async with registry_con.cursor() as cur:
         query = sql.SQL(
-            """INSERT INTO {} ({}, {}) VALUES (%s, %s) 
+            """INSERT INTO {} ({}, {}) VALUES (%s, %s)
 ON CONFLICT DO NOTHING;"""
         ).format(
             sql.Identifier(POSTGRES_TASK_DBS_TABLE),

icij-worker/icij_worker/tests/conftest.py

@@ -543,7 +543,6 @@ async def test_amqp_task_manager(
     rabbit_mq: str,
     test_async_app: AsyncApp,
 ) -> TestableAMQPTaskManager:
-
     task_manager = TestableAMQPTaskManager(
         test_async_app, fs_storage, management_client, broker_url=rabbit_mq
     )

icij-worker/icij_worker/tests/task_manager/test_amqp.py

@@ -242,3 +242,37 @@ async def test_task_manager_cancel(
     assert isinstance(datetime.fromisoformat(created_at), datetime)
     expected_json = {"@type": "CancelEvent", "requeue": requeue, "taskId": "some-id"}
     assert cancel_evt_json == expected_json
+
+
+async def test_get_health(
+    test_amqp_task_manager: TestableAMQPTaskManager,
+):
+    # Given
+    task_manager = test_amqp_task_manager
+    # When
+    async with task_manager:
+        health = await task_manager.get_health()
+    # Then
+    assert health == {"storage": True, "amqp": True}
+
+
+async def test_get_health_should_fail(
+    monkeypatch,
+    management_client: AMQPManagementClient,
+    fs_storage: TestableFSKeyValueStorage,
+    rabbit_mq: str,
+    test_async_app: AsyncApp,
+):
+    # Given
+    task_manager = TestableAMQPTaskManager(
+        test_async_app, fs_storage, management_client, broker_url=rabbit_mq
+    )
+
+    def _failing_len(self):
+        raise OSError("failing...")
+
+    monkeypatch.setattr("icij_worker.task_storage.fs.SqliteDict.__len__", _failing_len)
+    # When
+    health = await task_manager.get_health()
+    # Then
+    assert health == {"storage": False, "amqp": False}

icij-worker/icij_worker/tests/task_manager/test_neo4j.py

@@ -675,3 +675,34 @@ async def test_task_manager_shutdown_workers(
     event = ShutdownEvent.from_neo4j(recs[0])
     # Then
     assert event.created_at < datetime.now(timezone.utc)
+
+
+async def test_get_health(
+    neo4j_task_manager: TestableNeo4JTaskManager,
+):
+    # Given
+    task_manager = neo4j_task_manager
+    # When
+    async with task_manager:
+        health = await task_manager.get_health()
+    # Then
+    assert health
+
+
+async def test_get_health_should_fail(
+    monkeypatch,
+    neo4j_task_manager: TestableNeo4JTaskManager,
+):
+    # Given
+    task_manager = neo4j_task_manager
+
+    def _failing_session():
+        raise ConnectionError("failing...")
+
+    monkeypatch.setattr(
+        "icij_worker.task_storage.neo4j_.registry_db_session", _failing_session
+    )
+    # When
+    health = await task_manager.get_health()
+    # Then
+    assert not health

icij-worker/icij_worker/tests/task_storage/test_fs_storage.py

@@ -232,3 +232,27 @@ async def test_get_error(fs_storage: TestableFSKeyValueStorage):
     db_errors = await store.get_task_errors(task_id="task-0")
     # Then
     assert db_errors == [err]
+
+
+async def test_get_health(fs_storage: TestableFSKeyValueStorage):
+    # When
+    async with fs_storage:
+        health = await fs_storage.get_health()
+    # Then
+    assert health
+
+
+async def test_get_health_should_fail(
+    monkeypatch,
+    fs_storage: TestableFSKeyValueStorage,
+):
+    # Given
+    def _failing_len(self):
+        raise OSError("failing...")
+
+    monkeypatch.setattr("icij_worker.task_storage.fs.SqliteDict.__len__", _failing_len)
+    # When
+    async with fs_storage:
+        health = await fs_storage.get_health()
+    # Then
+    assert not health

icij-worker/icij_worker/tests/task_storage/test_postgres.py

@@ -654,3 +654,29 @@ async def test_migrate_rename_task_group_into_group_id(
         await cur.execute(group_id_col_query)
         args_cols = await cur.fetchone()
         assert args_cols is not None
+
+
+async def test_get_health(test_postgres_storage: PostgresStorage):
+    # When
+    health = await test_postgres_storage.get_health()
+    # Then
+    assert health
+
+
+async def test_get_health_should_fail(
+    monkeypatch, test_postgres_storage: PostgresStorage
+):
+    # Given
+    def _failing_get_connection(self, key):
+        # pylint: disable=unused-argument
+        raise OSError("failing...")
+
+    monkeypatch.setattr(
+        "icij_worker.task_storage.postgres.postgres.PoolManager.get_pool",
+        _failing_get_connection,
+    )
+
+    # When
+    health = await test_postgres_storage.get_health()
+    # Then
+    assert not health