From 86242b13d2a7431904c54569a9ff3c0d3dee3397 Mon Sep 17 00:00:00 2001
From: DingZQ <zqding@hotmail.com>
Date: Sat, 21 Oct 2023 00:10:07 +0800
Subject: [PATCH] Refine

---
 main.py | 20 ++++++++++++--------
 1 file changed, 12 insertions(+), 8 deletions(-)

diff --git a/main.py b/main.py
index bed84cb..c7fe776 100644
--- a/main.py
+++ b/main.py
@@ -11,6 +11,7 @@ from fastapi import FastAPI, File, UploadFile, Response, status, Request, Body,
 from fastapi.responses import PlainTextResponse
 from fastapi.middleware.gzip import GZipMiddleware
 from tjnetwork import *
+import asyncio
 
 JUNCTION = 0
 RESERVOIR = 1
@@ -32,7 +33,8 @@ if not os.path.exists(tmpDir):
     os.mkdir(tmpDir)
 
 app = FastAPI()
-app.state.is_simulation = False
+
+lock = asyncio.Lock()
 
 app.add_middleware(GZipMiddleware, minimum_size=1000)
 
@@ -147,14 +149,16 @@ async def fastapi_dump_inp(network: str, inp: str) -> bool:
 
 # 必须用这个PlainTextResponse，不然每个key都有引号
 @app.get("/runproject/", response_class = PlainTextResponse)
-def fastapi_run_project(network: str) -> str:
-    if not app.state.is_simulation:
-        app.state.is_simulation = True
-        result = run_project(network)
-        app.state.is_simulation = False
-        return result
+async def fastapi_run_project(network: str) -> str:
+    if lock.locked():
+        raise HTTPException(status_code=400, detail="is in simulation")
     else:
-        raise HTTPException(status_code=400, detail="Is in Simulation")
+        try:
+            async with lock:
+                result = run_project(network)
+                return result
+        except asyncio.CancelledError:
+            raise HTTPException(status_code=400, detail="is in simulation")
 
 # put in inp folder, name without extension
 @app.get("/runinp/")