From 6e9759cccaa6a216a186de6ec1c20243dd2df9f0 Mon Sep 17 00:00:00 2001 From: Vlad Lavrishko Date: Sun, 21 Sep 2025 01:06:54 +0300 Subject: [PATCH 1/7] feat(yaratelimit): base logic --- yaratelimit/yaratelimit.go | 165 +++++++++++++++++++++++++++++++++++++ 1 file changed, 165 insertions(+) create mode 100644 yaratelimit/yaratelimit.go diff --git a/yaratelimit/yaratelimit.go b/yaratelimit/yaratelimit.go new file mode 100644 index 0000000..cdc66f0 --- /dev/null +++ b/yaratelimit/yaratelimit.go @@ -0,0 +1,165 @@ +package yaratelimit + +import ( + "context" + "fmt" + "net/http" + "strconv" + "strings" + "time" + + "github.com/YaCodeDev/GoYaCodeDevUtils/yacache" + "github.com/YaCodeDev/GoYaCodeDevUtils/yaerrors" +) + +type IRateLimit interface { + Check( + ctx context.Context, + id uint64, + group string, + ) (bool, yaerrors.Error) + Refresh( + ctx context.Context, + id uint64, + group string, + ) yaerrors.Error + Increment( + ctx context.Context, + id uint64, + group string, + ) yaerrors.Error + Get( + ctx context.Context, + id uint64, + group string, + ) (*Storage, yaerrors.Error) +} + +type Storage struct { + Limit uint8 + FirstRequest int64 +} + +type RateLimit[Cache yacache.Container] struct { + Cache yacache.Cache[Cache] + Limit uint8 + Rate time.Duration +} + +func NewRateLimit[Cache yacache.Container]( + cache yacache.Cache[Cache], + limit uint8, + rate time.Duration, +) *RateLimit[Cache] { + return &RateLimit[Cache]{ + Limit: limit, + Rate: rate, + Cache: cache, + } +} + +func (r *RateLimit[Cache]) Check( + ctx context.Context, + id uint64, + group string, +) (bool, yaerrors.Error) { + storage, err := r.Get(ctx, id, group) + if err != nil { + err.Wrap("failed to check storage") + } + + if storage.Limit+1 >= r.Limit { + return true, nil + } + + return false, nil +} + +func (r *RateLimit[Cache]) Increment( + ctx context.Context, + id uint64, + group string, +) (bool, yaerrors.Error) { + storage, err := r.Get(ctx, id, group) + if err != nil { + if err := r.Refresh(ctx, id, group); err != nil { + return false, err.Wrap("failed to refresh") + } + } + + if time.Now().Add(-r.Rate).Before(time.Unix(storage.FirstRequest, 0)) { + if err := r.Cache.Set(ctx, formatKey(id, group), fmt.Sprintf("%d,%d", storage.Limit+1, storage.FirstRequest), 0); err != nil { + return false, err.Wrap("failed to increment storage") + } + } else { + if err := r.Refresh(ctx, id, group); err != nil { + return false, err.Wrap("failed to refresh") + } + + return false, nil + } + + if storage.Limit+1 >= r.Limit { + return true, nil + } + + return false, nil +} + +func (r *RateLimit[Cache]) Refresh( + ctx context.Context, + id uint64, + group string, +) yaerrors.Error { + if err := r.Cache.Set(ctx, formatKey(id, group), fmt.Sprintf("%d,%d", 1, time.Now().Unix()), 0); err != nil { + return err.Wrap("failed to set refreshed storage") + } + + return nil +} + +func (r *RateLimit[Cache]) Get( + ctx context.Context, + id uint64, + group string, +) (*Storage, yaerrors.Error) { + value, yaerr := r.Cache.Get(ctx, formatKey(id, group)) + if yaerr != nil { + return nil, yaerr.Wrap("failed to get storage") + } + + values := strings.Split(value, ",") + if len(values) != 2 { + return nil, yaerrors.FromString( + http.StatusInternalServerError, + "not compare storage", + ) + } + + limit, err := strconv.ParseUint(values[0], 10, 8) + if err != nil { + return nil, yaerrors.FromError( + http.StatusInternalServerError, + err, + "couldn't validate limit", + ) + } + + firstRequest, err := strconv.ParseInt(values[1], 10, 64) + if err != nil { + return nil, yaerrors.FromError( + http.StatusInternalServerError, + err, + "couldn't validate unix time", + ) + } + + return &Storage{ + Limit: uint8(limit), + FirstRequest: firstRequest, + }, nil +} + +func formatKey(id uint64, group string) string { + return fmt.Sprintf("rate-limit-%d-%s", id, group) +} From ba2b06ebd6210218a7fba6d439be0183875a0c8d Mon Sep 17 00:00:00 2001 From: Vlad Lavrishko Date: Sun, 21 Sep 2025 01:10:52 +0300 Subject: [PATCH 2/7] chore(yaratelimit): fix lint --- yaratelimit/yaratelimit.go | 13 ++++++++++--- 1 file changed, 10 insertions(+), 3 deletions(-) diff --git a/yaratelimit/yaratelimit.go b/yaratelimit/yaratelimit.go index cdc66f0..fdefccf 100644 --- a/yaratelimit/yaratelimit.go +++ b/yaratelimit/yaratelimit.go @@ -65,7 +65,7 @@ func (r *RateLimit[Cache]) Check( ) (bool, yaerrors.Error) { storage, err := r.Get(ctx, id, group) if err != nil { - err.Wrap("failed to check storage") + return false, err.Wrap("failed to check storage") } if storage.Limit+1 >= r.Limit { @@ -88,7 +88,12 @@ func (r *RateLimit[Cache]) Increment( } if time.Now().Add(-r.Rate).Before(time.Unix(storage.FirstRequest, 0)) { - if err := r.Cache.Set(ctx, formatKey(id, group), fmt.Sprintf("%d,%d", storage.Limit+1, storage.FirstRequest), 0); err != nil { + if err := r.Cache.Set( + ctx, + formatKey(id, group), + fmt.Sprintf("%d,%d", storage.Limit+1, storage.FirstRequest), + 0, + ); err != nil { return false, err.Wrap("failed to increment storage") } } else { @@ -128,8 +133,10 @@ func (r *RateLimit[Cache]) Get( return nil, yaerr.Wrap("failed to get storage") } + const separate = 2 + values := strings.Split(value, ",") - if len(values) != 2 { + if len(values) != separate { return nil, yaerrors.FromString( http.StatusInternalServerError, "not compare storage", From 251335b66d9873fec1108adce6bdfcbab394991f Mon Sep 17 00:00:00 2001 From: Vlad Lavrishko Date: Sun, 21 Sep 2025 12:03:13 +0300 Subject: [PATCH 3/7] feat(yaratelimit): unit tests --- yaratelimit/yaratelimit.go | 22 +++-- yaratelimit/yaratelimit_test.go | 157 ++++++++++++++++++++++++++++++++ 2 files changed, 173 insertions(+), 6 deletions(-) create mode 100644 yaratelimit/yaratelimit_test.go diff --git a/yaratelimit/yaratelimit.go b/yaratelimit/yaratelimit.go index fdefccf..a67dfd5 100644 --- a/yaratelimit/yaratelimit.go +++ b/yaratelimit/yaratelimit.go @@ -58,7 +58,7 @@ func NewRateLimit[Cache yacache.Container]( } } -func (r *RateLimit[Cache]) Check( +func (r *RateLimit[Cache]) CheckBanned( ctx context.Context, id uint64, group string, @@ -85,13 +85,19 @@ func (r *RateLimit[Cache]) Increment( if err := r.Refresh(ctx, id, group); err != nil { return false, err.Wrap("failed to refresh") } + + return false, nil + } + + if storage.Limit >= r.Limit { + return true, nil } if time.Now().Add(-r.Rate).Before(time.Unix(storage.FirstRequest, 0)) { if err := r.Cache.Set( ctx, - formatKey(id, group), - fmt.Sprintf("%d,%d", storage.Limit+1, storage.FirstRequest), + FormatKey(id, group), + FormatValue(storage.Limit+1, storage.FirstRequest), 0, ); err != nil { return false, err.Wrap("failed to increment storage") @@ -116,7 +122,7 @@ func (r *RateLimit[Cache]) Refresh( id uint64, group string, ) yaerrors.Error { - if err := r.Cache.Set(ctx, formatKey(id, group), fmt.Sprintf("%d,%d", 1, time.Now().Unix()), 0); err != nil { + if err := r.Cache.Set(ctx, FormatKey(id, group), fmt.Sprintf("%d,%d", 1, time.Now().Unix()), 0); err != nil { return err.Wrap("failed to set refreshed storage") } @@ -128,7 +134,7 @@ func (r *RateLimit[Cache]) Get( id uint64, group string, ) (*Storage, yaerrors.Error) { - value, yaerr := r.Cache.Get(ctx, formatKey(id, group)) + value, yaerr := r.Cache.Get(ctx, FormatKey(id, group)) if yaerr != nil { return nil, yaerr.Wrap("failed to get storage") } @@ -167,6 +173,10 @@ func (r *RateLimit[Cache]) Get( }, nil } -func formatKey(id uint64, group string) string { +func FormatKey(id uint64, group string) string { return fmt.Sprintf("rate-limit-%d-%s", id, group) } + +func FormatValue(limit uint8, firstRequest int64) string { + return fmt.Sprintf("%d,%d", limit, firstRequest) +} diff --git a/yaratelimit/yaratelimit_test.go b/yaratelimit/yaratelimit_test.go new file mode 100644 index 0000000..4792a39 --- /dev/null +++ b/yaratelimit/yaratelimit_test.go @@ -0,0 +1,157 @@ +package yaratelimit_test + +import ( + "context" + "testing" + "time" + + "github.com/YaCodeDev/GoYaCodeDevUtils/yacache" + "github.com/YaCodeDev/GoYaCodeDevUtils/yaratelimit" + "github.com/stretchr/testify/assert" +) + +func TestIncrementWorkFlow_Works(t *testing.T) { + ctx := context.Background() + + cache := yacache.NewCache(yacache.NewMemoryContainer()) + + t.Run("Increment works", func(t *testing.T) { + rate := yaratelimit.NewRateLimit(cache, 5, time.Second*400) + + userID, group := uint64(100), "party" + + _, _ = rate.Increment(ctx, userID, group) + + expected := yaratelimit.FormatValue(1, time.Now().Unix()) + + result, _ := cache.Get(ctx, yaratelimit.FormatKey(userID, group)) + + assert.Equal(t, expected, result) + }) + + t.Run("Refresh works", func(t *testing.T) { + rate := yaratelimit.NewRateLimit(cache, 3, time.Millisecond*5) + + userID, group := uint64(100), "party" + + rate.Increment(ctx, userID, group) + rate.Increment(ctx, userID, group) + rate.Increment(ctx, userID, group) + + time.Sleep(time.Millisecond * 5) + + expected := yaratelimit.FormatValue(1, time.Now().Unix()) + + _, _ = rate.Increment(ctx, userID, group) + + result, _ := cache.Get(ctx, yaratelimit.FormatKey(userID, group)) + + assert.Equal(t, expected, result) + }) + + t.Run("Overflow works", func(t *testing.T) { + rate := yaratelimit.NewRateLimit(cache, 3, time.Second) + + userID, group := uint64(100), "party" + + _, _ = rate.Increment(ctx, userID, group) + _, _ = rate.Increment(ctx, userID, group) + _, _ = rate.Increment(ctx, userID, group) + + expected := yaratelimit.FormatValue(3, time.Now().Unix()) + + _, _ = rate.Increment(ctx, userID, group) + _, _ = rate.Increment(ctx, userID, group) + _, _ = rate.Increment(ctx, userID, group) + + result, _ := cache.Get(ctx, yaratelimit.FormatKey(userID, group)) + + assert.Equal(t, expected, result) + }) + + t.Run("Ban works", func(t *testing.T) { + rate := yaratelimit.NewRateLimit(cache, 3, time.Second) + + userID, group := uint64(100), "party" + + _, _ = rate.Increment(ctx, userID, group) + _, _ = rate.Increment(ctx, userID, group) + _, _ = rate.Increment(ctx, userID, group) + + _, _ = rate.Increment(ctx, userID, group) + _, _ = rate.Increment(ctx, userID, group) + + result, _ := rate.Increment(ctx, userID, group) + + expected := true + + assert.Equal(t, expected, result) + }) +} + +func TestGet_Works(t *testing.T) { + ctx := context.Background() + + cache := yacache.NewCache(yacache.NewMemoryContainer()) + + rate := yaratelimit.NewRateLimit(cache, 5, time.Second*400) + + userID, group := uint64(100), "party" + + expected := &yaratelimit.Storage{ + Limit: 1, + FirstRequest: time.Now().Unix(), + } + + _, _ = rate.Increment(ctx, userID, group) + + result, _ := rate.Get(ctx, userID, group) + + assert.Equal(t, expected, result) +} + +func TestCheckBanned_Works(t *testing.T) { + ctx := context.Background() + + cache := yacache.NewCache(yacache.NewMemoryContainer()) + + rate := yaratelimit.NewRateLimit(cache, 5, time.Second*400) + + userID, group := uint64(100), "party" + + expected := true + + _, _ = rate.Increment(ctx, userID, group) + _, _ = rate.Increment(ctx, userID, group) + _, _ = rate.Increment(ctx, userID, group) + _, _ = rate.Increment(ctx, userID, group) + _, _ = rate.Increment(ctx, userID, group) + + result, _ := rate.CheckBanned(ctx, userID, group) + + assert.Equal(t, expected, result) +} + +func TestRefresh_Works(t *testing.T) { + ctx := context.Background() + + cache := yacache.NewCache(yacache.NewMemoryContainer()) + + rate := yaratelimit.NewRateLimit(cache, 5, time.Second*400) + + userID, group := uint64(100), "party" + + expected := false + + _, _ = rate.Increment(ctx, userID, group) + _, _ = rate.Increment(ctx, userID, group) + _, _ = rate.Increment(ctx, userID, group) + _, _ = rate.Increment(ctx, userID, group) + _, _ = rate.Increment(ctx, userID, group) + + _ = rate.Refresh(ctx, userID, group) + + result, _ := rate.CheckBanned(ctx, userID, group) + + assert.Equal(t, expected, result) +} From 6fb742e5a3ff7a3d484f26463c656abf186dbdee Mon Sep 17 00:00:00 2001 From: Vlad Lavrishko Date: Sun, 21 Sep 2025 12:05:46 +0300 Subject: [PATCH 4/7] chore(yaratelimit): fix lint consts --- yaratelimit/yaratelimit_test.go | 89 +++++++++++++++------------------ 1 file changed, 40 insertions(+), 49 deletions(-) diff --git a/yaratelimit/yaratelimit_test.go b/yaratelimit/yaratelimit_test.go index 4792a39..ca95fac 100644 --- a/yaratelimit/yaratelimit_test.go +++ b/yaratelimit/yaratelimit_test.go @@ -10,6 +10,11 @@ import ( "github.com/stretchr/testify/assert" ) +const ( + TestUserID = 100 + TestGroup = "sigma-life" +) + func TestIncrementWorkFlow_Works(t *testing.T) { ctx := context.Background() @@ -18,13 +23,11 @@ func TestIncrementWorkFlow_Works(t *testing.T) { t.Run("Increment works", func(t *testing.T) { rate := yaratelimit.NewRateLimit(cache, 5, time.Second*400) - userID, group := uint64(100), "party" - - _, _ = rate.Increment(ctx, userID, group) + _, _ = rate.Increment(ctx, TestUserID, TestGroup) expected := yaratelimit.FormatValue(1, time.Now().Unix()) - result, _ := cache.Get(ctx, yaratelimit.FormatKey(userID, group)) + result, _ := cache.Get(ctx, yaratelimit.FormatKey(TestUserID, TestGroup)) assert.Equal(t, expected, result) }) @@ -32,19 +35,17 @@ func TestIncrementWorkFlow_Works(t *testing.T) { t.Run("Refresh works", func(t *testing.T) { rate := yaratelimit.NewRateLimit(cache, 3, time.Millisecond*5) - userID, group := uint64(100), "party" - - rate.Increment(ctx, userID, group) - rate.Increment(ctx, userID, group) - rate.Increment(ctx, userID, group) + _, _ = rate.Increment(ctx, TestUserID, TestGroup) + _, _ = rate.Increment(ctx, TestUserID, TestGroup) + _, _ = rate.Increment(ctx, TestUserID, TestGroup) time.Sleep(time.Millisecond * 5) expected := yaratelimit.FormatValue(1, time.Now().Unix()) - _, _ = rate.Increment(ctx, userID, group) + _, _ = rate.Increment(ctx, TestUserID, TestGroup) - result, _ := cache.Get(ctx, yaratelimit.FormatKey(userID, group)) + result, _ := cache.Get(ctx, yaratelimit.FormatKey(TestUserID, TestGroup)) assert.Equal(t, expected, result) }) @@ -52,19 +53,17 @@ func TestIncrementWorkFlow_Works(t *testing.T) { t.Run("Overflow works", func(t *testing.T) { rate := yaratelimit.NewRateLimit(cache, 3, time.Second) - userID, group := uint64(100), "party" - - _, _ = rate.Increment(ctx, userID, group) - _, _ = rate.Increment(ctx, userID, group) - _, _ = rate.Increment(ctx, userID, group) + _, _ = rate.Increment(ctx, TestUserID, TestGroup) + _, _ = rate.Increment(ctx, TestUserID, TestGroup) + _, _ = rate.Increment(ctx, TestUserID, TestGroup) expected := yaratelimit.FormatValue(3, time.Now().Unix()) - _, _ = rate.Increment(ctx, userID, group) - _, _ = rate.Increment(ctx, userID, group) - _, _ = rate.Increment(ctx, userID, group) + _, _ = rate.Increment(ctx, TestUserID, TestGroup) + _, _ = rate.Increment(ctx, TestUserID, TestGroup) + _, _ = rate.Increment(ctx, TestUserID, TestGroup) - result, _ := cache.Get(ctx, yaratelimit.FormatKey(userID, group)) + result, _ := cache.Get(ctx, yaratelimit.FormatKey(TestUserID, TestGroup)) assert.Equal(t, expected, result) }) @@ -72,16 +71,14 @@ func TestIncrementWorkFlow_Works(t *testing.T) { t.Run("Ban works", func(t *testing.T) { rate := yaratelimit.NewRateLimit(cache, 3, time.Second) - userID, group := uint64(100), "party" + _, _ = rate.Increment(ctx, TestUserID, TestGroup) + _, _ = rate.Increment(ctx, TestUserID, TestGroup) + _, _ = rate.Increment(ctx, TestUserID, TestGroup) - _, _ = rate.Increment(ctx, userID, group) - _, _ = rate.Increment(ctx, userID, group) - _, _ = rate.Increment(ctx, userID, group) + _, _ = rate.Increment(ctx, TestUserID, TestGroup) + _, _ = rate.Increment(ctx, TestUserID, TestGroup) - _, _ = rate.Increment(ctx, userID, group) - _, _ = rate.Increment(ctx, userID, group) - - result, _ := rate.Increment(ctx, userID, group) + result, _ := rate.Increment(ctx, TestUserID, TestGroup) expected := true @@ -96,16 +93,14 @@ func TestGet_Works(t *testing.T) { rate := yaratelimit.NewRateLimit(cache, 5, time.Second*400) - userID, group := uint64(100), "party" - expected := &yaratelimit.Storage{ Limit: 1, FirstRequest: time.Now().Unix(), } - _, _ = rate.Increment(ctx, userID, group) + _, _ = rate.Increment(ctx, TestUserID, TestGroup) - result, _ := rate.Get(ctx, userID, group) + result, _ := rate.Get(ctx, TestUserID, TestGroup) assert.Equal(t, expected, result) } @@ -117,17 +112,15 @@ func TestCheckBanned_Works(t *testing.T) { rate := yaratelimit.NewRateLimit(cache, 5, time.Second*400) - userID, group := uint64(100), "party" - expected := true - _, _ = rate.Increment(ctx, userID, group) - _, _ = rate.Increment(ctx, userID, group) - _, _ = rate.Increment(ctx, userID, group) - _, _ = rate.Increment(ctx, userID, group) - _, _ = rate.Increment(ctx, userID, group) + _, _ = rate.Increment(ctx, TestUserID, TestGroup) + _, _ = rate.Increment(ctx, TestUserID, TestGroup) + _, _ = rate.Increment(ctx, TestUserID, TestGroup) + _, _ = rate.Increment(ctx, TestUserID, TestGroup) + _, _ = rate.Increment(ctx, TestUserID, TestGroup) - result, _ := rate.CheckBanned(ctx, userID, group) + result, _ := rate.CheckBanned(ctx, TestUserID, TestGroup) assert.Equal(t, expected, result) } @@ -139,19 +132,17 @@ func TestRefresh_Works(t *testing.T) { rate := yaratelimit.NewRateLimit(cache, 5, time.Second*400) - userID, group := uint64(100), "party" - expected := false - _, _ = rate.Increment(ctx, userID, group) - _, _ = rate.Increment(ctx, userID, group) - _, _ = rate.Increment(ctx, userID, group) - _, _ = rate.Increment(ctx, userID, group) - _, _ = rate.Increment(ctx, userID, group) + _, _ = rate.Increment(ctx, TestUserID, TestGroup) + _, _ = rate.Increment(ctx, TestUserID, TestGroup) + _, _ = rate.Increment(ctx, TestUserID, TestGroup) + _, _ = rate.Increment(ctx, TestUserID, TestGroup) + _, _ = rate.Increment(ctx, TestUserID, TestGroup) - _ = rate.Refresh(ctx, userID, group) + _ = rate.Refresh(ctx, TestUserID, TestGroup) - result, _ := rate.CheckBanned(ctx, userID, group) + result, _ := rate.CheckBanned(ctx, TestUserID, TestGroup) assert.Equal(t, expected, result) } From 19885d30d6b6a64d42a88db3466f6697325a7d73 Mon Sep 17 00:00:00 2001 From: Vlad Lavrishko Date: Sun, 21 Sep 2025 17:07:49 +0300 Subject: [PATCH 5/7] feat(yaratelimit): docs --- yaratelimit/yaratelimit.go | 126 ++++++++++++++++++++++++++++++++++++- 1 file changed, 124 insertions(+), 2 deletions(-) diff --git a/yaratelimit/yaratelimit.go b/yaratelimit/yaratelimit.go index a67dfd5..65bb3b3 100644 --- a/yaratelimit/yaratelimit.go +++ b/yaratelimit/yaratelimit.go @@ -1,3 +1,51 @@ +// Package yaratelimit implements a simple fixed-window rate limiter backed by +// a yacache.Cache. It stores a per-(id, group) counter alongside the unix +// timestamp of the first request in the current window. +// +// # Storage layout +// +// Each subject is addressed by a string key: +// +// rate-limit:- +// +// The cache value is a compact CSV tuple: +// +// "," +// +// For example: "3,1726860000" means 3 requests since unix time 1726860000. +// +// # Model +// +// The limiter uses a fixed window of size Rate (time.Duration). The value +// tracks the number of hits and the unix timestamp of the first hit within +// the active window. On each Increment: +// +// - If no record exists: Refresh() creates one with count=1, first=now. +// - If now - first < Rate: count++ (up to Limit). +// - If now - first >= Rate: Refresh() starts a new window with count=1. +// +// # Semantics +// +// - Increment(ctx, id, group) -> (banned bool, err) +// +// Increments the counter if inside the current window or refreshes the +// window if it has expired. Returns true if the subject should be treated +// as banned/over limit *after* this call (i.e., when the count reaches or +// exceeds Limit). +// +// - CheckBanned(ctx, id, group) -> (banned bool, err) +// +// Reads the current value and returns whether the next hit would be over +// the limit. (Useful to check prior to serving a request.) +// +// - Refresh(ctx, id, group) +// +// Resets the window (count=1, first=now). +// +// - Get(ctx, id, group) +// +// Returns the parsed Storage (count/limit used in this window and first +// request unix timestamp). package yaratelimit import ( @@ -12,22 +60,39 @@ import ( "github.com/YaCodeDev/GoYaCodeDevUtils/yaerrors" ) +// IRateLimit exposes the behaviour of a fixed-window rate limiter backed by a cache. +// +// Example: +// +// cache := yacache.NewCache(yacache.NewMemoryContainer()) +// rl := yaratelimit.NewRateLimit(cache, 5, time.Minute) +// banned, err := rl.Increment(ctx, 42, "signin") +// _ = banned; _ = err type IRateLimit interface { + // CheckBanned inspects the current window for (id, group) and returns true + // if the *next* request should be treated as banned (i.e., would reach/exceed Limit). Check( ctx context.Context, id uint64, group string, ) (bool, yaerrors.Error) + + // Refresh resets the window for (id, group) to count=1 at current timestamp. Refresh( ctx context.Context, id uint64, group string, ) yaerrors.Error + + // Increment applies a hit inside the current window or refreshes if expired. + // It returns true if the subject should be treated as banned after this hit. Increment( ctx context.Context, id uint64, group string, ) yaerrors.Error + + // Get returns the current storage tuple for (id, group). Get( ctx context.Context, id uint64, @@ -35,17 +100,33 @@ type IRateLimit interface { ) (*Storage, yaerrors.Error) } +// Storage is the parsed representation of the CSV value in the cache. type Storage struct { - Limit uint8 + // Limit is the count within the current window (despite the name, it stores the current usage). + Limit uint8 + // FirstRequest is the unix timestamp of the first request within the current window. FirstRequest int64 } +// RateLimit is a fixed-window limiter backed by a yacache.Cache. +// The zero value is not valid; use NewRateLimit. type RateLimit[Cache yacache.Container] struct { Cache yacache.Cache[Cache] + // Limit is the max allowed hits per window. Limit uint8 - Rate time.Duration + // Rate is the window size (duration). + Rate time.Duration } +// NewRateLimit wires dependencies and returns a ready-to-use limiter. +// +// - cache: any yacache implementation (memory, redis, etc.) +// - limit: maximum hits per window +// - rate : window duration +// +// Example: +// +// rl := yaratelimit.NewRateLimit(cache, 5, time.Minute) func NewRateLimit[Cache yacache.Container]( cache yacache.Cache[Cache], limit uint8, @@ -58,6 +139,15 @@ func NewRateLimit[Cache yacache.Container]( } } +// CheckBanned inspects the current window and returns true if the next call to +// Increment should be considered banned. Returns false if the subject has not +// reached the threshold yet or if no storage exists. +// +// Example: +// +// banned, err := rl.CheckBanned(ctx, userID, "signup") +// if err != nil { /* handle */ } +// if banned { /* throttle */ } func (r *RateLimit[Cache]) CheckBanned( ctx context.Context, id uint64, @@ -68,6 +158,7 @@ func (r *RateLimit[Cache]) CheckBanned( return false, err.Wrap("failed to check storage") } + // If the next increment would cross the limit, treat as banned. if storage.Limit+1 >= r.Limit { return true, nil } @@ -75,6 +166,15 @@ func (r *RateLimit[Cache]) CheckBanned( return false, nil } +// Increment records a hit for (id, group). +// If the window is still active, it increments the counter. +// If the window expired, it Refreshes the window (count=1). +// Returns true if the subject is now banned (count >= Limit). +// +// Example: +// +// banned, err := rl.Increment(ctx, userID, "api:v1") +// if banned { /* reject */ } func (r *RateLimit[Cache]) Increment( ctx context.Context, id uint64, @@ -117,6 +217,11 @@ func (r *RateLimit[Cache]) Increment( return false, nil } +// Refresh resets the window for (id, group) to count=1 at the current timestamp. +// +// Example: +// +// _ = rl.Refresh(ctx, 42, "password_reset") func (r *RateLimit[Cache]) Refresh( ctx context.Context, id uint64, @@ -129,6 +234,13 @@ func (r *RateLimit[Cache]) Refresh( return nil } +// Get fetches and parses the cache record for (id, group). +// +// Example: +// +// st, err := rl.Get(ctx, 42, "sms") +// if err != nil { /* handle */ } +// fmt.Println(st.Limit, st.FirstRequest) func (r *RateLimit[Cache]) Get( ctx context.Context, id uint64, @@ -173,10 +285,20 @@ func (r *RateLimit[Cache]) Get( }, nil } +// FormatKey constructs the cache key for (id, group). +// +// Example: +// +// k := yaratelimit.FormatKey(100, "signup") // "rate-limit-100-signup" func FormatKey(id uint64, group string) string { return fmt.Sprintf("rate-limit-%d-%s", id, group) } +// FormatValue serializes a (count, first_unix) tuple to cache string. +// +// Example: +// +// v := yaratelimit.FormatValue(2, 1726860000) // "2,1726860000" func FormatValue(limit uint8, firstRequest int64) string { return fmt.Sprintf("%d,%d", limit, firstRequest) } From 364f8a9f2b2794434212b44adf33140490b60286 Mon Sep 17 00:00:00 2001 From: Vlad Lavrishko Date: Sun, 21 Sep 2025 17:14:02 +0300 Subject: [PATCH 6/7] fix(yaratelimit): make correct API, logs --- yaratelimit/yaratelimit.go | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/yaratelimit/yaratelimit.go b/yaratelimit/yaratelimit.go index 65bb3b3..fad3904 100644 --- a/yaratelimit/yaratelimit.go +++ b/yaratelimit/yaratelimit.go @@ -71,7 +71,7 @@ import ( type IRateLimit interface { // CheckBanned inspects the current window for (id, group) and returns true // if the *next* request should be treated as banned (i.e., would reach/exceed Limit). - Check( + CheckBanned( ctx context.Context, id uint64, group string, @@ -90,7 +90,7 @@ type IRateLimit interface { ctx context.Context, id uint64, group string, - ) yaerrors.Error + ) (bool, yaerrors.Error) // Get returns the current storage tuple for (id, group). Get( @@ -158,7 +158,6 @@ func (r *RateLimit[Cache]) CheckBanned( return false, err.Wrap("failed to check storage") } - // If the next increment would cross the limit, treat as banned. if storage.Limit+1 >= r.Limit { return true, nil } @@ -257,7 +256,7 @@ func (r *RateLimit[Cache]) Get( if len(values) != separate { return nil, yaerrors.FromString( http.StatusInternalServerError, - "not compare storage", + "invalid storage format", ) } From c7737226eb63e7ddd7b28e10fc8ed939bc044624 Mon Sep 17 00:00:00 2001 From: Vlad Lavrishko Date: Mon, 29 Sep 2025 01:17:06 +0300 Subject: [PATCH 7/7] chore(yaratelimit): idiomatic return, increment --- yaratelimit/yaratelimit.go | 12 ++---------- yaratelimit/yaratelimit_test.go | 3 --- 2 files changed, 2 insertions(+), 13 deletions(-) diff --git a/yaratelimit/yaratelimit.go b/yaratelimit/yaratelimit.go index fad3904..537f0c2 100644 --- a/yaratelimit/yaratelimit.go +++ b/yaratelimit/yaratelimit.go @@ -158,11 +158,7 @@ func (r *RateLimit[Cache]) CheckBanned( return false, err.Wrap("failed to check storage") } - if storage.Limit+1 >= r.Limit { - return true, nil - } - - return false, nil + return storage.Limit+1 >= r.Limit, nil } // Increment records a hit for (id, group). @@ -209,11 +205,7 @@ func (r *RateLimit[Cache]) Increment( return false, nil } - if storage.Limit+1 >= r.Limit { - return true, nil - } - - return false, nil + return storage.Limit+1 >= r.Limit, nil } // Refresh resets the window for (id, group) to count=1 at the current timestamp. diff --git a/yaratelimit/yaratelimit_test.go b/yaratelimit/yaratelimit_test.go index ca95fac..a1547d4 100644 --- a/yaratelimit/yaratelimit_test.go +++ b/yaratelimit/yaratelimit_test.go @@ -75,9 +75,6 @@ func TestIncrementWorkFlow_Works(t *testing.T) { _, _ = rate.Increment(ctx, TestUserID, TestGroup) _, _ = rate.Increment(ctx, TestUserID, TestGroup) - _, _ = rate.Increment(ctx, TestUserID, TestGroup) - _, _ = rate.Increment(ctx, TestUserID, TestGroup) - result, _ := rate.Increment(ctx, TestUserID, TestGroup) expected := true