From bd4a75cfd0663b893bbb22731db5b59646839245 Mon Sep 17 00:00:00 2001 From: bootjp Date: Sat, 11 Jul 2026 13:39:34 +0900 Subject: [PATCH 01/10] backup: add live scan primitives --- kv/backup_scan.go | 93 +++++++++++++++ kv/leader_routed_store.go | 26 ++++ kv/shard_store.go | 186 ++++++++++++++++++++++++++++- kv/shard_store_test.go | 59 +++++++++ store/lsm_store.go | 72 +++++++++++ store/lsm_store_encryption_test.go | 27 +++++ store/mvcc_store.go | 38 ++++++ store/scan_keys_test.go | 81 +++++++++++++ store/store.go | 2 + 9 files changed, 579 insertions(+), 5 deletions(-) create mode 100644 kv/backup_scan.go create mode 100644 store/scan_keys_test.go diff --git a/kv/backup_scan.go b/kv/backup_scan.go new file mode 100644 index 000000000..9413ae185 --- /dev/null +++ b/kv/backup_scan.go @@ -0,0 +1,93 @@ +package kv + +import ( + "bytes" + "context" + + "github.com/bootjp/elastickv/store" +) + +const defaultBackupScanPageSize = 1024 + +// BackupScanner pages through ShardStore.ScanAt without holding store locks +// across pages. +type BackupScanner interface { + Next(ctx context.Context) (*store.KVPair, bool, error) + Close() error +} + +type backupScanner struct { + store *ShardStore + end []byte + ts uint64 + pageSize int + cursor []byte + page []*store.KVPair + index int + closed bool +} + +func NewBackupScanner(st *ShardStore, start []byte, end []byte, ts uint64, pageSize int) BackupScanner { + if pageSize <= 0 { + pageSize = defaultBackupScanPageSize + } + return &backupScanner{ + store: st, + cursor: bytes.Clone(start), + end: bytes.Clone(end), + ts: ts, + pageSize: pageSize, + } +} + +func (s *ShardStore) NewBackupScanner(start []byte, end []byte, ts uint64, pageSize int) BackupScanner { + return NewBackupScanner(s, start, end, ts, pageSize) +} + +func (s *backupScanner) Next(ctx context.Context) (*store.KVPair, bool, error) { + if s.closed || s.store == nil { + return nil, false, nil + } + for s.index >= len(s.page) { + if err := s.loadNextPage(ctx); err != nil { + return nil, false, err + } + if len(s.page) == 0 { + return nil, false, nil + } + } + kvp := s.page[s.index] + s.index++ + if kvp == nil { + return nil, true, nil + } + return &store.KVPair{ + Key: bytes.Clone(kvp.Key), + Value: bytes.Clone(kvp.Value), + }, true, nil +} + +func (s *backupScanner) Close() error { + s.closed = true + s.page = nil + return nil +} + +func (s *backupScanner) loadNextPage(ctx context.Context) error { + page, err := s.store.ScanAt(ctx, s.cursor, s.end, s.pageSize, s.ts) + if err != nil { + return err + } + s.page = page + s.index = 0 + if len(page) == 0 { + return nil + } + last := page[len(page)-1] + if last == nil || len(last.Key) == 0 { + s.cursor = nil + return nil + } + s.cursor = nextScanCursor(last.Key) + return nil +} diff --git a/kv/leader_routed_store.go b/kv/leader_routed_store.go index cdd4f2f4c..a9bec69d4 100644 --- a/kv/leader_routed_store.go +++ b/kv/leader_routed_store.go @@ -260,6 +260,32 @@ func (s *LeaderRoutedStore) ScanAt(ctx context.Context, start []byte, end []byte return s.proxyRawScanAt(ctx, start, end, limit, ts, false) } +func (s *LeaderRoutedStore) ScanKeysAt(ctx context.Context, start []byte, end []byte, limit int, ts uint64) ([][]byte, error) { + if s == nil || s.local == nil { + return [][]byte{}, nil + } + if limit <= 0 { + return [][]byte{}, nil + } + ok, fenceTS := s.leaderFenceTS(ctx, start) + if ok { + keys, err := s.local.ScanKeysAt(ctx, start, end, limit, max(ts, fenceTS)) + return keys, errors.WithStack(err) + } + kvs, err := s.proxyRawScanAt(ctx, start, end, limit, ts, false) + if err != nil { + return nil, err + } + keys := make([][]byte, 0, len(kvs)) + for _, kvp := range kvs { + if kvp == nil { + continue + } + keys = append(keys, bytes.Clone(kvp.Key)) + } + return keys, nil +} + func (s *LeaderRoutedStore) ReverseScanAt(ctx context.Context, start []byte, end []byte, limit int, ts uint64) ([]*store.KVPair, error) { if s == nil || s.local == nil { return []*store.KVPair{}, nil diff --git a/kv/shard_store.go b/kv/shard_store.go index e78381bd3..26bd0eefb 100644 --- a/kv/shard_store.go +++ b/kv/shard_store.go @@ -179,11 +179,11 @@ func tryEngineLinearizableFence(ctx context.Context, engine raftengine.LeaderVie return err == nil } -// ScanAt scans keys across shards at the given timestamp. Note: when the range -// spans multiple shards, each shard may have a different Raft apply position. -// This means the returned view is NOT a globally consistent snapshot — it is -// a best-effort point-in-time scan. Callers requiring cross-shard consistency -// should use a transaction or implement a cross-shard snapshot fence. +// ScanAt scans keys across shards at the given timestamp. When the caller has +// already fenced every group to applied_index >= f(ts), as BeginBackup does, +// the result is consistent across groups. Without that fence, ranges spanning +// multiple shards are best-effort because each shard may have a different Raft +// apply position. func (s *ShardStore) ScanAt(ctx context.Context, start []byte, end []byte, limit int, ts uint64) ([]*store.KVPair, error) { if limit <= 0 { return []*store.KVPair{}, nil @@ -204,6 +204,25 @@ func (s *ShardStore) ScanAt(ctx context.Context, start []byte, end []byte, limit return out, nil } +func (s *ShardStore) ScanKeysAt(ctx context.Context, start []byte, end []byte, limit int, ts uint64) ([][]byte, error) { + if limit <= 0 { + return [][]byte{}, nil + } + + routes, clampToRoutes := s.routesForScan(start, end) + out, err := s.scanKeyRoutesAt(ctx, routes, start, end, limit, ts, clampToRoutes) + if err != nil { + return nil, err + } + sort.Slice(out, func(i, j int) bool { + return bytes.Compare(out[i], out[j]) < 0 + }) + if len(out) > limit { + out = out[:limit] + } + return out, nil +} + // ScanGroupAt scans a range on the explicitly selected Raft group. // It is for keyspaces whose owner is resolved outside the byte-range // engine (for example SQS HT-FIFO's (queue, partition) resolver). @@ -284,6 +303,33 @@ func (s *ShardStore) scanRoutesAt(ctx context.Context, routes []distribution.Rou return out, nil } +func (s *ShardStore) scanKeyRoutesAt(ctx context.Context, routes []distribution.Route, start []byte, end []byte, limit int, ts uint64, clampToRoutes bool) ([][]byte, error) { + out := make([][]byte, 0) + for _, route := range routes { + scanStart := start + scanEnd := end + if clampToRoutes { + scanStart = clampScanStart(start, route.Start) + scanEnd = clampScanEnd(end, route.End) + } + + keys, err := s.scanKeyRouteAt(ctx, route, scanStart, scanEnd, limit, ts) + if err != nil { + return nil, err + } + if clampToRoutes { + out = append(out, keys...) + if len(out) >= limit { + out = out[:limit] + break + } + continue + } + out = mergeAndTrimScanKeys(out, keys, limit) + } + return out, nil +} + func (s *ShardStore) reverseScanRoutesAt( ctx context.Context, routes []distribution.Route, @@ -328,6 +374,85 @@ func (s *ShardStore) reverseScanRoutesAt( return out, nil } +func (s *ShardStore) scanKeyRouteAt( + ctx context.Context, + route distribution.Route, + start []byte, + end []byte, + limit int, + ts uint64, +) ([][]byte, error) { + g, ok := s.groupForID(route.GroupID) + if !ok || g == nil || g.Store == nil { + return nil, nil + } + + if engineForGroup(g) == nil { + return s.scanKeysRouteLocal(ctx, g, start, end, limit, ts) + } + + if isLinearizableRaftLeader(ctx, engineForGroup(g)) { + return s.scanKeysRouteAtLeader(ctx, g, start, end, limit, ts) + } + + return s.proxyScanKeysAt(ctx, g, start, end, limit, ts) +} + +func (s *ShardStore) scanKeysRouteLocal( + ctx context.Context, + g *ShardGroup, + start []byte, + end []byte, + limit int, + ts uint64, +) ([][]byte, error) { + keys, err := g.Store.ScanKeysAt(ctx, start, end, limit, ts) + if err != nil { + return nil, errors.WithStack(err) + } + return filterTxnInternalKeys(keys), nil +} + +func (s *ShardStore) scanKeysRouteAtLeader( + ctx context.Context, + g *ShardGroup, + start []byte, + end []byte, + limit int, + ts uint64, +) ([][]byte, error) { + keys, err := g.Store.ScanKeysAt(ctx, start, end, limit, ts) + if err != nil { + return nil, errors.WithStack(err) + } + keyKVs := kvPairsFromKeys(keys) + lockStart, lockEnd := scanLockBoundsForKVs(keyKVs, start, end, limit) + lockKVs, err := scanTxnLockRangeAt(ctx, g, lockStart, lockEnd, ts, limit) + if err != nil { + return nil, err + } + kvs, err := s.resolveScanLocks(ctx, g, keyKVs, lockKVs, ts) + if err != nil { + return nil, err + } + return keysFromKVs(kvs), nil +} + +func (s *ShardStore) proxyScanKeysAt( + ctx context.Context, + g *ShardGroup, + start []byte, + end []byte, + limit int, + ts uint64, +) ([][]byte, error) { + kvs, err := s.proxyRawScanAt(ctx, g, start, end, limit, ts, false, 0) + if err != nil { + return nil, err + } + return keysFromKVs(filterTxnInternalKVs(kvs)), nil +} + func (s *ShardStore) clampedReverseScanRouteAt( ctx context.Context, route distribution.Route, @@ -497,6 +622,21 @@ func mergeAndTrimScanResults(out []*store.KVPair, kvs []*store.KVPair, limit int return out[:limit] } +func mergeAndTrimScanKeys(out [][]byte, keys [][]byte, limit int) [][]byte { + if len(keys) == 0 { + return out + } + out = append(out, keys...) + if len(out) <= limit { + return out + } + sort.Slice(out, func(i, j int) bool { + return bytes.Compare(out[i], out[j]) < 0 + }) + clear(out[limit:]) + return out[:limit] +} + func mergeAndTrimReverseScanResults(out []*store.KVPair, kvs []*store.KVPair, limit int) []*store.KVPair { if len(kvs) == 0 { return out @@ -523,6 +663,42 @@ func countNonInternalKVs(kvs []*store.KVPair) int { return count } +func kvPairsFromKeys(keys [][]byte) []*store.KVPair { + kvs := make([]*store.KVPair, 0, len(keys)) + for _, key := range keys { + if len(key) == 0 { + continue + } + kvs = append(kvs, &store.KVPair{Key: bytes.Clone(key)}) + } + return kvs +} + +func keysFromKVs(kvs []*store.KVPair) [][]byte { + keys := make([][]byte, 0, len(kvs)) + for _, kvp := range kvs { + if kvp == nil || len(kvp.Key) == 0 { + continue + } + keys = append(keys, bytes.Clone(kvp.Key)) + } + return keys +} + +func filterTxnInternalKeys(keys [][]byte) [][]byte { + if len(keys) == 0 { + return keys + } + out := make([][]byte, 0, len(keys)) + for _, key := range keys { + if isTxnInternalKey(key) { + continue + } + out = append(out, key) + } + return out +} + func (s *ShardStore) groupForID(groupID uint64) (*ShardGroup, bool) { g, ok := s.groups[groupID] return g, ok diff --git a/kv/shard_store_test.go b/kv/shard_store_test.go index 4f0da38b7..5485b4be4 100644 --- a/kv/shard_store_test.go +++ b/kv/shard_store_test.go @@ -144,6 +144,65 @@ func TestShardStoreScanAt_IncludesS3ManifestKeysAcrossShards(t *testing.T) { require.Equal(t, k2, kvs[1].Key) } +func TestShardStoreScanKeysAt_IncludesKeysAcrossShards(t *testing.T) { + t.Parallel() + + ctx := context.Background() + engine := distribution.NewEngine() + engine.UpdateRoute([]byte(""), []byte("m"), 1) + engine.UpdateRoute([]byte("m"), nil, 2) + + groups := map[uint64]*ShardGroup{ + 1: {Store: store.NewMVCCStore()}, + 2: {Store: store.NewMVCCStore()}, + } + st := NewShardStore(engine, groups) + + require.NoError(t, st.PutAt(ctx, []byte("a"), []byte("va"), 1, 0)) + require.NoError(t, st.PutAt(ctx, []byte("b"), []byte("vb"), 2, 0)) + require.NoError(t, st.DeleteAt(ctx, []byte("b"), 3)) + require.NoError(t, st.PutAt(ctx, []byte("x"), []byte("vx"), 4, 0)) + require.NoError(t, st.PutAt(ctx, []byte("z"), []byte("vz"), 5, 0)) + + keys, err := st.ScanKeysAt(ctx, []byte(""), nil, 2, ^uint64(0)) + require.NoError(t, err) + require.Equal(t, [][]byte{[]byte("a"), []byte("x")}, keys) +} + +func TestBackupScannerPaging(t *testing.T) { + t.Parallel() + + ctx := context.Background() + engine := distribution.NewEngine() + engine.UpdateRoute([]byte(""), []byte("m"), 1) + engine.UpdateRoute([]byte("m"), nil, 2) + + groups := map[uint64]*ShardGroup{ + 1: {Store: store.NewMVCCStore()}, + 2: {Store: store.NewMVCCStore()}, + } + st := NewShardStore(engine, groups) + var commitTS uint64 = 1 + for _, key := range [][]byte{[]byte("a"), []byte("b"), []byte("c"), []byte("x"), []byte("z")} { + require.NoError(t, st.PutAt(ctx, key, []byte("v"), commitTS, 0)) + commitTS++ + } + + sc := st.NewBackupScanner([]byte(""), nil, ^uint64(0), 2) + defer sc.Close() + + var got [][]byte + for { + kvp, ok, err := sc.Next(ctx) + require.NoError(t, err) + if !ok { + break + } + got = append(got, kvp.Key) + } + require.Equal(t, [][]byte{[]byte("a"), []byte("b"), []byte("c"), []byte("x"), []byte("z")}, got) +} + func TestShardStoreScanAt_RoutesS3ManifestScansByLogicalObjectKey(t *testing.T) { t.Parallel() diff --git a/store/lsm_store.go b/store/lsm_store.go index 1bf9f4574..570f60aca 100644 --- a/store/lsm_store.go +++ b/store/lsm_store.go @@ -979,6 +979,20 @@ func (s *pebbleStore) processFoundValue(iter *pebble.Iterator, userKey []byte, t return nil, nil } +func (s *pebbleStore) foundValueVisible(iter *pebble.Iterator, ts uint64) (bool, error) { + valBytes := iter.Value() + sv, err := decodeValue(valBytes) + if err != nil { + return false, err + } + + // Keep key-only scans on the same authenticated visibility path as ScanAt. + if _, err := s.decryptForKey(iter.Key(), sv, sv.Value); err != nil { + return false, err + } + return !sv.Tombstone && (sv.ExpireAt == 0 || sv.ExpireAt > ts), nil +} + func (s *pebbleStore) seekToVisibleVersion(iter *pebble.Iterator, userKey []byte, currentVersion, ts uint64) bool { if currentVersion <= ts { return true @@ -1083,6 +1097,39 @@ func (s *pebbleStore) collectScanResults(iter *pebble.Iterator, start, end []byt return result, nil } +func (s *pebbleStore) collectScanKeys(iter *pebble.Iterator, start, end []byte, limit int, ts uint64) ([][]byte, error) { + result := make([][]byte, 0, boundedScanResultCapacity(limit)) + + for iter.SeekGE(encodeKey(start, math.MaxUint64)); iter.Valid() && len(result) < limit; { + userKey, version, ok := nextScannableUserKey(iter) + if !ok { + break + } + + if pastScanEnd(userKey, end) { + break + } + + if !s.seekToVisibleVersion(iter, userKey, version, ts) { + continue + } + + visible, err := s.foundValueVisible(iter, ts) + if err != nil { + return nil, err + } + if visible { + result = append(result, bytes.Clone(userKey)) + } + + if !s.skipToNextUserKey(iter, userKey) { + break + } + } + + return result, nil +} + func (s *pebbleStore) ScanAt(ctx context.Context, start []byte, end []byte, limit int, ts uint64) ([]*KVPair, error) { if limit <= 0 { return []*KVPair{}, nil @@ -1108,6 +1155,31 @@ func (s *pebbleStore) ScanAt(ctx context.Context, start []byte, end []byte, limi return s.collectScanResults(iter, start, end, limit, ts) } +func (s *pebbleStore) ScanKeysAt(ctx context.Context, start []byte, end []byte, limit int, ts uint64) ([][]byte, error) { + if limit <= 0 { + return [][]byte{}, nil + } + + // Acquire dbMu.RLock before reading effectiveMinRetainedTS (which takes + // mtx.RLock) to preserve the global lock ordering: dbMu before mtx. + s.dbMu.RLock() + defer s.dbMu.RUnlock() + + if readTSCompacted(ts, s.effectiveMinRetainedTS()) { + return nil, ErrReadTSCompacted + } + + iter, err := s.db.NewIter(&pebble.IterOptions{ + LowerBound: encodeKey(start, math.MaxUint64), + }) + if err != nil { + return nil, errors.WithStack(err) + } + defer iter.Close() + + return s.collectScanKeys(iter, start, end, limit, ts) +} + func (s *pebbleStore) collectReverseScanResults( iter *pebble.Iterator, start []byte, diff --git a/store/lsm_store_encryption_test.go b/store/lsm_store_encryption_test.go index c5bfdb5de..5f1a6680c 100644 --- a/store/lsm_store_encryption_test.go +++ b/store/lsm_store_encryption_test.go @@ -686,6 +686,33 @@ func TestEncryption_ValueHeaderTamperRejected(t *testing.T) { } } +func TestEncryptionScanKeysAtHeaderTamperRejected(t *testing.T) { + t.Parallel() + f := newEncryptedStoreFixture(t, 1059) + ctx := context.Background() + const writeTS uint64 = 100 + const readTS uint64 = 200 + if err := f.mvcc.PutAt(ctx, []byte("scan-key"), []byte("payload"), writeTS, 0); err != nil { + t.Fatalf("PutAt: %v", err) + } + keys, err := f.mvcc.ScanKeysAt(ctx, []byte("scan-"), []byte("scan."), 10, readTS) + if err != nil { + t.Fatalf("ScanKeysAt before tamper: %v", err) + } + if len(keys) != 1 || !bytes.Equal(keys[0], []byte("scan-key")) { + t.Fatalf("ScanKeysAt before tamper = %q, want scan-key", keys) + } + + f.tamperPebbleValue(t, []byte("scan-key"), writeTS, func(raw []byte) []byte { + raw[0] |= tombstoneMask + return raw + }) + _, err = f.mvcc.ScanKeysAt(ctx, []byte("scan-"), []byte("scan."), 10, readTS) + if !errors.Is(err, ErrEncryptedReadIntegrity) { + t.Fatalf("ScanKeysAt over a tampered encrypted entry should fail integrity, got %v", err) + } +} + // TestEncryption_DeletePrefixHeaderTamperRejected covers the PR742 // claude[bot] round-5 follow-up: isVisibleLiveKey is the write-side // counterpart to readVisibleVersion / processFoundValue (read paths diff --git a/store/mvcc_store.go b/store/mvcc_store.go index 1886fbf3b..0c4a02624 100644 --- a/store/mvcc_store.go +++ b/store/mvcc_store.go @@ -328,6 +328,25 @@ func (s *mvccStore) collectScanResults(it *treemap.Iterator, end []byte, limit, return result } +func (s *mvccStore) collectScanKeys(it *treemap.Iterator, end []byte, limit, capHint int, ts uint64) [][]byte { + result := make([][]byte, 0, capHint) + for ok := true; ok && len(result) < limit; ok = it.Next() { + k, keyOK := it.Key().([]byte) + if !keyOK { + continue + } + if end != nil && bytes.Compare(k, end) >= 0 { + break + } + versions, _ := it.Value().([]VersionedValue) + if _, visible := visibleValue(versions, ts); !visible { + continue + } + result = append(result, bytes.Clone(k)) + } + return result +} + func (s *mvccStore) ScanAt(_ context.Context, start []byte, end []byte, limit int, ts uint64) ([]*KVPair, error) { s.mtx.RLock() defer s.mtx.RUnlock() @@ -347,6 +366,25 @@ func (s *mvccStore) ScanAt(_ context.Context, start []byte, end []byte, limit in return s.collectScanResults(&it, end, limit, capHint, ts), nil } +func (s *mvccStore) ScanKeysAt(_ context.Context, start []byte, end []byte, limit int, ts uint64) ([][]byte, error) { + s.mtx.RLock() + defer s.mtx.RUnlock() + if readTSCompacted(ts, s.minRetainedTS) { + return nil, ErrReadTSCompacted + } + + if limit <= 0 { + return [][]byte{}, nil + } + + capHint := computeScanCapHint(s.tree.Size(), limit) + it := s.tree.Iterator() + if !seekForwardIteratorStart(s.tree, &it, start) { + return make([][]byte, 0, capHint), nil + } + return s.collectScanKeys(&it, end, limit, capHint, ts), nil +} + // seekForwardIteratorStart positions the iterator at the first key >= start. // Returns true if such a position was found, false if no elements satisfy the bound. func seekForwardIteratorStart(tree *treemap.Map, it *treemap.Iterator, start []byte) bool { diff --git a/store/scan_keys_test.go b/store/scan_keys_test.go new file mode 100644 index 000000000..e32babf8e --- /dev/null +++ b/store/scan_keys_test.go @@ -0,0 +1,81 @@ +package store + +import ( + "context" + "testing" + + "github.com/stretchr/testify/require" +) + +func TestMVCCStoreScanKeysAtMatchesScanAtVisibility(t *testing.T) { + t.Parallel() + ctx := context.Background() + st := newTestMVCCStore(t) + + require.NoError(t, st.PutAt(ctx, []byte("a"), []byte("va"), 10, 0)) + require.NoError(t, st.PutAt(ctx, []byte("b"), []byte("vb"), 20, 0)) + require.NoError(t, st.DeleteAt(ctx, []byte("b"), 30)) + require.NoError(t, st.PutAt(ctx, []byte("c"), []byte("vc"), 15, 40)) + require.NoError(t, st.PutAt(ctx, []byte("d"), []byte("vd"), 50, 0)) + + keys, err := st.ScanKeysAt(ctx, []byte("a"), []byte("z"), 2, 35) + require.NoError(t, err) + require.Equal(t, [][]byte{[]byte("a"), []byte("c")}, keys) + + kvs, err := st.ScanAt(ctx, []byte("a"), []byte("z"), 10, 35) + require.NoError(t, err) + require.Equal(t, keysFromStoreKVs(kvs), keys) + + keys, err = st.ScanKeysAt(ctx, []byte("a"), []byte("z"), 10, 45) + require.NoError(t, err) + require.Equal(t, [][]byte{[]byte("a")}, keys) +} + +func TestMVCCStoreScanKeysAtReturnsCompactionError(t *testing.T) { + t.Parallel() + ctx := context.Background() + st := newTestMVCCStore(t) + + require.NoError(t, st.PutAt(ctx, []byte("a"), []byte("va"), 10, 0)) + st.SetMinRetainedTS(20) + + _, err := st.ScanKeysAt(ctx, []byte("a"), nil, 10, 19) + require.ErrorIs(t, err, ErrReadTSCompacted) +} + +func TestPebbleStoreScanKeysAtMatchesScanAtVisibility(t *testing.T) { + t.Parallel() + ctx := context.Background() + st, err := NewPebbleStore(t.TempDir()) + require.NoError(t, err) + defer st.Close() + + require.NoError(t, st.PutAt(ctx, []byte("a"), []byte("va"), 10, 0)) + require.NoError(t, st.PutAt(ctx, []byte("b"), []byte("vb"), 20, 0)) + require.NoError(t, st.DeleteAt(ctx, []byte("b"), 30)) + require.NoError(t, st.PutAt(ctx, []byte("c"), []byte("vc"), 15, 40)) + require.NoError(t, st.PutAt(ctx, []byte("d"), []byte("vd"), 50, 0)) + + keys, err := st.ScanKeysAt(ctx, []byte("a"), []byte("z"), 2, 35) + require.NoError(t, err) + require.Equal(t, [][]byte{[]byte("a"), []byte("c")}, keys) + + kvs, err := st.ScanAt(ctx, []byte("a"), []byte("z"), 10, 35) + require.NoError(t, err) + require.Equal(t, keysFromStoreKVs(kvs), keys) + + keys, err = st.ScanKeysAt(ctx, []byte("a"), []byte("z"), 10, 45) + require.NoError(t, err) + require.Equal(t, [][]byte{[]byte("a")}, keys) +} + +func keysFromStoreKVs(kvs []*KVPair) [][]byte { + keys := make([][]byte, 0, len(kvs)) + for _, kvp := range kvs { + if kvp == nil { + continue + } + keys = append(keys, kvp.Key) + } + return keys +} diff --git a/store/store.go b/store/store.go index 3991da19d..600c27cf7 100644 --- a/store/store.go +++ b/store/store.go @@ -115,6 +115,8 @@ type MVCCStore interface { ExistsAt(ctx context.Context, key []byte, ts uint64) (bool, error) // ScanAt returns versions visible at the given timestamp. ScanAt(ctx context.Context, start []byte, end []byte, limit int, ts uint64) ([]*KVPair, error) + // ScanKeysAt returns keys with visible versions at the given timestamp. + ScanKeysAt(ctx context.Context, start []byte, end []byte, limit int, ts uint64) ([][]byte, error) // ReverseScanAt returns visible versions in descending key order for keys in [start, end). ReverseScanAt(ctx context.Context, start []byte, end []byte, limit int, ts uint64) ([]*KVPair, error) // PutAt commits a value at the provided commit timestamp and optional expireAt. From a9f258151e53963268c6fa80b27214a620443866 Mon Sep 17 00:00:00 2001 From: bootjp Date: Sat, 11 Jul 2026 13:45:14 +0900 Subject: [PATCH 02/10] backup: fix scanner cursor exhaustion --- kv/backup_scan.go | 31 ++++++++++++++++++------------- kv/leader_routed_store.go | 2 +- kv/shard_store.go | 4 ++-- kv/shard_store_test.go | 27 +++++++++++++++++++++++++++ store/lsm_store.go | 2 +- 5 files changed, 49 insertions(+), 17 deletions(-) diff --git a/kv/backup_scan.go b/kv/backup_scan.go index 9413ae185..08d1edab6 100644 --- a/kv/backup_scan.go +++ b/kv/backup_scan.go @@ -17,14 +17,15 @@ type BackupScanner interface { } type backupScanner struct { - store *ShardStore - end []byte - ts uint64 - pageSize int - cursor []byte - page []*store.KVPair - index int - closed bool + store *ShardStore + end []byte + ts uint64 + pageSize int + cursor []byte + page []*store.KVPair + index int + closed bool + exhausted bool } func NewBackupScanner(st *ShardStore, start []byte, end []byte, ts uint64, pageSize int) BackupScanner { @@ -61,19 +62,22 @@ func (s *backupScanner) Next(ctx context.Context) (*store.KVPair, bool, error) { if kvp == nil { return nil, true, nil } - return &store.KVPair{ - Key: bytes.Clone(kvp.Key), - Value: bytes.Clone(kvp.Value), - }, true, nil + return kvp, true, nil } func (s *backupScanner) Close() error { s.closed = true + s.exhausted = true s.page = nil return nil } func (s *backupScanner) loadNextPage(ctx context.Context) error { + if s.exhausted { + s.page = nil + s.index = 0 + return nil + } page, err := s.store.ScanAt(ctx, s.cursor, s.end, s.pageSize, s.ts) if err != nil { return err @@ -81,11 +85,12 @@ func (s *backupScanner) loadNextPage(ctx context.Context) error { s.page = page s.index = 0 if len(page) == 0 { + s.exhausted = true return nil } last := page[len(page)-1] if last == nil || len(last.Key) == 0 { - s.cursor = nil + s.exhausted = true return nil } s.cursor = nextScanCursor(last.Key) diff --git a/kv/leader_routed_store.go b/kv/leader_routed_store.go index a9bec69d4..99ac1d213 100644 --- a/kv/leader_routed_store.go +++ b/kv/leader_routed_store.go @@ -281,7 +281,7 @@ func (s *LeaderRoutedStore) ScanKeysAt(ctx context.Context, start []byte, end [] if kvp == nil { continue } - keys = append(keys, bytes.Clone(kvp.Key)) + keys = append(keys, kvp.Key) } return keys, nil } diff --git a/kv/shard_store.go b/kv/shard_store.go index 26bd0eefb..e535c15ba 100644 --- a/kv/shard_store.go +++ b/kv/shard_store.go @@ -669,7 +669,7 @@ func kvPairsFromKeys(keys [][]byte) []*store.KVPair { if len(key) == 0 { continue } - kvs = append(kvs, &store.KVPair{Key: bytes.Clone(key)}) + kvs = append(kvs, &store.KVPair{Key: key}) } return kvs } @@ -680,7 +680,7 @@ func keysFromKVs(kvs []*store.KVPair) [][]byte { if kvp == nil || len(kvp.Key) == 0 { continue } - keys = append(keys, bytes.Clone(kvp.Key)) + keys = append(keys, kvp.Key) } return keys } diff --git a/kv/shard_store_test.go b/kv/shard_store_test.go index 5485b4be4..c28384db5 100644 --- a/kv/shard_store_test.go +++ b/kv/shard_store_test.go @@ -203,6 +203,33 @@ func TestBackupScannerPaging(t *testing.T) { require.Equal(t, [][]byte{[]byte("a"), []byte("b"), []byte("c"), []byte("x"), []byte("z")}, got) } +func TestBackupScannerEmptyKeyTerminatesAfterPage(t *testing.T) { + t.Parallel() + + ctx := context.Background() + engine := distribution.NewEngine() + engine.UpdateRoute([]byte(""), nil, 1) + groups := map[uint64]*ShardGroup{ + 1: {Store: store.NewMVCCStore()}, + } + st := NewShardStore(engine, groups) + require.NoError(t, st.PutAt(ctx, []byte(""), []byte("empty"), 1, 0)) + + sc := st.NewBackupScanner(nil, []byte("\x00"), ^uint64(0), 1) + defer sc.Close() + + kvp, ok, err := sc.Next(ctx) + require.NoError(t, err) + require.True(t, ok) + require.Equal(t, []byte(""), kvp.Key) + require.Equal(t, []byte("empty"), kvp.Value) + + kvp, ok, err = sc.Next(ctx) + require.NoError(t, err) + require.False(t, ok) + require.Nil(t, kvp) +} + func TestShardStoreScanAt_RoutesS3ManifestScansByLogicalObjectKey(t *testing.T) { t.Parallel() diff --git a/store/lsm_store.go b/store/lsm_store.go index 570f60aca..5d85b2699 100644 --- a/store/lsm_store.go +++ b/store/lsm_store.go @@ -1119,7 +1119,7 @@ func (s *pebbleStore) collectScanKeys(iter *pebble.Iterator, start, end []byte, return nil, err } if visible { - result = append(result, bytes.Clone(userKey)) + result = append(result, userKey) } if !s.skipToNextUserKey(iter, userKey) { From 4b859556f9ad714eaf55dbde10c77f5701233db5 Mon Sep 17 00:00:00 2001 From: bootjp Date: Sat, 11 Jul 2026 14:01:12 +0900 Subject: [PATCH 03/10] backup: fix key scan pagination edge cases --- kv/backup_scan.go | 2 +- kv/shard_store.go | 13 ++++++++---- kv/shard_store_test.go | 46 ++++++++++++++++++++++++++++++++++++++++-- 3 files changed, 54 insertions(+), 7 deletions(-) diff --git a/kv/backup_scan.go b/kv/backup_scan.go index 08d1edab6..66eda9d83 100644 --- a/kv/backup_scan.go +++ b/kv/backup_scan.go @@ -89,7 +89,7 @@ func (s *backupScanner) loadNextPage(ctx context.Context) error { return nil } last := page[len(page)-1] - if last == nil || len(last.Key) == 0 { + if last == nil { s.exhausted = true return nil } diff --git a/kv/shard_store.go b/kv/shard_store.go index e535c15ba..7b907d1d2 100644 --- a/kv/shard_store.go +++ b/kv/shard_store.go @@ -305,12 +305,18 @@ func (s *ShardStore) scanRoutesAt(ctx context.Context, routes []distribution.Rou func (s *ShardStore) scanKeyRoutesAt(ctx context.Context, routes []distribution.Route, start []byte, end []byte, limit int, ts uint64, clampToRoutes bool) ([][]byte, error) { out := make([][]byte, 0) + seenGroups := make(map[uint64]struct{}) for _, route := range routes { scanStart := start scanEnd := end if clampToRoutes { scanStart = clampScanStart(start, route.Start) scanEnd = clampScanEnd(end, route.End) + } else { + if _, seen := seenGroups[route.GroupID]; seen { + continue + } + seenGroups[route.GroupID] = struct{}{} } keys, err := s.scanKeyRouteAt(ctx, route, scanStart, scanEnd, limit, ts) @@ -318,9 +324,8 @@ func (s *ShardStore) scanKeyRoutesAt(ctx context.Context, routes []distribution. return nil, err } if clampToRoutes { - out = append(out, keys...) + out = mergeAndTrimScanKeys(out, keys, limit) if len(out) >= limit { - out = out[:limit] break } continue @@ -666,7 +671,7 @@ func countNonInternalKVs(kvs []*store.KVPair) int { func kvPairsFromKeys(keys [][]byte) []*store.KVPair { kvs := make([]*store.KVPair, 0, len(keys)) for _, key := range keys { - if len(key) == 0 { + if key == nil { continue } kvs = append(kvs, &store.KVPair{Key: key}) @@ -677,7 +682,7 @@ func kvPairsFromKeys(keys [][]byte) []*store.KVPair { func keysFromKVs(kvs []*store.KVPair) [][]byte { keys := make([][]byte, 0, len(kvs)) for _, kvp := range kvs { - if kvp == nil || len(kvp.Key) == 0 { + if kvp == nil || kvp.Key == nil { continue } keys = append(keys, kvp.Key) diff --git a/kv/shard_store_test.go b/kv/shard_store_test.go index c28384db5..cc3669bda 100644 --- a/kv/shard_store_test.go +++ b/kv/shard_store_test.go @@ -169,6 +169,41 @@ func TestShardStoreScanKeysAt_IncludesKeysAcrossShards(t *testing.T) { require.Equal(t, [][]byte{[]byte("a"), []byte("x")}, keys) } +func TestShardStoreScanKeysAt_DeduplicatesUnclampedGroups(t *testing.T) { + t.Parallel() + + ctx := context.Background() + engine := distribution.NewEngine() + engine.UpdateRoute([]byte(""), []byte("m"), 1) + engine.UpdateRoute([]byte("m"), nil, 1) + + groups := map[uint64]*ShardGroup{ + 1: {Store: store.NewMVCCStore()}, + } + st := NewShardStore(engine, groups) + + require.NoError(t, st.PutAt(ctx, []byte("a"), []byte("va"), 1, 0)) + require.NoError(t, st.PutAt(ctx, []byte("z"), []byte("vz"), 2, 0)) + + keys, err := st.ScanKeysAt(ctx, []byte(""), nil, 10, ^uint64(0)) + require.NoError(t, err) + require.Equal(t, [][]byte{[]byte("a"), []byte("z")}, keys) +} + +func TestKeyOnlyScanHelpersPreserveEmptyKey(t *testing.T) { + t.Parallel() + + keys := keysFromKVs(kvPairsFromKeys([][]byte{nil, []byte(""), []byte("a")})) + require.Equal(t, [][]byte{[]byte(""), []byte("a")}, keys) +} + +func TestMergeAndTrimScanKeysSortsBeforeTruncating(t *testing.T) { + t.Parallel() + + keys := mergeAndTrimScanKeys(nil, [][]byte{[]byte("c"), []byte("d"), []byte("b")}, 2) + require.Equal(t, [][]byte{[]byte("b"), []byte("c")}, keys) +} + func TestBackupScannerPaging(t *testing.T) { t.Parallel() @@ -203,7 +238,7 @@ func TestBackupScannerPaging(t *testing.T) { require.Equal(t, [][]byte{[]byte("a"), []byte("b"), []byte("c"), []byte("x"), []byte("z")}, got) } -func TestBackupScannerEmptyKeyTerminatesAfterPage(t *testing.T) { +func TestBackupScannerEmptyKeyContinuesAfterPage(t *testing.T) { t.Parallel() ctx := context.Background() @@ -214,8 +249,9 @@ func TestBackupScannerEmptyKeyTerminatesAfterPage(t *testing.T) { } st := NewShardStore(engine, groups) require.NoError(t, st.PutAt(ctx, []byte(""), []byte("empty"), 1, 0)) + require.NoError(t, st.PutAt(ctx, []byte("a"), []byte("later"), 2, 0)) - sc := st.NewBackupScanner(nil, []byte("\x00"), ^uint64(0), 1) + sc := st.NewBackupScanner(nil, nil, ^uint64(0), 1) defer sc.Close() kvp, ok, err := sc.Next(ctx) @@ -224,6 +260,12 @@ func TestBackupScannerEmptyKeyTerminatesAfterPage(t *testing.T) { require.Equal(t, []byte(""), kvp.Key) require.Equal(t, []byte("empty"), kvp.Value) + kvp, ok, err = sc.Next(ctx) + require.NoError(t, err) + require.True(t, ok) + require.Equal(t, []byte("a"), kvp.Key) + require.Equal(t, []byte("later"), kvp.Value) + kvp, ok, err = sc.Next(ctx) require.NoError(t, err) require.False(t, ok) From 582b2eba2ef18e894626c3701e6a0032368557e9 Mon Sep 17 00:00:00 2001 From: bootjp Date: Sat, 11 Jul 2026 14:39:10 +0900 Subject: [PATCH 04/10] backup: harden scan pagination edge cases --- kv/leader_routed_store_test.go | 5 +- kv/shard_store.go | 158 ++++++++++++++++++++++++++------- kv/shard_store_test.go | 141 +++++++++++++++++++++++++++++ store/lsm_store.go | 23 ++++- store/scan_keys_test.go | 37 ++++++++ 5 files changed, 328 insertions(+), 36 deletions(-) diff --git a/kv/leader_routed_store_test.go b/kv/leader_routed_store_test.go index 1d83ead89..c6abd562a 100644 --- a/kv/leader_routed_store_test.go +++ b/kv/leader_routed_store_test.go @@ -85,6 +85,8 @@ type fakeRawKVServer struct { scanCalls int latestCalls int + lastScanGroupID uint64 + getResp *pb.RawGetResponse scanResp *pb.RawScanAtResponse latestResp *pb.RawLatestCommitTSResponse @@ -100,10 +102,11 @@ func (f *fakeRawKVServer) RawGet(context.Context, *pb.RawGetRequest) (*pb.RawGet return &pb.RawGetResponse{}, nil } -func (f *fakeRawKVServer) RawScanAt(context.Context, *pb.RawScanAtRequest) (*pb.RawScanAtResponse, error) { +func (f *fakeRawKVServer) RawScanAt(_ context.Context, req *pb.RawScanAtRequest) (*pb.RawScanAtResponse, error) { f.mu.Lock() defer f.mu.Unlock() f.scanCalls++ + f.lastScanGroupID = req.GetGroupId() if f.scanResp != nil { return f.scanResp, nil } diff --git a/kv/shard_store.go b/kv/shard_store.go index 7b907d1d2..21f4479e2 100644 --- a/kv/shard_store.go +++ b/kv/shard_store.go @@ -400,7 +400,7 @@ func (s *ShardStore) scanKeyRouteAt( return s.scanKeysRouteAtLeader(ctx, g, start, end, limit, ts) } - return s.proxyScanKeysAt(ctx, g, start, end, limit, ts) + return s.proxyScanKeysAt(ctx, g, start, end, limit, ts, route.GroupID) } func (s *ShardStore) scanKeysRouteLocal( @@ -411,11 +411,13 @@ func (s *ShardStore) scanKeysRouteLocal( limit int, ts uint64, ) ([][]byte, error) { - keys, err := g.Store.ScanKeysAt(ctx, start, end, limit, ts) - if err != nil { - return nil, errors.WithStack(err) - } - return filterTxnInternalKeys(keys), nil + return scanKeysWithRefill(start, end, limit, func(cursor []byte, pageLimit int) ([][]byte, error) { + keys, err := g.Store.ScanKeysAt(ctx, cursor, end, pageLimit, ts) + if err != nil { + return nil, errors.WithStack(err) + } + return keys, nil + }) } func (s *ShardStore) scanKeysRouteAtLeader( @@ -426,21 +428,40 @@ func (s *ShardStore) scanKeysRouteAtLeader( limit int, ts uint64, ) ([][]byte, error) { - keys, err := g.Store.ScanKeysAt(ctx, start, end, limit, ts) - if err != nil { - return nil, errors.WithStack(err) - } - keyKVs := kvPairsFromKeys(keys) - lockStart, lockEnd := scanLockBoundsForKVs(keyKVs, start, end, limit) - lockKVs, err := scanTxnLockRangeAt(ctx, g, lockStart, lockEnd, ts, limit) - if err != nil { - return nil, err + if limit <= 0 { + return [][]byte{}, nil } - kvs, err := s.resolveScanLocks(ctx, g, keyKVs, lockKVs, ts) - if err != nil { - return nil, err + + out := make([][]byte, 0, limit) + cursor := start + for len(out) < limit { + keys, err := g.Store.ScanKeysAt(ctx, cursor, end, limit, ts) + if err != nil { + return nil, errors.WithStack(err) + } + if len(keys) == 0 { + break + } + + keyKVs := kvPairsFromKeys(keys) + lockStart, lockEnd := scanLockBoundsForKVs(keyKVs, cursor, end, limit) + lockKVs, err := scanTxnLockRangeAt(ctx, g, lockStart, lockEnd, ts, limit) + if err != nil { + return nil, err + } + kvs, err := s.resolveScanLocks(ctx, g, keyKVs, lockKVs, ts) + if err != nil { + return nil, err + } + out = mergeAndTrimScanKeys(out, filterTxnInternalKeys(keysFromKVs(kvs)), limit) + + nextCursor, ok := nextKeyScanCursor(keys, end, limit) + if !ok { + break + } + cursor = nextCursor } - return keysFromKVs(kvs), nil + return out, nil } func (s *ShardStore) proxyScanKeysAt( @@ -450,12 +471,68 @@ func (s *ShardStore) proxyScanKeysAt( end []byte, limit int, ts uint64, + groupID uint64, ) ([][]byte, error) { - kvs, err := s.proxyRawScanAt(ctx, g, start, end, limit, ts, false, 0) - if err != nil { - return nil, err + return scanKeysWithRefill(start, end, limit, func(cursor []byte, pageLimit int) ([][]byte, error) { + kvs, err := s.proxyRawScanAt(ctx, g, cursor, end, pageLimit, ts, false, groupID) + if err != nil { + return nil, err + } + return keysFromKVs(kvs), nil + }) +} + +func scanKeysWithRefill( + start []byte, + end []byte, + limit int, + scan func(cursor []byte, pageLimit int) ([][]byte, error), +) ([][]byte, error) { + if limit <= 0 { + return [][]byte{}, nil + } + + out := make([][]byte, 0, limit) + cursor := start + for len(out) < limit { + keys, err := scan(cursor, limit) + if err != nil { + return nil, err + } + if len(keys) == 0 { + break + } + + out = mergeAndTrimScanKeys(out, filterTxnInternalKeys(keys), limit) + + nextCursor, ok := nextKeyScanCursor(keys, end, limit) + if !ok { + break + } + cursor = nextCursor } - return keysFromKVs(filterTxnInternalKVs(kvs)), nil + return out, nil +} + +func nextKeyScanCursor(keys [][]byte, end []byte, limit int) ([]byte, bool) { + lastKey := lastScanKey(keys) + if len(keys) < limit || lastKey == nil { + return nil, false + } + nextCursor := nextScanCursor(lastKey) + if end != nil && bytes.Compare(nextCursor, end) >= 0 { + return nil, false + } + return nextCursor, true +} + +func lastScanKey(keys [][]byte) []byte { + for i := len(keys) - 1; i >= 0; i-- { + if keys[i] != nil { + return keys[i] + } + } + return nil } func (s *ShardStore) clampedReverseScanRouteAt( @@ -584,26 +661,28 @@ func scanLockBoundsForKVs(kvs []*store.KVPair, scanStart []byte, scanEnd []byte, func observedScanUserBounds(kvs []*store.KVPair) ([]byte, []byte, bool) { var minKey []byte var maxKey []byte + seen := false for _, kvp := range kvs { userKey, ok := scanUserKey(kvp) if !ok { continue } - if minKey == nil || bytes.Compare(userKey, minKey) < 0 { + if !seen || bytes.Compare(userKey, minKey) < 0 { minKey = userKey } - if maxKey == nil || bytes.Compare(userKey, maxKey) > 0 { + if !seen || bytes.Compare(userKey, maxKey) > 0 { maxKey = userKey } + seen = true } - if len(minKey) == 0 || len(maxKey) == 0 { + if !seen { return nil, nil, false } return minKey, maxKey, true } func scanUserKey(kvp *store.KVPair) ([]byte, bool) { - if kvp == nil || len(kvp.Key) == 0 { + if kvp == nil || kvp.Key == nil { return nil, false } if !isTxnInternalKey(kvp.Key) { @@ -632,12 +711,25 @@ func mergeAndTrimScanKeys(out [][]byte, keys [][]byte, limit int) [][]byte { return out } out = append(out, keys...) - if len(out) <= limit { - return out - } sort.Slice(out, func(i, j int) bool { return bytes.Compare(out[i], out[j]) < 0 }) + write := 0 + for _, key := range out { + if key == nil { + continue + } + if write > 0 && bytes.Equal(out[write-1], key) { + continue + } + out[write] = key + write++ + } + clear(out[write:]) + out = out[:write] + if len(out) <= limit { + return out + } clear(out[limit:]) return out[:limit] } @@ -660,7 +752,7 @@ func mergeAndTrimReverseScanResults(out []*store.KVPair, kvs []*store.KVPair, li func countNonInternalKVs(kvs []*store.KVPair) int { count := 0 for _, kvp := range kvs { - if kvp == nil || isTxnInternalKey(kvp.Key) { + if kvp == nil || kvp.Key == nil || isTxnInternalKey(kvp.Key) { continue } count++ @@ -696,7 +788,7 @@ func filterTxnInternalKeys(keys [][]byte) [][]byte { } out := make([][]byte, 0, len(keys)) for _, key := range keys { - if isTxnInternalKey(key) { + if key == nil || isTxnInternalKey(key) { continue } out = append(out, key) @@ -1064,7 +1156,7 @@ func appendScanLockResolutionBatch(plan *scanLockPlan, txnKey lockTxnKey, phase } func appendScanItem(plan *scanLockPlan, kvp *store.KVPair, locked bool) { - if kvp == nil || len(kvp.Key) == 0 { + if kvp == nil || kvp.Key == nil { return } keyID := string(kvp.Key) diff --git a/kv/shard_store_test.go b/kv/shard_store_test.go index cc3669bda..51c5294bb 100644 --- a/kv/shard_store_test.go +++ b/kv/shard_store_test.go @@ -5,11 +5,56 @@ import ( "testing" "github.com/bootjp/elastickv/distribution" + "github.com/bootjp/elastickv/internal/raftengine" "github.com/bootjp/elastickv/internal/s3keys" + pb "github.com/bootjp/elastickv/proto" "github.com/bootjp/elastickv/store" "github.com/stretchr/testify/require" ) +type followerProxyEngine struct { + leader string +} + +func (e *followerProxyEngine) Propose(context.Context, []byte) (*raftengine.ProposalResult, error) { + return nil, ErrLeaderNotFound +} + +func (e *followerProxyEngine) ProposeAdmin(context.Context, []byte) (*raftengine.ProposalResult, error) { + return nil, ErrLeaderNotFound +} + +func (e *followerProxyEngine) State() raftengine.State { + return raftengine.StateFollower +} + +func (e *followerProxyEngine) Leader() raftengine.LeaderInfo { + return raftengine.LeaderInfo{Address: e.leader} +} + +func (e *followerProxyEngine) VerifyLeader(context.Context) error { + return ErrLeaderNotFound +} + +func (e *followerProxyEngine) LinearizableRead(context.Context) (uint64, error) { + return 0, ErrLeaderNotFound +} + +func (e *followerProxyEngine) Status() raftengine.Status { + return raftengine.Status{ + State: raftengine.StateFollower, + Leader: raftengine.LeaderInfo{Address: e.leader}, + } +} + +func (e *followerProxyEngine) Configuration(context.Context) (raftengine.Configuration, error) { + return raftengine.Configuration{}, nil +} + +func (e *followerProxyEngine) Close() error { + return nil +} + func TestShardStoreScanAt_IncludesListKeysAcrossShards(t *testing.T) { t.Parallel() @@ -190,6 +235,66 @@ func TestShardStoreScanKeysAt_DeduplicatesUnclampedGroups(t *testing.T) { require.Equal(t, [][]byte{[]byte("a"), []byte("z")}, keys) } +func TestShardStoreScanKeysRouteAtLeaderRefillsAfterTxnInternalKeys(t *testing.T) { + t.Parallel() + + ctx := context.Background() + g := &ShardGroup{Store: store.NewMVCCStore()} + st := NewShardStore(distribution.NewEngine(), map[uint64]*ShardGroup{1: g}) + + require.NoError(t, g.Store.PutAt(ctx, txnCommitKey([]byte("primary"), 10), []byte("commit"), 1, 0)) + require.NoError(t, g.Store.PutAt(ctx, []byte("a"), []byte("va"), 2, 0)) + + keys, err := st.scanKeysRouteAtLeader(ctx, g, []byte(""), nil, 1, ^uint64(0)) + require.NoError(t, err) + require.Equal(t, [][]byte{[]byte("a")}, keys) +} + +func TestShardStoreScanKeysRouteAtLeaderPreservesEmptyKey(t *testing.T) { + t.Parallel() + + ctx := context.Background() + g := &ShardGroup{Store: store.NewMVCCStore()} + st := NewShardStore(distribution.NewEngine(), map[uint64]*ShardGroup{1: g}) + + require.NoError(t, g.Store.PutAt(ctx, []byte(""), []byte("empty"), 1, 0)) + require.NoError(t, g.Store.PutAt(ctx, []byte("a"), []byte("va"), 2, 0)) + + keys, err := st.scanKeysRouteAtLeader(ctx, g, nil, nil, 2, ^uint64(0)) + require.NoError(t, err) + require.Equal(t, [][]byte{[]byte(""), []byte("a")}, keys) +} + +func TestShardStoreProxyScanKeysAtUsesSelectedGroup(t *testing.T) { + t.Parallel() + + fake := &fakeRawKVServer{ + scanResp: &pb.RawScanAtResponse{ + Kv: []*pb.RawKVPair{{Key: []byte("k"), Value: []byte("v")}}, + }, + } + addr, stop := startRawKVServer(t, fake) + t.Cleanup(stop) + + ctx := context.Background() + engine := distribution.NewEngine() + engine.UpdateRoute([]byte(""), nil, 42) + g := &ShardGroup{ + Engine: &followerProxyEngine{leader: addr}, + Store: store.NewMVCCStore(), + } + st := NewShardStore(engine, map[uint64]*ShardGroup{42: g}) + + keys, err := st.scanKeyRouteAt(ctx, distribution.Route{GroupID: 42}, []byte(""), nil, 10, ^uint64(0)) + require.NoError(t, err) + require.Equal(t, [][]byte{[]byte("k")}, keys) + + fake.mu.Lock() + defer fake.mu.Unlock() + require.Equal(t, 1, fake.scanCalls) + require.Equal(t, uint64(42), fake.lastScanGroupID) +} + func TestKeyOnlyScanHelpersPreserveEmptyKey(t *testing.T) { t.Parallel() @@ -272,6 +377,42 @@ func TestBackupScannerEmptyKeyContinuesAfterPage(t *testing.T) { require.Nil(t, kvp) } +func TestBackupScannerPebblePagingDoesNotRepeatLastKey(t *testing.T) { + t.Parallel() + + ctx := context.Background() + engine := distribution.NewEngine() + engine.UpdateRoute([]byte(""), nil, 1) + pebbleStore, err := store.NewPebbleStore(t.TempDir()) + require.NoError(t, err) + t.Cleanup(func() { require.NoError(t, pebbleStore.Close()) }) + + groups := map[uint64]*ShardGroup{ + 1: {Store: pebbleStore}, + } + st := NewShardStore(engine, groups) + require.NoError(t, pebbleStore.PutAt(ctx, []byte("a"), []byte("va"), 1, 0)) + require.NoError(t, pebbleStore.PutAt(ctx, []byte("b"), []byte("vb"), 2, 0)) + + sc := st.NewBackupScanner(nil, nil, ^uint64(0), 1) + defer sc.Close() + + kvp, ok, err := sc.Next(ctx) + require.NoError(t, err) + require.True(t, ok) + require.Equal(t, []byte("a"), kvp.Key) + + kvp, ok, err = sc.Next(ctx) + require.NoError(t, err) + require.True(t, ok) + require.Equal(t, []byte("b"), kvp.Key) + + kvp, ok, err = sc.Next(ctx) + require.NoError(t, err) + require.False(t, ok) + require.Nil(t, kvp) +} + func TestShardStoreScanAt_RoutesS3ManifestScansByLogicalObjectKey(t *testing.T) { t.Parallel() diff --git a/store/lsm_store.go b/store/lsm_store.go index 5d85b2699..d31c73b29 100644 --- a/store/lsm_store.go +++ b/store/lsm_store.go @@ -1022,6 +1022,10 @@ func pastScanEnd(userKey, end []byte) bool { return end != nil && bytes.Compare(userKey, end) >= 0 } +func beforeScanStart(userKey, start []byte) bool { + return start != nil && bytes.Compare(userKey, start) < 0 +} + func nextScannableUserKey(iter *pebble.Iterator) ([]byte, uint64, bool) { for iter.Valid() { rawKey := iter.Key() @@ -1043,6 +1047,21 @@ func nextScannableUserKey(iter *pebble.Iterator) ([]byte, uint64, bool) { return nil, 0, false } +func (s *pebbleStore) nextForwardScanUserKey(iter *pebble.Iterator, start []byte) ([]byte, uint64, bool) { + for { + userKey, version, ok := nextScannableUserKey(iter) + if !ok { + return nil, 0, false + } + if !beforeScanStart(userKey, start) { + return userKey, version, true + } + if !s.skipToNextUserKey(iter, userKey) { + return nil, 0, false + } + } +} + func prevScannableUserKey(iter *pebble.Iterator) ([]byte, bool) { for iter.Valid() { rawKey := iter.Key() @@ -1068,7 +1087,7 @@ func (s *pebbleStore) collectScanResults(iter *pebble.Iterator, start, end []byt result := make([]*KVPair, 0, boundedScanResultCapacity(limit)) for iter.SeekGE(encodeKey(start, math.MaxUint64)); iter.Valid() && len(result) < limit; { - userKey, version, ok := nextScannableUserKey(iter) + userKey, version, ok := s.nextForwardScanUserKey(iter, start) if !ok { break } @@ -1101,7 +1120,7 @@ func (s *pebbleStore) collectScanKeys(iter *pebble.Iterator, start, end []byte, result := make([][]byte, 0, boundedScanResultCapacity(limit)) for iter.SeekGE(encodeKey(start, math.MaxUint64)); iter.Valid() && len(result) < limit; { - userKey, version, ok := nextScannableUserKey(iter) + userKey, version, ok := s.nextForwardScanUserKey(iter, start) if !ok { break } diff --git a/store/scan_keys_test.go b/store/scan_keys_test.go index e32babf8e..527e77581 100644 --- a/store/scan_keys_test.go +++ b/store/scan_keys_test.go @@ -69,6 +69,43 @@ func TestPebbleStoreScanKeysAtMatchesScanAtVisibility(t *testing.T) { require.Equal(t, [][]byte{[]byte("a")}, keys) } +func TestPebbleStoreScanAtCursorAfterKeyDoesNotRepeat(t *testing.T) { + t.Parallel() + ctx := context.Background() + st, err := NewPebbleStore(t.TempDir()) + require.NoError(t, err) + defer st.Close() + + require.NoError(t, st.PutAt(ctx, []byte("a"), []byte("va"), 10, 0)) + require.NoError(t, st.PutAt(ctx, []byte("b"), []byte("vb"), 20, 0)) + + kvs, err := st.ScanAt(ctx, cursorAfterTestKey([]byte("a")), nil, 1, 20) + require.NoError(t, err) + require.Len(t, kvs, 1) + require.Equal(t, []byte("b"), kvs[0].Key) +} + +func TestPebbleStoreScanKeysAtCursorAfterKeyDoesNotRepeat(t *testing.T) { + t.Parallel() + ctx := context.Background() + st, err := NewPebbleStore(t.TempDir()) + require.NoError(t, err) + defer st.Close() + + require.NoError(t, st.PutAt(ctx, []byte(""), []byte("empty"), 10, 0)) + require.NoError(t, st.PutAt(ctx, []byte("a"), []byte("va"), 20, 0)) + + keys, err := st.ScanKeysAt(ctx, cursorAfterTestKey([]byte("")), nil, 1, 20) + require.NoError(t, err) + require.Equal(t, [][]byte{[]byte("a")}, keys) +} + +func cursorAfterTestKey(key []byte) []byte { + cursor := make([]byte, len(key)+1) + copy(cursor, key) + return cursor +} + func keysFromStoreKVs(kvs []*KVPair) [][]byte { keys := make([][]byte, 0, len(kvs)) for _, kvp := range kvs { From aa99e90709cb1f100b92ff43af043e61cf0fb0a9 Mon Sep 17 00:00:00 2001 From: bootjp Date: Sat, 11 Jul 2026 16:47:35 +0900 Subject: [PATCH 05/10] backup: preserve scan progress across filtered pages --- kv/backup_scan.go | 41 +++++--- kv/shard_store.go | 169 +++++++++++++++++++++++++++++--- kv/shard_store_test.go | 92 +++++++++++++++++ kv/shard_store_txn_lock_test.go | 62 ++++++++++++ 4 files changed, 337 insertions(+), 27 deletions(-) diff --git a/kv/backup_scan.go b/kv/backup_scan.go index 66eda9d83..9f8d514f5 100644 --- a/kv/backup_scan.go +++ b/kv/backup_scan.go @@ -4,6 +4,7 @@ import ( "bytes" "context" + "github.com/bootjp/elastickv/distribution" "github.com/bootjp/elastickv/store" ) @@ -17,27 +18,37 @@ type BackupScanner interface { } type backupScanner struct { - store *ShardStore - end []byte - ts uint64 - pageSize int - cursor []byte - page []*store.KVPair - index int - closed bool - exhausted bool + store *ShardStore + routes []distribution.Route + clampToRoutes bool + end []byte + ts uint64 + pageSize int + cursor []byte + page []*store.KVPair + index int + closed bool + exhausted bool } func NewBackupScanner(st *ShardStore, start []byte, end []byte, ts uint64, pageSize int) BackupScanner { if pageSize <= 0 { pageSize = defaultBackupScanPageSize } + var routes []distribution.Route + var clampToRoutes bool + if st != nil { + routes, clampToRoutes = st.routesForScan(start, end) + routes = append([]distribution.Route(nil), routes...) + } return &backupScanner{ - store: st, - cursor: bytes.Clone(start), - end: bytes.Clone(end), - ts: ts, - pageSize: pageSize, + store: st, + routes: routes, + clampToRoutes: clampToRoutes, + cursor: bytes.Clone(start), + end: bytes.Clone(end), + ts: ts, + pageSize: pageSize, } } @@ -78,7 +89,7 @@ func (s *backupScanner) loadNextPage(ctx context.Context) error { s.index = 0 return nil } - page, err := s.store.ScanAt(ctx, s.cursor, s.end, s.pageSize, s.ts) + page, err := s.store.scanRoutesAtSorted(ctx, s.routes, s.cursor, s.end, s.pageSize, s.ts, s.clampToRoutes) if err != nil { return err } diff --git a/kv/shard_store.go b/kv/shard_store.go index 21f4479e2..a09110b5d 100644 --- a/kv/shard_store.go +++ b/kv/shard_store.go @@ -190,18 +190,7 @@ func (s *ShardStore) ScanAt(ctx context.Context, start []byte, end []byte, limit } routes, clampToRoutes := s.routesForScan(start, end) - out, err := s.scanRoutesAt(ctx, routes, start, end, limit, ts, clampToRoutes) - if err != nil { - return nil, err - } - - sort.Slice(out, func(i, j int) bool { - return bytes.Compare(out[i].Key, out[j].Key) < 0 - }) - if len(out) > limit { - out = out[:limit] - } - return out, nil + return s.scanRoutesAtSorted(ctx, routes, start, end, limit, ts, clampToRoutes) } func (s *ShardStore) ScanKeysAt(ctx context.Context, start []byte, end []byte, limit int, ts uint64) ([][]byte, error) { @@ -303,6 +292,20 @@ func (s *ShardStore) scanRoutesAt(ctx context.Context, routes []distribution.Rou return out, nil } +func (s *ShardStore) scanRoutesAtSorted(ctx context.Context, routes []distribution.Route, start []byte, end []byte, limit int, ts uint64, clampToRoutes bool) ([]*store.KVPair, error) { + out, err := s.scanRoutesAt(ctx, routes, start, end, limit, ts, clampToRoutes) + if err != nil { + return nil, err + } + sort.Slice(out, func(i, j int) bool { + return bytes.Compare(out[i].Key, out[j].Key) < 0 + }) + if len(out) > limit { + out = out[:limit] + } + return out, nil +} + func (s *ShardStore) scanKeyRoutesAt(ctx context.Context, routes []distribution.Route, start []byte, end []byte, limit int, ts uint64, clampToRoutes bool) ([][]byte, error) { out := make([][]byte, 0) seenGroups := make(map[uint64]struct{}) @@ -440,6 +443,11 @@ func (s *ShardStore) scanKeysRouteAtLeader( return nil, errors.WithStack(err) } if len(keys) == 0 { + keys, err := s.scanLockOnlyKeysAtLeader(ctx, g, cursor, end, ts, limit) + if err != nil { + return nil, err + } + out = mergeAndTrimScanKeys(out, keys, limit) break } @@ -464,6 +472,28 @@ func (s *ShardStore) scanKeysRouteAtLeader( return out, nil } +func (s *ShardStore) scanLockOnlyKeysAtLeader( + ctx context.Context, + g *ShardGroup, + start []byte, + end []byte, + ts uint64, + limit int, +) ([][]byte, error) { + lockKVs, err := scanTxnLockRangeAt(ctx, g, start, end, ts, limit) + if err != nil { + return nil, err + } + if len(lockKVs) == 0 { + return nil, nil + } + kvs, err := s.resolveScanLocks(ctx, g, nil, lockKVs, ts) + if err != nil { + return nil, err + } + return filterTxnInternalKeys(keysFromKVs(kvs)), nil +} + func (s *ShardStore) proxyScanKeysAt( ctx context.Context, g *ShardGroup, @@ -572,6 +602,10 @@ func (s *ShardStore) scanRouteAtDirection( return nil, nil } + if !reverse { + return s.scanRouteAtForward(ctx, route, g, start, end, limit, ts) + } + if engineForGroup(g) == nil { kvs, err := s.scanRouteLocal(ctx, g, start, end, limit, ts, reverse) if err != nil { @@ -597,6 +631,107 @@ func (s *ShardStore) scanRouteAtDirection( return filterTxnInternalKVs(kvs), nil } +type scanRoutePage struct { + kvs []*store.KVPair + advanceKey []byte + full bool +} + +func (s *ShardStore) scanRouteAtForward( + ctx context.Context, + route distribution.Route, + g *ShardGroup, + start []byte, + end []byte, + limit int, + ts uint64, +) ([]*store.KVPair, error) { + if limit <= 0 { + return []*store.KVPair{}, nil + } + + out := make([]*store.KVPair, 0, limit) + cursor := start + for len(out) < limit { + page, err := s.scanRouteAtForwardPage(ctx, route, g, cursor, end, limit, ts) + if err != nil { + return nil, err + } + out = mergeAndTrimScanResults(out, page.kvs, limit) + if !page.full || page.advanceKey == nil { + break + } + nextCursor := nextScanCursor(page.advanceKey) + if end != nil && bytes.Compare(nextCursor, end) >= 0 { + break + } + cursor = nextCursor + } + + sort.Slice(out, func(i, j int) bool { + return bytes.Compare(out[i].Key, out[j].Key) < 0 + }) + if len(out) > limit { + out = out[:limit] + } + return out, nil +} + +func (s *ShardStore) scanRouteAtForwardPage( + ctx context.Context, + route distribution.Route, + g *ShardGroup, + start []byte, + end []byte, + limit int, + ts uint64, +) (scanRoutePage, error) { + engine := engineForGroup(g) + if engine == nil { + raw, err := s.scanRouteLocal(ctx, g, start, end, limit, ts, false) + if err != nil { + return scanRoutePage{}, errors.WithStack(err) + } + return scanRoutePage{ + kvs: filterTxnInternalKVs(raw), + advanceKey: lastKVKey(raw), + full: len(raw) >= limit, + }, nil + } + + if isLinearizableRaftLeader(ctx, engine) { + raw, err := g.Store.ScanAt(ctx, start, end, limit, ts) + if err != nil { + return scanRoutePage{}, errors.WithStack(err) + } + lockStart, lockEnd := scanLockBoundsForKVs(raw, start, end, limit) + lockKVs, err := scanTxnLockRangeAt(ctx, g, lockStart, lockEnd, ts, limit) + if err != nil { + return scanRoutePage{}, err + } + kvs, err := s.resolveScanLocks(ctx, g, raw, lockKVs, ts) + if err != nil { + return scanRoutePage{}, err + } + return scanRoutePage{ + kvs: filterTxnInternalKVs(kvs), + advanceKey: lastKVKey(raw), + full: len(raw) >= limit, + }, nil + } + + kvs, err := s.proxyRawScanAt(ctx, g, start, end, limit, ts, false, route.GroupID) + if err != nil { + return scanRoutePage{}, err + } + kvs = filterTxnInternalKVs(kvs) + return scanRoutePage{ + kvs: kvs, + advanceKey: lastKVKey(kvs), + full: len(kvs) >= limit, + }, nil +} + func (s *ShardStore) scanRouteLocal( ctx context.Context, g *ShardGroup, @@ -782,6 +917,16 @@ func keysFromKVs(kvs []*store.KVPair) [][]byte { return keys } +func lastKVKey(kvs []*store.KVPair) []byte { + for i := len(kvs) - 1; i >= 0; i-- { + if kvs[i] == nil || kvs[i].Key == nil { + continue + } + return kvs[i].Key + } + return nil +} + func filterTxnInternalKeys(keys [][]byte) [][]byte { if len(keys) == 0 { return keys diff --git a/kv/shard_store_test.go b/kv/shard_store_test.go index 51c5294bb..5a7eabd6d 100644 --- a/kv/shard_store_test.go +++ b/kv/shard_store_test.go @@ -295,6 +295,37 @@ func TestShardStoreProxyScanKeysAtUsesSelectedGroup(t *testing.T) { require.Equal(t, uint64(42), fake.lastScanGroupID) } +func TestShardStoreProxyScanAtUsesSelectedGroup(t *testing.T) { + t.Parallel() + + fake := &fakeRawKVServer{ + scanResp: &pb.RawScanAtResponse{ + Kv: []*pb.RawKVPair{{Key: []byte("k"), Value: []byte("v")}}, + }, + } + addr, stop := startRawKVServer(t, fake) + t.Cleanup(stop) + + ctx := context.Background() + engine := distribution.NewEngine() + engine.UpdateRoute([]byte(""), nil, 42) + g := &ShardGroup{ + Engine: &followerProxyEngine{leader: addr}, + Store: store.NewMVCCStore(), + } + st := NewShardStore(engine, map[uint64]*ShardGroup{42: g}) + + kvs, err := st.ScanAt(ctx, []byte(""), nil, 10, ^uint64(0)) + require.NoError(t, err) + require.Len(t, kvs, 1) + require.Equal(t, []byte("k"), kvs[0].Key) + + fake.mu.Lock() + defer fake.mu.Unlock() + require.Equal(t, 1, fake.scanCalls) + require.Equal(t, uint64(42), fake.lastScanGroupID) +} + func TestKeyOnlyScanHelpersPreserveEmptyKey(t *testing.T) { t.Parallel() @@ -343,6 +374,67 @@ func TestBackupScannerPaging(t *testing.T) { require.Equal(t, [][]byte{[]byte("a"), []byte("b"), []byte("c"), []byte("x"), []byte("z")}, got) } +func TestBackupScannerPreservesFullRoutingAfterListKeyCursor(t *testing.T) { + t.Parallel() + + ctx := context.Background() + engine := distribution.NewEngine() + engine.UpdateRoute([]byte(""), []byte("m"), 1) + engine.UpdateRoute([]byte("m"), nil, 2) + + groups := map[uint64]*ShardGroup{ + 1: {Store: store.NewMVCCStore()}, + 2: {Store: store.NewMVCCStore()}, + } + st := NewShardStore(engine, groups) + + listKey := store.ListItemKey([]byte("x"), 0) + require.NoError(t, st.PutAt(ctx, listKey, []byte("list"), 1, 0)) + require.NoError(t, st.PutAt(ctx, []byte("a"), []byte("plain"), 2, 0)) + + sc := st.NewBackupScanner([]byte(""), nil, ^uint64(0), 1) + defer sc.Close() + + var got [][]byte + for { + kvp, ok, err := sc.Next(ctx) + require.NoError(t, err) + if !ok { + break + } + got = append(got, kvp.Key) + } + require.Equal(t, [][]byte{listKey, []byte("a")}, got) +} + +func TestBackupScannerContinuesPastTxnInternalOnlyPage(t *testing.T) { + t.Parallel() + + ctx := context.Background() + engine := distribution.NewEngine() + engine.UpdateRoute([]byte(""), nil, 1) + groups := map[uint64]*ShardGroup{ + 1: {Store: store.NewMVCCStore()}, + } + st := NewShardStore(engine, groups) + + require.NoError(t, groups[1].Store.PutAt(ctx, txnCommitKey([]byte("primary"), 10), []byte("commit"), 1, 0)) + require.NoError(t, groups[1].Store.PutAt(ctx, []byte("a"), []byte("visible"), 2, 0)) + + sc := st.NewBackupScanner([]byte(""), nil, ^uint64(0), 1) + defer sc.Close() + + kvp, ok, err := sc.Next(ctx) + require.NoError(t, err) + require.True(t, ok) + require.Equal(t, []byte("a"), kvp.Key) + + kvp, ok, err = sc.Next(ctx) + require.NoError(t, err) + require.False(t, ok) + require.Nil(t, kvp) +} + func TestBackupScannerEmptyKeyContinuesAfterPage(t *testing.T) { t.Parallel() diff --git a/kv/shard_store_txn_lock_test.go b/kv/shard_store_txn_lock_test.go index 8c50ff3a4..71a793afc 100644 --- a/kv/shard_store_txn_lock_test.go +++ b/kv/shard_store_txn_lock_test.go @@ -281,6 +281,68 @@ func TestShardStoreScanAt_ReturnsTxnLockedForPendingLockWithoutCommittedValue(t require.True(t, errors.Is(err, ErrTxnLocked), "expected ErrTxnLocked, got %v", err) } +func TestShardStoreScanKeysAt_ReturnsTxnLockedForPendingLockWithoutCommittedValue(t *testing.T) { + t.Parallel() + + ctx := context.Background() + + engine := distribution.NewEngine() + engine.UpdateRoute([]byte(""), nil, 1) + + st1 := store.NewMVCCStore() + r1, stop1 := newSingleRaft(t, "g1", NewKvFSMWithHLC(st1, NewHLC())) + defer stop1() + + groups := map[uint64]*ShardGroup{ + 1: {Engine: r1, Store: st1, Txn: NewLeaderProxyWithEngine(r1)}, + } + shardStore := NewShardStore(engine, groups) + + key := []byte("k") + startTS := uint64(1) + _, err := groups[1].Txn.Commit(context.Background(), []*pb.Request{makePrepareRequest(startTS, key, []byte("v"), key)}) + require.NoError(t, err) + + _, err = shardStore.ScanKeysAt(ctx, []byte("k"), []byte("l"), 100, ^uint64(0)) + require.Error(t, err) + require.True(t, errors.Is(err, ErrTxnLocked), "expected ErrTxnLocked, got %v", err) +} + +func TestShardStoreScanKeysAt_ResolvesCommittedSecondaryLockWithoutCommittedValue(t *testing.T) { + t.Parallel() + + ctx := context.Background() + + shardStore, groups, cleanup := setupTwoShardStore(t) + defer cleanup() + + startTS := uint64(1) + commitTS := uint64(2) + primaryKey := []byte("b") + secondaryKey := []byte("x") + + _, err := groups[1].Txn.Commit(context.Background(), []*pb.Request{makePrepareRequest(startTS, primaryKey, []byte("v1"), primaryKey)}) + require.NoError(t, err) + _, err = groups[2].Txn.Commit(context.Background(), []*pb.Request{makePrepareRequest(startTS, secondaryKey, []byte("v2"), primaryKey)}) + require.NoError(t, err) + + commitPrimary := &pb.Request{ + IsTxn: true, + Phase: pb.Phase_COMMIT, + Ts: startTS, + Mutations: []*pb.Mutation{ + {Op: pb.Op_PUT, Key: []byte(txnMetaPrefix), Value: EncodeTxnMeta(TxnMeta{PrimaryKey: primaryKey, LockTTLms: 0, CommitTS: commitTS})}, + {Op: pb.Op_PUT, Key: primaryKey}, + }, + } + _, err = groups[1].Txn.Commit(context.Background(), []*pb.Request{commitPrimary}) + require.NoError(t, err) + + keys, err := shardStore.ScanKeysAt(ctx, []byte("x"), []byte("z"), 100, commitTS) + require.NoError(t, err) + require.Equal(t, [][]byte{secondaryKey}, keys) +} + func TestShardStoreScanAt_ReturnsTxnLockedWhenPendingLockExceedsUserLimit(t *testing.T) { t.Parallel() From 9adece85e3f70d95ffd0953bb2f7943192813af8 Mon Sep 17 00:00:00 2001 From: bootjp Date: Sat, 11 Jul 2026 17:15:16 +0900 Subject: [PATCH 06/10] backup: fix scan pagination edge cases --- kv/leader_routed_store.go | 9 +---- kv/shard_store.go | 18 ++++++--- kv/shard_store_test.go | 84 +++++++++++++++++++++++++++++++++++++++ store/lsm_store.go | 21 ++++++---- store/scan_keys_test.go | 36 +++++++++++++++++ 5 files changed, 147 insertions(+), 21 deletions(-) diff --git a/kv/leader_routed_store.go b/kv/leader_routed_store.go index 99ac1d213..d27e802da 100644 --- a/kv/leader_routed_store.go +++ b/kv/leader_routed_store.go @@ -276,14 +276,7 @@ func (s *LeaderRoutedStore) ScanKeysAt(ctx context.Context, start []byte, end [] if err != nil { return nil, err } - keys := make([][]byte, 0, len(kvs)) - for _, kvp := range kvs { - if kvp == nil { - continue - } - keys = append(keys, kvp.Key) - } - return keys, nil + return keysFromKVs(kvs), nil } func (s *LeaderRoutedStore) ReverseScanAt(ctx context.Context, start []byte, end []byte, limit int, ts uint64) ([]*store.KVPair, error) { diff --git a/kv/shard_store.go b/kv/shard_store.go index a09110b5d..4bd9a76f3 100644 --- a/kv/shard_store.go +++ b/kv/shard_store.go @@ -267,12 +267,18 @@ func (s *ShardStore) routesForScan(start []byte, end []byte) ([]distribution.Rou func (s *ShardStore) scanRoutesAt(ctx context.Context, routes []distribution.Route, start []byte, end []byte, limit int, ts uint64, clampToRoutes bool) ([]*store.KVPair, error) { out := make([]*store.KVPair, 0) + seenGroups := make(map[uint64]struct{}) for _, route := range routes { scanStart := start scanEnd := end if clampToRoutes { scanStart = clampScanStart(start, route.Start) scanEnd = clampScanEnd(end, route.End) + } else { + if _, seen := seenGroups[route.GroupID]; seen { + continue + } + seenGroups[route.GroupID] = struct{}{} } kvs, err := s.scanRouteAtDirection(ctx, route, scanStart, scanEnd, limit, ts, false, false) @@ -658,6 +664,9 @@ func (s *ShardStore) scanRouteAtForward( return nil, err } out = mergeAndTrimScanResults(out, page.kvs, limit) + if len(out) >= limit { + break + } if !page.full || page.advanceKey == nil { break } @@ -720,15 +729,14 @@ func (s *ShardStore) scanRouteAtForwardPage( }, nil } - kvs, err := s.proxyRawScanAt(ctx, g, start, end, limit, ts, false, route.GroupID) + raw, err := s.proxyRawScanAt(ctx, g, start, end, limit, ts, false, route.GroupID) if err != nil { return scanRoutePage{}, err } - kvs = filterTxnInternalKVs(kvs) return scanRoutePage{ - kvs: kvs, - advanceKey: lastKVKey(kvs), - full: len(kvs) >= limit, + kvs: filterTxnInternalKVs(raw), + advanceKey: lastKVKey(raw), + full: len(raw) >= limit, }, nil } diff --git a/kv/shard_store_test.go b/kv/shard_store_test.go index 5a7eabd6d..cc1daf8e7 100644 --- a/kv/shard_store_test.go +++ b/kv/shard_store_test.go @@ -235,6 +235,60 @@ func TestShardStoreScanKeysAt_DeduplicatesUnclampedGroups(t *testing.T) { require.Equal(t, [][]byte{[]byte("a"), []byte("z")}, keys) } +func TestShardStoreScanAt_DeduplicatesUnclampedGroups(t *testing.T) { + t.Parallel() + + ctx := context.Background() + engine := distribution.NewEngine() + engine.UpdateRoute([]byte(""), []byte("m"), 1) + engine.UpdateRoute([]byte("m"), nil, 1) + + groups := map[uint64]*ShardGroup{ + 1: {Store: store.NewMVCCStore()}, + } + st := NewShardStore(engine, groups) + + require.NoError(t, st.PutAt(ctx, []byte("a"), []byte("va"), 1, 0)) + require.NoError(t, st.PutAt(ctx, []byte("z"), []byte("vz"), 2, 0)) + + kvs, err := st.ScanAt(ctx, []byte(""), nil, 10, ^uint64(0)) + require.NoError(t, err) + require.Len(t, kvs, 2) + require.Equal(t, []byte("a"), kvs[0].Key) + require.Equal(t, []byte("z"), kvs[1].Key) +} + +func TestBackupScannerDeduplicatesUnclampedGroups(t *testing.T) { + t.Parallel() + + ctx := context.Background() + engine := distribution.NewEngine() + engine.UpdateRoute([]byte(""), []byte("m"), 1) + engine.UpdateRoute([]byte("m"), nil, 1) + + groups := map[uint64]*ShardGroup{ + 1: {Store: store.NewMVCCStore()}, + } + st := NewShardStore(engine, groups) + + require.NoError(t, st.PutAt(ctx, []byte("a"), []byte("va"), 1, 0)) + require.NoError(t, st.PutAt(ctx, []byte("z"), []byte("vz"), 2, 0)) + + sc := st.NewBackupScanner([]byte(""), nil, ^uint64(0), 10) + defer sc.Close() + + var got [][]byte + for { + kvp, ok, err := sc.Next(ctx) + require.NoError(t, err) + if !ok { + break + } + got = append(got, kvp.Key) + } + require.Equal(t, [][]byte{[]byte("a"), []byte("z")}, got) +} + func TestShardStoreScanKeysRouteAtLeaderRefillsAfterTxnInternalKeys(t *testing.T) { t.Parallel() @@ -326,6 +380,36 @@ func TestShardStoreProxyScanAtUsesSelectedGroup(t *testing.T) { require.Equal(t, uint64(42), fake.lastScanGroupID) } +func TestShardStoreProxyForwardPageAdvancesFromRawPage(t *testing.T) { + t.Parallel() + + internalKey := txnCommitKey([]byte("primary"), 10) + fake := &fakeRawKVServer{ + scanResp: &pb.RawScanAtResponse{ + Kv: []*pb.RawKVPair{ + {Key: []byte("a"), Value: []byte("va")}, + {Key: internalKey, Value: []byte("commit")}, + }, + }, + } + addr, stop := startRawKVServer(t, fake) + t.Cleanup(stop) + + ctx := context.Background() + g := &ShardGroup{ + Engine: &followerProxyEngine{leader: addr}, + Store: store.NewMVCCStore(), + } + st := NewShardStore(distribution.NewEngine(), map[uint64]*ShardGroup{42: g}) + + page, err := st.scanRouteAtForwardPage(ctx, distribution.Route{GroupID: 42}, g, []byte(""), nil, 2, ^uint64(0)) + require.NoError(t, err) + require.True(t, page.full) + require.Equal(t, internalKey, page.advanceKey) + require.Len(t, page.kvs, 1) + require.Equal(t, []byte("a"), page.kvs[0].Key) +} + func TestKeyOnlyScanHelpersPreserveEmptyKey(t *testing.T) { t.Parallel() diff --git a/store/lsm_store.go b/store/lsm_store.go index d31c73b29..6ec9d548b 100644 --- a/store/lsm_store.go +++ b/store/lsm_store.go @@ -1007,15 +1007,20 @@ func (s *pebbleStore) seekToVisibleVersion(iter *pebble.Iterator, userKey []byte } func (s *pebbleStore) skipToNextUserKey(iter *pebble.Iterator, userKey []byte) bool { - if !iter.SeekGE(encodeKey(userKey, 0)) { - return false - } - k := iter.Key() - u, _ := decodeKeyView(k) - if bytes.Equal(u, userKey) { - return iter.Next() + for iter.Next() { + rawKey := iter.Key() + if isPebbleMetaKey(rawKey) { + continue + } + nextUserKey, _ := decodeKeyView(rawKey) + if nextUserKey == nil { + continue + } + if !bytes.Equal(nextUserKey, userKey) { + return true + } } - return true + return false } func pastScanEnd(userKey, end []byte) bool { diff --git a/store/scan_keys_test.go b/store/scan_keys_test.go index 527e77581..9031370fa 100644 --- a/store/scan_keys_test.go +++ b/store/scan_keys_test.go @@ -100,6 +100,42 @@ func TestPebbleStoreScanKeysAtCursorAfterKeyDoesNotRepeat(t *testing.T) { require.Equal(t, [][]byte{[]byte("a")}, keys) } +func TestPebbleStoreScanAtPreservesPrefixExtensionKeys(t *testing.T) { + t.Parallel() + + ctx := context.Background() + st, err := NewPebbleStore(t.TempDir()) + require.NoError(t, err) + defer st.Close() + + highTS := ^uint64(0) - 1 + require.NoError(t, st.PutAt(ctx, []byte("a"), []byte("va"), highTS, 0)) + require.NoError(t, st.PutAt(ctx, []byte("a\x00"), []byte("vax"), 10, 0)) + require.NoError(t, st.PutAt(ctx, []byte("b"), []byte("vb"), 20, 0)) + + kvs, err := st.ScanAt(ctx, []byte("a"), []byte("c"), 3, ^uint64(0)) + require.NoError(t, err) + require.Equal(t, [][]byte{[]byte("a"), []byte("a\x00"), []byte("b")}, keysFromStoreKVs(kvs)) +} + +func TestPebbleStoreScanKeysAtPreservesPrefixExtensionKeys(t *testing.T) { + t.Parallel() + + ctx := context.Background() + st, err := NewPebbleStore(t.TempDir()) + require.NoError(t, err) + defer st.Close() + + highTS := ^uint64(0) - 1 + require.NoError(t, st.PutAt(ctx, []byte("a"), []byte("va"), highTS, 0)) + require.NoError(t, st.PutAt(ctx, []byte("a\x00"), []byte("vax"), 10, 0)) + require.NoError(t, st.PutAt(ctx, []byte("b"), []byte("vb"), 20, 0)) + + keys, err := st.ScanKeysAt(ctx, []byte("a"), []byte("c"), 3, ^uint64(0)) + require.NoError(t, err) + require.Equal(t, [][]byte{[]byte("a"), []byte("a\x00"), []byte("b")}, keys) +} + func cursorAfterTestKey(key []byte) []byte { cursor := make([]byte, len(key)+1) copy(cursor, key) From d0e8b04df12288deda2c9508d5976fbb8945d8b5 Mon Sep 17 00:00:00 2001 From: bootjp Date: Fri, 17 Jul 2026 11:48:59 +0900 Subject: [PATCH 07/10] backup: preserve key scan logical ordering --- store/lsm_store.go | 66 ++++++++++++++++++++++++++++++++++++++++- store/scan_keys_test.go | 22 ++++++++++++++ 2 files changed, 87 insertions(+), 1 deletion(-) diff --git a/store/lsm_store.go b/store/lsm_store.go index b0139010c..26e674369 100644 --- a/store/lsm_store.go +++ b/store/lsm_store.go @@ -12,6 +12,7 @@ import ( "math" "os" "path/filepath" + "sort" "strconv" "strings" "sync" @@ -1166,6 +1167,7 @@ func (s *pebbleStore) collectForwardScanKV(iter *pebble.Iterator, userKey []byte func (s *pebbleStore) collectScanKeys(ctx context.Context, iter *pebble.Iterator, start, end []byte, limit int, ts uint64) ([][]byte, error) { result := make([][]byte, 0, boundedScanResultCapacity(limit)) + seen := make(map[string]struct{}, boundedScanResultCapacity(limit)) for iter.SeekGE(encodeKey(start, math.MaxUint64)); iter.Valid() && len(result) < limit; { userKey, version, ok, err := s.nextForwardScanUserKeyContext(ctx, iter, start) @@ -1185,7 +1187,11 @@ func (s *pebbleStore) collectScanKeys(ctx context.Context, iter *pebble.Iterator return nil, err } if visible { - result = append(result, userKey) + result = appendScanKeyResult(result, seen, userKey, limit) + result, err = s.appendVisibleProperPrefixScanKeys(ctx, result, seen, userKey, start, end, limit, ts) + if err != nil { + return nil, err + } } if !s.skipToNextUserKey(iter, userKey) { @@ -1196,6 +1202,64 @@ func (s *pebbleStore) collectScanKeys(ctx context.Context, iter *pebble.Iterator return result, nil } +func appendScanKeyResult(result [][]byte, seen map[string]struct{}, key []byte, limit int) [][]byte { + if _, ok := seen[string(key)]; ok { + return result + } + keyCopy := make([]byte, len(key)) + copy(keyCopy, key) + seen[string(keyCopy)] = struct{}{} + result = append(result, keyCopy) + sort.Slice(result, func(i, j int) bool { + return bytes.Compare(result[i], result[j]) < 0 + }) + for len(result) > limit { + delete(seen, string(result[len(result)-1])) + result = result[:len(result)-1] + } + return result +} + +func (s *pebbleStore) appendVisibleProperPrefixScanKeys( + ctx context.Context, + result [][]byte, + seen map[string]struct{}, + userKey []byte, + start []byte, + end []byte, + limit int, + ts uint64, +) ([][]byte, error) { + for i := 0; i < len(userKey); i++ { + prefix := userKey[:i] + if _, ok := seen[string(prefix)]; ok { + continue + } + if beforeScanStart(prefix, start) || pastScanEnd(prefix, end) { + continue + } + visible, err := s.visibleExactKeyAt(ctx, prefix, ts) + if err != nil { + return nil, err + } + if visible { + result = appendScanKeyResult(result, seen, prefix, limit) + } + } + return result, nil +} + +func (s *pebbleStore) visibleExactKeyAt(ctx context.Context, key []byte, ts uint64) (bool, error) { + _, err := s.getAt(ctx, key, ts) + if errors.Is(err, ErrKeyNotFound) { + return false, nil + } + if err != nil { + return false, err + } + return true, nil +} + func (s *pebbleStore) ScanAt(ctx context.Context, start []byte, end []byte, limit int, ts uint64) ([]*KVPair, error) { if limit <= 0 { return []*KVPair{}, nil diff --git a/store/scan_keys_test.go b/store/scan_keys_test.go index 9031370fa..8ba7c5479 100644 --- a/store/scan_keys_test.go +++ b/store/scan_keys_test.go @@ -136,6 +136,28 @@ func TestPebbleStoreScanKeysAtPreservesPrefixExtensionKeys(t *testing.T) { require.Equal(t, [][]byte{[]byte("a"), []byte("a\x00"), []byte("b")}, keys) } +func TestPebbleStoreScanKeysAtSortsPrefixBeforeLimit(t *testing.T) { + t.Parallel() + + ctx := context.Background() + st, err := NewPebbleStore(t.TempDir()) + require.NoError(t, err) + defer st.Close() + + require.NoError(t, st.PutAt(ctx, []byte(""), []byte("empty"), 10, 0)) + require.NoError(t, st.PutAt(ctx, []byte("a"), []byte("va"), 20, 0)) + require.NoError(t, st.PutAt(ctx, []byte("\x00"), []byte("prefix"), 30, 0)) + require.NoError(t, st.PutAt(ctx, []byte("\x00a"), []byte("extension"), 40, 0)) + + keys, err := st.ScanKeysAt(ctx, nil, nil, 1, 40) + require.NoError(t, err) + require.Equal(t, [][]byte{[]byte("")}, keys) + + keys, err = st.ScanKeysAt(ctx, cursorAfterTestKey([]byte("")), nil, 1, 40) + require.NoError(t, err) + require.Equal(t, [][]byte{[]byte("\x00")}, keys) +} + func cursorAfterTestKey(key []byte) []byte { cursor := make([]byte, len(key)+1) copy(cursor, key) From a9b83546267cac77aba653554f9994163495a180 Mon Sep 17 00:00:00 2001 From: bootjp Date: Fri, 17 Jul 2026 12:30:02 +0900 Subject: [PATCH 08/10] backup: page scans by logical keys --- adapter/grpc.go | 85 ++++++++++++++++++++++++++++--- adapter/grpc_test.go | 36 +++++++++++++ kv/backup_scan.go | 25 ++++++--- kv/leader_routed_store.go | 45 +++++++++++++++-- kv/leader_routed_store_test.go | 4 +- kv/shard_store.go | 92 ++++++++++++++++++++++++++-------- kv/shard_store_test.go | 52 +++++++++++++++++++ proto/service.pb.go | 15 ++++-- proto/service.proto | 1 + store/lsm_store.go | 43 ++++++++++++++-- store/scan_keys_test.go | 20 ++++++++ 11 files changed, 370 insertions(+), 48 deletions(-) diff --git a/adapter/grpc.go b/adapter/grpc.go index 712efd911..3e1dcfab4 100644 --- a/adapter/grpc.go +++ b/adapter/grpc.go @@ -44,6 +44,10 @@ type rawGroupScanner interface { ScanGroupAt(ctx context.Context, groupID uint64, start []byte, end []byte, limit int, ts uint64) ([]*store.KVPair, error) } +type rawGroupKeyScanner interface { + ScanGroupKeysAt(ctx context.Context, groupID uint64, start []byte, end []byte, limit int, ts uint64) ([][]byte, error) +} + func WithCloseStore() GRPCServerOption { return func(s *GRPCServer) { s.closeStore = true @@ -159,14 +163,32 @@ func (r *GRPCServer) RawScanAt(ctx context.Context, req *pb.RawScanAtRequest) (* readTS = globalSnapshotTS(ctx, r.clock(), r.store) } + if req.GetKeysOnly() { + keys, err := r.rawScanKeysAt(ctx, req, limit, readTS) + if err != nil { + return rawScanErrorResponse(err) + } + return &pb.RawScanAtResponse{Kv: rawKeyPairs(keys)}, nil + } + + res, err := r.rawScanValuesAt(ctx, req, limit, readTS) + if err != nil { + return rawScanErrorResponse(err) + } + + return &pb.RawScanAtResponse{Kv: rawKvPairs(res)}, nil +} + +func (r *GRPCServer) rawScanValuesAt(ctx context.Context, req *pb.RawScanAtRequest, limit int, readTS uint64) ([]*store.KVPair, error) { var res []*store.KVPair + var err error if groupID := req.GetGroupId(); groupID != 0 { if req.GetReverse() { - return &pb.RawScanAtResponse{Kv: nil}, errors.WithStack(status.Error(codes.InvalidArgument, "raw scan with explicit group does not support reverse scans")) + return nil, errors.WithStack(status.Error(codes.InvalidArgument, "raw scan with explicit group does not support reverse scans")) } groupScanner, ok := r.store.(rawGroupScanner) if !ok { - return &pb.RawScanAtResponse{Kv: nil}, errors.WithStack(status.Error(codes.FailedPrecondition, "raw scan with explicit group requires a group-aware store")) + return nil, errors.WithStack(status.Error(codes.FailedPrecondition, "raw scan with explicit group requires a group-aware store")) } res, err = groupScanner.ScanGroupAt(ctx, groupID, req.StartKey, req.EndKey, limit, readTS) } else if req.GetReverse() { @@ -175,13 +197,42 @@ func (r *GRPCServer) RawScanAt(ctx context.Context, req *pb.RawScanAtRequest) (* res, err = r.store.ScanAt(ctx, req.StartKey, req.EndKey, limit, readTS) } if err != nil { - if errors.Is(err, store.ErrReadTSCompacted) { - return &pb.RawScanAtResponse{Kv: nil}, errors.WithStack(status.Error(codes.FailedPrecondition, store.ErrReadTSCompacted.Error())) + return nil, errors.WithStack(err) + } + return res, nil +} + +func (r *GRPCServer) rawScanKeysAt(ctx context.Context, req *pb.RawScanAtRequest, limit int, readTS uint64) ([][]byte, error) { + if groupID := req.GetGroupId(); groupID != 0 { + if req.GetReverse() { + return nil, errors.WithStack(status.Error(codes.InvalidArgument, "raw scan with explicit group does not support reverse scans")) } - return &pb.RawScanAtResponse{Kv: nil}, errors.WithStack(err) + groupScanner, ok := r.store.(rawGroupKeyScanner) + if !ok { + return nil, errors.WithStack(status.Error(codes.FailedPrecondition, "raw key scan with explicit group requires a group-aware store")) + } + keys, err := groupScanner.ScanGroupKeysAt(ctx, groupID, req.StartKey, req.EndKey, limit, readTS) + if err != nil { + return nil, errors.WithStack(err) + } + return keys, nil } + if req.GetReverse() { + kvs, err := r.store.ReverseScanAt(ctx, req.StartKey, req.EndKey, limit, readTS) + if err != nil { + return nil, errors.WithStack(err) + } + return storeKeysFromKVPairs(kvs), nil + } + keys, err := r.store.ScanKeysAt(ctx, req.StartKey, req.EndKey, limit, readTS) + return keys, errors.WithStack(err) +} - return &pb.RawScanAtResponse{Kv: rawKvPairs(res)}, nil +func rawScanErrorResponse(err error) (*pb.RawScanAtResponse, error) { + if errors.Is(err, store.ErrReadTSCompacted) { + return &pb.RawScanAtResponse{Kv: nil}, errors.WithStack(status.Error(codes.FailedPrecondition, store.ErrReadTSCompacted.Error())) + } + return &pb.RawScanAtResponse{Kv: nil}, errors.WithStack(err) } func rawScanLimit(limit64 int64) (int, error) { @@ -225,6 +276,28 @@ func rawKvPairs(res []*store.KVPair) []*pb.RawKVPair { return out } +func rawKeyPairs(keys [][]byte) []*pb.RawKVPair { + out := make([]*pb.RawKVPair, 0, len(keys)) + for _, key := range keys { + if key == nil { + continue + } + out = append(out, &pb.RawKVPair{Key: key}) + } + return out +} + +func storeKeysFromKVPairs(kvs []*store.KVPair) [][]byte { + keys := make([][]byte, 0, len(kvs)) + for _, kvp := range kvs { + if kvp == nil { + continue + } + keys = append(keys, kvp.Key) + } + return keys +} + func (r *GRPCServer) RawPut(ctx context.Context, req *pb.RawPutRequest) (*pb.RawPutResponse, error) { m, err := r.grpcTranscoder.RawPutToRequest(req) if err != nil { diff --git a/adapter/grpc_test.go b/adapter/grpc_test.go index fd1393cd5..8d5d73c9a 100644 --- a/adapter/grpc_test.go +++ b/adapter/grpc_test.go @@ -142,6 +142,7 @@ type recordingRawGroupStore struct { scanGroupID uint64 scanStart []byte scanEnd []byte + keyScanGroup bool fallbackGet bool fallbackScan bool } @@ -169,6 +170,14 @@ func (s *recordingRawGroupStore) ScanGroupAt(ctx context.Context, groupID uint64 return s.MVCCStore.ScanAt(ctx, start, end, limit, ts) } +func (s *recordingRawGroupStore) ScanGroupKeysAt(ctx context.Context, groupID uint64, start []byte, end []byte, limit int, ts uint64) ([][]byte, error) { + s.scanGroupID = groupID + s.scanStart = append([]byte(nil), start...) + s.scanEnd = append([]byte(nil), end...) + s.keyScanGroup = true + return s.ScanKeysAt(ctx, start, end, limit, ts) +} + func TestGRPCServer_RawGet_UsesExplicitGroup(t *testing.T) { t.Parallel() @@ -209,6 +218,33 @@ func TestGRPCServer_RawScanAt_UsesExplicitGroup(t *testing.T) { require.Equal(t, []byte("z"), st.scanEnd) } +func TestGRPCServer_RawScanAt_KeysOnlyUsesExplicitGroup(t *testing.T) { + t.Parallel() + + ctx := context.Background() + st := &recordingRawGroupStore{MVCCStore: store.NewMVCCStore()} + require.NoError(t, st.PutAt(ctx, []byte("a"), []byte("large-value"), 9, 0)) + s := NewGRPCServer(st, nil) + + resp, err := s.RawScanAt(ctx, &pb.RawScanAtRequest{ + StartKey: []byte("a"), + EndKey: []byte("z"), + Limit: 10, + Ts: 9, + GroupId: 42, + KeysOnly: true, + }) + require.NoError(t, err) + require.Len(t, resp.GetKv(), 1) + require.Equal(t, []byte("a"), resp.GetKv()[0].GetKey()) + require.Empty(t, resp.GetKv()[0].GetValue()) + require.True(t, st.keyScanGroup) + require.False(t, st.fallbackScan) + require.Equal(t, uint64(42), st.scanGroupID) + require.Equal(t, []byte("a"), st.scanStart) + require.Equal(t, []byte("z"), st.scanEnd) +} + func TestGRPCServer_Scan_RejectsOversizedLimit(t *testing.T) { t.Parallel() diff --git a/kv/backup_scan.go b/kv/backup_scan.go index 9f8d514f5..efb3da80f 100644 --- a/kv/backup_scan.go +++ b/kv/backup_scan.go @@ -6,12 +6,13 @@ import ( "github.com/bootjp/elastickv/distribution" "github.com/bootjp/elastickv/store" + "github.com/cockroachdb/errors" ) const defaultBackupScanPageSize = 1024 -// BackupScanner pages through ShardStore.ScanAt without holding store locks -// across pages. +// BackupScanner pages through ShardStore.ScanKeysAt without holding store locks +// across pages, then materializes each key at the pinned read timestamp. type BackupScanner interface { Next(ctx context.Context) (*store.KVPair, bool, error) Close() error @@ -89,21 +90,31 @@ func (s *backupScanner) loadNextPage(ctx context.Context) error { s.index = 0 return nil } - page, err := s.store.scanRoutesAtSorted(ctx, s.routes, s.cursor, s.end, s.pageSize, s.ts, s.clampToRoutes) + keys, err := s.store.scanKeyRoutesAt(ctx, s.routes, s.cursor, s.end, s.pageSize, s.ts, s.clampToRoutes) if err != nil { return err } - s.page = page + s.page = s.page[:0] + for _, key := range keys { + val, err := s.store.GetAt(ctx, key, s.ts) + if errors.Is(err, store.ErrKeyNotFound) { + continue + } + if err != nil { + return err + } + s.page = append(s.page, &store.KVPair{Key: bytes.Clone(key), Value: bytes.Clone(val)}) + } s.index = 0 - if len(page) == 0 { + if len(keys) == 0 { s.exhausted = true return nil } - last := page[len(page)-1] + last := lastScanKey(keys) if last == nil { s.exhausted = true return nil } - s.cursor = nextScanCursor(last.Key) + s.cursor = nextScanCursor(last) return nil } diff --git a/kv/leader_routed_store.go b/kv/leader_routed_store.go index dd264ac05..47e4bc6b1 100644 --- a/kv/leader_routed_store.go +++ b/kv/leader_routed_store.go @@ -153,6 +153,45 @@ func (s *LeaderRoutedStore) proxyRawScanAt( return out, nil } +func (s *LeaderRoutedStore) proxyRawScanKeysAt( + ctx context.Context, + start []byte, + end []byte, + limit int, + ts uint64, +) ([][]byte, error) { + addr := s.leaderAddrForKey(start) + if addr == "" { + return nil, errors.WithStack(ErrLeaderNotFound) + } + + conn, err := s.connCache.ConnFor(addr) + if err != nil { + return nil, err + } + + cli := pb.NewRawKVClient(conn) + resp, err := cli.RawScanAt(ctx, &pb.RawScanAtRequest{ + StartKey: start, + EndKey: end, + Limit: int64(limit), + Ts: ts, + KeysOnly: true, + }) + if err != nil { + return nil, errors.WithStack(err) + } + + out := make([][]byte, 0, len(resp.Kv)) + for _, kvp := range resp.Kv { + if kvp == nil { + continue + } + out = append(out, bytes.Clone(kvp.Key)) + } + return out, nil +} + func (s *LeaderRoutedStore) GetAt(ctx context.Context, key []byte, ts uint64) ([]byte, error) { if s == nil || s.local == nil { return nil, store.ErrKeyNotFound @@ -272,11 +311,7 @@ func (s *LeaderRoutedStore) ScanKeysAt(ctx context.Context, start []byte, end [] keys, err := s.local.ScanKeysAt(ctx, start, end, limit, max(ts, fenceTS)) return keys, errors.WithStack(err) } - kvs, err := s.proxyRawScanAt(ctx, start, end, limit, ts, false) - if err != nil { - return nil, err - } - return keysFromKVs(kvs), nil + return s.proxyRawScanKeysAt(ctx, start, end, limit, ts) } func (s *LeaderRoutedStore) ScanAtPhysicalLimit(ctx context.Context, start []byte, end []byte, visibleLimit, physicalLimit int, ts uint64) ([]*store.KVPair, bool, error) { diff --git a/kv/leader_routed_store_test.go b/kv/leader_routed_store_test.go index c6abd562a..22ddc9535 100644 --- a/kv/leader_routed_store_test.go +++ b/kv/leader_routed_store_test.go @@ -85,7 +85,8 @@ type fakeRawKVServer struct { scanCalls int latestCalls int - lastScanGroupID uint64 + lastScanGroupID uint64 + lastScanKeysOnly bool getResp *pb.RawGetResponse scanResp *pb.RawScanAtResponse @@ -107,6 +108,7 @@ func (f *fakeRawKVServer) RawScanAt(_ context.Context, req *pb.RawScanAtRequest) defer f.mu.Unlock() f.scanCalls++ f.lastScanGroupID = req.GetGroupId() + f.lastScanKeysOnly = req.GetKeysOnly() if f.scanResp != nil { return f.scanResp, nil } diff --git a/kv/shard_store.go b/kv/shard_store.go index 3068af12a..ca0fdece5 100644 --- a/kv/shard_store.go +++ b/kv/shard_store.go @@ -236,6 +236,15 @@ func (s *ShardStore) ScanGroupAt(ctx context.Context, groupID uint64, start []by return s.scanRouteAtDirection(ctx, distribution.Route{GroupID: groupID}, start, end, limit, ts, false, true) } +// ScanGroupKeysAt scans keys on the explicitly selected Raft group without +// materializing values over proxy links. +func (s *ShardStore) ScanGroupKeysAt(ctx context.Context, groupID uint64, start []byte, end []byte, limit int, ts uint64) ([][]byte, error) { + if limit <= 0 { + return [][]byte{}, nil + } + return s.scanKeyRouteAt(ctx, distribution.Route{GroupID: groupID}, start, end, limit, ts) +} + func (s *ShardStore) ReverseScanAt(ctx context.Context, start []byte, end []byte, limit int, ts uint64) ([]*store.KVPair, error) { if limit <= 0 { return []*store.KVPair{}, nil @@ -534,11 +543,7 @@ func (s *ShardStore) proxyScanKeysAt( groupID uint64, ) ([][]byte, error) { return scanKeysWithRefill(start, end, limit, func(cursor []byte, pageLimit int) ([][]byte, error) { - kvs, err := s.proxyRawScanAt(ctx, g, cursor, end, pageLimit, ts, false, groupID) - if err != nil { - return nil, err - } - return keysFromKVs(kvs), nil + return s.proxyRawScanKeysAt(ctx, g, cursor, end, pageLimit, ts, groupID) }) } @@ -911,20 +916,26 @@ func (s *ShardStore) scanRouteAtLeader( } func scanLockBoundsForKVs(kvs []*store.KVPair, scanStart []byte, scanEnd []byte, limit int) ([]byte, []byte) { - if countNonInternalKVs(kvs) < limit { + if len(kvs) < limit { return scanStart, scanEnd } _, lastUserKey, ok := observedScanUserBounds(kvs) - if !ok { - return scanStart, scanEnd + if ok { + return scanStart, boundScanEnd(scanEnd, nextScanCursor(lastUserKey)) } - bound := nextScanCursor(lastUserKey) - if scanEnd == nil || bytes.Compare(bound, scanEnd) < 0 { - scanEnd = bound + if lastKey := lastKVKey(kvs); lastKey != nil { + return scanStart, boundScanEnd(scanEnd, nextScanCursor(lastKey)) } return scanStart, scanEnd } +func boundScanEnd(scanEnd []byte, bound []byte) []byte { + if scanEnd == nil || bytes.Compare(bound, scanEnd) < 0 { + return bound + } + return scanEnd +} + func observedScanUserBounds(kvs []*store.KVPair) ([]byte, []byte, bool) { var minKey []byte var maxKey []byte @@ -1016,17 +1027,6 @@ func mergeAndTrimReverseScanResults(out []*store.KVPair, kvs []*store.KVPair, li return out[:limit] } -func countNonInternalKVs(kvs []*store.KVPair) int { - count := 0 - for _, kvp := range kvs { - if kvp == nil || kvp.Key == nil || isTxnInternalKey(kvp.Key) { - continue - } - count++ - } - return count -} - func kvPairsFromKeys(keys [][]byte) []*store.KVPair { kvs := make([]*store.KVPair, 0, len(keys)) for _, key := range keys { @@ -2142,4 +2142,52 @@ func (s *ShardStore) proxyRawScanAt( return out, nil } +func (s *ShardStore) proxyRawScanKeysAt( + ctx context.Context, + g *ShardGroup, + start []byte, + end []byte, + limit int, + ts uint64, + groupID uint64, +) ([][]byte, error) { + engine := engineForGroup(g) + if engine == nil { + return nil, store.ErrNotSupported + } + addr := leaderAddrFromEngine(engine) + if addr == "" { + return nil, errors.WithStack(ErrLeaderNotFound) + } + + conn, err := s.connCache.ConnFor(addr) + if err != nil { + return nil, err + } + + ctx, cancel := context.WithTimeout(ctx, proxyForwardTimeout) + defer cancel() + cli := pb.NewRawKVClient(conn) + resp, err := cli.RawScanAt(ctx, &pb.RawScanAtRequest{ + StartKey: start, + EndKey: end, + Limit: int64(limit), + Ts: ts, + GroupId: groupID, + KeysOnly: true, + }) + if err != nil { + return nil, errors.WithStack(err) + } + + out := make([][]byte, 0, len(resp.Kv)) + for _, kvp := range resp.Kv { + if kvp == nil { + continue + } + out = append(out, bytes.Clone(kvp.Key)) + } + return out, nil +} + var _ store.MVCCStore = (*ShardStore)(nil) diff --git a/kv/shard_store_test.go b/kv/shard_store_test.go index cc1daf8e7..d873a6bdb 100644 --- a/kv/shard_store_test.go +++ b/kv/shard_store_test.go @@ -347,6 +347,7 @@ func TestShardStoreProxyScanKeysAtUsesSelectedGroup(t *testing.T) { defer fake.mu.Unlock() require.Equal(t, 1, fake.scanCalls) require.Equal(t, uint64(42), fake.lastScanGroupID) + require.True(t, fake.lastScanKeysOnly) } func TestShardStoreProxyScanAtUsesSelectedGroup(t *testing.T) { @@ -589,6 +590,44 @@ func TestBackupScannerPebblePagingDoesNotRepeatLastKey(t *testing.T) { require.Nil(t, kvp) } +func TestBackupScannerPagesPebbleKeysInLogicalOrder(t *testing.T) { + t.Parallel() + + ctx := context.Background() + engine := distribution.NewEngine() + engine.UpdateRoute([]byte(""), nil, 1) + pebbleStore, err := store.NewPebbleStore(t.TempDir()) + require.NoError(t, err) + defer pebbleStore.Close() + + groups := map[uint64]*ShardGroup{ + 1: {Store: pebbleStore}, + } + st := NewShardStore(engine, groups) + require.NoError(t, pebbleStore.PutAt(ctx, []byte("a"), []byte("va"), 10, 0)) + require.NoError(t, pebbleStore.PutAt(ctx, []byte("a\x80"), []byte("vax"), 20, 0)) + + sc := st.NewBackupScanner(nil, nil, 20, 1) + defer sc.Close() + + kvp, ok, err := sc.Next(ctx) + require.NoError(t, err) + require.True(t, ok) + require.Equal(t, []byte("a"), kvp.Key) + require.Equal(t, []byte("va"), kvp.Value) + + kvp, ok, err = sc.Next(ctx) + require.NoError(t, err) + require.True(t, ok) + require.Equal(t, []byte("a\x80"), kvp.Key) + require.Equal(t, []byte("vax"), kvp.Value) + + kvp, ok, err = sc.Next(ctx) + require.NoError(t, err) + require.False(t, ok) + require.Nil(t, kvp) +} + func TestShardStoreScanAt_RoutesS3ManifestScansByLogicalObjectKey(t *testing.T) { t.Parallel() @@ -854,3 +893,16 @@ func TestScanLockBoundsForKVs_EmptyUsesOriginalRange(t *testing.T) { require.Equal(t, []byte("a"), lockStart) require.Equal(t, []byte("z"), lockEnd) } + +func TestScanLockBoundsForKVs_FullInternalPageUsesRawPageBound(t *testing.T) { + t.Parallel() + + internalKey := txnCommitKey([]byte("primary"), 10) + kvs := []*store.KVPair{ + {Key: internalKey, Value: []byte("commit")}, + } + + lockStart, lockEnd := scanLockBoundsForKVs(kvs, []byte(""), nil, 1) + require.Equal(t, []byte(""), lockStart) + require.Equal(t, nextScanCursor(internalKey), lockEnd) +} diff --git a/proto/service.pb.go b/proto/service.pb.go index a869e0c1d..00000ae2f 100644 --- a/proto/service.pb.go +++ b/proto/service.pb.go @@ -499,7 +499,8 @@ type RawScanAtRequest struct { Limit int64 `protobuf:"varint,3,opt,name=limit,proto3" json:"limit,omitempty"` // validated against host int size; large values may be rejected Ts uint64 `protobuf:"varint,4,opt,name=ts,proto3" json:"ts,omitempty"` // optional read timestamp; if zero, server uses current HLC Reverse bool `protobuf:"varint,5,opt,name=reverse,proto3" json:"reverse,omitempty"` - GroupId uint64 `protobuf:"varint,6,opt,name=group_id,json=groupId,proto3" json:"group_id,omitempty"` // optional explicit Raft group for non-range-owned keyspaces + GroupId uint64 `protobuf:"varint,6,opt,name=group_id,json=groupId,proto3" json:"group_id,omitempty"` // optional explicit Raft group for non-range-owned keyspaces + KeysOnly bool `protobuf:"varint,7,opt,name=keys_only,json=keysOnly,proto3" json:"keys_only,omitempty"` // when true, response kv entries omit values unknownFields protoimpl.UnknownFields sizeCache protoimpl.SizeCache } @@ -576,6 +577,13 @@ func (x *RawScanAtRequest) GetGroupId() uint64 { return 0 } +func (x *RawScanAtRequest) GetKeysOnly() bool { + if x != nil { + return x.KeysOnly + } + return false +} + type RawKVPair struct { state protoimpl.MessageState `protogen:"open.v1"` Key []byte `protobuf:"bytes,1,opt,name=key,proto3" json:"key,omitempty"` @@ -2279,14 +2287,15 @@ const file_service_proto_rawDesc = "" + "\x03key\x18\x01 \x01(\fR\x03key\"C\n" + "\x19RawLatestCommitTSResponse\x12\x0e\n" + "\x02ts\x18\x01 \x01(\x04R\x02ts\x12\x16\n" + - "\x06exists\x18\x02 \x01(\bR\x06exists\"\xa3\x01\n" + + "\x06exists\x18\x02 \x01(\bR\x06exists\"\xc0\x01\n" + "\x10RawScanAtRequest\x12\x1b\n" + "\tstart_key\x18\x01 \x01(\fR\bstartKey\x12\x17\n" + "\aend_key\x18\x02 \x01(\fR\x06endKey\x12\x14\n" + "\x05limit\x18\x03 \x01(\x03R\x05limit\x12\x0e\n" + "\x02ts\x18\x04 \x01(\x04R\x02ts\x12\x18\n" + "\areverse\x18\x05 \x01(\bR\areverse\x12\x19\n" + - "\bgroup_id\x18\x06 \x01(\x04R\agroupId\"3\n" + + "\bgroup_id\x18\x06 \x01(\x04R\agroupId\x12\x1b\n" + + "\tkeys_only\x18\a \x01(\bR\bkeysOnly\"3\n" + "\tRawKVPair\x12\x10\n" + "\x03key\x18\x01 \x01(\fR\x03key\x12\x14\n" + "\x05value\x18\x02 \x01(\fR\x05value\"/\n" + diff --git a/proto/service.proto b/proto/service.proto index b1ac9105f..68d12d5f7 100644 --- a/proto/service.proto +++ b/proto/service.proto @@ -78,6 +78,7 @@ message RawScanAtRequest { uint64 ts = 4; // optional read timestamp; if zero, server uses current HLC bool reverse = 5; uint64 group_id = 6; // optional explicit Raft group for non-range-owned keyspaces + bool keys_only = 7; // when true, response kv entries omit values } message RawKVPair { diff --git a/store/lsm_store.go b/store/lsm_store.go index 26e674369..41b0a9eb4 100644 --- a/store/lsm_store.go +++ b/store/lsm_store.go @@ -1250,14 +1250,49 @@ func (s *pebbleStore) appendVisibleProperPrefixScanKeys( } func (s *pebbleStore) visibleExactKeyAt(ctx context.Context, key []byte, ts uint64) (bool, error) { - _, err := s.getAt(ctx, key, ts) + iter, err := s.db.NewIter(&pebble.IterOptions{ + LowerBound: encodeKey(key, math.MaxUint64), + UpperBound: keyUpperBound(key), + }) + if err != nil { + return false, errors.WithStack(err) + } + defer iter.Close() + + for iter.SeekGE(encodeKey(key, math.MaxUint64)); iter.Valid(); iter.Next() { + if err := ctx.Err(); err != nil { + return false, errors.WithStack(err) + } + visible, stop, err := s.visibleExactIteratorEntry(iter, key, ts) + if err != nil { + return false, err + } + if stop { + return visible, nil + } + } + return false, nil +} + +func (s *pebbleStore) visibleExactIteratorEntry(iter *pebble.Iterator, key []byte, ts uint64) (bool, bool, error) { + userKey, version := decodeKeyView(iter.Key()) + if userKey == nil { + return false, false, nil + } + if len(key) > 0 && !bytes.HasPrefix(userKey, key) { + return false, true, nil + } + if !bytes.Equal(userKey, key) || version > ts { + return false, false, nil + } + _, err := s.readVisibleVersion(iter, key, ts) if errors.Is(err, ErrKeyNotFound) { - return false, nil + return false, true, nil } if err != nil { - return false, err + return false, true, err } - return true, nil + return true, true, nil } func (s *pebbleStore) ScanAt(ctx context.Context, start []byte, end []byte, limit int, ts uint64) ([]*KVPair, error) { diff --git a/store/scan_keys_test.go b/store/scan_keys_test.go index 8ba7c5479..1756beb24 100644 --- a/store/scan_keys_test.go +++ b/store/scan_keys_test.go @@ -158,6 +158,26 @@ func TestPebbleStoreScanKeysAtSortsPrefixBeforeLimit(t *testing.T) { require.Equal(t, [][]byte{[]byte("\x00")}, keys) } +func TestPebbleStoreScanKeysAtProbesExactPrefixPastBinaryExtension(t *testing.T) { + t.Parallel() + + ctx := context.Background() + st, err := NewPebbleStore(t.TempDir()) + require.NoError(t, err) + defer st.Close() + + require.NoError(t, st.PutAt(ctx, []byte("a"), []byte("va"), 10, 0)) + require.NoError(t, st.PutAt(ctx, []byte("a\x80"), []byte("vax"), 20, 0)) + + keys, err := st.ScanKeysAt(ctx, nil, nil, 1, 20) + require.NoError(t, err) + require.Equal(t, [][]byte{[]byte("a")}, keys) + + keys, err = st.ScanKeysAt(ctx, cursorAfterTestKey([]byte("a")), nil, 1, 20) + require.NoError(t, err) + require.Equal(t, [][]byte{[]byte("a\x80")}, keys) +} + func cursorAfterTestKey(key []byte) []byte { cursor := make([]byte, len(key)+1) copy(cursor, key) From f709355173c2605447d5c2980e9845edcb27c97e Mon Sep 17 00:00:00 2001 From: bootjp Date: Fri, 17 Jul 2026 13:26:07 +0900 Subject: [PATCH 09/10] store: keep prefix keys before scan end --- store/lsm_store.go | 79 +++++++++++++++++++++++++++++++++-------- store/scan_keys_test.go | 16 +++++++++ 2 files changed, 80 insertions(+), 15 deletions(-) diff --git a/store/lsm_store.go b/store/lsm_store.go index 41b0a9eb4..0b7d60c2c 100644 --- a/store/lsm_store.go +++ b/store/lsm_store.go @@ -1174,27 +1174,21 @@ func (s *pebbleStore) collectScanKeys(ctx context.Context, iter *pebble.Iterator if err != nil { return nil, err } - if forwardScanDone(userKey, end, ok) { - break + var done bool + result, done, err = s.appendEndPrefixScanKeys(ctx, result, seen, userKey, start, end, limit, ts, ok) + if err != nil { + return nil, err } - - if !s.seekToVisibleVersion(iter, userKey, version, ts) { - continue + if done { + break } - visible, err := s.foundValueVisible(iter, ts) + var advance bool + result, advance, err = s.collectVisibleScanKey(ctx, iter, result, seen, userKey, version, start, end, limit, ts) if err != nil { return nil, err } - if visible { - result = appendScanKeyResult(result, seen, userKey, limit) - result, err = s.appendVisibleProperPrefixScanKeys(ctx, result, seen, userKey, start, end, limit, ts) - if err != nil { - return nil, err - } - } - - if !s.skipToNextUserKey(iter, userKey) { + if !advance { break } } @@ -1202,6 +1196,61 @@ func (s *pebbleStore) collectScanKeys(ctx context.Context, iter *pebble.Iterator return result, nil } +func (s *pebbleStore) appendEndPrefixScanKeys( + ctx context.Context, + result [][]byte, + seen map[string]struct{}, + userKey []byte, + start []byte, + end []byte, + limit int, + ts uint64, + ok bool, +) ([][]byte, bool, error) { + if !ok { + return result, true, nil + } + if !pastScanEnd(userKey, end) { + return result, false, nil + } + result, err := s.appendVisibleProperPrefixScanKeys(ctx, result, seen, userKey, start, end, limit, ts) + if err != nil { + return nil, true, err + } + return result, true, nil +} + +func (s *pebbleStore) collectVisibleScanKey( + ctx context.Context, + iter *pebble.Iterator, + result [][]byte, + seen map[string]struct{}, + userKey []byte, + version uint64, + start []byte, + end []byte, + limit int, + ts uint64, +) ([][]byte, bool, error) { + if !s.seekToVisibleVersion(iter, userKey, version, ts) { + return result, true, nil + } + + visible, err := s.foundValueVisible(iter, ts) + if err != nil { + return nil, false, err + } + if visible { + result = appendScanKeyResult(result, seen, userKey, limit) + result, err = s.appendVisibleProperPrefixScanKeys(ctx, result, seen, userKey, start, end, limit, ts) + if err != nil { + return nil, false, err + } + } + + return result, s.skipToNextUserKey(iter, userKey), nil +} + func appendScanKeyResult(result [][]byte, seen map[string]struct{}, key []byte, limit int) [][]byte { if _, ok := seen[string(key)]; ok { return result diff --git a/store/scan_keys_test.go b/store/scan_keys_test.go index 1756beb24..4775ae127 100644 --- a/store/scan_keys_test.go +++ b/store/scan_keys_test.go @@ -178,6 +178,22 @@ func TestPebbleStoreScanKeysAtProbesExactPrefixPastBinaryExtension(t *testing.T) require.Equal(t, [][]byte{[]byte("a\x80")}, keys) } +func TestPebbleStoreScanKeysAtProbesExactPrefixBeforeBoundedEnd(t *testing.T) { + t.Parallel() + + ctx := context.Background() + st, err := NewPebbleStore(t.TempDir()) + require.NoError(t, err) + defer st.Close() + + require.NoError(t, st.PutAt(ctx, []byte("a"), []byte("va"), 10, 0)) + require.NoError(t, st.PutAt(ctx, []byte("a\x80"), []byte("vax"), 20, 0)) + + keys, err := st.ScanKeysAt(ctx, []byte("a"), []byte("a\x80"), 1, 20) + require.NoError(t, err) + require.Equal(t, [][]byte{[]byte("a")}, keys) +} + func cursorAfterTestKey(key []byte) []byte { cursor := make([]byte, len(key)+1) copy(cursor, key) From d6d8f4d152e80a4a553be320f2e4174277cc8d4d Mon Sep 17 00:00:00 2001 From: bootjp Date: Fri, 17 Jul 2026 13:51:23 +0900 Subject: [PATCH 10/10] store: stabilize key-only prefix scans --- adapter/grpc.go | 3 --- adapter/grpc_test.go | 11 +++++++++ store/lsm_store.go | 52 ++++++++++++++++++++++++++++++----------- store/scan_keys_test.go | 17 ++++++++++++++ 4 files changed, 66 insertions(+), 17 deletions(-) diff --git a/adapter/grpc.go b/adapter/grpc.go index 3e1dcfab4..589f28e3c 100644 --- a/adapter/grpc.go +++ b/adapter/grpc.go @@ -279,9 +279,6 @@ func rawKvPairs(res []*store.KVPair) []*pb.RawKVPair { func rawKeyPairs(keys [][]byte) []*pb.RawKVPair { out := make([]*pb.RawKVPair, 0, len(keys)) for _, key := range keys { - if key == nil { - continue - } out = append(out, &pb.RawKVPair{Key: key}) } return out diff --git a/adapter/grpc_test.go b/adapter/grpc_test.go index 8d5d73c9a..0ae8d321b 100644 --- a/adapter/grpc_test.go +++ b/adapter/grpc_test.go @@ -1,6 +1,7 @@ package adapter import ( + "bytes" "context" "strconv" "sync" @@ -16,6 +17,16 @@ import ( _ "google.golang.org/grpc/health" ) +func TestRawKeyPairsPreservesNilAndEmptyKeys(t *testing.T) { + t.Parallel() + + pairs := rawKeyPairs([][]byte{nil, {}, []byte("a")}) + require.Len(t, pairs, 3) + require.Nil(t, pairs[0].Key) + require.True(t, bytes.Equal([]byte{}, pairs[1].Key)) + require.Equal(t, []byte("a"), pairs[2].Key) +} + func Test_value_can_be_deleted(t *testing.T) { t.Parallel() nodes, adders, _ := createNode(t, 3) diff --git a/store/lsm_store.go b/store/lsm_store.go index 0b7d60c2c..6a41c3e86 100644 --- a/store/lsm_store.go +++ b/store/lsm_store.go @@ -1165,9 +1165,14 @@ func (s *pebbleStore) collectForwardScanKV(iter *pebble.Iterator, userKey []byte return kv, s.skipToNextUserKey(iter, userKey), nil } -func (s *pebbleStore) collectScanKeys(ctx context.Context, iter *pebble.Iterator, start, end []byte, limit int, ts uint64) ([][]byte, error) { +type pebbleIteratorProvider interface { + NewIter(*pebble.IterOptions) (*pebble.Iterator, error) +} + +func (s *pebbleStore) collectScanKeys(ctx context.Context, reader pebbleIteratorProvider, iter *pebble.Iterator, start, end []byte, limit int, ts uint64) ([][]byte, error) { result := make([][]byte, 0, boundedScanResultCapacity(limit)) seen := make(map[string]struct{}, boundedScanResultCapacity(limit)) + probedPrefixes := make(map[string]struct{}, boundedScanResultCapacity(limit)) for iter.SeekGE(encodeKey(start, math.MaxUint64)); iter.Valid() && len(result) < limit; { userKey, version, ok, err := s.nextForwardScanUserKeyContext(ctx, iter, start) @@ -1175,7 +1180,7 @@ func (s *pebbleStore) collectScanKeys(ctx context.Context, iter *pebble.Iterator return nil, err } var done bool - result, done, err = s.appendEndPrefixScanKeys(ctx, result, seen, userKey, start, end, limit, ts, ok) + result, done, err = s.appendEndPrefixScanKeys(ctx, reader, result, seen, probedPrefixes, userKey, start, end, limit, ts, ok) if err != nil { return nil, err } @@ -1184,7 +1189,7 @@ func (s *pebbleStore) collectScanKeys(ctx context.Context, iter *pebble.Iterator } var advance bool - result, advance, err = s.collectVisibleScanKey(ctx, iter, result, seen, userKey, version, start, end, limit, ts) + result, advance, err = s.collectVisibleScanKey(ctx, reader, iter, result, seen, probedPrefixes, userKey, version, start, end, limit, ts) if err != nil { return nil, err } @@ -1198,8 +1203,10 @@ func (s *pebbleStore) collectScanKeys(ctx context.Context, iter *pebble.Iterator func (s *pebbleStore) appendEndPrefixScanKeys( ctx context.Context, + reader pebbleIteratorProvider, result [][]byte, seen map[string]struct{}, + probedPrefixes map[string]struct{}, userKey []byte, start []byte, end []byte, @@ -1213,7 +1220,7 @@ func (s *pebbleStore) appendEndPrefixScanKeys( if !pastScanEnd(userKey, end) { return result, false, nil } - result, err := s.appendVisibleProperPrefixScanKeys(ctx, result, seen, userKey, start, end, limit, ts) + result, err := s.appendVisibleProperPrefixScanKeys(ctx, reader, result, seen, probedPrefixes, userKey, start, end, limit, ts) if err != nil { return nil, true, err } @@ -1222,9 +1229,11 @@ func (s *pebbleStore) appendEndPrefixScanKeys( func (s *pebbleStore) collectVisibleScanKey( ctx context.Context, + reader pebbleIteratorProvider, iter *pebble.Iterator, result [][]byte, seen map[string]struct{}, + probedPrefixes map[string]struct{}, userKey []byte, version uint64, start []byte, @@ -1232,17 +1241,22 @@ func (s *pebbleStore) collectVisibleScanKey( limit int, ts uint64, ) ([][]byte, bool, error) { - if !s.seekToVisibleVersion(iter, userKey, version, ts) { - return result, true, nil + if _, ok := seen[string(userKey)]; ok { + return result, s.skipToNextUserKey(iter, userKey), nil + } + var visible bool + var err error + if version <= ts { + visible, err = s.foundValueVisible(iter, ts) + } else { + visible, err = s.visibleExactKeyAt(ctx, reader, userKey, ts) } - - visible, err := s.foundValueVisible(iter, ts) if err != nil { return nil, false, err } if visible { result = appendScanKeyResult(result, seen, userKey, limit) - result, err = s.appendVisibleProperPrefixScanKeys(ctx, result, seen, userKey, start, end, limit, ts) + result, err = s.appendVisibleProperPrefixScanKeys(ctx, reader, result, seen, probedPrefixes, userKey, start, end, limit, ts) if err != nil { return nil, false, err } @@ -1271,8 +1285,10 @@ func appendScanKeyResult(result [][]byte, seen map[string]struct{}, key []byte, func (s *pebbleStore) appendVisibleProperPrefixScanKeys( ctx context.Context, + reader pebbleIteratorProvider, result [][]byte, seen map[string]struct{}, + probedPrefixes map[string]struct{}, userKey []byte, start []byte, end []byte, @@ -1284,10 +1300,15 @@ func (s *pebbleStore) appendVisibleProperPrefixScanKeys( if _, ok := seen[string(prefix)]; ok { continue } + probeKey := string(prefix) + if _, ok := probedPrefixes[probeKey]; ok { + continue + } + probedPrefixes[probeKey] = struct{}{} if beforeScanStart(prefix, start) || pastScanEnd(prefix, end) { continue } - visible, err := s.visibleExactKeyAt(ctx, prefix, ts) + visible, err := s.visibleExactKeyAt(ctx, reader, prefix, ts) if err != nil { return nil, err } @@ -1298,8 +1319,8 @@ func (s *pebbleStore) appendVisibleProperPrefixScanKeys( return result, nil } -func (s *pebbleStore) visibleExactKeyAt(ctx context.Context, key []byte, ts uint64) (bool, error) { - iter, err := s.db.NewIter(&pebble.IterOptions{ +func (s *pebbleStore) visibleExactKeyAt(ctx context.Context, reader pebbleIteratorProvider, key []byte, ts uint64) (bool, error) { + iter, err := reader.NewIter(&pebble.IterOptions{ LowerBound: encodeKey(key, math.MaxUint64), UpperBound: keyUpperBound(key), }) @@ -1423,7 +1444,10 @@ func (s *pebbleStore) ScanKeysAt(ctx context.Context, start []byte, end []byte, return nil, ErrReadTSCompacted } - iter, err := s.db.NewIter(&pebble.IterOptions{ + snap := s.db.NewSnapshot() + defer snap.Close() + + iter, err := snap.NewIter(&pebble.IterOptions{ LowerBound: encodeKey(start, math.MaxUint64), }) if err != nil { @@ -1431,7 +1455,7 @@ func (s *pebbleStore) ScanKeysAt(ctx context.Context, start []byte, end []byte, } defer iter.Close() - return s.collectScanKeys(ctx, iter, start, end, limit, ts) + return s.collectScanKeys(ctx, snap, iter, start, end, limit, ts) } func (s *pebbleStore) collectReverseScanResults( diff --git a/store/scan_keys_test.go b/store/scan_keys_test.go index 4775ae127..47aa435ef 100644 --- a/store/scan_keys_test.go +++ b/store/scan_keys_test.go @@ -194,6 +194,23 @@ func TestPebbleStoreScanKeysAtProbesExactPrefixBeforeBoundedEnd(t *testing.T) { require.Equal(t, [][]byte{[]byte("a")}, keys) } +func TestPebbleStoreScanKeysAtPreservesExtensionAfterInvisiblePrefix(t *testing.T) { + t.Parallel() + + ctx := context.Background() + st, err := NewPebbleStore(t.TempDir()) + require.NoError(t, err) + defer st.Close() + + highTS := ^uint64(0) - 1 + require.NoError(t, st.PutAt(ctx, []byte("a"), []byte("future"), highTS, 0)) + require.NoError(t, st.PutAt(ctx, []byte("a\x80"), []byte("visible-extension"), 10, 0)) + + keys, err := st.ScanKeysAt(ctx, []byte("a"), []byte("b"), 2, 10) + require.NoError(t, err) + require.Equal(t, [][]byte{[]byte("a\x80")}, keys) +} + func cursorAfterTestKey(key []byte) []byte { cursor := make([]byte, len(key)+1) copy(cursor, key)