package model import ( "strconv" "sync" "github.com/QuantumNous/new-api/common" ) var ( userChannelRatioCache = make(map[string]float64) // key: "userId:modelName:channelId" -> ratio userChannelRatioCacheLock sync.RWMutex ) // UserChannelRatio 用户-模型-渠道倍率表 type UserChannelRatio struct { Id int `json:"id" gorm:"primaryKey"` UserId int `json:"user_id" gorm:"not null;uniqueIndex:idx_user_model_channel,priority:1"` ModelName string `json:"model_name" gorm:"size:128;not null;uniqueIndex:idx_user_model_channel,priority:2"` ChannelId int `json:"channel_id" gorm:"not null;uniqueIndex:idx_user_model_channel,priority:3"` Ratio float64 `json:"ratio" gorm:"default:1"` CreatedAt int64 `json:"created_at" gorm:"bigint"` UpdatedAt int64 `json:"updated_at" gorm:"bigint"` } func getUserChannelRatioCacheKey(userId int, modelName string, channelId int) string { return strconv.Itoa(userId) + ":" + modelName + ":" + strconv.Itoa(channelId) } func setUserChannelRatioCache(key string, ratio float64) { userChannelRatioCacheLock.Lock() userChannelRatioCache[key] = ratio userChannelRatioCacheLock.Unlock() } func removeUserChannelRatioCache(key string) { userChannelRatioCacheLock.Lock() delete(userChannelRatioCache, key) userChannelRatioCacheLock.Unlock() } // GetUserChannelRatio 获取用户在指定模型+渠道的倍率(纯内存读) // 未命中返回 1.0(不影响计费) func GetUserChannelRatio(userId int, modelName string, channelId int) float64 { key := getUserChannelRatioCacheKey(userId, modelName, channelId) userChannelRatioCacheLock.RLock() ratio, ok := userChannelRatioCache[key] userChannelRatioCacheLock.RUnlock() if !ok { return 1.0 } return ratio } // LoadUserChannelRatioCache 全量加载到内存(启动时调用) func LoadUserChannelRatioCache() { var records []*UserChannelRatio if err := DB.Find(&records).Error; err != nil { common.SysError("[UserChannelRatio] LoadCache failed: " + err.Error()) return } userChannelRatioCacheLock.Lock() userChannelRatioCache = make(map[string]float64, len(records)) for _, r := range records { key := getUserChannelRatioCacheKey(r.UserId, r.ModelName, r.ChannelId) userChannelRatioCache[key] = r.Ratio } userChannelRatioCacheLock.Unlock() common.SysLog("[UserChannelRatio] cache loaded " + strconv.Itoa(len(records)) + " records") } func (ucr *UserChannelRatio) Insert() error { ucr.CreatedAt = common.GetTimestamp() ucr.UpdatedAt = common.GetTimestamp() err := DB.Create(ucr).Error if err == nil { setUserChannelRatioCache(getUserChannelRatioCacheKey(ucr.UserId, ucr.ModelName, ucr.ChannelId), ucr.Ratio) } return err } func (ucr *UserChannelRatio) Update() error { var existing UserChannelRatio if err := DB.First(&existing, ucr.Id).Error; err != nil { return err } ucr.UpdatedAt = common.GetTimestamp() err := DB.Model(&UserChannelRatio{}).Where("id = ?", ucr.Id). Select("ratio", "updated_at"). Updates(ucr).Error if err == nil { setUserChannelRatioCache(getUserChannelRatioCacheKey(existing.UserId, existing.ModelName, existing.ChannelId), ucr.Ratio) } return err } func DeleteUserChannelRatioById(id int) error { var existing UserChannelRatio if err := DB.First(&existing, id).Error; err != nil { return err } err := DB.Delete(&existing).Error if err == nil { removeUserChannelRatioCache(getUserChannelRatioCacheKey(existing.UserId, existing.ModelName, existing.ChannelId)) } return err } // GetUserChannelRatiosByUserId 获取指定用户的所有倍率记录 func GetUserChannelRatiosByUserId(userId int) ([]*UserChannelRatio, error) { var list []*UserChannelRatio err := DB.Where("user_id = ?", userId).Find(&list).Error return list, err }