From e2ebd42a8cf830b157450604d188aad244cd8ada Mon Sep 17 00:00:00 2001 From: CaIon Date: Wed, 4 Feb 2026 17:36:07 +0800 Subject: [PATCH] feat(cache): enhance disk cache management with concurrency control and cleanup optimizations --- common/disk_cache.go | 6 +- common/disk_cache_config.go | 8 +- controller/performance.go | 25 ++--- service/file_service.go | 106 +++++++++++------- service/token_counter.go | 10 +- types/file_source.go | 48 +++++--- web/src/i18n/locales/zh.json | 3 + .../Performance/SettingsPerformance.jsx | 6 +- 8 files changed, 128 insertions(+), 84 deletions(-) diff --git a/common/disk_cache.go b/common/disk_cache.go index b41fdcb6a..bea3de044 100644 --- a/common/disk_cache.go +++ b/common/disk_cache.go @@ -127,7 +127,11 @@ func CleanupOldDiskCacheFiles(maxAge time.Duration) error { continue } if now.Sub(info.ModTime()) > maxAge { - os.Remove(filepath.Join(dir, entry.Name())) + // 注意：后台清理任务删除文件时，由于无法得知原始 base64Size， + // 只能按磁盘文件大小扣减。这在目前 base64 存储模式下是准确的。 + if err := os.Remove(filepath.Join(dir, entry.Name())); err == nil { + DecrementDiskFiles(info.Size()) + } } } return nil diff --git a/common/disk_cache_config.go b/common/disk_cache_config.go index ea0b1e1ad..b629c9ce7 100644 --- a/common/disk_cache_config.go +++ b/common/disk_cache_config.go @@ -113,8 +113,12 @@ func IncrementDiskFiles(size int64) { // DecrementDiskFiles 减少磁盘文件计数 func DecrementDiskFiles(size int64) { - atomic.AddInt64(&diskCacheStats.ActiveDiskFiles, -1) - atomic.AddInt64(&diskCacheStats.CurrentDiskUsageBytes, -size) + if atomic.AddInt64(&diskCacheStats.ActiveDiskFiles, -1) < 0 { + atomic.StoreInt64(&diskCacheStats.ActiveDiskFiles, 0) + } + if atomic.AddInt64(&diskCacheStats.CurrentDiskUsageBytes, -size) < 0 { + atomic.StoreInt64(&diskCacheStats.CurrentDiskUsageBytes, 0) + } } // IncrementMemoryBuffers 增加内存缓存计数 diff --git a/controller/performance.go b/controller/performance.go index 8a261ad82..44887c2db 100644 --- a/controller/performance.go +++ b/controller/performance.go @@ -4,6 +4,7 @@ import ( "net/http" "os" "runtime" + "time" "github.com/QuantumNous/new-api/common" "github.com/gin-gonic/gin" @@ -77,10 +78,8 @@ type PerformanceConfig struct { // GetPerformanceStats 获取性能统计信息 func GetPerformanceStats(c *gin.Context) { - // 先同步磁盘缓存统计，确保显示准确 - common.SyncDiskCacheStats() - - // 获取缓存统计 + // 不再每次获取统计都全量扫描磁盘，依赖原子计数器保证性能 + // 仅在系统启动或显式清理时同步 cacheStats := common.GetDiskCacheStats() // 获取内存统计 @@ -123,25 +122,19 @@ func GetPerformanceStats(c *gin.Context) { }) } -// ClearDiskCache 清理磁盘缓存 +// ClearDiskCache 清理不活跃的磁盘缓存 func ClearDiskCache(c *gin.Context) { - // 使用统一的缓存目录 - dir := common.GetDiskCacheDir() - - // 删除缓存目录 - err := os.RemoveAll(dir) - if err != nil && !os.IsNotExist(err) { + // 清理超过 10 分钟未使用的缓存文件 + // 10 分钟是一个安全的阈值，确保正在进行的请求不会被误删 + err := common.CleanupOldDiskCacheFiles(10 * time.Minute) + if err != nil { common.ApiError(c, err) return } - // 重置统计（包括命中次数和使用量） - common.ResetDiskCacheStats() - common.ResetDiskCacheUsage() - c.JSON(http.StatusOK, gin.H{ "success": true, - "message": "磁盘缓存已清理", + "message": "不活跃的磁盘缓存已清理", }) } diff --git a/service/file_service.go b/service/file_service.go index a42a42bf2..c592aa475 100644 --- a/service/file_service.go +++ b/service/file_service.go @@ -36,11 +36,44 @@ func LoadFileSource(c *gin.Context, source *types.FileSource, reason ...string) return nil, fmt.Errorf("file source is nil") } - // 如果已有缓存，直接返回 + if common.DebugEnabled { + logger.LogDebug(c, fmt.Sprintf("LoadFileSource starting for: %s", source.GetIdentifier())) + } + + // 1. 快速检查内部缓存 if source.HasCache() { + // 即使命中内部缓存，也要确保注册到清理列表（如果尚未注册） + if c != nil { + registerSourceForCleanup(c, source) + } return source.GetCache(), nil } + // 2. 加锁保护加载过程 + source.Mu().Lock() + defer source.Mu().Unlock() + + // 3. 双重检查 + if source.HasCache() { + if c != nil { + registerSourceForCleanup(c, source) + } + return source.GetCache(), nil + } + + // 4. 如果是 URL，检查 Context 缓存 + var contextKey string + if source.IsURL() && c != nil { + contextKey = getContextCacheKey(source.URL) + if cachedData, exists := c.Get(contextKey); exists { + data := cachedData.(*types.CachedFileData) + source.SetCache(data) + registerSourceForCleanup(c, source) + return data, nil + } + } + + // 5. 执行加载逻辑 var cachedData *types.CachedFileData var err error @@ -54,10 +87,13 @@ func LoadFileSource(c *gin.Context, source *types.FileSource, reason ...string) return nil, err } - // 设置缓存 + // 6. 设置缓存 source.SetCache(cachedData) + if contextKey != "" && c != nil { + c.Set(contextKey, cachedData) + } - // 注册到 context 以便请求结束时自动清理 + // 7. 注册到 context 以便请求结束时自动清理 if c != nil { registerSourceForCleanup(c, source) } @@ -67,6 +103,10 @@ func LoadFileSource(c *gin.Context, source *types.FileSource, reason ...string) // registerSourceForCleanup 注册 FileSource 到 context 以便请求结束时清理 func registerSourceForCleanup(c *gin.Context, source *types.FileSource) { + if source.IsRegistered() { + return + } + key := string(constant.ContextKeyFileSourcesToCleanup) var sources []*types.FileSource if existing, exists := c.Get(key); exists { @@ -74,6 +114,7 @@ func registerSourceForCleanup(c *gin.Context, source *types.FileSource) { } sources = append(sources, source) c.Set(key, sources) + source.SetRegistered(true) } // CleanupFileSources 清理请求中所有注册的 FileSource @@ -83,9 +124,6 @@ func CleanupFileSources(c *gin.Context) { if sources, exists := c.Get(key); exists { for _, source := range sources.([]*types.FileSource) { if cache := source.GetCache(); cache != nil { - if cache.IsDisk() { - common.DecrementDiskFiles(cache.Size) - } cache.Close() } } @@ -94,21 +132,13 @@ func CleanupFileSources(c *gin.Context) { } // loadFromURL 从 URL 加载文件 -// 支持磁盘缓存：当文件大小超过阈值且磁盘缓存可用时，将数据存储到磁盘 func loadFromURL(c *gin.Context, url string, reason ...string) (*types.CachedFileData, error) { - contextKey := getContextCacheKey(url) - - // 检查 context 缓存 - if cachedData, exists := c.Get(contextKey); exists { - if common.DebugEnabled { - logger.LogDebug(c, fmt.Sprintf("Using cached file data for URL: %s", url)) - } - return cachedData.(*types.CachedFileData), nil - } - // 下载文件 var maxFileSize = constant.MaxFileDownloadMB * 1024 * 1024 + if common.DebugEnabled { + logger.LogDebug(c, "loadFromURL: initiating download") + } resp, err := DoDownloadRequest(url, reason...) if err != nil { return nil, fmt.Errorf("failed to download file from %s: %w", url, err) @@ -120,6 +150,9 @@ func loadFromURL(c *gin.Context, url string, reason ...string) (*types.CachedFil } // 读取文件内容（限制大小） + if common.DebugEnabled { + logger.LogDebug(c, "loadFromURL: reading response body") + } fileBytes, err := io.ReadAll(io.LimitReader(resp.Body, int64(maxFileSize+1))) if err != nil { return nil, fmt.Errorf("failed to read file content: %w", err) @@ -147,6 +180,10 @@ func loadFromURL(c *gin.Context, url string, reason ...string) (*types.CachedFil cachedData = types.NewMemoryCachedData(base64Data, mimeType, int64(len(fileBytes))) } else { cachedData = types.NewDiskCachedData(diskPath, mimeType, int64(len(fileBytes))) + cachedData.DiskSize = base64Size + cachedData.OnClose = func(size int64) { + common.DecrementDiskFiles(size) + } common.IncrementDiskFiles(base64Size) if common.DebugEnabled { logger.LogDebug(c, fmt.Sprintf("File cached to disk: %s, size: %d bytes", diskPath, base64Size)) @@ -159,6 +196,9 @@ func loadFromURL(c *gin.Context, url string, reason ...string) (*types.CachedFil // 如果是图片，尝试获取图片配置 if strings.HasPrefix(mimeType, "image/") { + if common.DebugEnabled { + logger.LogDebug(c, "loadFromURL: decoding image config") + } config, format, err := decodeImageConfig(fileBytes) if err == nil { cachedData.ImageConfig = &config @@ -170,9 +210,6 @@ func loadFromURL(c *gin.Context, url string, reason ...string) (*types.CachedFil } } - // 存入 context 缓存 - c.Set(contextKey, cachedData) - return cachedData, nil } @@ -187,7 +224,6 @@ func writeToDiskCache(base64Data string) (string, error) { } // smartDetectMimeType 智能检测 MIME 类型 -// 优先级：Content-Type header > Content-Disposition filename > URL 路径 > 内容嗅探 > 图片解码 func smartDetectMimeType(resp *http.Response, url string, fileBytes []byte) string { // 1. 尝试从 Content-Type header 获取 mimeType := resp.Header.Get("Content-Type") @@ -259,13 +295,11 @@ func loadFromBase64(base64String string, providedMimeType string) (*types.Cached // 处理 data: 前缀 if strings.HasPrefix(base64String, "data:") { - // 格式: data:mime/type;base64,xxxxx idx := strings.Index(base64String, ",") if idx != -1 { header := base64String[:idx] cleanBase64 = base64String[idx+1:] - // 从 header 提取 MIME 类型 if strings.Contains(header, ":") && strings.Contains(header, ";") { mimeStart := strings.Index(header, ":") + 1 mimeEnd := strings.Index(header, ";") @@ -280,36 +314,34 @@ func loadFromBase64(base64String string, providedMimeType string) (*types.Cached cleanBase64 = base64String } - // 使用提供的 MIME 类型（如果有） if providedMimeType != "" { mimeType = providedMimeType } - // 解码 base64 decodedData, err := base64.StdEncoding.DecodeString(cleanBase64) if err != nil { return nil, fmt.Errorf("failed to decode base64 data: %w", err) } - // 判断是否使用磁盘缓存（对于 base64 内联数据也支持磁盘缓存） base64Size := int64(len(cleanBase64)) var cachedData *types.CachedFileData if shouldUseDiskCache(base64Size) { - // 使用磁盘缓存 diskPath, err := writeToDiskCache(cleanBase64) if err != nil { - // 磁盘缓存失败，回退到内存 cachedData = types.NewMemoryCachedData(cleanBase64, mimeType, int64(len(decodedData))) } else { cachedData = types.NewDiskCachedData(diskPath, mimeType, int64(len(decodedData))) + cachedData.DiskSize = base64Size + cachedData.OnClose = func(size int64) { + common.DecrementDiskFiles(size) + } common.IncrementDiskFiles(base64Size) } } else { cachedData = types.NewMemoryCachedData(cleanBase64, mimeType, int64(len(decodedData))) } - // 如果是图片或 MIME 类型未知，尝试解码图片获取更多信息 if mimeType == "" || strings.HasPrefix(mimeType, "image/") { config, format, err := decodeImageConfig(decodedData) if err == nil { @@ -324,8 +356,7 @@ func loadFromBase64(base64String string, providedMimeType string) (*types.Cached return cachedData, nil } -// GetImageConfig 获取图片配置（宽高等信息） -// 会自动处理缓存，避免重复下载/解码 +// GetImageConfig 获取图片配置 func GetImageConfig(c *gin.Context, source *types.FileSource) (image.Config, string, error) { cachedData, err := LoadFileSource(c, source, "get_image_config") if err != nil { @@ -336,7 +367,6 @@ func GetImageConfig(c *gin.Context, source *types.FileSource) (image.Config, str return *cachedData.ImageConfig, cachedData.ImageFormat, nil } - // 如果缓存中没有图片配置，尝试解码 base64Str, err := cachedData.GetBase64Data() if err != nil { return image.Config{}, "", fmt.Errorf("failed to get base64 data: %w", err) @@ -351,7 +381,6 @@ func GetImageConfig(c *gin.Context, source *types.FileSource) (image.Config, str return image.Config{}, "", err } - // 更新缓存 cachedData.ImageConfig = &config cachedData.ImageFormat = format @@ -359,8 +388,6 @@ func GetImageConfig(c *gin.Context, source *types.FileSource) (image.Config, str } // GetBase64Data 获取 base64 编码的数据 -// 会自动处理缓存，避免重复下载 -// 支持内存缓存和磁盘缓存 func GetBase64Data(c *gin.Context, source *types.FileSource, reason ...string) (string, string, error) { cachedData, err := LoadFileSource(c, source, reason...) if err != nil { @@ -375,12 +402,10 @@ func GetBase64Data(c *gin.Context, source *types.FileSource, reason ...string) ( // GetMimeType 获取文件的 MIME 类型 func GetMimeType(c *gin.Context, source *types.FileSource) (string, error) { - // 如果已经有缓存，直接返回 if source.HasCache() { return source.GetCache().MimeType, nil } - // 如果是 URL，尝试只获取 header 而不下载完整文件 if source.IsURL() { mimeType, err := GetFileTypeFromUrl(c, source.URL, "get_mime_type") if err == nil && mimeType != "" && mimeType != "application/octet-stream" { @@ -388,7 +413,6 @@ func GetMimeType(c *gin.Context, source *types.FileSource) (string, error) { } } - // 否则加载完整数据 cachedData, err := LoadFileSource(c, source, "get_mime_type") if err != nil { return "", err @@ -396,7 +420,7 @@ func GetMimeType(c *gin.Context, source *types.FileSource) (string, error) { return cachedData.MimeType, nil } -// DetectFileType 检测文件类型（image/audio/video/file） +// DetectFileType 检测文件类型 func DetectFileType(mimeType string) types.FileType { if strings.HasPrefix(mimeType, "image/") { return types.FileTypeImage @@ -414,13 +438,11 @@ func DetectFileType(mimeType string) types.FileType { func decodeImageConfig(data []byte) (image.Config, string, error) { reader := bytes.NewReader(data) - // 尝试标准格式 config, format, err := image.DecodeConfig(reader) if err == nil { return config, format, nil } - // 尝试 webp reader.Seek(0, io.SeekStart) config, err = webp.DecodeConfig(reader) if err == nil { @@ -432,13 +454,11 @@ func decodeImageConfig(data []byte) (image.Config, string, error) { // guessMimeTypeFromURL 从 URL 猜测 MIME 类型 func guessMimeTypeFromURL(url string) string { - // 移除查询参数 cleanedURL := url if q := strings.Index(cleanedURL, "?"); q != -1 { cleanedURL = cleanedURL[:q] } - // 获取最后一段 if slash := strings.LastIndex(cleanedURL, "/"); slash != -1 && slash+1 < len(cleanedURL) { last := cleanedURL[slash+1:] if dot := strings.LastIndex(last, "."); dot != -1 && dot+1 < len(last) { diff --git a/service/token_counter.go b/service/token_counter.go index 2020845e3..7d648d77c 100644 --- a/service/token_counter.go +++ b/service/token_counter.go @@ -258,16 +258,18 @@ func EstimateRequestToken(c *gin.Context, meta *types.TokenCountMeta, info *rela // 如果文件类型未知且需要获取，通过 MIME 类型检测 if file.FileType == "" || (file.Source.IsURL() && shouldFetchFiles) { - mimeType, err := GetMimeType(c, file.Source) + // 注意：这里我们直接调用 LoadFileSource 而不是 GetMimeType + // 因为 GetMimeType 内部可能会调用 GetFileTypeFromUrl (HEAD 请求) + // 而我们这里既然要计算 token，通常需要完整数据 + cachedData, err := LoadFileSource(c, file.Source, "token_counter") if err != nil { if shouldFetchFiles { return 0, fmt.Errorf("error getting file type: %v", err) } - // 如果不需要获取，使用默认类型 continue } - file.MimeType = mimeType - file.FileType = DetectFileType(mimeType) + file.MimeType = cachedData.MimeType + file.FileType = DetectFileType(cachedData.MimeType) } } diff --git a/types/file_source.go b/types/file_source.go index d2d217fd4..c52062d78 100644 --- a/types/file_source.go +++ b/types/file_source.go @@ -25,8 +25,14 @@ type FileSource struct { // 内部缓存（不导出，不序列化） cachedData *CachedFileData - cacheMu sync.RWMutex cacheLoaded bool + registered bool // 是否已注册到清理列表 + mu sync.Mutex // 保护加载过程 +} + +// Mu 获取内部锁 +func (f *FileSource) Mu() *sync.Mutex { + return &f.mu } // CachedFileData 缓存的文件数据 @@ -35,14 +41,19 @@ type CachedFileData struct { base64Data string // 内存中的 base64 数据（小文件） MimeType string // MIME 类型 Size int64 // 文件大小（字节） + DiskSize int64 // 磁盘缓存实际占用大小（字节，通常是 base64 长度） ImageConfig *image.Config // 图片配置（如果是图片） ImageFormat string // 图片格式（如果是图片） // 磁盘缓存相关 - diskPath string // 磁盘缓存文件路径（大文件） - isDisk bool // 是否使用磁盘缓存 - diskMu sync.Mutex // 磁盘操作锁 - diskClosed bool // 是否已关闭/清理 + diskPath string // 磁盘缓存文件路径（大文件） + isDisk bool // 是否使用磁盘缓存 + diskMu sync.Mutex // 磁盘操作锁（保护磁盘文件的读取和删除） + diskClosed bool // 是否已关闭/清理 + statDecremented bool // 是否已扣减统计 + + // 统计回调，避免循环依赖 + OnClose func(size int64) } // NewMemoryCachedData 创建内存缓存的数据 @@ -114,7 +125,13 @@ func (c *CachedFileData) Close() error { c.diskClosed = true if c.diskPath != "" { - return os.Remove(c.diskPath) + err := os.Remove(c.diskPath) + // 只有在删除成功且未扣减过统计时，才执行回调 + if err == nil && !c.statDecremented && c.OnClose != nil { + c.OnClose(c.DiskSize) + c.statDecremented = true + } + return err } return nil } @@ -170,31 +187,32 @@ func (f *FileSource) GetRawData() string { // SetCache 设置缓存数据 func (f *FileSource) SetCache(data *CachedFileData) { - f.cacheMu.Lock() - defer f.cacheMu.Unlock() f.cachedData = data f.cacheLoaded = true } +// IsRegistered 是否已注册到清理列表 +func (f *FileSource) IsRegistered() bool { + return f.registered +} + +// SetRegistered 设置注册状态 +func (f *FileSource) SetRegistered(registered bool) { + f.registered = registered +} + // GetCache 获取缓存数据 func (f *FileSource) GetCache() *CachedFileData { - f.cacheMu.RLock() - defer f.cacheMu.RUnlock() return f.cachedData } // HasCache 是否有缓存 func (f *FileSource) HasCache() bool { - f.cacheMu.RLock() - defer f.cacheMu.RUnlock() return f.cacheLoaded && f.cachedData != nil } // ClearCache 清除缓存，释放内存和磁盘文件 func (f *FileSource) ClearCache() { - f.cacheMu.Lock() - defer f.cacheMu.Unlock() - // 如果有缓存数据，先关闭它（会清理磁盘文件） if f.cachedData != nil { f.cachedData.Close() diff --git a/web/src/i18n/locales/zh.json b/web/src/i18n/locales/zh.json index 7ff6c35a0..ed4f921e9 100644 --- a/web/src/i18n/locales/zh.json +++ b/web/src/i18n/locales/zh.json @@ -442,6 +442,9 @@ "兑换人ID": "兑换人ID", "兑换成功！": "兑换成功！", "兑换码充值": "兑换码充值", + "确认清理不活跃的磁盘缓存？": "确认清理不活跃的磁盘缓存？", + "这将删除超过 10 分钟未使用的临时缓存文件": "这将删除超过 10 分钟未使用的临时缓存文件", + "清理不活跃缓存": "清理不活跃缓存", "兑换码创建成功": "兑换码创建成功", "兑换码创建成功，是否下载兑换码？": "兑换码创建成功，是否下载兑换码？", "兑换码创建成功！": "兑换码创建成功！", diff --git a/web/src/pages/Setting/Performance/SettingsPerformance.jsx b/web/src/pages/Setting/Performance/SettingsPerformance.jsx index 7f32e9873..1669a43f0 100644 --- a/web/src/pages/Setting/Performance/SettingsPerformance.jsx +++ b/web/src/pages/Setting/Performance/SettingsPerformance.jsx @@ -291,11 +291,11 @@ export default function SettingsPerformance(props) {
- +