From 8694da0ff12984895fdc8003a7487bf267067b9e Mon Sep 17 00:00:00 2001 From: kaffa Date: Sun, 1 Feb 2026 14:08:31 +0000 Subject: [PATCH] feat: Complete remaining improvements MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit MEDIUM PRIORITY: 4. restore operations - use haproxy_cmd_checked() - restore_servers_from_config() now validates commands - haproxy_restore_state() logs warnings for individual failures 5. container health check - add exception logging - logger.warning() for unexpected exceptions 6. haproxy_add_domain - optimize map file reading - Read map contents once, reuse for domain check and pool lookup - Eliminates redundant file reads LOW PRIORITY: 7. Remove unused IP_PATTERN constant - Was unused since validate_ip() uses ipaddress module 8. haproxy_add_server - auto-select slot - slot=0 or slot=-1 finds first available slot - Returns error if all slots in use 9. haproxy_wait_drain - new tool - Wait for connections to drain before maintenance - Polls HAProxy stats until scur=0 or timeout MCP Tools: 21 → 22 Co-Authored-By: Claude Opus 4.5 --- mcp/server.py | 143 +++++++++++++++++++++++++++++++++++++++++--------- 1 file changed, 118 insertions(+), 25 deletions(-) diff --git a/mcp/server.py b/mcp/server.py index 6f40b5b..c50bc52 100644 --- a/mcp/server.py +++ b/mcp/server.py @@ -69,7 +69,6 @@ DOMAIN_PATTERN = re.compile( r'^[a-zA-Z0-9]([a-zA-Z0-9\-]{0,61}[a-zA-Z0-9])?' r'(\.[a-zA-Z0-9]([a-zA-Z0-9\-]{0,61}[a-zA-Z0-9])?)*$' ) -IP_PATTERN = re.compile(r'^(\d{1,3}\.){3}\d{1,3}$') # Backend/server names: alphanumeric, underscore, hyphen only BACKEND_NAME_PATTERN = re.compile(r'^[a-zA-Z0-9_-]+$') # Pattern for converting domain to backend name @@ -673,14 +672,14 @@ def restore_servers_from_config() -> int: logger.warning("Invalid port for %s slot %d, skipping", domain, slot) continue - try: - for suffix, port in get_server_suffixes(http_port): - server = f"{server_prefix}{suffix}_{slot}" - haproxy_cmd(f"set server {backend}/{server} addr {ip} port {port}") - haproxy_cmd(f"set server {backend}/{server} state ready") - restored += 1 - except HaproxyError as e: - logger.warning("Failed to restore %s slot %d: %s", domain, slot, e) + for suffix, port in get_server_suffixes(http_port): + server = f"{server_prefix}{suffix}_{slot}" + try: + haproxy_cmd_checked(f"set server {backend}/{server} addr {ip} port {port}") + haproxy_cmd_checked(f"set server {backend}/{server} state ready") + restored += 1 + except HaproxyError as e: + logger.warning("Failed to restore %s/%s: %s", backend, server, e) return restored @@ -777,21 +776,33 @@ def haproxy_add_domain(domain: str, ip: str = "", http_port: int = 80) -> str: if not (1 <= http_port <= 65535): return "Error: Port must be between 1 and 65535" - # Check if domain already exists - existing_backend = get_domain_backend(domain) - if existing_backend: - return f"Error: Domain {domain} already exists (mapped to {existing_backend})" + # Read map contents once for both existence check and pool lookup + entries = get_map_contents() - # Find available pool - try: - pool = find_available_pool() - except NoAvailablePoolError as e: - return f"Error: {e}" + # Check if domain already exists (using cached entries) + for domain_entry, backend in entries: + if domain_entry == domain: + return f"Error: Domain {domain} already exists (mapped to {backend})" + + # Find available pool (using cached entries) + used_pools: Set[str] = set() + for _, backend in entries: + if backend.startswith("pool_"): + used_pools.add(backend) + + pool = None + for i in range(1, POOL_COUNT + 1): + pool_name = f"pool_{i}" + if pool_name not in used_pools: + pool = pool_name + break + if not pool: + return f"Error: All {POOL_COUNT} pool backends are in use" try: # Save to disk first (atomic write for persistence) # If HAProxy update fails after this, state will be correct on restart - entries = get_map_contents() + # Note: We already have 'entries' from the map contents read above entries.append((domain, pool)) entries.append((f".{domain}", pool)) try: @@ -948,7 +959,7 @@ def haproxy_add_server(domain: str, slot: int, ip: str, http_port: int = 80) -> Args: domain: The domain name to add the server to - slot: Server slot number (1-10). Use different slots for multiple servers + slot: Server slot number (1-10), or 0/-1 for auto-select first available slot ip: IP address of the server (required) http_port: HTTP port (default: 80) @@ -962,6 +973,9 @@ def haproxy_add_server(domain: str, slot: int, ip: str, http_port: int = 80) -> # Add a third server haproxy_add_server("api.example.com", slot=3, ip="10.0.0.3", http_port=8080) + + # Auto-select next available slot + haproxy_add_server("api.example.com", slot=0, ip="10.0.0.4", http_port=8080) """ if not validate_domain(domain): return "Error: Invalid domain format" @@ -969,14 +983,35 @@ def haproxy_add_server(domain: str, slot: int, ip: str, http_port: int = 80) -> return "Error: IP address is required" if not validate_ip(ip): return "Error: Invalid IP address format" - if not (1 <= slot <= MAX_SLOTS): - return f"Error: Slot must be between 1 and {MAX_SLOTS}" if not (1 <= http_port <= 65535): return "Error: Port must be between 1 and 65535" try: backend, server_prefix = get_backend_and_prefix(domain) + # Auto-select slot if slot <= 0 + if slot <= 0: + state = haproxy_cmd("show servers state") + used_slots: Set[int] = set() + for line in state.split("\n"): + parts = line.split() + if len(parts) >= STATE_MIN_COLUMNS and parts[StateField.BE_NAME] == backend: + if parts[StateField.SRV_ADDR] != "0.0.0.0": + # Extract slot number from server name (e.g., pool_1_3 -> 3) + server_name = parts[StateField.SRV_NAME] + try: + used_slots.add(int(server_name.rsplit("_", 1)[1])) + except (ValueError, IndexError): + pass + for s in range(1, MAX_SLOTS + 1): + if s not in used_slots: + slot = s + break + else: + return f"Error: No available slots (all {MAX_SLOTS} slots in use)" + elif not (1 <= slot <= MAX_SLOTS): + return f"Error: Slot must be between 1 and {MAX_SLOTS}, or 0/-1 for auto-select" + # Save to persistent config FIRST (disk-first pattern) add_server_to_config(domain, slot, ip, http_port) @@ -1242,6 +1277,60 @@ def haproxy_set_domain_state(domain: str, state: str) -> str: return result +@mcp.tool() +def haproxy_wait_drain(domain: str, timeout: int = 30) -> str: + """Wait for all active connections to drain from a domain's servers. + + Useful after setting servers to 'drain' state before maintenance. + Polls every second until all servers have 0 current connections or timeout. + + Args: + domain: The domain name to wait for + timeout: Maximum seconds to wait (default: 30, max: 300) + + Returns: + Success message or timeout error + + Example: + haproxy_set_domain_state("api.example.com", "drain") + haproxy_wait_drain("api.example.com", timeout=60) + # Now safe to perform maintenance + """ + if not validate_domain(domain): + return "Error: Invalid domain format" + if not (1 <= timeout <= 300): + return "Error: Timeout must be between 1 and 300 seconds" + + try: + backend, _ = get_backend_and_prefix(domain) + except ValueError as e: + return f"Error: {e}" + + start_time = time.time() + while time.time() - start_time < timeout: + try: + stats = haproxy_cmd("show stat") + total_connections = 0 + for line in stats.split("\n"): + parts = line.split(",") + if len(parts) > StatField.SCUR and parts[0] == backend and parts[1] not in ["FRONTEND", "BACKEND", ""]: + try: + scur = int(parts[StatField.SCUR]) if parts[StatField.SCUR] else 0 + total_connections += scur + except ValueError: + pass + + if total_connections == 0: + elapsed = int(time.time() - start_time) + return f"All connections drained for {domain} ({elapsed}s)" + + time.sleep(1) + except HaproxyError as e: + return f"Error checking connections: {e}" + + return f"Timeout: Connections still active after {timeout}s" + + @mcp.tool() def haproxy_stats() -> str: """Get HAProxy status and statistics. @@ -1327,6 +1416,7 @@ def haproxy_health() -> str: result["components"]["container"] = {"status": "timeout"} result["status"] = "unhealthy" except Exception as e: + logger.warning("Container health check failed: %s", e) result["components"]["container"] = {"status": "error", "error": str(e)} # Check configuration files @@ -1730,9 +1820,12 @@ def haproxy_restore_state() -> str: skipped += 1 continue - haproxy_cmd(f"set server {backend}/{server} addr {addr} port {port}") - haproxy_cmd(f"set server {backend}/{server} state ready") - restored += 1 + try: + haproxy_cmd_checked(f"set server {backend}/{server} addr {addr} port {port}") + haproxy_cmd_checked(f"set server {backend}/{server} state ready") + restored += 1 + except HaproxyError as e: + logger.warning("Failed to restore %s/%s: %s", backend, server, e) result = f"Server state restored ({restored} servers)" if skipped: