205 lines
8.0 KiB
Python
205 lines
8.0 KiB
Python
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)))
|