393 lines
13 KiB
Python
393 lines
13 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
from collections.abc import Awaitable, Callable
|
|
from dataclasses import dataclass
|
|
import hashlib
|
|
import json
|
|
import time
|
|
from typing import Any
|
|
|
|
from fastapi import Response
|
|
|
|
from app.core.cache import _RedisClient
|
|
from app.core.config import settings
|
|
from app.core.logging import get_logger
|
|
|
|
|
|
logger = get_logger(__name__, service="earth_layer_cache")
|
|
|
|
EARTH_LAYER_CACHE_PREFIX = "earth:layer:v1"
|
|
EARTH_LAYER_LOCK_PREFIX = "earth:layer:lock:v1"
|
|
DEFAULT_LOCK_TTL_SECONDS = 10
|
|
DEFAULT_LOCK_WAIT_SECONDS = 0.2
|
|
DEFAULT_MAX_FEATURES = 5000
|
|
DEFAULT_MAX_BYTES = 5 * 1024 * 1024
|
|
DEFAULT_BBOX_PRECISION_DEGREES = 0.1
|
|
DEV_CACHE_KEY_HEADER = {"development", "dev", "test", "testing", "local"}
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class EarthLayerCachePolicy:
|
|
fresh_ttl_seconds: int
|
|
stale_ttl_seconds: int
|
|
max_features: int = DEFAULT_MAX_FEATURES
|
|
max_bytes: int = DEFAULT_MAX_BYTES
|
|
lock_ttl_seconds: int = DEFAULT_LOCK_TTL_SECONDS
|
|
lock_wait_seconds: float = DEFAULT_LOCK_WAIT_SECONDS
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class EarthLayerCacheResult:
|
|
payload: dict[str, Any]
|
|
state: str
|
|
key: str
|
|
features: int
|
|
bytes: int
|
|
|
|
|
|
class EarthLayerCache:
|
|
def __init__(self) -> None:
|
|
self._client = None
|
|
|
|
@property
|
|
def client(self):
|
|
if self._client is None:
|
|
self._client = _RedisClient.get_client()
|
|
return self._client
|
|
|
|
@staticmethod
|
|
def key(layer: str, **params: Any) -> str:
|
|
parts = [EARTH_LAYER_CACHE_PREFIX, _safe_key_part(layer)]
|
|
for name in sorted(params):
|
|
value = params[name]
|
|
if value is None:
|
|
value = "none"
|
|
parts.append(f"{_safe_key_part(name)}:{_safe_key_part(value)}")
|
|
return ":".join(parts)
|
|
|
|
@staticmethod
|
|
def stale_key(key: str) -> str:
|
|
return f"{key}:stale"
|
|
|
|
@staticmethod
|
|
def lock_key(key: str) -> str:
|
|
digest = hashlib.sha256(key.encode("utf-8")).hexdigest()[:32]
|
|
return f"{EARTH_LAYER_LOCK_PREFIX}:{digest}"
|
|
|
|
def get_json(self, key: str) -> dict[str, Any] | None:
|
|
raw = self.client.get(key)
|
|
if not raw:
|
|
return None
|
|
value = json.loads(raw)
|
|
return value if isinstance(value, dict) else None
|
|
|
|
def set_json(self, key: str, payload: dict[str, Any], ttl_seconds: int) -> None:
|
|
self.client.setex(key, ttl_seconds, json.dumps(payload, ensure_ascii=False, default=str))
|
|
|
|
def acquire_lock(self, key: str, ttl_seconds: int) -> bool:
|
|
return bool(self.client.set(self.lock_key(key), "1", nx=True, ex=ttl_seconds))
|
|
|
|
def release_lock(self, key: str) -> None:
|
|
try:
|
|
self.client.delete(self.lock_key(key))
|
|
except Exception:
|
|
pass
|
|
|
|
def delete_pattern(self, pattern: str = f"{EARTH_LAYER_CACHE_PREFIX}:*") -> int:
|
|
keys = list(self.client.scan_iter(match=pattern))
|
|
if not keys:
|
|
return 0
|
|
return int(self.client.delete(*keys))
|
|
|
|
def status(self) -> dict[str, Any]:
|
|
keys = list(self.client.scan_iter(match=f"{EARTH_LAYER_CACHE_PREFIX}:*"))
|
|
by_layer: dict[str, dict[str, Any]] = {}
|
|
total_memory = 0
|
|
for key in keys:
|
|
key_str = key.decode("utf-8") if isinstance(key, bytes) else str(key)
|
|
layer = _layer_from_key(key_str)
|
|
entry = by_layer.setdefault(layer, {"keys": 0, "stale_keys": 0, "memory_bytes": 0})
|
|
entry["keys"] += 1
|
|
if key_str.endswith(":stale"):
|
|
entry["stale_keys"] += 1
|
|
try:
|
|
memory = int(self.client.memory_usage(key) or 0)
|
|
except Exception:
|
|
memory = 0
|
|
entry["memory_bytes"] += memory
|
|
total_memory += memory
|
|
return {
|
|
"prefix": EARTH_LAYER_CACHE_PREFIX,
|
|
"key_count": len(keys),
|
|
"memory_bytes": total_memory,
|
|
"layers": by_layer,
|
|
}
|
|
|
|
|
|
earth_layer_cache = EarthLayerCache()
|
|
|
|
|
|
def quantize_bbox(
|
|
bbox: tuple[float, float, float, float],
|
|
*,
|
|
precision: float = DEFAULT_BBOX_PRECISION_DEGREES,
|
|
) -> tuple[float, float, float, float]:
|
|
return tuple(round(value / precision) * precision for value in bbox) # type: ignore[return-value]
|
|
|
|
|
|
def format_bbox_key(bbox: tuple[float, float, float, float]) -> str:
|
|
return ",".join(f"{value:.1f}" for value in bbox)
|
|
|
|
|
|
def apply_cache_headers(response: Response | None, result: EarthLayerCacheResult) -> None:
|
|
if response is None:
|
|
return
|
|
response.headers["X-Planet-Cache"] = result.state
|
|
response.headers["X-Planet-Cache-Features"] = str(result.features)
|
|
response.headers["X-Planet-Cache-Bytes"] = str(result.bytes)
|
|
env_name = str(getattr(settings, "ENVIRONMENT", "") or "development").lower()
|
|
if env_name in DEV_CACHE_KEY_HEADER:
|
|
response.headers["X-Planet-Cache-Key"] = result.key
|
|
|
|
|
|
async def get_or_build_layer_payload(
|
|
*,
|
|
key: str,
|
|
policy: EarthLayerCachePolicy,
|
|
builder: Callable[[], Awaitable[dict[str, Any]]],
|
|
response: Response | None = None,
|
|
) -> dict[str, Any]:
|
|
result = await resolve_layer_payload(key=key, policy=policy, builder=builder)
|
|
apply_cache_headers(response, result)
|
|
return result.payload
|
|
|
|
|
|
async def resolve_layer_payload(
|
|
*,
|
|
key: str,
|
|
policy: EarthLayerCachePolicy,
|
|
builder: Callable[[], Awaitable[dict[str, Any]]],
|
|
) -> EarthLayerCacheResult:
|
|
started = time.perf_counter()
|
|
try:
|
|
cached = earth_layer_cache.get_json(key)
|
|
if cached is not None:
|
|
return _result(cached, state="hit", key=key)
|
|
|
|
lock_acquired = earth_layer_cache.acquire_lock(key, policy.lock_ttl_seconds)
|
|
if lock_acquired:
|
|
try:
|
|
payload = await _build_budgeted_payload(builder, policy)
|
|
_write_fresh_and_stale(key, payload, policy)
|
|
_log_cache_event("refresh", key, payload, started)
|
|
return _result(payload, state="refresh", key=key)
|
|
except Exception as exc:
|
|
stale = _read_stale(key)
|
|
if stale is not None:
|
|
logger.warning_event(
|
|
"Earth layer cache builder failed; returning stale payload",
|
|
event="earth_layer_cache.stale_after_builder_error",
|
|
context={"key": key, "error": str(exc)},
|
|
)
|
|
return _result(stale, state="stale", key=key)
|
|
raise
|
|
finally:
|
|
earth_layer_cache.release_lock(key)
|
|
|
|
stale = _read_stale(key)
|
|
if stale is not None:
|
|
return _result(stale, state="stale", key=key)
|
|
|
|
await asyncio.sleep(policy.lock_wait_seconds)
|
|
cached_after_wait = earth_layer_cache.get_json(key)
|
|
if cached_after_wait is not None:
|
|
return _result(cached_after_wait, state="hit", key=key)
|
|
|
|
payload = await _build_budgeted_payload(builder, policy)
|
|
_log_cache_event("miss", key, payload, started)
|
|
return _result(payload, state="miss", key=key)
|
|
except Exception as exc:
|
|
try:
|
|
payload = await _build_budgeted_payload(builder, policy)
|
|
except Exception:
|
|
raise exc
|
|
logger.warning_event(
|
|
"Earth layer cache bypassed",
|
|
event="earth_layer_cache.bypass",
|
|
context={"key": key, "error": str(exc)},
|
|
)
|
|
return _result(payload, state="bypass", key=key)
|
|
|
|
|
|
def apply_payload_budget(payload: dict[str, Any], policy: EarthLayerCachePolicy) -> dict[str, Any]:
|
|
budgeted = _truncate_features(payload, policy.max_features, "feature_budget")
|
|
size = _payload_size(budgeted)
|
|
if size <= policy.max_bytes:
|
|
return budgeted
|
|
|
|
features = budgeted.get("features")
|
|
if not isinstance(features, list):
|
|
return _with_budget_diagnostics(
|
|
budgeted,
|
|
truncated=True,
|
|
reason="byte_budget",
|
|
bytes_before=size,
|
|
bytes_after=size,
|
|
)
|
|
|
|
low = 0
|
|
high = len(features)
|
|
best = []
|
|
best_size = _payload_size({**budgeted, "features": best})
|
|
while low <= high:
|
|
mid = (low + high) // 2
|
|
candidate_features = features[:mid]
|
|
candidate = _with_budget_diagnostics(
|
|
{**budgeted, "features": candidate_features},
|
|
truncated=mid < len(features),
|
|
reason="byte_budget",
|
|
bytes_before=size,
|
|
bytes_after=0,
|
|
)
|
|
candidate_size = _payload_size(candidate)
|
|
if candidate_size <= policy.max_bytes:
|
|
best = candidate_features
|
|
best_size = candidate_size
|
|
low = mid + 1
|
|
else:
|
|
high = mid - 1
|
|
|
|
return _with_budget_diagnostics(
|
|
{**budgeted, "features": best},
|
|
truncated=True,
|
|
reason="byte_budget",
|
|
bytes_before=size,
|
|
bytes_after=best_size,
|
|
)
|
|
|
|
|
|
def invalidate_earth_layer_cache_for_source(source: str) -> int:
|
|
source_key = str(source or "").strip()
|
|
patterns = {
|
|
"arcgis_cables": ["cables*", "landing-points*", "summary*"],
|
|
"arcgis_landing_points": ["landing-points*", "summary*"],
|
|
"arcgis_cable_landing_relation": ["landing-points*", "summary*"],
|
|
"celestrak_tle": ["satellites*", "summary*"],
|
|
"spacetrack_tle": ["satellites*", "summary*"],
|
|
"top500": ["compute-centers*", "summary*"],
|
|
"epoch_ai_gpu": ["compute-centers*", "summary*"],
|
|
"ris_live_bgp": ["bgp*", "summary*"],
|
|
"bgpstream_bgp": ["bgp*", "summary*"],
|
|
}.get(source_key, [])
|
|
deleted = 0
|
|
for layer_pattern in patterns:
|
|
deleted += earth_layer_cache.delete_pattern(f"{EARTH_LAYER_CACHE_PREFIX}:{layer_pattern}")
|
|
return deleted
|
|
|
|
|
|
async def _build_budgeted_payload(
|
|
builder: Callable[[], Awaitable[dict[str, Any]]],
|
|
policy: EarthLayerCachePolicy,
|
|
) -> dict[str, Any]:
|
|
payload = await builder()
|
|
return apply_payload_budget(payload, policy)
|
|
|
|
|
|
def _write_fresh_and_stale(key: str, payload: dict[str, Any], policy: EarthLayerCachePolicy) -> None:
|
|
earth_layer_cache.set_json(key, payload, policy.fresh_ttl_seconds)
|
|
earth_layer_cache.set_json(earth_layer_cache.stale_key(key), payload, policy.stale_ttl_seconds)
|
|
|
|
|
|
def _read_stale(key: str) -> dict[str, Any] | None:
|
|
try:
|
|
return earth_layer_cache.get_json(earth_layer_cache.stale_key(key))
|
|
except Exception:
|
|
return None
|
|
|
|
|
|
def _truncate_features(payload: dict[str, Any], max_features: int, reason: str) -> dict[str, Any]:
|
|
features = payload.get("features")
|
|
if not isinstance(features, list) or len(features) <= max_features:
|
|
return payload
|
|
return _with_budget_diagnostics(
|
|
{**payload, "features": features[:max_features]},
|
|
truncated=True,
|
|
reason=reason,
|
|
original_feature_count=len(features),
|
|
)
|
|
|
|
|
|
def _with_budget_diagnostics(
|
|
payload: dict[str, Any],
|
|
*,
|
|
truncated: bool,
|
|
reason: str,
|
|
original_feature_count: int | None = None,
|
|
bytes_before: int | None = None,
|
|
bytes_after: int | None = None,
|
|
) -> dict[str, Any]:
|
|
diagnostics = dict(payload.get("diagnostics") or {})
|
|
diagnostics.update(
|
|
{
|
|
"truncated": bool(truncated or diagnostics.get("truncated")),
|
|
"limit_reason": reason,
|
|
}
|
|
)
|
|
if original_feature_count is not None:
|
|
diagnostics["original_feature_count"] = original_feature_count
|
|
if bytes_before is not None:
|
|
diagnostics["bytes_before_budget"] = bytes_before
|
|
if bytes_after is not None:
|
|
diagnostics["bytes_after_budget"] = bytes_after
|
|
return {**payload, "diagnostics": diagnostics}
|
|
|
|
|
|
def _result(payload: dict[str, Any], *, state: str, key: str) -> EarthLayerCacheResult:
|
|
return EarthLayerCacheResult(
|
|
payload=payload,
|
|
state=state,
|
|
key=key,
|
|
features=_feature_count(payload),
|
|
bytes=_payload_size(payload),
|
|
)
|
|
|
|
|
|
def _feature_count(payload: dict[str, Any]) -> int:
|
|
features = payload.get("features")
|
|
if isinstance(features, list):
|
|
return len(features)
|
|
count = payload.get("count")
|
|
return int(count) if isinstance(count, int) else 0
|
|
|
|
|
|
def _payload_size(payload: dict[str, Any]) -> int:
|
|
return len(json.dumps(payload, ensure_ascii=False, default=str).encode("utf-8"))
|
|
|
|
|
|
def _safe_key_part(value: Any) -> str:
|
|
raw = str(value).strip().lower()
|
|
return "".join(char if char.isalnum() or char in {"-", "_", ".", ","} else "_" for char in raw)[:160]
|
|
|
|
|
|
def _layer_from_key(key: str) -> str:
|
|
prefix = f"{EARTH_LAYER_CACHE_PREFIX}:"
|
|
if not key.startswith(prefix):
|
|
return "unknown"
|
|
remainder = key[len(prefix):]
|
|
return remainder.split(":", 1)[0]
|
|
|
|
|
|
def _log_cache_event(state: str, key: str, payload: dict[str, Any], started: float) -> None:
|
|
logger.info_event(
|
|
"Earth layer cache resolved",
|
|
event="earth_layer_cache.resolved",
|
|
context={
|
|
"state": state,
|
|
"key": key,
|
|
"features": _feature_count(payload),
|
|
"bytes": _payload_size(payload),
|
|
"duration_ms": round((time.perf_counter() - started) * 1000, 2),
|
|
},
|
|
)
|