diff --git a/Makefile b/Makefile index 4e16ca44..a2aac632 100644 --- a/Makefile +++ b/Makefile @@ -5,7 +5,7 @@ LDFLAGS?=-s -w TESTFILE=_testok # go tools versions -GOLANGCI=github.com/golangci/golangci-lint/cmd/golangci-lint@v1.61.0 +GOLANGCI=github.com/golangci/golangci-lint/cmd/golangci-lint@v1.63.4 gofumpt=mvdan.cc/gofumpt@latest govulncheck=golang.org/x/vuln/cmd/govulncheck@latest goimports=golang.org/x/tools/cmd/goimports@latest diff --git a/go.mod b/go.mod index ce456857..dc5830dd 100644 --- a/go.mod +++ b/go.mod @@ -1,6 +1,6 @@ module github.com/rudderlabs/rudder-go-kit -go 1.23.1 +go 1.23.5 replace github.com/gocql/gocql => github.com/scylladb/gocql v1.14.2 diff --git a/throttling/options.go b/throttling/options.go index 46b78a40..29541531 100644 --- a/throttling/options.go +++ b/throttling/options.go @@ -36,9 +36,12 @@ func WithRedisSortedSet(rc *redis.Client) Option { } } -// WithStatsCollector allows to setup a stats collector for the limiter +// WithStatsCollector allows to set up a stats collector for the limiter func WithStatsCollector(sc statsCollector) Option { - return func(l *Limiter) { - l.statsCollector = sc - } + return func(l *Limiter) { l.statsCollector = sc } +} + +// WithStatsTagger allows to set up a stats tagger for the limiter +func WithStatsTagger(st statsTagger) Option { + return func(l *Limiter) { l.statsTagger = st } } diff --git a/throttling/throttling.go b/throttling/throttling.go index 5ef57895..d8aa95d7 100644 --- a/throttling/throttling.go +++ b/throttling/throttling.go @@ -42,6 +42,8 @@ type statsCollector interface { NewTaggedStat(name, statType string, tags stats.Tags) stats.Measurement } +type statsTagger func(key, algo string, rate, window int64) stats.Tags + type Limiter struct { // for Redis configurations // a default redisSpeaker should always be provided for Redis configurations @@ -56,10 +58,20 @@ type Limiter struct { // metrics statsCollector statsCollector + statsTagger statsTagger } func New(options ...Option) (*Limiter, error) { - rl := &Limiter{} + rl := &Limiter{ + statsTagger: func(key, algo string, rate, window int64) stats.Tags { + return stats.Tags{ + "key": key, + "algo": algo, + "rate": strconv.FormatInt(rate, 10), + "window": strconv.FormatInt(window, 10), + } + }, + } for i := range options { options[i](rl) } @@ -227,12 +239,7 @@ func (l *Limiter) gcraLimit(ctx context.Context, cost, rate, window int64, key s } func (l *Limiter) getTimer(key, algo string, rate, window int64) func() { - m := l.statsCollector.NewTaggedStat("throttling", stats.TimerType, stats.Tags{ - "key": key, - "algo": algo, - "rate": strconv.FormatInt(rate, 10), - "window": strconv.FormatInt(window, 10), - }) + m := l.statsCollector.NewTaggedStat("throttling", stats.TimerType, l.statsTagger(key, algo, rate, window)) start := time.Now() return func() { m.Since(start) diff --git a/throttling/throttling_test.go b/throttling/throttling_test.go index 6365a91e..57f97b04 100644 --- a/throttling/throttling_test.go +++ b/throttling/throttling_test.go @@ -11,6 +11,8 @@ import ( "github.com/ory/dockertest/v3" "github.com/stretchr/testify/require" + "github.com/rudderlabs/rudder-go-kit/stats" + "github.com/rudderlabs/rudder-go-kit/stats/memstats" "github.com/rudderlabs/rudder-go-kit/testhelper/rand" ) @@ -407,6 +409,33 @@ func TestRetryAfter(t *testing.T) { } } +func TestTags(t *testing.T) { + ms, err := memstats.New() + require.NoError(t, err) + + l, err := New( + WithStatsCollector(ms), + WithStatsTagger(func(key, algo string, rate, window int64) stats.Tags { + return stats.Tags{ + "algo": algo, + "custom": "bar", + } + }), + WithInMemoryGCRA(1), + ) + require.NoError(t, err) + allowed, _, err := l.Allow(context.Background(), 1, 1, 1, "foo") + require.NoError(t, err) + require.True(t, allowed) + + metrics := ms.GetByName("throttling") + require.Len(t, metrics, 1) + require.Contains(t, metrics[0].Tags, "algo") + require.Equal(t, "gcra", metrics[0].Tags["algo"]) + require.Contains(t, metrics[0].Tags, "custom") + require.Equal(t, "bar", metrics[0].Tags["custom"]) +} + func testName(name string, rate, window int64) string { return fmt.Sprintf("%s/%d tokens per %ds", name, rate, window) }