支持内存的流控实现

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