From bd4a75cfd0663b893bbb22731db5b59646839245 Mon Sep 17 00:00:00 2001 From: bootjp Date: Sat, 11 Jul 2026 13:39:34 +0900 Subject: [PATCH 1/6] 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 2/6] 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 3/6] 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 4/6] 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 5/6] 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 6/6] 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)