From 843a237fad698e48b8cf0285fbc1bce22ca2f643 Mon Sep 17 00:00:00 2001
From: Lan <hhxk666@gmail.com>
Date: Fri, 13 Feb 2026 23:30:48 +0800
Subject: [PATCH] fix: implement async database startup lock and refactor db
 configuration

---
 core/database.py | 101 ++++++++++++++++++++++++++++++++++++++---------
 main.py          |  19 +++------
 2 files changed, 89 insertions(+), 31 deletions(-)

diff --git a/core/database.py b/core/database.py
index 4421bf6..1b04119 100644
--- a/core/database.py
+++ b/core/database.py
@@ -1,6 +1,9 @@
+import asyncio
 import glob
 import importlib
 import os
+from contextlib import asynccontextmanager
+from typing import IO
 
 from tortoise import Tortoise
 
@@ -8,29 +11,91 @@ from core.logger import logger
 from core.settings import data_root
 
 
+_DB_FILE = os.path.join(data_root, "filecodebox.db")
+_STARTUP_LOCK_FILE = os.path.join(data_root, "filecodebox.startup.lock")
+
+
+def get_db_config() -> dict:
+    return {
+        "connections": {
+            "default": {
+                "engine": "tortoise.backends.sqlite",
+                "credentials": {
+                    "file_path": _DB_FILE,
+                    "journal_mode": "WAL",
+                    "busy_timeout": 10000,
+                    "foreign_keys": "ON",
+                },
+            }
+        },
+        "apps": {
+            "models": {
+                "models": ["apps.base.models"],
+                "default_connection": "default",
+            }
+        },
+        "use_tz": False,
+        "timezone": "Asia/Shanghai",
+    }
+
+
+def _lock_file(file_obj: IO[str]) -> None:
+    if os.name == "nt":
+        import msvcrt
+
+        # Windows 需要锁定至少 1 字节
+        if os.fstat(file_obj.fileno()).st_size == 0:
+            file_obj.write("0")
+            file_obj.flush()
+        msvcrt.locking(file_obj.fileno(), msvcrt.LK_LOCK, 1)
+    else:
+        import fcntl
+
+        fcntl.flock(file_obj.fileno(), fcntl.LOCK_EX)
+
+
+def _unlock_file(file_obj: IO[str]) -> None:
+    if os.name == "nt":
+        import msvcrt
+
+        msvcrt.locking(file_obj.fileno(), msvcrt.LK_UNLCK, 1)
+    else:
+        import fcntl
+
+        fcntl.flock(file_obj.fileno(), fcntl.LOCK_UN)
+
+
+@asynccontextmanager
+async def db_startup_lock():
+    os.makedirs(data_root, exist_ok=True)
+    lock_file = open(_STARTUP_LOCK_FILE, "a+", encoding="utf-8")
+    try:
+        await asyncio.to_thread(_lock_file, lock_file)
+        yield
+    finally:
+        await asyncio.to_thread(_unlock_file, lock_file)
+        lock_file.close()
+
+
 async def init_db():
     try:
-        # 使用正确的Tortoise初始化配置格式
-        db_config = {
-            "db_url": f"sqlite://{data_root}/filecodebox.db",
-            "modules": {"models": ["apps.base.models"]},
-            "use_tz": False,
-            "timezone": "Asia/Shanghai"
-        }
+        db_config = get_db_config()
 
-        await Tortoise.init(**db_config)
+        if not Tortoise._inited:
+            await Tortoise.init(config=db_config)
 
-        # 创建migrations表
-        await Tortoise.get_connection("default").execute_script("""
-            CREATE TABLE IF NOT EXISTS migrates (
-                id INTEGER PRIMARY KEY AUTOINCREMENT,
-                migration_file VARCHAR(255) NOT NULL UNIQUE,
-                executed_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
-            )
-        """)
+        async with db_startup_lock():
+            # 创建migrations表
+            await Tortoise.get_connection("default").execute_script("""
+                CREATE TABLE IF NOT EXISTS migrates (
+                    id INTEGER PRIMARY KEY AUTOINCREMENT,
+                    migration_file VARCHAR(255) NOT NULL UNIQUE,
+                    executed_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
+                )
+            """)
 
-        # 执行迁移
-        await execute_migrations()
+            # 执行迁移
+            await execute_migrations()
 
     except Exception as e:
         logger.error(f"数据库初始化失败: {str(e)}")
diff --git a/main.py b/main.py
index 051bb4a..b3ec45f 100644
--- a/main.py
+++ b/main.py
@@ -17,10 +17,10 @@ from apps.admin.views import admin_api
 from apps.base.models import KeyValue
 from apps.base.utils import ip_limit
 from apps.base.views import share_api, chunk_api, presign_api
-from core.database import init_db
+from core.database import db_startup_lock, get_db_config, init_db
 from core.logger import logger
 from core.response import APIResponse
-from core.settings import data_root, settings, BASE_DIR, DEFAULT_CONFIG
+from core.settings import settings, BASE_DIR, DEFAULT_CONFIG
 from core.tasks import delete_expire_files, clean_incomplete_uploads
 from core.utils import hash_password, is_password_hashed
 
@@ -31,8 +31,9 @@ async def lifespan(app: FastAPI):
     # 初始化数据库
     await init_db()
 
-    # 加载配置
-    await load_config()
+    # 加载配置（多进程下串行化启动写操作）
+    async with db_startup_lock():
+        await load_config()
     app.mount(
         "/assets",
         StaticFiles(directory=f"./{settings.themesSelect}/assets"),
@@ -97,15 +98,7 @@ app.add_middleware(
 # 使用 register_tortoise 来添加异常处理器
 register_tortoise(
     app,
-    config={
-        "connections": {"default": f"sqlite://{data_root}/filecodebox.db"},
-        "apps": {
-            "models": {
-                "models": ["apps.base.models"],
-                "default_connection": "default",
-            },
-        },
-    },
+    config=get_db_config(),
     generate_schemas=False,
     add_exception_handlers=True,
 )