ratelimiter

package module
v0.2.0 Latest Latest
Warning

This package is not in the latest version of its module.

 Go to latest Published: Aug 19, 2024 License: MIT Imports: 11 Imported by: 0

README

ratelimiter

Currently only supports Redis / GORM as the driver now.

package ratelimiter_test

import (
	"context"
	"fmt"
	"time"

	"github.com/theplant/ratelimiter"
)

func runExample(limiter *ratelimiter.RateLimiter, key string) {
	// every 10 min , burst 5
	durationPerToken := 10 * time.Minute
	burst := 5
	now := time.Now()

	ctx := context.Background()

	try := func(delta time.Duration) bool {
		reserveReq := &ratelimiter.ReserveRequest{
			Key:              key,
			DurationPerToken: durationPerToken,
			Burst:            burst,
			Tokens:           1,
			MaxFutureReserve: 0,
		}
		advancedNow := now.Add(delta)
		r, err := limiter.Reserve(
			// only for test, you should not use this in production !!
			ratelimiter.WithNowFuncForTest(ctx, func() time.Time {
				return advancedNow
			}),
			reserveReq,
		)
		if err != nil {
			panic(err)
		}

		if r.OK {
			fmt.Printf("%v: allowed: %t\n", delta, r.OK)
			return true
		}

		fmt.Printf("%v: allowed: %t , you can retry after %v\n", delta, false, r.RetryAfterFrom(advancedNow))
		return false
	}

	for i := 0; i < int(25); i++ {
		delta := time.Duration(i) * time.Minute
		try(delta)
	}

	fmt.Printf("--- Sleep 20 minutes ---\n")

	for i := 45; i < int(55); i++ {
		delta := time.Duration(i) * time.Minute
		try(delta)
	}

	fmt.Printf("--- Sleep 100 minutes ---\n")

	for i := 155; i < int(165); i++ {
		delta := time.Duration(i) * time.Minute
		try(delta)
	}
}


func ExampleInitRedisDriver() {
	d, err := ratelimiter.InitRedisDriver(context.Background(), redisCli)
	if err != nil {
		panic(err)
	}
	limiter := ratelimiter.New(d)
	runExample(limiter, "ExampleInitRedisDriver")
	// Output:
	// 0s: allowed: true
	// 1m0s: allowed: true
	// 2m0s: allowed: true
	// 3m0s: allowed: true
	// 4m0s: allowed: true
	// 5m0s: allowed: false , you can retry after 5m0s
	// 6m0s: allowed: false , you can retry after 4m0s
	// 7m0s: allowed: false , you can retry after 3m0s
	// 8m0s: allowed: false , you can retry after 2m0s
	// 9m0s: allowed: false , you can retry after 1m0s
	// 10m0s: allowed: true
	// 11m0s: allowed: false , you can retry after 9m0s
	// 12m0s: allowed: false , you can retry after 8m0s
	// 13m0s: allowed: false , you can retry after 7m0s
	// 14m0s: allowed: false , you can retry after 6m0s
	// 15m0s: allowed: false , you can retry after 5m0s
	// 16m0s: allowed: false , you can retry after 4m0s
	// 17m0s: allowed: false , you can retry after 3m0s
	// 18m0s: allowed: false , you can retry after 2m0s
	// 19m0s: allowed: false , you can retry after 1m0s
	// 20m0s: allowed: true
	// 21m0s: allowed: false , you can retry after 9m0s
	// 22m0s: allowed: false , you can retry after 8m0s
	// 23m0s: allowed: false , you can retry after 7m0s
	// 24m0s: allowed: false , you can retry after 6m0s
	// --- Sleep 20 minutes ---
	// 45m0s: allowed: true
	// 46m0s: allowed: true
	// 47m0s: allowed: false , you can retry after 3m0s
	// 48m0s: allowed: false , you can retry after 2m0s
	// 49m0s: allowed: false , you can retry after 1m0s
	// 50m0s: allowed: true
	// 51m0s: allowed: false , you can retry after 9m0s
	// 52m0s: allowed: false , you can retry after 8m0s
	// 53m0s: allowed: false , you can retry after 7m0s
	// 54m0s: allowed: false , you can retry after 6m0s
	// --- Sleep 100 minutes ---
	// 2h35m0s: allowed: true
	// 2h36m0s: allowed: true
	// 2h37m0s: allowed: true
	// 2h38m0s: allowed: true
	// 2h39m0s: allowed: true
	// 2h40m0s: allowed: false , you can retry after 5m0s
	// 2h41m0s: allowed: false , you can retry after 4m0s
	// 2h42m0s: allowed: false , you can retry after 3m0s
	// 2h43m0s: allowed: false , you can retry after 2m0s
	// 2h44m0s: allowed: false , you can retry after 1m0s
}
Benchmark
goos: darwin
goarch: arm64
pkg: github.com/theplant/ratelimiter
BenchmarkDriverRedis_Reserve
BenchmarkDriverRedis_Reserve/Key1_Duration10ms_Burst5
BenchmarkDriverRedis_Reserve/Key1_Duration10ms_Burst5-12         	    4045	    274446 ns/op	     692 B/op	      16 allocs/op
BenchmarkDriverRedis_Reserve/Key2_Duration20ms_Burst10
BenchmarkDriverRedis_Reserve/Key2_Duration20ms_Burst10-12        	    4638	    277619 ns/op	     688 B/op	      16 allocs/op
BenchmarkDriverRedis_Reserve/Key3_Duration50ms_Burst3
BenchmarkDriverRedis_Reserve/Key3_Duration50ms_Burst3-12         	    4406	    274867 ns/op	     688 B/op	      16 allocs/op

BenchmarkDriverGORM_Reserve
BenchmarkDriverGORM_Reserve/Key1_Duration10ms_Burst5
BenchmarkDriverGORM_Reserve/Key1_Duration10ms_Burst5-12         	    1108	    909553 ns/op	   10346 B/op	     139 allocs/op
BenchmarkDriverGORM_Reserve/Key2_Duration20ms_Burst10
BenchmarkDriverGORM_Reserve/Key2_Duration20ms_Burst10-12        	    1152	   1061490 ns/op	   12973 B/op	     166 allocs/op
BenchmarkDriverGORM_Reserve/Key3_Duration50ms_Burst3
BenchmarkDriverGORM_Reserve/Key3_Duration50ms_Burst3-12         	    1156	   1064777 ns/op	   12999 B/op	     165 allocs/op

Documentation

Index

Examples

Constants

This section is empty.

Variables

View Source 
var ErrInvalidParameters = errors.New("ratelimiter: invalid parameters")
View Source 
var Test = false

Functions

func NowFuncFromContextForTest added in v0.2.0 

func NowFuncFromContextForTest(ctx context.Context) (func() time.Time, bool)

func WithNowFuncForTest added in v0.2.0 

func WithNowFuncForTest(ctx context.Context, nowFunc func() time.Time) context.Context

Types

type AllowRequest 

type AllowRequest struct {
	Key              string
	DurationPerToken time.Duration
	Burst            int
	Tokens           int
}

type Driver 

type Driver interface {
	Reserve(ctx context.Context, req *ReserveRequest) (*Reservation, error)
}

type DriverFunc 

type DriverFunc func(ctx context.Context, req *ReserveRequest) (*Reservation, error)

func (DriverFunc) Reserve 

func (f DriverFunc) Reserve(ctx context.Context, req *ReserveRequest) (*Reservation, error)

type GormDriver added in v0.2.0 

type GormDriver struct {
	// contains filtered or unexported fields
}

func InitGormDriver added in v0.2.0 

func InitGormDriver(ctx context.Context, db *gorm.DB) (*GormDriver, error)

InitGormDriver initializes a GormDriver with the provided Gorm DB. Sometimes you may not need to auto migrate the KV table, you can use `NewGormDriver` instead.

func NewGormDriver added in v0.2.0 

func NewGormDriver(db *gorm.DB) *GormDriver

NewGormDriver returns a Driver that uses Gorm as the storage. Sometimes you may need to auto migrate the KV table, you can use `InitGormDriver` instead.

Example 
limiter := New(
	NewGormDriver(db),
)
runExample(limiter, "ExampleNewGormDriver")

Output:

0s: allowed: true
1m0s: allowed: true
2m0s: allowed: true
3m0s: allowed: true
4m0s: allowed: true
5m0s: allowed: false , you can retry after 5m0s
6m0s: allowed: false , you can retry after 4m0s
7m0s: allowed: false , you can retry after 3m0s
8m0s: allowed: false , you can retry after 2m0s
9m0s: allowed: false , you can retry after 1m0s
10m0s: allowed: true
11m0s: allowed: false , you can retry after 9m0s
12m0s: allowed: false , you can retry after 8m0s
13m0s: allowed: false , you can retry after 7m0s
14m0s: allowed: false , you can retry after 6m0s
15m0s: allowed: false , you can retry after 5m0s
16m0s: allowed: false , you can retry after 4m0s
17m0s: allowed: false , you can retry after 3m0s
18m0s: allowed: false , you can retry after 2m0s
19m0s: allowed: false , you can retry after 1m0s
20m0s: allowed: true
21m0s: allowed: false , you can retry after 9m0s
22m0s: allowed: false , you can retry after 8m0s
23m0s: allowed: false , you can retry after 7m0s
24m0s: allowed: false , you can retry after 6m0s
--- Sleep 20 minutes ---
45m0s: allowed: true
46m0s: allowed: true
47m0s: allowed: false , you can retry after 3m0s
48m0s: allowed: false , you can retry after 2m0s
49m0s: allowed: false , you can retry after 1m0s
50m0s: allowed: true
51m0s: allowed: false , you can retry after 9m0s
52m0s: allowed: false , you can retry after 8m0s
53m0s: allowed: false , you can retry after 7m0s
54m0s: allowed: false , you can retry after 6m0s
--- Sleep 100 minutes ---
2h35m0s: allowed: true
2h36m0s: allowed: true
2h37m0s: allowed: true
2h38m0s: allowed: true
2h39m0s: allowed: true
2h40m0s: allowed: false , you can retry after 5m0s
2h41m0s: allowed: false , you can retry after 4m0s
2h42m0s: allowed: false , you can retry after 3m0s
2h43m0s: allowed: false , you can retry after 2m0s
2h44m0s: allowed: false , you can retry after 1m0s

func (*GormDriver) Reserve added in v0.2.0 

func (d *GormDriver) Reserve(ctx context.Context, req *ReserveRequest) (*Reservation, error)

type KV 

type KV struct {
	Key   string `json:"key" gorm:"primaryKey;not null;"`
	Value string `json:"value" gorm:"not null;"`
}

type RateLimiter 

type RateLimiter struct {
	// contains filtered or unexported fields
}

func New 

func New(driver Driver) *RateLimiter

func (*RateLimiter) Allow 

func (lim *RateLimiter) Allow(ctx context.Context, req *AllowRequest) (bool, error)

func (*RateLimiter) Reserve 

func (lim *RateLimiter) Reserve(ctx context.Context, req *ReserveRequest) (*Reservation, error)

type RedisDriver added in v0.1.0 

type RedisDriver struct {
	// contains filtered or unexported fields
}

func InitRedisDriver added in v0.1.0 

func InitRedisDriver(ctx context.Context, client *redis.Client) (*RedisDriver, error)
Example 
d, err := InitRedisDriver(context.Background(), redisCli)
if err != nil {
	panic(err)
}
limiter := New(d)
runExample(limiter, "ExampleInitRedisDriver")

Output:

0s: allowed: true
1m0s: allowed: true
2m0s: allowed: true
3m0s: allowed: true
4m0s: allowed: true
5m0s: allowed: false , you can retry after 5m0s
6m0s: allowed: false , you can retry after 4m0s
7m0s: allowed: false , you can retry after 3m0s
8m0s: allowed: false , you can retry after 2m0s
9m0s: allowed: false , you can retry after 1m0s
10m0s: allowed: true
11m0s: allowed: false , you can retry after 9m0s
12m0s: allowed: false , you can retry after 8m0s
13m0s: allowed: false , you can retry after 7m0s
14m0s: allowed: false , you can retry after 6m0s
15m0s: allowed: false , you can retry after 5m0s
16m0s: allowed: false , you can retry after 4m0s
17m0s: allowed: false , you can retry after 3m0s
18m0s: allowed: false , you can retry after 2m0s
19m0s: allowed: false , you can retry after 1m0s
20m0s: allowed: true
21m0s: allowed: false , you can retry after 9m0s
22m0s: allowed: false , you can retry after 8m0s
23m0s: allowed: false , you can retry after 7m0s
24m0s: allowed: false , you can retry after 6m0s
--- Sleep 20 minutes ---
45m0s: allowed: true
46m0s: allowed: true
47m0s: allowed: false , you can retry after 3m0s
48m0s: allowed: false , you can retry after 2m0s
49m0s: allowed: false , you can retry after 1m0s
50m0s: allowed: true
51m0s: allowed: false , you can retry after 9m0s
52m0s: allowed: false , you can retry after 8m0s
53m0s: allowed: false , you can retry after 7m0s
54m0s: allowed: false , you can retry after 6m0s
--- Sleep 100 minutes ---
2h35m0s: allowed: true
2h36m0s: allowed: true
2h37m0s: allowed: true
2h38m0s: allowed: true
2h39m0s: allowed: true
2h40m0s: allowed: false , you can retry after 5m0s
2h41m0s: allowed: false , you can retry after 4m0s
2h42m0s: allowed: false , you can retry after 3m0s
2h43m0s: allowed: false , you can retry after 2m0s
2h44m0s: allowed: false , you can retry after 1m0s

func (*RedisDriver) Reserve added in v0.1.0 

func (d *RedisDriver) Reserve(ctx context.Context, req *ReserveRequest) (*Reservation, error)

type Reservation 

type Reservation struct {
	*ReserveRequest
	OK        bool
	TimeToAct time.Time
	Now       time.Time
}

func (*Reservation) Delay 

func (r *Reservation) Delay() time.Duration

func (*Reservation) DelayFrom 

func (r *Reservation) DelayFrom(t time.Time) time.Duration

func (*Reservation) RetryAfter 

func (r *Reservation) RetryAfter() time.Duration

func (*Reservation) RetryAfterFrom 

func (r *Reservation) RetryAfterFrom(t time.Time) time.Duration

type ReserveRequest 

type ReserveRequest struct {
	Key              string
	DurationPerToken time.Duration
	Burst            int
	Tokens           int
	MaxFutureReserve time.Duration
}

Source Files

View all Source files

Jump to

Keyboard shortcuts

? : This menu
/ : Search site
f or F : Jump to
y or Y : Canonical URL
go.dev uses cookies from Google to deliver and enhance the quality of its services and to analyze traffic. Learn more.