Initial

theplant · Aug 1, 2024 · 592180c · 592180c
1 parent 10aacb0
commit 592180c
Show file tree

Hide file tree

Showing 7 changed files with 1,006 additions and 0 deletions.
diff --git a/README.md b/README.md
@@ -0,0 +1,125 @@
+# ratelimiter
+
+Currently only supports GORM as the driver, with plans to integrate Redis in the future.
+
+```
+package ratelimiter_test
+
+import (
+	"context"
+	"fmt"
+	"time"
+
+	"github.com/theplant/ratelimiter"
+)
+
+func runExample(limiter *ratelimiter.RateLimiter) {
+	// every 10 min , burst 5
+	durationPerToken := 10 * time.Minute
+	burst := int64(5)
+	now := time.Now()
+	key := "test_allow"
+
+	ctx := context.Background()
+
+	try := func(delta time.Duration) bool {
+		reserveReq := &ratelimiter.ReserveRequest{
+			Key:              key,
+			DurationPerToken: durationPerToken,
+			Burst:            burst,
+			Now:              now.Add(delta),
+			Tokens:           1,
+			MaxFutureReserve: 0,
+		}
+		r, err := limiter.Reserve(ctx, 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(reserveReq.Now))
+		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 ExampleDriverGORM() {
+	resetDB()
+
+	limiter := ratelimiter.New(
+		ratelimiter.DriverGORM(db),
+	)
+	runExample(limiter)
+	// 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
+}
+
+```
diff --git a/driver_gorm.go b/driver_gorm.go
@@ -0,0 +1,88 @@
+package ratelimiter
+
+import (
+	"context"
+	"fmt"
+	"strconv"
+	"time"
+
+	"github.com/pkg/errors"
+	"gorm.io/gorm"
+	"gorm.io/gorm/clause"
+)
+
+type KV struct {
+	Key   string `json:"key" gorm:"primaryKey;not null;"`
+	Value string `json:"value" gorm:"not null;"`
+}
+
+func DriverGORM(db *gorm.DB) Driver {
+	return DriverFunc(func(ctx context.Context, req *ReserveRequest) (*Reservation, error) {
+		if req.Key == "" || req.Now.IsZero() || req.DurationPerToken <= 0 || req.Burst <= 0 || req.Tokens <= 0 || req.Tokens > req.Burst {
+			return nil, errors.Errorf("ratelimiter: invalid parameters: %v", req)
+		}
+
+		select {
+		case <-ctx.Done():
+			return nil, errors.Wrap(ctx.Err(), "ratelimiter: context done")
+		default:
+		}
+
+		resetValue := req.Now.Add(-time.Duration(req.Burst) * req.DurationPerToken)
+
+		var baseTime time.Time
+		var timeToAct time.Time
+		var ok bool
+
+		db := db.WithContext(ctx)
+		err := db.Transaction(func(tx *gorm.DB) error {
+			var kv KV
+
+			if err := tx.Clauses(clause.Locking{Strength: "UPDATE"}).First(&kv, "key = ?", req.Key).Error; err != nil {
+				if !errors.Is(err, gorm.ErrRecordNotFound) {
+					return err
+				}
+
+				baseTime = resetValue
+				kv = KV{Key: req.Key, Value: fmt.Sprintf("%d", baseTime.UnixNano())}
+				if err := tx.Create(&kv).Error; err != nil {
+					return err
+				}
+			} else {
+				baseTimeUnix, err := strconv.ParseInt(kv.Value, 10, 64)
+				if err != nil {
+					return errors.Wrap(err, "ratelimiter: failed to parse base time")
+				}
+				baseTime = time.Unix(0, baseTimeUnix)
+
+				if baseTime.Before(resetValue) {
+					baseTime = resetValue
+				}
+			}
+
+			tokensDuration := req.DurationPerToken * time.Duration(req.Tokens)
+			timeToAct = baseTime.Add(tokensDuration)
+
+			if timeToAct.After(req.Now.Add(req.MaxFutureReserve)) {
+				ok = false
+				return nil
+			}
+
+			kv.Value = fmt.Sprintf("%d", timeToAct.UnixNano())
+			if err := tx.Save(&kv).Error; err != nil {
+				return errors.Wrap(err, "ratelimiter: failed to save time to act")
+			}
+			ok = true
+			return nil
+		})
+		if err != nil {
+			return nil, err
+		}
+
+		return &Reservation{
+			ReserveRequest: req,
+			OK:             ok,
+			TimeToAct:      timeToAct,
+		}, nil
+	})
+}
diff --git a/example_test.go b/example_test.go
@@ -0,0 +1,118 @@
+package ratelimiter_test
+
+import (
+	"context"
+	"fmt"
+	"time"
+
+	"github.com/theplant/ratelimiter"
+)
+
+func runExample(limiter *ratelimiter.RateLimiter) {
+	// every 10 min , burst 5
+	durationPerToken := 10 * time.Minute
+	burst := int64(5)
+	now := time.Now()
+	key := "test_allow"
+
+	ctx := context.Background()
+
+	try := func(delta time.Duration) bool {
+		reserveReq := &ratelimiter.ReserveRequest{
+			Key:              key,
+			DurationPerToken: durationPerToken,
+			Burst:            burst,
+			Now:              now.Add(delta),
+			Tokens:           1,
+			MaxFutureReserve: 0,
+		}
+		r, err := limiter.Reserve(ctx, 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(reserveReq.Now))
+		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 ExampleDriverGORM() {
+	resetDB()
+
+	limiter := ratelimiter.New(
+		ratelimiter.DriverGORM(db),
+	)
+	runExample(limiter)
+	// 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
+}