Use redis to restrict corrurent access
This commit is contained in:
27
main.py
27
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/")
|
||||
|
||||
Reference in New Issue
Block a user