diff --git a/common/body_storage.go b/common/body_storage.go new file mode 100644 index 000000000..13062bd06 --- /dev/null +++ b/common/body_storage.go @@ -0,0 +1,365 @@ +package common + +import ( + "bytes" + "fmt" + "io" + "os" + "path/filepath" + "sync" + "sync/atomic" + "time" + + "github.com/google/uuid" +) + +// BodyStorage 请求体存储接口 +type BodyStorage interface { + io.ReadSeeker + io.Closer + // Bytes 获取全部内容 + Bytes() ([]byte, error) + // Size 获取数据大小 + Size() int64 + // IsDisk 是否是磁盘存储 + IsDisk() bool +} + +// ErrStorageClosed 存储已关闭错误 +var ErrStorageClosed = fmt.Errorf("body storage is closed") + +// memoryStorage 内存存储实现 +type memoryStorage struct { + data []byte + reader *bytes.Reader + size int64 + closed int32 + mu sync.Mutex +} + +func newMemoryStorage(data []byte) *memoryStorage { + size := int64(len(data)) + IncrementMemoryBuffers(size) + return &memoryStorage{ + data: data, + reader: bytes.NewReader(data), + size: size, + } +} + +func (m *memoryStorage) Read(p []byte) (n int, err error) { + m.mu.Lock() + defer m.mu.Unlock() + if atomic.LoadInt32(&m.closed) == 1 { + return 0, ErrStorageClosed + } + return m.reader.Read(p) +} + +func (m *memoryStorage) Seek(offset int64, whence int) (int64, error) { + m.mu.Lock() + defer m.mu.Unlock() + if atomic.LoadInt32(&m.closed) == 1 { + return 0, ErrStorageClosed + } + return m.reader.Seek(offset, whence) +} + +func (m *memoryStorage) Close() error { + m.mu.Lock() + defer m.mu.Unlock() + if atomic.CompareAndSwapInt32(&m.closed, 0, 1) { + DecrementMemoryBuffers(m.size) + } + return nil +} + +func (m *memoryStorage) Bytes() ([]byte, error) { + m.mu.Lock() + defer m.mu.Unlock() + if atomic.LoadInt32(&m.closed) == 1 { + return nil, ErrStorageClosed + } + return m.data, nil +} + +func (m *memoryStorage) Size() int64 { + return m.size +} + +func (m *memoryStorage) IsDisk() bool { + return false +} + +// diskStorage 磁盘存储实现 +type diskStorage struct { + file *os.File + filePath string + size int64 + closed int32 + mu sync.Mutex +} + +func newDiskStorage(data []byte, cachePath string) (*diskStorage, error) { + // 确定缓存目录 + dir := cachePath + if dir == "" { + dir = os.TempDir() + } + dir = filepath.Join(dir, "new-api-body-cache") + + // 确保目录存在 + if err := os.MkdirAll(dir, 0755); err != nil { + return nil, fmt.Errorf("failed to create cache directory: %w", err) + } + + // 创建临时文件 + filename := fmt.Sprintf("body-%s-%d.tmp", uuid.New().String()[:8], time.Now().UnixNano()) + filePath := filepath.Join(dir, filename) + + file, err := os.OpenFile(filePath, os.O_CREATE|os.O_RDWR|os.O_EXCL, 0600) + if err != nil { + return nil, fmt.Errorf("failed to create temp file: %w", err) + } + + // 写入数据 + n, err := file.Write(data) + if err != nil { + file.Close() + os.Remove(filePath) + return nil, fmt.Errorf("failed to write to temp file: %w", err) + } + + // 重置文件指针 + if _, err := file.Seek(0, io.SeekStart); err != nil { + file.Close() + os.Remove(filePath) + return nil, fmt.Errorf("failed to seek temp file: %w", err) + } + + size := int64(n) + IncrementDiskFiles(size) + + return &diskStorage{ + file: file, + filePath: filePath, + size: size, + }, nil +} + +func newDiskStorageFromReader(reader io.Reader, maxBytes int64, cachePath string) (*diskStorage, error) { + // 确定缓存目录 + dir := cachePath + if dir == "" { + dir = os.TempDir() + } + dir = filepath.Join(dir, "new-api-body-cache") + + // 确保目录存在 + if err := os.MkdirAll(dir, 0755); err != nil { + return nil, fmt.Errorf("failed to create cache directory: %w", err) + } + + // 创建临时文件 + filename := fmt.Sprintf("body-%s-%d.tmp", uuid.New().String()[:8], time.Now().UnixNano()) + filePath := filepath.Join(dir, filename) + + file, err := os.OpenFile(filePath, os.O_CREATE|os.O_RDWR|os.O_EXCL, 0600) + if err != nil { + return nil, fmt.Errorf("failed to create temp file: %w", err) + } + + // 从 reader 读取并写入文件 + written, err := io.Copy(file, io.LimitReader(reader, maxBytes+1)) + if err != nil { + file.Close() + os.Remove(filePath) + return nil, fmt.Errorf("failed to write to temp file: %w", err) + } + + if written > maxBytes { + file.Close() + os.Remove(filePath) + return nil, ErrRequestBodyTooLarge + } + + // 重置文件指针 + if _, err := file.Seek(0, io.SeekStart); err != nil { + file.Close() + os.Remove(filePath) + return nil, fmt.Errorf("failed to seek temp file: %w", err) + } + + IncrementDiskFiles(written) + + return &diskStorage{ + file: file, + filePath: filePath, + size: written, + }, nil +} + +func (d *diskStorage) Read(p []byte) (n int, err error) { + d.mu.Lock() + defer d.mu.Unlock() + if atomic.LoadInt32(&d.closed) == 1 { + return 0, ErrStorageClosed + } + return d.file.Read(p) +} + +func (d *diskStorage) Seek(offset int64, whence int) (int64, error) { + d.mu.Lock() + defer d.mu.Unlock() + if atomic.LoadInt32(&d.closed) == 1 { + return 0, ErrStorageClosed + } + return d.file.Seek(offset, whence) +} + +func (d *diskStorage) Close() error { + d.mu.Lock() + defer d.mu.Unlock() + if atomic.CompareAndSwapInt32(&d.closed, 0, 1) { + d.file.Close() + os.Remove(d.filePath) + DecrementDiskFiles(d.size) + } + return nil +} + +func (d *diskStorage) Bytes() ([]byte, error) { + d.mu.Lock() + defer d.mu.Unlock() + + if atomic.LoadInt32(&d.closed) == 1 { + return nil, ErrStorageClosed + } + + // 保存当前位置 + currentPos, err := d.file.Seek(0, io.SeekCurrent) + if err != nil { + return nil, err + } + + // 移动到开头 + if _, err := d.file.Seek(0, io.SeekStart); err != nil { + return nil, err + } + + // 读取全部内容 + data := make([]byte, d.size) + _, err = io.ReadFull(d.file, data) + if err != nil { + return nil, err + } + + // 恢复位置 + if _, err := d.file.Seek(currentPos, io.SeekStart); err != nil { + return nil, err + } + + return data, nil +} + +func (d *diskStorage) Size() int64 { + return d.size +} + +func (d *diskStorage) IsDisk() bool { + return true +} + +// CreateBodyStorage 根据数据大小创建合适的存储 +func CreateBodyStorage(data []byte) (BodyStorage, error) { + size := int64(len(data)) + threshold := GetDiskCacheThresholdBytes() + + // 检查是否应该使用磁盘缓存 + if IsDiskCacheEnabled() && + size >= threshold && + IsDiskCacheAvailable(size) { + storage, err := newDiskStorage(data, GetDiskCachePath()) + if err != nil { + // 如果磁盘存储失败，回退到内存存储 + SysError(fmt.Sprintf("failed to create disk storage, falling back to memory: %v", err)) + return newMemoryStorage(data), nil + } + return storage, nil + } + + return newMemoryStorage(data), nil +} + +// CreateBodyStorageFromReader 从 Reader 创建存储（用于大请求的流式处理） +func CreateBodyStorageFromReader(reader io.Reader, contentLength int64, maxBytes int64) (BodyStorage, error) { + threshold := GetDiskCacheThresholdBytes() + + // 如果启用了磁盘缓存且内容长度超过阈值，直接使用磁盘存储 + if IsDiskCacheEnabled() && + contentLength > 0 && + contentLength >= threshold && + IsDiskCacheAvailable(contentLength) { + storage, err := newDiskStorageFromReader(reader, maxBytes, GetDiskCachePath()) + if err != nil { + if IsRequestBodyTooLargeError(err) { + return nil, err + } + // 磁盘存储失败，reader 已被消费，无法安全回退 + // 直接返回错误而非尝试回退（因为 reader 数据已丢失） + return nil, fmt.Errorf("disk storage creation failed: %w", err) + } + IncrementDiskCacheHits() + return storage, nil + } + + // 使用内存读取 + data, err := io.ReadAll(io.LimitReader(reader, maxBytes+1)) + if err != nil { + return nil, err + } + if int64(len(data)) > maxBytes { + return nil, ErrRequestBodyTooLarge + } + + storage, err := CreateBodyStorage(data) + if err != nil { + return nil, err + } + // 如果最终使用内存存储，记录内存缓存命中 + if !storage.IsDisk() { + IncrementMemoryCacheHits() + } else { + IncrementDiskCacheHits() + } + return storage, nil +} + +// CleanupOldCacheFiles 清理旧的缓存文件（用于启动时清理残留） +func CleanupOldCacheFiles() { + cachePath := GetDiskCachePath() + if cachePath == "" { + cachePath = os.TempDir() + } + dir := filepath.Join(cachePath, "new-api-body-cache") + + entries, err := os.ReadDir(dir) + if err != nil { + return // 目录不存在或无法读取 + } + + now := time.Now() + for _, entry := range entries { + if entry.IsDir() { + continue + } + info, err := entry.Info() + if err != nil { + continue + } + // 删除超过 5 分钟的旧文件 + if now.Sub(info.ModTime()) > 5*time.Minute { + os.Remove(filepath.Join(dir, entry.Name())) + } + } +} diff --git a/common/disk_cache_config.go b/common/disk_cache_config.go new file mode 100644 index 000000000..416ec94f4 --- /dev/null +++ b/common/disk_cache_config.go @@ -0,0 +1,156 @@ +package common + +import ( + "sync" + "sync/atomic" +) + +// DiskCacheConfig 磁盘缓存配置（由 performance_setting 包更新） +type DiskCacheConfig struct { + // Enabled 是否启用磁盘缓存 + Enabled bool + // ThresholdMB 触发磁盘缓存的请求体大小阈值（MB） + ThresholdMB int + // MaxSizeMB 磁盘缓存最大总大小（MB） + MaxSizeMB int + // Path 磁盘缓存目录 + Path string +} + +// 全局磁盘缓存配置 +var diskCacheConfig = DiskCacheConfig{ + Enabled: false, + ThresholdMB: 10, + MaxSizeMB: 1024, + Path: "", +} +var diskCacheConfigMu sync.RWMutex + +// GetDiskCacheConfig 获取磁盘缓存配置 +func GetDiskCacheConfig() DiskCacheConfig { + diskCacheConfigMu.RLock() + defer diskCacheConfigMu.RUnlock() + return diskCacheConfig +} + +// SetDiskCacheConfig 设置磁盘缓存配置 +func SetDiskCacheConfig(config DiskCacheConfig) { + diskCacheConfigMu.Lock() + defer diskCacheConfigMu.Unlock() + diskCacheConfig = config +} + +// IsDiskCacheEnabled 是否启用磁盘缓存 +func IsDiskCacheEnabled() bool { + diskCacheConfigMu.RLock() + defer diskCacheConfigMu.RUnlock() + return diskCacheConfig.Enabled +} + +// GetDiskCacheThresholdBytes 获取磁盘缓存阈值（字节） +func GetDiskCacheThresholdBytes() int64 { + diskCacheConfigMu.RLock() + defer diskCacheConfigMu.RUnlock() + return int64(diskCacheConfig.ThresholdMB) << 20 +} + +// GetDiskCacheMaxSizeBytes 获取磁盘缓存最大大小（字节） +func GetDiskCacheMaxSizeBytes() int64 { + diskCacheConfigMu.RLock() + defer diskCacheConfigMu.RUnlock() + return int64(diskCacheConfig.MaxSizeMB) << 20 +} + +// GetDiskCachePath 获取磁盘缓存目录 +func GetDiskCachePath() string { + diskCacheConfigMu.RLock() + defer diskCacheConfigMu.RUnlock() + return diskCacheConfig.Path +} + +// DiskCacheStats 磁盘缓存统计信息 +type DiskCacheStats struct { + // 当前活跃的磁盘缓存文件数 + ActiveDiskFiles int64 `json:"active_disk_files"` + // 当前磁盘缓存总大小（字节） + CurrentDiskUsageBytes int64 `json:"current_disk_usage_bytes"` + // 当前内存缓存数量 + ActiveMemoryBuffers int64 `json:"active_memory_buffers"` + // 当前内存缓存总大小（字节） + CurrentMemoryUsageBytes int64 `json:"current_memory_usage_bytes"` + // 磁盘缓存命中次数 + DiskCacheHits int64 `json:"disk_cache_hits"` + // 内存缓存命中次数 + MemoryCacheHits int64 `json:"memory_cache_hits"` + // 磁盘缓存最大限制（字节） + DiskCacheMaxBytes int64 `json:"disk_cache_max_bytes"` + // 磁盘缓存阈值（字节） + DiskCacheThresholdBytes int64 `json:"disk_cache_threshold_bytes"` +} + +var diskCacheStats DiskCacheStats + +// GetDiskCacheStats 获取缓存统计信息 +func GetDiskCacheStats() DiskCacheStats { + stats := DiskCacheStats{ + ActiveDiskFiles: atomic.LoadInt64(&diskCacheStats.ActiveDiskFiles), + CurrentDiskUsageBytes: atomic.LoadInt64(&diskCacheStats.CurrentDiskUsageBytes), + ActiveMemoryBuffers: atomic.LoadInt64(&diskCacheStats.ActiveMemoryBuffers), + CurrentMemoryUsageBytes: atomic.LoadInt64(&diskCacheStats.CurrentMemoryUsageBytes), + DiskCacheHits: atomic.LoadInt64(&diskCacheStats.DiskCacheHits), + MemoryCacheHits: atomic.LoadInt64(&diskCacheStats.MemoryCacheHits), + DiskCacheMaxBytes: GetDiskCacheMaxSizeBytes(), + DiskCacheThresholdBytes: GetDiskCacheThresholdBytes(), + } + return stats +} + +// IncrementDiskFiles 增加磁盘文件计数 +func IncrementDiskFiles(size int64) { + atomic.AddInt64(&diskCacheStats.ActiveDiskFiles, 1) + atomic.AddInt64(&diskCacheStats.CurrentDiskUsageBytes, size) +} + +// DecrementDiskFiles 减少磁盘文件计数 +func DecrementDiskFiles(size int64) { + atomic.AddInt64(&diskCacheStats.ActiveDiskFiles, -1) + atomic.AddInt64(&diskCacheStats.CurrentDiskUsageBytes, -size) +} + +// IncrementMemoryBuffers 增加内存缓存计数 +func IncrementMemoryBuffers(size int64) { + atomic.AddInt64(&diskCacheStats.ActiveMemoryBuffers, 1) + atomic.AddInt64(&diskCacheStats.CurrentMemoryUsageBytes, size) +} + +// DecrementMemoryBuffers 减少内存缓存计数 +func DecrementMemoryBuffers(size int64) { + atomic.AddInt64(&diskCacheStats.ActiveMemoryBuffers, -1) + atomic.AddInt64(&diskCacheStats.CurrentMemoryUsageBytes, -size) +} + +// IncrementDiskCacheHits 增加磁盘缓存命中次数 +func IncrementDiskCacheHits() { + atomic.AddInt64(&diskCacheStats.DiskCacheHits, 1) +} + +// IncrementMemoryCacheHits 增加内存缓存命中次数 +func IncrementMemoryCacheHits() { + atomic.AddInt64(&diskCacheStats.MemoryCacheHits, 1) +} + +// ResetDiskCacheStats 重置统计信息（不重置当前使用量） +func ResetDiskCacheStats() { + atomic.StoreInt64(&diskCacheStats.DiskCacheHits, 0) + atomic.StoreInt64(&diskCacheStats.MemoryCacheHits, 0) +} + +// IsDiskCacheAvailable 检查是否可以创建新的磁盘缓存 +func IsDiskCacheAvailable(requestSize int64) bool { + if !IsDiskCacheEnabled() { + return false + } + maxBytes := GetDiskCacheMaxSizeBytes() + currentUsage := atomic.LoadInt64(&diskCacheStats.CurrentDiskUsageBytes) + return currentUsage+requestSize <= maxBytes +} diff --git a/common/gin.go b/common/gin.go index 9acb34938..1a8a2b31c 100644 --- a/common/gin.go +++ b/common/gin.go @@ -18,6 +18,7 @@ import ( ) const KeyRequestBody = "key_request_body" +const KeyBodyStorage = "key_body_storage" var ErrRequestBodyTooLarge = errors.New("request body too large") @@ -33,42 +34,99 @@ func IsRequestBodyTooLargeError(err error) bool { } func GetRequestBody(c *gin.Context) ([]byte, error) { + // 首先检查是否有 BodyStorage 缓存 + if storage, exists := c.Get(KeyBodyStorage); exists && storage != nil { + if bs, ok := storage.(BodyStorage); ok { + if _, err := bs.Seek(0, io.SeekStart); err != nil { + return nil, fmt.Errorf("failed to seek body storage: %w", err) + } + return bs.Bytes() + } + } + + // 检查旧的缓存方式 cached, exists := c.Get(KeyRequestBody) if exists && cached != nil { if b, ok := cached.([]byte); ok { return b, nil } } + maxMB := constant.MaxRequestBodyMB if maxMB <= 0 { - // no limit - body, err := io.ReadAll(c.Request.Body) - _ = c.Request.Body.Close() - if err != nil { - return nil, err - } - c.Set(KeyRequestBody, body) - return body, nil + maxMB = 128 // 默认 128MB } maxBytes := int64(maxMB) << 20 - limited := io.LimitReader(c.Request.Body, maxBytes+1) - body, err := io.ReadAll(limited) + contentLength := c.Request.ContentLength + + // 使用新的存储系统 + storage, err := CreateBodyStorageFromReader(c.Request.Body, contentLength, maxBytes) + _ = c.Request.Body.Close() + if err != nil { - _ = c.Request.Body.Close() if IsRequestBodyTooLargeError(err) { return nil, errors.Wrap(ErrRequestBodyTooLarge, fmt.Sprintf("request body exceeds %d MB", maxMB)) } return nil, err } - _ = c.Request.Body.Close() - if int64(len(body)) > maxBytes { - return nil, errors.Wrap(ErrRequestBodyTooLarge, fmt.Sprintf("request body exceeds %d MB", maxMB)) + + // 缓存存储对象 + c.Set(KeyBodyStorage, storage) + + // 获取字节数据 + body, err := storage.Bytes() + if err != nil { + return nil, err } + + // 同时设置旧的缓存键以保持兼容性 c.Set(KeyRequestBody, body) + return body, nil } +// GetBodyStorage 获取请求体存储对象（用于需要多次读取的场景） +func GetBodyStorage(c *gin.Context) (BodyStorage, error) { + // 检查是否已有存储 + if storage, exists := c.Get(KeyBodyStorage); exists && storage != nil { + if bs, ok := storage.(BodyStorage); ok { + if _, err := bs.Seek(0, io.SeekStart); err != nil { + return nil, fmt.Errorf("failed to seek body storage: %w", err) + } + return bs, nil + } + } + + // 如果没有，调用 GetRequestBody 创建存储 + _, err := GetRequestBody(c) + if err != nil { + return nil, err + } + + // 再次获取存储 + if storage, exists := c.Get(KeyBodyStorage); exists && storage != nil { + if bs, ok := storage.(BodyStorage); ok { + if _, err := bs.Seek(0, io.SeekStart); err != nil { + return nil, fmt.Errorf("failed to seek body storage: %w", err) + } + return bs, nil + } + } + + return nil, errors.New("failed to get body storage") +} + +// CleanupBodyStorage 清理请求体存储（应在请求结束时调用） +func CleanupBodyStorage(c *gin.Context) { + if storage, exists := c.Get(KeyBodyStorage); exists && storage != nil { + if bs, ok := storage.(BodyStorage); ok { + bs.Close() + } + c.Set(KeyBodyStorage, nil) + } +} + func UnmarshalBodyReusable(c *gin.Context, v any) error { requestBody, err := GetRequestBody(c) if err != nil { diff --git a/controller/misc.go b/controller/misc.go index 4d299fc81..6219676e7 100644 --- a/controller/misc.go +++ b/controller/misc.go @@ -115,6 +115,7 @@ func GetStatus(c *gin.Context) { "user_agreement_enabled": legalSetting.UserAgreement != "", "privacy_policy_enabled": legalSetting.PrivacyPolicy != "", "checkin_enabled": operation_setting.GetCheckinSetting().Enabled, + "_qn": "new-api", } // 根据启用状态注入可选内容 diff --git a/controller/performance.go b/controller/performance.go new file mode 100644 index 000000000..c7e853548 --- /dev/null +++ b/controller/performance.go @@ -0,0 +1,202 @@ +package controller + +import ( + "net/http" + "os" + "path/filepath" + "runtime" + + "github.com/QuantumNous/new-api/common" + "github.com/gin-gonic/gin" +) + +// PerformanceStats 性能统计信息 +type PerformanceStats struct { + // 缓存统计 + CacheStats common.DiskCacheStats `json:"cache_stats"` + // 系统内存统计 + MemoryStats MemoryStats `json:"memory_stats"` + // 磁盘缓存目录信息 + DiskCacheInfo DiskCacheInfo `json:"disk_cache_info"` + // 磁盘空间信息 + DiskSpaceInfo DiskSpaceInfo `json:"disk_space_info"` + // 配置信息 + Config PerformanceConfig `json:"config"` +} + +// MemoryStats 内存统计 +type MemoryStats struct { + // 已分配内存（字节） + Alloc uint64 `json:"alloc"` + // 总分配内存（字节） + TotalAlloc uint64 `json:"total_alloc"` + // 系统内存（字节） + Sys uint64 `json:"sys"` + // GC 次数 + NumGC uint32 `json:"num_gc"` + // Goroutine 数量 + NumGoroutine int `json:"num_goroutine"` +} + +// DiskCacheInfo 磁盘缓存目录信息 +type DiskCacheInfo struct { + // 缓存目录路径 + Path string `json:"path"` + // 目录是否存在 + Exists bool `json:"exists"` + // 文件数量 + FileCount int `json:"file_count"` + // 总大小（字节） + TotalSize int64 `json:"total_size"` +} + +// DiskSpaceInfo 磁盘空间信息 +type DiskSpaceInfo struct { + // 总空间（字节） + Total uint64 `json:"total"` + // 可用空间（字节） + Free uint64 `json:"free"` + // 已用空间（字节） + Used uint64 `json:"used"` + // 使用百分比 + UsedPercent float64 `json:"used_percent"` +} + +// PerformanceConfig 性能配置 +type PerformanceConfig struct { + // 是否启用磁盘缓存 + DiskCacheEnabled bool `json:"disk_cache_enabled"` + // 磁盘缓存阈值（MB） + DiskCacheThresholdMB int `json:"disk_cache_threshold_mb"` + // 磁盘缓存最大大小（MB） + DiskCacheMaxSizeMB int `json:"disk_cache_max_size_mb"` + // 磁盘缓存路径 + DiskCachePath string `json:"disk_cache_path"` + // 是否在容器中运行 + IsRunningInContainer bool `json:"is_running_in_container"` +} + +// GetPerformanceStats 获取性能统计信息 +func GetPerformanceStats(c *gin.Context) { + // 获取缓存统计 + cacheStats := common.GetDiskCacheStats() + + // 获取内存统计 + var memStats runtime.MemStats + runtime.ReadMemStats(&memStats) + + // 获取磁盘缓存目录信息 + diskCacheInfo := getDiskCacheInfo() + + // 获取配置信息 + diskConfig := common.GetDiskCacheConfig() + config := PerformanceConfig{ + DiskCacheEnabled: diskConfig.Enabled, + DiskCacheThresholdMB: diskConfig.ThresholdMB, + DiskCacheMaxSizeMB: diskConfig.MaxSizeMB, + DiskCachePath: diskConfig.Path, + IsRunningInContainer: common.IsRunningInContainer(), + } + + // 获取磁盘空间信息 + diskSpaceInfo := getDiskSpaceInfo() + + stats := PerformanceStats{ + CacheStats: cacheStats, + MemoryStats: MemoryStats{ + Alloc: memStats.Alloc, + TotalAlloc: memStats.TotalAlloc, + Sys: memStats.Sys, + NumGC: memStats.NumGC, + NumGoroutine: runtime.NumGoroutine(), + }, + DiskCacheInfo: diskCacheInfo, + DiskSpaceInfo: diskSpaceInfo, + Config: config, + } + + c.JSON(http.StatusOK, gin.H{ + "success": true, + "data": stats, + }) +} + +// ClearDiskCache 清理磁盘缓存 +func ClearDiskCache(c *gin.Context) { + cachePath := common.GetDiskCachePath() + if cachePath == "" { + cachePath = os.TempDir() + } + dir := filepath.Join(cachePath, "new-api-body-cache") + + // 删除缓存目录 + err := os.RemoveAll(dir) + if err != nil && !os.IsNotExist(err) { + common.ApiError(c, err) + return + } + + // 重置统计 + common.ResetDiskCacheStats() + + c.JSON(http.StatusOK, gin.H{ + "success": true, + "message": "磁盘缓存已清理", + }) +} + +// ResetPerformanceStats 重置性能统计 +func ResetPerformanceStats(c *gin.Context) { + common.ResetDiskCacheStats() + + c.JSON(http.StatusOK, gin.H{ + "success": true, + "message": "统计信息已重置", + }) +} + +// ForceGC 强制执行 GC +func ForceGC(c *gin.Context) { + runtime.GC() + + c.JSON(http.StatusOK, gin.H{ + "success": true, + "message": "GC 已执行", + }) +} + +// getDiskCacheInfo 获取磁盘缓存目录信息 +func getDiskCacheInfo() DiskCacheInfo { + cachePath := common.GetDiskCachePath() + if cachePath == "" { + cachePath = os.TempDir() + } + dir := filepath.Join(cachePath, "new-api-body-cache") + + info := DiskCacheInfo{ + Path: dir, + Exists: false, + } + + entries, err := os.ReadDir(dir) + if err != nil { + return info + } + + info.Exists = true + info.FileCount = 0 + info.TotalSize = 0 + + for _, entry := range entries { + if entry.IsDir() { + continue + } + info.FileCount++ + if fileInfo, err := entry.Info(); err == nil { + info.TotalSize += fileInfo.Size() + } + } + + return info +} + diff --git a/controller/performance_unix.go b/controller/performance_unix.go new file mode 100644 index 000000000..3421b6acf --- /dev/null +++ b/controller/performance_unix.go @@ -0,0 +1,37 @@ +//go:build !windows + +package controller + +import ( + "os" + + "github.com/QuantumNous/new-api/common" + "golang.org/x/sys/unix" +) + +// getDiskSpaceInfo 获取缓存目录所在磁盘的空间信息 (Unix/Linux/macOS) +func getDiskSpaceInfo() DiskSpaceInfo { + cachePath := common.GetDiskCachePath() + if cachePath == "" { + cachePath = os.TempDir() + } + + info := DiskSpaceInfo{} + + var stat unix.Statfs_t + err := unix.Statfs(cachePath, &stat) + if err != nil { + return info + } + + // 计算磁盘空间 + info.Total = stat.Blocks * uint64(stat.Bsize) + info.Free = stat.Bavail * uint64(stat.Bsize) + info.Used = info.Total - stat.Bfree*uint64(stat.Bsize) + + if info.Total > 0 { + info.UsedPercent = float64(info.Used) / float64(info.Total) * 100 + } + + return info +} diff --git a/controller/performance_windows.go b/controller/performance_windows.go new file mode 100644 index 000000000..d4c4b6e9c --- /dev/null +++ b/controller/performance_windows.go @@ -0,0 +1,52 @@ +//go:build windows + +package controller + +import ( + "os" + "syscall" + "unsafe" + + "github.com/QuantumNous/new-api/common" +) + +// getDiskSpaceInfo 获取缓存目录所在磁盘的空间信息 (Windows) +func getDiskSpaceInfo() DiskSpaceInfo { + cachePath := common.GetDiskCachePath() + if cachePath == "" { + cachePath = os.TempDir() + } + + info := DiskSpaceInfo{} + + kernel32 := syscall.NewLazyDLL("kernel32.dll") + getDiskFreeSpaceEx := kernel32.NewProc("GetDiskFreeSpaceExW") + + var freeBytesAvailable, totalBytes, totalFreeBytes uint64 + + pathPtr, err := syscall.UTF16PtrFromString(cachePath) + if err != nil { + return info + } + + ret, _, _ := getDiskFreeSpaceEx.Call( + uintptr(unsafe.Pointer(pathPtr)), + uintptr(unsafe.Pointer(&freeBytesAvailable)), + uintptr(unsafe.Pointer(&totalBytes)), + uintptr(unsafe.Pointer(&totalFreeBytes)), + ) + + if ret == 0 { + return info + } + + info.Total = totalBytes + info.Free = freeBytesAvailable + info.Used = totalBytes - totalFreeBytes + + if info.Total > 0 { + info.UsedPercent = float64(info.Used) / float64(info.Total) * 100 + } + + return info +} diff --git a/main.go b/main.go index 1326b1227..ae391ac3c 100644 --- a/main.go +++ b/main.go @@ -20,6 +20,7 @@ import ( "github.com/QuantumNous/new-api/router" "github.com/QuantumNous/new-api/service" "github.com/QuantumNous/new-api/setting/ratio_setting" + _ "github.com/QuantumNous/new-api/setting/performance_setting" // 注册性能设置 "github.com/bytedance/gopkg/util/gopool" "github.com/gin-contrib/sessions" @@ -146,6 +147,7 @@ func main() { // This will cause SSE not to work!!! //server.Use(gzip.Gzip(gzip.DefaultCompression)) server.Use(middleware.RequestId()) + server.Use(middleware.PoweredBy()) middleware.SetUpLogger(server) // Initialize session store store := cookie.NewStore([]byte(common.SessionSecret)) @@ -252,6 +254,9 @@ func InitResources() error { // Initialize options, should after model.InitDB() model.InitOptionMap() + // 清理旧的磁盘缓存文件 + common.CleanupOldCacheFiles() + // 初始化模型 model.GetPricing() diff --git a/middleware/body_cleanup.go b/middleware/body_cleanup.go new file mode 100644 index 000000000..5d06726f7 --- /dev/null +++ b/middleware/body_cleanup.go @@ -0,0 +1,18 @@ +package middleware + +import ( + "github.com/QuantumNous/new-api/common" + "github.com/gin-gonic/gin" +) + +// BodyStorageCleanup 请求体存储清理中间件 +// 在请求处理完成后自动清理磁盘/内存缓存 +func BodyStorageCleanup() gin.HandlerFunc { + return func(c *gin.Context) { + // 处理请求 + c.Next() + + // 请求结束后清理存储 + common.CleanupBodyStorage(c) + } +} diff --git a/middleware/cors.go b/middleware/cors.go index d2a109abe..6aaa15d73 100644 --- a/middleware/cors.go +++ b/middleware/cors.go @@ -1,6 +1,7 @@ package middleware import ( + "github.com/QuantumNous/new-api/common" "github.com/gin-contrib/cors" "github.com/gin-gonic/gin" ) @@ -13,3 +14,10 @@ func CORS() gin.HandlerFunc { config.AllowHeaders = []string{"*"} return cors.New(config) } + +func PoweredBy() gin.HandlerFunc { + return func(c *gin.Context) { + c.Header("X-New-Api-Version", common.Version) + c.Next() + } +} diff --git a/model/option.go b/model/option.go index e268cf577..b738db4dc 100644 --- a/model/option.go +++ b/model/option.go @@ -9,6 +9,7 @@ import ( "github.com/QuantumNous/new-api/setting" "github.com/QuantumNous/new-api/setting/config" "github.com/QuantumNous/new-api/setting/operation_setting" + "github.com/QuantumNous/new-api/setting/performance_setting" "github.com/QuantumNous/new-api/setting/ratio_setting" "github.com/QuantumNous/new-api/setting/system_setting" ) @@ -480,5 +481,11 @@ func handleConfigUpdate(key, value string) bool { } config.UpdateConfigFromMap(cfg, configMap) + // 特定配置的后处理 + if configName == "performance_setting" { + // 同步磁盘缓存配置到 common 包 + performance_setting.UpdateAndSync() + } + return true // 已处理 } diff --git a/router/api-router.go b/router/api-router.go index 39e43eee9..973684958 100644 --- a/router/api-router.go +++ b/router/api-router.go @@ -11,6 +11,7 @@ import ( func SetApiRouter(router *gin.Engine) { apiRouter := router.Group("/api") apiRouter.Use(gzip.Gzip(gzip.DefaultCompression)) + apiRouter.Use(middleware.BodyStorageCleanup()) // 清理请求体存储 apiRouter.Use(middleware.GlobalAPIRateLimit()) { apiRouter.GET("/setup", controller.GetSetup) @@ -128,6 +129,14 @@ func SetApiRouter(router *gin.Engine) { optionRoute.POST("/rest_model_ratio", controller.ResetModelRatio) optionRoute.POST("/migrate_console_setting", controller.MigrateConsoleSetting) // 用于迁移检测的旧键，下个版本会删除 } + performanceRoute := apiRouter.Group("/performance") + performanceRoute.Use(middleware.RootAuth()) + { + performanceRoute.GET("/stats", controller.GetPerformanceStats) + performanceRoute.DELETE("/disk_cache", controller.ClearDiskCache) + performanceRoute.POST("/reset_stats", controller.ResetPerformanceStats) + performanceRoute.POST("/gc", controller.ForceGC) + } ratioSyncRoute := apiRouter.Group("/ratio_sync") ratioSyncRoute.Use(middleware.RootAuth()) { diff --git a/router/relay-router.go b/router/relay-router.go index 3dc4f5c16..3f0329b36 100644 --- a/router/relay-router.go +++ b/router/relay-router.go @@ -13,6 +13,7 @@ import ( func SetRelayRouter(router *gin.Engine) { router.Use(middleware.CORS()) router.Use(middleware.DecompressRequestMiddleware()) + router.Use(middleware.BodyStorageCleanup()) // 清理请求体存储 router.Use(middleware.StatsMiddleware()) // https://platform.openai.com/docs/api-reference/introduction modelsRouter := router.Group("/v1/models") diff --git a/setting/performance_setting/config.go b/setting/performance_setting/config.go new file mode 100644 index 000000000..47ce9d45e --- /dev/null +++ b/setting/performance_setting/config.go @@ -0,0 +1,64 @@ +package performance_setting + +import ( + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/setting/config" +) + +// PerformanceSetting 性能设置配置 +type PerformanceSetting struct { + // DiskCacheEnabled 是否启用磁盘缓存（磁盘换内存） + DiskCacheEnabled bool `json:"disk_cache_enabled"` + // DiskCacheThresholdMB 触发磁盘缓存的请求体大小阈值（MB） + DiskCacheThresholdMB int `json:"disk_cache_threshold_mb"` + // DiskCacheMaxSizeMB 磁盘缓存最大总大小（MB） + DiskCacheMaxSizeMB int `json:"disk_cache_max_size_mb"` + // DiskCachePath 磁盘缓存目录 + DiskCachePath string `json:"disk_cache_path"` +} + +// 默认配置 +var performanceSetting = PerformanceSetting{ + DiskCacheEnabled: false, + DiskCacheThresholdMB: 10, // 超过 10MB 使用磁盘缓存 + DiskCacheMaxSizeMB: 1024, // 最大 1GB 磁盘缓存 + DiskCachePath: "", // 空表示使用系统临时目录 +} + +func init() { + // 注册到全局配置管理器 + config.GlobalConfig.Register("performance_setting", &performanceSetting) + // 同步初始配置到 common 包 + syncToCommon() +} + +// syncToCommon 将配置同步到 common 包 +func syncToCommon() { + common.SetDiskCacheConfig(common.DiskCacheConfig{ + Enabled: performanceSetting.DiskCacheEnabled, + ThresholdMB: performanceSetting.DiskCacheThresholdMB, + MaxSizeMB: performanceSetting.DiskCacheMaxSizeMB, + Path: performanceSetting.DiskCachePath, + }) +} + +// GetPerformanceSetting 获取性能设置 +func GetPerformanceSetting() *PerformanceSetting { + return &performanceSetting +} + +// UpdateAndSync 更新配置并同步到 common 包 +// 当配置从数据库加载后，需要调用此函数同步 +func UpdateAndSync() { + syncToCommon() +} + +// GetCacheStats 获取缓存统计信息（代理到 common 包） +func GetCacheStats() common.DiskCacheStats { + return common.GetDiskCacheStats() +} + +// ResetStats 重置统计信息 +func ResetStats() { + common.ResetDiskCacheStats() +} diff --git a/web/index.html b/web/index.html index b1bcf1713..3946f6619 100644 --- a/web/index.html +++ b/web/index.html @@ -9,6 +9,7 @@ name="description" content="OpenAI 接口聚合管理，支持多种渠道包括 Azure，可用于二次分发管理 key，仅单可执行文件，已打包好 Docker 镜像，一键部署，开箱即用" /> + New API diff --git a/web/src/components/settings/PerformanceSetting.jsx b/web/src/components/settings/PerformanceSetting.jsx new file mode 100644 index 000000000..6f1a547ac --- /dev/null +++ b/web/src/components/settings/PerformanceSetting.jsx @@ -0,0 +1,80 @@ +/* +Copyright (C) 2025 QuantumNous + +This program is free software: you can redistribute it and/or modify +it under the terms of the GNU Affero General Public License as +published by the Free Software Foundation, either version 3 of the +License, or (at your option) any later version. + +This program is distributed in the hope that it will be useful, +but WITHOUT ANY WARRANTY; without even the implied warranty of +MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +GNU Affero General Public License for more details. + +You should have received a copy of the GNU Affero General Public License +along with this program. If not, see . + +For commercial licensing, please contact support@quantumnous.com +*/ + +import React, { useEffect, useState } from 'react'; +import { Card, Spin } from '@douyinfe/semi-ui'; +import SettingsPerformance from '../../pages/Setting/Performance/SettingsPerformance'; +import { API, showError, toBoolean } from '../../helpers'; + +const PerformanceSetting = () => { + let [inputs, setInputs] = useState({ + 'performance_setting.disk_cache_enabled': false, + 'performance_setting.disk_cache_threshold_mb': 10, + 'performance_setting.disk_cache_max_size_mb': 1024, + 'performance_setting.disk_cache_path': '', + }); + + let [loading, setLoading] = useState(false); + + const getOptions = async () => { + const res = await API.get('/api/option/'); + const { success, message, data } = res.data; + if (success) { + let newInputs = {}; + data.forEach((item) => { + if (typeof inputs[item.key] === 'boolean') { + newInputs[item.key] = toBoolean(item.value); + } else { + newInputs[item.key] = item.value; + } + }); + setInputs(newInputs); + } else { + showError(message); + } + }; + + async function onRefresh() { + try { + setLoading(true); + await getOptions(); + } catch (error) { + showError('刷新失败'); + } finally { + setLoading(false); + } + } + + useEffect(() => { + onRefresh(); + }, []); + + return ( + <> + + {/* 性能设置 */} + + + + + + ); +}; + +export default PerformanceSetting; diff --git a/web/src/pages/Setting/Performance/SettingsPerformance.jsx b/web/src/pages/Setting/Performance/SettingsPerformance.jsx new file mode 100644 index 000000000..d9caf42af --- /dev/null +++ b/web/src/pages/Setting/Performance/SettingsPerformance.jsx @@ -0,0 +1,382 @@ +/* +Copyright (C) 2025 QuantumNous + +This program is free software: you can redistribute it and/or modify +it under the terms of the GNU Affero General Public License as +published by the Free Software Foundation, either version 3 of the +License, or (at your option) any later version. + +This program is distributed in the hope that it will be useful, +but WITHOUT ANY WARRANTY; without even the implied warranty of +MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +GNU Affero General Public License for more details. + +You should have received a copy of the GNU Affero General Public License +along with this program. If not, see . + +For commercial licensing, please contact support@quantumnous.com +*/ + +import React, { useEffect, useState, useRef } from 'react'; +import { + Banner, + Button, + Col, + Form, + Row, + Spin, + Progress, + Descriptions, + Tag, + Popconfirm, + Typography, +} from '@douyinfe/semi-ui'; +import { + compareObjects, + API, + showError, + showSuccess, + showWarning, +} from '../../../helpers'; +import { useTranslation } from 'react-i18next'; + +const { Text } = Typography; + +// 格式化字节大小 +function formatBytes(bytes, decimals = 2) { + if (bytes === null || bytes === undefined || isNaN(bytes)) return '0 Bytes'; + if (bytes === 0) return '0 Bytes'; + if (bytes < 0) return '-' + formatBytes(-bytes, decimals); + const k = 1024; + const dm = decimals < 0 ? 0 : decimals; + const sizes = ['Bytes', 'KB', 'MB', 'GB', 'TB']; + const i = Math.floor(Math.log(bytes) / Math.log(k)); + if (i < 0 || i >= sizes.length) return bytes + ' Bytes'; + return parseFloat((bytes / Math.pow(k, i)).toFixed(dm)) + ' ' + sizes[i]; +} + +export default function SettingsPerformance(props) { + const { t } = useTranslation(); + const [loading, setLoading] = useState(false); + const [statsLoading, setStatsLoading] = useState(false); + const [stats, setStats] = useState(null); + const [inputs, setInputs] = useState({ + 'performance_setting.disk_cache_enabled': false, + 'performance_setting.disk_cache_threshold_mb': 10, + 'performance_setting.disk_cache_max_size_mb': 1024, + 'performance_setting.disk_cache_path': '', + }); + const refForm = useRef(); + const [inputsRow, setInputsRow] = useState(inputs); + + function handleFieldChange(fieldName) { + return (value) => { + setInputs((inputs) => ({ ...inputs, [fieldName]: value })); + }; + } + + function onSubmit() { + const updateArray = compareObjects(inputs, inputsRow); + if (!updateArray.length) return showWarning(t('你似乎并没有修改什么')); + const requestQueue = updateArray.map((item) => { + let value = ''; + if (typeof inputs[item.key] === 'boolean') { + value = String(inputs[item.key]); + } else { + value = String(inputs[item.key]); + } + return API.put('/api/option/', { + key: item.key, + value, + }); + }); + setLoading(true); + Promise.all(requestQueue) + .then((res) => { + if (requestQueue.length === 1) { + if (res.includes(undefined)) return; + } else if (requestQueue.length > 1) { + if (res.includes(undefined)) + return showError(t('部分保存失败，请重试')); + } + showSuccess(t('保存成功')); + props.refresh(); + fetchStats(); + }) + .catch(() => { + showError(t('保存失败，请重试')); + }) + .finally(() => { + setLoading(false); + }); + } + + async function fetchStats() { + setStatsLoading(true); + try { + const res = await API.get('/api/performance/stats'); + if (res.data.success) { + setStats(res.data.data); + } + } catch (error) { + console.error('Failed to fetch performance stats:', error); + } finally { + setStatsLoading(false); + } + } + + async function clearDiskCache() { + try { + const res = await API.delete('/api/performance/disk_cache'); + if (res.data.success) { + showSuccess(t('磁盘缓存已清理')); + fetchStats(); + } else { + showError(res.data.message || t('清理失败')); + } + } catch (error) { + showError(t('清理失败')); + } + } + + async function resetStats() { + try { + const res = await API.post('/api/performance/reset_stats'); + if (res.data.success) { + showSuccess(t('统计已重置')); + fetchStats(); + } + } catch (error) { + showError(t('重置失败')); + } + } + + async function forceGC() { + try { + const res = await API.post('/api/performance/gc'); + if (res.data.success) { + showSuccess(t('GC 已执行')); + fetchStats(); + } + } catch (error) { + showError(t('GC 执行失败')); + } + } + + useEffect(() => { + const currentInputs = {}; + for (let key in props.options) { + if (Object.keys(inputs).includes(key)) { + if (typeof inputs[key] === 'boolean') { + currentInputs[key] = props.options[key] === 'true' || props.options[key] === true; + } else if (typeof inputs[key] === 'number') { + currentInputs[key] = parseInt(props.options[key]) || inputs[key]; + } else { + currentInputs[key] = props.options[key]; + } + } + } + setInputs({ ...inputs, ...currentInputs }); + setInputsRow({ ...inputs, ...currentInputs }); + if (refForm.current) { + refForm.current.setValues({ ...inputs, ...currentInputs }); + } + fetchStats(); + }, [props.options]); + + const diskCacheUsagePercent = stats?.cache_stats?.disk_cache_max_bytes > 0 + ? (stats.cache_stats.current_disk_usage_bytes / stats.cache_stats.disk_cache_max_bytes * 100).toFixed(1) + : 0; + + return ( + <> + +
(refForm.current = formAPI)} + style={{ marginBottom: 15 }} + > + + + + + + + + + + + 0 + ? t('可用空间: {{free}} / 总空间: {{total}}', { + free: formatBytes(stats.disk_space_info.free), + total: formatBytes(stats.disk_space_info.total), + }) + : t('磁盘缓存占用的最大空间') + } + min={100} + max={102400} + onChange={handleFieldChange('performance_setting.disk_cache_max_size_mb')} + disabled={!inputs['performance_setting.disk_cache_enabled']} + /> + + {/* 只在非容器环境显示缓存目录配置 */} + {!stats?.config?.is_running_in_container && ( + + + + )} + + + + + +
+ + + {/* 性能统计 */} + + + + +
+ + + + + + +
+ + + + {stats && ( + <> + {/* 缓存使用情况 */} + + +
+ {t('请求体磁盘缓存')} + 80 ? 'var(--semi-color-danger)' : 'var(--semi-color-primary)'} + /> +
+ + {formatBytes(stats.cache_stats.current_disk_usage_bytes)} / {formatBytes(stats.cache_stats.disk_cache_max_bytes)} + + + {t('活跃文件')}: {stats.cache_stats.active_disk_files} + +
+
+ {t('磁盘命中')}: {stats.cache_stats.disk_cache_hits} +
+
+ + +
+ {t('请求体内存缓存')} +
+ {t('当前缓存大小')}: {formatBytes(stats.cache_stats.current_memory_usage_bytes)} + {t('活跃缓存数')}: {stats.cache_stats.active_memory_buffers} +
+
+ {t('内存命中')}: {stats.cache_stats.memory_cache_hits} +
+
+ + + + {/* 缓存目录磁盘空间 */} + {stats.disk_space_info?.total > 0 && ( + + +
+ {t('缓存目录磁盘空间')} + 90 ? 'var(--semi-color-danger)' : stats.disk_space_info.used_percent > 70 ? 'var(--semi-color-warning)' : 'var(--semi-color-primary)'} + /> +
+ {t('已用')}: {formatBytes(stats.disk_space_info.used)} + {t('可用')}: {formatBytes(stats.disk_space_info.free)} + {t('总计')}: {formatBytes(stats.disk_space_info.total)} +
+ {stats.disk_space_info.free < inputs['performance_setting.disk_cache_max_size_mb'] * 1024 * 1024 && ( + + )} +
+ + + )} + + {/* 系统内存统计 */} + + + + + + + )} + + + + ); +} diff --git a/web/src/pages/Setting/index.jsx b/web/src/pages/Setting/index.jsx index 058b355e9..5f537945d 100644 --- a/web/src/pages/Setting/index.jsx +++ b/web/src/pages/Setting/index.jsx @@ -33,6 +33,7 @@ import { Palette, CreditCard, Server, + Activity, } from 'lucide-react'; import SystemSetting from '../../components/settings/SystemSetting'; @@ -47,6 +48,7 @@ import ChatsSetting from '../../components/settings/ChatsSetting'; import DrawingSetting from '../../components/settings/DrawingSetting'; import PaymentSetting from '../../components/settings/PaymentSetting'; import ModelDeploymentSetting from '../../components/settings/ModelDeploymentSetting'; +import PerformanceSetting from '../../components/settings/PerformanceSetting'; const Setting = () => { const { t } = useTranslation(); @@ -146,6 +148,16 @@ const Setting = () => { content: , itemKey: 'model-deployment', }); + panes.push({ + tab: ( + + + {t('性能设置')} + + ), + content: , + itemKey: 'performance', + }); panes.push({ tab: (