Browse Source

refactor(channel-pricing): 结构体新增扩展字段 + 缓存重写为全量加载写穿

- ChannelPricing 结构体新增 5 个扩展计费字段: cache_ratio, cache_creation_ratio, image_ratio, audio_ratio, audio_completion_ratio
- 缓存机制从 TTL+惰性加载重写为全量加载+写穿模式,消除 DB 查询延迟
- Insert/Update/Delete 改为写穿缓存(直接更新内存,不再整表失效)
- 新增 LoadChannelPricingCache 启动时全量加载函数
- 删除 RefreshChannelPricingCache/InvalidateChannelPricingCache(不再需要)
- BatchUpsertChannelPricing DoUpdates 列表同步新增 5 字段
- GetEffectivePricing 签名改为 (*ChannelPricing, bool)(调用方适配在后续 Task)

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
master
fengsilin 3 weeks ago
parent
commit
2d4c73d3aa
1 changed files with 48 additions and 68 deletions
  1. +48
    -68
      model/channel_pricing.go

+ 48
- 68
model/channel_pricing.go View File

@@ -5,7 +5,6 @@ import (
"strconv"
"strings"
"sync"
"time"

"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/setting/ratio_setting"
@@ -17,8 +16,6 @@ import (
var (
channelPricingCache = make(map[string]*ChannelPricing) // key: "modelName:channelId"
channelPricingCacheLock sync.RWMutex
channelPricingCacheTime time.Time
channelPricingCacheTTL = time.Minute * 5 // 缓存5分钟
)

// QuotaType 计费类型
@@ -41,6 +38,12 @@ type ChannelPricing struct {
CreatedTime int64 `json:"created_time" gorm:"bigint"`
UpdatedTime int64 `json:"updated_time" gorm:"bigint"`
DeletedAt gorm.DeletedAt `json:"-" gorm:"index"`
// === 新增字段(0 = 未设置,回退全局值) ===
CacheRatio float64 `json:"cache_ratio" gorm:"default:0"`
CacheCreationRatio float64 `json:"cache_creation_ratio" gorm:"default:0"`
ImageRatio float64 `json:"image_ratio" gorm:"default:0"`
AudioRatio float64 `json:"audio_ratio" gorm:"default:0"`
AudioCompletionRatio float64 `json:"audio_completion_ratio" gorm:"default:0"`
}

func (cp *ChannelPricing) Insert() error {
@@ -49,31 +52,41 @@ func (cp *ChannelPricing) Insert() error {
cp.UpdatedTime = now
err := DB.Create(cp).Error
if err == nil {
InvalidateChannelPricingCache()
key := getChannelPricingCacheKey(cp.ModelName, cp.ChannelId)
channelPricingCacheLock.Lock()
channelPricingCache[key] = cp
channelPricingCacheLock.Unlock()
}
return err
}

func (cp *ChannelPricing) Update() error {
cp.UpdatedTime = common.GetTimestamp()
err := DB.Model(&ChannelPricing{}).Where("id = ?", cp.Id).Updates(map[string]interface{}{
"quota_type": cp.QuotaType,
"model_ratio": cp.ModelRatio,
"completion_ratio": cp.CompletionRatio,
"model_price": cp.ModelPrice,
"tag_ids": cp.TagIds,
"updated_time": cp.UpdatedTime,
}).Error
err := DB.Model(&ChannelPricing{}).Where("id = ?", cp.Id).
Select("quota_type", "model_ratio", "completion_ratio", "model_price",
"tag_ids", "cache_ratio", "cache_creation_ratio", "image_ratio",
"audio_ratio", "audio_completion_ratio", "updated_time").
Updates(cp).Error
if err == nil {
InvalidateChannelPricingCache()
key := getChannelPricingCacheKey(cp.ModelName, cp.ChannelId)
channelPricingCacheLock.Lock()
channelPricingCache[key] = cp
channelPricingCacheLock.Unlock()
}
return err
}

func (cp *ChannelPricing) Delete() error {
var existing ChannelPricing
if err := DB.First(&existing, cp.Id).Error; err != nil {
return err
}
err := DB.Delete(cp).Error
if err == nil {
InvalidateChannelPricingCache()
key := getChannelPricingCacheKey(existing.ModelName, existing.ChannelId)
channelPricingCacheLock.Lock()
delete(channelPricingCache, key)
channelPricingCacheLock.Unlock()
}
return err
}
@@ -132,6 +145,11 @@ func BatchUpsertChannelPricing(pricings []*ChannelPricing) error {
"completion_ratio",
"model_price",
"tag_ids",
"cache_ratio",
"cache_creation_ratio",
"image_ratio",
"audio_ratio",
"audio_completion_ratio",
"updated_time",
}),
}).Create(&pricings).Error
@@ -142,71 +160,33 @@ func getChannelPricingCacheKey(modelName string, channelId int) string {
return fmt.Sprintf("%s:%d", modelName, channelId)
}

// GetEffectivePricing 获取有效定价(优先渠道定价,回退全局定价)
// 返回: modelRatio, completionRatio, modelPrice, usePrice, found
func GetEffectivePricing(modelName string, channelId int) (modelRatio, completionRatio, modelPrice float64, usePrice, found bool) {
cacheKey := getChannelPricingCacheKey(modelName, channelId)

// 首先检查缓存
// GetEffectivePricing 获取有效定价(纯内存查找)
func GetEffectivePricing(modelName string, channelId int) (*ChannelPricing, bool) {
key := getChannelPricingCacheKey(modelName, channelId)
channelPricingCacheLock.RLock()
// 检查缓存是否过期
if time.Since(channelPricingCacheTime) < channelPricingCacheTTL {
if cp, ok := channelPricingCache[cacheKey]; ok {
channelPricingCacheLock.RUnlock()
return cp.ModelRatio, cp.CompletionRatio, cp.ModelPrice, cp.QuotaType == QuotaTypeByCall, true
}
}
cp, ok := channelPricingCache[key]
channelPricingCacheLock.RUnlock()

// 缓存未命中或已过期,查询数据库
var cp ChannelPricing
err := DB.Where("model_name = ? AND channel_id = ?", modelName, channelId).First(&cp).Error
if err != nil {
// 未找到渠道定价,返回 false 让调用者使用全局定价
return 0, 0, 0, false, false
}

// 更新缓存
channelPricingCacheLock.Lock()
if channelPricingCacheTime.IsZero() || time.Since(channelPricingCacheTime) >= channelPricingCacheTTL {
// 缓存过期,清空并更新时间
channelPricingCache = make(map[string]*ChannelPricing)
channelPricingCacheTime = time.Now()
if !ok {
return nil, false
}
channelPricingCache[cacheKey] = &cp
channelPricingCacheLock.Unlock()

return cp.ModelRatio, cp.CompletionRatio, cp.ModelPrice, cp.QuotaType == QuotaTypeByCall, true
return cp, true
}

// RefreshChannelPricingCache 刷新渠道定价缓存
func RefreshChannelPricingCache() {
channelPricingCacheLock.Lock()
defer channelPricingCacheLock.Unlock()

// 清空缓存
channelPricingCache = make(map[string]*ChannelPricing)
channelPricingCacheTime = time.Now()

// 预加载所有渠道定价
// LoadChannelPricingCache 全量加载渠道定价到内存(启动时调用)
func LoadChannelPricingCache() {
var pricings []*ChannelPricing
if err := DB.Find(&pricings).Error; err != nil {
common.SysError("[ChannelPricing] LoadChannelPricingCache failed: " + err.Error())
return
}

channelPricingCacheLock.Lock()
channelPricingCache = make(map[string]*ChannelPricing, len(pricings))
for _, cp := range pricings {
cacheKey := getChannelPricingCacheKey(cp.ModelName, cp.ChannelId)
channelPricingCache[cacheKey] = cp
key := getChannelPricingCacheKey(cp.ModelName, cp.ChannelId)
channelPricingCache[key] = cp
}
}

// InvalidateChannelPricingCache 使渠道定价缓存失效
func InvalidateChannelPricingCache() {
channelPricingCacheLock.Lock()
defer channelPricingCacheLock.Unlock()

channelPricingCache = make(map[string]*ChannelPricing)
channelPricingCacheTime = time.Time{} // 重置为零值
channelPricingCacheLock.Unlock()
common.SysLog(fmt.Sprintf("[ChannelPricing] cache loaded %d records", len(pricings)))
}

// ChannelPricingWithChannel 带渠道信息的定价响应


Loading…
Cancel
Save