diff --git a/kv/backup_scan.go b/kv/backup_scan.go new file mode 100644 index 000000000..9f8d514f5 --- /dev/null +++ b/kv/backup_scan.go @@ -0,0 +1,109 @@ +package kv + +import ( + "bytes" + "context" + + "github.com/bootjp/elastickv/distribution" + "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 + 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, + routes: routes, + clampToRoutes: clampToRoutes, + 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 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.scanRoutesAtSorted(ctx, s.routes, s.cursor, s.end, s.pageSize, s.ts, s.clampToRoutes) + if err != nil { + return err + } + s.page = page + s.index = 0 + if len(page) == 0 { + s.exhausted = true + return nil + } + last := page[len(page)-1] + if last == nil { + s.exhausted = true + return nil + } + s.cursor = nextScanCursor(last.Key) + return nil +} diff --git a/kv/leader_routed_store.go b/kv/leader_routed_store.go index d484431a1..dd264ac05 100644 --- a/kv/leader_routed_store.go +++ b/kv/leader_routed_store.go @@ -260,6 +260,25 @@ 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 + } + return keysFromKVs(kvs), nil +} + func (s *LeaderRoutedStore) ScanAtPhysicalLimit(ctx context.Context, start []byte, end []byte, visibleLimit, physicalLimit int, ts uint64) ([]*store.KVPair, bool, error) { return s.scanAtPhysicalLimit(ctx, start, end, visibleLimit, physicalLimit, ts, false) } 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 bb0c15035..3068af12a 100644 --- a/kv/shard_store.go +++ b/kv/shard_store.go @@ -179,24 +179,32 @@ 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 } routes, clampToRoutes := s.routesForScan(start, end) - out, err := s.scanRoutesAt(ctx, routes, start, end, limit, ts, clampToRoutes) + 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) { + 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].Key, out[j].Key) < 0 + return bytes.Compare(out[i], out[j]) < 0 }) if len(out) > limit { out = out[:limit] @@ -283,12 +291,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) @@ -308,6 +322,52 @@ 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{}) + 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) + if err != nil { + return nil, err + } + if clampToRoutes { + out = mergeAndTrimScanKeys(out, keys, limit) + if len(out) >= limit { + break + } + continue + } + out = mergeAndTrimScanKeys(out, keys, limit) + } + return out, nil +} + func (s *ShardStore) reverseScanRoutesAt( ctx context.Context, routes []distribution.Route, @@ -352,6 +412,189 @@ 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, route.GroupID) +} + +func (s *ShardStore) scanKeysRouteLocal( + ctx context.Context, + g *ShardGroup, + start []byte, + end []byte, + limit int, + ts uint64, +) ([][]byte, error) { + 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( + ctx context.Context, + g *ShardGroup, + start []byte, + end []byte, + limit int, + ts uint64, +) ([][]byte, error) { + if limit <= 0 { + return [][]byte{}, nil + } + + 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 { + keys, err := s.scanLockOnlyKeysAtLeader(ctx, g, cursor, end, ts, limit) + if err != nil { + return nil, err + } + out = mergeAndTrimScanKeys(out, keys, limit) + 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 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, + start []byte, + end []byte, + limit int, + ts uint64, + 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 + }) +} + +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 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( ctx context.Context, route distribution.Route, @@ -389,6 +632,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 { @@ -414,6 +661,109 @@ 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 len(out) >= limit { + break + } + 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 + } + + raw, err := s.proxyRawScanAt(ctx, g, start, end, limit, ts, false, route.GroupID) + if err != nil { + return scanRoutePage{}, err + } + return scanRoutePage{ + kvs: filterTxnInternalKVs(raw), + advanceKey: lastKVKey(raw), + full: len(raw) >= limit, + }, nil +} + type physicalLimitedStore interface { ScanAtPhysicalLimit(ctx context.Context, start []byte, end []byte, visibleLimit, physicalLimit int, ts uint64) ([]*store.KVPair, bool, error) ReverseScanAtPhysicalLimit(ctx context.Context, start []byte, end []byte, visibleLimit, physicalLimit int, ts uint64) ([]*store.KVPair, bool, error) @@ -578,26 +928,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) { @@ -621,6 +973,34 @@ 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...) + 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] +} + func mergeAndTrimReverseScanResults(out []*store.KVPair, kvs []*store.KVPair, limit int) []*store.KVPair { if len(kvs) == 0 { return out @@ -639,7 +1019,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++ @@ -647,6 +1027,52 @@ 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 key == nil { + continue + } + kvs = append(kvs, &store.KVPair{Key: key}) + } + return kvs +} + +func keysFromKVs(kvs []*store.KVPair) [][]byte { + keys := make([][]byte, 0, len(kvs)) + for _, kvp := range kvs { + if kvp == nil || kvp.Key == nil { + continue + } + keys = append(keys, kvp.Key) + } + 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 + } + out := make([][]byte, 0, len(keys)) + for _, key := range keys { + if key == nil || 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 @@ -1007,7 +1433,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 4f0da38b7..cc1daf8e7 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() @@ -144,6 +189,406 @@ 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 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 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() + + 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 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 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() + + 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() + + 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 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() + + 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)) + require.NoError(t, st.PutAt(ctx, []byte("a"), []byte("later"), 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(""), 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) + 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/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() diff --git a/store/lsm_store.go b/store/lsm_store.go index d818b7bde..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" @@ -979,6 +980,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 @@ -993,21 +1008,30 @@ 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 { 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() @@ -1029,6 +1053,21 @@ func nextScannableUserKey(iter *pebble.Iterator) ([]byte, uint64, bool) { return nil, 0, false } +func (s *pebbleStore) nextForwardScanUserKeyContext(ctx context.Context, iter *pebble.Iterator, start []byte) ([]byte, uint64, bool, error) { + for { + userKey, version, ok, err := nextScannableUserKeyContext(ctx, iter) + if err != nil || !ok { + return nil, 0, false, err + } + if !beforeScanStart(userKey, start) { + return userKey, version, true, nil + } + if !s.skipToNextUserKey(iter, userKey) { + return nil, 0, false, nil + } + } +} + func nextScannableUserKeyContext(ctx context.Context, iter *pebble.Iterator) ([]byte, uint64, bool, error) { for iter.Valid() { if err := ctx.Err(); err != nil { @@ -1084,7 +1123,7 @@ func (s *pebbleStore) collectScanResultsWithPhysicalLimit(ctx context.Context, i visited := 0 for iter.SeekGE(encodeKey(start, math.MaxUint64)); iter.Valid() && len(result) < limit; { - userKey, version, ok, err := nextScannableUserKeyContext(ctx, iter) + userKey, version, ok, err := s.nextForwardScanUserKeyContext(ctx, iter, start) if err != nil { return nil, false, err } @@ -1126,6 +1165,101 @@ 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) { + 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) + if err != nil { + return nil, err + } + if forwardScanDone(userKey, end, ok) { + break + } + + if !s.seekToVisibleVersion(iter, userKey, version, ts) { + continue + } + + visible, err := s.foundValueVisible(iter, 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) { + break + } + } + + 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 @@ -1185,6 +1319,37 @@ func (s *pebbleStore) ScanAtPhysicalLimit(ctx context.Context, start []byte, end return s.collectScanResultsWithPhysicalLimit(ctx, iter, start, end, visibleLimit, physicalLimit, ts) } +func (s *pebbleStore) ScanKeysAt(ctx context.Context, start []byte, end []byte, limit int, ts uint64) ([][]byte, error) { + if limit <= 0 { + return [][]byte{}, nil + } + if err := ctx.Err(); err != nil { + return nil, errors.WithStack(err) + } + + // 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 err := ctx.Err(); err != nil { + return nil, errors.WithStack(err) + } + + 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(ctx, 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 abbda3240..158847d16 100644 --- a/store/mvcc_store.go +++ b/store/mvcc_store.go @@ -341,6 +341,28 @@ func collectScanResultsWithPhysicalLimit(ctx context.Context, it *treemap.Iterat return result, false, nil } +func (s *mvccStore) collectScanKeys(ctx context.Context, it *treemap.Iterator, end []byte, limit, capHint int, ts uint64) ([][]byte, error) { + result := make([][]byte, 0, capHint) + for ok := true; ok && len(result) < limit; ok = it.Next() { + if err := ctx.Err(); err != nil { + return nil, errors.WithStack(err) + } + 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, nil +} + func (s *mvccStore) ScanAt(ctx context.Context, start []byte, end []byte, limit int, ts uint64) ([]*KVPair, error) { if err := ctx.Err(); err != nil { return nil, errors.WithStack(err) @@ -390,6 +412,31 @@ func (s *mvccStore) ScanAtPhysicalLimit(ctx context.Context, start []byte, end [ return collectScanResultsWithPhysicalLimit(ctx, &it, end, visibleLimit, physicalLimit, capHint, ts) } +func (s *mvccStore) ScanKeysAt(ctx context.Context, start []byte, end []byte, limit int, ts uint64) ([][]byte, error) { + if err := ctx.Err(); err != nil { + return nil, errors.WithStack(err) + } + s.mtx.RLock() + defer s.mtx.RUnlock() + if err := ctx.Err(); err != nil { + return nil, errors.WithStack(err) + } + 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(ctx, &it, end, limit, capHint, ts) +} + // 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..8ba7c5479 --- /dev/null +++ b/store/scan_keys_test.go @@ -0,0 +1,176 @@ +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 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 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 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) + return cursor +} + +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.