| 123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910911912913914915916917918919920921922923924925926927928929930931932933934935936937938939940941942943944945946947948949950951952953954955956957958959960961962963964965966967968969970971972973974975976977978979980981982983984985986987988989990991992993994995 |
- package model
- import (
- "errors"
- "fmt"
- "math/rand/v2"
- "strings"
- "time"
- "github.com/labring/aiproxy/core/common"
- "github.com/labring/aiproxy/core/common/config"
- "github.com/labring/aiproxy/core/common/conv"
- log "github.com/sirupsen/logrus"
- "gorm.io/gorm"
- "gorm.io/gorm/clause"
- )
- const (
- ErrTokenNotFound = "token"
- )
- const (
- PeriodTypeDaily = "daily"
- PeriodTypeWeekly = "weekly"
- PeriodTypeMonthly = "monthly"
- )
- const (
- TokenStatusEnabled = 1
- TokenStatusDisabled = 2
- )
- type Token struct {
- CreatedAt time.Time `json:"created_at"`
- Group *Group `json:"-" gorm:"foreignKey:GroupID"`
- Key string `json:"key" gorm:"type:char(48);uniqueIndex"`
- Name EmptyNullString `json:"name" gorm:"size:32;index;uniqueIndex:idx_group_name;not null"`
- GroupID string `json:"group" gorm:"size:64;index;uniqueIndex:idx_group_name"`
- Subnets []string `json:"subnets" gorm:"serializer:fastjson;type:text"`
- Models []string `json:"models" gorm:"serializer:fastjson;type:text"`
- Status int `json:"status" gorm:"default:1;index"`
- ID int `json:"id" gorm:"primaryKey"`
- UsedAmount float64 `json:"used_amount" gorm:"index"`
- RequestCount int `json:"request_count" gorm:"index"`
- Quota float64 `json:"quota"`
- PeriodQuota float64 `json:"period_quota"`
- PeriodType EmptyNullString `json:"period_type" gorm:"size:20"` // daily, weekly, monthly, default is monthly
- PeriodLastUpdateTime time.Time `json:"period_last_update_time"` // Last time period was reset
- PeriodLastUpdateAmount float64 `json:"period_last_update_amount"` // Total usage at last period reset
- }
- func (t *Token) BeforeCreate(_ *gorm.DB) error {
- if t.Key == "" || len(t.Key) != 48 {
- t.Key = generateKey()
- }
- return nil
- }
- func (t *Token) BeforeSave(_ *gorm.DB) error {
- if len(t.Name) > 32 {
- return errors.New("token name is too long")
- }
- return nil
- }
- // GetEffectiveQuotaStatus returns the effective quota status for token
- func (t *Token) GetEffectiveQuotaStatus() (totalExceeded, periodExceeded bool, err error) {
- // Check total quota (if set)
- if t.Quota > 0 && t.UsedAmount >= t.Quota {
- totalExceeded = true
- }
- if t.PeriodQuota > 0 {
- // Check if we need to reset period usage
- if needsReset, err := t.NeedsPeriodReset(); err != nil {
- return false, false, err
- } else if needsReset {
- // Period usage should be considered as reset (0) but we don't modify the struct here
- // The actual database reset should be handled separately
- periodExceeded = false // Consider as reset, so no period limit exceeded
- // Trigger async period reset - don't wait for it to complete
- go func() {
- if err := ResetTokenPeriodUsage(t.ID); err != nil {
- log.Error("failed to reset token period usage: " + err.Error())
- }
- }()
- } else {
- // Period is still valid, check against current usage
- // Calculate period usage: current total - last recorded total at period reset
- periodUsage := t.UsedAmount - t.PeriodLastUpdateAmount
- if periodUsage >= t.PeriodQuota {
- periodExceeded = true
- }
- }
- }
- return totalExceeded, periodExceeded, nil
- }
- // NeedsPeriodReset checks if the period usage should be reset
- // Uses PeriodLastUpdateTime to determine when the last period reset occurred
- func (t *Token) NeedsPeriodReset() (bool, error) {
- // If never been reset, use PeriodStartTime as baseline
- baseTime := t.PeriodLastUpdateTime
- if baseTime.IsZero() {
- return true, nil // Never initialized
- }
- now := time.Now()
- switch t.PeriodType {
- case "", PeriodTypeMonthly:
- // Check if we've crossed a month boundary since last reset
- baseMonth := baseTime.Month()
- baseYear := baseTime.Year()
- currentMonth := now.Month()
- currentYear := now.Year()
- if currentYear > baseYear {
- return true, nil
- }
- if currentYear == baseYear && currentMonth > baseMonth {
- return true, nil
- }
- return false, nil
- case PeriodTypeWeekly:
- // Check if we've passed 7 days since last reset
- return now.Sub(baseTime) >= 7*24*time.Hour, nil
- case PeriodTypeDaily:
- // Check if we've crossed to a new day since last reset
- baseDate := baseTime.Truncate(24 * time.Hour)
- currentDate := now.Truncate(24 * time.Hour)
- return currentDate.After(baseDate), nil
- default:
- return false, fmt.Errorf("unknown period type: %s", t.PeriodType)
- }
- }
- const (
- keyChars = "0123456789abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ"
- )
- func generateKey() string {
- key := make([]byte, 48)
- for i := range key {
- key[i] = keyChars[rand.IntN(len(keyChars))]
- }
- return conv.BytesToString(key)
- }
- func getTokenOrder(order string) string {
- prefix, suffix, _ := strings.Cut(order, "-")
- switch prefix {
- case "name", "expired_at", "group", "used_amount", "request_count", "id", "created_at":
- switch suffix {
- case "asc":
- return prefix + " asc"
- default:
- return prefix + " desc"
- }
- default:
- return "id desc"
- }
- }
- func InsertToken(token *Token, autoCreateGroup, ignoreExist bool) error {
- if autoCreateGroup {
- group := &Group{
- ID: token.GroupID,
- }
- if err := OnConflictDoNothing().Create(group).Error; err != nil {
- return err
- }
- }
- maxTokenNum := config.GetGroupMaxTokenNum()
- err := DB.Transaction(func(tx *gorm.DB) error {
- if maxTokenNum > 0 {
- var count int64
- err := tx.Model(&Token{}).Where("group_id = ?", token.GroupID).Count(&count).Error
- if err != nil {
- return err
- }
- if count >= maxTokenNum {
- return errors.New("group max token num reached")
- }
- }
- if ignoreExist {
- return tx.
- Where("group_id = ? and name = ?", token.GroupID, token.Name).
- FirstOrCreate(token).Error
- }
- return tx.Create(token).Error
- })
- if err != nil {
- if errors.Is(err, gorm.ErrDuplicatedKey) {
- if ignoreExist {
- return nil
- }
- return errors.New("token name already exists in this group")
- }
- return err
- }
- return nil
- }
- func GetTokens(
- group string,
- page, perPage int,
- order string,
- status int,
- ) (tokens []*Token, total int64, err error) {
- tx := DB.Model(&Token{})
- if group != "" {
- tx = tx.Where("group_id = ?", group)
- }
- if status != 0 {
- tx = tx.Where("status = ?", status)
- }
- err = tx.Count(&total).Error
- if err != nil {
- return nil, 0, err
- }
- if total <= 0 {
- return nil, 0, nil
- }
- limit, offset := toLimitOffset(page, perPage)
- err = tx.Order(getTokenOrder(order)).Limit(limit).Offset(offset).Find(&tokens).Error
- return tokens, total, err
- }
- func SearchTokens(
- group, keyword string,
- page, perPage int,
- order string,
- status int,
- name, key string,
- ) (tokens []*Token, total int64, err error) {
- tx := DB.Model(&Token{})
- if group != "" {
- tx = tx.Where("group_id = ?", group)
- }
- if status != 0 {
- tx = tx.Where("status = ?", status)
- }
- if name != "" {
- tx = tx.Where("name = ?", name)
- }
- if key != "" {
- tx = tx.Where("key = ?", key)
- }
- if keyword != "" {
- var (
- conditions []string
- values []any
- )
- if group == "" {
- if !common.UsingSQLite {
- conditions = append(conditions, "group_id ILIKE ?")
- } else {
- conditions = append(conditions, "group_id LIKE ?")
- }
- values = append(values, "%"+keyword+"%")
- }
- if name == "" {
- if !common.UsingSQLite {
- conditions = append(conditions, "name ILIKE ?")
- } else {
- conditions = append(conditions, "name LIKE ?")
- }
- values = append(values, "%"+keyword+"%")
- }
- if key == "" {
- if !common.UsingSQLite {
- conditions = append(conditions, "key ILIKE ?")
- } else {
- conditions = append(conditions, "key LIKE ?")
- }
- values = append(values, "%"+keyword+"%")
- }
- if !common.UsingSQLite {
- conditions = append(conditions, "models ILIKE ?")
- } else {
- conditions = append(conditions, "models LIKE ?")
- }
- values = append(values, "%"+keyword+"%")
- if len(conditions) > 0 {
- tx = tx.Where(fmt.Sprintf("(%s)", strings.Join(conditions, " OR ")), values...)
- }
- }
- err = tx.Count(&total).Error
- if err != nil {
- return nil, 0, err
- }
- if total <= 0 {
- return nil, 0, nil
- }
- limit, offset := toLimitOffset(page, perPage)
- err = tx.Order(getTokenOrder(order)).Limit(limit).Offset(offset).Find(&tokens).Error
- return tokens, total, err
- }
- func SearchGroupTokens(
- group, keyword string,
- page, perPage int,
- order string,
- status int,
- name, key string,
- ) (tokens []*Token, total int64, err error) {
- if group == "" {
- return nil, 0, errors.New("group is empty")
- }
- tx := DB.Model(&Token{}).
- Where("group_id = ?", group)
- if name != "" {
- tx = tx.Where("name = ?", name)
- }
- if key != "" {
- tx = tx.Where("key = ?", key)
- }
- if status != 0 {
- tx = tx.Where("status = ?", status)
- }
- if keyword != "" {
- var (
- conditions []string
- values []any
- )
- if name == "" {
- if !common.UsingSQLite {
- conditions = append(conditions, "name ILIKE ?")
- } else {
- conditions = append(conditions, "name LIKE ?")
- }
- values = append(values, "%"+keyword+"%")
- }
- if key == "" {
- if !common.UsingSQLite {
- conditions = append(conditions, "key ILIKE ?")
- } else {
- conditions = append(conditions, "key LIKE ?")
- }
- values = append(values, "%"+keyword+"%")
- }
- if !common.UsingSQLite {
- conditions = append(conditions, "models ILIKE ?")
- } else {
- conditions = append(conditions, "models LIKE ?")
- }
- values = append(values, "%"+keyword+"%")
- if len(conditions) > 0 {
- tx = tx.Where(fmt.Sprintf("(%s)", strings.Join(conditions, " OR ")), values...)
- }
- }
- err = tx.Count(&total).Error
- if err != nil {
- return nil, 0, err
- }
- if total <= 0 {
- return nil, 0, nil
- }
- limit, offset := toLimitOffset(page, perPage)
- err = tx.Order(getTokenOrder(order)).Limit(limit).Offset(offset).Find(&tokens).Error
- return tokens, total, err
- }
- func GetTokenByKey(key string) (*Token, error) {
- if key == "" {
- return nil, errors.New("key is empty")
- }
- var token Token
- err := DB.Where("key = ?", key).First(&token).Error
- return &token, HandleNotFound(err, ErrTokenNotFound)
- }
- // GetAndValidateToken validates a token and checks quota limits
- // This function is safe for concurrent use and handles period resets atomically
- func GetAndValidateToken(key string) (token *TokenCache, err error) {
- if key == "" {
- return nil, errors.New("no token provided")
- }
- token, err = CacheGetTokenByKey(key)
- if err != nil {
- if errors.Is(err, gorm.ErrRecordNotFound) {
- return nil, errors.New("invalid token")
- }
- log.Error("get token from cache failed: " + err.Error())
- return nil, errors.New("token validation failed")
- }
- if token.Status == TokenStatusDisabled {
- return nil, fmt.Errorf("token (%s[%d]) is disabled", token.Name, token.ID)
- }
- // Convert TokenCache to Token for quota checking
- tokenModel := Token{
- ID: token.ID,
- Quota: token.Quota,
- UsedAmount: token.UsedAmount,
- PeriodQuota: token.PeriodQuota,
- PeriodType: EmptyNullString(token.PeriodType),
- PeriodLastUpdateTime: time.Time(token.PeriodLastUpdateTime),
- PeriodLastUpdateAmount: token.PeriodLastUpdateAmount,
- }
- totalExceeded, periodExceeded, err := tokenModel.GetEffectiveQuotaStatus()
- if err != nil {
- return nil, fmt.Errorf("token (%s[%d]) quota check failed: %w", token.Name, token.ID, err)
- }
- if totalExceeded {
- return nil, fmt.Errorf("token (%s[%d]) total quota is exhausted", token.Name, token.ID)
- }
- if periodExceeded {
- return nil, fmt.Errorf("token (%s[%d]) period quota is exhausted", token.Name, token.ID)
- }
- return token, nil
- }
- func GetGroupTokenByID(group string, id int) (*Token, error) {
- if id == 0 || group == "" {
- return nil, errors.New("id or group is empty")
- }
- token := Token{}
- err := DB.
- Where("id = ? and group_id = ?", id, group).
- First(&token).Error
- return &token, HandleNotFound(err, ErrTokenNotFound)
- }
- func GetTokenByID(id int) (*Token, error) {
- if id == 0 {
- return nil, errors.New("id is empty")
- }
- token := Token{ID: id}
- err := DB.First(&token, "id = ?", id).Error
- return &token, HandleNotFound(err, ErrTokenNotFound)
- }
- func UpdateTokenStatus(id, status int) (err error) {
- token := Token{ID: id}
- defer func() {
- if err == nil {
- if err := CacheUpdateTokenStatus(token.Key, status); err != nil {
- log.Error("update token status in cache failed: " + err.Error())
- }
- }
- }()
- result := DB.
- Model(&token).
- Clauses(clause.Returning{
- Columns: []clause.Column{
- {Name: "key"},
- },
- }).
- Where("id = ?", id).
- Updates(
- map[string]any{
- "status": status,
- },
- )
- return HandleUpdateResult(result, ErrTokenNotFound)
- }
- func UpdateGroupTokenStatus(group string, id, status int) (err error) {
- if id == 0 || group == "" {
- return errors.New("id or group is empty")
- }
- token := Token{}
- defer func() {
- if err == nil {
- if err := CacheUpdateTokenStatus(token.Key, status); err != nil {
- log.Error("update token status in cache failed: " + err.Error())
- }
- }
- }()
- result := DB.
- Model(&token).
- Clauses(clause.Returning{
- Columns: []clause.Column{
- {Name: "key"},
- },
- }).
- Where("id = ? and group_id = ?", id, group).
- Updates(
- map[string]any{
- "status": status,
- },
- )
- return HandleUpdateResult(result, ErrTokenNotFound)
- }
- func DeleteGroupTokenByID(groupID string, id int) (err error) {
- if id == 0 || groupID == "" {
- return errors.New("id or group is empty")
- }
- token := Token{ID: id, GroupID: groupID}
- defer func() {
- if err == nil {
- if err := CacheDeleteToken(token.Key); err != nil {
- log.Error("delete token from cache failed: " + err.Error())
- }
- }
- }()
- result := DB.
- Clauses(clause.Returning{
- Columns: []clause.Column{
- {Name: "key"},
- },
- }).
- Where(token).
- Delete(&token)
- return HandleUpdateResult(result, ErrTokenNotFound)
- }
- func DeleteGroupTokensByIDs(group string, ids []int) (err error) {
- if group == "" {
- return errors.New("group is empty")
- }
- if len(ids) == 0 {
- return nil
- }
- tokens := make([]Token, len(ids))
- defer func() {
- if err == nil {
- for _, token := range tokens {
- if err := CacheDeleteToken(token.Key); err != nil {
- log.Error("delete token from cache failed: " + err.Error())
- }
- }
- }
- }()
- return DB.Transaction(func(tx *gorm.DB) error {
- return tx.
- Clauses(clause.Returning{
- Columns: []clause.Column{
- {Name: "key"},
- },
- }).
- Where("group_id = ?", group).
- Where("id IN (?)", ids).
- Delete(&tokens).
- Error
- })
- }
- func DeleteTokenByID(id int) (err error) {
- if id == 0 {
- return errors.New("id is empty")
- }
- token := Token{ID: id}
- defer func() {
- if err == nil {
- if err := CacheDeleteToken(token.Key); err != nil {
- log.Error("delete token from cache failed: " + err.Error())
- }
- }
- }()
- result := DB.
- Clauses(clause.Returning{
- Columns: []clause.Column{
- {Name: "key"},
- },
- }).
- Where(token).
- Delete(&token)
- return HandleUpdateResult(result, ErrTokenNotFound)
- }
- func DeleteTokensByIDs(ids []int) (err error) {
- if len(ids) == 0 {
- return nil
- }
- tokens := make([]Token, len(ids))
- defer func() {
- if err == nil {
- for _, token := range tokens {
- if err := CacheDeleteToken(token.Key); err != nil {
- log.Error("delete token from cache failed: " + err.Error())
- }
- }
- }
- }()
- return DB.Transaction(func(tx *gorm.DB) error {
- return tx.
- Clauses(clause.Returning{
- Columns: []clause.Column{
- {Name: "key"},
- },
- }).
- Where("id IN (?)", ids).
- Delete(&tokens).
- Error
- })
- }
- type UpdateTokenRequest struct {
- Name *string `json:"name"`
- Subnets *[]string `json:"subnets"`
- Models *[]string `json:"models"`
- Status int `json:"status"`
- // Quota system
- Quota *float64 `json:"quota"`
- PeriodQuota *float64 `json:"period_quota"`
- PeriodType *string `json:"period_type"`
- PeriodLastUpdateTime *int64 `json:"period_last_update_time"`
- }
- func UpdateToken(id int, update UpdateTokenRequest) (token *Token, err error) {
- if id == 0 {
- return nil, errors.New("id is empty")
- }
- token = &Token{
- ID: id,
- Status: update.Status,
- }
- defer func() {
- if err == nil {
- if err := CacheDeleteToken(token.Key); err != nil {
- log.Error("delete token from cache failed: " + err.Error())
- }
- }
- }()
- selects := []string{}
- if update.Name != nil && *update.Name != "" {
- token.Name = EmptyNullString(*update.Name)
- selects = append(selects, "name")
- }
- if update.Quota != nil {
- token.Quota = *update.Quota
- selects = append(selects, "quota")
- }
- if update.PeriodQuota != nil {
- token.PeriodQuota = *update.PeriodQuota
- selects = append(selects, "period_quota")
- }
- if update.PeriodType != nil {
- token.PeriodType = EmptyNullString(*update.PeriodType)
- selects = append(selects, "period_type")
- }
- if update.PeriodLastUpdateTime != nil {
- token.PeriodLastUpdateTime = time.UnixMilli(*update.PeriodLastUpdateTime)
- selects = append(selects, "period_last_update_time")
- }
- if update.Subnets != nil {
- token.Subnets = *update.Subnets
- selects = append(selects, "subnets")
- }
- if update.Models != nil {
- token.Models = *update.Models
- selects = append(selects, "models")
- }
- if update.Status != 0 {
- selects = append(selects, "status")
- }
- if len(selects) == 0 {
- return nil, errors.New("empty update request")
- }
- result := DB.
- Select(selects).
- Where("id = ?", id).
- Clauses(clause.Returning{}).
- Updates(token)
- if result.Error != nil {
- if errors.Is(result.Error, gorm.ErrDuplicatedKey) {
- return nil, errors.New("token name already exists in this group")
- }
- }
- return token, HandleUpdateResult(result, ErrTokenNotFound)
- }
- func UpdateGroupToken(
- id int,
- group string,
- update UpdateTokenRequest,
- ) (token *Token, err error) {
- if id == 0 || group == "" {
- return nil, errors.New("id or group is empty")
- }
- token = &Token{
- ID: id,
- GroupID: group,
- Status: update.Status,
- }
- defer func() {
- if err == nil {
- if err := CacheDeleteToken(token.Key); err != nil {
- log.Error("delete token from cache failed: " + err.Error())
- }
- }
- }()
- selects := []string{}
- if update.Name != nil && *update.Name != "" {
- token.Name = EmptyNullString(*update.Name)
- selects = append(selects, "name")
- }
- if update.Quota != nil {
- token.Quota = *update.Quota
- selects = append(selects, "quota")
- }
- if update.PeriodQuota != nil {
- token.PeriodQuota = *update.PeriodQuota
- selects = append(selects, "period_quota")
- }
- if update.PeriodType != nil {
- token.PeriodType = EmptyNullString(*update.PeriodType)
- selects = append(selects, "period_type")
- }
- if update.PeriodLastUpdateTime != nil {
- token.PeriodLastUpdateTime = time.UnixMilli(*update.PeriodLastUpdateTime)
- selects = append(selects, "period_last_update_time")
- }
- if update.Subnets != nil {
- token.Subnets = *update.Subnets
- selects = append(selects, "subnets")
- }
- if update.Models != nil {
- token.Models = *update.Models
- selects = append(selects, "models")
- }
- if update.Status != 0 {
- selects = append(selects, "status")
- }
- if len(selects) == 0 {
- return nil, errors.New("empty update request")
- }
- result := DB.
- Select(selects).
- Where("id = ? and group_id = ?", id, group).
- Clauses(clause.Returning{}).
- Updates(token)
- if result.Error != nil {
- if errors.Is(result.Error, gorm.ErrDuplicatedKey) {
- return nil, errors.New("token name already exists in this group")
- }
- }
- return token, HandleUpdateResult(result, ErrTokenNotFound)
- }
- func UpdateTokenUsedAmount(id int, amount float64, requestCount int) (err error) {
- token := &Token{}
- defer func() {
- if amount > 0 && err == nil && (token.Quota > 0 || token.PeriodQuota > 0) {
- if err := CacheUpdateTokenUsedAmountOnlyIncrease(token.Key, token.UsedAmount); err != nil {
- log.Error("update token used amount in cache failed: " + err.Error())
- }
- }
- }()
- result := DB.
- Model(token).
- Clauses(clause.Returning{
- Columns: []clause.Column{
- {Name: "key"},
- {Name: "quota"},
- {Name: "used_amount"},
- {Name: "period_quota"},
- },
- }).
- Where("id = ?", id).
- Updates(
- map[string]any{
- "used_amount": gorm.Expr("used_amount + ?", amount),
- "request_count": gorm.Expr("request_count + ?", requestCount),
- },
- )
- return HandleUpdateResult(result, ErrTokenNotFound)
- }
- // ResetTokenPeriodUsage resets the period usage for a token with concurrency safety
- // This updates PeriodLastUpdateTime and PeriodLastUpdateAmount to current values
- func ResetTokenPeriodUsage(id int) error {
- token := &Token{}
- // Use database transaction with optimistic locking to prevent concurrent resets
- err := DB.Transaction(func(tx *gorm.DB) error {
- // First, read the current state with FOR UPDATE lock
- if err := tx.Clauses(clause.Locking{Strength: "UPDATE"}).
- Where("id = ?", id).
- First(token).Error; err != nil {
- return err
- }
- // Check if period still needs reset (another concurrent request might have already reset it)
- needsReset, err := token.NeedsPeriodReset()
- if err != nil {
- return err
- }
- // If period no longer needs reset, skip the update
- if !needsReset {
- return nil
- }
- // Perform the reset with the lock held - update period last update time and amount
- result := tx.
- Model(token).
- Clauses(clause.Returning{
- Columns: []clause.Column{
- {Name: "key"},
- },
- }).
- Where("id = ?", id).
- Updates(
- map[string]any{
- "period_last_update_time": time.Now(),
- "period_last_update_amount": gorm.Expr(
- "used_amount",
- ), // Set to current total usage
- },
- )
- return HandleUpdateResult(result, ErrTokenNotFound)
- })
- // Update cache only if database update succeeded
- if err == nil && token.Key != "" {
- if cacheErr := CacheResetTokenPeriodUsage(token.Key, time.Now(), token.UsedAmount); cacheErr != nil {
- log.Error("reset token period usage in cache failed: " + cacheErr.Error())
- }
- }
- return err
- }
- func UpdateTokenName(id int, name string) (err error) {
- token := &Token{ID: id}
- defer func() {
- if err == nil {
- if err := CacheUpdateTokenName(token.Key, name); err != nil {
- log.Error("update token name in cache failed: " + err.Error())
- }
- }
- }()
- result := DB.
- Model(token).
- Clauses(clause.Returning{
- Columns: []clause.Column{
- {Name: "key"},
- },
- }).
- Where("id = ?", id).
- Update("name", name)
- if result.Error != nil && errors.Is(result.Error, gorm.ErrDuplicatedKey) {
- return errors.New("token name already exists in this group")
- }
- return HandleUpdateResult(result, ErrTokenNotFound)
- }
- func UpdateGroupTokenName(group string, id int, name string) (err error) {
- token := &Token{ID: id, GroupID: group}
- defer func() {
- if err == nil {
- if err := CacheUpdateTokenName(token.Key, name); err != nil {
- log.Error("update token name in cache failed: " + err.Error())
- }
- }
- }()
- result := DB.
- Model(token).
- Clauses(clause.Returning{
- Columns: []clause.Column{
- {Name: "key"},
- },
- }).
- Where("id = ? and group_id = ?", id, group).
- Update("name", name)
- if result.Error != nil && errors.Is(result.Error, gorm.ErrDuplicatedKey) {
- return errors.New("token name already exists in this group")
- }
- return HandleUpdateResult(result, ErrTokenNotFound)
- }
|