55 lines
1.2 KiB
Go
55 lines
1.2 KiB
Go
// Package rate_limit ...
|
|
//
|
|
// Description : rate_limit ...
|
|
//
|
|
// Author : go_developer@163.com<白茶清欢>
|
|
//
|
|
// Date : 2024-06-20 20:40
|
|
package rate_limit
|
|
|
|
import (
|
|
"context"
|
|
"git.zhangdeman.cn/zhangdeman/rate_limit/define"
|
|
"golang.org/x/time/rate"
|
|
"sync"
|
|
"time"
|
|
)
|
|
|
|
var (
|
|
MemoryClient = &Memory{
|
|
lock: &sync.RWMutex{},
|
|
limiterTable: make(map[string]*rate.Limiter),
|
|
}
|
|
)
|
|
|
|
// Memory ...
|
|
//
|
|
// Author : go_developer@163.com<白茶清欢>
|
|
//
|
|
// Date : 20:40 2024/6/20
|
|
type Memory struct {
|
|
lock *sync.RWMutex
|
|
limiterTable map[string]*rate.Limiter
|
|
}
|
|
|
|
func (m *Memory) AllowN(ctx context.Context, limitCfg *define.LimitConfig, tokenCnt int) (bool, error) {
|
|
limiter := m.getLimiter(limitCfg)
|
|
return limiter.AllowN(time.Now(), tokenCnt), nil
|
|
}
|
|
|
|
func (m *Memory) Reset(ctx context.Context, limitCfg *define.LimitConfig) error {
|
|
m.lock.Lock()
|
|
defer m.lock.Unlock()
|
|
delete(m.limiterTable, limitCfg.Key)
|
|
return nil
|
|
}
|
|
|
|
func (m *Memory) getLimiter(limitCfg *define.LimitConfig) *rate.Limiter {
|
|
m.lock.Lock()
|
|
defer m.lock.Unlock()
|
|
if _, exist := m.limiterTable[limitCfg.Key]; !exist {
|
|
m.limiterTable[limitCfg.Key] = rate.NewLimiter(rate.Limit(limitCfg.Rate), limitCfg.Total)
|
|
}
|
|
return m.limiterTable[limitCfg.Key]
|
|
}
|