From d509496d72a55f101e40b1260d98bcea6ea9e3d1 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E7=99=BD=E8=8C=B6=E6=B8=85=E6=AC=A2?=
 <go_developer@163.com>
Date: Thu, 20 Jun 2024 20:58:15 +0800
Subject: [PATCH] =?UTF-8?q?=E6=94=AF=E6=8C=81=E5=86=85=E5=AD=98=E7=9A=84?=
 =?UTF-8?q?=E6=B5=81=E6=8E=A7=E5=AE=9E=E7=8E=B0?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 go.mod         |  1 +
 go.sum         |  2 ++
 memory.go      | 54 ++++++++++++++++++++++++++++++++++++++++++++++++++
 memory_test.go | 27 +++++++++++++++++++++++++
 4 files changed, 84 insertions(+)
 create mode 100644 memory.go
 create mode 100644 memory_test.go

diff --git a/go.mod b/go.mod
index b3e5537..926c626 100644
--- a/go.mod
+++ b/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
 )
diff --git a/go.sum b/go.sum
index 9664f02..69ef50f 100644
--- a/go.sum
+++ b/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=
diff --git a/memory.go b/memory.go
new file mode 100644
index 0000000..c61dc1c
--- /dev/null
+++ b/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]
+}
diff --git a/memory_test.go b/memory_test.go
new file mode 100644
index 0000000..9dfdb94
--- /dev/null
+++ b/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))
+	}
+}