Fix delete key endpoint

Author: kgantsov

pkg/http_server/handlers.go

@@ -18,10 +18,8 @@ type (
 )
 
 func (h *Handler) Reduce(ctx context.Context, input *ReduceRequest) (*ReduceResponse, error) {
-	key := input.Key
-
 	val, err := h.RateLimiter.Reduce(
-		key, input.MaxTokens, input.RefillTime, input.RefillAmount, input.Tokens,
+		input.Key, input.MaxTokens, input.RefillTime, input.RefillAmount, input.Tokens,
 	)
 
 	if err != nil {
@@ -40,9 +38,9 @@ func (h *Handler) Reduce(ctx context.Context, input *ReduceRequest) (*ReduceResp
 }
 
 func (h *Handler) Remove(ctx context.Context, input *RemoveRequest) (*RemoveResponse, error) {
-	key := input.Key
-
-	h.RateLimiter.Remove(key)
+	h.RateLimiter.Remove(
+		input.Key, input.MaxTokens, input.RefillTime, input.RefillAmount,
+	)
 
 	res := &RemoveResponse{Status: http.StatusOK}
 	res.Body.Status = "OK"

pkg/http_server/http_server.go

@@ -96,7 +96,7 @@ func (h *Handler) RegisterRoutes(api huma.API) {
 		huma.Operation{
 			OperationID: "rate-limiter-delete",
 			Method:      http.MethodDelete,
-			Path:        "/API/v1/rate-limiters/{key}",
+			Path:        "/API/v1/rate-limiters/{key}/{max_tokens}/{refill_time}/{refill_amount}/",
 			Summary:     "Delete the key",
 			Description: "Delete the key from the rate limiter",
 			Tags:        []string{"Rate Limiter"},
@@ -108,7 +108,7 @@ func (h *Handler) RegisterRoutes(api huma.API) {
 		huma.Operation{
 			OperationID: "rate-limiter-get-stats",
 			Method:      http.MethodGet,
-			Path:        "/API/v1/rate-limiters/stats",
+			Path:        "/API/v1/rate-limiters/stats/",
 			Summary:     "Stats of the rate limiter",
 			Description: "Return the stats of the rate limiter service",
 			Tags:        []string{"Rate Limiter"},

pkg/http_server/schemas.go

@@ -20,7 +20,10 @@ type ReduceResponse struct {
 }
 
 type RemoveRequest struct {
-	Key string `path:"key" maxLength:"1024" example:"user:1" doc:"Key for the lock"`
+	Key          string `path:"key" maxLength:"1024" example:"user:1" doc:"Key for the lock"`
+	MaxTokens    int64  `path:"max_tokens" minimum:"1" example:"1000" doc:"Maximum number of tokens in the bucket"`
+	RefillTime   int64  `path:"refill_time" minimum:"1" example:"60" doc:"Time in seconds to refill the bucket"`
+	RefillAmount int64  `path:"refill_amount" minimum:"1" example:"100" doc:"Number of tokens to refill the bucket"`
 }
 
 type RemoveResponse struct {

pkg/limiter/limiter.go

@@ -51,10 +51,12 @@ func (l *RateLimiter) GetShard(key int64) *Shard {
 	return l.shards[uint64(key)%SHARDS]
 }
 
-func (l *RateLimiter) Reduce(key string, maxTokens int64, refillTime int64, refillAmount int64, tokens int64) (int64, error) {
-	defer TimeTrack(time.Now(), "RateLimiter.Reduce")
+func (l *RateLimiter) getKey(key string, maxTokens int64, refillTime int64, refillAmount int64) string {
+	return fmt.Sprintf("%s:%d:%d:%d", key, maxTokens, refillTime, refillAmount)
+}
 
-	key = fmt.Sprintf("%s:%d:%d:%d", key, maxTokens, refillTime, refillAmount)
+func (l *RateLimiter) Reduce(key string, maxTokens int64, refillTime int64, refillAmount int64, tokens int64) (int64, error) {
+	key = l.getKey(key, maxTokens, refillTime, refillAmount)
 
 	h := int64(hash(key))
 
@@ -67,8 +69,6 @@ func (l *RateLimiter) Reduce(key string, maxTokens int64, refillTime int64, refi
 
 	now := time.Now().Unix()
 
-	log.Debug().Msgf("Bucket for key: %s found: %t now %d: %+v", key, ok, now, bucket)
-
 	if !ok {
 		value := maxTokens - tokens
 		tokensNeeded := maxTokens - value
@@ -119,7 +119,6 @@ func (l *RateLimiter) Len() int64 {
 }
 
 func (l *RateLimiter) CleanUpFullBuckets() {
-
 	for _, shard := range l.shards {
 		shard.mu.Lock()
 		for key, bucket := range shard.Buckets {
@@ -143,7 +142,9 @@ func (l *RateLimiter) StartCleanUpFullBuckets() {
 	}
 }
 
-func (l *RateLimiter) Remove(key string) {
+func (l *RateLimiter) Remove(key string, maxTokens int64, refillTime int64, refillAmount int64) {
+	key = l.getKey(key, maxTokens, refillTime, refillAmount)
+
 	h := int64(hash(key))
 	shard := l.GetShard(h)
 

pkg/limiter/limiter_test.go

@@ -74,6 +74,24 @@ func TestRateLimiterWithManyKeys(t *testing.T) {
 	}
 }
 
+func TestRateLimiterReuseTheSameKey(t *testing.T) {
+
+	rl := NewRateLimiter()
+
+	val, _ := rl.Reduce("user:123", 1000, 1, 50, 1)
+	assert.Equal(t, int64(999), val)
+	val, _ = rl.Reduce("user:123", 1000, 1, 50, 1)
+	assert.Equal(t, int64(998), val)
+	val, _ = rl.Reduce("user:123", 1000, 1, 50, 1)
+	assert.Equal(t, int64(997), val)
+
+	val, _ = rl.Reduce("user:123", 50, 1, 50, 1)
+	assert.Equal(t, int64(49), val)
+
+	val, _ = rl.Reduce("user:123", 1000, 1, 50, 1)
+	assert.Equal(t, int64(996), val)
+}
+
 func TestRateLimiterD(t *testing.T) {
 	rl := NewRateLimiter()