Merge remote-tracking branch 'origin/upstream/main'

2025-11-12 03:13:41 +08:00 · 2024-03-15 09:49:49 +00:00
parent d953da1ff2 b204f6d82b
commit 41afad713e
49 changed files with 623 additions and 205 deletions
--- a/model/cache.go
+++ b/model/cache.go
@@ -1,6 +1,7 @@
 package model

 import (
+	"context"
 	"encoding/json"
 	"fmt"
 	"github.com/Laisky/errors/v2"
@@ -71,31 +72,42 @@ func CacheGetUserGroup(id int) (group string, err error) {
 	return group, err
 }

-func CacheGetUserQuota(id int) (quota int, err error) {
+func fetchAndUpdateUserQuota(ctx context.Context, id int) (quota int64, err error) {
+	quota, err = GetUserQuota(id)
+	if err != nil {
+		return 0, err
+	}
+	err = common.RedisSet(fmt.Sprintf("user_quota:%d", id), fmt.Sprintf("%d", quota), time.Duration(UserId2QuotaCacheSeconds)*time.Second)
+	if err != nil {
+		logger.Error(ctx, "Redis set user quota error: "+err.Error())
+	}
+	return
+}
+
+func CacheGetUserQuota(ctx context.Context, id int) (quota int64, err error) {
 	if !common.RedisEnabled {
 		return GetUserQuota(id)
 	}
 	quotaString, err := common.RedisGet(fmt.Sprintf("user_quota:%d", id))
 	if err != nil {
-		quota, err = GetUserQuota(id)
-		if err != nil {
-			return 0, err
-		}
-		err = common.RedisSet(fmt.Sprintf("user_quota:%d", id), fmt.Sprintf("%d", quota), time.Duration(UserId2QuotaCacheSeconds)*time.Second)
-		if err != nil {
-			logger.SysError("Redis set user quota error: " + err.Error())
-		}
-		return quota, err
+		return fetchAndUpdateUserQuota(ctx, id)
 	}
-	quota, err = strconv.Atoi(quotaString)
-	return quota, err
+	quota, err = strconv.ParseInt(quotaString, 10, 64)
+	if err != nil {
+		return 0, nil
+	}
+	if quota <= config.PreConsumedQuota { // when user's quota is less than pre-consumed quota, we need to fetch from db
+		logger.Infof(ctx, "user %d's cached quota is too low: %d, refreshing from db", quota, id)
+		return fetchAndUpdateUserQuota(ctx, id)
+	}
+	return quota, nil
 }

-func CacheUpdateUserQuota(id int) error {
+func CacheUpdateUserQuota(ctx context.Context, id int) error {
 	if !common.RedisEnabled {
 		return nil
 	}
-	quota, err := CacheGetUserQuota(id)
+	quota, err := CacheGetUserQuota(ctx, id)
 	if err != nil {
 		return err
 	}
@@ -103,7 +115,7 @@ func CacheUpdateUserQuota(id int) error {
 	return err
 }

-func CacheDecreaseUserQuota(id int, quota int) error {
+func CacheDecreaseUserQuota(id int, quota int64) error {
 	if !common.RedisEnabled {
 		return nil
 	}
--- a/model/channel.go
+++ b/model/channel.go
@@ -178,7 +178,7 @@ func UpdateChannelStatusById(id int, status int) {
 	}
 }

-func UpdateChannelUsedQuota(id int, quota int) {
+func UpdateChannelUsedQuota(id int, quota int64) {
 	if config.BatchUpdateEnabled {
 		addNewRecord(BatchUpdateTypeChannelUsedQuota, id, quota)
 		return
@@ -186,7 +186,7 @@ func UpdateChannelUsedQuota(id int, quota int) {
 	updateChannelUsedQuota(id, quota)
 }

-func updateChannelUsedQuota(id int, quota int) {
+func updateChannelUsedQuota(id int, quota int64) {
 	err := DB.Model(&Channel{}).Where("id = ?", id).Update("used_quota", gorm.Expr("used_quota + ?", quota)).Error
 	if err != nil {
 		logger.SysError("failed to update channel used quota: " + err.Error())
--- a/model/log.go
+++ b/model/log.go
@@ -45,13 +45,13 @@ func RecordLog(userId int, logType int, content string) {
 		Type:      logType,
 		Content:   content,
 	}
-	err := DB.Create(log).Error
+	err := LOG_DB.Create(log).Error
 	if err != nil {
 		logger.SysError("failed to record log: " + err.Error())
 	}
 }

-func RecordConsumeLog(ctx context.Context, userId int, channelId int, promptTokens int, completionTokens int, modelName string, tokenName string, quota int, content string) {
+func RecordConsumeLog(ctx context.Context, userId int, channelId int, promptTokens int, completionTokens int, modelName string, tokenName string, quota int64, content string) {
 	logger.Info(ctx, fmt.Sprintf("record consume log: userId=%d, channelId=%d, promptTokens=%d, completionTokens=%d, modelName=%s, tokenName=%s, quota=%d, content=%s", userId, channelId, promptTokens, completionTokens, modelName, tokenName, quota, content))
 	if !config.LogConsumeEnabled {
 		return
@@ -66,10 +66,10 @@ func RecordConsumeLog(ctx context.Context, userId int, channelId int, promptToke
 		CompletionTokens: completionTokens,
 		TokenName:        tokenName,
 		ModelName:        modelName,
-		Quota:            quota,
+		Quota:            int(quota),
 		ChannelId:        channelId,
 	}
-	err := DB.Create(log).Error
+	err := LOG_DB.Create(log).Error
 	if err != nil {
 		logger.Error(ctx, "failed to record log: "+err.Error())
 	}
@@ -78,9 +78,9 @@ func RecordConsumeLog(ctx context.Context, userId int, channelId int, promptToke
 func GetAllLogs(logType int, startTimestamp int64, endTimestamp int64, modelName string, username string, tokenName string, startIdx int, num int, channel int) (logs []*Log, err error) {
 	var tx *gorm.DB
 	if logType == LogTypeUnknown {
-		tx = DB
+		tx = LOG_DB
 	} else {
-		tx = DB.Where("type = ?", logType)
+		tx = LOG_DB.Where("type = ?", logType)
 	}
 	if modelName != "" {
 		tx = tx.Where("model_name = ?", modelName)
@@ -107,9 +107,9 @@ func GetAllLogs(logType int, startTimestamp int64, endTimestamp int64, modelName
 func GetUserLogs(userId int, logType int, startTimestamp int64, endTimestamp int64, modelName string, tokenName string, startIdx int, num int) (logs []*Log, err error) {
 	var tx *gorm.DB
 	if logType == LogTypeUnknown {
-		tx = DB.Where("user_id = ?", userId)
+		tx = LOG_DB.Where("user_id = ?", userId)
 	} else {
-		tx = DB.Where("user_id = ? and type = ?", userId, logType)
+		tx = LOG_DB.Where("user_id = ? and type = ?", userId, logType)
 	}
 	if modelName != "" {
 		tx = tx.Where("model_name = ?", modelName)
@@ -128,17 +128,17 @@ func GetUserLogs(userId int, logType int, startTimestamp int64, endTimestamp int
 }

 func SearchAllLogs(keyword string) (logs []*Log, err error) {
-	err = DB.Where("type = ? or content LIKE ?", keyword, keyword+"%").Order("id desc").Limit(config.MaxRecentItems).Find(&logs).Error
+	err = LOG_DB.Where("type = ? or content LIKE ?", keyword, keyword+"%").Order("id desc").Limit(config.MaxRecentItems).Find(&logs).Error
 	return logs, err
 }

 func SearchUserLogs(userId int, keyword string) (logs []*Log, err error) {
-	err = DB.Where("user_id = ? and type = ?", userId, keyword).Order("id desc").Limit(config.MaxRecentItems).Omit("id").Find(&logs).Error
+	err = LOG_DB.Where("user_id = ? and type = ?", userId, keyword).Order("id desc").Limit(config.MaxRecentItems).Omit("id").Find(&logs).Error
 	return logs, err
 }

-func SumUsedQuota(logType int, startTimestamp int64, endTimestamp int64, modelName string, username string, tokenName string, channel int) (quota int) {
-	tx := DB.Table("logs").Select("ifnull(sum(quota),0)")
+func SumUsedQuota(logType int, startTimestamp int64, endTimestamp int64, modelName string, username string, tokenName string, channel int) (quota int64) {
+	tx := LOG_DB.Table("logs").Select("ifnull(sum(quota),0)")
 	if username != "" {
 		tx = tx.Where("username = ?", username)
 	}
@@ -162,7 +162,7 @@ func SumUsedQuota(logType int, startTimestamp int64, endTimestamp int64, modelNa
 }

 func SumUsedToken(logType int, startTimestamp int64, endTimestamp int64, modelName string, username string, tokenName string) (token int) {
-	tx := DB.Table("logs").Select("ifnull(sum(prompt_tokens),0) + ifnull(sum(completion_tokens),0)")
+	tx := LOG_DB.Table("logs").Select("ifnull(sum(prompt_tokens),0) + ifnull(sum(completion_tokens),0)")
 	if username != "" {
 		tx = tx.Where("username = ?", username)
 	}
@@ -183,7 +183,7 @@ func SumUsedToken(logType int, startTimestamp int64, endTimestamp int64, modelNa
 }

 func DeleteOldLog(targetTimestamp int64) (int64, error) {
-	result := DB.Where("created_at < ?", targetTimestamp).Delete(&Log{})
+	result := LOG_DB.Where("created_at < ?", targetTimestamp).Delete(&Log{})
 	return result.RowsAffected, result.Error
 }

@@ -207,7 +207,7 @@ func SearchLogsByDayAndModel(userId, start, end int) (LogStatistics []*LogStatis
 		groupSelect = "strftime('%Y-%m-%d', datetime(created_at, 'unixepoch')) as day"
 	}

-	err = DB.Raw(`
+	err = LOG_DB.Raw(`
 		SELECT `+groupSelect+`,
 		model_name, count(1) as request_count,
 		sum(quota) as quota,
--- a/model/main.go
+++ b/model/main.go
@@ -9,6 +9,7 @@ import (
 	"github.com/pkg/errors"
 	"github.com/songquanpeng/one-api/common"
 	"github.com/songquanpeng/one-api/common/config"
+	"github.com/songquanpeng/one-api/common/env"
 	"github.com/songquanpeng/one-api/common/helper"
 	"github.com/songquanpeng/one-api/common/logger"
 	"gorm.io/driver/mysql"
@@ -18,8 +19,9 @@ import (
 )

 var DB *gorm.DB
+var LOG_DB *gorm.DB

-func createRootAccountIfNeed() error {
+func CreateRootAccountIfNeed() error {
 	var user User
 	//if user.Status != util.UserStatusEnabled {
 	if err := DB.First(&user).Error; err != nil {
@@ -42,9 +44,9 @@ func createRootAccountIfNeed() error {
 	return nil
 }

-func chooseDB() (*gorm.DB, error) {
-	if os.Getenv("SQL_DSN") != "" {
-		dsn := os.Getenv("SQL_DSN")
+func chooseDB(envName string) (*gorm.DB, error) {
+	if os.Getenv(envName) != "" {
+		dsn := os.Getenv(envName)
 		if strings.HasPrefix(dsn, "postgres://") {
 			// Use PostgreSQL
 			logger.SysLog("using PostgreSQL as database")
@@ -72,23 +74,22 @@ func chooseDB() (*gorm.DB, error) {
 	})
 }

-func InitDB() (err error) {
-	db, err := chooseDB()
+func InitDB(envName string) (db *gorm.DB, err error) {
+	db, err = chooseDB(envName)
 	if err == nil {
 		if config.DebugSQLEnabled {
 			db = db.Debug()
 		}
-		DB = db
-		sqlDB, err := DB.DB()
+		sqlDB, err := db.DB()
 		if err != nil {
-			return errors.WithStack(err)
+			return nil, err
 		}
-		sqlDB.SetMaxIdleConns(helper.GetOrDefaultEnvInt("SQL_MAX_IDLE_CONNS", 100))
-		sqlDB.SetMaxOpenConns(helper.GetOrDefaultEnvInt("SQL_MAX_OPEN_CONNS", 1000))
-		sqlDB.SetConnMaxLifetime(time.Second * time.Duration(helper.GetOrDefaultEnvInt("SQL_MAX_LIFETIME", 60)))
+		sqlDB.SetMaxIdleConns(env.Int("SQL_MAX_IDLE_CONNS", 100))
+		sqlDB.SetMaxOpenConns(env.Int("SQL_MAX_OPEN_CONNS", 1000))
+		sqlDB.SetConnMaxLifetime(time.Second * time.Duration(env.Int("SQL_MAX_LIFETIME", 60)))

 		if !config.IsMasterNode {
-			return nil
+			return db, err
 		}
 		if common.UsingMySQL {
 			_, _ = sqlDB.Exec("DROP INDEX idx_channels_key ON channels;") // TODO: delete this line when most users have upgraded
@@ -96,46 +97,55 @@ func InitDB() (err error) {
 		logger.SysLog("database migration started")
 		err = db.AutoMigrate(&Channel{})
 		if err != nil {
-			return errors.WithStack(err)
+			return nil, err
 		}
 		err = db.AutoMigrate(&Token{})
 		if err != nil {
-			return errors.WithStack(err)
+			return nil, err
 		}
 		err = db.AutoMigrate(&User{})
 		if err != nil {
-			return errors.WithStack(err)
+			return nil, err
 		}
 		err = db.AutoMigrate(&Option{})
 		if err != nil {
-			return errors.WithStack(err)
+			return nil, err
 		}
 		err = db.AutoMigrate(&Redemption{})
 		if err != nil {
-			return errors.WithStack(err)
+			return nil, err
 		}
 		err = db.AutoMigrate(&Ability{})
 		if err != nil {
-			return errors.WithStack(err)
+			return nil, err
 		}
 		err = db.AutoMigrate(&Log{})
 		if err != nil {
-			return errors.WithStack(err)
+			return nil, err
 		}
 		logger.SysLog("database migrated")
-		err = createRootAccountIfNeed()
-		return errors.WithStack(err)
+		return db, err
 	} else {
 		logger.FatalLog(err)
 	}
-	return errors.WithStack(err)
+	return db, err
 }

-func CloseDB() error {
-	sqlDB, err := DB.DB()
+func closeDB(db *gorm.DB) error {
+	sqlDB, err := db.DB()
 	if err != nil {
 		return errors.WithStack(err)
 	}
 	err = sqlDB.Close()
 	return errors.WithStack(err)
 }
+
+func CloseDB() error {
+	if LOG_DB != DB {
+		err := closeDB(LOG_DB)
+		if err != nil {
+			return err
+		}
+	}
+	return closeDB(DB)
+}
--- a/model/option.go
+++ b/model/option.go
@@ -61,11 +61,11 @@ func InitOptionMap() {
 	config.OptionMap["MessagePusherToken"] = ""
 	config.OptionMap["TurnstileSiteKey"] = ""
 	config.OptionMap["TurnstileSecretKey"] = ""
-	config.OptionMap["QuotaForNewUser"] = strconv.Itoa(config.QuotaForNewUser)
-	config.OptionMap["QuotaForInviter"] = strconv.Itoa(config.QuotaForInviter)
-	config.OptionMap["QuotaForInvitee"] = strconv.Itoa(config.QuotaForInvitee)
-	config.OptionMap["QuotaRemindThreshold"] = strconv.Itoa(config.QuotaRemindThreshold)
-	config.OptionMap["PreConsumedQuota"] = strconv.Itoa(config.PreConsumedQuota)
+	config.OptionMap["QuotaForNewUser"] = strconv.FormatInt(config.QuotaForNewUser, 10)
+	config.OptionMap["QuotaForInviter"] = strconv.FormatInt(config.QuotaForInviter, 10)
+	config.OptionMap["QuotaForInvitee"] = strconv.FormatInt(config.QuotaForInvitee, 10)
+	config.OptionMap["QuotaRemindThreshold"] = strconv.FormatInt(config.QuotaRemindThreshold, 10)
+	config.OptionMap["PreConsumedQuota"] = strconv.FormatInt(config.PreConsumedQuota, 10)
 	config.OptionMap["ModelRatio"] = common.ModelRatio2JSONString()
 	config.OptionMap["GroupRatio"] = common.GroupRatio2JSONString()
 	config.OptionMap["CompletionRatio"] = common.CompletionRatio2JSONString()
@@ -193,15 +193,15 @@ func updateOptionMap(key string, value string) (err error) {
 	case "TurnstileSecretKey":
 		config.TurnstileSecretKey = value
 	case "QuotaForNewUser":
-		config.QuotaForNewUser, _ = strconv.Atoi(value)
+		config.QuotaForNewUser, _ = strconv.ParseInt(value, 10, 64)
 	case "QuotaForInviter":
-		config.QuotaForInviter, _ = strconv.Atoi(value)
+		config.QuotaForInviter, _ = strconv.ParseInt(value, 10, 64)
 	case "QuotaForInvitee":
-		config.QuotaForInvitee, _ = strconv.Atoi(value)
+		config.QuotaForInvitee, _ = strconv.ParseInt(value, 10, 64)
 	case "QuotaRemindThreshold":
-		config.QuotaRemindThreshold, _ = strconv.Atoi(value)
+		config.QuotaRemindThreshold, _ = strconv.ParseInt(value, 10, 64)
 	case "PreConsumedQuota":
-		config.PreConsumedQuota, _ = strconv.Atoi(value)
+		config.PreConsumedQuota, _ = strconv.ParseInt(value, 10, 64)
 	case "RetryTimes":
 		config.RetryTimes, _ = strconv.Atoi(value)
 	case "ModelRatio":
--- a/model/redemption.go
+++ b/model/redemption.go
@@ -14,7 +14,7 @@ type Redemption struct {
 	Key          string `json:"key" gorm:"type:char(32);uniqueIndex"`
 	Status       int    `json:"status" gorm:"default:1"`
 	Name         string `json:"name" gorm:"index"`
-	Quota        int    `json:"quota" gorm:"default:100"`
+	Quota        int64  `json:"quota" gorm:"default:100"`
 	CreatedTime  int64  `json:"created_time" gorm:"bigint"`
 	RedeemedTime int64  `json:"redeemed_time" gorm:"bigint"`
 	Count        int    `json:"count" gorm:"-:all"` // only for api request
@@ -42,7 +42,7 @@ func GetRedemptionById(id int) (*Redemption, error) {
 	return &redemption, err
 }

-func Redeem(key string, userId int) (quota int, err error) {
+func Redeem(key string, userId int) (quota int64, err error) {
 	if key == "" {
 		return 0, errors.New("未提供兑换码")
 	}
--- a/model/token.go
+++ b/model/token.go
@@ -21,9 +21,9 @@ type Token struct {
 	CreatedTime    int64  `json:"created_time" gorm:"bigint"`
 	AccessedTime   int64  `json:"accessed_time" gorm:"bigint"`
 	ExpiredTime    int64  `json:"expired_time" gorm:"bigint;default:-1"` // -1 means never expired
-	RemainQuota    int    `json:"remain_quota" gorm:"default:0"`
+	RemainQuota    int64  `json:"remain_quota" gorm:"default:0"`
 	UnlimitedQuota bool   `json:"unlimited_quota" gorm:"default:false"`
-	UsedQuota      int    `json:"used_quota" gorm:"default:0"` // used quota
+	UsedQuota      int64  `json:"used_quota" gorm:"default:0"` // used quota
 }

 func GetAllUserTokens(userId int, startIdx int, num int) ([]*Token, error) {
@@ -141,7 +141,7 @@ func DeleteTokenById(id int, userId int) (err error) {
 	return token.Delete()
 }

-func IncreaseTokenQuota(id int, quota int) (err error) {
+func IncreaseTokenQuota(id int, quota int64) (err error) {
 	if quota < 0 {
 		return errors.New("quota 不能为负数！")
 	}
@@ -152,7 +152,7 @@ func IncreaseTokenQuota(id int, quota int) (err error) {
 	return increaseTokenQuota(id, quota)
 }

-func increaseTokenQuota(id int, quota int) (err error) {
+func increaseTokenQuota(id int, quota int64) (err error) {
 	err = DB.Model(&Token{}).Where("id = ?", id).Updates(
 		map[string]interface{}{
 			"remain_quota":  gorm.Expr("remain_quota + ?", quota),
@@ -163,7 +163,7 @@ func increaseTokenQuota(id int, quota int) (err error) {
 	return err
 }

-func DecreaseTokenQuota(id int, quota int) (err error) {
+func DecreaseTokenQuota(id int, quota int64) (err error) {
 	if quota < 0 {
 		return errors.New("quota 不能为负数！")
 	}
@@ -174,7 +174,7 @@ func DecreaseTokenQuota(id int, quota int) (err error) {
 	return decreaseTokenQuota(id, quota)
 }

-func decreaseTokenQuota(id int, quota int) (err error) {
+func decreaseTokenQuota(id int, quota int64) (err error) {
 	err = DB.Model(&Token{}).Where("id = ?", id).Updates(
 		map[string]interface{}{
 			"remain_quota":  gorm.Expr("remain_quota - ?", quota),
@@ -185,7 +185,7 @@ func decreaseTokenQuota(id int, quota int) (err error) {
 	return err
 }

-func PreConsumeTokenQuota(tokenId int, quota int) (err error) {
+func PreConsumeTokenQuota(tokenId int, quota int64) (err error) {
 	if quota < 0 {
 		return errors.New("quota 不能为负数！")
 	}
@@ -235,7 +235,7 @@ func PreConsumeTokenQuota(tokenId int, quota int) (err error) {
 	return err
 }

-func PostConsumeTokenQuota(tokenId int, quota int) (err error) {
+func PostConsumeTokenQuota(tokenId int, quota int64) (err error) {
 	token, err := GetTokenById(tokenId)
 	if quota > 0 {
 		err = DecreaseUserQuota(token.UserId, quota)
--- a/model/user.go
+++ b/model/user.go
@@ -2,6 +2,8 @@ package model

 import (
 	"fmt"
+	"strings"
+
 	"github.com/Laisky/errors/v2"
 	"github.com/songquanpeng/one-api/common"
 	"github.com/songquanpeng/one-api/common/blacklist"
@@ -9,7 +11,6 @@ import (
 	"github.com/songquanpeng/one-api/common/helper"
 	"github.com/songquanpeng/one-api/common/logger"
 	"gorm.io/gorm"
-	"strings"
 )

 // User if you add sensitive fields, don't forget to clean them in setupLogin function.
@@ -26,10 +27,10 @@ type User struct {
 	WeChatId         string `json:"wechat_id" gorm:"column:wechat_id;index"`
 	VerificationCode string `json:"verification_code" gorm:"-:all"`                                    // this field is only for Email verification, don't save it to database!
 	AccessToken      string `json:"access_token" gorm:"type:char(32);column:access_token;uniqueIndex"` // this token is for system management
-	Quota            int    `json:"quota" gorm:"column:quota;type:int;default:0"`
-	UsedQuota        int    `json:"used_quota" gorm:"type:int;default:0;column:used_quota"`        // used quota
-	RequestCount     int    `json:"request_count" gorm:"column:request_count;type:int;default:0;"` // request number
-	Group            string `json:"group" gorm:"column:group;type:varchar(32);default:'default'"`
+	Quota            int64  `json:"quota" gorm:"type:int;default:0"`
+	UsedQuota        int64  `json:"used_quota" gorm:"type:int;default:0;column:used_quota"` // used quota
+	RequestCount     int    `json:"request_count" gorm:"type:int;default:0;"`               // request number
+	Group            string `json:"group" gorm:"type:varchar(32);default:'default'"`
 	AffCode          string `json:"aff_code" gorm:"type:varchar(32);column:aff_code;uniqueIndex"`
 	InviterId        int    `json:"inviter_id" gorm:"type:int;column:inviter_id;index"`
 }
@@ -274,12 +275,12 @@ func ValidateAccessToken(token string) (user *User) {
 	return nil
 }

-func GetUserQuota(id int) (quota int, err error) {
+func GetUserQuota(id int) (quota int64, err error) {
 	err = DB.Model(&User{}).Where("id = ?", id).Select("quota").Find(&quota).Error
 	return quota, err
 }

-func GetUserUsedQuota(id int) (quota int, err error) {
+func GetUserUsedQuota(id int) (quota int64, err error) {
 	err = DB.Model(&User{}).Where("id = ?", id).Select("used_quota").Find(&quota).Error
 	return quota, err
 }
@@ -299,7 +300,7 @@ func GetUserGroup(id int) (group string, err error) {
 	return group, err
 }

-func IncreaseUserQuota(id int, quota int) (err error) {
+func IncreaseUserQuota(id int, quota int64) (err error) {
 	if quota < 0 {
 		return errors.New("quota 不能为负数！")
 	}
@@ -310,12 +311,12 @@ func IncreaseUserQuota(id int, quota int) (err error) {
 	return increaseUserQuota(id, quota)
 }

-func increaseUserQuota(id int, quota int) (err error) {
+func increaseUserQuota(id int, quota int64) (err error) {
 	err = DB.Model(&User{}).Where("id = ?", id).Update("quota", gorm.Expr("quota + ?", quota)).Error
 	return err
 }

-func DecreaseUserQuota(id int, quota int) (err error) {
+func DecreaseUserQuota(id int, quota int64) (err error) {
 	if quota < 0 {
 		return errors.New("quota 不能为负数！")
 	}
@@ -326,7 +327,7 @@ func DecreaseUserQuota(id int, quota int) (err error) {
 	return decreaseUserQuota(id, quota)
 }

-func decreaseUserQuota(id int, quota int) (err error) {
+func decreaseUserQuota(id int, quota int64) (err error) {
 	err = DB.Model(&User{}).Where("id = ?", id).Update("quota", gorm.Expr("quota - ?", quota)).Error
 	return err
 }
@@ -336,7 +337,7 @@ func GetRootUserEmail() (email string) {
 	return email
 }

-func UpdateUserUsedQuotaAndRequestCount(id int, quota int) {
+func UpdateUserUsedQuotaAndRequestCount(id int, quota int64) {
 	if config.BatchUpdateEnabled {
 		addNewRecord(BatchUpdateTypeUsedQuota, id, quota)
 		addNewRecord(BatchUpdateTypeRequestCount, id, 1)
@@ -345,7 +346,7 @@ func UpdateUserUsedQuotaAndRequestCount(id int, quota int) {
 	updateUserUsedQuotaAndRequestCount(id, quota, 1)
 }

-func updateUserUsedQuotaAndRequestCount(id int, quota int, count int) {
+func updateUserUsedQuotaAndRequestCount(id int, quota int64, count int) {
 	err := DB.Model(&User{}).Where("id = ?", id).Updates(
 		map[string]interface{}{
 			"used_quota":    gorm.Expr("used_quota + ?", quota),
@@ -357,7 +358,7 @@ func updateUserUsedQuotaAndRequestCount(id int, quota int, count int) {
 	}
 }

-func updateUserUsedQuota(id int, quota int) {
+func updateUserUsedQuota(id int, quota int64) {
 	err := DB.Model(&User{}).Where("id = ?", id).Updates(
 		map[string]interface{}{
 			"used_quota": gorm.Expr("used_quota + ?", quota),
--- a/model/utils.go
+++ b/model/utils.go
@@ -16,12 +16,12 @@ const (
 	BatchUpdateTypeCount // if you add a new type, you need to add a new map and a new lock
 )

-var batchUpdateStores []map[int]int
+var batchUpdateStores []map[int]int64
 var batchUpdateLocks []sync.Mutex

 func init() {
 	for i := 0; i < BatchUpdateTypeCount; i++ {
-		batchUpdateStores = append(batchUpdateStores, make(map[int]int))
+		batchUpdateStores = append(batchUpdateStores, make(map[int]int64))
 		batchUpdateLocks = append(batchUpdateLocks, sync.Mutex{})
 	}
 }
@@ -35,7 +35,7 @@ func InitBatchUpdater() {
 	}()
 }

-func addNewRecord(type_ int, id int, value int) {
+func addNewRecord(type_ int, id int, value int64) {
 	batchUpdateLocks[type_].Lock()
 	defer batchUpdateLocks[type_].Unlock()
 	if _, ok := batchUpdateStores[type_][id]; !ok {
@@ -50,7 +50,7 @@ func batchUpdate() {
 	for i := 0; i < BatchUpdateTypeCount; i++ {
 		batchUpdateLocks[i].Lock()
 		store := batchUpdateStores[i]
-		batchUpdateStores[i] = make(map[int]int)
+		batchUpdateStores[i] = make(map[int]int64)
 		batchUpdateLocks[i].Unlock()
 		// TODO: maybe we can combine updates with same key?
 		for key, value := range store {
@@ -68,7 +68,7 @@ func batchUpdate() {
 			case BatchUpdateTypeUsedQuota:
 				updateUserUsedQuota(key, value)
 			case BatchUpdateTypeRequestCount:
-				updateUserRequestCount(key, value)
+				updateUserRequestCount(key, int(value))
 			case BatchUpdateTypeChannelUsedQuota:
 				updateChannelUsedQuota(key, value)
 			}