diff --git a/main.py b/main.py index 976cc3f..c679f18 100644 --- a/main.py +++ b/main.py @@ -18,6 +18,7 @@ from multiprocessing import Value from fastapi.middleware.cors import CORSMiddleware import random import logging +import redis JUNCTION = 0 RESERVOIR = 1 @@ -40,6 +41,10 @@ if not os.path.exists(tmpDir): app = FastAPI() +# 初始化 Redis 连接 +# 用redis 来限制并发访问 +redis_client = redis.Redis(host="localhost", port=6379, db=0) + # 配置日志记录器 logging.basicConfig(level=logging.INFO) logger = logging.getLogger(__name__) @@ -169,17 +174,21 @@ async def fastapi_dump_inp(network: str, inp: str) -> bool: # 必须用这个PlainTextResponse,不然每个key都有引号 @app.get("/runproject/", response_class = PlainTextResponse) async def fastapi_run_project(network: str) -> str: - filename = 'c:/lock.simulation' - filename2 = 'c:/lock.simulation2' - if os.path.exists(filename2): - print('file exists') + lock_key = "exclusive_api_lock" + timeout = 120 # 锁自动过期时间(秒) + + # 尝试获取锁(NX=True: 不存在时设置,EX=timeout: 过期时间) + acquired = redis_client.set(lock_key, "locked", nx=True, ex=timeout) + + if not acquired: raise HTTPException(status_code=409, detail="is in simulation") else: - print('file doesnt exists') - #os.rename(filename, filename2) - result = run_project(network) - #os.rename(filename2, filename) - return result + try: + result = run_project(network) + return result + finally: + # 手动释放锁(可选,依赖过期时间自动释放更安全) + redis_client.delete(lock_key) # put in inp folder, name without extension @app.get("/runinp/")