Refactor memory store ScanKeys implementation

Author: bootjp

store/bolt_store.go

@@ -95,6 +95,34 @@ func (s *boltStore) Scan(ctx context.Context, start []byte, end []byte, limit in
 	return res, errors.WithStack(err)
 }
 
+func (s *boltStore) ScanKeys(ctx context.Context, start []byte, end []byte, limit int) ([][]byte, error) {
+	s.log.InfoContext(ctx, "ScanKeys",
+		slog.String("start", string(start)),
+		slog.String("end", string(end)),
+		slog.Int("limit", limit),
+	)
+
+	var res [][]byte
+
+	err := s.bbolt.View(func(tx *bbolt.Tx) error {
+		b := tx.Bucket(defaultBucket)
+		if b == nil {
+			return nil
+		}
+
+		c := b.Cursor()
+		for k, _ := c.Seek(start); k != nil && (end == nil || bytes.Compare(k, end) < 0); k, _ = c.Next() {
+			res = append(res, k)
+			if len(res) >= limit {
+				break
+			}
+		}
+		return nil
+	})
+
+	return res, errors.WithStack(err)
+}
+
 func (s *boltStore) Put(ctx context.Context, key []byte, value []byte) error {
 	s.log.InfoContext(ctx, "put",
 		slog.String("key", string(key)),

store/bolt_store_test.go

@@ -83,6 +83,45 @@ func TestBoltStore_Scan(t *testing.T) {
 	assert.Equal(t, 100, cnt)
 }
 
+func TestBoltStore_ScanKeys(t *testing.T) {
+	ctx := context.Background()
+	t.Parallel()
+	st := mustStore(NewBoltStore(t.TempDir() + "/bolt.db"))
+
+	for i := 0; i < 999; i++ {
+		keyStr := "prefix " + strconv.Itoa(i) + "foo"
+		key := []byte(keyStr)
+		b := make([]byte, 8)
+		binary.PutVarint(b, int64(i))
+		err := st.Put(ctx, key, b)
+		assert.NoError(t, err)
+	}
+
+	res, err := st.ScanKeys(ctx, []byte("prefix"), []byte("z"), 100)
+	assert.NoError(t, err)
+	assert.Equal(t, 100, len(res))
+
+	sortedKeys := make([][]byte, 999)
+
+	for _, k := range res {
+		str := string(k)
+		i, err := strconv.Atoi(str[7 : len(str)-3])
+		assert.NoError(t, err)
+		sortedKeys[i] = k
+	}
+
+	cnt := 0
+	for i, k := range sortedKeys {
+		if k == nil {
+			continue
+		}
+		cnt++
+		assert.Equal(t, []byte("prefix "+strconv.Itoa(i)+"foo"), k)
+	}
+
+	assert.Equal(t, 100, cnt)
+}
+
 func TestBoltStore_Txn(t *testing.T) {
 	t.Parallel()
 	t.Run("success", func(t *testing.T) {

store/rb_memory_store.go

@@ -126,6 +126,40 @@ func (s *rbMemoryStore) Scan(ctx context.Context, start []byte, end []byte, limi
 	return result, nil
 }
 
+func (s *rbMemoryStore) ScanKeys(ctx context.Context, start []byte, end []byte, limit int) ([][]byte, error) {
+	s.mtx.RLock()
+	defer s.mtx.RUnlock()
+
+	var result [][]byte
+
+	it := s.tree.Iterator()
+
+	var ok bool
+	if start != nil {
+		it.Begin()
+		ok = it.NextTo(func(key, _ interface{}) bool {
+			k, _ := key.([]byte)
+			return bytes.Compare(k, start) >= 0
+		})
+	} else {
+		ok = it.First()
+	}
+
+	for ; ok && len(result) < limit; ok = it.Next() {
+		k, _ := it.Key().([]byte)
+
+		if end != nil && bytes.Compare(k, end) > 0 {
+			break
+		}
+
+		keyCopy := make([]byte, len(k))
+		copy(keyCopy, k)
+		result = append(result, keyCopy)
+	}
+
+	return result, nil
+}
+
 func (s *rbMemoryStore) Put(ctx context.Context, key []byte, value []byte) error {
 	s.mtx.Lock()
 	defer s.mtx.Unlock()

store/rb_memory_store_test.go

@@ -85,6 +85,64 @@ func TestRbMemoryStore_Scan(t *testing.T) {
 	assert.Equal(t, 100, cnt)
 }
 
+func TestRbMemoryStore_ScanKeys(t *testing.T) {
+	ctx := context.Background()
+	t.Parallel()
+	st := NewRbMemoryStore()
+
+	for i := 0; i < 9999; i++ {
+		keyStr := "prefix " + strconv.Itoa(i) + "foo"
+		key := []byte(keyStr)
+		b := make([]byte, 8)
+		binary.PutVarint(b, int64(i))
+		err := st.Put(ctx, key, b)
+		assert.NoError(t, err)
+	}
+
+	res, err := st.ScanKeys(ctx, []byte("prefix"), []byte("z"), 100)
+	assert.NoError(t, err)
+	assert.Equal(t, 100, len(res))
+
+	sortedKeys := make([][]byte, 9999)
+
+	for _, k := range res {
+		str := string(k)
+		i, err := strconv.Atoi(str[7 : len(str)-3])
+		assert.NoError(t, err)
+		sortedKeys[i] = k
+	}
+
+	cnt := 0
+	for i, k := range sortedKeys {
+		if k == nil {
+			continue
+		}
+		cnt++
+		assert.Equal(t, []byte("prefix "+strconv.Itoa(i)+"foo"), k)
+	}
+
+	assert.Equal(t, 100, cnt)
+}
+
+func TestRbMemoryStore_ScanKeysReturnsCopies(t *testing.T) {
+	ctx := context.Background()
+	st := NewRbMemoryStore()
+
+	assert.NoError(t, st.Put(ctx, []byte("foo"), []byte("bar")))
+
+	res, err := st.ScanKeys(ctx, nil, nil, 10)
+	assert.NoError(t, err)
+	if len(res) == 0 {
+		t.Fatalf("expected keys, got none")
+	}
+
+	res[0][0] = 'x'
+
+	res2, err := st.ScanKeys(ctx, nil, nil, 10)
+	assert.NoError(t, err)
+	assert.Equal(t, []byte("foo"), res2[0])
+}
+
 func TestRbMemoryStore_Txn(t *testing.T) {
 	t.Parallel()
 	t.Run("success", func(t *testing.T) {

store/store.go

@@ -32,6 +32,7 @@ type Store interface {
 type ScanStore interface {
 	Store
 	Scan(ctx context.Context, start []byte, end []byte, limit int) ([]*KVPair, error)
+	ScanKeys(ctx context.Context, start []byte, end []byte, limit int) ([][]byte, error)
 }
 
 type TTLStore interface {
@@ -56,6 +57,7 @@ type Txn interface {
 type ScanTxn interface {
 	Txn
 	Scan(ctx context.Context, start []byte, end []byte, limit int) ([]*KVPair, error)
+	ScanKeys(ctx context.Context, start []byte, end []byte, limit int) ([][]byte, error)
 }
 
 type TTLTxn interface {