diff --git a/backend/monitor.py b/backend/monitor.py index d1fbb8c..34d941a 100644 --- a/backend/monitor.py +++ b/backend/monitor.py @@ -9,13 +9,20 @@ async def check_node_status(node_id: str, node: dict, redis_client): async with httpx.AsyncClient(timeout=3) as client: resp = await client.get(url) if resp.status_code == 200: - data = json.loads(await redis_client.hget("nodes", node_id)) - data["status"] = "online" + node_json = await redis_client.hget("nodes", node_id) + if not node_json: + return + data = json.loads(node_json) + if data.get("status") != "busy": + data["status"] = "online" await redis_client.hset("nodes", node_id, json.dumps(data)) else: raise Exception() except Exception: - data = json.loads(await redis_client.hget("nodes", node_id)) + node_json = await redis_client.hget("nodes", node_id) + if not node_json: + return + data = json.loads(node_json) data["status"] = "offline" await redis_client.hset("nodes", node_id, json.dumps(data)) diff --git a/backend/node_manager.py b/backend/node_manager.py index fa1efaf..3221cb5 100644 --- a/backend/node_manager.py +++ b/backend/node_manager.py @@ -11,6 +11,15 @@ async def register_node(node: Dict): node_id = node.get("id", str(uuid.uuid4())) node["id"] = node_id + + # Ensure max_sessions and active_sessions are tracked numerically + max_sessions = int(node.get("max_sessions", 1)) + active_sessions = int(node.get("active_sessions", 0)) + + node["max_sessions"] = max_sessions + node["active_sessions"] = active_sessions + node.setdefault("status", "online") + await redis_client.hset("nodes", node_id, json.dumps(node)) return {"message": "Node registered", "id": node_id} diff --git a/backend/proxy_router.py b/backend/proxy_router.py index eef1b04..77b19b7 100644 --- a/backend/proxy_router.py +++ b/backend/proxy_router.py @@ -3,28 +3,105 @@ import httpx import redis.asyncio as aioredis import json +from typing import Optional router = APIRouter() redis_client = aioredis.from_url("redis://10.160.13.16:6379/0", decode_responses=True) +SESSION_MAP_KEY = "session_map" + + +def _extract_session_id_from_path(path: str) -> Optional[str]: + parts = [part for part in path.split("/") if part] + if len(parts) >= 2 and parts[0] == "session": + return parts[1] + return None + + +def _extract_session_id_from_response(content: bytes) -> Optional[str]: + try: + payload = json.loads(content.decode()) + except Exception: + return None + + if not isinstance(payload, dict): + return None + + if "sessionId" in payload: + return payload.get("sessionId") + + value = payload.get("value") + if isinstance(value, dict): + return value.get("sessionId") or value.get("session_id") + + return None + + +async def _update_node_session_count(node_id: str, delta: int): + node_json = await redis_client.hget("nodes", node_id) + if not node_json: + return + + node = json.loads(node_json) + active_sessions = int(node.get("active_sessions", 0)) + delta + if active_sessions < 0: + active_sessions = 0 + + node["active_sessions"] = active_sessions + + max_sessions = int(node.get("max_sessions", 1)) + if active_sessions >= max_sessions: + node["status"] = "busy" + elif node.get("status") == "busy": + node["status"] = "online" + + await redis_client.hset("nodes", node_id, json.dumps(node)) + + +async def _cleanup_session(session_id: str, node_id: Optional[str]): + await redis_client.hdel(SESSION_MAP_KEY, session_id) + if node_id: + await _update_node_session_count(node_id, -1) + + async def forward_request(request: Request, path: str): body = await request.body() headers = dict(request.headers) - nodes = await redis_client.hgetall("nodes") - if not nodes: - raise HTTPException(status_code=503, detail="No nodes available") + session_id = _extract_session_id_from_path(path) - # Pick first available online node target_node = None - for node_data in nodes.values(): - node = json.loads(node_data) - if node.get("status") == "online": - target_node = node - break + target_node_id = None - if not target_node: - raise HTTPException(status_code=503, detail="No online nodes available") + if session_id: + target_node_id = await redis_client.hget(SESSION_MAP_KEY, session_id) + if not target_node_id: + raise HTTPException(status_code=404, detail="Session not found") + + node_json = await redis_client.hget("nodes", target_node_id) + if not node_json: + await redis_client.hdel(SESSION_MAP_KEY, session_id) + raise HTTPException(status_code=503, detail="Node unavailable for session") + + target_node = json.loads(node_json) + else: + nodes = await redis_client.hgetall("nodes") + if not nodes: + raise HTTPException(status_code=503, detail="No nodes available") + + for node_id, node_data in nodes.items(): + node = json.loads(node_data) + status = node.get("status") + max_sessions = int(node.get("max_sessions", 1)) + active_sessions = int(node.get("active_sessions", 0)) + + if status == "online" and active_sessions < max_sessions: + target_node = node + target_node_id = node_id + break + + if not target_node: + raise HTTPException(status_code=503, detail="No nodes available for new session") target_url = f"http://{target_node['host']}:{target_node['port']}/wd/hub/{path}" @@ -33,22 +110,58 @@ async def forward_request(request: Request, path: str): request.method, target_url, headers=headers, content=body ) - return resp.content, resp.status_code, resp.headers + return resp.content, resp.status_code, resp.headers, target_node_id or target_node.get("id"), session_id + -# ✅ Add explicit route for session creation @router.api_route("/wd/hub/session", methods=["POST", "OPTIONS"]) async def create_session(request: Request): - content, status, headers = await forward_request(request, "session") + content, status, headers, node_id, _ = await forward_request(request, "session") + + if node_id and 200 <= status < 300: + session_id = _extract_session_id_from_response(content) + if session_id: + await redis_client.hset(SESSION_MAP_KEY, session_id, node_id) + await _update_node_session_count(node_id, 1) + return Response(content=content, status_code=status, headers=dict(headers)) -# ✅ Add route for generic proxy handling (any other Appium path) + @router.api_route("/wd/hub/{path:path}", methods=["GET", "POST", "DELETE", "PUT", "PATCH", "OPTIONS"]) async def proxy_generic(request: Request, path: str): - content, status, headers = await forward_request(request, path) + content, status, headers, node_id, session_id = await forward_request(request, path) + + if ( + request.method == "DELETE" + and session_id + and ( + 200 <= status < 300 + or status == 404 + or status == 410 + ) + ): + await _cleanup_session(session_id, node_id) + return Response(content=content, status_code=status, headers=dict(headers)) -# ✅ Add Selenium-style root route (no wd/hub prefix) + @router.api_route("/session", methods=["POST", "OPTIONS"]) async def selenium_create_session(request: Request): - content, status, headers = await forward_request(request, "session") + content, status, headers, node_id, _ = await forward_request(request, "session") + + if node_id and 200 <= status < 300: + session_id = _extract_session_id_from_response(content) + if session_id: + await redis_client.hset(SESSION_MAP_KEY, session_id, node_id) + await _update_node_session_count(node_id, 1) + return Response(content=content, status_code=status, headers=dict(headers)) + + +@router.delete("/sessions/{session_id}") +async def release_session(session_id: str): + node_id = await redis_client.hget(SESSION_MAP_KEY, session_id) + if not node_id: + raise HTTPException(status_code=404, detail="Session not found") + + await _cleanup_session(session_id, node_id) + return {"message": f"Session {session_id} terminated"}