fix: guard reserved sessions and bounded response state
This commit is contained in:
parent
0e85c399bb
commit
745a29d979
3 changed files with 62 additions and 3 deletions
|
|
@ -270,7 +270,8 @@ def create_app(
|
||||||
async def lifespan(_app: FastAPI):
|
async def lifespan(_app: FastAPI):
|
||||||
yield
|
yield
|
||||||
for session in list(sessions.values()):
|
for session in list(sessions.values()):
|
||||||
session.close()
|
if session is not None:
|
||||||
|
session.close()
|
||||||
engine.shutdown()
|
engine.shutdown()
|
||||||
|
|
||||||
app = FastAPI(
|
app = FastAPI(
|
||||||
|
|
@ -1111,6 +1112,11 @@ def main() -> None:
|
||||||
LOGGER.info("Loading model and catalog (device=%s, imgsz=%s)...",
|
LOGGER.info("Loading model and catalog (device=%s, imgsz=%s)...",
|
||||||
engine_settings.device or "auto", engine_settings.imgsz)
|
engine_settings.device or "auto", engine_settings.imgsz)
|
||||||
engine = build_default_engine(engine_settings)
|
engine = build_default_engine(engine_settings)
|
||||||
|
if (engine_settings.identity_backend or "").lower() == "gallery":
|
||||||
|
if engine.identity_gallery() is None:
|
||||||
|
raise RuntimeError(
|
||||||
|
"gallery identity backend requires local DINOv2 weights; refusing to start"
|
||||||
|
)
|
||||||
warm_ms = engine.warm_up()
|
warm_ms = engine.warm_up()
|
||||||
LOGGER.info(
|
LOGGER.info(
|
||||||
"Ready: %d SKUs, catalog %s, warmup %.0f ms, classify workers %d",
|
"Ready: %d SKUs, catalog %s, warmup %.0f ms, classify workers %d",
|
||||||
|
|
|
||||||
|
|
@ -196,6 +196,31 @@ class RestContractTests(unittest.TestCase):
|
||||||
self.assertEqual(sent[0]["status"], 413)
|
self.assertEqual(sent[0]["status"], 413)
|
||||||
self.assertIn(b"PAYLOAD_TOO_LARGE", sent[1]["body"])
|
self.assertIn(b"PAYLOAD_TOO_LARGE", sent[1]["body"])
|
||||||
|
|
||||||
|
def test_overflow_after_response_start_sends_no_second_start(self):
|
||||||
|
"""A body that overflows after headers began must not start again."""
|
||||||
|
import anyio
|
||||||
|
|
||||||
|
from transport import BodyLimit
|
||||||
|
|
||||||
|
async def app(scope, receive, send):
|
||||||
|
await send({"type": "http.response.start", "status": 200, "headers": []})
|
||||||
|
await receive()
|
||||||
|
await send({"type": "http.response.body", "body": b"ok"})
|
||||||
|
|
||||||
|
sent = []
|
||||||
|
|
||||||
|
async def receive():
|
||||||
|
return {"type": "http.request", "body": b"x" * 40, "more_body": False}
|
||||||
|
|
||||||
|
async def send(message):
|
||||||
|
sent.append(message)
|
||||||
|
|
||||||
|
anyio.run(BodyLimit(app, 8), {"type": "http"}, receive, send)
|
||||||
|
starts = [message for message in sent if message.get("type") == "http.response.start"]
|
||||||
|
self.assertEqual(len(starts), 1)
|
||||||
|
self.assertEqual(starts[0]["status"], 200)
|
||||||
|
self.assertFalse(any(message.get("status") == 413 for message in sent))
|
||||||
|
|
||||||
def test_classify_and_inspect_roundtrip(self):
|
def test_classify_and_inspect_roundtrip(self):
|
||||||
recognized = self.client.post(
|
recognized = self.client.post(
|
||||||
"/v1/recognize?classify=false", content=jpeg_bytes(),
|
"/v1/recognize?classify=false", content=jpeg_bytes(),
|
||||||
|
|
@ -360,6 +385,16 @@ class WebSocketProtocolTests(unittest.TestCase):
|
||||||
finally:
|
finally:
|
||||||
engine.shutdown()
|
engine.shutdown()
|
||||||
|
|
||||||
|
def test_lifespan_shutdown_skips_a_reserved_slot(self):
|
||||||
|
"""A reservation is None until setup finishes. Shutdown must not call it."""
|
||||||
|
client, engine = make_app()
|
||||||
|
client.app.state.sessions["reserved"] = None
|
||||||
|
try:
|
||||||
|
with client:
|
||||||
|
pass
|
||||||
|
finally:
|
||||||
|
engine.shutdown()
|
||||||
|
|
||||||
|
|
||||||
class LegacyRouteTests(unittest.TestCase):
|
class LegacyRouteTests(unittest.TestCase):
|
||||||
def test_legacy_sessions_still_respond(self):
|
def test_legacy_sessions_still_respond(self):
|
||||||
|
|
|
||||||
|
|
@ -104,7 +104,10 @@ class BodyLimit:
|
||||||
|
|
||||||
A handler that checks Content-Length, or reads request.body(), runs after
|
A handler that checks Content-Length, or reads request.body(), runs after
|
||||||
the server has already buffered the stream. This counts bytes at receive
|
the server has already buffered the stream. This counts bytes at receive
|
||||||
and returns 413 without handing the rest to the app. WebSocket frames are
|
and returns 413 without handing the rest to the app. Content-Length is
|
||||||
|
not trusted: the stream is what counts, and reading stops at the cap.
|
||||||
|
If the app already started the response, the connection is aborted
|
||||||
|
instead of sending a second http.response.start. WebSocket frames are
|
||||||
not HTTP bodies and are left alone.
|
not HTTP bodies and are left alone.
|
||||||
"""
|
"""
|
||||||
|
|
||||||
|
|
@ -116,10 +119,13 @@ class BodyLimit:
|
||||||
if scope.get("type") != "http":
|
if scope.get("type") != "http":
|
||||||
await self.app(scope, receive, send)
|
await self.app(scope, receive, send)
|
||||||
return
|
return
|
||||||
|
|
||||||
class TooLarge(Exception):
|
class TooLarge(Exception):
|
||||||
pass
|
pass
|
||||||
|
|
||||||
seen = 0
|
seen = 0
|
||||||
|
started = False
|
||||||
|
|
||||||
async def limited() -> Dict[str, Any]:
|
async def limited() -> Dict[str, Any]:
|
||||||
nonlocal seen
|
nonlocal seen
|
||||||
message = await receive()
|
message = await receive()
|
||||||
|
|
@ -128,12 +134,24 @@ class BodyLimit:
|
||||||
body = message.get("body") or b""
|
body = message.get("body") or b""
|
||||||
seen += len(body)
|
seen += len(body)
|
||||||
if seen > self.max_bytes:
|
if seen > self.max_bytes:
|
||||||
|
# The app must not see this chunk or any later one.
|
||||||
raise TooLarge
|
raise TooLarge
|
||||||
return message
|
return message
|
||||||
|
|
||||||
|
async def tracked(message: Dict[str, Any]) -> None:
|
||||||
|
nonlocal started
|
||||||
|
if message.get("type") == "http.response.start":
|
||||||
|
started = True
|
||||||
|
await send(message)
|
||||||
|
|
||||||
try:
|
try:
|
||||||
await self.app(scope, limited, send)
|
await self.app(scope, limited, tracked)
|
||||||
except TooLarge:
|
except TooLarge:
|
||||||
|
if started:
|
||||||
|
# The status line is already on the wire. Returning leaves the
|
||||||
|
# response unfinished so the server closes the connection
|
||||||
|
# instead of accepting a second http.response.start.
|
||||||
|
return
|
||||||
await send(
|
await send(
|
||||||
{
|
{
|
||||||
"type": "http.response.start",
|
"type": "http.response.start",
|
||||||
|
|
|
||||||
Loading…
Add table
Add a link
Reference in a new issue