add lock

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: