feat: Added Reset to TokenBucket

tokenbucket.go

@@ -122,6 +122,33 @@ func (t *TokenBucket) Limit(ctx context.Context) (time.Duration, error) {
 	return t.Take(ctx, 1)
 }
 
+// Reset the bucket to zero state
+func (t *TokenBucket) Reset(ctx context.Context) (time.Duration, error) {
+	t.mu.Lock()
+	defer t.mu.Unlock()
+	if err := t.locker.Lock(ctx); err != nil {
+		return 0, err
+	}
+	defer func() {
+		if err := t.locker.Unlock(ctx); err != nil {
+			t.logger.Log(err)
+		}
+	}()
+	state, err := t.backend.State(ctx)
+	if err != nil {
+		return 0, err
+	}
+
+	// Reset bucket to initial state
+	state.Available = 0
+	state.Last = 0
+
+	if err = t.backend.SetState(ctx, state); err != nil {
+		return 0, err
+	}
+	return 0, nil
+}
+
 // TokenBucketInMemory is an in-memory implementation of TokenBucketStateBackend.
 //
 // The state is not shared nor persisted so it won't survive restarts or failures.

tokenbucket_test.go

@@ -205,6 +205,50 @@ func (s *LimitersTestSuite) TestTokenBucketRefill() {
 	}
 }
 
+func (s *LimitersTestSuite) TestTokenBucketReset() {
+	for name, backend := range s.tokenBucketBackends() {
+		s.Run(name, func() {
+			clock := newFakeClock()
+
+			bucketCapacity := 2
+			bucket := l.NewTokenBucket(int64(bucketCapacity), time.Second*1, l.NewLockNoop(), backend, clock, s.logger)
+
+			// Check while Bucket is full, partially full and empty
+			noOfAccess := 4
+
+			for access := range noOfAccess {
+				clock.reset()
+				for i := range access {
+					_, err := bucket.Limit(context.TODO())
+
+					// Handle err when bucket capacity is reached
+					if i >= bucketCapacity {
+						s.Require().Equal(l.ErrLimitExhausted, err)
+					} else {
+						s.Require().NoError(err)
+					}
+				}
+
+				_, err := bucket.Reset(context.TODO())
+				s.Require().NoError(err)
+
+				// Since bucket state is zero valued, Limit needs to be invoked
+				wait, err := bucket.Limit(context.TODO())
+				s.Require().NoError(err)
+				s.Equal(time.Duration(0), wait)
+
+				state, err := backend.State(context.TODO())
+				s.Require().NoError(err, "unable to retrieve backend state")
+				s.Require().Equal(int64(bucketCapacity-1), state.Available)
+
+				// Cleanup, Set Bucket to initial state
+				_, err = bucket.Reset(context.TODO())
+				s.Require().NoError(err)
+			}
+		})
+	}
+}
+
 func BenchmarkTokenBuckets(b *testing.B) {
 	s := new(LimitersTestSuite)
 	s.SetT(&testing.T{})