ez-api/internal/cron/quota_reset.go

package cron

import (
	"context"
	"log/slog"
	"strings"
	"time"

	"github.com/ez-api/ez-api/internal/model"
	"github.com/ez-api/ez-api/internal/service"
	"gorm.io/gorm"
)

type QuotaResetter struct {
	db       *gorm.DB
	sync     *service.SyncService
	interval time.Duration
}

func NewQuotaResetter(db *gorm.DB, sync *service.SyncService, interval time.Duration) *QuotaResetter {
	if interval <= 0 {
		interval = 5 * time.Minute
	}
	return &QuotaResetter{db: db, sync: sync, interval: interval}
}

func (q *QuotaResetter) Start(ctx context.Context) {
	if q == nil || q.db == nil {
		return
	}
	if ctx == nil {
		ctx = context.Background()
	}
	ticker := time.NewTicker(q.interval)
	defer ticker.Stop()

	for {
		select {
		case <-ctx.Done():
			return
		case <-ticker.C:
			if err := q.resetOnce(ctx); err != nil {
				slog.Default().Warn("quota reset failed", "err", err)
			}
		}
	}
}

func (q *QuotaResetter) resetOnce(ctx context.Context) error {
	if q == nil || q.db == nil {
		return nil
	}
	now := time.Now().UTC()
	var keys []model.Key
	if err := q.db.Where("quota_reset_type IN ? AND (quota_reset_at IS NULL OR quota_reset_at <= ?)", []string{"daily", "monthly"}, now).Find(&keys).Error; err != nil {
		return err
	}
	for i := range keys {
		resetType := strings.ToLower(strings.TrimSpace(keys[i].QuotaResetType))
		nextAt, ok := nextQuotaReset(now, resetType)
		if !ok {
			continue
		}
		if err := q.db.Model(&keys[i]).Updates(map[string]any{
			"quota_used":     0,
			"quota_reset_at": nextAt,
		}).Error; err != nil {
			slog.Default().Warn("quota reset update failed", "key_id", keys[i].ID, "err", err)
			continue
		}
		keys[i].QuotaUsed = 0
		keys[i].QuotaResetAt = &nextAt
		if q.sync != nil {
			_ = q.sync.SyncKey(&keys[i])
		}
	}
	return nil
}

func nextQuotaReset(now time.Time, resetType string) (time.Time, bool) {
	now = now.UTC()
	switch resetType {
	case "daily":
		next := time.Date(now.Year(), now.Month(), now.Day(), 0, 0, 0, 0, time.UTC).AddDate(0, 0, 1)
		return next, true
	case "monthly":
		next := time.Date(now.Year(), now.Month(), 1, 0, 0, 0, 0, time.UTC).AddDate(0, 1, 0)
		return next, true
	default:
		return time.Time{}, false
	}
}