feat: bound aggregate Gitea transport traffic (#246)

src/gitea_proxy.py

@@ -2,6 +2,7 @@ import asyncio
 import os
 import re
 import shlex
+from contextlib import asynccontextmanager
 from typing import Any
 from urllib.parse import urlsplit
 
@@ -12,7 +13,116 @@ GITEA_TOKEN = os.getenv("GITEA_TOKEN", "")
 REVIEW_DIFF_MAX_BYTES = 64 * 1024
 REVIEW_DIFF_MAX_LINES = 400
 AVAILABLE_ISSUE_PAGE_CONCURRENCY = 3
-_client: httpx.AsyncClient | None = None
+_client: "GiteaTransport | None" = None
+
+
+class GiteaOverloadedError(RuntimeError):
+    """Raised when bounded transport admission expires before capacity is available."""
+
+
+class GiteaTransport:
+    """Application-lifetime HTTP transport with single-flight concurrent reads."""
+
+    def __init__(
+        self,
+        *,
+        max_concurrency: int = 8,
+        admission_timeout: float = 0.25,
+        **kwargs,
+    ) -> None:
+        max_concurrency = max(2, max_concurrency)
+        admission_timeout = max(0.001, admission_timeout)
+        self._http = httpx.AsyncClient(base_url=GITEA_URL, timeout=10, **kwargs)
+        self.max_concurrency = max_concurrency
+        self.admission_timeout = admission_timeout
+        self._request_slots = asyncio.Semaphore(max(2, max_concurrency))
+        self._read_slots = asyncio.Semaphore(max(1, max_concurrency - 1))
+        self._reads: dict[tuple, asyncio.Task[httpx.Response]] = {}
+
+    @property
+    def is_closed(self) -> bool:
+        return self._http.is_closed
+
+    def _read_key(self, url: str, kwargs: dict) -> tuple:
+        params = tuple(httpx.QueryParams(kwargs.get("params", {})).multi_items())
+        headers = tuple(sorted(httpx.Headers(kwargs.get("headers", {})).multi_items()))
+        return url, params, headers
+
+    async def _acquire(self, semaphore: asyncio.Semaphore) -> None:
+        try:
+            await asyncio.wait_for(semaphore.acquire(), timeout=self.admission_timeout)
+        except TimeoutError as exc:
+            raise GiteaOverloadedError("Gitea transport is saturated") from exc
+
+    async def _perform_get(self, url: str, kwargs: dict) -> httpx.Response:
+        await self._acquire(self._read_slots)
+        try:
+            await self._acquire(self._request_slots)
+            try:
+                return await self._http.get(url, **kwargs)
+            finally:
+                self._request_slots.release()
+        finally:
+            self._read_slots.release()
+
+    async def get(self, url: str, **kwargs) -> httpx.Response:
+        key = self._read_key(url, kwargs)
+        task = self._reads.get(key)
+        if task is None:
+            task = asyncio.create_task(self._perform_get(url, kwargs))
+            self._reads[key] = task
+            task.add_done_callback(
+                lambda completed, request_key=key: (
+                    self._reads.pop(request_key, None)
+                    if self._reads.get(request_key) is completed else None
+                )
+            )
+        return await asyncio.shield(task)
+
+    async def _mutate(self, method: str, url: str, kwargs: dict) -> httpx.Response:
+        await self._acquire(self._request_slots)
+        try:
+            return await self._http.request(method, url, **kwargs)
+        finally:
+            self._request_slots.release()
+
+    async def post(self, url: str, **kwargs) -> httpx.Response:
+        return await self._mutate("POST", url, kwargs)
+
+    async def patch(self, url: str, **kwargs) -> httpx.Response:
+        return await self._mutate("PATCH", url, kwargs)
+
+    async def put(self, url: str, **kwargs) -> httpx.Response:
+        return await self._mutate("PUT", url, kwargs)
+
+    async def delete(self, url: str, **kwargs) -> httpx.Response:
+        return await self._mutate("DELETE", url, kwargs)
+
+    @asynccontextmanager
+    async def stream(self, method: str, url: str, **kwargs):
+        await self._acquire(self._read_slots)
+        try:
+            await self._acquire(self._request_slots)
+            try:
+                async with self._http.stream(method, url, **kwargs) as response:
+                    yield response
+            finally:
+                self._request_slots.release()
+        finally:
+            self._read_slots.release()
+
+    def __getattr__(self, name: str):
+        return getattr(self._http, name)
+
+    async def aclose(self) -> None:
+        tasks = list(self._reads.values())
+        for task in tasks:
+            if not task.done():
+                task.cancel()
+        if tasks:
+            await asyncio.gather(*tasks, return_exceptions=True)
+        self._reads.clear()
+        await self._http.aclose()
 
 
 class WorkItems(list[dict]):
@@ -58,14 +168,19 @@ def _auth() -> dict[str, str]:
     return headers
 
 
-def start_client(**kwargs) -> httpx.AsyncClient:
+def start_client(**kwargs) -> GiteaTransport:
     """Create the application-lifetime Gitea transport."""
     global _client
-    _client = httpx.AsyncClient(base_url=GITEA_URL, timeout=10, **kwargs)
+    kwargs.setdefault("max_concurrency", int(os.getenv("GITEA_MAX_CONCURRENCY", "8")))
+    kwargs.setdefault(
+        "admission_timeout",
+        float(os.getenv("GITEA_ADMISSION_TIMEOUT_SECONDS", "0.25")),
+    )
+    _client = GiteaTransport(**kwargs)
     return _client
 
 
-def _get_client() -> httpx.AsyncClient:
+def _get_client() -> GiteaTransport:
     if _client is None or _client.is_closed:
         return start_client()
     return _client