fix(api): address review on the event-loop streams

- Per-user SSE cap uses per-connection leases in a sorted set instead of
  a shared INCR/DECR counter with a TTL. A stream that outlived the TTL
  could let the counter expire, then decrement another stream's slot or
  drive the count negative. Leases refresh while a stream sends frames
  and age out when a stream dies without cleanup.
- Shielded cleanup is bounded per step: stream close and on_close in
  ClosingStreamingResponse, unsubscribe and close in AsyncTopic, lease
  release, and the replay-budget check, so a dead Redis connection can't
  hold a request or a graceful shutdown.
- Artifact downloads send an ASCII filename plus an RFC 5987 filename*
  for non-ASCII names, build headers before opening the file, and close
  disk handles off the event loop.
- Type hints and docstrings on the new helpers.
This commit is contained in:
Alex committed 2026-09-14 10:19:59 +01:00
1 parent 54b540ea42
commit e889392eba
17 files changed
+911 -255

No files matched your search

+32 -3
View File
@@ -11,9 +11,12 @@ from __future__ import annotations
import io
import logging
import os
import re
import unicodedata
from functools import partial
from typing import AsyncIterator, BinaryIO, Mapping, Optional, Tuple
from urllib.parse import quote
import anyio
from starlette.requests import Request
@@ -47,6 +50,25 @@ def _sanitize_header_filename(filename: Optional[str], fallback: str) -> str:
return cleaned or fallback
def _ascii_fold(value: str) -> str:
return unicodedata.normalize("NFKD", value).encode("ascii", "ignore").decode("ascii")
def _content_disposition(filename: Optional[str], artifact_id: str) -> str:
"""Build an ``attachment`` Content-Disposition value that any filename can travel in.
Header values are Latin-1, so a non-ASCII name gets an ASCII ``filename``
fallback plus an RFC 5987 ``filename*`` carrying the real name.
"""
name = _sanitize_header_filename(filename, f"artifact-{artifact_id}")
if name.isascii():
return f'attachment; filename="{name}"'
stem, ext = os.path.splitext(name)
ascii_stem = _ascii_fold(stem).strip()
fallback = f"{ascii_stem or f'artifact-{artifact_id}'}{_ascii_fold(ext)}"
return f"attachment; filename=\"{fallback}\"; filename*=UTF-8''{quote(name, safe='')}"
def _load_version(
artifact_id: str, decoded_token: Optional[dict], query: Mapping[str, str]
) -> Tuple[Optional[dict], Optional[Tuple[str, int]]]:
@@ -90,9 +112,14 @@ async def _read_chunks(file_obj: BinaryIO) -> AsyncIterator[bytes]:
async def _close_file(file_obj: BinaryIO) -> None:
"""Close the handle; closing anything but an in-memory buffer runs in a worker thread."""
close = getattr(file_obj, "close", None)
if callable(close):
if not callable(close):
return
if isinstance(file_obj, io.BytesIO):
close()
else:
await anyio.to_thread.run_sync(close)
async def download_artifact(request: Request) -> Response:
@@ -118,7 +145,6 @@ async def download_artifact(request: Request) -> Response:
if not storage_path:
return json_error("No file for this version", 404)
filename = _sanitize_header_filename(version_row.get("filename"), f"artifact-{artifact_id}")
mime_type = version_row.get("mime_type") or "application/octet-stream"
storage = await anyio.to_thread.run_sync(StorageCreator.get_storage)
@@ -157,6 +183,9 @@ async def download_artifact(request: Request) -> Response:
)
return RedirectResponse(url, status_code=302)
# Build the headers before opening the file, so nothing between the
# open and the response can fail and leave the handle unclosed.
headers = {"Content-Disposition": _content_disposition(version_row.get("filename"), artifact_id)}
# Stream the bytes in chunks instead of buffering the whole object in
# memory (artifacts can be many MB); the handle closes when the
# response ends, including on a client disconnect.
@@ -164,7 +193,7 @@ async def download_artifact(request: Request) -> Response:
return ClosingStreamingResponse(
_read_chunks(file_obj),
media_type=mime_type,
headers={"Content-Disposition": f'attachment; filename="{filename}"'},
headers=headers,
on_close=partial(_close_file, file_obj),
)
except FileNotFoundError: