configurable max concurrent health check batch size

Signed-off-by: Keval Mahajan <mahajankeval23@gmail.com>

Author: kevalmahajan

.env.example

@@ -733,26 +733,41 @@ PROMPT_CACHE_SIZE=100
 MAX_PROMPT_SIZE=102400
 PROMPT_RENDER_TIMEOUT=10
 
+#####################################
+# MCP server Health Check COnfigurations
+#####################################
+
 # Health Check Configuration
 HEALTH_CHECK_INTERVAL=60
 # Health check timeout in seconds (default: 10, matches config.py)
 HEALTH_CHECK_TIMEOUT=10
 UNHEALTHY_THRESHOLD=5
 # Gateway URL validation timeout in seconds (default: 5, matches config.py)
 GATEWAY_VALIDATION_TIMEOUT=5
+# Maximum number of concurrent health checks per worker. Prevents resource exhaustion during health check operations (default: 20, matches config.py)
+MAX_CONCURRENT_HEALTH_CHECKS=20
 
 # File lock name for gateway service leader election
 # Used to coordinate multiple gateway instances when running in cluster mode
 # Default: "gateway_service_leader.lock"
 FILELOCK_NAME=gateway_service_leader.lock
 
+
+#####################################
+# Default Root Paths
+#####################################
+
 # Default root paths (JSON array)
 # List of default root paths for resource resolution
 # Example: ["/api/v1", "/mcp"]
 # Default: []
 DEFAULT_ROOTS=[]
 
+
+#####################################
 # OpenTelemetry Observability Configuration
+#####################################
+
 # Enable distributed tracing and metrics collection
 # Options: true (default), false
 OTEL_ENABLE_OBSERVABILITY=true

README.md

@@ -1996,6 +1996,7 @@ ENABLE_METRICS=false
 | `UNHEALTHY_THRESHOLD`   | Fail-count before peer deactivation,      | `3`     | int > 0 |
 |                         | Set to -1 if deactivation is not needed.  |         |         |
 | `GATEWAY_VALIDATION_TIMEOUT` | Gateway URL validation timeout (secs) | `5`     | int > 0 |
+| `MAX_CONCURRENT_HEALTH_CHECKS` | Max Concurrent health checks        | `20`     | int > 0 |
 | `FILELOCK_NAME`         | File lock for leader election             | `gateway_service_leader.lock` | string |
 | `DEFAULT_ROOTS`         | Default root paths for resources          | `[]`    | JSON array |
 

mcpgateway/config.py

@@ -925,7 +925,8 @@ def parse_issuers(cls, v: Any) -> list[str]:
     health_check_interval: int = 60  # seconds
     health_check_timeout: int = 10  # seconds
     unhealthy_threshold: int = 5  # after this many failures, mark as Offline
-
+    max_concurrent_health_checks: int = 20  # maximum concurrent health checks per worker
+ 
     # Validation Gateway URL
     gateway_validation_timeout: int = 5  # seconds
     gateway_max_redirects: int = 5

mcpgateway/services/gateway_service.py

@@ -2141,22 +2141,32 @@ async def check_health_of_gateways(self, db: Session, gateways: List[DbGateway],
             True
         """
         start_time = time.monotonic()
+        concurrency_limit = min(settings.max_concurrent_health_checks, max(10, os.cpu_count() * 5))  # adaptive concurrency
+        semaphore = asyncio.Semaphore(concurrency_limit)
+
+        async def limited_check(gateway: DbGateway):
+            async with semaphore:
+                await self._check_single_gateway_health(db, gateway, user_email)
 
         # Create trace span for health check batch
         with create_span("gateway.health_check_batch", {"gateway.count": len(gateways), "check.type": "health"}) as batch_span:
-            # Create tasks for concurrent health checks
-            tasks = []
-            for gateway in gateways:
-                if gateway.auth_type == "one_time_auth":
-                    continue  # Skip health check for one_time auth gateways
-                tasks.append(self._check_single_gateway_health(db, gateway, user_email))
-
-            # Execute all health checks concurrently
-            if tasks:
+            # Chunk processing to avoid overload
+            if not gateways:
+                return True
+            chunk_size = concurrency_limit
+            for i in range(0, len(gateways), chunk_size):
+                # batch will be a sublist of gateways from index i to i + chunk_size
+                batch = gateways[i:i + chunk_size]
+
+                # Each task is a health check for a gateway in the batch, excluding those with auth_type == "one_time_auth"
+                tasks = [limited_check(gw) for gw in batch if gw.auth_type != "one_time_auth"]
+
+                # Execute all health checks concurrently
                 await asyncio.gather(*tasks, return_exceptions=True)
+                await asyncio.sleep(0.05)  # small pause prevents network saturation
 
             elapsed = time.monotonic() - start_time
-
+            
             if batch_span:
                 batch_span.set_attribute("check.duration_ms", int(elapsed * 1000))
                 batch_span.set_attribute("check.completed", True)