Nie możesz wybrać więcej, niż 25 tematów Tematy muszą się zaczynać od litery lub cyfry, mogą zawierać myślniki ('-') i mogą mieć do 35 znaków.
 
 
 

130 wiersze
4.2 KiB

  1. package model
  2. import (
  3. "strconv"
  4. "strings"
  5. "sync"
  6. "github.com/QuantumNous/new-api/common"
  7. )
  8. var (
  9. userChannelRatioCache = make(map[string]float64) // key: "userId:modelName:channelId" -> ratio
  10. userChannelRatioCacheLock sync.RWMutex
  11. )
  12. // UserChannelRatio 用户-模型-渠道倍率表
  13. type UserChannelRatio struct {
  14. Id int `json:"id" gorm:"primaryKey"`
  15. UserId int `json:"user_id" gorm:"not null;uniqueIndex:idx_user_model_channel,priority:1"`
  16. ModelName string `json:"model_name" gorm:"size:128;not null;uniqueIndex:idx_user_model_channel,priority:2"`
  17. ChannelId int `json:"channel_id" gorm:"not null;uniqueIndex:idx_user_model_channel,priority:3"`
  18. Ratio float64 `json:"ratio" gorm:"default:1"`
  19. CreatedAt int64 `json:"created_at" gorm:"bigint"`
  20. UpdatedAt int64 `json:"updated_at" gorm:"bigint"`
  21. }
  22. func getUserChannelRatioCacheKey(userId int, modelName string, channelId int) string {
  23. return strconv.Itoa(userId) + ":" + modelName + ":" + strconv.Itoa(channelId)
  24. }
  25. func setUserChannelRatioCache(key string, ratio float64) {
  26. userChannelRatioCacheLock.Lock()
  27. userChannelRatioCache[key] = ratio
  28. userChannelRatioCacheLock.Unlock()
  29. }
  30. func removeUserChannelRatioCache(key string) {
  31. userChannelRatioCacheLock.Lock()
  32. delete(userChannelRatioCache, key)
  33. userChannelRatioCacheLock.Unlock()
  34. }
  35. // GetUserChannelRatio 获取用户在指定模型+渠道的倍率(纯内存读)
  36. // 未命中返回 1.0(不影响计费)
  37. func GetUserChannelRatio(userId int, modelName string, channelId int) float64 {
  38. key := getUserChannelRatioCacheKey(userId, modelName, channelId)
  39. userChannelRatioCacheLock.RLock()
  40. ratio, ok := userChannelRatioCache[key]
  41. userChannelRatioCacheLock.RUnlock()
  42. if !ok {
  43. return 1.0
  44. }
  45. return ratio
  46. }
  47. // LoadUserChannelRatioCache 全量加载到内存(启动时调用)
  48. func LoadUserChannelRatioCache() {
  49. var records []*UserChannelRatio
  50. if err := DB.Find(&records).Error; err != nil {
  51. common.SysError("[UserChannelRatio] LoadCache failed: " + err.Error())
  52. return
  53. }
  54. userChannelRatioCacheLock.Lock()
  55. userChannelRatioCache = make(map[string]float64, len(records))
  56. for _, r := range records {
  57. key := getUserChannelRatioCacheKey(r.UserId, r.ModelName, r.ChannelId)
  58. userChannelRatioCache[key] = r.Ratio
  59. }
  60. userChannelRatioCacheLock.Unlock()
  61. common.SysLog("[UserChannelRatio] cache loaded " + strconv.Itoa(len(records)) + " records")
  62. }
  63. func (ucr *UserChannelRatio) Insert() error {
  64. ucr.CreatedAt = common.GetTimestamp()
  65. ucr.UpdatedAt = common.GetTimestamp()
  66. err := DB.Create(ucr).Error
  67. if err == nil {
  68. setUserChannelRatioCache(getUserChannelRatioCacheKey(ucr.UserId, ucr.ModelName, ucr.ChannelId), ucr.Ratio)
  69. }
  70. return err
  71. }
  72. func (ucr *UserChannelRatio) Update() error {
  73. var existing UserChannelRatio
  74. if err := DB.First(&existing, ucr.Id).Error; err != nil {
  75. return err
  76. }
  77. ucr.UpdatedAt = common.GetTimestamp()
  78. err := DB.Model(&UserChannelRatio{}).Where("id = ?", ucr.Id).
  79. Select("ratio", "updated_at").
  80. Updates(ucr).Error
  81. if err == nil {
  82. setUserChannelRatioCache(getUserChannelRatioCacheKey(existing.UserId, existing.ModelName, existing.ChannelId), ucr.Ratio)
  83. }
  84. return err
  85. }
  86. func DeleteUserChannelRatioById(id int) error {
  87. var existing UserChannelRatio
  88. if err := DB.First(&existing, id).Error; err != nil {
  89. return err
  90. }
  91. err := DB.Delete(&existing).Error
  92. if err == nil {
  93. removeUserChannelRatioCache(getUserChannelRatioCacheKey(existing.UserId, existing.ModelName, existing.ChannelId))
  94. }
  95. return err
  96. }
  97. // GetUserChannelRatiosByUserId 获取指定用户的所有倍率记录
  98. func GetUserChannelRatiosByUserId(userId int) ([]*UserChannelRatio, error) {
  99. var list []*UserChannelRatio
  100. err := DB.Where("user_id = ?", userId).Find(&list).Error
  101. return list, err
  102. }
  103. // GetBestUserChannelRatio 从内存缓存中查找指定用户+模型的最低渠道倍率(未命中返回 1.0)
  104. func GetBestUserChannelRatio(userId int, modelName string) float64 {
  105. best := 1.0
  106. prefix := strconv.Itoa(userId) + ":" + modelName + ":"
  107. userChannelRatioCacheLock.RLock()
  108. for key, ratio := range userChannelRatioCache {
  109. if len(key) > len(prefix) && strings.HasPrefix(key, prefix) && ratio < best {
  110. best = ratio
  111. }
  112. }
  113. userChannelRatioCacheLock.RUnlock()
  114. return best
  115. }