Merge pull request #43 from projectdiscovery/feat-can-take

Adding optimistic CanTake method

Author: Mzack9999

keyratelimit.go

@@ -81,6 +81,15 @@ func (m *MultiLimiter) Take(key string) error {
 	return nil
 }
 
+// CanTake checks if the rate limiter with the given key has any token
+func (m *MultiLimiter) CanTake(key string) bool {
+	limiter, err := m.get(key)
+	if err != nil {
+		return false
+	}
+	return limiter.CanTake()
+}
+
 // AddAndTake adds key if not present and then takes token from bucket
 func (m *MultiLimiter) AddAndTake(opts *Options) {
 	if limiter, err := m.get(opts.Key); err == nil {

ratelimit.go

@@ -49,6 +49,11 @@ func (limiter *Limiter) Take() {
 	<-limiter.tokens
 }
 
+// CanTake checks if the rate limiter has any token
+func (limiter *Limiter) CanTake() bool {
+	return limiter.count.Load() > 0
+}
+
 // GetLimit returns current rate limit per given duration
 func (limiter *Limiter) GetLimit() uint {
 	return uint(limiter.maxCount)

ratelimit_test.go

@@ -85,4 +85,14 @@ func TestRateLimit(t *testing.T) {
 		expected := time.Duration(6) * time.Second
 		require.GreaterOrEqualf(t, timetaken.Nanoseconds(), expected.Nanoseconds(), "more tokens sent than expected with ratelimit")
 	})
+
+	t.Run("Test Take and CanTake", func(t *testing.T) {
+		limiter := New(context.TODO(), 3, time.Hour)
+
+		require.True(t, limiter.CanTake())
+		limiter.Take()
+		limiter.Take()
+		limiter.Take()
+		require.False(t, limiter.CanTake())
+	})
 }