Prefix Scans for Storage Keys

auth.go

@@ -69,17 +69,15 @@ func (a *Auth) CheckPermission(scope, userID string) error {
 func (a *Auth) Users() ([]string, error) {
 	a.logger.Debug("Retrieving all user IDs from storage")
 
-	keys, err := a.store.Keys()
+	keys, err := a.store.KeysWithPrefix(permissionKeyPrefix)
 	if err != nil {
 		return nil, errors.Wrap(err, "failed to load permissions")
 	}
 
 	var userIDs []string
 	for _, key := range keys {
-		if strings.HasPrefix(key, permissionKeyPrefix) {
-			userID := strings.TrimPrefix(key, permissionKeyPrefix)
-			userIDs = append(userIDs, userID)
-		}
+		userID := strings.TrimPrefix(key, permissionKeyPrefix)
+		userIDs = append(userIDs, userID)
 	}
 
 	return userIDs, nil

storage.go

@@ -3,6 +3,7 @@ package joe
 import (
 	"encoding/json"
 	"sort"
+	"strings"
 	"sync"
 
 	"github.com/pkg/errors"
@@ -34,6 +35,14 @@ type Memory interface {
 	Close() error
 }
 
+// PrefixAwareMemory is an optional interface for Memory to implement if the
+// backing storage has a higher performance way to search keys matching a given
+// prefix. If the Memory does not implement this interface, we will iterate
+// over the keys given by Keys().
+type PrefixAwareMemory interface {
+	KeysWithPrefix(prefix string) ([]string, error)
+}
+
 // A MemoryEncoder is used to encode and decode any values that are stored in
 // the Memory. The default implementation that is used by the Storage uses a
 // JSON encoding.
@@ -83,6 +92,31 @@ func (s *Storage) Keys() ([]string, error) {
 	return keys, err
 }
 
+// KeysWithPrefix returns all keys matching a given prefix.
+func (s *Storage) KeysWithPrefix(prefix string) ([]string, error) {
+	if m, ok := s.memory.(PrefixAwareMemory); ok {
+		s.mu.RLock()
+		keys, err := m.KeysWithPrefix(prefix)
+		s.mu.RUnlock()
+
+		sort.Strings(keys)
+		return keys, err
+	}
+
+	keys, err := s.memory.Keys()
+	if err != nil {
+		return nil, err
+	}
+	var results []string
+	for _, k := range keys {
+		if strings.HasPrefix(k, prefix) {
+			results = append(results, k)
+		}
+	}
+	sort.Strings(results)
+	return results, nil
+}
+
 // Set encodes the given data and stores it in the Memory that is managed by the
 // Storage.
 func (s *Storage) Set(key string, value interface{}) error {

storage_test.go

@@ -3,6 +3,9 @@ package joe
 import (
 	"bytes"
 	"encoding/gob"
+	"go.uber.org/zap"
+	"sort"
+	"strings"
 	"testing"
 
 	"github.com/pkg/errors"
@@ -121,3 +124,62 @@ func (e gobEncoder) Decode(data []byte, target interface{}) error {
 	enc := gob.NewDecoder(bytes.NewBuffer(data))
 	return enc.Decode(target)
 }
+
+type prefixAwareMemory struct {
+	*inMemory
+	prefixUsed bool
+}
+
+func (m *prefixAwareMemory) KeysWithPrefix(prefix string) ([]string, error) {
+	m.prefixUsed = true
+	keys, err := m.Keys()
+	if err != nil {
+		return nil, err
+	}
+	var results []string
+	for _, k := range keys {
+		if strings.HasPrefix(k, prefix) {
+			results = append(results, k)
+		}
+	}
+	sort.Strings(results)
+	return results, nil
+}
+
+func newPrefixAwareStorage(logger *zap.Logger) *Storage {
+	return &Storage{
+		logger:  logger,
+		memory:  &prefixAwareMemory{newInMemory(), false},
+		encoder: new(jsonEncoder),
+	}
+}
+
+var _ PrefixAwareMemory = &prefixAwareMemory{nil, false}
+
+func TestStorage_PrefixAware(t *testing.T) {
+	logger := zaptest.NewLogger(t)
+	store := newPrefixAwareStorage(logger)
+
+	testEntries := []string{
+		"test.k3",
+		"test.k1",
+		"non-matching",
+		"test.k2",
+	}
+
+	for _, k := range testEntries {
+		err := store.Set(k, nil)
+		require.NoError(t, err)
+	}
+
+	expectedKeys := []string{
+		"test.k1",
+		"test.k2",
+		"test.k3",
+	}
+	actualKeys, err := store.KeysWithPrefix("test.")
+	require.NoError(t, err)
+
+	assert.Equal(t, expectedKeys, actualKeys)
+	assert.True(t, store.memory.(*prefixAwareMemory).prefixUsed)
+}