125 lines
5.0 KiB
Python
125 lines
5.0 KiB
Python
"""Bounded, identity-free cache for the embedded WeKnora application's assets."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import gzip
|
|
import hashlib
|
|
import time
|
|
from collections import OrderedDict
|
|
from collections.abc import Awaitable, Callable
|
|
from dataclasses import dataclass
|
|
|
|
from fastapi import Request
|
|
from fastapi.responses import Response
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class StaticAsset:
|
|
content: bytes
|
|
compressed: bytes | None
|
|
headers: dict[str, str]
|
|
etag: str
|
|
compressed_etag: str
|
|
ttl: float
|
|
|
|
@property
|
|
def size(self) -> int:
|
|
return len(self.content) + len(self.compressed or b"")
|
|
|
|
def response(self, request: Request) -> Response:
|
|
compress = self.compressed is not None and _accepts_gzip(request.headers.get("accept-encoding", ""))
|
|
content = self.compressed if compress else self.content
|
|
etag = self.compressed_etag if compress else self.etag
|
|
headers = {**self.headers, "etag": etag, "vary": "Accept-Encoding"}
|
|
if compress:
|
|
headers["content-encoding"] = "gzip"
|
|
if _etag_matches(request.headers.get("if-none-match", ""), etag):
|
|
return Response(status_code=304, headers=headers)
|
|
headers["content-length"] = str(len(content))
|
|
return Response(content=b"" if request.method == "HEAD" else content, headers=headers)
|
|
|
|
|
|
def _accepts_gzip(value: str) -> bool:
|
|
qualities: dict[str, float] = {}
|
|
for item in value.split(","):
|
|
encoding, *parameters = item.strip().lower().split(";")
|
|
quality = 1.0
|
|
for parameter in parameters:
|
|
if parameter.strip().startswith("q="):
|
|
try:
|
|
quality = float(parameter.strip()[2:])
|
|
except ValueError:
|
|
quality = 0.0
|
|
qualities[encoding] = quality
|
|
return qualities.get("gzip", qualities.get("*", 0.0)) > 0
|
|
|
|
|
|
def _etag_matches(value: str, etag: str) -> bool:
|
|
return any(candidate.strip().removeprefix("W/") in {etag, "*"} for candidate in value.split(","))
|
|
|
|
|
|
def prepare_static_asset(content: bytes, *, content_type: str, cache_control: str, ttl: float) -> StaticAsset:
|
|
"""Hash and compress rewritten bytes; call from a worker thread."""
|
|
compressed = None
|
|
if len(content) >= 1024 and any(kind in content_type.lower() for kind in ("text/", "javascript", "json", "svg", "xml", "manifest")):
|
|
candidate = gzip.compress(content, compresslevel=5, mtime=0)
|
|
if len(candidate) < len(content):
|
|
compressed = candidate
|
|
return StaticAsset(
|
|
content=content,
|
|
compressed=compressed,
|
|
headers={"content-type": content_type, "cache-control": cache_control, "x-content-type-options": "nosniff"},
|
|
etag=f'"{hashlib.sha256(content).hexdigest()}"',
|
|
compressed_etag=f'"{hashlib.sha256(compressed).hexdigest()}"' if compressed is not None else "",
|
|
ttl=ttl,
|
|
)
|
|
|
|
|
|
class StaticAssetCache:
|
|
"""LRU bounded by bytes and entries; coalesces an asset's in-flight reads."""
|
|
|
|
def __init__(self, *, max_bytes: int = 64 * 1024 * 1024, max_asset_bytes: int = 8 * 1024 * 1024, max_entries: int = 256) -> None:
|
|
self.max_bytes = max_bytes
|
|
self.max_asset_bytes = max_asset_bytes
|
|
self.max_entries = max_entries
|
|
self._entries: OrderedDict[tuple[str, ...], tuple[float, StaticAsset]] = OrderedDict()
|
|
self._pending: dict[tuple[str, ...], asyncio.Task[StaticAsset]] = {}
|
|
self._bytes = 0
|
|
|
|
async def get(self, key: tuple[str, ...], load: Callable[[], Awaitable[StaticAsset]], *, refresh: bool = False) -> StaticAsset:
|
|
cached = self._entries.get(key)
|
|
if cached is not None and not refresh:
|
|
if cached[0] > time.monotonic():
|
|
self._entries.move_to_end(key)
|
|
return cached[1]
|
|
self._discard(key)
|
|
task = self._pending.get(key)
|
|
if task is None:
|
|
task = asyncio.create_task(self._load(key, load))
|
|
self._pending[key] = task
|
|
|
|
def clear(completed: asyncio.Task[StaticAsset]) -> None:
|
|
if self._pending.get(key) is completed:
|
|
self._pending.pop(key, None)
|
|
if not completed.cancelled():
|
|
completed.exception()
|
|
|
|
task.add_done_callback(clear)
|
|
return await asyncio.shield(task)
|
|
|
|
async def _load(self, key: tuple[str, ...], load: Callable[[], Awaitable[StaticAsset]]) -> StaticAsset:
|
|
asset = await load()
|
|
self._discard(key)
|
|
if asset.ttl > 0 and asset.size <= min(self.max_asset_bytes, self.max_bytes):
|
|
self._entries[key] = (time.monotonic() + asset.ttl, asset)
|
|
self._bytes += asset.size
|
|
while self._bytes > self.max_bytes or len(self._entries) > self.max_entries:
|
|
self._discard(next(iter(self._entries)))
|
|
return asset
|
|
|
|
def _discard(self, key: tuple[str, ...]) -> None:
|
|
previous = self._entries.pop(key, None)
|
|
if previous is not None:
|
|
self._bytes -= previous[1].size
|