new-api/model/channel_quota.go

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
	}
}