first commit

pkg/limiter/rate.go

@@ -0,0 +1,71 @@
+package limiter
+
+import (
+	"git.bvbej.com/bvbej/base-golang/pkg/ticker"
+	"golang.org/x/time/rate"
+	"sync"
+	"time"
+)
+
+var _ RateLimiter = (*rateLimiter)(nil)
+
+type item struct {
+	lastTime time.Time
+	limiter  *rate.Limiter
+}
+
+type RateLimiter interface {
+	set(key string) *rate.Limiter
+	get(key string) *rate.Limiter
+	Allow(key string) bool
+}
+
+type rateLimiter struct {
+	limit   rate.Limit
+	burst   int
+	list    *sync.Map
+	recycle ticker.Ticker
+}
+
+func NewRateLimiter(limit rate.Limit, burst int) RateLimiter {
+	list := new(sync.Map)
+	t := ticker.New(time.Minute)
+	t.Process(func() {
+		list.Range(func(key, value any) bool {
+			if value.(*item).lastTime.Before(time.Now().Add(-time.Hour)) {
+				list.Delete(key)
+			}
+			return true
+		})
+	})
+	return &rateLimiter{
+		list:    list,
+		limit:   limit,
+		recycle: t,
+		burst:   burst,
+	}
+}
+
+func (i *rateLimiter) set(key string) *rate.Limiter {
+	store := &item{
+		lastTime: time.Now(),
+		limiter:  rate.NewLimiter(i.limit, i.burst),
+	}
+	i.list.Store(key, store)
+	return store.limiter
+}
+
+func (i *rateLimiter) get(key string) *rate.Limiter {
+	value, ok := i.list.Load(key)
+	if !ok {
+		return i.set(key)
+	}
+	value.(*item).lastTime = time.Now()
+	i.list.Store(key, value)
+	return value.(*item).limiter
+}
+
+func (i *rateLimiter) Allow(key string) bool {
+	limiter := i.get(key)
+	return limiter.Allow()
+}