You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
 
 
 

115 lines
3.7 KiB

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