Adding semaphore usage on the '/run' endpoint (#8526)

### What problem does this PR solve?

Switching threading.Lock() to asyncio.Lock(), since threading.Lock() is
blocking.

### Type of change

- [x] Performance Improvement

sandbox/executor_manager/api/handlers.py

@@ -22,23 +22,23 @@ from models.schemas import CodeExecutionRequest, CodeExecutionResult
 from services.execution import execute_code
 from services.limiter import limiter
 from services.security import analyze_code_security
-
+from core.container import _CONTAINER_EXECUTION_SEMAPHORES
 
 async def healthz_handler():
     return {"status": "ok"}
 
-
 @limiter.limit("5/second")
 async def run_code_handler(req: CodeExecutionRequest, request: Request):
     logger.info("🟢 Received /run request")
 
-    code = base64.b64decode(req.code_b64).decode("utf-8")
-    is_safe, issues = analyze_code_security(code, language=req.language)
-    if not is_safe:
-        issue_details = "\n".join([f"Line {lineno}: {issue}" for issue, lineno in issues])
-        return CodeExecutionResult(status=ResultStatus.PROGRAM_RUNNER_ERROR, stdout="", stderr=issue_details, exit_code=-999, detail="Code is unsafe")
+    async with _CONTAINER_EXECUTION_SEMAPHORES[req.language]:
+        code = base64.b64decode(req.code_b64).decode("utf-8")
+        is_safe, issues = analyze_code_security(code, language=req.language)
+        if not is_safe:
+            issue_details = "\n".join([f"Line {lineno}: {issue}" for issue, lineno in issues])
+            return CodeExecutionResult(status=ResultStatus.PROGRAM_RUNNER_ERROR, stdout="", stderr=issue_details, exit_code=-999, detail="Code is unsafe")
 
-    try:
-        return await execute_code(req)
-    except Exception as e:
-        return CodeExecutionResult(status=ResultStatus.PROGRAM_RUNNER_ERROR, stdout="", stderr=str(e), exit_code=-999, detail="unhandled_exception")
+        try:
+            return await execute_code(req)
+        except Exception as e:
+            return CodeExecutionResult(status=ResultStatus.PROGRAM_RUNNER_ERROR, stdout="", stderr=str(e), exit_code=-999, detail="unhandled_exception")

sandbox/executor_manager/api/routes.py

@@ -20,4 +20,4 @@ from api.handlers import healthz_handler, run_code_handler
 router = APIRouter()
 
 router.get("/healthz")(healthz_handler)
-router.post("/run")(run_code_handler)
+router.post("/run")(run_code_handler)

sandbox/executor_manager/core/container.py

@@ -17,7 +17,6 @@ import asyncio
 import contextlib
 import os
 from queue import Empty, Queue
-from threading import Lock
 
 from models.enums import SupportLanguage
 from util import env_setting_enabled, is_valid_memory_limit
@@ -26,18 +25,22 @@ from utils.common import async_run_command
 from core.logger import logger
 
 _CONTAINER_QUEUES: dict[SupportLanguage, Queue] = {}
-_CONTAINER_LOCK: Lock = Lock()
+_CONTAINER_LOCK: asyncio.Lock = asyncio.Lock()
+_CONTAINER_EXECUTION_SEMAPHORES:dict[SupportLanguage,asyncio.Semaphore] = {}
 
 
 async def init_containers(size: int) -> tuple[int, int]:
     global _CONTAINER_QUEUES
     _CONTAINER_QUEUES = {SupportLanguage.PYTHON: Queue(), SupportLanguage.NODEJS: Queue()}
 
-    with _CONTAINER_LOCK:
+    async with _CONTAINER_LOCK:
         while not _CONTAINER_QUEUES[SupportLanguage.PYTHON].empty():
             _CONTAINER_QUEUES[SupportLanguage.PYTHON].get_nowait()
         while not _CONTAINER_QUEUES[SupportLanguage.NODEJS].empty():
             _CONTAINER_QUEUES[SupportLanguage.NODEJS].get_nowait()
+    
+    for language in SupportLanguage:
+        _CONTAINER_EXECUTION_SEMAPHORES[language] = asyncio.Semaphore(size)
 
     create_tasks = []
     for i in range(size):
@@ -56,7 +59,7 @@ async def init_containers(size: int) -> tuple[int, int]:
 
 
 async def teardown_containers():
-    with _CONTAINER_LOCK:
+    async with _CONTAINER_LOCK:
         while not _CONTAINER_QUEUES[SupportLanguage.PYTHON].empty():
             name = _CONTAINER_QUEUES[SupportLanguage.PYTHON].get_nowait()
             await async_run_command("docker", "rm", "-f", name, timeout=5)
@@ -151,7 +154,7 @@ async def recreate_container(name: str, language: SupportLanguage) -> bool:
 
 async def release_container(name: str, language: SupportLanguage):
     """Asynchronously release a container"""
-    with _CONTAINER_LOCK:
+    async with _CONTAINER_LOCK:
         if await container_is_running(name):
             _CONTAINER_QUEUES[language].put(name)
             logger.info(f"🟢 Released container: {name} (remaining available: {_CONTAINER_QUEUES[language].qsize()})")
@@ -168,7 +171,7 @@ async def allocate_container_blocking(language: SupportLanguage, timeout=10) ->
     while asyncio.get_running_loop().time() - start_time < timeout:
         try:
             name = _CONTAINER_QUEUES[language].get_nowait()
-            with _CONTAINER_LOCK:
+            async with _CONTAINER_LOCK:
                 if not await container_is_running(name) and not await recreate_container(name, language):
                     continue
 

sandbox/executor_manager/requirements.txt

@@ -1,3 +1,3 @@
 fastapi
 uvicorn
-slowapi
+slowapi