package model import ( "errors" "fmt" "strings" "time" "github.com/QuantumNous/new-api/common" "gorm.io/gorm" ) const ( ChannelResetNever = "never" ChannelResetDaily = "daily" ChannelResetWeekly = "weekly" ChannelResetMonthly = "monthly" ChannelResetCustom = "custom" ) func NormalizeChannelResetPeriod(period string) string { switch strings.TrimSpace(period) { case ChannelResetDaily, ChannelResetWeekly, ChannelResetMonthly, ChannelResetCustom: return strings.TrimSpace(period) default: return ChannelResetNever } } func calcChannelNextResetTime(base time.Time, period string, customSeconds int) int64 { period = NormalizeChannelResetPeriod(period) if period == ChannelResetNever { return 0 } var next time.Time switch period { case ChannelResetDaily: next = time.Date(base.Year(), base.Month(), base.Day(), 0, 0, 0, 0, base.Location()). AddDate(0, 0, 1) case ChannelResetWeekly: weekday := int(base.Weekday()) if weekday == 0 { weekday = 7 } daysUntil := 8 - weekday next = time.Date(base.Year(), base.Month(), base.Day(), 0, 0, 0, 0, base.Location()). AddDate(0, 0, daysUntil) case ChannelResetMonthly: next = time.Date(base.Year(), base.Month(), 1, 0, 0, 0, 0, base.Location()). AddDate(0, 1, 0) case ChannelResetCustom: if customSeconds <= 0 { return 0 } next = base.Add(time.Duration(customSeconds) * time.Second) default: return 0 } return next.Unix() } func maybeResetChannelQuotaTx(tx *gorm.DB, channel *Channel, now int64) error { if tx == nil || channel == nil { return errors.New("invalid reset args") } if channel.NextResetTime > 0 && channel.NextResetTime > now { return nil } if NormalizeChannelResetPeriod(channel.QuotaResetPeriod) == ChannelResetNever { return nil } baseUnix := channel.LastResetTime if baseUnix <= 0 { baseUnix = channel.CreatedTime } base := time.Unix(baseUnix, 0) next := calcChannelNextResetTime(base, channel.QuotaResetPeriod, channel.QuotaResetCustomSeconds) advanced := false for next > 0 && next <= now { advanced = true base = time.Unix(next, 0) next = calcChannelNextResetTime(base, channel.QuotaResetPeriod, channel.QuotaResetCustomSeconds) } if !advanced { if channel.NextResetTime == 0 && next > 0 { return tx.Model(&Channel{}).Where("id = ?", channel.Id).Updates(map[string]interface{}{ "next_reset_time": next, "last_reset_time": base.Unix(), }).Error } return nil } return tx.Model(&Channel{}).Where("id = ?", channel.Id).Updates(map[string]interface{}{ "used_quota": 0, "last_reset_time": base.Unix(), "next_reset_time": next, }).Error } func ResetDueChannels(limit int) (int, error) { if limit <= 0 { limit = 200 } now := GetDBTimestamp() var channels []Channel if err := DB.Where("next_reset_time > 0 AND next_reset_time <= ? AND status = ? AND quota_limit > 0", now, common.ChannelStatusEnabled). Order("next_reset_time asc"). Limit(limit). Find(&channels).Error; err != nil { return 0, err } if len(channels) == 0 { return 0, nil } resetCount := 0 for _, ch := range channels { chCopy := ch err := DB.Transaction(func(tx *gorm.DB) error { var locked Channel if err := tx.Set("gorm:query_option", "FOR UPDATE"). Where("id = ? AND next_reset_time > 0 AND next_reset_time <= ?", chCopy.Id, now). First(&locked).Error; err != nil { return nil } if err := maybeResetChannelQuotaTx(tx, &locked, now); err != nil { return err } resetCount++ return nil }) if err != nil { return resetCount, err } } return resetCount, nil } func CheckChannelQuota(channelId int) error { var channel Channel if err := DB.Select("id, used_quota, quota_limit, quota_reset_period, next_reset_time"). Where("id = ?", channelId).First(&channel).Error; err != nil { return err } if channel.QuotaLimit <= 0 { return nil } if channel.NextResetTime > 0 { now := GetDBTimestamp() if channel.NextResetTime <= now { _ = DB.Transaction(func(tx *gorm.DB) error { var locked Channel if err := tx.Set("gorm:query_option", "FOR UPDATE"). Where("id = ? AND next_reset_time > 0 AND next_reset_time <= ?", channelId, now). First(&locked).Error; err != nil { return nil } return maybeResetChannelQuotaTx(tx, &locked, now) }) if err := DB.Select("used_quota, quota_limit"). Where("id = ?", channelId).First(&channel).Error; err != nil { return err } } } if channel.UsedQuota >= channel.QuotaLimit { return fmt.Errorf("channel quota exceeded (used=%d, limit=%d)", channel.UsedQuota, channel.QuotaLimit) } return nil } func InitChannelNextResetTime(channel *Channel) { if channel.QuotaLimit <= 0 { channel.NextResetTime = 0 channel.LastResetTime = 0 return } period := NormalizeChannelResetPeriod(channel.QuotaResetPeriod) if period == ChannelResetNever { channel.NextResetTime = 0 channel.LastResetTime = 0 return } baseUnix := channel.LastResetTime if baseUnix <= 0 { baseUnix = channel.CreatedTime } if baseUnix <= 0 { baseUnix = GetDBTimestamp() } base := time.Unix(baseUnix, 0) next := calcChannelNextResetTime(base, channel.QuotaResetPeriod, channel.QuotaResetCustomSeconds) channel.NextResetTime = next if channel.LastResetTime <= 0 { channel.LastResetTime = baseUnix } }