支持内存的流控实现

go.mod

@@ -29,5 +29,6 @@ require (
 	github.com/tidwall/pretty v1.2.1 // indirect
 	go.uber.org/multierr v1.11.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
 )

go.sum

@@ -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/zap v1.27.0 h1:aJMhYGrd5QSmlpLMr2MftRKl7t8J8PTZPA732ud/XR8=
 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/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
 gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA=

memory.go

@@ -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]
+}

memory_test.go

@@ -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))
+	}
+}