限速器

gin/middleware/sync_ratelimiter/sync_ratelimiter.go

+package sync_ratelimiter
+
+import (
+	"fmt"
+	"github.com/RussellLuo/slidingwindow"
+	"github.com/go-redis/redis"
+	"strconv"
+	"sync"
+	"time"
+)
+
+type RedisDatastore struct {
+	client redis.Cmdable
+	ttl    time.Duration
+}
+
+func (d *RedisDatastore) fullKey(key string, start int64) string {
+	return fmt.Sprintf("%s@%d", key, start)
+}
+
+func (d *RedisDatastore) Add(key string, start, value int64) (int64, error) {
+	k := d.fullKey(key, start)
+	c, err := d.client.IncrBy(k, value).Result()
+	if err != nil {
+		return 0, err
+	}
+	_, _ = d.client.Expire(k, d.ttl).Result()
+	return c, err
+}
+
+func (d *RedisDatastore) Get(key string, start int64) (int64, error) {
+	k := d.fullKey(key, start)
+	value, err := d.client.Get(k).Result()
+	if err != nil {
+		if err == redis.Nil {
+			err = nil
+		}
+		return 0, err
+	}
+	return strconv.ParseInt(value, 10, 64)
+}
+
+func NewRedisDatastore(client redis.Cmdable, ttl time.Duration) *RedisDatastore {
+	return &RedisDatastore{client: client, ttl: ttl}
+}
+
+type IPRateLimiter struct {
+	mu      sync.Mutex
+	options *IPRateLimiterOptions
+	store   *RedisDatastore
+	ips     map[string]*Limiter
+}
+
+type IPRateLimiterOptions struct {
+	windowSize    time.Duration
+	maxRequests   int64
+	redisAddr     string
+	redisDB       int
+	redisPassword string
+}
+
+func defaultOptions() *IPRateLimiterOptions {
+	return &IPRateLimiterOptions{
+		windowSize:    time.Minute,
+		maxRequests:   60 * 10,
+		redisAddr:     "localhost:6379",
+		redisPassword: "",
+		redisDB:       0,
+	}
+}
+
+func (x *IPRateLimiterOptions) apply(opts ...Option) {
+	for _, opt := range opts {
+		opt(x)
+	}
+}
+
+func WithWindowSize(windowSize time.Duration) Option {
+	return func(o *IPRateLimiterOptions) {
+		o.windowSize = windowSize
+	}
+}
+
+func WithMaxRequests(maxRequests int64) Option {
+	return func(o *IPRateLimiterOptions) {
+		o.maxRequests = maxRequests
+	}
+}
+
+func WithRedisAddr(redisAddr string) Option {
+	return func(o *IPRateLimiterOptions) {
+		o.redisAddr = redisAddr
+	}
+}
+func WithRedisDB(redisDB int) Option {
+	return func(o *IPRateLimiterOptions) {
+		o.redisDB = redisDB
+	}
+}
+
+func WithRedisPassword(pwd string) Option {
+	return func(o *IPRateLimiterOptions) {
+		o.redisPassword = pwd
+	}
+}
+
+type Option func(o *IPRateLimiterOptions)
+
+func NewIPRateLimiter(opts ...Option) *IPRateLimiter {
+	options := defaultOptions()
+	options.apply(opts...)
+	return &IPRateLimiter{
+		mu:      sync.Mutex{},
+		options: options,
+		store: NewRedisDatastore(
+			redis.NewClient(&redis.Options{
+				Addr:     options.redisAddr,
+				Password: options.redisPassword,
+				DB:       options.redisDB,
+			}),
+			2*options.windowSize, // twice of window-size is just enough.
+		),
+		ips: make(map[string]*Limiter),
+	}
+}
+
+type Limiter struct {
+	L    *slidingwindow.Limiter
+	Stop func()
+}
+
+func (x *IPRateLimiter) Allow(ip string) bool {
+	x.mu.Lock()
+	defer x.mu.Unlock()
+	if v, ok := x.ips[ip]; ok {
+		return v.L.Allow()
+	} else {
+		lim, stop := slidingwindow.NewLimiter(x.options.windowSize, x.options.maxRequests, func() (slidingwindow.Window, slidingwindow.StopFunc) {
+			return slidingwindow.NewSyncWindow(fmt.Sprintf("SyncRateLimiter_%s", ip), slidingwindow.NewBlockingSynchronizer(x.store, 500*time.Millisecond))
+		})
+		x.ips[ip] = &Limiter{
+			L:    lim,
+			Stop: stop,
+		}
+		return lim.Allow()
+	}
+}
+
+func (x *IPRateLimiter) Stop() {
+	x.mu.Lock()
+	defer x.mu.Unlock()
+	for _, l := range x.ips {
+		l.Stop()
+	}
+}

gin/middleware/sync_ratelimiter/sync_ratelimiter_test.go

+package sync_ratelimiter
+
+import (
+	"gitlab.wanzhuangkj.com/tush/xpkg/utils/xtime"
+	"testing"
+	"time"
+)
+
+import (
+	"fmt"
+)
+
+func TestRateLimiter(t *testing.T) {
+	rl := NewIPRateLimiter(
+		WithWindowSize(1*time.Second),
+		WithMaxRequests(4),
+		WithRedisAddr("localhost:6379"),
+		WithRedisDB(0),
+		WithRedisPassword(""),
+	)
+	ip := "192.168.1.1"
+	for i := 0; i < 100; i++ {
+		allowed := rl.Allow(ip)
+		fmt.Printf("%s Request %d: %t\n", time.Now().Format(xtime.Layout_DateTime), i+1, allowed)
+		time.Sleep(100 * time.Millisecond)
+	}
+}

go.mod

 module gitlab.wanzhuangkj.com/tush/xpkg
 
-go 1.22.5
+go 1.23.0
+
+toolchain go1.23.9
 
 require (
 	github.com/DATA-DOG/go-sqlmock v1.5.0
@@ -54,10 +56,10 @@ require (
 	go.opentelemetry.io/otel/sdk v1.24.0
 	go.opentelemetry.io/otel/trace v1.24.0
 	go.uber.org/zap v1.24.0
-	golang.org/x/crypto v0.31.0
-	golang.org/x/sync v0.10.0
+	golang.org/x/crypto v0.36.0
+	golang.org/x/sync v0.12.0
 	google.golang.org/grpc v1.67.1
-	google.golang.org/protobuf v1.34.2
+	google.golang.org/protobuf v1.36.5
 	gopkg.in/yaml.v3 v3.0.1
 	gorm.io/driver/mysql v1.5.2
 	gorm.io/driver/postgres v1.5.4
@@ -67,14 +69,19 @@ require (
 )
 
 require (
+	github.com/RussellLuo/slidingwindow v0.0.0-20200528002341-535bb99d338b
 	github.com/bwmarrin/snowflake v0.3.0
 	github.com/duke-git/lancet/v2 v2.3.4
+	github.com/go-redis/redis v6.15.9+incompatible
 	github.com/mojocn/base64Captcha v1.3.8
 	gopkg.in/gomail.v2 v2.0.0-20160411212932-81ebce5c23df
 )
 
 require (
 	github.com/golang/freetype v0.0.0-20170609003504-e2365dfdc4a0 // indirect
+	github.com/google/go-cmp v0.7.0 // indirect
+	github.com/onsi/ginkgo v1.16.5 // indirect
+	github.com/onsi/gomega v1.37.0 // indirect
 	github.com/rogpeppe/go-internal v1.9.0 // indirect
 	golang.org/x/image v0.23.0 // indirect
 	gopkg.in/alexcesaro/quotedprintable.v3 v3.0.0-20150716171945-2caba252f4dc // indirect
@@ -122,7 +129,7 @@ require (
 	github.com/fatih/color v1.13.0 // indirect
 	github.com/gabriel-vasile/mimetype v1.4.3 // indirect
 	github.com/gin-contrib/sse v0.1.0 // indirect
-	github.com/go-logr/logr v1.4.1 // indirect
+	github.com/go-logr/logr v1.4.2 // indirect
 	github.com/go-logr/stdr v1.2.2 // indirect
 	github.com/go-ole/go-ole v1.2.6 // indirect
 	github.com/go-openapi/jsonpointer v0.19.5 // indirect
@@ -136,7 +143,7 @@ require (
 	github.com/golang/glog v1.2.2 // indirect
 	github.com/golang/mock v1.6.0 // indirect
 	github.com/golang/protobuf v1.5.4 // indirect
-	github.com/google/pprof v0.0.0-20211214055906-6f57359322fd // indirect
+	github.com/google/pprof v0.0.0-20241210010833-40e02aabc2ad // indirect
 	github.com/google/uuid v1.6.0
 	github.com/hashicorp/errwrap v1.1.0 // indirect
 	github.com/hashicorp/go-cleanhttp v0.5.2 // indirect
@@ -217,12 +224,12 @@ require (
 	go.uber.org/atomic v1.9.0 // indirect
 	go.uber.org/multierr v1.9.0 // indirect
 	golang.org/x/arch v0.7.0 // indirect
-	golang.org/x/net v0.30.0 // indirect
+	golang.org/x/net v0.37.0 // indirect
 	golang.org/x/oauth2 v0.22.0 // indirect
-	golang.org/x/sys v0.28.0 // indirect
-	golang.org/x/text v0.21.0 // indirect
+	golang.org/x/sys v0.31.0 // indirect
+	golang.org/x/text v0.23.0 // indirect
 	golang.org/x/time v0.1.0 // indirect
-	golang.org/x/tools v0.21.1-0.20240508182429-e35e4ccd0d2d // indirect
+	golang.org/x/tools v0.30.0 // indirect
 	google.golang.org/appengine v1.6.8 // indirect
 	google.golang.org/genproto/googleapis/api v0.0.0-20240814211410-ddb44dafa142 // indirect
 	google.golang.org/genproto/googleapis/rpc v0.0.0-20240814211410-ddb44dafa142 // indirect