fusenapi/utils/check/limit.go

// Timeup是一个限制某个值的重复频率的工具
// 通过内部计时器实现限制同一个值的访问频率不超过指定时间间隔
package check

import (
	"sync"
	"time"
)

// TimeLimit结构体包含需要限频的映射表,以及时间间隔
// 使用sync.Mutex保证对字典的线程安全访问
type TimeLimit[T comparable] struct {
	mu   sync.Mutex
	dict map[T]struct{}
	dur  time.Duration
}

// NewTimelimit构造函数,接收限频的时间间隔
// 并初始化内部字典和间隔字段
func NewTimelimit[T comparable](dur time.Duration) *TimeLimit[T] {
	return &TimeLimit[T]{
		dict: make(map[T]struct{}),
		dur:  dur,
	}
}

// WithTime方法用于更新限频的时间间隔
func (tup *TimeLimit[T]) WithTime(dur time.Duration) *TimeLimit[T] {
	tup.mu.Lock()
	defer tup.mu.Unlock()
	tup.dur = dur
	return tup
}

// Is方法检查传入值是否是一个新的值
// 首先会查询字典,如果存在则表示在间隔内已经访问过
// 否则将其添加到字典,并启动一个定时器在间隔后删除
// 返回true表示新值,false表示重复值
func (tup *TimeLimit[T]) Is(v T) bool {
	tup.mu.Lock()
	defer tup.mu.Unlock()

	if _, ok := tup.dict[v]; ok {
		return false
	}
	tup.dict[v] = struct{}{}
	time.AfterFunc(tup.dur, func() {
		tup.mu.Lock()
		defer tup.mu.Unlock()
		delete(tup.dict, v)
	})
	return true
}