支持内存的流控实现

This commit is contained in:
白茶清欢 2024-06-20 20:58:15 +08:00
parent 7c31a1111e
commit d509496d72
4 changed files with 84 additions and 0 deletions

1
go.mod
View File

@ -29,5 +29,6 @@ require (
github.com/tidwall/pretty v1.2.1 // indirect github.com/tidwall/pretty v1.2.1 // indirect
go.uber.org/multierr v1.11.0 // indirect go.uber.org/multierr v1.11.0 // indirect
go.uber.org/zap v1.27.0 // indirect go.uber.org/zap v1.27.0 // indirect
golang.org/x/time v0.5.0 // indirect
gopkg.in/yaml.v3 v3.0.1 // indirect gopkg.in/yaml.v3 v3.0.1 // indirect
) )

2
go.sum
View File

@ -57,6 +57,8 @@ go.uber.org/multierr v1.11.0 h1:blXXJkSxSSfBVBlC76pxqeO+LN3aDfLQo+309xJstO0=
go.uber.org/multierr v1.11.0/go.mod h1:20+QtiLqy0Nd6FdQB9TLXag12DsQkrbs3htMFfDN80Y= go.uber.org/multierr v1.11.0/go.mod h1:20+QtiLqy0Nd6FdQB9TLXag12DsQkrbs3htMFfDN80Y=
go.uber.org/zap v1.27.0 h1:aJMhYGrd5QSmlpLMr2MftRKl7t8J8PTZPA732ud/XR8= go.uber.org/zap v1.27.0 h1:aJMhYGrd5QSmlpLMr2MftRKl7t8J8PTZPA732ud/XR8=
go.uber.org/zap v1.27.0/go.mod h1:GB2qFLM7cTU87MWRP2mPIjqfIDnGu+VIO4V/SdhGo2E= go.uber.org/zap v1.27.0/go.mod h1:GB2qFLM7cTU87MWRP2mPIjqfIDnGu+VIO4V/SdhGo2E=
golang.org/x/time v0.5.0 h1:o7cqy6amK/52YcAKIPlM3a+Fpj35zvRj2TP+e1xFSfk=
golang.org/x/time v0.5.0/go.mod h1:3BpzKBy/shNhVucY/MWOyx10tF3SFh9QdLuxbVysPQM=
gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+SDo693bJlVdllGtEeKM= gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+SDo693bJlVdllGtEeKM=
gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA= gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA=

54
memory.go Normal file
View File

@ -0,0 +1,54 @@
// 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]
}

27
memory_test.go Normal file
View File

@ -0,0 +1,27 @@
// Package rate_limit ...
//
// Description : rate_limit ...
//
// Author : go_developer@163.com<白茶清欢>
//
// Date : 2024-06-20 20:55
package rate_limit
import (
"context"
"fmt"
"git.zhangdeman.cn/zhangdeman/rate_limit/define"
"testing"
)
func TestMemory_AllowN(t *testing.T) {
for i := 0; i < 10; i++ {
fmt.Println(MemoryClient.AllowN(context.Background(), &define.LimitConfig{
Key: "deman",
Total: 15,
Rate: 10,
TimeInterval: 10,
}, 8))
}
}