|
- 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
- }
|