from __future__ import annotations import asyncio import contextvars import os import socket import ssl from datetime import UTC, datetime from time import perf_counter from uuid import uuid4 import httpx from .models import HealthStatus, ProbeResult, Service, ServiceKind _probe_metadata: contextvars.ContextVar[dict[str, object] | None] = contextvars.ContextVar( "servicemap_probe_metadata", default=None ) _service_locks: dict[str, asyncio.Lock] = {} def _lock_for(service_id: str) -> asyncio.Lock: lock = _service_locks.get(service_id) if lock is None: lock = asyncio.Lock() _service_locks[service_id] = lock return lock def _http_status(status_code: int, service: Service) -> HealthStatus: minimum = service.expected_status_min maximum = service.expected_status_max if minimum is not None or maximum is not None: if (minimum is not None and status_code < minimum) or (maximum is not None and status_code > maximum): return HealthStatus.DOWN if status_code >= 500 else HealthStatus.DEGRADED return HealthStatus.HEALTHY if status_code < 400: return HealthStatus.HEALTHY if status_code < 500: return HealthStatus.DEGRADED return HealthStatus.DOWN async def _probe_once_with_metadata(service: Service, timeout_seconds: float) -> tuple[HealthStatus, str, dict[str, object]]: if service.kind == ServiceKind.HTTP: async with httpx.AsyncClient(timeout=timeout_seconds, follow_redirects=True) as client: response = await client.get(service.url or "") status = _http_status(response.status_code, service) return status, f"HTTP {response.status_code}", {"status_code": response.status_code} if service.kind == ServiceKind.DNS: started = perf_counter() infos = await asyncio.wait_for( asyncio.get_running_loop().getaddrinfo(service.host, None, type=socket.SOCK_STREAM), timeout=timeout_seconds, ) addresses = sorted({str(info[4][0]) for info in infos if info[4]}) if not addresses: return HealthStatus.DOWN, f"DNS {service.host} returned no addresses", {"addresses": []} return ( HealthStatus.HEALTHY, f"DNS {service.host} resolved ({round((perf_counter() - started) * 1000, 2)} ms)", {"addresses": addresses[:20]}, ) if service.kind == ServiceKind.TLS: context = ssl.create_default_context() reader, writer = await asyncio.wait_for( asyncio.open_connection(service.host, service.port, ssl=context, server_hostname=service.host), timeout=timeout_seconds, ) del reader ssl_object = writer.get_extra_info("ssl_object") metadata: dict[str, object] = {"hostname": service.host, "port": service.port} detail = f"TLS {service.host}:{service.port} certificate reachable" status = HealthStatus.HEALTHY if ssl_object is not None: certificate = ssl_object.getpeercert() or {} not_after = certificate.get("notAfter") if isinstance(certificate, dict) else None if not_after: expires_at = datetime.fromtimestamp(ssl.cert_time_to_seconds(not_after), UTC) days_remaining = (expires_at - datetime.now(UTC)).total_seconds() / 86400 metadata.update({"expires_at": expires_at.isoformat(), "days_remaining": round(days_remaining, 2)}) warning_days = max(1, int(os.getenv("SERVICEMAP_TLS_WARNING_DAYS", "30"))) if days_remaining <= 0: status = HealthStatus.DOWN detail = f"TLS certificate expired on {expires_at.date().isoformat()}" elif days_remaining <= warning_days: status = HealthStatus.DEGRADED detail = f"TLS certificate expires in {round(days_remaining, 1)} days" writer.close() await writer.wait_closed() return status, detail, metadata _, writer = await asyncio.wait_for( asyncio.open_connection(service.host, service.port), timeout=timeout_seconds, ) writer.close() await writer.wait_closed() return HealthStatus.HEALTHY, f"TCP {service.host}:{service.port} reachable", {} async def _probe_once(service: Service, timeout_seconds: float) -> tuple[HealthStatus, str]: status, detail, metadata = await _probe_once_with_metadata(service, timeout_seconds) _probe_metadata.set(metadata) return status, detail async def probe_service( service: Service, timeout_seconds: float | None = 5.0, retries: int | None = 1, *, trigger: str = "manual", default_timeout_seconds: float = 5.0, default_retries: int = 1, ) -> ProbeResult: if not service.enabled or not service.monitoring_enabled: return ProbeResult( id=str(uuid4()), service_id=service.id, status=HealthStatus.UNKNOWN, latency_ms=None, detail="Monitoring disabled for service", checked_at=datetime.now(UTC), trigger=trigger, ) timeout = service.timeout_seconds if timeout_seconds is None else timeout_seconds timeout = timeout if timeout is not None else default_timeout_seconds retry_count = service.retries if retries is None else retries retry_count = retry_count if retry_count is not None else default_retries attempts = max(1, retry_count + 1) status = HealthStatus.DOWN detail = "Probe failed" latency_ms: float | None = None metadata: dict[str, object] = {} async with _lock_for(service.id): for attempt in range(1, attempts + 1): _probe_metadata.set({}) started = perf_counter() try: status, detail = await _probe_once(service, timeout) metadata = _probe_metadata.get() or {} latency_ms = round((perf_counter() - started) * 1000, 2) if attempt > 1: detail = f"{detail} after {attempt} attempts" break except (TimeoutError, httpx.HTTPError, OSError, ssl.SSLError, socket.gaierror) as exc: latency_ms = round((perf_counter() - started) * 1000, 2) detail = f"{type(exc).__name__}: {exc}" if isinstance(exc, ssl.CertificateError): status = HealthStatus.DOWN detail = f"TLS hostname mismatch: {exc}" if attempt < attempts: await asyncio.sleep(min(0.2 * attempt, 1.0)) return ProbeResult( id=str(uuid4()), service_id=service.id, status=status, latency_ms=latency_ms, detail=detail[:500], checked_at=datetime.now(UTC), trigger=trigger, metadata=metadata, ) async def probe_many( services: list[Service], *, timeout_seconds: float | None = 5.0, retries: int | None = 1, concurrency: int = 10, trigger: str = "manual", default_timeout_seconds: float = 5.0, default_retries: int = 1, ) -> list[ProbeResult]: semaphore = asyncio.Semaphore(max(1, concurrency)) async def run(service: Service) -> ProbeResult: async with semaphore: if trigger == "manual" and default_timeout_seconds == 5.0 and default_retries == 1: return await probe_service(service, timeout_seconds=timeout_seconds, retries=retries) if trigger == "manual": return await probe_service( service, timeout_seconds=timeout_seconds, retries=retries, default_timeout_seconds=default_timeout_seconds, default_retries=default_retries, ) return await probe_service( service, timeout_seconds=timeout_seconds, retries=retries, trigger=trigger, default_timeout_seconds=default_timeout_seconds, default_retries=default_retries, ) return list(await asyncio.gather(*(run(service) for service in services)))