Use redis to restrict corrurent access

This commit is contained in:
DingZQ
2025-01-30 22:54:29 +08:00
parent 97cd59cfae
commit 8e4574b167

27
main.py
View File

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