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