Merge pull request #88 from IBM/federated_gateway_health_check

fixed gateway healthcheck for multiple and single workers

Author: crivetimihai

.env.example

@@ -214,6 +214,12 @@ HEALTH_CHECK_TIMEOUT=10
 # Number of failed checks before marking peer unhealthy
 UNHEALTHY_THRESHOLD=3
 
+#####################################
+# Lock file Settings
+#####################################
+FILELOCK_PATH=/tmp/gateway_healthcheck_init.lock
+
+
 #####################################
 # Development Settings
 #####################################

.gitignore

@@ -197,4 +197,7 @@ downloads/
 # db_path
 db_path/
 
+# filelock path
+tmp/
+
 .continue

README.md

@@ -546,11 +546,12 @@ You can get started by copying the provided `.env.examples` to `.env` and making
 
 ### Health Checks
 
-| Setting                 | Description                   | Default | Options |
-| ----------------------- | ----------------------------- | ------- | ------- |
-| `HEALTH_CHECK_INTERVAL` | Health poll interval (secs)   | `60`    | int > 0 |
-| `HEALTH_CHECK_TIMEOUT`  | Health request timeout (secs) | `10`    | int > 0 |
-| `UNHEALTHY_THRESHOLD`   | Fail-count before unhealthy   | `3`     | int > 0 |
+| Setting                 | Description                               | Default | Options |
+| ----------------------- | ----------------------------------------- | ------- | ------- |
+| `HEALTH_CHECK_INTERVAL` | Health poll interval (secs)               | `60`    | int > 0 |
+| `HEALTH_CHECK_TIMEOUT`  | Health request timeout (secs)             | `10`    | int > 0 |
+| `UNHEALTHY_THRESHOLD`   | Fail-count before peer deactivation,      | `3`     | int > 0 |
+|                         | Set to -1 if deactivation is not needed.  |         |         |
 
 ### Database
 

mcpgateway/config.py

@@ -171,7 +171,9 @@ def _parse_federation_peers(cls, v):
     # Health Checks
     health_check_interval: int = 60  # seconds
     health_check_timeout: int = 10  # seconds
-    unhealthy_threshold: int = 3
+    unhealthy_threshold: int = 10
+
+    filelock_path: str = "tmp/gateway_service_leader.lock"
 
     # Default Roots
     default_roots: List[str] = []

mcpgateway/main.py

@@ -185,15 +185,7 @@ async def lifespan(_app: FastAPI) -> AsyncIterator[None]:
         await tool_service.initialize()
         await resource_service.initialize()
         await prompt_service.initialize()
-        try:
-            # Try to create the file exclusively
-            fd = os.open(settings.lock_file_path, os.O_CREAT | os.O_EXCL | os.O_WRONLY, 0o644)
-        except FileExistsError:
-            logger.info("Gateway already initialized by another worker")
-        else:
-            with os.fdopen(fd, "w") as lock_file:
-                lock_file.write("initialized")
-            await gateway_service.initialize()
+        await gateway_service.initialize()
         await root_service.initialize()
         await completion_service.initialize()
         await logging_service.initialize()

mcpgateway/services/gateway_service.py

@@ -16,10 +16,12 @@
 
 import asyncio
 import logging
+import uuid
 from datetime import datetime, timezone
 from typing import Any, AsyncGenerator, Dict, List, Optional, Set
 
 import httpx
+from filelock import FileLock, Timeout
 from mcp import ClientSession
 from mcp.client.sse import sse_client
 from sqlalchemy import select
@@ -33,9 +35,21 @@
 from mcpgateway.services.tool_service import ToolService
 from mcpgateway.utils.services_auth import decode_auth
 
+try:
+    import redis
+    REDIS_AVAILABLE = True
+except ImportError:
+    REDIS_AVAILABLE = False
+    logging.info("Redis is not utilized in this environment.")
+
+# logging.getLogger("httpx").setLevel(logging.WARNING)  # Disables httpx logs for regular health checks
 logger = logging.getLogger(__name__)
 
 
+GW_FAILURE_THRESHOLD = settings.unhealthy_threshold
+GW_HEALTH_CHECK_INTERVAL = settings.health_check_interval
+
+
 class GatewayError(Exception):
     """Base class for gateway-related errors."""
 
@@ -83,17 +97,45 @@ def __init__(self):
         """Initialize the gateway service."""
         self._event_subscribers: List[asyncio.Queue] = []
         self._http_client = httpx.AsyncClient(timeout=settings.federation_timeout, verify=not settings.skip_ssl_verify)
-        self._health_check_interval = 60  # seconds
+        self._health_check_interval = GW_HEALTH_CHECK_INTERVAL 
         self._health_check_task: Optional[asyncio.Task] = None
         self._active_gateways: Set[str] = set()  # Track active gateway URLs
         self._stream_response = None
         self._pending_responses = {}
         self.tool_service = ToolService()
+        self._gateway_failure_counts: dict[str, int] = {}
+
+        # For health checks, we determine the leader instance.
+        self.redis_url = settings.redis_url if settings.cache_type == "redis" else None
+
+        if self.redis_url and REDIS_AVAILABLE:
+            self._redis_client = redis.from_url(self.redis_url)
+            self._instance_id = str(uuid.uuid4())             # Unique ID for this process
+            self._leader_key = "gateway_service_leader"
+            self._leader_ttl = 40                             # seconds
+        elif settings.cache_type != "none":
+            # Fallback: File-based lock
+            self._redis_client = None
+            self._lock_path = settings.filelock_path
+            self._file_lock = FileLock(self._lock_path)
 
     async def initialize(self) -> None:
-        """Initialize the service."""
+        """Initialize the service and start health check if this instance is the leader."""
         logger.info("Initializing gateway service")
-        self._health_check_task = asyncio.create_task(self._run_health_checks())
+
+        if self._redis_client:
+            # Check if Redis is available
+            pong = self._redis_client.ping()
+            if not pong:
+                raise ConnectionError("Redis ping failed.")
+
+            is_leader = self._redis_client.set(self._leader_key, self._instance_id, ex=self._leader_ttl, nx=True)
+            if is_leader:
+                logger.info("Acquired Redis leadership. Starting health check task.")
+                self._health_check_task = asyncio.create_task(self._run_health_checks())
+        else:
+            # Always create the health check task in filelock mode; leader check is handled inside.
+            self._health_check_task = asyncio.create_task(self._run_health_checks())
 
     async def shutdown(self) -> None:
         """Shutdown the service."""
@@ -474,6 +516,30 @@ async def forward_request(self, gateway: DbGateway, method: str, params: Optiona
         except Exception as e:
             raise GatewayConnectionError(f"Failed to forward request to {gateway.name}: {str(e)}")
 
+    async def _handle_gateway_failure(self, gateway: str) -> None:
+        """
+        Tracks and handles gateway failures during health checks.
+        If the failure count exceeds the threshold, the gateway is deactivated.
+
+        Args:
+            gateway (str): The gateway object that failed its health check.
+
+        Returns:
+            None
+        """
+        if GW_FAILURE_THRESHOLD == -1:
+            return  # Gateway failure action disabled
+        count = self._gateway_failure_counts.get(gateway.id, 0) + 1
+        self._gateway_failure_counts[gateway.id] = count
+
+        logger.warning(f"Gateway {gateway.name} failed health check {count} time(s).")
+
+        if count >= GW_FAILURE_THRESHOLD:
+            logger.error(f"Gateway {gateway.name} failed {GW_FAILURE_THRESHOLD} times. Deactivating...")
+            with SessionLocal() as db:
+                await self.toggle_gateway_status(db, gateway.id, False)
+                self._gateway_failure_counts[gateway.id] = 0  # Reset after deactivation
+
     async def check_health_of_gateways(self, gateways: List[DbGateway]) -> bool:
         """Health check for a list of gateways.
 
@@ -506,9 +572,8 @@ async def check_health_of_gateways(self, gateways: List[DbGateway]) -> bool:
                     gateway.last_seen = datetime.utcnow()
 
                 except Exception:
-                    with SessionLocal() as db:
-                        await self.toggle_gateway_status(db=db, gateway_id=gateway.id, activate=False)
-
+                    await self._handle_gateway_failure(gateway)
+                    
         # All gateways passed
         return True
 
@@ -620,19 +685,68 @@ def _get_active_gateways(self) -> list[DbGateway]:
             return db.execute(select(DbGateway).where(DbGateway.is_active)).scalars().all()
 
     async def _run_health_checks(self) -> None:
-        """Run health checks with sync Session in async code."""
+        """Run health checks periodically,
+        Uses Redis or FileLock - for multiple workers.
+        Uses simple health check for single worker mode."""
+
         while True:
             try:
-                # Run sync database code in a thread
-                gateways = await asyncio.to_thread(self._get_active_gateways)
+                if self._redis_client and settings.cache_type == "redis":
+                    # Redis-based leader check
+                    current_leader = self._redis_client.get(self._leader_key)
+                    if current_leader != self._instance_id.encode():
+                        return
+                    self._redis_client.expire(self._leader_key, self._leader_ttl)
 
-                if len(gateways) > 0:
-                    # Async health checks (non-blocking)
-                    await self.check_health_of_gateways(gateways)
-            except Exception as e:
-                logger.error(f"Health check run failed: {str(e)}")
+                    # Run health checks
+                    gateways = await asyncio.to_thread(self._get_active_gateways)
+                    if gateways:
+                        await self.check_health_of_gateways(gateways)
+
+                    await asyncio.sleep(self._health_check_interval)
+
+                elif settings.cache_type == "none":
+                    try:
+                        # For single worker mode, run health checks directly
+                        gateways = await asyncio.to_thread(self._get_active_gateways)
 
-            await asyncio.sleep(self._health_check_interval)
+                        if gateways:
+                            await self.check_health_of_gateways(gateways)
+                    except Exception as e:
+                        logger.error(f"Health check run failed: {str(e)}")
+
+                    await asyncio.sleep(self._health_check_interval)
+                
+                else:
+                    # FileLock-based leader fallback
+                    try:
+                        self._file_lock.acquire(timeout=0)
+                        logger.info("File lock acquired. Running health checks.")
+
+                        while True:
+                            gateways = await asyncio.to_thread(self._get_active_gateways)
+                            if gateways:
+                                await self.check_health_of_gateways(gateways)
+                            await asyncio.sleep(self._health_check_interval)
+
+                    except Timeout:
+                        logger.debug("File lock already held. Retrying later.")
+                        await asyncio.sleep(self._health_check_interval)
+
+                    except Exception as e:
+                        logger.error(f"FileLock health check failed: {str(e)}")
+
+                    finally:
+                        if self._file_lock.is_locked:
+                            try:
+                                self._file_lock.release()
+                                logger.info("Released file lock.")
+                            except Exception as e:
+                                logger.warning(f"Failed to release file lock: {str(e)}")
+
+            except Exception as e:
+                logger.error(f"Unexpected error in health check loop: {str(e)}")
+                await asyncio.sleep(self._health_check_interval)
 
     def _get_auth_headers(self) -> Dict[str, str]:
         """

pyproject.toml

@@ -65,6 +65,7 @@ dependencies = [
     "starlette>=0.46.2",
     "uvicorn>=0.34.3",
     "zeroconf>=0.147.0",
+    "filelock>=3.18.0",
 ]
 
 # ----------------------------------------------------------------