diff --git a/manga_translator/share.py b/manga_translator/share.py index 2902c0c6..5a2daa3b 100644 --- a/manga_translator/share.py +++ b/manga_translator/share.py @@ -1,5 +1,6 @@ import asyncio import pickle +from threading import Lock import uvicorn from fastapi import FastAPI, HTTPException, Path, Request, Response @@ -19,18 +20,27 @@ def __init__(self, params: dict = None): super().__init__(params) self.host = params.get('host', '127.0.0.1') self.port = int(params.get('port', '5003')) + self.lock = Lock() async def listen(self, translation_params: dict = None): app = FastAPI() + @app.get("/is_locked") + async def is_locked(): + if self.lock.locked(): + return {"locked": True} + return {"locked": False} + @app.post("/execute/{method_name}") async def execute_method(request: Request, method_name: str = Path(...)): - attributes_bytes = await request.body() + if not self.lock.acquire(blocking=False): + raise HTTPException(status_code=429, detail="some Method is already being executed.") if method_name == "listen" or method_name.startswith("__"): raise HTTPException(status_code=403, detail="These functions are not allowed to be executed remotely") method = getattr(self, method_name, None) if not method: raise HTTPException(status_code=404, detail="Method not found") + attributes_bytes = await request.body() attributes = pickle.loads(attributes_bytes) sig = inspect.signature(method) expected_args = set(sig.parameters.keys()) @@ -44,7 +54,6 @@ async def execute_method(request: Request, method_name: str = Path(...)): result = await method(**attributes) else: result = method(**attributes) - print(result) result_bytes = pickle.dumps(result) return Response(content=result_bytes, media_type="application/octet-stream") except Exception as e: