Implement pass-through route registration and removal to prevent duplicates and memory leaks

This commit is contained in:
Jugal Bhatt 2025-08-13 14:56:01 -07:00
parent 758fe45401
commit fd7808e0f3

View file

@ -57,6 +57,9 @@ router = APIRouter()
pass_through_endpoint_logging = PassThroughEndpointLogging()
# Global registry to track registered pass-through routes and prevent memory leaks
_registered_pass_through_routes: Dict[str, Dict[str, str]] = {}
def get_response_body(response: httpx.Response) -> Optional[dict]:
try:
@ -952,8 +955,19 @@ class InitPassThroughEndpointHelpers:
merge_query_params: Optional[bool],
dependencies: Optional[List],
cost_per_request: Optional[float],
endpoint_id: str,
):
"""Add exact path route for pass-through endpoint"""
route_key = f"{endpoint_id}:exact:{path}"
# Check if this exact route is already registered
if route_key in _registered_pass_through_routes:
verbose_proxy_logger.debug(
"Skipping duplicate exact pass through endpoint: %s (already registered)",
path,
)
return
verbose_proxy_logger.debug(
"adding exact pass through endpoint: %s, dependencies: %s",
path,
@ -974,6 +988,13 @@ class InitPassThroughEndpointHelpers:
methods=["GET", "POST", "PUT", "DELETE", "PATCH"],
dependencies=dependencies,
)
# Register the route to prevent duplicates
_registered_pass_through_routes[route_key] = {
"endpoint_id": endpoint_id,
"path": path,
"type": "exact"
}
@staticmethod
def add_subpath_route(
@ -985,9 +1006,20 @@ class InitPassThroughEndpointHelpers:
merge_query_params: Optional[bool],
dependencies: Optional[List],
cost_per_request: Optional[float],
endpoint_id: str,
):
"""Add wildcard route for sub-paths"""
wildcard_path = f"{path}/{{subpath:path}}"
route_key = f"{endpoint_id}:subpath:{path}"
# Check if this subpath route is already registered
if route_key in _registered_pass_through_routes:
verbose_proxy_logger.debug(
"Skipping duplicate wildcard pass through endpoint: %s (already registered)",
wildcard_path,
)
return
verbose_proxy_logger.debug(
"adding wildcard pass through endpoint: %s, dependencies: %s",
wildcard_path,
@ -1009,6 +1041,26 @@ class InitPassThroughEndpointHelpers:
methods=["GET", "POST", "PUT", "DELETE", "PATCH"],
dependencies=dependencies,
)
# Register the route to prevent duplicates
_registered_pass_through_routes[route_key] = {
"endpoint_id": endpoint_id,
"path": path,
"type": "subpath"
}
@staticmethod
def remove_endpoint_routes(endpoint_id: str):
"""Remove all routes for a specific endpoint ID from the registry"""
keys_to_remove = [
key for key, value in _registered_pass_through_routes.items()
if value["endpoint_id"] == endpoint_id
]
for key in keys_to_remove:
del _registered_pass_through_routes[key]
verbose_proxy_logger.debug(
"Removed pass-through route from registry: %s", key
)
async def initialize_pass_through_endpoints(
@ -1072,6 +1124,7 @@ async def initialize_pass_through_endpoints(
merge_query_params=_merge_query_params,
dependencies=_dependencies,
cost_per_request=endpoint.get("cost_per_request", None),
endpoint_id=endpoint.get("id"),
)
# Add wildcard route for sub-paths
@ -1085,6 +1138,7 @@ async def initialize_pass_through_endpoints(
merge_query_params=_merge_query_params,
dependencies=_dependencies,
cost_per_request=endpoint.get("cost_per_request", None),
endpoint_id=endpoint.get("id"),
)
verbose_proxy_logger.debug("Added new pass through endpoint: %s (ID: %s)", _path, endpoint.get("id"))
@ -1229,6 +1283,9 @@ async def update_pass_through_endpoints(
# Update the list
pass_through_endpoint_data[endpoint_index] = endpoint_dict
# Remove old routes from registry before they get re-registered
InitPassThroughEndpointHelpers.remove_endpoint_routes(endpoint_id)
## Update db
updated_data = ConfigFieldUpdate(
field_name="pass_through_endpoints",
@ -1369,6 +1426,9 @@ async def delete_pass_through_endpoints(
pass_through_endpoint_data.pop(endpoint_index)
response_obj = found_endpoint
# Remove routes from registry
InitPassThroughEndpointHelpers.remove_endpoint_routes(endpoint_id)
## Update db
updated_data = ConfigFieldUpdate(
field_name="pass_through_endpoints",