mirror of
https://github.com/tiennm99/DocsGPT.git
synced 2026-10-14 22:17:00 +00:00
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:
1 parent
54b540ea42
commit
e889392eba
17 files changed
+911
-255
No files matched your search
@@ -0,0 +1,92 @@
|
||||
"""Tests for ``docsgpt/api/asgi_stream.py`` — cleanup that can't stall a worker.
|
||||
|
||||
The cleanup runs shielded from the cancellation that ended the response, so
|
||||
each step needs its own deadline: a close or release waiting on a dead Redis
|
||||
connection must not hold the request task, or a graceful shutdown, forever.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import anyio
|
||||
import pytest
|
||||
|
||||
from docsgpt.api import asgi_stream
|
||||
from docsgpt.api.asgi_stream import ClosingStreamingResponse
|
||||
from tests.asgi_stream import stream_then_disconnect
|
||||
|
||||
|
||||
class _StuckCloseIterator:
|
||||
"""Yields one chunk, then ends; its ``aclose`` never returns."""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self._sent = False
|
||||
self.aclose_started = False
|
||||
|
||||
def __aiter__(self):
|
||||
return self
|
||||
|
||||
async def __anext__(self):
|
||||
if self._sent:
|
||||
raise StopAsyncIteration
|
||||
self._sent = True
|
||||
return b"chunk"
|
||||
|
||||
async def aclose(self):
|
||||
self.aclose_started = True
|
||||
await anyio.sleep_forever()
|
||||
|
||||
|
||||
async def _one_chunk():
|
||||
yield b"chunk"
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
@pytest.mark.asyncio
|
||||
class TestBoundedCleanup:
|
||||
async def test_stuck_stream_close_is_abandoned_and_on_close_still_runs(self, monkeypatch):
|
||||
monkeypatch.setattr(asgi_stream, "_CLEANUP_TIMEOUT_SECONDS", 0.05)
|
||||
body = _StuckCloseIterator()
|
||||
released = {"ran": False}
|
||||
|
||||
async def on_close():
|
||||
released["ran"] = True
|
||||
|
||||
response = ClosingStreamingResponse(body, media_type="text/plain", on_close=on_close)
|
||||
with anyio.fail_after(2):
|
||||
status, _, chunks = await stream_then_disconnect(
|
||||
response, "/", chunks_before_disconnect=10**6
|
||||
)
|
||||
assert status == 200
|
||||
assert chunks == [b"chunk"]
|
||||
assert body.aclose_started
|
||||
assert released["ran"] is True
|
||||
|
||||
async def test_stuck_on_close_is_abandoned(self, monkeypatch):
|
||||
monkeypatch.setattr(asgi_stream, "_CLEANUP_TIMEOUT_SECONDS", 0.05)
|
||||
|
||||
async def on_close():
|
||||
await anyio.sleep_forever()
|
||||
|
||||
response = ClosingStreamingResponse(_one_chunk(), media_type="text/plain", on_close=on_close)
|
||||
with anyio.fail_after(2):
|
||||
status, _, chunks = await stream_then_disconnect(
|
||||
response, "/", chunks_before_disconnect=10**6
|
||||
)
|
||||
assert status == 200
|
||||
assert chunks == [b"chunk"]
|
||||
|
||||
async def test_on_close_runs_when_client_leaves_before_first_frame(self):
|
||||
released = {"count": 0}
|
||||
|
||||
async def on_close():
|
||||
released["count"] += 1
|
||||
|
||||
async def endless():
|
||||
while True:
|
||||
await anyio.sleep(0.01)
|
||||
yield b"tick"
|
||||
|
||||
response = ClosingStreamingResponse(endless(), media_type="text/plain", on_close=on_close)
|
||||
with anyio.fail_after(2):
|
||||
await stream_then_disconnect(response, "/", chunks_before_disconnect=0)
|
||||
assert released["count"] == 1
|
||||
Reference in new issue
Block a user