Optimize performance with caching and startup improvements

apps/base/dependencies.py

@@ -1,6 +1,7 @@
+import time
+from collections import OrderedDict
 from ipaddress import ip_address, ip_network
-from typing import Dict, Iterable, TypedDict
-from datetime import datetime, timedelta
+from typing import Iterable, TypedDict
 from fastapi import HTTPException, Request
 
 from core.settings import settings
@@ -67,41 +68,66 @@ def get_client_ip(request: Request) -> str:
 
 class _IPRecord(TypedDict):
     count: int
-    time: datetime
+    last_seen: float
+
+
+# Upper bound for tracked source IPs per limiter. Without it a scan using
+# spoofed/rotating addresses grows the dict forever on a small VPS.
+DEFAULT_MAX_TRACKED_IPS = 10_000
 
 
 class IPRateLimit:
-    def __init__(self, count: int, minutes: int):
-        self.ips: Dict[str, _IPRecord] = {}
+    def __init__(
+        self,
+        count: int,
+        minutes: int,
+        max_ips: int = DEFAULT_MAX_TRACKED_IPS,
+    ):
+        # Ordered by last activity (oldest first) so capacity eviction and the
+        # periodic cleanup are both O(1)/O(expired) instead of a full rebuild.
+        self.ips: OrderedDict[str, _IPRecord] = OrderedDict()
         self.count = count
         self.minutes = minutes
+        self.max_ips = max(1, max_ips)
+
+    def _window_seconds(self) -> float:
+        # minutes is mutable at runtime via apps.base.config._sync_ip_limits.
+        return float(self.minutes) * 60
 
     def check_ip(self, ip: str) -> bool:
-        if ip in self.ips:
-            ip_info = self.ips[ip]
-            if ip_info["count"] >= self.count:
-                if ip_info["time"] + timedelta(minutes=self.minutes) > datetime.now():
-                    return False
-                self.ips.pop(ip)
+        record = self.ips.get(ip)
+        if record is None:
+            return True
+        if record["count"] >= self.count:
+            if record["last_seen"] + self._window_seconds() > time.monotonic():
+                return False
+            self.ips.pop(ip, None)
         return True
 
     def add_ip(self, ip: str) -> int:
-        ip_info = self.ips.setdefault(
-            ip, {"count": 0, "time": datetime.now()}
-        )
-        ip_info["count"] += 1
-        ip_info["time"] = datetime.now()
-        self.ips[ip] = ip_info
-        return ip_info["count"]
+        now = time.monotonic()
+        record = self.ips.get(ip)
+        if record is None:
+            record = {"count": 0, "last_seen": now}
+            self.ips[ip] = record
+        else:
+            self.ips.move_to_end(ip)
+        record["count"] += 1
+        record["last_seen"] = now
+        while len(self.ips) > self.max_ips:
+            # Evict the least recently active IP; fail-open for the evicted
+            # address is preferable to unbounded memory growth.
+            self.ips.popitem(last=False)
+        return record["count"]
 
     async def remove_expired_ip(self) -> None:
-        now = datetime.now()
-        expiration = timedelta(minutes=self.minutes)
-        self.ips = {
-            ip: info
-            for ip, info in self.ips.items()
-            if info["time"] + expiration >= now
-        }
+        now = time.monotonic()
+        window = self._window_seconds()
+        while self.ips:
+            _, record = next(iter(self.ips.items()))
+            if record["last_seen"] + window >= now:
+                break
+            self.ips.popitem(last=False)
 
     def __call__(self, request: Request) -> str:
         ip = get_client_ip(request)

apps/base/pages.py

@@ -1,6 +1,9 @@
 """Public-facing pages and routes: setup wizard, theme assets, index,
 robots.txt, and the public config endpoints."""
 import html
+import re
+from functools import lru_cache
+from pathlib import Path
 
 from fastapi import APIRouter, HTTPException, Request
 from fastapi.responses import FileResponse, HTMLResponse, JSONResponse, RedirectResponse
@@ -83,9 +86,58 @@ def resolve_theme_file(*parts: str):
     return file_path
 
 
+# Vite 产物带内容指纹（如 index-BX2s1QqP.js），内容变化时文件名必然变化，
+# 可以安全地长缓存；public/ 下的未哈希资源原地升级时名字不变，只给短缓存，
+# 过期后靠 FileResponse 自带的 ETag/Last-Modified 走 304。
+_HASHED_ASSET_RE = re.compile(
+    r"-[A-Za-z0-9_-]{8,}\.(?:js|mjs|css|map|woff2?|ttf|otf|eot|svg|png|jpe?g|gif|webp|avif|ico)$",
+    re.IGNORECASE,
+)
+_IMMUTABLE_ASSET_CACHE = "public, max-age=31536000, immutable"
+_DEFAULT_ASSET_CACHE = "public, max-age=86400"
+
+
+def _asset_cache_control(asset_path: str) -> str:
+    if _HASHED_ASSET_RE.search(asset_path):
+        return _IMMUTABLE_ASSET_CACHE
+    return _DEFAULT_ASSET_CACHE
+
+
+@lru_cache(maxsize=8)
+def _render_index_html(
+    template_path: str,
+    inode: int,
+    mtime_ns: int,
+    size: int,
+    title: str,
+    description: str,
+    keywords: str,
+    opacity: str,
+    background: str,
+) -> str:
+    """Render index.html once per (theme file version, site config) pair.
+
+    The stat fields keep the cache honest when a theme is replaced in place;
+    config values are part of the key so admin edits are visible immediately
+    without any cross-module invalidation hook.
+    """
+    return (
+        Path(template_path)
+        .read_text(encoding="utf-8")
+        .replace("{{title}}", html.escape(title))
+        .replace("{{description}}", html.escape(description))
+        .replace("{{keywords}}", html.escape(keywords))
+        .replace("{{opacity}}", html.escape(opacity))
+        .replace("{{background}}", html.escape(background))
+    )
+
+
 @router.get("/assets/{asset_path:path}", include_in_schema=False)
 async def theme_asset(asset_path: str):
-    return FileResponse(resolve_theme_file("assets", asset_path))
+    return FileResponse(
+        resolve_theme_file("assets", asset_path),
+        headers={"Cache-Control": _asset_cache_control(asset_path)},
+    )
 
 
 async def not_found_handler(request, exc=None):
@@ -115,14 +167,26 @@ async def index(request=None, exc=None):
     # Site config is admin input (and during the setup window anyone can claim it);
     # always escape before injecting into the theme template to prevent stored XSS
     # (mirrors the setup page).
+    template_path = resolve_theme_file("index.html")
+    try:
+        stat_result = template_path.stat()
+    except OSError:
+        # resolve_theme_file 与 stat 之间文件被删除/替换时仍按 404 处理，
+        # 不能退回缓存内容。
+        raise HTTPException(status_code=404, detail="资源不存在")
+    content = _render_index_html(
+        str(template_path),
+        stat_result.st_ino,
+        stat_result.st_mtime_ns,
+        stat_result.st_size,
+        str(settings.name),
+        str(settings.description),
+        str(settings.keywords),
+        str(settings.opacity),
+        str(settings.background),
+    )
     return HTMLResponse(
-        content=resolve_theme_file("index.html")
-        .read_text(encoding="utf-8")
-        .replace("{{title}}", html.escape(str(settings.name)))
-        .replace("{{description}}", html.escape(str(settings.description)))
-        .replace("{{keywords}}", html.escape(str(settings.keywords)))
-        .replace("{{opacity}}", html.escape(str(settings.opacity)))
-        .replace("{{background}}", html.escape(str(settings.background))),
+        content=content,
         media_type="text/html",
         headers={"Cache-Control": "no-cache"},
     )

main.py

@@ -10,7 +10,7 @@
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi.responses import JSONResponse
 from tortoise import Tortoise
-from tortoise.contrib.fastapi import register_tortoise
+from tortoise.contrib.fastapi import tortoise_exception_handlers
 
 from apps.admin.views import admin_api
 from apps.base.config import (
@@ -29,7 +29,7 @@
 )
 from apps.base.views import share_api, chunk_api, presign_api
 from apps.base.mcp import router as mcp_router
-from core.database import db_startup_lock, get_db_config, init_db
+from core.database import db_startup_lock, init_db
 from core.errors import StorageError
 from core.logger import get_log_level_name, is_access_log_enabled, logger
 from core.settings import settings
@@ -89,7 +89,14 @@ async def load_config():
         )
 
 
-app = FastAPI(lifespan=lifespan, version=APP_VERSION)
+# Tortoise 的 ORM 异常映射（DoesNotExist → 404、IntegrityError → 422）。
+# 只注册处理器；连接初始化统一由 lifespan 里的 init_db() 负责，
+# 避免 register_tortoise 合并 lifespan 造成的双重初始化。
+app = FastAPI(
+    lifespan=lifespan,
+    version=APP_VERSION,
+    exception_handlers=tortoise_exception_handlers(),
+)
 
 
 @app.exception_handler(StorageError)
@@ -126,14 +133,6 @@ async def refresh_settings_middleware(request, call_next):
     allow_headers=["*"],
 )
 
-# 使用 register_tortoise 来添加异常处理器
-register_tortoise(
-    app,
-    config=get_db_config(),
-    generate_schemas=False,
-    add_exception_handlers=True,
-)
-
 app.include_router(share_api)
 app.include_router(chunk_api)
 app.include_router(presign_api)

tests/test_index_cache.py

@@ -0,0 +1,82 @@
+"""index.html render cache must stay correct across config/theme changes."""
+import asyncio
+import os
+import tempfile
+import unittest
+from pathlib import Path
+
+from fastapi import HTTPException
+
+import apps.base.pages as pages
+from core.settings import settings
+
+
+class IndexRenderCacheTests(unittest.TestCase):
+    def setUp(self):
+        self._original_user_config = dict(settings.user_config)
+        self._original_resolve_theme_file = pages.resolve_theme_file
+        self._tmpdir = tempfile.TemporaryDirectory()
+        self.template = Path(self._tmpdir.name) / "index.html"
+        self.template.write_text(
+            "<title>{{title}}</title>"
+            '<meta name="description" content="{{description}}">'
+            "<div>{{opacity}}</div>",
+            encoding="utf-8",
+        )
+        pages.resolve_theme_file = lambda *args, **kwargs: self.template
+        pages._render_index_html.cache_clear()
+
+    def tearDown(self):
+        pages.resolve_theme_file = self._original_resolve_theme_file
+        settings.user_config = self._original_user_config
+        pages._render_index_html.cache_clear()
+        self._tmpdir.cleanup()
+
+    def _render(self) -> str:
+        return asyncio.run(pages.index()).body.decode("utf-8")
+
+    def test_repeated_requests_hit_the_render_cache(self):
+        first = self._render()
+        second = self._render()
+
+        self.assertEqual(first, second)
+        info = pages._render_index_html.cache_info()
+        self.assertEqual(info.misses, 1)
+        self.assertEqual(info.hits, 1)
+
+    def test_site_config_change_invalidates_cache(self):
+        self._render()
+
+        settings.name = "New Site Name"
+        html = self._render()
+
+        self.assertIn("<title>New Site Name</title>", html)
+        self.assertEqual(pages._render_index_html.cache_info().misses, 2)
+
+    def test_template_replacement_invalidates_cache(self):
+        self._render()
+        old_stat = self.template.stat()
+
+        self.template.write_text(
+            "<title>{{title}}</title><p>rebuilt</p>", encoding="utf-8"
+        )
+        os.utime(
+            self.template,
+            ns=(old_stat.st_atime_ns, old_stat.st_mtime_ns + 1_000_000),
+        )
+
+        html = self._render()
+        self.assertIn("<p>rebuilt</p>", html)
+        self.assertEqual(pages._render_index_html.cache_info().misses, 2)
+
+    def test_missing_template_still_404s_instead_of_serving_cache(self):
+        self._render()
+        self.template.unlink()
+
+        with self.assertRaises(HTTPException) as ctx:
+            asyncio.run(pages.index())
+        self.assertEqual(ctx.exception.status_code, 404)
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/test_index_template_escaping.py

@@ -18,6 +18,7 @@ class IndexTemplateEscapingTests(unittest.TestCase):
     def setUp(self):
         self._original_user_config = dict(settings.user_config)
         self._original_resolve_theme_file = pages.resolve_theme_file
+        pages._render_index_html.cache_clear()
         tmp = tempfile.NamedTemporaryFile(
             "w", suffix=".html", delete=False, encoding="utf-8"
         )
@@ -32,6 +33,7 @@ def setUp(self):
     def tearDown(self):
         settings.user_config = self._original_user_config
         pages.resolve_theme_file = self._original_resolve_theme_file
+        pages._render_index_html.cache_clear()
         self._template_path.unlink()
 
     def _patch_template(self):