Add Get interface for retrieving current values (#20)

memorystore/store.go

@@ -143,6 +143,25 @@ func (s *store) Take(ctx context.Context, key string) (uint64, uint64, uint64, b
 	return b.take()
 }
 
+// Get retrieves the information about the key, if any exists.
+func (s *store) Get(ctx context.Context, key string) (uint64, uint64, error) {
+	// If the store is stopped, all requests are rejected.
+	if atomic.LoadUint32(&s.stopped) == 1 {
+		return 0, 0, limiter.ErrStopped
+	}
+
+	// Acquire a read lock first - this allows other to concurrently check limits
+	// without taking a full lock.
+	s.dataLock.RLock()
+	if b, ok := s.data[key]; ok {
+		s.dataLock.RUnlock()
+		return b.get()
+	}
+	s.dataLock.RUnlock()
+
+	return 0, 0, nil
+}
+
 // Set configures the bucket-specific tokens and interval.
 func (s *store) Set(ctx context.Context, key string, tokens uint64, interval time.Duration) error {
 	s.dataLock.Lock()
@@ -261,6 +280,16 @@ func newBucket(tokens uint64, interval time.Duration) *bucket {
 	return b
 }
 
+// get returns information about the bucket.
+func (b *bucket) get() (tokens uint64, remaining uint64, retErr error) {
+	b.lock.Lock()
+	defer b.lock.Unlock()
+
+	tokens = b.maxTokens
+	remaining = b.availableTokens
+	return
+}
+
 // take attempts to remove a token from the bucket. If there are no tokens
 // available and the clock has ticked forward, it recalculates the number of
 // tokens and retries. It returns the limit, remaining tokens, time until

memorystore/store_test.go

@@ -45,6 +45,21 @@ func TestStore_Exercise(t *testing.T) {
 
 	key := testKey(t)
 
+	// Get when no config exists
+	{
+		limit, remaining, err := s.(*store).Get(ctx, key)
+		if err != nil {
+			t.Fatal(err)
+		}
+
+		if got, want := limit, uint64(0); got != want {
+			t.Errorf("expected %v to be %v", got, want)
+		}
+		if got, want := remaining, uint64(0); got != want {
+			t.Errorf("expected %v to be %v", got, want)
+		}
+	}
+
 	// Take with no key configuration - this should use the default values
 	{
 		limit, remaining, reset, ok, err := s.Take(ctx, key)
@@ -65,13 +80,41 @@ func TestStore_Exercise(t *testing.T) {
 		}
 	}
 
+	// Get the value
+	{
+		limit, remaining, err := s.(*store).Get(ctx, key)
+		if err != nil {
+			t.Fatal(err)
+		}
+		if got, want := limit, uint64(5); got != want {
+			t.Errorf("expected %v to be %v", got, want)
+		}
+		if got, want := remaining, uint64(4); got != want {
+			t.Errorf("expected %v to be %v", got, want)
+		}
+	}
+
 	// Now set a value
 	{
 		if err := s.Set(ctx, key, 11, 5*time.Second); err != nil {
 			t.Fatal(err)
 		}
 	}
 
+	// Get the value again
+	{
+		limit, remaining, err := s.(*store).Get(ctx, key)
+		if err != nil {
+			t.Fatal(err)
+		}
+		if got, want := limit, uint64(11); got != want {
+			t.Errorf("expected %v to be %v", got, want)
+		}
+		if got, want := remaining, uint64(11); got != want {
+			t.Errorf("expected %v to be %v", got, want)
+		}
+	}
+
 	// Take again, this should use the new values
 	{
 		limit, remaining, reset, ok, err := s.Take(ctx, key)
@@ -92,6 +135,20 @@ func TestStore_Exercise(t *testing.T) {
 		}
 	}
 
+	// Get the value again
+	{
+		limit, remaining, err := s.(*store).Get(ctx, key)
+		if err != nil {
+			t.Fatal(err)
+		}
+		if got, want := limit, uint64(11); got != want {
+			t.Errorf("expected %v to be %v", got, want)
+		}
+		if got, want := remaining, uint64(10); got != want {
+			t.Errorf("expected %v to be %v", got, want)
+		}
+	}
+
 	// Burst and take
 	{
 		if err := s.Burst(ctx, key, 5); err != nil {
@@ -115,6 +172,20 @@ func TestStore_Exercise(t *testing.T) {
 			t.Errorf("expected %v to less than %v", got, want)
 		}
 	}
+
+	// Get the value one final time
+	{
+		limit, remaining, err := s.(*store).Get(ctx, key)
+		if err != nil {
+			t.Fatal(err)
+		}
+		if got, want := limit, uint64(11); got != want {
+			t.Errorf("expected %v to be %v", got, want)
+		}
+		if got, want := remaining, uint64(14); got != want {
+			t.Errorf("expected %v to be %v", got, want)
+		}
+	}
 }
 
 func TestStore_Take(t *testing.T) {

noopstore/store.go

@@ -22,6 +22,11 @@ func (s *store) Take(_ context.Context, _ string) (uint64, uint64, uint64, bool,
 	return 0, 0, 0, true, nil
 }
 
+// Get does nothing.
+func (s *store) Get(_ context.Context, _ string) (uint64, uint64, error) {
+	return 0, 0, nil
+}
+
 // Set does nothing.
 func (s *store) Set(_ context.Context, _ string, _ uint64, _ time.Duration) error {
 	return nil

store.go

@@ -37,6 +37,10 @@ type Store interface {
 	// See the note about keys on the interface documentation.
 	Take(ctx context.Context, key string) (tokens, remaining, reset uint64, ok bool, err error)
 
+	// Get gets the current limit and remaining tokens for the provided key. It
+	// does not change any of the values.
+	Get(ctx context.Context, key string) (tokens, remaining uint64, err error)
+
 	// Set configures the limit at the provided key. If a limit already exists, it
 	// is overwritten. This also sets the number of tokens in the bucket to the
 	// limit.