From 5dc7ec2bef6ebdb304cf4ee6fa5abcd50dbe5020 Mon Sep 17 00:00:00 2001 From: "Yoshiaki Ueda (bootjp)" Date: Sun, 14 Sep 2025 03:51:45 +0900 Subject: [PATCH] Refactor memory store ScanKeys implementation --- store/bolt_store.go | 28 +++++++++++++++++ store/bolt_store_test.go | 39 +++++++++++++++++++++++ store/rb_memory_store.go | 34 ++++++++++++++++++++ store/rb_memory_store_test.go | 58 +++++++++++++++++++++++++++++++++++ store/store.go | 2 ++ 5 files changed, 161 insertions(+) diff --git a/store/bolt_store.go b/store/bolt_store.go index 1f1f20a..2289b8e 100644 --- a/store/bolt_store.go +++ b/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)), diff --git a/store/bolt_store_test.go b/store/bolt_store_test.go index 43826d2..8830106 100644 --- a/store/bolt_store_test.go +++ b/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) { diff --git a/store/rb_memory_store.go b/store/rb_memory_store.go index 2c96f89..04139fa 100644 --- a/store/rb_memory_store.go +++ b/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() diff --git a/store/rb_memory_store_test.go b/store/rb_memory_store_test.go index 79d2113..3df8e9f 100644 --- a/store/rb_memory_store_test.go +++ b/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) { diff --git a/store/store.go b/store/store.go index b3d8c82..0208779 100644 --- a/store/store.go +++ b/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 {