diff --git a/adapter/grpc.go b/adapter/grpc.go index 712efd911..589f28e3c 100644 --- a/adapter/grpc.go +++ b/adapter/grpc.go @@ -44,6 +44,10 @@ type rawGroupScanner interface { ScanGroupAt(ctx context.Context, groupID uint64, start []byte, end []byte, limit int, ts uint64) ([]*store.KVPair, error) } +type rawGroupKeyScanner interface { + ScanGroupKeysAt(ctx context.Context, groupID uint64, start []byte, end []byte, limit int, ts uint64) ([][]byte, error) +} + func WithCloseStore() GRPCServerOption { return func(s *GRPCServer) { s.closeStore = true @@ -159,14 +163,32 @@ func (r *GRPCServer) RawScanAt(ctx context.Context, req *pb.RawScanAtRequest) (* readTS = globalSnapshotTS(ctx, r.clock(), r.store) } + if req.GetKeysOnly() { + keys, err := r.rawScanKeysAt(ctx, req, limit, readTS) + if err != nil { + return rawScanErrorResponse(err) + } + return &pb.RawScanAtResponse{Kv: rawKeyPairs(keys)}, nil + } + + res, err := r.rawScanValuesAt(ctx, req, limit, readTS) + if err != nil { + return rawScanErrorResponse(err) + } + + return &pb.RawScanAtResponse{Kv: rawKvPairs(res)}, nil +} + +func (r *GRPCServer) rawScanValuesAt(ctx context.Context, req *pb.RawScanAtRequest, limit int, readTS uint64) ([]*store.KVPair, error) { var res []*store.KVPair + var err error if groupID := req.GetGroupId(); groupID != 0 { if req.GetReverse() { - return &pb.RawScanAtResponse{Kv: nil}, errors.WithStack(status.Error(codes.InvalidArgument, "raw scan with explicit group does not support reverse scans")) + return nil, errors.WithStack(status.Error(codes.InvalidArgument, "raw scan with explicit group does not support reverse scans")) } groupScanner, ok := r.store.(rawGroupScanner) if !ok { - return &pb.RawScanAtResponse{Kv: nil}, errors.WithStack(status.Error(codes.FailedPrecondition, "raw scan with explicit group requires a group-aware store")) + return nil, errors.WithStack(status.Error(codes.FailedPrecondition, "raw scan with explicit group requires a group-aware store")) } res, err = groupScanner.ScanGroupAt(ctx, groupID, req.StartKey, req.EndKey, limit, readTS) } else if req.GetReverse() { @@ -175,13 +197,42 @@ func (r *GRPCServer) RawScanAt(ctx context.Context, req *pb.RawScanAtRequest) (* res, err = r.store.ScanAt(ctx, req.StartKey, req.EndKey, limit, readTS) } if err != nil { - if errors.Is(err, store.ErrReadTSCompacted) { - return &pb.RawScanAtResponse{Kv: nil}, errors.WithStack(status.Error(codes.FailedPrecondition, store.ErrReadTSCompacted.Error())) + return nil, errors.WithStack(err) + } + return res, nil +} + +func (r *GRPCServer) rawScanKeysAt(ctx context.Context, req *pb.RawScanAtRequest, limit int, readTS uint64) ([][]byte, error) { + if groupID := req.GetGroupId(); groupID != 0 { + if req.GetReverse() { + return nil, errors.WithStack(status.Error(codes.InvalidArgument, "raw scan with explicit group does not support reverse scans")) + } + groupScanner, ok := r.store.(rawGroupKeyScanner) + if !ok { + return nil, errors.WithStack(status.Error(codes.FailedPrecondition, "raw key scan with explicit group requires a group-aware store")) } - return &pb.RawScanAtResponse{Kv: nil}, errors.WithStack(err) + keys, err := groupScanner.ScanGroupKeysAt(ctx, groupID, req.StartKey, req.EndKey, limit, readTS) + if err != nil { + return nil, errors.WithStack(err) + } + return keys, nil + } + if req.GetReverse() { + kvs, err := r.store.ReverseScanAt(ctx, req.StartKey, req.EndKey, limit, readTS) + if err != nil { + return nil, errors.WithStack(err) + } + return storeKeysFromKVPairs(kvs), nil } + keys, err := r.store.ScanKeysAt(ctx, req.StartKey, req.EndKey, limit, readTS) + return keys, errors.WithStack(err) +} - return &pb.RawScanAtResponse{Kv: rawKvPairs(res)}, nil +func rawScanErrorResponse(err error) (*pb.RawScanAtResponse, error) { + if errors.Is(err, store.ErrReadTSCompacted) { + return &pb.RawScanAtResponse{Kv: nil}, errors.WithStack(status.Error(codes.FailedPrecondition, store.ErrReadTSCompacted.Error())) + } + return &pb.RawScanAtResponse{Kv: nil}, errors.WithStack(err) } func rawScanLimit(limit64 int64) (int, error) { @@ -225,6 +276,25 @@ func rawKvPairs(res []*store.KVPair) []*pb.RawKVPair { return out } +func rawKeyPairs(keys [][]byte) []*pb.RawKVPair { + out := make([]*pb.RawKVPair, 0, len(keys)) + for _, key := range keys { + out = append(out, &pb.RawKVPair{Key: key}) + } + return out +} + +func storeKeysFromKVPairs(kvs []*store.KVPair) [][]byte { + keys := make([][]byte, 0, len(kvs)) + for _, kvp := range kvs { + if kvp == nil { + continue + } + keys = append(keys, kvp.Key) + } + return keys +} + func (r *GRPCServer) RawPut(ctx context.Context, req *pb.RawPutRequest) (*pb.RawPutResponse, error) { m, err := r.grpcTranscoder.RawPutToRequest(req) if err != nil { diff --git a/adapter/grpc_test.go b/adapter/grpc_test.go index fd1393cd5..0ae8d321b 100644 --- a/adapter/grpc_test.go +++ b/adapter/grpc_test.go @@ -1,6 +1,7 @@ package adapter import ( + "bytes" "context" "strconv" "sync" @@ -16,6 +17,16 @@ import ( _ "google.golang.org/grpc/health" ) +func TestRawKeyPairsPreservesNilAndEmptyKeys(t *testing.T) { + t.Parallel() + + pairs := rawKeyPairs([][]byte{nil, {}, []byte("a")}) + require.Len(t, pairs, 3) + require.Nil(t, pairs[0].Key) + require.True(t, bytes.Equal([]byte{}, pairs[1].Key)) + require.Equal(t, []byte("a"), pairs[2].Key) +} + func Test_value_can_be_deleted(t *testing.T) { t.Parallel() nodes, adders, _ := createNode(t, 3) @@ -142,6 +153,7 @@ type recordingRawGroupStore struct { scanGroupID uint64 scanStart []byte scanEnd []byte + keyScanGroup bool fallbackGet bool fallbackScan bool } @@ -169,6 +181,14 @@ func (s *recordingRawGroupStore) ScanGroupAt(ctx context.Context, groupID uint64 return s.MVCCStore.ScanAt(ctx, start, end, limit, ts) } +func (s *recordingRawGroupStore) ScanGroupKeysAt(ctx context.Context, groupID uint64, start []byte, end []byte, limit int, ts uint64) ([][]byte, error) { + s.scanGroupID = groupID + s.scanStart = append([]byte(nil), start...) + s.scanEnd = append([]byte(nil), end...) + s.keyScanGroup = true + return s.ScanKeysAt(ctx, start, end, limit, ts) +} + func TestGRPCServer_RawGet_UsesExplicitGroup(t *testing.T) { t.Parallel() @@ -209,6 +229,33 @@ func TestGRPCServer_RawScanAt_UsesExplicitGroup(t *testing.T) { require.Equal(t, []byte("z"), st.scanEnd) } +func TestGRPCServer_RawScanAt_KeysOnlyUsesExplicitGroup(t *testing.T) { + t.Parallel() + + ctx := context.Background() + st := &recordingRawGroupStore{MVCCStore: store.NewMVCCStore()} + require.NoError(t, st.PutAt(ctx, []byte("a"), []byte("large-value"), 9, 0)) + s := NewGRPCServer(st, nil) + + resp, err := s.RawScanAt(ctx, &pb.RawScanAtRequest{ + StartKey: []byte("a"), + EndKey: []byte("z"), + Limit: 10, + Ts: 9, + GroupId: 42, + KeysOnly: true, + }) + require.NoError(t, err) + require.Len(t, resp.GetKv(), 1) + require.Equal(t, []byte("a"), resp.GetKv()[0].GetKey()) + require.Empty(t, resp.GetKv()[0].GetValue()) + require.True(t, st.keyScanGroup) + require.False(t, st.fallbackScan) + require.Equal(t, uint64(42), st.scanGroupID) + require.Equal(t, []byte("a"), st.scanStart) + require.Equal(t, []byte("z"), st.scanEnd) +} + func TestGRPCServer_Scan_RejectsOversizedLimit(t *testing.T) { t.Parallel() diff --git a/kv/backup_scan.go b/kv/backup_scan.go new file mode 100644 index 000000000..8c132eb78 --- /dev/null +++ b/kv/backup_scan.go @@ -0,0 +1,137 @@ +package kv + +import ( + "bytes" + "context" + + "github.com/bootjp/elastickv/distribution" + "github.com/bootjp/elastickv/store" + "github.com/cockroachdb/errors" +) + +const defaultBackupScanPageSize = 1024 + +// BackupScanner pages through ShardStore.ScanKeysAt without holding store locks +// across pages, then materializes each key at the pinned read timestamp. +type BackupScanner interface { + Next(ctx context.Context) (*store.KVPair, bool, error) + Close() error +} + +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 + } + keys, err := s.store.scanKeyRoutesAt(ctx, s.routes, s.cursor, s.end, s.pageSize, s.ts, s.clampToRoutes) + if err != nil { + return err + } + s.page = s.page[:0] + for _, key := range keys { + val, err := s.getAtCapturedRoute(ctx, key) + if errors.Is(err, store.ErrKeyNotFound) { + continue + } + if err != nil { + return err + } + s.page = append(s.page, &store.KVPair{Key: bytes.Clone(key), Value: bytes.Clone(val)}) + } + s.index = 0 + if len(keys) == 0 { + s.exhausted = true + return nil + } + last := lastScanKey(keys) + if last == nil { + s.exhausted = true + return nil + } + s.cursor = nextScanCursor(last) + return nil +} + +func (s *backupScanner) getAtCapturedRoute(ctx context.Context, key []byte) ([]byte, error) { + normalized := routeKey(key) + for _, route := range s.routes { + if routeContainsKey(route, normalized) { + return s.store.getRouteAt(ctx, route, key, s.ts) + } + } + return nil, store.ErrKeyNotFound +} + +func routeContainsKey(route distribution.Route, key []byte) bool { + if bytes.Compare(key, route.Start) < 0 { + return false + } + return route.End == nil || bytes.Compare(key, route.End) < 0 +} diff --git a/kv/leader_routed_store.go b/kv/leader_routed_store.go index d484431a1..47e4bc6b1 100644 --- a/kv/leader_routed_store.go +++ b/kv/leader_routed_store.go @@ -153,6 +153,45 @@ func (s *LeaderRoutedStore) proxyRawScanAt( return out, nil } +func (s *LeaderRoutedStore) proxyRawScanKeysAt( + ctx context.Context, + start []byte, + end []byte, + limit int, + ts uint64, +) ([][]byte, error) { + addr := s.leaderAddrForKey(start) + if addr == "" { + return nil, errors.WithStack(ErrLeaderNotFound) + } + + conn, err := s.connCache.ConnFor(addr) + if err != nil { + return nil, err + } + + cli := pb.NewRawKVClient(conn) + resp, err := cli.RawScanAt(ctx, &pb.RawScanAtRequest{ + StartKey: start, + EndKey: end, + Limit: int64(limit), + Ts: ts, + KeysOnly: true, + }) + if err != nil { + return nil, errors.WithStack(err) + } + + out := make([][]byte, 0, len(resp.Kv)) + for _, kvp := range resp.Kv { + if kvp == nil { + continue + } + out = append(out, bytes.Clone(kvp.Key)) + } + return out, nil +} + func (s *LeaderRoutedStore) GetAt(ctx context.Context, key []byte, ts uint64) ([]byte, error) { if s == nil || s.local == nil { return nil, store.ErrKeyNotFound @@ -260,6 +299,21 @@ 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) + } + return s.proxyRawScanKeysAt(ctx, start, end, limit, ts) +} + func (s *LeaderRoutedStore) ScanAtPhysicalLimit(ctx context.Context, start []byte, end []byte, visibleLimit, physicalLimit int, ts uint64) ([]*store.KVPair, bool, error) { 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..22ddc9535 100644 --- a/kv/leader_routed_store_test.go +++ b/kv/leader_routed_store_test.go @@ -85,6 +85,9 @@ type fakeRawKVServer struct { scanCalls int latestCalls int + lastScanGroupID uint64 + lastScanKeysOnly bool + getResp *pb.RawGetResponse scanResp *pb.RawScanAtResponse latestResp *pb.RawLatestCommitTSResponse @@ -100,10 +103,12 @@ 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() + f.lastScanKeysOnly = req.GetKeysOnly() if f.scanResp != nil { return f.scanResp, nil } diff --git a/kv/shard_store.go b/kv/shard_store.go index bb0c15035..581913379 100644 --- a/kv/shard_store.go +++ b/kv/shard_store.go @@ -62,6 +62,18 @@ func (s *ShardStore) GetGroupAt(ctx context.Context, groupID uint64, key []byte, return nil, store.ErrKeyNotFound } + return s.getGroupAt(ctx, g, key, ts, groupID) +} + +func (s *ShardStore) getRouteAt(ctx context.Context, route distribution.Route, key []byte, ts uint64) ([]byte, error) { + g, ok := s.groupForID(route.GroupID) + if !ok || g.Store == nil { + return nil, store.ErrKeyNotFound + } + return s.getGroupAt(ctx, g, key, ts, route.GroupID) +} + +func (s *ShardStore) getGroupAt(ctx context.Context, g *ShardGroup, key []byte, ts uint64, groupID uint64) ([]byte, error) { if engineForGroup(g) == nil { return s.localGetAt(ctx, g, key, ts) } @@ -179,24 +191,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] @@ -228,6 +248,15 @@ func (s *ShardStore) ScanGroupAt(ctx context.Context, groupID uint64, start []by return s.scanRouteAtDirection(ctx, distribution.Route{GroupID: groupID}, start, end, limit, ts, false, true) } +// ScanGroupKeysAt scans keys on the explicitly selected Raft group without +// materializing values over proxy links. +func (s *ShardStore) ScanGroupKeysAt(ctx context.Context, groupID uint64, start []byte, end []byte, limit int, ts uint64) ([][]byte, error) { + if limit <= 0 { + return [][]byte{}, nil + } + return s.scanKeyRouteAt(ctx, distribution.Route{GroupID: groupID}, start, end, limit, ts) +} + func (s *ShardStore) ReverseScanAt(ctx context.Context, start []byte, end []byte, limit int, ts uint64) ([]*store.KVPair, error) { if limit <= 0 { return []*store.KVPair{}, nil @@ -283,12 +312,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 +343,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 +433,185 @@ 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) { + return s.proxyRawScanKeysAt(ctx, g, cursor, end, pageLimit, ts, groupID) + }) +} + +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 +649,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 +678,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) @@ -561,43 +928,51 @@ func (s *ShardStore) scanRouteAtLeader( } func scanLockBoundsForKVs(kvs []*store.KVPair, scanStart []byte, scanEnd []byte, limit int) ([]byte, []byte) { - if countNonInternalKVs(kvs) < limit { + if len(kvs) < limit { return scanStart, scanEnd } _, lastUserKey, ok := observedScanUserBounds(kvs) - if !ok { - return scanStart, scanEnd + if ok { + return scanStart, boundScanEnd(scanEnd, nextScanCursor(lastUserKey)) } - bound := nextScanCursor(lastUserKey) - if scanEnd == nil || bytes.Compare(bound, scanEnd) < 0 { - scanEnd = bound + if lastKey := lastKVKey(kvs); lastKey != nil { + return scanStart, boundScanEnd(scanEnd, nextScanCursor(lastKey)) } return scanStart, scanEnd } +func boundScanEnd(scanEnd []byte, bound []byte) []byte { + if scanEnd == nil || bytes.Compare(bound, scanEnd) < 0 { + return bound + } + return scanEnd +} + func observedScanUserBounds(kvs []*store.KVPair) ([]byte, []byte, bool) { var minKey []byte var maxKey []byte + 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 +996,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 @@ -636,15 +1039,50 @@ func mergeAndTrimReverseScanResults(out []*store.KVPair, kvs []*store.KVPair, li return out[:limit] } -func countNonInternalKVs(kvs []*store.KVPair) int { - count := 0 +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 || isTxnInternalKey(kvp.Key) { + 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 } - count++ + out = append(out, key) } - return count + return out } func (s *ShardStore) groupForID(groupID uint64) (*ShardGroup, bool) { @@ -1007,7 +1445,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) @@ -1716,4 +2154,52 @@ func (s *ShardStore) proxyRawScanAt( return out, nil } +func (s *ShardStore) proxyRawScanKeysAt( + ctx context.Context, + g *ShardGroup, + start []byte, + end []byte, + limit int, + ts uint64, + groupID uint64, +) ([][]byte, error) { + engine := engineForGroup(g) + if engine == nil { + return nil, store.ErrNotSupported + } + addr := leaderAddrFromEngine(engine) + if addr == "" { + return nil, errors.WithStack(ErrLeaderNotFound) + } + + conn, err := s.connCache.ConnFor(addr) + if err != nil { + return nil, err + } + + ctx, cancel := context.WithTimeout(ctx, proxyForwardTimeout) + defer cancel() + cli := pb.NewRawKVClient(conn) + resp, err := cli.RawScanAt(ctx, &pb.RawScanAtRequest{ + StartKey: start, + EndKey: end, + Limit: int64(limit), + Ts: ts, + GroupId: groupID, + KeysOnly: true, + }) + if err != nil { + return nil, errors.WithStack(err) + } + + out := make([][]byte, 0, len(resp.Kv)) + for _, kvp := range resp.Kv { + if kvp == nil { + continue + } + out = append(out, bytes.Clone(kvp.Key)) + } + return out, nil +} + var _ store.MVCCStore = (*ShardStore)(nil) diff --git a/kv/shard_store_test.go b/kv/shard_store_test.go index 4f0da38b7..f34999ab0 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,481 @@ 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) + require.True(t, fake.lastScanKeysOnly) +} + +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 TestBackupScannerMaterializesFromCapturedRoute(t *testing.T) { + t.Parallel() + + ctx := context.Background() + engine := distribution.NewEngine() + engine.UpdateRoute([]byte(""), nil, 1) + + groups := map[uint64]*ShardGroup{ + 1: {Store: store.NewMVCCStore()}, + 2: {Store: store.NewMVCCStore()}, + } + st := NewShardStore(engine, groups) + require.NoError(t, groups[1].Store.PutAt(ctx, []byte("a"), []byte("old-owner"), 1, 0)) + + sc := st.NewBackupScanner([]byte(""), nil, ^uint64(0), 1) + defer sc.Close() + + require.NoError(t, engine.ApplySnapshot(distribution.CatalogSnapshot{ + Version: 1, + Routes: []distribution.RouteDescriptor{ + {RouteID: 1, Start: []byte(""), GroupID: 2, State: distribution.RouteStateActive}, + }, + })) + + 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("old-owner"), kvp.Value) + + kvp, ok, err = sc.Next(ctx) + require.NoError(t, err) + require.False(t, ok) + require.Nil(t, kvp) +} + +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 TestBackupScannerPagesPebbleKeysInLogicalOrder(t *testing.T) { + t.Parallel() + + ctx := context.Background() + engine := distribution.NewEngine() + engine.UpdateRoute([]byte(""), nil, 1) + pebbleStore, err := store.NewPebbleStore(t.TempDir()) + require.NoError(t, err) + defer pebbleStore.Close() + + groups := map[uint64]*ShardGroup{ + 1: {Store: pebbleStore}, + } + st := NewShardStore(engine, groups) + require.NoError(t, pebbleStore.PutAt(ctx, []byte("a"), []byte("va"), 10, 0)) + require.NoError(t, pebbleStore.PutAt(ctx, []byte("a\x80"), []byte("vax"), 20, 0)) + + sc := st.NewBackupScanner(nil, nil, 20, 1) + defer sc.Close() + + kvp, ok, err := sc.Next(ctx) + require.NoError(t, err) + require.True(t, ok) + require.Equal(t, []byte("a"), kvp.Key) + require.Equal(t, []byte("va"), kvp.Value) + + kvp, ok, err = sc.Next(ctx) + require.NoError(t, err) + require.True(t, ok) + require.Equal(t, []byte("a\x80"), kvp.Key) + require.Equal(t, []byte("vax"), kvp.Value) + + kvp, ok, err = sc.Next(ctx) + require.NoError(t, err) + require.False(t, ok) + require.Nil(t, kvp) +} + func TestShardStoreScanAt_RoutesS3ManifestScansByLogicalObjectKey(t *testing.T) { t.Parallel() @@ -409,3 +929,16 @@ func TestScanLockBoundsForKVs_EmptyUsesOriginalRange(t *testing.T) { require.Equal(t, []byte("a"), lockStart) require.Equal(t, []byte("z"), lockEnd) } + +func TestScanLockBoundsForKVs_FullInternalPageUsesRawPageBound(t *testing.T) { + t.Parallel() + + internalKey := txnCommitKey([]byte("primary"), 10) + kvs := []*store.KVPair{ + {Key: internalKey, Value: []byte("commit")}, + } + + lockStart, lockEnd := scanLockBoundsForKVs(kvs, []byte(""), nil, 1) + require.Equal(t, []byte(""), lockStart) + require.Equal(t, nextScanCursor(internalKey), lockEnd) +} diff --git a/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/proto/service.pb.go b/proto/service.pb.go index a869e0c1d..00000ae2f 100644 --- a/proto/service.pb.go +++ b/proto/service.pb.go @@ -499,7 +499,8 @@ type RawScanAtRequest struct { Limit int64 `protobuf:"varint,3,opt,name=limit,proto3" json:"limit,omitempty"` // validated against host int size; large values may be rejected Ts uint64 `protobuf:"varint,4,opt,name=ts,proto3" json:"ts,omitempty"` // optional read timestamp; if zero, server uses current HLC Reverse bool `protobuf:"varint,5,opt,name=reverse,proto3" json:"reverse,omitempty"` - GroupId uint64 `protobuf:"varint,6,opt,name=group_id,json=groupId,proto3" json:"group_id,omitempty"` // optional explicit Raft group for non-range-owned keyspaces + GroupId uint64 `protobuf:"varint,6,opt,name=group_id,json=groupId,proto3" json:"group_id,omitempty"` // optional explicit Raft group for non-range-owned keyspaces + KeysOnly bool `protobuf:"varint,7,opt,name=keys_only,json=keysOnly,proto3" json:"keys_only,omitempty"` // when true, response kv entries omit values unknownFields protoimpl.UnknownFields sizeCache protoimpl.SizeCache } @@ -576,6 +577,13 @@ func (x *RawScanAtRequest) GetGroupId() uint64 { return 0 } +func (x *RawScanAtRequest) GetKeysOnly() bool { + if x != nil { + return x.KeysOnly + } + return false +} + type RawKVPair struct { state protoimpl.MessageState `protogen:"open.v1"` Key []byte `protobuf:"bytes,1,opt,name=key,proto3" json:"key,omitempty"` @@ -2279,14 +2287,15 @@ const file_service_proto_rawDesc = "" + "\x03key\x18\x01 \x01(\fR\x03key\"C\n" + "\x19RawLatestCommitTSResponse\x12\x0e\n" + "\x02ts\x18\x01 \x01(\x04R\x02ts\x12\x16\n" + - "\x06exists\x18\x02 \x01(\bR\x06exists\"\xa3\x01\n" + + "\x06exists\x18\x02 \x01(\bR\x06exists\"\xc0\x01\n" + "\x10RawScanAtRequest\x12\x1b\n" + "\tstart_key\x18\x01 \x01(\fR\bstartKey\x12\x17\n" + "\aend_key\x18\x02 \x01(\fR\x06endKey\x12\x14\n" + "\x05limit\x18\x03 \x01(\x03R\x05limit\x12\x0e\n" + "\x02ts\x18\x04 \x01(\x04R\x02ts\x12\x18\n" + "\areverse\x18\x05 \x01(\bR\areverse\x12\x19\n" + - "\bgroup_id\x18\x06 \x01(\x04R\agroupId\"3\n" + + "\bgroup_id\x18\x06 \x01(\x04R\agroupId\x12\x1b\n" + + "\tkeys_only\x18\a \x01(\bR\bkeysOnly\"3\n" + "\tRawKVPair\x12\x10\n" + "\x03key\x18\x01 \x01(\fR\x03key\x12\x14\n" + "\x05value\x18\x02 \x01(\fR\x05value\"/\n" + diff --git a/proto/service.proto b/proto/service.proto index b1ac9105f..68d12d5f7 100644 --- a/proto/service.proto +++ b/proto/service.proto @@ -78,6 +78,7 @@ message RawScanAtRequest { uint64 ts = 4; // optional read timestamp; if zero, server uses current HLC bool reverse = 5; uint64 group_id = 6; // optional explicit Raft group for non-range-owned keyspaces + bool keys_only = 7; // when true, response kv entries omit values } message RawKVPair { diff --git a/store/lsm_store.go b/store/lsm_store.go index d818b7bde..078c042fe 100644 --- a/store/lsm_store.go +++ b/store/lsm_store.go @@ -12,6 +12,7 @@ import ( "math" "os" "path/filepath" + "sort" "strconv" "strings" "sync" @@ -475,21 +476,10 @@ func writeValueHeaderBytes(dst []byte, tombstone bool, expireAt uint64, encState } func decodeValue(data []byte) (storedValue, error) { - if len(data) < valueHeaderSize { - return storedValue{}, errors.New("invalid value length") - } - flags := data[0] - if flags&encStateReservedMask != 0 { - return storedValue{}, errors.Wrapf(ErrEncryptedValueReservedState, - "value header byte = %#08b", flags) - } - encState := (flags & encStateMask) >> encStateShift - if encState != encStateCleartext && encState != encStateEncrypted { - return storedValue{}, errors.Wrapf(ErrEncryptedValueReservedState, - "encryption_state=%#x is reserved", encState) + tombstone, expireAt, encState, err := decodeValueHeader(data) + if err != nil { + return storedValue{}, err } - tombstone := (flags & tombstoneMask) != 0 - expireAt := binary.LittleEndian.Uint64(data[1:]) val := make([]byte, len(data)-valueHeaderSize) copy(val, data[valueHeaderSize:]) @@ -501,6 +491,23 @@ func decodeValue(data []byte) (storedValue, error) { }, nil } +func decodeValueHeader(data []byte) (bool, uint64, byte, error) { + if len(data) < valueHeaderSize { + return false, 0, 0, errors.New("invalid value length") + } + flags := data[0] + if flags&encStateReservedMask != 0 { + return false, 0, 0, errors.Wrapf(ErrEncryptedValueReservedState, + "value header byte = %#08b", flags) + } + encState := (flags & encStateMask) >> encStateShift + if encState != encStateCleartext && encState != encStateEncrypted { + return false, 0, 0, errors.Wrapf(ErrEncryptedValueReservedState, + "encryption_state=%#x is reserved", encState) + } + return flags&tombstoneMask != 0, binary.LittleEndian.Uint64(data[1:]), encState, nil +} + type pebbleUint64Getter interface { Get(key []byte) ([]byte, io.Closer, error) } @@ -979,6 +986,28 @@ 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() + tombstone, expireAt, encState, err := decodeValueHeader(valBytes) + if err != nil { + return false, err + } + + // Keep key-only scans on the same authenticated visibility path as ScanAt. + if encState == encStateEncrypted { + sv := storedValue{ + Value: valBytes[valueHeaderSize:], + Tombstone: tombstone, + ExpireAt: expireAt, + EncState: encState, + } + if _, err := s.decryptForKey(iter.Key(), sv, sv.Value); err != nil { + return false, err + } + } + return !tombstone && (expireAt == 0 || expireAt > ts), nil +} + func (s *pebbleStore) seekToVisibleVersion(iter *pebble.Iterator, userKey []byte, currentVersion, ts uint64) bool { if currentVersion <= ts { return true @@ -993,21 +1022,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 +1067,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 +1137,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 +1179,207 @@ func (s *pebbleStore) collectForwardScanKV(iter *pebble.Iterator, userKey []byte return kv, s.skipToNextUserKey(iter, userKey), nil } +type pebbleIteratorProvider interface { + NewIter(*pebble.IterOptions) (*pebble.Iterator, error) +} + +func (s *pebbleStore) collectScanKeys(ctx context.Context, reader pebbleIteratorProvider, iter *pebble.Iterator, start, end []byte, limit int, ts uint64) ([][]byte, error) { + result := make([][]byte, 0, boundedScanResultCapacity(limit)) + seen := make(map[string]struct{}, boundedScanResultCapacity(limit)) + probedPrefixes := make(map[string]struct{}, boundedScanResultCapacity(limit)) + + for iter.SeekGE(encodeKey(start, math.MaxUint64)); iter.Valid() && len(result) < limit; { + userKey, version, ok, err := s.nextForwardScanUserKeyContext(ctx, iter, start) + if err != nil { + return nil, err + } + var done bool + result, done, err = s.appendEndPrefixScanKeys(ctx, reader, result, seen, probedPrefixes, userKey, start, end, limit, ts, ok) + if err != nil { + return nil, err + } + if done { + break + } + + var advance bool + result, advance, err = s.collectVisibleScanKey(ctx, reader, iter, result, seen, probedPrefixes, userKey, version, start, end, limit, ts) + if err != nil { + return nil, err + } + if !advance { + break + } + } + + return result, nil +} + +func (s *pebbleStore) appendEndPrefixScanKeys( + ctx context.Context, + reader pebbleIteratorProvider, + result [][]byte, + seen map[string]struct{}, + probedPrefixes map[string]struct{}, + userKey []byte, + start []byte, + end []byte, + limit int, + ts uint64, + ok bool, +) ([][]byte, bool, error) { + if !ok { + return result, true, nil + } + if !pastScanEnd(userKey, end) { + return result, false, nil + } + result, err := s.appendVisibleProperPrefixScanKeys(ctx, reader, result, seen, probedPrefixes, userKey, start, end, limit, ts) + if err != nil { + return nil, true, err + } + return result, true, nil +} + +func (s *pebbleStore) collectVisibleScanKey( + ctx context.Context, + reader pebbleIteratorProvider, + iter *pebble.Iterator, + result [][]byte, + seen map[string]struct{}, + probedPrefixes map[string]struct{}, + userKey []byte, + version uint64, + start []byte, + end []byte, + limit int, + ts uint64, +) ([][]byte, bool, error) { + if _, ok := seen[string(userKey)]; ok { + return result, s.skipToNextUserKey(iter, userKey), nil + } + var visible bool + var err error + if version <= ts { + visible, err = s.foundValueVisible(iter, ts) + } else { + visible, err = s.visibleExactKeyAt(ctx, reader, userKey, ts) + } + if err != nil { + return nil, false, err + } + if visible { + result = appendScanKeyResult(result, seen, userKey, limit) + result, err = s.appendVisibleProperPrefixScanKeys(ctx, reader, result, seen, probedPrefixes, userKey, start, end, limit, ts) + if err != nil { + return nil, false, err + } + } + + return result, s.skipToNextUserKey(iter, userKey), nil +} + +func appendScanKeyResult(result [][]byte, seen map[string]struct{}, key []byte, limit int) [][]byte { + if _, ok := seen[string(key)]; ok { + return result + } + 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, + reader pebbleIteratorProvider, + result [][]byte, + seen map[string]struct{}, + probedPrefixes 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 + } + probeKey := string(prefix) + if _, ok := probedPrefixes[probeKey]; ok { + continue + } + probedPrefixes[probeKey] = struct{}{} + if beforeScanStart(prefix, start) || pastScanEnd(prefix, end) { + continue + } + visible, err := s.visibleExactKeyAt(ctx, reader, 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, reader pebbleIteratorProvider, key []byte, ts uint64) (bool, error) { + upperBound := keyUpperBound(key) + iter, err := reader.NewIter(&pebble.IterOptions{ + LowerBound: encodeKey(key, math.MaxUint64), + UpperBound: upperBound, + }) + if err != nil { + return false, errors.WithStack(err) + } + defer iter.Close() + + for iter.SeekGE(encodeKey(key, math.MaxUint64)); iter.Valid(); iter.Next() { + if err := ctx.Err(); err != nil { + return false, errors.WithStack(err) + } + visible, stop, err := s.visibleExactIteratorEntry(iter, key, ts, upperBound != nil) + if err != nil { + return false, err + } + if stop { + return visible, nil + } + } + return false, nil +} + +func (s *pebbleStore) visibleExactIteratorEntry(iter *pebble.Iterator, key []byte, ts uint64, stopOnPrefixExit bool) (bool, bool, error) { + userKey, version := decodeKeyView(iter.Key()) + if userKey == nil { + return false, false, nil + } + if stopOnPrefixExit && len(key) > 0 && !bytes.HasPrefix(userKey, key) { + return false, true, nil + } + if !bytes.Equal(userKey, key) || version > ts { + return false, false, nil + } + _, err := s.readVisibleVersion(iter, key, ts) + if errors.Is(err, ErrKeyNotFound) { + return false, true, nil + } + if err != nil { + return false, true, err + } + return true, 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 +1439,40 @@ 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 + } + + snap := s.db.NewSnapshot() + defer snap.Close() + + iter, err := snap.NewIter(&pebble.IterOptions{ + LowerBound: encodeKey(start, math.MaxUint64), + }) + if err != nil { + return nil, errors.WithStack(err) + } + defer iter.Close() + + return s.collectScanKeys(ctx, snap, 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..895e7df01 --- /dev/null +++ b/store/scan_keys_test.go @@ -0,0 +1,251 @@ +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 TestPebbleStoreScanKeysAtProbesExactPrefixPastBinaryExtension(t *testing.T) { + t.Parallel() + + ctx := context.Background() + st, err := NewPebbleStore(t.TempDir()) + require.NoError(t, err) + defer st.Close() + + require.NoError(t, st.PutAt(ctx, []byte("a"), []byte("va"), 10, 0)) + require.NoError(t, st.PutAt(ctx, []byte("a\x80"), []byte("vax"), 20, 0)) + + keys, err := st.ScanKeysAt(ctx, nil, nil, 1, 20) + require.NoError(t, err) + require.Equal(t, [][]byte{[]byte("a")}, keys) + + keys, err = st.ScanKeysAt(ctx, cursorAfterTestKey([]byte("a")), nil, 1, 20) + require.NoError(t, err) + require.Equal(t, [][]byte{[]byte("a\x80")}, keys) +} + +func TestPebbleStoreScanKeysAtProbesExactPrefixBeforeBoundedEnd(t *testing.T) { + t.Parallel() + + ctx := context.Background() + st, err := NewPebbleStore(t.TempDir()) + require.NoError(t, err) + defer st.Close() + + require.NoError(t, st.PutAt(ctx, []byte("a"), []byte("va"), 10, 0)) + require.NoError(t, st.PutAt(ctx, []byte("a\x80"), []byte("vax"), 20, 0)) + + keys, err := st.ScanKeysAt(ctx, []byte("a"), []byte("a\x80"), 1, 20) + require.NoError(t, err) + require.Equal(t, [][]byte{[]byte("a")}, keys) +} + +func TestPebbleStoreScanKeysAtPreservesExtensionAfterInvisiblePrefix(t *testing.T) { + t.Parallel() + + ctx := context.Background() + st, err := NewPebbleStore(t.TempDir()) + require.NoError(t, err) + defer st.Close() + + highTS := ^uint64(0) - 1 + require.NoError(t, st.PutAt(ctx, []byte("a"), []byte("future"), highTS, 0)) + require.NoError(t, st.PutAt(ctx, []byte("a\x80"), []byte("visible-extension"), 10, 0)) + + keys, err := st.ScanKeysAt(ctx, []byte("a"), []byte("b"), 2, 10) + require.NoError(t, err) + require.Equal(t, [][]byte{[]byte("a\x80")}, keys) +} + +func TestPebbleStoreScanKeysAtKeepsAllFFKeyWithoutFiniteUpperBound(t *testing.T) { + t.Parallel() + + ctx := context.Background() + st, err := NewPebbleStore(t.TempDir()) + require.NoError(t, err) + defer st.Close() + + key := []byte{0xff, 0xff} + require.NoError(t, st.PutAt(ctx, key, []byte("visible"), 18, 0)) + require.NoError(t, st.PutAt(ctx, []byte(""), []byte("future-empty"), 45, 0)) + require.NoError(t, st.DeleteAt(ctx, key, 47)) + + val, err := st.GetAt(ctx, key, 30) + require.NoError(t, err) + require.Equal(t, []byte("visible"), val) + + keys, err := st.ScanKeysAt(ctx, nil, nil, 2, 30) + require.NoError(t, err) + require.Equal(t, [][]byte{key}, 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.