chore: throttling tagger (#730)

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

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
 

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 }
 }

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)

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)
 }