Browse Source

🐛 Fix race condition in _IncludedRouter lazy cache rebuild

pull/15975/head
uditDewan 4 weeks ago
parent
commit
3b11c0b78f
  1. 78
      fastapi/routing.py
  2. 50
      tests/test_included_router_concurrent_rebuild.py

78
fastapi/routing.py

@ -7,6 +7,7 @@ import inspect
import json import json
import os import os
import stat import stat
import threading
import types import types
from collections.abc import ( from collections.abc import (
AsyncIterator, AsyncIterator,
@ -1579,49 +1580,60 @@ class _IncludedRouter(BaseRoute):
default_factory=list default_factory=list
) )
_effective_low_priority_routes_version: int | None = None _effective_low_priority_routes_version: int | None = None
_rebuild_lock: threading.Lock = field(
default_factory=threading.Lock, repr=False, compare=False
)
def effective_candidates(self) -> list["_EffectiveRouteContext | _IncludedRouter"]: def effective_candidates(self) -> list["_EffectiveRouteContext | _IncludedRouter"]:
routes_version = self.original_router._get_routes_version() routes_version = self.original_router._get_routes_version()
if routes_version == self._effective_candidates_version: if routes_version == self._effective_candidates_version:
return self._effective_candidates return self._effective_candidates
self._effective_candidates = [] with self._rebuild_lock:
candidates = self.original_router.routes if routes_version == self._effective_candidates_version:
for route in candidates: return self._effective_candidates
if isinstance(route, _IncludedRouter): effective_candidates: list[_EffectiveRouteContext | _IncludedRouter] = []
child_context = self.include_context.combine(route.include_context) candidates = self.original_router.routes
child_branch = _IncludedRouter( for route in candidates:
original_router=route.original_router, if isinstance(route, _IncludedRouter):
include_context=child_context, child_context = self.include_context.combine(route.include_context)
) child_branch = _IncludedRouter(
self._effective_candidates.append(child_branch) original_router=route.original_router,
continue include_context=child_context,
route_context = self._build_effective_context(route) )
if route_context is not None: effective_candidates.append(child_branch)
self._effective_candidates.append(route_context) continue
self._effective_candidates_version = routes_version route_context = self._build_effective_context(route)
return self._effective_candidates if route_context is not None:
effective_candidates.append(route_context)
self._effective_candidates = effective_candidates
self._effective_candidates_version = routes_version
return effective_candidates
def effective_low_priority_routes(self) -> list["_EffectiveRouteContext"]: def effective_low_priority_routes(self) -> list["_EffectiveRouteContext"]:
routes_version = self.original_router._get_routes_version() routes_version = self.original_router._get_routes_version()
if routes_version == self._effective_low_priority_routes_version: if routes_version == self._effective_low_priority_routes_version:
return self._effective_low_priority_routes return self._effective_low_priority_routes
self._effective_low_priority_routes = [] with self._rebuild_lock:
for route in self.original_router._low_priority_routes: if routes_version == self._effective_low_priority_routes_version:
route_context = self._build_effective_context(route) return self._effective_low_priority_routes
if route_context is not None: low_priority_routes: list[_EffectiveRouteContext] = []
self._effective_low_priority_routes.append(route_context) for route in self.original_router._low_priority_routes:
for route in self.original_router.routes: route_context = self._build_effective_context(route)
if isinstance(route, _IncludedRouter): if route_context is not None:
child_context = self.include_context.combine(route.include_context) low_priority_routes.append(route_context)
child_branch = _IncludedRouter( for route in self.original_router.routes:
original_router=route.original_router, if isinstance(route, _IncludedRouter):
include_context=child_context, child_context = self.include_context.combine(route.include_context)
) child_branch = _IncludedRouter(
self._effective_low_priority_routes.extend( original_router=route.original_router,
child_branch.effective_low_priority_routes() include_context=child_context,
) )
self._effective_low_priority_routes_version = routes_version low_priority_routes.extend(
return self._effective_low_priority_routes child_branch.effective_low_priority_routes()
)
self._effective_low_priority_routes = low_priority_routes
self._effective_low_priority_routes_version = routes_version
return low_priority_routes
def _build_effective_context( def _build_effective_context(
self, route: BaseRoute self, route: BaseRoute

50
tests/test_included_router_concurrent_rebuild.py

@ -0,0 +1,50 @@
import sys
import threading
from fastapi import APIRouter, FastAPI
from fastapi.routing import _IncludedRouter
from fastapi.testclient import TestClient
N_ROUTES = 120
N_THREADS = 6
def build_app() -> FastAPI:
app = FastAPI()
router = APIRouter()
for i in range(N_ROUTES):
def endpoint(i: int = i) -> dict[str, int]:
return {"i": i}
router.add_api_route(f"/r{i}", endpoint, methods=["GET"])
app.include_router(router)
return app
def test_concurrent_first_requests_do_not_corrupt_candidate_cache() -> None:
app = build_app()
old_interval = sys.getswitchinterval()
sys.setswitchinterval(1e-5)
try:
barrier = threading.Barrier(N_THREADS)
statuses: list[int] = []
def worker() -> None:
client = TestClient(app)
barrier.wait()
statuses.append(client.get(f"/r{N_ROUTES - 1}").status_code)
threads = [threading.Thread(target=worker) for _ in range(N_THREADS)]
for thread in threads:
thread.start()
for thread in threads:
thread.join()
finally:
sys.setswitchinterval(old_interval)
assert statuses == [200] * N_THREADS
included = next(
route for route in app.router.routes if isinstance(route, _IncludedRouter)
)
assert len(included.effective_candidates()) == N_ROUTES
Loading…
Cancel
Save