From ab0e674f8ce6219ad55dec85757139c8caa89eec Mon Sep 17 00:00:00 2001
From: DingZQ <ding@aecforge.com>
Date: Thu, 30 Jan 2025 23:05:02 +0800
Subject: [PATCH] Use redis

---
 main.py | 30 +++++++++++++++++++++---------
 1 file changed, 21 insertions(+), 9 deletions(-)

diff --git a/main.py b/main.py
index 976cc3f..4185d1f 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 限制并发访u
+redis_client = redis.Redis(host="localhost", port=6379, db=0)
+
 # 配置日志记录器
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
@@ -169,17 +174,24 @@ 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')
+    # 初始化 Redis 连接
+    # 用redis 来限制并发访问
+    redis_client = redis.Redis(host="localhost", port=6379, db=0)
+    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/")