limti

go.mod

@@ -4,6 +4,7 @@ go 1.25
 
 require (
 	github.com/fogleman/gg v1.3.0
+	github.com/go-redis/redis v6.15.9+incompatible
 	github.com/sirupsen/logrus v1.9.3
 	github.com/speps/go-hashids v2.0.0+incompatible
 	gorm.io/gorm v1.31.0
@@ -13,7 +14,9 @@ require (
 	github.com/golang/freetype v0.0.0-20170609003504-e2365dfdc4a0 // indirect
 	github.com/jinzhu/inflection v1.0.0 // indirect
 	github.com/jinzhu/now v1.1.5 // indirect
+	github.com/onsi/ginkgo v1.16.5 // indirect
+	github.com/onsi/gomega v1.40.0 // indirect
 	golang.org/x/image v0.31.0 // indirect
 	golang.org/x/sys v0.0.0-20220715151400-c0bba94af5f8 // indirect
-	golang.org/x/text v0.29.0 // indirect
+	golang.org/x/text v0.33.0 // indirect
 )

limit.go

@@ -0,0 +1,92 @@
+package goutil
+
+import (
+	"fmt"
+	"math/rand"
+	"time"
+
+	"github.com/go-redis/redis"
+)
+
+type WindowsLimitOption struct {
+	keyPrefix  string        //redis可以前缀，会拼接上succ和fail
+	limitCount int64         //限制数量
+	scope      time.Duration //限制时间区间，默认一分钟
+	succKey    string
+	failKey    string
+}
+type WindowsLimit struct {
+	client *redis.Client
+	option *WindowsLimitOption
+}
+
+func NewWindowsLimitOption() *WindowsLimitOption {
+	return &WindowsLimitOption{
+		keyPrefix:  fmt.Sprintf("windows:limit:%s:%d", time.Now().Format("20060102"), time.Now().UnixMilli()),
+		limitCount: 1000,
+		scope:      time.Second * 60,
+	}
+}
+
+func NewWindowsLimit(client *redis.Client, option *WindowsLimitOption) *WindowsLimit {
+	return &WindowsLimit{
+		client: client,
+		option: option,
+	}
+}
+
+func (w *WindowsLimit) tidy() {
+	now := time.Now().UnixMilli()
+	start := now - w.option.scope.Milliseconds()
+
+	pipe := w.client.Pipeline()
+	pipe.ZRemRangeByScore(w.option.failKey, "0", fmt.Sprintf("%d", start))
+	pipe.ZRemRangeByScore(w.option.succKey, "0", fmt.Sprintf("%d", start))
+	pipe.Exec()
+}
+
+func (w *WindowsLimit) Count() (int64, int64, error) {
+	now := time.Now().UnixMilli()
+	start := now - w.option.scope.Milliseconds()
+
+	pipe := w.client.Pipeline()
+	succCmd := pipe.ZCount(w.option.succKey, fmt.Sprintf("%d", start), fmt.Sprintf("%d", now))
+	failCmd := pipe.ZCount(w.option.failKey, fmt.Sprintf("%d", start), fmt.Sprintf("%d", now))
+	_, err := pipe.Exec()
+	if err != nil {
+		return 0, 0, err
+	}
+	return succCmd.Val(), failCmd.Val(), nil
+}
+
+func (w *WindowsLimit) Allow() bool {
+	now := time.Now().UnixMilli()
+	start := now - w.option.scope.Milliseconds()
+	pipe := w.client.Pipeline()
+	succCmd := w.client.ZCount(w.option.succKey, fmt.Sprintf("%d", start), fmt.Sprintf("%d", now))
+	_, err := pipe.Exec()
+	if err != nil {
+		return false
+	}
+
+	if succCmd.Val() >= w.option.limitCount {
+		return false
+	}
+	return true
+}
+
+func (w *WindowsLimit) Add(succ bool) error {
+	now := time.Now().UnixMilli()
+	w.tidy()
+	pipe := w.client.Pipeline()
+	if succ {
+		pipe.ZAdd(w.option.succKey, redis.Z{Score: float64(now), Member: fmt.Sprintf("%d:%d", now, rand.Int())})
+	} else {
+		pipe.ZAdd(w.option.failKey, redis.Z{Score: float64(now), Member: fmt.Sprintf("%d:%d", now, rand.Int())})
+	}
+
+	pipe.Expire(w.option.failKey, w.option.scope)
+	pipe.Expire(w.option.succKey, w.option.scope)
+	_, err := pipe.Exec()
+	return err
+}