aiforge/services/task/limiter.go

package task

import (
	"code.gitea.io/gitea/models"
	"time"
)

var LimiterMap = map[string]Limiter{
	models.TaskConfigRefreshRateNotCycle: new(NoCycleLimiter),
	models.TaskConfigRefreshRateDaily:    new(DailyLimiter),
}

type Limiter interface {
	GetCurrentPeriod() *models.LimiterPeriod
}

type NoCycleLimiter struct {
}

func (l *NoCycleLimiter) GetCurrentPeriod() *models.LimiterPeriod {
	return nil
}

type DailyLimiter struct {
}

func (l *DailyLimiter) GetCurrentPeriod() *models.LimiterPeriod {
	t := time.Now()
	startTime := time.Date(t.Year(), t.Month(), t.Day(), 0, 0, 0, 0, t.Location())
	endTime := startTime.Add(24 * time.Hour)
	return &models.LimiterPeriod{
		StartTime: startTime,
		EndTime:   endTime,
	}
}

func GetLimiter(refreshRateype string) Limiter {
	return LimiterMap[refreshRateype]
}