diff --git a/adapter/distribution_server.go b/adapter/distribution_server.go index 2e594a614..550cfc38b 100644 --- a/adapter/distribution_server.go +++ b/adapter/distribution_server.go @@ -123,6 +123,7 @@ var ( errDistributionNotLeader = errors.New("not leader for distribution catalog") errDistributionCoordinatorRequired = errors.New("distribution coordinator is not configured") errDistributionEngineNotConfigured = errors.New("distribution engine is not configured") + errDistributionCatalogVersionNotFound = errors.New("route catalog version not found") errDistributionCatalogMutationInvalid = errors.New("catalog store mutation is invalid") ) @@ -167,7 +168,7 @@ func (s *DistributionServer) GetRoute(ctx context.Context, req *pb.GetRouteReque if err := s.requireReadReady(); err != nil { return nil, err } - r, ok := s.engine.GetRoute(kv.RouteKey(req.Key)) + r, ok := s.engine.GetRoute(kv.RouteOwnershipKey(req.Key)) if !ok { return &pb.GetRouteResponse{}, nil } @@ -200,6 +201,67 @@ func (s *DistributionServer) ListRoutes(ctx context.Context, req *pb.ListRoutesR }, nil } +func (s *DistributionServer) GetRouteOwnership(ctx context.Context, req *pb.GetRouteOwnershipRequest) (*pb.GetRouteOwnershipResponse, error) { + if err := s.requireReadReady(); err != nil { + return nil, err + } + snapshot, err := s.routeSnapshotAt(req.GetCatalogVersion()) + if err != nil { + return nil, err + } + // Normalized exactly like GetRoute above. An internal storage key -- a + // filesystem chunk, a Redis collection row -- routes by its logical key, + // so looking the raw bytes up in the snapshot answers with the owner of + // the raw family prefix instead of the group that actually owned the key + // at that catalog version. + route, ok := snapshot.RouteOf(kv.RouteOwnershipKey(req.GetKey())) + if !ok { + return &pb.GetRouteOwnershipResponse{ + CatalogVersion: snapshot.Version(), + Found: false, + }, nil + } + return &pb.GetRouteOwnershipResponse{ + Route: toProtoRoute(route), + CatalogVersion: snapshot.Version(), + Found: true, + }, nil +} + +func (s *DistributionServer) GetIntersectingRoutes(ctx context.Context, req *pb.GetIntersectingRoutesRequest) (*pb.GetIntersectingRoutesResponse, error) { + if err := s.requireReadReady(); err != nil { + return nil, err + } + snapshot, err := s.routeSnapshotAt(req.GetCatalogVersion()) + if err != nil { + return nil, err + } + end := req.GetEnd() + if len(end) == 0 { + end = nil + } + routes := snapshot.IntersectingRoutes(req.GetStart(), end) + out := make([]*pb.RouteDescriptor, 0, len(routes)) + for _, route := range routes { + out = append(out, toProtoRoute(route)) + } + return &pb.GetIntersectingRoutesResponse{ + Routes: out, + CatalogVersion: snapshot.Version(), + }, nil +} + +func (s *DistributionServer) routeSnapshotAt(version uint64) (distribution.RouteHistorySnapshot, error) { + if s.engine == nil { + return distribution.RouteHistorySnapshot{}, grpcStatusError(codes.FailedPrecondition, errDistributionEngineNotConfigured.Error()) + } + snapshot, ok := s.engine.SnapshotAt(version) + if !ok { + return distribution.RouteHistorySnapshot{}, grpcStatusErrorf(codes.NotFound, "%s: %d", errDistributionCatalogVersionNotFound, version) + } + return snapshot, nil +} + // GetCatalogCapabilities negotiates the durable delta-watch protocol. func (s *DistributionServer) GetCatalogCapabilities(ctx context.Context, _ *pb.CatalogCapabilitiesRequest) (*pb.CatalogCapabilitiesResponse, error) { if s.catalog == nil { @@ -899,6 +961,19 @@ func toProtoRouteDescriptor(route distribution.RouteDescriptor) *pb.RouteDescrip } } +func toProtoRoute(route distribution.Route) *pb.RouteDescriptor { + return &pb.RouteDescriptor{ + RouteId: route.RouteID, + Start: distribution.CloneBytes(route.Start), + End: distribution.CloneBytes(route.End), + RaftGroupId: route.GroupID, + State: toProtoRouteState(route.State), + StagedVisibilityActive: route.StagedVisibilityActive, + MigrationJobId: route.MigrationJobID, + MinWriteTsExclusive: route.MinWriteTSExclusive, + } +} + func toProtoRouteState(state distribution.RouteState) pb.RouteState { switch state { case distribution.RouteStateActive: diff --git a/adapter/distribution_server_test.go b/adapter/distribution_server_test.go index ab524d279..d8418a2d3 100644 --- a/adapter/distribution_server_test.go +++ b/adapter/distribution_server_test.go @@ -9,6 +9,7 @@ import ( "github.com/bootjp/elastickv/distribution" "github.com/bootjp/elastickv/internal/fskeys" + "github.com/bootjp/elastickv/internal/s3keys" "github.com/bootjp/elastickv/kv" pb "github.com/bootjp/elastickv/proto" "github.com/bootjp/elastickv/store" @@ -59,11 +60,37 @@ func TestDistributionServerGetRoute_NormalizesFilesystemChunkKeys(t *testing.T) require.Equal(t, uint64(2), resp.RaftGroupId) } +func TestDistributionServerGetRoute_NormalizesS3BucketAuxiliaryKeys(t *testing.T) { + t.Parallel() + + bucket := "bucket-a" + routeKey := s3keys.RoutePrefixForBucketAnyGeneration(bucket) + engine := distribution.NewEngine() + engine.UpdateRoute([]byte(""), routeKey, 1) + engine.UpdateRoute(routeKey, nil, 2) + + s := NewDistributionServer(engine, nil) + for _, key := range [][]byte{ + s3keys.BucketMetaKey(bucket), + s3keys.BucketGenerationKey(bucket), + } { + resp, err := s.GetRoute(context.Background(), &pb.GetRouteRequest{Key: key}) + require.NoError(t, err) + require.Equal(t, routeKey, resp.Start) + require.Equal(t, uint64(2), resp.RaftGroupId) + } +} + func TestDistributionServerRouteReadsHonorStartupGate(t *testing.T) { t.Parallel() engine := distribution.NewEngine() - engine.UpdateRoute([]byte("a"), nil, 1) + require.NoError(t, engine.ApplySnapshot(distribution.CatalogSnapshot{ + Version: 1, + Routes: []distribution.RouteDescriptor{ + {RouteID: 1, Start: []byte("a"), End: nil, GroupID: 1, State: distribution.RouteStateActive}, + }, + })) catalog := distribution.NewCatalogStore(store.NewMVCCStore()) _, err := catalog.Save(context.Background(), 0, []distribution.RouteDescriptor{ {RouteID: 1, Start: []byte("a"), End: nil, GroupID: 1, State: distribution.RouteStateActive}, @@ -81,11 +108,37 @@ func TestDistributionServerRouteReadsHonorStartupGate(t *testing.T) { require.Error(t, err) require.Equal(t, codes.Unavailable, status.Code(err)) + _, err = s.GetRouteOwnership(context.Background(), &pb.GetRouteOwnershipRequest{ + Key: []byte("a"), + CatalogVersion: engine.Version(), + }) + require.Error(t, err) + require.Equal(t, codes.Unavailable, status.Code(err)) + + _, err = s.GetIntersectingRoutes(context.Background(), &pb.GetIntersectingRoutesRequest{ + Start: []byte("a"), + End: []byte("z"), + CatalogVersion: engine.Version(), + }) + require.Error(t, err) + require.Equal(t, codes.Unavailable, status.Code(err)) + blocked = false _, err = s.GetRoute(context.Background(), &pb.GetRouteRequest{Key: []byte("a")}) require.NoError(t, err) _, err = s.ListRoutes(context.Background(), &pb.ListRoutesRequest{}) require.NoError(t, err) + _, err = s.GetRouteOwnership(context.Background(), &pb.GetRouteOwnershipRequest{ + Key: []byte("a"), + CatalogVersion: engine.Version(), + }) + require.NoError(t, err) + _, err = s.GetIntersectingRoutes(context.Background(), &pb.GetIntersectingRoutesRequest{ + Start: []byte("a"), + End: []byte("z"), + CatalogVersion: engine.Version(), + }) + require.NoError(t, err) } func TestDistributionServerGetTimestamp_IsMonotonic(t *testing.T) { @@ -182,6 +235,130 @@ func TestDistributionServerListRoutes_RequiresCatalog(t *testing.T) { require.ErrorContains(t, err, errDistributionCatalogNotConfigured.Error()) } +func TestDistributionServerGetRouteOwnership_UsesExactVersionSnapshot(t *testing.T) { + t.Parallel() + + engine := distribution.NewEngine() + require.NoError(t, engine.ApplySnapshot(distribution.CatalogSnapshot{ + Version: 7, + Routes: []distribution.RouteDescriptor{ + {RouteID: 1, Start: []byte("a"), End: []byte("m"), GroupID: 1, State: distribution.RouteStateActive}, + { + RouteID: 2, + Start: []byte("m"), + End: nil, + GroupID: 2, + State: distribution.RouteStateMigratingTarget, + StagedVisibilityActive: true, + MigrationJobID: 44, + MinWriteTSExclusive: 55, + }, + }, + })) + + s := NewDistributionServer(engine, nil) + resp, err := s.GetRouteOwnership(context.Background(), &pb.GetRouteOwnershipRequest{ + Key: []byte("t"), + CatalogVersion: 7, + }) + require.NoError(t, err) + require.True(t, resp.Found) + require.Equal(t, uint64(7), resp.CatalogVersion) + require.Equal(t, uint64(2), resp.Route.RouteId) + require.Equal(t, uint64(2), resp.Route.RaftGroupId) + require.Equal(t, pb.RouteState_ROUTE_STATE_MIGRATING_TARGET, resp.Route.State) + require.True(t, resp.Route.StagedVisibilityActive) + require.Equal(t, uint64(44), resp.Route.MigrationJobId) + require.Equal(t, uint64(55), resp.Route.MinWriteTsExclusive) + + miss, err := s.GetRouteOwnership(context.Background(), &pb.GetRouteOwnershipRequest{ + Key: []byte("0"), + CatalogVersion: 7, + }) + require.NoError(t, err) + require.False(t, miss.Found) + require.Equal(t, uint64(7), miss.CatalogVersion) + require.Nil(t, miss.Route) +} + +func TestDistributionServerGetIntersectingRoutes_UsesExactVersionSnapshot(t *testing.T) { + t.Parallel() + + engine := distribution.NewEngine() + require.NoError(t, engine.ApplySnapshot(distribution.CatalogSnapshot{ + Version: 9, + Routes: []distribution.RouteDescriptor{ + {RouteID: 1, Start: []byte(""), End: []byte("g"), GroupID: 1, State: distribution.RouteStateActive}, + {RouteID: 2, Start: []byte("g"), End: []byte("m"), GroupID: 2, State: distribution.RouteStateWriteFenced}, + {RouteID: 3, Start: []byte("m"), End: nil, GroupID: 3, State: distribution.RouteStateActive}, + }, + })) + + s := NewDistributionServer(engine, nil) + resp, err := s.GetIntersectingRoutes(context.Background(), &pb.GetIntersectingRoutesRequest{ + Start: []byte("f"), + End: []byte("z"), + CatalogVersion: 9, + }) + require.NoError(t, err) + require.Equal(t, uint64(9), resp.CatalogVersion) + require.Len(t, resp.Routes, 3) + require.Equal(t, []uint64{1, 2, 3}, []uint64{resp.Routes[0].RouteId, resp.Routes[1].RouteId, resp.Routes[2].RouteId}) + + rightOpen, err := s.GetIntersectingRoutes(context.Background(), &pb.GetIntersectingRoutesRequest{ + Start: []byte("m"), + End: nil, + CatalogVersion: 9, + }) + require.NoError(t, err) + require.Len(t, rightOpen.Routes, 1) + require.Equal(t, uint64(3), rightOpen.Routes[0].RouteId) +} + +func TestDistributionServerOwnershipRPCs_RejectUnknownCatalogVersion(t *testing.T) { + t.Parallel() + + engine := distribution.NewEngine() + require.NoError(t, engine.ApplySnapshot(distribution.CatalogSnapshot{ + Version: 1, + Routes: []distribution.RouteDescriptor{ + {RouteID: 1, Start: []byte(""), End: nil, GroupID: 1, State: distribution.RouteStateActive}, + }, + })) + + s := NewDistributionServer(engine, nil) + _, err := s.GetRouteOwnership(context.Background(), &pb.GetRouteOwnershipRequest{ + Key: []byte("a"), + CatalogVersion: 2, + }) + require.Error(t, err) + require.Equal(t, codes.NotFound, status.Code(err)) + require.ErrorContains(t, err, errDistributionCatalogVersionNotFound.Error()) + + _, err = s.GetIntersectingRoutes(context.Background(), &pb.GetIntersectingRoutesRequest{ + Start: []byte(""), + CatalogVersion: 2, + }) + require.Error(t, err) + require.Equal(t, codes.NotFound, status.Code(err)) + require.ErrorContains(t, err, errDistributionCatalogVersionNotFound.Error()) +} + +func TestDistributionServerOwnershipRPCs_RequireEngine(t *testing.T) { + t.Parallel() + + s := NewDistributionServer(nil, nil) + _, err := s.GetRouteOwnership(context.Background(), &pb.GetRouteOwnershipRequest{CatalogVersion: 1}) + require.Error(t, err) + require.Equal(t, codes.FailedPrecondition, status.Code(err)) + require.ErrorContains(t, err, errDistributionEngineNotConfigured.Error()) + + _, err = s.GetIntersectingRoutes(context.Background(), &pb.GetIntersectingRoutesRequest{CatalogVersion: 1}) + require.Error(t, err) + require.Equal(t, codes.FailedPrecondition, status.Code(err)) + require.ErrorContains(t, err, errDistributionEngineNotConfigured.Error()) +} + func TestDistributionServerSplitRange_Success(t *testing.T) { t.Parallel() @@ -1250,3 +1427,67 @@ type recordingDistributionFilesystemObserver struct { func (o *recordingDistributionFilesystemObserver) ObserveFilePinnedHotspot(reason string) { o.reasons = append(o.reasons, reason) } + +// GetRouteOwnership answers the historical owner of a key, so it has to +// normalize the same way GetRoute does. An internal storage key -- here a +// filesystem chunk -- routes by its logical !fs|route|chk| key, and the raw +// !fs|chk| bytes sort into a different route entirely. Without normalization +// the RPC reports the raw-prefix owner rather than the group that owned the +// key at that catalog version. +func TestDistributionServerGetRouteOwnership_NormalizesFilesystemChunkKeys(t *testing.T) { + t.Parallel() + + home := uint64(11) + inode := uint64(22) + routeKey := fskeys.ChunkRouteKey(home, inode) + engine := distribution.NewEngine() + require.NoError(t, engine.ApplySnapshot(distribution.CatalogSnapshot{ + Version: 3, + Routes: []distribution.RouteDescriptor{ + {RouteID: 1, Start: []byte(""), End: routeKey, GroupID: 1, State: distribution.RouteStateActive}, + {RouteID: 2, Start: routeKey, End: nil, GroupID: 2, State: distribution.RouteStateActive}, + }, + })) + + s := NewDistributionServer(engine, nil) + resp, err := s.GetRouteOwnership(context.Background(), &pb.GetRouteOwnershipRequest{ + Key: fskeys.ChunkKey(home, inode, 99), + CatalogVersion: 3, + }) + require.NoError(t, err) + require.True(t, resp.Found) + require.Equal(t, uint64(2), resp.Route.RaftGroupId, + "the chunk must resolve to its logical route owner, not the raw-prefix owner") + require.Equal(t, uint64(2), resp.Route.RouteId) +} + +func TestDistributionServerGetRouteOwnership_NormalizesS3BucketAuxiliaryKeys(t *testing.T) { + t.Parallel() + + bucket := "bucket-a" + routeKey := s3keys.RoutePrefixForBucketAnyGeneration(bucket) + engine := distribution.NewEngine() + require.NoError(t, engine.ApplySnapshot(distribution.CatalogSnapshot{ + Version: 3, + Routes: []distribution.RouteDescriptor{ + {RouteID: 1, Start: []byte(""), End: routeKey, GroupID: 1, State: distribution.RouteStateActive}, + {RouteID: 2, Start: routeKey, End: nil, GroupID: 2, State: distribution.RouteStateActive}, + }, + })) + + s := NewDistributionServer(engine, nil) + for _, key := range [][]byte{ + s3keys.BucketMetaKey(bucket), + s3keys.BucketGenerationKey(bucket), + } { + resp, err := s.GetRouteOwnership(context.Background(), &pb.GetRouteOwnershipRequest{ + Key: key, + CatalogVersion: 3, + }) + require.NoError(t, err) + require.True(t, resp.Found) + require.Equal(t, uint64(2), resp.Route.RaftGroupId, + "bucket auxiliary keys must resolve to the bucket route owner, not the raw-prefix owner") + require.Equal(t, uint64(2), resp.Route.RouteId) + } +} diff --git a/adapter/grpc.go b/adapter/grpc.go index 8fe8e9c5d..f08c8cfa1 100644 --- a/adapter/grpc.go +++ b/adapter/grpc.go @@ -76,6 +76,10 @@ type rawVersionPresenceReader interface { VersionExistsAtOrBeforeGroupWithReadFence(ctx context.Context, key []byte, groupID uint64, ts uint64, readRouteVersion uint64) (bool, bool, error) } +type rawVersionPresenceBatchReader interface { + VersionsExistAtOrBeforeGroupWithReadFence(ctx context.Context, keys [][]byte, groupID uint64, ts uint64, readRouteVersion uint64) ([]bool, bool, error) +} + type rawGroupReverseScanner interface { ReverseScanGroupAt(ctx context.Context, groupID uint64, start []byte, end []byte, limit int, ts uint64) ([]*store.KVPair, error) } @@ -189,6 +193,21 @@ func (r *GRPCServer) RawLatestCommitTS(ctx context.Context, req *pb.RawLatestCom if err := r.requireReadReady(); err != nil { return nil, err } + readRouteVersion := r.readRouteVersion(req.GetReadRouteVersion()) + if len(req.GetKeyBatch()) > 0 { + visible, visibleSupported, err := r.rawVersionsVisibleAt(ctx, req, readRouteVersion) + if err != nil { + return nil, errors.WithStack(err) + } + return &pb.RawLatestCommitTSResponse{ + VersionVisibleResults: visible, + VersionVisibleSupported: visibleSupported, + }, nil + } + return r.rawLatestCommitTSSingle(ctx, req, readRouteVersion) +} + +func (r *GRPCServer) rawLatestCommitTSSingle(ctx context.Context, req *pb.RawLatestCommitTSRequest, readRouteVersion uint64) (*pb.RawLatestCommitTSResponse, error) { key := req.GetKey() if len(key) == 0 { // No key: return the store's global last-committed watermark. @@ -204,7 +223,6 @@ func (r *GRPCServer) RawLatestCommitTS(ctx context.Context, req *pb.RawLatestCom var ts uint64 var exists bool var err error - readRouteVersion := r.readRouteVersion(req.GetReadRouteVersion()) if groupID := req.GetGroupId(); groupID != 0 { groupReader, ok := r.store.(rawGroupCommitTSReader) if !ok { @@ -233,6 +251,37 @@ func (r *GRPCServer) RawLatestCommitTS(ctx context.Context, req *pb.RawLatestCom }, nil } +func (r *GRPCServer) rawVersionsVisibleAt(ctx context.Context, req *pb.RawLatestCommitTSRequest, readRouteVersion uint64) ([]bool, bool, error) { + keys, err := pb.DecodeRawLatestCommitTSKeyBatch(req.GetKeyBatch(), maxGRPCScanLimit) + if err != nil { + return nil, false, errors.WithStack(status.Error(codes.InvalidArgument, err.Error())) + } + out := make([]bool, len(keys)) + at := req.GetVersionVisibleAtTs() + if at == 0 { + return out, false, nil + } + if reader, ok := r.store.(rawVersionPresenceBatchReader); ok { + visible, supported, err := reader.VersionsExistAtOrBeforeGroupWithReadFence(ctx, keys, req.GetGroupId(), at, readRouteVersion) + return visible, supported, errors.WithStack(err) + } + reader, ok := r.store.(rawVersionPresenceReader) + if !ok { + return out, false, nil + } + for i, key := range keys { + visible, supported, err := reader.VersionExistsAtOrBeforeGroupWithReadFence(ctx, key, req.GetGroupId(), at, readRouteVersion) + if err != nil { + return nil, false, errors.WithStack(err) + } + if !supported { + return out, false, nil + } + out[i] = visible + } + return out, true, nil +} + // rawVersionVisibleAt answers the optional version_visible_at_ts probe. The // second bool tells the caller whether this server answered it at all, so a // store that cannot check presence never looks like "no version exists". diff --git a/adapter/grpc_test.go b/adapter/grpc_test.go index a8b0e06fb..b80534040 100644 --- a/adapter/grpc_test.go +++ b/adapter/grpc_test.go @@ -37,6 +37,8 @@ const ( grpcSequenceShortIterations = 256 ) +var _ rawGroupCommitTSReader = (*kvstore.ShardStore)(nil) + func grpcSequenceIterations(t testing.TB) int { t.Helper() if testing.Short() { @@ -403,6 +405,39 @@ func TestGRPCServer_RawLatestCommitTS_UsesExplicitGroup(t *testing.T) { require.Zero(t, st.latestReadRouteVersion) } +func TestGRPCServer_RawLatestCommitTS_ExplicitGroupShardStoreVersionProbe(t *testing.T) { + t.Parallel() + + ctx := context.Background() + engine := distribution.NewEngine() + require.NoError(t, engine.ApplySnapshot(distribution.CatalogSnapshot{ + Version: 1, + Routes: []distribution.RouteDescriptor{ + {RouteID: 1, Start: []byte(""), End: nil, GroupID: 42, State: distribution.RouteStateActive}, + }, + })) + groupStore := store.NewMVCCStore() + t.Cleanup(func() { require.NoError(t, groupStore.Close()) }) + key := []byte("!dist|migstage|probe|k") + require.NoError(t, groupStore.PutAt(ctx, key, []byte("v"), 10, 0)) + shards := kvstore.NewShardStore(engine, map[uint64]*kvstore.ShardGroup{ + 42: {Store: groupStore}, + }) + server := NewGRPCServer(shards, nil) + + resp, err := server.RawLatestCommitTS(ctx, &pb.RawLatestCommitTSRequest{ + Key: key, + GroupId: 42, + ReadRouteVersion: 1, + VersionVisibleAtTs: 10, + }) + require.NoError(t, err) + require.True(t, resp.GetExists()) + require.Equal(t, uint64(10), resp.GetTs()) + require.True(t, resp.GetVersionVisibleSupported()) + require.True(t, resp.GetVersionVisible()) +} + func TestGRPCServer_RawReadFenceHelpersKeepCallerRouteVersion(t *testing.T) { t.Parallel() @@ -859,6 +894,49 @@ func TestGRPCServer_RawScanAt_KeysOnlyUsesExplicitGroup(t *testing.T) { require.Equal(t, []byte("z"), st.scanEnd) } +func TestGRPCServer_RawScanAt_KeysOnlyExplicitGroupMergesStagedVisibility(t *testing.T) { + t.Parallel() + + ctx := context.Background() + engine := distribution.NewEngine() + require.NoError(t, engine.ApplySnapshot(distribution.CatalogSnapshot{ + Version: 1, + Routes: []distribution.RouteDescriptor{ + { + RouteID: 1, + Start: []byte("a"), + End: []byte("z"), + GroupID: 1, + State: distribution.RouteStateActive, + StagedVisibilityActive: true, + MigrationJobID: 9, + }, + }, + })) + group := &kvstore.ShardGroup{Store: store.NewMVCCStore()} + shards := kvstore.NewShardStore(engine, map[uint64]*kvstore.ShardGroup{1: group}) + t.Cleanup(func() { require.NoError(t, shards.Close()) }) + + require.NoError(t, group.Store.PutAt(ctx, []byte("b"), []byte("live-b"), 10, 0)) + require.NoError(t, group.Store.PutAt(ctx, distribution.MigrationStagedDataKey(9, []byte("b")), []byte("staged-b"), 20, 0)) + require.NoError(t, group.Store.PutAt(ctx, distribution.MigrationStagedDataKey(9, []byte("c")), []byte("staged-c"), 30, 0)) + + s := NewGRPCServer(shards, nil) + resp, err := s.RawScanAt(ctx, &pb.RawScanAtRequest{ + StartKey: []byte("a"), + EndKey: []byte("z"), + Limit: 10, + Ts: 35, + GroupId: 1, + KeysOnly: true, + }) + require.NoError(t, err) + require.Equal(t, []*pb.RawKVPair{ + {Key: []byte("b")}, + {Key: []byte("c")}, + }, resp.GetKv()) +} + func TestGRPCServer_RawScanAt_ReverseKeysOnlyUsesExplicitGroup(t *testing.T) { t.Parallel() @@ -1143,6 +1221,41 @@ func (s *recordingVersionPresenceStore) VersionExistsAtOrBeforeGroupWithReadFenc return s.visible, s.supported, nil } +type recordingVersionPresenceBatchStore struct { + store.MVCCStore + + visible map[string]bool + supported bool + calls int + lastKeys [][]byte + lastGroup uint64 + lastTS uint64 + lastReadRouteVersion uint64 +} + +func cloneBytes2D(keys [][]byte) [][]byte { + out := make([][]byte, 0, len(keys)) + for _, key := range keys { + out = append(out, append([]byte(nil), key...)) + } + return out +} + +func (s *recordingVersionPresenceBatchStore) VersionsExistAtOrBeforeGroupWithReadFence( + _ context.Context, keys [][]byte, groupID uint64, ts uint64, readRouteVersion uint64, +) ([]bool, bool, error) { + s.calls++ + s.lastKeys = cloneBytes2D(keys) + s.lastGroup = groupID + s.lastTS = ts + s.lastReadRouteVersion = readRouteVersion + out := make([]bool, len(keys)) + for i, key := range keys { + out[i] = s.visible[string(key)] + } + return out, s.supported, nil +} + // version_visible_at_ts is optional: only a request that asks gets an answer, // and a store that cannot answer must not look like "no version exists". func TestGRPCServer_RawLatestCommitTS_VersionVisibleProbe(t *testing.T) { @@ -1219,3 +1332,54 @@ func TestGRPCServer_RawLatestCommitTS_VersionVisibleProbe(t *testing.T) { }) } } + +func TestGRPCServer_RawLatestCommitTS_BatchVersionVisibleProbe(t *testing.T) { + t.Parallel() + + st := &recordingVersionPresenceBatchStore{ + MVCCStore: store.NewMVCCStore(), + visible: map[string]bool{"a": true, "b": false}, + supported: true, + } + t.Cleanup(func() { _ = st.Close() }) + s := NewGRPCServer(st, nil) + + resp, err := s.RawLatestCommitTS(context.Background(), &pb.RawLatestCommitTSRequest{ + KeyBatch: pb.EncodeRawLatestCommitTSKeyBatch([][]byte{[]byte("a"), []byte("b")}), + GroupId: 42, + ReadRouteVersion: 77, + VersionVisibleAtTs: 100, + }) + require.NoError(t, err) + require.Equal(t, []bool{true, false}, resp.GetVersionVisibleResults()) + require.True(t, resp.GetVersionVisibleSupported()) + require.Equal(t, 1, st.calls) + require.Equal(t, [][]byte{[]byte("a"), []byte("b")}, st.lastKeys) + require.Equal(t, uint64(42), st.lastGroup) + require.Equal(t, uint64(100), st.lastTS) + require.Equal(t, uint64(77), st.lastReadRouteVersion) +} + +func TestGRPCServer_RawLatestCommitTS_RejectsOversizedBatchBeforeProbe(t *testing.T) { + t.Parallel() + + keys := make([][]byte, maxGRPCScanLimit+1) + for i := range keys { + keys[i] = []byte("k") + } + st := &recordingVersionPresenceBatchStore{ + MVCCStore: store.NewMVCCStore(), + visible: map[string]bool{}, + supported: true, + } + t.Cleanup(func() { _ = st.Close() }) + s := NewGRPCServer(st, nil) + + _, err := s.RawLatestCommitTS(context.Background(), &pb.RawLatestCommitTSRequest{ + KeyBatch: pb.EncodeRawLatestCommitTSKeyBatch(keys), + VersionVisibleAtTs: 100, + }) + require.Error(t, err) + require.Equal(t, codes.InvalidArgument, status.Code(err)) + require.Zero(t, st.calls) +} diff --git a/adapter/internal.go b/adapter/internal.go index 6bc5d9407..0f7b58773 100644 --- a/adapter/internal.go +++ b/adapter/internal.go @@ -3,10 +3,17 @@ package adapter import ( "bytes" "context" + "os" + "strings" + "github.com/bootjp/elastickv/distribution" + "github.com/bootjp/elastickv/internal" + "github.com/bootjp/elastickv/internal/fskeys" "github.com/bootjp/elastickv/internal/raftengine" + "github.com/bootjp/elastickv/internal/s3keys" "github.com/bootjp/elastickv/kv" pb "github.com/bootjp/elastickv/proto" + "github.com/bootjp/elastickv/store" "github.com/cockroachdb/errors" "google.golang.org/grpc/codes" "google.golang.org/grpc/status" @@ -28,6 +35,37 @@ func WithInternalTimestampAllocator(alloc kv.TimestampAllocator) InternalOption } } +func WithInternalStore(st store.MVCCStore) InternalOption { + return func(i *Internal) { + i.store = st + } +} + +func WithInternalMigrationProposer(proposer raftengine.Proposer) InternalOption { + return func(i *Internal) { + i.migrationProposer = proposer + } +} + +func WithInternalMigrationImportGate(gate func(context.Context) error) InternalOption { + return func(i *Internal) { + i.migrationImportGate = gate + } +} + +func WithInternalMigrationPromoteGate(gate func(context.Context) error) InternalOption { + return func(i *Internal) { + i.migrationPromoteGate = gate + } +} + +func WithInternalMigrationExportRouting(groupID uint64, resolver kv.PartitionResolver) InternalOption { + return func(i *Internal) { + i.migrationExportGroupID = groupID + i.migrationExportResolver = resolver + } +} + func WithInternalAdminProposer(proposer raftengine.Proposer) InternalOption { return func(i *Internal) { i.adminProposer = proposer @@ -53,10 +91,12 @@ func WithInternalForwardWriteObserver(observer ForwardWriteObserver) InternalOpt func NewInternalWithEngine(txm kv.Transactional, leader raftengine.LeaderView, clock *kv.HLC, relay *RedisPubSubRelay, opts ...InternalOption) *Internal { i := &Internal{ - leader: leader, - transactionManager: txm, - clock: clock, - relay: relay, + leader: leader, + transactionManager: txm, + clock: clock, + relay: relay, + migrationImportGate: defaultMigrationImportGate, + migrationPromoteGate: defaultMigrationPromoteGate, } for _, opt := range opts { opt(i) @@ -65,15 +105,21 @@ func NewInternalWithEngine(txm kv.Transactional, leader raftengine.LeaderView, c } type Internal struct { - leader raftengine.LeaderView - transactionManager kv.Transactional - clock *kv.HLC - tsAllocator kv.TimestampAllocator - adminProposer raftengine.Proposer - lastCommitTimestamp func() uint64 - relay *RedisPubSubRelay - writeGate kv.MutationWriteGate - forwardWriteObserver ForwardWriteObserver + leader raftengine.LeaderView + transactionManager kv.Transactional + clock *kv.HLC + tsAllocator kv.TimestampAllocator + adminProposer raftengine.Proposer + lastCommitTimestamp func() uint64 + relay *RedisPubSubRelay + store store.MVCCStore + migrationProposer raftengine.Proposer + migrationImportGate func(context.Context) error + migrationPromoteGate func(context.Context) error + migrationExportGroupID uint64 + migrationExportResolver kv.PartitionResolver + writeGate kv.MutationWriteGate + forwardWriteObserver ForwardWriteObserver pb.UnimplementedInternalServer } @@ -84,6 +130,30 @@ var ErrNotLeader = errors.New("not leader") var ErrLeaderNotFound = errors.New("leader not found") var ErrTxnTimestampOverflow = errors.New("txn timestamp overflow") +const ( + defaultMigrationExportChunkBytes = 4 << 20 + defaultMigrationExportScanFactor = 4 + defaultMigrationExportMaxVersions = 1024 + + // Hard server-side ceilings. The defaults above only fill in a bound the + // request left unset, so without these a caller could ask one + // ExportVersions call to scan and decode an arbitrary portion of the source + // store before producing its next streamed response -- for a sparse family + // or route filter that accepts few rows, math.MaxUint64 removes the only + // work bound there is and the serving leader's I/O and CPU go with it. A + // clamped caller simply advances through more cursor rounds, which the + // response already supports. Same shape as the promotion path's clamp. + maxMigrationExportChunkBytes = 32 << 20 + maxMigrationExportScanBytes = maxMigrationExportChunkBytes * defaultMigrationExportScanFactor + + // Headroom for proto framing over the payload bytes a page carries: field + // tags and length prefixes for at most defaultMigrationExportMaxVersions + // entries come to tens of kilobytes, so a mebibyte covers them with room + // to spare. + migrationExportPageFramingHeadroom = 1 << 20 + migrationExportPageByteBudget = internal.GRPCMaxMessageBytes - migrationExportPageFramingHeadroom +) + func (i *Internal) Forward(ctx context.Context, req *pb.ForwardRequest) (*pb.ForwardResponse, error) { if i.leader == nil || i.leader.State() != raftengine.StateLeader { return nil, errors.WithStack(ErrNotLeader) @@ -184,6 +254,538 @@ func (i *Internal) RelayPublish(_ context.Context, req *pb.RelayPublishRequest) }, nil } +func (i *Internal) ExportRangeVersions(req *pb.ExportRangeVersionsRequest, stream pb.Internal_ExportRangeVersionsServer) error { + if err := i.validateExportRangeVersionsRequest(req); err != nil { + return err + } + if err := i.verifyInternalLeaderApplied(stream.Context()); err != nil { + return err + } + return i.streamExportRangeVersions(req, stream) +} + +func (i *Internal) validateExportRangeVersionsRequest(req *pb.ExportRangeVersionsRequest) error { + if req == nil { + return errors.WithStack(status.Error(codes.InvalidArgument, "export range versions request is nil")) + } + if i.store == nil { + return errors.WithStack(status.Error(codes.FailedPrecondition, "migration export store is not configured")) + } + if req.GetMaxCommitTs() == 0 { + return errors.WithStack(status.Error(codes.InvalidArgument, "migration export max_commit_ts is required")) + } + if req.GetKeyFamily() == 0 { + return errors.WithStack(status.Error(codes.InvalidArgument, "migration export key_family is required")) + } + if exportRangeVersionsRequestRouteUnbounded(req) { + return errors.WithStack(status.Error(codes.InvalidArgument, "migration export requires route bounds")) + } + return nil +} + +func exportRangeVersionsRequestRouteUnbounded(req *pb.ExportRangeVersionsRequest) bool { + return len(req.GetRouteStart()) == 0 && + len(req.GetRouteEnd()) == 0 +} + +func (i *Internal) streamExportRangeVersions(req *pb.ExportRangeVersionsRequest, stream pb.Internal_ExportRangeVersionsServer) error { + opts := i.exportRangeVersionsOptions(req) + for { + result, err := i.store.ExportVersions(stream.Context(), opts) + if err != nil { + return errors.WithStack(err) + } + if !result.Done && bytes.Equal(opts.Cursor, result.NextCursor) { + return errors.WithStack(status.Error(codes.Internal, "migration export cursor did not progress")) + } + resp := &pb.ExportRangeVersionsResponse{ + Versions: protoMVCCVersionsFromStore(result.Versions), + NextCursor: result.NextCursor, + Done: result.Done, + } + if err := checkMigrationExportPageSize(resp); err != nil { + return err + } + if err := stream.Send(resp); err != nil { + return errors.WithStack(err) + } + if result.Done { + return nil + } + opts.Cursor = result.NextCursor + } +} + +// checkMigrationExportPageSize refuses to send a page the transport cannot +// carry. The store's MaxBytes budget is only consulted after a version has +// been appended, and storage accepts values far above the message limit +// (store.maxSnapshotValueSize is 256 MiB against internal.GRPCMaxMessageBytes' +// 64 MiB), so a single oversized row can land in an otherwise bounded page. +// Sending it fails with ResourceExhausted, and so does every retry of the same +// cursor -- the bracket never completes and the error says nothing about which +// row is responsible. Naming the row instead keeps the failure diagnosable. +// Carrying values that large needs a chunked migration wire format, which is a +// protocol change rather than a fix here. +func checkMigrationExportPageSize(resp *pb.ExportRangeVersionsResponse) error { + if migrationExportPagePayloadBytes(resp) <= migrationExportPageByteBudget { + return nil + } + widest := widestMigrationExportVersion(resp.GetVersions()) + return errors.WithStack(status.Errorf(codes.FailedPrecondition, + "migration export page exceeds the %d byte message limit; key %q holds %d value bytes", + internal.GRPCMaxMessageBytes, widest.GetKey(), len(widest.GetValue()))) +} + +// migrationExportPagePayloadBytes sums the caller-controlled bytes in a page. +// It deliberately undercounts the proto framing rather than calling +// proto.Size, which caches its result inside the message; the budget carries +// enough headroom to cover the framing of a full page. +func migrationExportPagePayloadBytes(resp *pb.ExportRangeVersionsResponse) int { + total := len(resp.GetNextCursor()) + for _, version := range resp.GetVersions() { + total += len(version.GetKey()) + len(version.GetValue()) + } + return total +} + +func widestMigrationExportVersion(versions []*pb.MVCCVersion) *pb.MVCCVersion { + var widest *pb.MVCCVersion + for _, version := range versions { + if widest == nil || len(version.GetValue()) > len(widest.GetValue()) { + widest = version + } + } + return widest +} + +func (i *Internal) ImportRangeVersions(ctx context.Context, req *pb.ImportRangeVersionsRequest) (*pb.ImportRangeVersionsResponse, error) { + if req == nil { + return nil, errors.WithStack(status.Error(codes.InvalidArgument, "import range versions request is nil")) + } + if err := validateImportRangeVersionsRequest(req); err != nil { + return nil, err + } + if i.migrationProposer == nil { + return nil, errors.WithStack(status.Error(codes.FailedPrecondition, "migration import proposer is not configured")) + } + if err := i.verifyInternalLeader(ctx); err != nil { + return nil, err + } + if err := i.verifyMigrationImportEnabled(ctx); err != nil { + return nil, err + } + result, err := i.proposeMigrationImport(ctx, req) + if err != nil { + return nil, errors.WithStack(err) + } + if i.clock != nil && result.MaxImportedTS > 0 { + i.clock.Observe(result.MaxImportedTS) + } + return &pb.ImportRangeVersionsResponse{AckedCursor: result.AckedCursor}, nil +} + +func validateImportRangeVersionsRequest(req *pb.ImportRangeVersionsRequest) error { + if req.GetJobId() == 0 { + return errors.WithStack(status.Error(codes.InvalidArgument, "import range versions job_id is required")) + } + if req.GetBracketId() == 0 { + return errors.WithStack(status.Error(codes.InvalidArgument, "import range versions bracket_id is required")) + } + return nil +} + +func (i *Internal) verifyMigrationImportEnabled(ctx context.Context) error { + if i.migrationImportGate == nil { + return nil + } + if err := i.migrationImportGate(ctx); err != nil { + return errors.WithStack(err) + } + return nil +} + +func (i *Internal) PromoteStagedVersions(ctx context.Context, req *pb.PromoteStagedVersionsRequest) (*pb.PromoteStagedVersionsResponse, error) { + if req == nil { + return nil, errors.WithStack(status.Error(codes.InvalidArgument, "promote staged versions request is nil")) + } + if req.GetJobId() == 0 { + return nil, errors.WithStack(status.Error(codes.InvalidArgument, "promote staged versions job_id is required")) + } + if i.migrationProposer == nil { + return nil, errors.WithStack(status.Error(codes.FailedPrecondition, "migration promote proposer is not configured")) + } + if err := i.verifyInternalLeader(ctx); err != nil { + return nil, err + } + if err := i.verifyMigrationPromoteEnabled(ctx); err != nil { + return nil, err + } + if err := validatePromoteStagedVersionsRequest(req); err != nil { + return nil, errors.WithStack(err) + } + result, err := i.proposeMigrationPromote(ctx, req) + if err != nil { + return nil, errors.WithStack(err) + } + if i.clock != nil && result.MaxPromotedTS > 0 { + i.clock.Observe(result.MaxPromotedTS) + } + return &pb.PromoteStagedVersionsResponse{ + NextCursor: result.NextCursor, + Done: result.Done, + PromotedRows: result.PromotedRows, + MaxPromotedTs: result.MaxPromotedTS, + }, nil +} + +func (i *Internal) verifyMigrationPromoteEnabled(ctx context.Context) error { + if i.migrationPromoteGate == nil { + return nil + } + if err := i.migrationPromoteGate(ctx); err != nil { + return errors.WithStack(err) + } + return nil +} + +func validatePromoteStagedVersionsRequest(req *pb.PromoteStagedVersionsRequest) error { + prefix := distribution.MigrationStagedDataKeyPrefix(req.GetJobId()) + if err := store.ValidatePromotionCursorForRange(req.GetCursor(), prefix, store.PrefixScanEnd(prefix)); err != nil { + if errors.Is(err, store.ErrInvalidExportCursor) { + return errors.WithStack(status.Error(codes.InvalidArgument, store.ErrInvalidExportCursor.Error())) + } + return errors.WithStack(status.Errorf(codes.Internal, "validate promote cursor: %v", err)) + } + return nil +} + +func defaultMigrationImportGate(context.Context) error { + if migrationImportOpcodeEnabledFromEnv() { + return nil + } + return errors.WithStack(status.Error(codes.FailedPrecondition, "migration import opcode is disabled; enable after every voter is running a build that supports migration import")) +} + +func migrationImportOpcodeEnabledFromEnv() bool { + return envFlagEnabled("ELASTICKV_ENABLE_MIGRATION_IMPORT_OPCODE") +} + +func defaultMigrationPromoteGate(context.Context) error { + if migrationPromoteOpcodeEnabledFromEnv() { + return nil + } + return errors.WithStack(status.Error(codes.FailedPrecondition, "migration promote opcode is disabled; enable after every voter is running a build that supports migration promotion")) +} + +func migrationPromoteOpcodeEnabledFromEnv() bool { + return envFlagEnabled("ELASTICKV_ENABLE_MIGRATION_PROMOTE_OPCODE") +} + +func envFlagEnabled(name string) bool { + switch strings.ToLower(strings.TrimSpace(os.Getenv(name))) { + case "1", "true", "yes", "on": + return true + default: + return false + } +} + +func (i *Internal) verifyInternalLeader(ctx context.Context) error { + if i.leader == nil { + return errors.WithStack(ErrNotLeader) + } + if i.leader.State() != raftengine.StateLeader { + return errors.WithStack(ErrNotLeader) + } + if err := i.leader.VerifyLeader(ctx); err != nil { + return errors.WithStack(ErrNotLeader) + } + return nil +} + +func (i *Internal) verifyInternalLeaderApplied(ctx context.Context) error { + if i.leader == nil { + return errors.WithStack(ErrNotLeader) + } + if i.leader.State() != raftengine.StateLeader { + return errors.WithStack(ErrNotLeader) + } + _, err := i.leader.LinearizableRead(ctx) + return errors.WithStack(err) +} + +func (i *Internal) proposeMigrationImport(ctx context.Context, req *pb.ImportRangeVersionsRequest) (store.ImportVersionsResult, error) { + cmd, err := kv.MarshalMigrationImportCommand(req) + if err != nil { + return store.ImportVersionsResult{}, errors.WithStack(err) + } + resp, err := i.proposeMigrationCommand(ctx, cmd, "migration import") + if err != nil { + return store.ImportVersionsResult{}, errors.WithStack(err) + } + switch resp := resp.(type) { + case store.ImportVersionsResult: + return resp, nil + case *store.ImportVersionsResult: + if resp == nil { + return store.ImportVersionsResult{}, errors.New("migration import apply returned nil result") + } + return *resp, nil + case error: + return store.ImportVersionsResult{}, errors.WithStack(resp) + default: + return store.ImportVersionsResult{}, errors.WithStack(errors.Newf("unexpected migration import apply response type %T", resp)) + } +} + +func (i *Internal) proposeMigrationPromote(ctx context.Context, req *pb.PromoteStagedVersionsRequest) (store.PromoteVersionsResult, error) { + cmd, err := kv.MarshalMigrationPromoteCommand(req) + if err != nil { + return store.PromoteVersionsResult{}, errors.WithStack(err) + } + resp, err := i.proposeMigrationCommand(ctx, cmd, "migration promote") + if err != nil { + return store.PromoteVersionsResult{}, errors.WithStack(err) + } + switch resp := resp.(type) { + case store.PromoteVersionsResult: + return resp, nil + case *store.PromoteVersionsResult: + if resp == nil { + return store.PromoteVersionsResult{}, errors.New("migration promote apply returned nil result") + } + return *resp, nil + case error: + return store.PromoteVersionsResult{}, errors.WithStack(resp) + default: + return store.PromoteVersionsResult{}, errors.WithStack(errors.Newf("unexpected migration promote apply response type %T", resp)) + } +} + +func (i *Internal) proposeMigrationCommand(ctx context.Context, cmd []byte, label string) (any, error) { + result, err := i.migrationProposer.Propose(ctx, cmd) + if err != nil { + return nil, errors.WithStack(err) + } + if result == nil { + return nil, errors.WithStack(errors.Newf("%s proposal returned nil result", label)) + } + return result.Response, nil +} + +func (i *Internal) exportRangeVersionsOptions(req *pb.ExportRangeVersionsRequest) store.ExportVersionsOptions { + chunkBytes := clampMigrationExportBound( + uint64(req.GetChunkBytes()), defaultMigrationExportChunkBytes, maxMigrationExportChunkBytes) + maxScannedBytes := clampMigrationExportBound( + req.GetMaxScannedBytes(), chunkBytes*defaultMigrationExportScanFactor, maxMigrationExportScanBytes) + startKey, endKey := migrationExportScanBounds(req) + opts := store.ExportVersionsOptions{ + StartKey: startKey, + EndKey: endKey, + MinCommitTSExclusive: req.GetMinCommitTs(), + MaxCommitTSInclusive: req.GetMaxCommitTs(), + Cursor: req.GetCursor(), + MaxVersions: defaultMigrationExportMaxVersions, + MaxBytes: chunkBytes, + MaxScannedBytes: maxScannedBytes, + KeyFamily: req.GetKeyFamily(), + AcceptKey: i.migrationExportFilter(req), + AcceptVersion: i.migrationExportVersionFilter(req), + } + return opts +} + +func migrationExportScanBounds(req *pb.ExportRangeVersionsRequest) ([]byte, []byte) { + start := bytes.Clone(req.GetRangeStart()) + end := bytes.Clone(req.GetRangeEnd()) + switch req.GetKeyFamily() { + case distribution.MigrationFamilyFilesystemChunk: + if bytes.HasPrefix(start, fskeys.ChunkAllPrefix()) { + return start, end + } + if scanStart, scanEnd, ok := filesystemChunkExportScanBounds(req.GetRouteStart(), req.GetRouteEnd()); ok { + return scanStart, scanEnd + } + return fskeys.ChunkAllPrefix(), prefixScanEnd(fskeys.ChunkAllPrefix()) + case distribution.MigrationFamilyFilesystemUsage: + if bytes.HasPrefix(start, fskeys.UsageRouteAllPrefix()) { + return start, end + } + return filesystemUsageExportScanBounds(req.GetRouteStart(), req.GetRouteEnd()) + default: + return start, end + } +} + +func filesystemChunkExportScanBounds(routeStart, routeEnd []byte) ([]byte, []byte, bool) { + routePrefix := fskeys.ChunkRouteAllPrefix() + routeDomainEnd := prefixScanEnd(routePrefix) + if !rangesIntersect(routeStart, routeEnd, routePrefix, routeDomainEnd) { + return fskeys.ChunkAllPrefix(), fskeys.ChunkAllPrefix(), true + } + rawPrefix := fskeys.ChunkAllPrefix() + rawDomainEnd := prefixScanEnd(rawPrefix) + start := rawPrefix + if len(routeStart) > 0 && bytes.Compare(routeStart, routePrefix) > 0 { + if !bytes.HasPrefix(routeStart, routePrefix) { + return nil, nil, false + } + start = append(bytes.Clone(rawPrefix), routeStart[len(routePrefix):]...) + } + end := rawDomainEnd + if len(routeEnd) > 0 && bytes.Compare(routeEnd, routeDomainEnd) < 0 { + if bytes.Compare(routeEnd, routePrefix) <= 0 { + return rawPrefix, rawPrefix, true + } + if !bytes.HasPrefix(routeEnd, routePrefix) { + return nil, nil, false + } + end = append(bytes.Clone(rawPrefix), routeEnd[len(routePrefix):]...) + } + return bytes.Clone(start), bytes.Clone(end), true +} + +func filesystemUsageExportScanBounds(routeStart, routeEnd []byte) ([]byte, []byte) { + start := fskeys.UsageRouteAllPrefix() + if len(routeStart) > 0 { + start = fskeys.UsageRouteKey(routeStart) + } + end := prefixScanEnd(fskeys.UsageRouteAllPrefix()) + if len(routeEnd) > 0 { + end = fskeys.UsageRouteKey(routeEnd) + } + return start, end +} + +// clampMigrationExportBound resolves one export bound: unset takes the default, +// anything above the hard ceiling is clamped down to it. +func clampMigrationExportBound(requested, fallback, ceiling uint64) uint64 { + if requested == 0 { + return fallback + } + return min(requested, ceiling) +} + +func (i *Internal) migrationExportFilter(req *pb.ExportRangeVersionsRequest) func([]byte) bool { + bracket := migrationExportBracket(req) + if req.GetKeyFamily() == distribution.MigrationFamilyLegacyListMetaDelta { + return bracket.ContainsRawKey + } + routeFilter := i.migrationExportRouteFilter(req) + if migrationFamilyRequiresDecodedS3(req.GetKeyFamily()) { + routeFilter = decodedS3BucketRouteFilter(req.GetKeyFamily(), req.GetRouteStart(), req.GetRouteEnd()) + } + return func(rawKey []byte) bool { + return bracket.ContainsRawKey(rawKey) && routeFilter(rawKey) + } +} + +func (i *Internal) migrationExportVersionFilter(req *pb.ExportRangeVersionsRequest) func([]byte, []byte) bool { + if req.GetKeyFamily() != distribution.MigrationFamilyLegacyListMetaDelta { + return nil + } + bracket := migrationExportBracket(req) + return func(rawKey, value []byte) bool { + return bracket.ContainsRoutedVersion(rawKey, value, req.GetRouteStart(), req.GetRouteEnd(), nil) + } +} + +func migrationExportBracket(req *pb.ExportRangeVersionsRequest) distribution.MigrationBracket { + excludeKnownInternal := req.GetExcludeKnownInternal() || req.GetKeyFamily() == distribution.MigrationFamilyUser + start, end := migrationExportScanBounds(req) + return distribution.MigrationBracket{ + Family: req.GetKeyFamily(), + Start: start, + End: end, + ExcludeKnownInternal: excludeKnownInternal, + ExcludePrefixes: cloneByteSlices(req.GetExcludePrefixes()), + } +} + +func (i *Internal) migrationExportRouteFilter(req *pb.ExportRangeVersionsRequest) func([]byte) bool { + if i != nil && i.migrationExportGroupID != 0 && i.migrationExportResolver != nil { + return kv.RouteKeyFilterForGroup(req.GetRouteStart(), req.GetRouteEnd(), i.migrationExportGroupID, i.migrationExportResolver) + } + return kv.RouteKeyFilter(req.GetRouteStart(), req.GetRouteEnd()) +} + +func migrationFamilyRequiresDecodedS3(family uint32) bool { + return family == distribution.MigrationFamilyS3BucketMeta || + family == distribution.MigrationFamilyS3BucketGeneration +} + +func decodedS3BucketRouteFilter(family uint32, routeStart, routeEnd []byte) func([]byte) bool { + return func(rawKey []byte) bool { + bucket, ok := decodedS3BucketName(family, rawKey) + if !ok { + return false + } + if keyInRouteRange(rawKey, routeStart, routeEnd) { + return true + } + return decodedS3BucketRouteSelected(bucket, routeStart, routeEnd) + } +} + +func decodedS3BucketRouteSelected(bucket string, routeStart, routeEnd []byte) bool { + bucketRouteStart := s3keys.RoutePrefixForBucketAnyGeneration(bucket) + return keyInRouteRange(bucketRouteStart, routeStart, routeEnd) +} + +func rangesIntersect(aStart, aEnd, bStart, bEnd []byte) bool { + if len(aEnd) > 0 && bytes.Compare(aEnd, bStart) <= 0 { + return false + } + if len(bEnd) > 0 && bytes.Compare(bEnd, aStart) <= 0 { + return false + } + return true +} + +func keyInRouteRange(key, start, end []byte) bool { + if bytes.Compare(key, start) < 0 { + return false + } + return len(end) == 0 || bytes.Compare(key, end) < 0 +} + +func decodedS3BucketName(family uint32, rawKey []byte) (string, bool) { + switch family { + case distribution.MigrationFamilyS3BucketMeta: + return s3keys.ParseBucketMetaKey(rawKey) + case distribution.MigrationFamilyS3BucketGeneration: + return s3keys.ParseBucketGenerationKey(rawKey) + default: + return "", false + } +} + +func cloneByteSlices(in [][]byte) [][]byte { + if len(in) == 0 { + return nil + } + out := make([][]byte, len(in)) + for i := range in { + out[i] = bytes.Clone(in[i]) + } + return out +} + +func protoMVCCVersionsFromStore(in []store.MVCCVersion) []*pb.MVCCVersion { + out := make([]*pb.MVCCVersion, 0, len(in)) + for _, version := range in { + out = append(out, &pb.MVCCVersion{ + Key: bytes.Clone(version.Key), + CommitTs: version.CommitTS, + Tombstone: version.Tombstone, + Value: bytes.Clone(version.Value), + KeyFamily: version.KeyFamily, + ExpireAt: version.ExpireAt, + }) + } + return out +} + func (i *Internal) stampTimestamps(ctx context.Context, req *pb.ForwardRequest) (uint64, error) { if req == nil { return 0, nil diff --git a/adapter/internal_migration_test.go b/adapter/internal_migration_test.go new file mode 100644 index 000000000..47448aaf9 --- /dev/null +++ b/adapter/internal_migration_test.go @@ -0,0 +1,890 @@ +package adapter + +import ( + "bytes" + "context" + "encoding/binary" + "testing" + + "github.com/bootjp/elastickv/distribution" + internalutil "github.com/bootjp/elastickv/internal" + "github.com/bootjp/elastickv/internal/fskeys" + "github.com/bootjp/elastickv/internal/raftengine" + "github.com/bootjp/elastickv/internal/s3keys" + "github.com/bootjp/elastickv/kv" + pb "github.com/bootjp/elastickv/proto" + "github.com/bootjp/elastickv/store" + "github.com/stretchr/testify/require" + "google.golang.org/grpc" + "google.golang.org/grpc/codes" + "google.golang.org/grpc/status" +) + +type mockInternalLeader struct { + raftengine.LeaderView +} + +func (mockInternalLeader) State() raftengine.State { + return raftengine.StateLeader +} + +func (mockInternalLeader) Leader() raftengine.LeaderInfo { + return raftengine.LeaderInfo{ID: "n1", Address: "127.0.0.1:50051"} +} + +func (mockInternalLeader) VerifyLeader(context.Context) error { + return nil +} + +func (mockInternalLeader) LinearizableRead(context.Context) (uint64, error) { + return 1, nil +} + +type recordingInternalLeader struct { + mockInternalLeader + linearizableReadErr error + linearizableReadCalls int + verifyLeaderCalls int +} + +func (l *recordingInternalLeader) VerifyLeader(context.Context) error { + l.verifyLeaderCalls++ + return nil +} + +func (l *recordingInternalLeader) LinearizableRead(context.Context) (uint64, error) { + l.linearizableReadCalls++ + return 7, l.linearizableReadErr +} + +type applyingMigrationProposer struct { + fsm raftengine.StateMachine + calls uint64 +} + +func (p *applyingMigrationProposer) Propose(_ context.Context, data []byte) (*raftengine.ProposalResult, error) { + p.calls++ + return &raftengine.ProposalResult{ + CommitIndex: p.calls, + Response: p.fsm.Apply(data), + }, nil +} + +func (p *applyingMigrationProposer) ProposeAdmin(ctx context.Context, data []byte) (*raftengine.ProposalResult, error) { + return p.Propose(ctx, data) +} + +type captureExportRangeVersionsStream struct { + grpc.ServerStream + ctx context.Context + responses []*pb.ExportRangeVersionsResponse +} + +const ( + testExportCursorTagEmitted byte = iota + testExportCursorTagScanned + testExportCursorTagPrunedKey + testExportCursorTagSkippedKey +) + +func (s *captureExportRangeVersionsStream) Context() context.Context { + if s.ctx != nil { + return s.ctx + } + return context.Background() +} + +func (s *captureExportRangeVersionsStream) Send(resp *pb.ExportRangeVersionsResponse) error { + s.responses = append(s.responses, resp) + return nil +} + +func encodeTestExportCursor(key []byte, commitTS uint64, tag byte) []byte { + var out []byte + out = binary.AppendUvarint(out, uint64(len(key))) + out = append(out, key...) + out = binary.AppendUvarint(out, commitTS) + out = append(out, tag) + return out +} + +func testPrefixScanEnd(prefix []byte) []byte { + out := append([]byte(nil), prefix...) + for i := len(out) - 1; i >= 0; i-- { + if out[i] != 0xFF { + out[i]++ + return out[:i+1] + } + } + return nil +} + +func TestInternalExportRangeVersionsUsesStoreAndRouteFilter(t *testing.T) { + t.Parallel() + + ctx := context.Background() + st := store.NewMVCCStore() + require.NoError(t, st.PutAt(ctx, []byte("a"), []byte("va"), 10, 0)) + require.NoError(t, st.PutAt(ctx, []byte("z"), []byte("vz"), 10, 0)) + require.NoError(t, st.PutAt(ctx, []byte("!txn|int|a"), []byte("intent"), 10, 0)) + internal := NewInternalWithEngine(nil, mockInternalLeader{}, nil, nil, WithInternalStore(st)) + stream := &captureExportRangeVersionsStream{ctx: ctx} + + err := internal.ExportRangeVersions(&pb.ExportRangeVersionsRequest{ + MaxCommitTs: 20, + RouteStart: []byte("a"), + RouteEnd: []byte("b"), + KeyFamily: distribution.MigrationFamilyUser, + ExcludeKnownInternal: true, + RangeStart: []byte(""), + RangeEnd: []byte("z"), + MaxScannedBytes: 1 << 20, + ExcludePrefixes: [][]byte{[]byte("!custom|")}, + }, stream) + require.NoError(t, err) + require.Len(t, stream.responses, 1) + require.True(t, stream.responses[0].GetDone()) + require.Empty(t, stream.responses[0].GetNextCursor()) + require.Equal(t, []*pb.MVCCVersion{ + {Key: []byte("a"), CommitTs: 10, Value: []byte("va"), KeyFamily: distribution.MigrationFamilyUser}, + }, stream.responses[0].GetVersions()) +} + +func TestInternalExportRangeVersionsExcludesPeerLocalChunkBlobsFromUserFamily(t *testing.T) { + t.Parallel() + + ctx := context.Background() + st := store.NewMVCCStore() + var digest [32]byte + digest[0] = 0xab + blobKey := s3keys.ChunkBlobKey(digest) + require.NoError(t, st.PutAt(ctx, blobKey, []byte("peer-local"), 10, 0)) + internal := NewInternalWithEngine(nil, mockInternalLeader{}, nil, nil, WithInternalStore(st)) + stream := &captureExportRangeVersionsStream{ctx: ctx} + + err := internal.ExportRangeVersions(&pb.ExportRangeVersionsRequest{ + MaxCommitTs: 20, + RouteStart: []byte(s3keys.ChunkBlobPrefix), + RouteEnd: []byte("!s4|"), + KeyFamily: distribution.MigrationFamilyUser, + RangeStart: []byte(s3keys.ChunkBlobPrefix), + RangeEnd: testPrefixScanEnd([]byte(s3keys.ChunkBlobPrefix)), + MaxScannedBytes: 1 << 20, + }, stream) + require.NoError(t, err) + require.Len(t, stream.responses, 1) + require.True(t, stream.responses[0].GetDone()) + require.Empty(t, stream.responses[0].GetVersions()) +} + +func TestInternalExportRangeVersionsUsesValueAwareLegacyListDeltaRouteFilter(t *testing.T) { + t.Parallel() + + ctx := context.Background() + st := store.NewMVCCStore() + userKey := []byte("target-list") + key := legacyListMetaDeltaKey(userKey, 10) + value := store.MarshalListMetaDelta(store.ListMetaDelta{LenDelta: 1}) + require.NoError(t, st.PutAt(ctx, key, value, 10, 0)) + internal := NewInternalWithEngine(nil, mockInternalLeader{}, nil, nil, WithInternalStore(st)) + stream := &captureExportRangeVersionsStream{ctx: ctx} + + err := internal.ExportRangeVersions(&pb.ExportRangeVersionsRequest{ + MaxCommitTs: 20, + RouteStart: []byte("target"), + RouteEnd: []byte("target-list\x00"), + KeyFamily: distribution.MigrationFamilyLegacyListMetaDelta, + RangeStart: []byte(store.LegacyListMetaDeltaPrefix), + RangeEnd: testPrefixScanEnd([]byte(store.LegacyListMetaDeltaPrefix)), + MaxScannedBytes: 1 << 20, + }, stream) + require.NoError(t, err) + require.Len(t, stream.responses, 1) + require.Equal(t, []*pb.MVCCVersion{ + {Key: key, CommitTs: 10, Value: value, KeyFamily: distribution.MigrationFamilyLegacyListMetaDelta}, + }, stream.responses[0].GetVersions()) +} + +func TestInternalExportRangeVersionsUsesAppliedReadFence(t *testing.T) { + t.Parallel() + + ctx := context.Background() + st := store.NewMVCCStore() + require.NoError(t, st.PutAt(ctx, []byte("a"), []byte("va"), 10, 0)) + leader := &recordingInternalLeader{} + internal := NewInternalWithEngine(nil, leader, nil, nil, WithInternalStore(st)) + stream := &captureExportRangeVersionsStream{ctx: ctx} + + err := internal.ExportRangeVersions(&pb.ExportRangeVersionsRequest{ + MaxCommitTs: 20, + RouteStart: []byte("a"), + RouteEnd: []byte("b"), + KeyFamily: distribution.MigrationFamilyUser, + RangeStart: []byte("a"), + RangeEnd: []byte("b"), + MaxScannedBytes: 1 << 20, + }, stream) + require.NoError(t, err) + require.Equal(t, 1, leader.linearizableReadCalls) + require.Zero(t, leader.verifyLeaderCalls) + require.Len(t, stream.responses, 1) + require.True(t, stream.responses[0].GetDone()) +} + +func TestInternalExportRangeVersionsFailsClosedWhenAppliedReadFenceFails(t *testing.T) { + t.Parallel() + + leader := &recordingInternalLeader{linearizableReadErr: context.Canceled} + internal := NewInternalWithEngine(nil, leader, nil, nil, WithInternalStore(store.NewMVCCStore())) + stream := &captureExportRangeVersionsStream{ctx: context.Background()} + + err := internal.ExportRangeVersions(&pb.ExportRangeVersionsRequest{ + MaxCommitTs: 20, + RouteStart: []byte("a"), + RouteEnd: []byte("b"), + KeyFamily: distribution.MigrationFamilyUser, + RangeStart: []byte("a"), + RangeEnd: []byte("b"), + MaxScannedBytes: 1 << 20, + }, stream) + require.ErrorIs(t, err, context.Canceled) + require.Equal(t, 1, leader.linearizableReadCalls) + require.Empty(t, stream.responses) +} + +func TestInternalExportRangeVersionsRejectsUnboundedExport(t *testing.T) { + t.Parallel() + + internal := NewInternalWithEngine(nil, mockInternalLeader{}, nil, nil, WithInternalStore(store.NewMVCCStore())) + stream := &captureExportRangeVersionsStream{ctx: context.Background()} + + err := internal.ExportRangeVersions(&pb.ExportRangeVersionsRequest{ + KeyFamily: distribution.MigrationFamilyUser, + }, stream) + require.Error(t, err) + require.Equal(t, codes.InvalidArgument, status.Code(err)) + + stream = &captureExportRangeVersionsStream{ctx: context.Background()} + err = internal.ExportRangeVersions(&pb.ExportRangeVersionsRequest{ + MaxCommitTs: 20, + KeyFamily: distribution.MigrationFamilyUser, + }, stream) + require.Error(t, err) + require.Equal(t, codes.InvalidArgument, status.Code(err)) + + stream = &captureExportRangeVersionsStream{ctx: context.Background()} + err = internal.ExportRangeVersions(&pb.ExportRangeVersionsRequest{ + MaxCommitTs: 20, + KeyFamily: distribution.MigrationFamilyUser, + RangeStart: []byte("a"), + RangeEnd: []byte("b"), + }, stream) + require.Error(t, err) + require.Equal(t, codes.InvalidArgument, status.Code(err)) +} + +func TestInternalExportRangeVersionsUsesDecodedS3BucketRouteFilter(t *testing.T) { + t.Parallel() + + ctx := context.Background() + st := store.NewMVCCStore() + internal := NewInternalWithEngine(nil, mockInternalLeader{}, nil, nil, WithInternalStore(st)) + + for _, tc := range []struct { + name string + family uint32 + prefix string + keyFor func(string) []byte + value []byte + routeStart []byte + routeEnd []byte + }{ + { + name: "bucket meta", + family: distribution.MigrationFamilyS3BucketMeta, + prefix: s3keys.BucketMetaPrefix, + keyFor: s3keys.BucketMetaKey, + value: []byte("meta"), + routeStart: s3keys.RoutePrefixForBucketAnyGeneration("bucket-b"), + routeEnd: testPrefixScanEnd(s3keys.RoutePrefixForBucketAnyGeneration("bucket-b")), + }, + { + name: "bucket generation", + family: distribution.MigrationFamilyS3BucketGeneration, + prefix: s3keys.BucketGenerationPrefix, + keyFor: s3keys.BucketGenerationKey, + value: []byte("generation"), + routeStart: s3keys.RoutePrefixForBucketAnyGeneration("bucket-b"), + routeEnd: testPrefixScanEnd(s3keys.RoutePrefixForBucketAnyGeneration("bucket-b")), + }, + } { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + + inRouteKey := tc.keyFor("bucket-b") + outRouteKey := tc.keyFor("bucket-a") + require.NoError(t, st.PutAt(ctx, inRouteKey, tc.value, 10, 0)) + require.NoError(t, st.PutAt(ctx, outRouteKey, []byte("skip"), 10, 0)) + + stream := &captureExportRangeVersionsStream{ctx: ctx} + err := internal.ExportRangeVersions(&pb.ExportRangeVersionsRequest{ + MaxCommitTs: 20, + RouteStart: tc.routeStart, + RouteEnd: tc.routeEnd, + KeyFamily: tc.family, + RangeStart: []byte(tc.prefix), + RangeEnd: testPrefixScanEnd([]byte(tc.prefix)), + MaxScannedBytes: 1 << 20, + }, stream) + require.NoError(t, err) + require.Len(t, stream.responses, 1) + require.Equal(t, []*pb.MVCCVersion{ + {Key: inRouteKey, CommitTs: 10, Value: tc.value, KeyFamily: tc.family}, + }, stream.responses[0].GetVersions()) + }) + } +} + +func TestInternalExportRangeVersionsDecodedS3MixedRoutePreservesRawRangeMatches(t *testing.T) { + t.Parallel() + + ctx := context.Background() + st := store.NewMVCCStore() + internal := NewInternalWithEngine(nil, mockInternalLeader{}, nil, nil, WithInternalStore(st)) + + inRouteKey := s3keys.BucketMetaKey("bucket-z") + rawRangeKey := s3keys.BucketMetaKey("bucket-a") + require.NoError(t, st.PutAt(ctx, inRouteKey, []byte("meta-z"), 10, 0)) + require.NoError(t, st.PutAt(ctx, rawRangeKey, []byte("meta-a"), 10, 0)) + + stream := &captureExportRangeVersionsStream{ctx: ctx} + err := internal.ExportRangeVersions(&pb.ExportRangeVersionsRequest{ + MaxCommitTs: 20, + RouteStart: s3keys.RoutePrefixForBucketAnyGeneration("bucket-z"), + RouteEnd: []byte{}, + KeyFamily: distribution.MigrationFamilyS3BucketMeta, + RangeStart: []byte(s3keys.BucketMetaPrefix), + RangeEnd: testPrefixScanEnd([]byte(s3keys.BucketMetaPrefix)), + MaxScannedBytes: 1 << 20, + }, stream) + require.NoError(t, err) + require.Len(t, stream.responses, 1) + require.Equal(t, []*pb.MVCCVersion{ + {Key: rawRangeKey, CommitTs: 10, Value: []byte("meta-a"), KeyFamily: distribution.MigrationFamilyS3BucketMeta}, + {Key: inRouteKey, CommitTs: 10, Value: []byte("meta-z"), KeyFamily: distribution.MigrationFamilyS3BucketMeta}, + }, stream.responses[0].GetVersions()) +} + +func TestInternalExportRangeVersionsSkipsS3BucketAuxiliaryForNonOwnerRouteSlice(t *testing.T) { + t.Parallel() + + ctx := context.Background() + st := store.NewMVCCStore() + internal := NewInternalWithEngine(nil, mockInternalLeader{}, nil, nil, WithInternalStore(st)) + + const bucket = "bucket-b" + for _, tc := range []struct { + name string + family uint32 + prefix string + key []byte + value []byte + }{ + { + name: "bucket meta", + family: distribution.MigrationFamilyS3BucketMeta, + prefix: s3keys.BucketMetaPrefix, + key: s3keys.BucketMetaKey(bucket), + value: []byte("meta"), + }, + { + name: "bucket generation", + family: distribution.MigrationFamilyS3BucketGeneration, + prefix: s3keys.BucketGenerationPrefix, + key: s3keys.BucketGenerationKey(bucket), + value: []byte("generation"), + }, + } { + t.Run(tc.name, func(t *testing.T) { + require.NoError(t, st.PutAt(ctx, tc.key, tc.value, 10, 0)) + + stream := &captureExportRangeVersionsStream{ctx: ctx} + err := internal.ExportRangeVersions(&pb.ExportRangeVersionsRequest{ + MaxCommitTs: 20, + RouteStart: s3keys.RouteKey(bucket, 7, "m"), + RouteEnd: s3keys.RouteKey(bucket, 7, "z"), + KeyFamily: tc.family, + RangeStart: []byte(tc.prefix), + RangeEnd: testPrefixScanEnd([]byte(tc.prefix)), + MaxScannedBytes: 1 << 20, + }, stream) + require.NoError(t, err) + require.Len(t, stream.responses, 1) + require.Empty(t, stream.responses[0].GetVersions()) + }) + } +} + +func TestInternalExportRangeVersionsPreservesS3BucketRawRouteMatches(t *testing.T) { + t.Parallel() + + ctx := context.Background() + st := store.NewMVCCStore() + internal := NewInternalWithEngine(nil, mockInternalLeader{}, nil, nil, WithInternalStore(st)) + + key := s3keys.BucketMetaKey("bucket-raw") + require.NoError(t, st.PutAt(ctx, key, []byte("meta"), 10, 0)) + + stream := &captureExportRangeVersionsStream{ctx: ctx} + err := internal.ExportRangeVersions(&pb.ExportRangeVersionsRequest{ + MaxCommitTs: 20, + RouteStart: []byte("!s3|"), + RouteEnd: nil, + KeyFamily: distribution.MigrationFamilyS3BucketMeta, + RangeStart: []byte(s3keys.BucketMetaPrefix), + RangeEnd: testPrefixScanEnd([]byte(s3keys.BucketMetaPrefix)), + MaxScannedBytes: 1 << 20, + }, stream) + require.NoError(t, err) + require.Len(t, stream.responses, 1) + require.Equal(t, []*pb.MVCCVersion{ + {Key: key, CommitTs: 10, Value: []byte("meta"), KeyFamily: distribution.MigrationFamilyS3BucketMeta}, + }, stream.responses[0].GetVersions()) +} + +func TestInternalExportRangeVersionsUsesPartitionResolverGroup(t *testing.T) { + t.Parallel() + + ctx := context.Background() + st := store.NewMVCCStore() + resolver := NewSQSPartitionResolver(map[string][]uint64{ + "orders.fifo": {10, 11}, + }) + internal := NewInternalWithEngine(nil, mockInternalLeader{}, nil, nil, + WithInternalStore(st), + WithInternalMigrationExportRouting(11, resolver), + ) + + p0 := sqsPartitionedMsgDataKey("orders.fifo", 0, 1, "msg-0") + p1 := sqsPartitionedMsgDataKey("orders.fifo", 1, 1, "msg-1") + unknown := sqsPartitionedMsgDataKey("unknown.fifo", 0, 1, "msg-unknown") + require.NoError(t, st.PutAt(ctx, p0, []byte("p0"), 10, 0)) + require.NoError(t, st.PutAt(ctx, p1, []byte("p1"), 10, 0)) + require.NoError(t, st.PutAt(ctx, unknown, []byte("unknown"), 10, 0)) + + stream := &captureExportRangeVersionsStream{ctx: ctx} + prefix := []byte(SqsPartitionedMsgDataPrefix) + err := internal.ExportRangeVersions(&pb.ExportRangeVersionsRequest{ + MaxCommitTs: 20, + KeyFamily: distribution.MigrationFamilySQSPartitionedMessageData, + RouteStart: []byte("!sqs|route|global"), + RouteEnd: testPrefixScanEnd([]byte("!sqs|route|global")), + RangeStart: prefix, + RangeEnd: testPrefixScanEnd(prefix), + MaxScannedBytes: 1 << 20, + }, stream) + require.NoError(t, err) + require.Len(t, stream.responses, 1) + require.Equal(t, []*pb.MVCCVersion{ + {Key: p1, CommitTs: 10, Value: []byte("p1"), KeyFamily: distribution.MigrationFamilySQSPartitionedMessageData}, + }, stream.responses[0].GetVersions()) +} + +func TestInternalExportRangeVersionsResolvesTxnWrappedPartitionedSQSKeysByEmbeddedOwner(t *testing.T) { + t.Parallel() + + wrapWithStartTS := func(prefix string, userKey []byte) []byte { + var raw [8]byte + binary.BigEndian.PutUint64(raw[:], 77) + key := append([]byte(prefix), userKey...) + return append(key, raw[:]...) + } + cases := []struct { + name string + family uint32 + prefix []byte + wrap func([]byte) []byte + }{ + { + name: "intent", + family: distribution.MigrationFamilyTxnIntent, + prefix: []byte("!txn|int|"), + wrap: func(userKey []byte) []byte { + return append([]byte("!txn|int|"), userKey...) + }, + }, + { + name: "commit", + family: distribution.MigrationFamilyTxnCommit, + prefix: []byte("!txn|cmt|"), + wrap: func(userKey []byte) []byte { + return wrapWithStartTS("!txn|cmt|", userKey) + }, + }, + { + name: "rollback", + family: distribution.MigrationFamilyTxnRollback, + prefix: []byte("!txn|rb|"), + wrap: func(userKey []byte) []byte { + return wrapWithStartTS("!txn|rb|", userKey) + }, + }, + { + name: "success", + family: distribution.MigrationFamilyTxnSuccess, + prefix: []byte("!txn|ok|"), + wrap: func(userKey []byte) []byte { + return kv.TxnSuccessMarkerKey(userKey, 77, 88, []byte("primary")) + }, + }, + { + name: "meta", + family: distribution.MigrationFamilyTxnMeta, + prefix: []byte(kv.TxnMetaPrefix), + wrap: func(userKey []byte) []byte { + return append([]byte(kv.TxnMetaPrefix), userKey...) + }, + }, + } + + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + + ctx := context.Background() + st := store.NewMVCCStore() + resolver := NewSQSPartitionResolver(map[string][]uint64{ + "orders.fifo": {10, 11}, + }) + internal := NewInternalWithEngine(nil, mockInternalLeader{}, nil, nil, + WithInternalStore(st), + WithInternalMigrationExportRouting(11, resolver), + ) + + p0 := tc.wrap(sqsPartitionedMsgDataKey("orders.fifo", 0, 1, "msg-0")) + p1 := tc.wrap(sqsPartitionedMsgDataKey("orders.fifo", 1, 1, "msg-1")) + unknown := tc.wrap(sqsPartitionedMsgDataKey("unknown.fifo", 0, 1, "msg-unknown")) + require.NoError(t, st.PutAt(ctx, p0, []byte("p0"), 10, 0)) + require.NoError(t, st.PutAt(ctx, p1, []byte("p1"), 10, 0)) + require.NoError(t, st.PutAt(ctx, unknown, []byte("unknown"), 10, 0)) + + stream := &captureExportRangeVersionsStream{ctx: ctx} + err := internal.ExportRangeVersions(&pb.ExportRangeVersionsRequest{ + MaxCommitTs: 20, + KeyFamily: tc.family, + RouteStart: []byte("!sqs|route|global"), + RouteEnd: testPrefixScanEnd([]byte("!sqs|route|global")), + RangeStart: tc.prefix, + RangeEnd: testPrefixScanEnd(tc.prefix), + MaxScannedBytes: 1 << 20, + }, stream) + require.NoError(t, err) + require.Len(t, stream.responses, 1) + require.Equal(t, []*pb.MVCCVersion{ + {Key: p1, CommitTs: 10, Value: []byte("p1"), KeyFamily: tc.family}, + }, stream.responses[0].GetVersions()) + }) + } +} + +func TestInternalExportRangeVersionsDerivesFilesystemChunkScanBoundsFromRouteBounds(t *testing.T) { + t.Parallel() + + ctx := context.Background() + st := store.NewMVCCStore() + internal := NewInternalWithEngine(nil, mockInternalLeader{}, nil, nil, WithInternalStore(st)) + + routeKey := fskeys.ChunkRouteKey(10, 20) + inRouteChunk := fskeys.ChunkKey(10, 20, 3) + outRouteChunk := fskeys.ChunkKey(10, 21, 3) + require.NoError(t, st.PutAt(ctx, inRouteChunk, []byte("chunk"), 10, 0)) + require.NoError(t, st.PutAt(ctx, outRouteChunk, []byte("skip"), 10, 0)) + + stream := &captureExportRangeVersionsStream{ctx: ctx} + err := internal.ExportRangeVersions(&pb.ExportRangeVersionsRequest{ + MaxCommitTs: 20, + RouteStart: routeKey, + RouteEnd: testPrefixScanEnd(routeKey), + KeyFamily: distribution.MigrationFamilyFilesystemChunk, + RangeStart: routeKey, + RangeEnd: testPrefixScanEnd(routeKey), + MaxScannedBytes: 1 << 20, + }, stream) + require.NoError(t, err) + require.Len(t, stream.responses, 1) + require.Equal(t, []*pb.MVCCVersion{ + {Key: inRouteChunk, CommitTs: 10, Value: []byte("chunk"), KeyFamily: distribution.MigrationFamilyFilesystemChunk}, + }, stream.responses[0].GetVersions()) +} + +func TestInternalExportRangeVersionsDerivesFilesystemUsageScanBoundsFromRouteBounds(t *testing.T) { + t.Parallel() + + ctx := context.Background() + st := store.NewMVCCStore() + internal := NewInternalWithEngine(nil, mockInternalLeader{}, nil, nil, WithInternalStore(st)) + + routeKey := fskeys.ChunkRouteKey(10, 20) + inRouteUsage := fskeys.UsageRouteKey(routeKey) + outRouteUsage := fskeys.UsageRouteKey(fskeys.ChunkRouteKey(10, 21)) + require.NoError(t, st.PutAt(ctx, inRouteUsage, []byte("usage"), 10, 0)) + require.NoError(t, st.PutAt(ctx, outRouteUsage, []byte("skip"), 10, 0)) + + stream := &captureExportRangeVersionsStream{ctx: ctx} + err := internal.ExportRangeVersions(&pb.ExportRangeVersionsRequest{ + MaxCommitTs: 20, + RouteStart: routeKey, + RouteEnd: testPrefixScanEnd(routeKey), + KeyFamily: distribution.MigrationFamilyFilesystemUsage, + RangeStart: routeKey, + RangeEnd: testPrefixScanEnd(routeKey), + MaxScannedBytes: 1 << 20, + }, stream) + require.NoError(t, err) + require.Len(t, stream.responses, 1) + require.Equal(t, []*pb.MVCCVersion{ + {Key: inRouteUsage, CommitTs: 10, Value: []byte("usage"), KeyFamily: distribution.MigrationFamilyFilesystemUsage}, + }, stream.responses[0].GetVersions()) +} + +func TestInternalImportRangeVersionsAppliesStoreBatch(t *testing.T) { + t.Parallel() + + ctx := context.Background() + st := store.NewMVCCStore() + clock := kv.NewHLC() + proposer := &applyingMigrationProposer{ + fsm: kv.NewKvFSMWithHLC(st, clock), + } + internal := NewInternalWithEngine(nil, mockInternalLeader{}, clock, nil, + WithInternalStore(st), + WithInternalMigrationProposer(proposer), + WithInternalMigrationImportGate(func(context.Context) error { return nil }), + ) + + resp, err := internal.ImportRangeVersions(ctx, &pb.ImportRangeVersionsRequest{ + JobId: 7, + BracketId: 3, + BatchSeq: 1, + Cursor: []byte("cursor-1"), + Versions: []*pb.MVCCVersion{ + {Key: []byte("k"), CommitTs: 30, Value: []byte("v"), ExpireAt: 100}, + }, + }) + require.NoError(t, err) + require.Equal(t, []byte("cursor-1"), resp.GetAckedCursor()) + require.Equal(t, uint64(1), proposer.calls) + + staged := distribution.MigrationStagedDataKey(7, []byte("k")) + got, err := st.GetAt(ctx, staged, 30) + require.NoError(t, err) + require.Equal(t, []byte("v"), got) + _, err = st.GetAt(ctx, []byte("k"), 30) + require.ErrorIs(t, err, store.ErrKeyNotFound) + floor, err := st.MigrationHLCFloor(ctx, 7) + require.NoError(t, err) + require.Equal(t, uint64(30), floor) + require.GreaterOrEqual(t, clock.Current(), uint64(30)) +} + +func TestInternalImportRangeVersionsRejectsWhenOpcodeGateClosed(t *testing.T) { + t.Parallel() + + ctx := context.Background() + st := store.NewMVCCStore() + clock := kv.NewHLC() + proposer := &applyingMigrationProposer{ + fsm: kv.NewKvFSMWithHLC(st, clock), + } + internal := NewInternalWithEngine(nil, mockInternalLeader{}, clock, nil, + WithInternalStore(st), + WithInternalMigrationProposer(proposer), + WithInternalMigrationImportGate(func(context.Context) error { + return status.Error(codes.FailedPrecondition, "migration import disabled for test") + }), + ) + + resp, err := internal.ImportRangeVersions(ctx, &pb.ImportRangeVersionsRequest{ + JobId: 7, + BracketId: 3, + BatchSeq: 1, + Cursor: []byte("cursor-1"), + Versions: []*pb.MVCCVersion{ + {Key: []byte("k"), CommitTs: 30, Value: []byte("v")}, + }, + }) + require.Nil(t, resp) + require.Error(t, err) + require.Equal(t, codes.FailedPrecondition, status.Code(err)) + require.Equal(t, uint64(0), proposer.calls) +} + +func TestInternalImportRangeVersionsRejectsMissingIdentifiers(t *testing.T) { + t.Parallel() + + internal := NewInternalWithEngine(nil, mockInternalLeader{}, nil, nil, + WithInternalMigrationProposer(&applyingMigrationProposer{}), + ) + + _, err := internal.ImportRangeVersions(context.Background(), &pb.ImportRangeVersionsRequest{ + BracketId: 1, + BatchSeq: 1, + }) + require.Error(t, err) + require.Equal(t, codes.InvalidArgument, status.Code(err)) + + _, err = internal.ImportRangeVersions(context.Background(), &pb.ImportRangeVersionsRequest{ + JobId: 1, + BatchSeq: 1, + }) + require.Error(t, err) + require.Equal(t, codes.InvalidArgument, status.Code(err)) +} + +func TestInternalPromoteStagedVersionsRejectsWhenOpcodeGateClosed(t *testing.T) { + t.Parallel() + + ctx := context.Background() + st := store.NewMVCCStore() + clock := kv.NewHLC() + proposer := &applyingMigrationProposer{ + fsm: kv.NewKvFSMWithHLC(st, clock), + } + internal := NewInternalWithEngine(nil, mockInternalLeader{}, clock, nil, + WithInternalStore(st), + WithInternalMigrationProposer(proposer), + WithInternalMigrationPromoteGate(func(context.Context) error { + return status.Error(codes.FailedPrecondition, "migration promote disabled for test") + }), + ) + + resp, err := internal.PromoteStagedVersions(ctx, &pb.PromoteStagedVersionsRequest{ + JobId: 7, + MaxVersions: 10, + }) + require.Nil(t, resp) + require.Error(t, err) + require.Equal(t, codes.FailedPrecondition, status.Code(err)) + require.Equal(t, uint64(0), proposer.calls) +} + +func TestInternalPromoteStagedVersionsRejectsInvalidCursorBeforePropose(t *testing.T) { + t.Parallel() + + ctx := context.Background() + for _, tc := range []struct { + name string + cursor []byte + }{ + {name: "malformed cursor", cursor: []byte{0xff}}, + { + name: "cursor outside job staged prefix", + cursor: encodeTestExportCursor(distribution.MigrationStagedDataKey(8, []byte("k")), 30, testExportCursorTagEmitted), + }, + { + name: "scanned cursor inside staged prefix", + cursor: encodeTestExportCursor(distribution.MigrationStagedDataKey(7, []byte("k")), 31, testExportCursorTagScanned), + }, + { + name: "pruned-key cursor inside staged prefix", + cursor: encodeTestExportCursor(distribution.MigrationStagedDataKey(7, []byte("k")), 32, testExportCursorTagPrunedKey), + }, + { + name: "skipped-key cursor inside staged prefix", + cursor: encodeTestExportCursor(distribution.MigrationStagedDataKey(7, []byte("k")), 33, testExportCursorTagSkippedKey), + }, + } { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + + st := store.NewMVCCStore() + clock := kv.NewHLC() + proposer := &applyingMigrationProposer{ + fsm: kv.NewKvFSMWithHLC(st, clock), + } + internal := NewInternalWithEngine(nil, mockInternalLeader{}, clock, nil, + WithInternalStore(st), + WithInternalMigrationProposer(proposer), + WithInternalMigrationPromoteGate(func(context.Context) error { return nil }), + ) + + resp, err := internal.PromoteStagedVersions(ctx, &pb.PromoteStagedVersionsRequest{ + JobId: 7, + Cursor: tc.cursor, + MaxVersions: 10, + }) + require.Nil(t, resp) + require.Error(t, err) + require.Equal(t, codes.InvalidArgument, status.Code(err)) + require.ErrorContains(t, err, store.ErrInvalidExportCursor.Error()) + require.Equal(t, uint64(0), proposer.calls) + }) + } +} + +func TestInternalPromoteStagedVersionsAppliesStoreBatch(t *testing.T) { + t.Parallel() + + ctx := context.Background() + st := store.NewMVCCStore() + clock := kv.NewHLC() + proposer := &applyingMigrationProposer{ + fsm: kv.NewKvFSMWithHLC(st, clock), + } + internal := NewInternalWithEngine(nil, mockInternalLeader{}, clock, nil, + WithInternalStore(st), + WithInternalMigrationProposer(proposer), + WithInternalMigrationPromoteGate(func(context.Context) error { return nil }), + ) + + staged := distribution.MigrationStagedDataKey(7, []byte("k")) + require.NoError(t, st.PutAt(ctx, staged, []byte("v"), 30, 0)) + + resp, err := internal.PromoteStagedVersions(ctx, &pb.PromoteStagedVersionsRequest{ + JobId: 7, + MaxVersions: 10, + }) + require.NoError(t, err) + require.True(t, resp.GetDone()) + require.Equal(t, uint64(1), resp.GetPromotedRows()) + require.Equal(t, uint64(30), resp.GetMaxPromotedTs()) + require.Equal(t, uint64(1), proposer.calls) + + got, err := st.GetAt(ctx, []byte("k"), 30) + require.NoError(t, err) + require.Equal(t, []byte("v"), got) + _, err = st.GetAt(ctx, staged, 30) + require.ErrorIs(t, err, store.ErrKeyNotFound) + require.GreaterOrEqual(t, clock.Current(), uint64(30)) +} + +// Storage accepts values well past the internal gRPC message limit, and the +// store's byte budget is only consulted after a version has been appended, so +// one oversized row can land in an otherwise bounded export page. Sending it +// fails with ResourceExhausted, and so does every retry of the same cursor: +// the bracket stalls forever on an error that never names the row. The export +// must refuse the page and say which key is too large. +func TestInternalExportRangeVersionsRefusesOversizedPage(t *testing.T) { + t.Parallel() + + ctx := context.Background() + st := store.NewMVCCStore() + huge := bytes.Repeat([]byte{'x'}, internalutil.GRPCMaxMessageBytes+1) + require.NoError(t, st.PutAt(ctx, []byte("a"), huge, 10, 0)) + srv := NewInternalWithEngine(nil, mockInternalLeader{}, nil, nil, WithInternalStore(st)) + stream := &captureExportRangeVersionsStream{ctx: ctx} + + err := srv.ExportRangeVersions(&pb.ExportRangeVersionsRequest{ + MaxCommitTs: 20, + RouteStart: []byte("a"), + RouteEnd: []byte("b"), + KeyFamily: distribution.MigrationFamilyUser, + RangeStart: []byte("a"), + RangeEnd: []byte("b"), + }, stream) + + require.Error(t, err) + require.Equal(t, codes.FailedPrecondition, status.Code(err)) + require.Contains(t, err.Error(), `"a"`) + require.Empty(t, stream.responses, "an unsendable page must not be sent") +} diff --git a/adapter/internal_test.go b/adapter/internal_test.go index c5fe3a2ef..299e42406 100644 --- a/adapter/internal_test.go +++ b/adapter/internal_test.go @@ -463,3 +463,41 @@ func (t *forwardObserverTxn) Commit(_ context.Context, reqs []*pb.Request) (*kv. func (t *forwardObserverTxn) Abort(context.Context, []*pb.Request) (*kv.TransactionResponse, error) { return &kv.TransactionResponse{}, nil } + +// The export defaults only fill in a bound the request left unset, so without a +// ceiling one ExportVersions call can scan and decode an arbitrary portion of +// the source store before producing its next streamed response. For a sparse +// family or route filter that accepts few rows, math.MaxUint64 removes the only +// work bound there is. +func TestExportRangeVersionsOptionsClampOversizedBounds(t *testing.T) { + t.Parallel() + + i := &Internal{} + opts := i.exportRangeVersionsOptions(&pb.ExportRangeVersionsRequest{ + ChunkBytes: ^uint32(0), + MaxScannedBytes: ^uint64(0), + }) + + require.Equal(t, uint64(maxMigrationExportChunkBytes), opts.MaxBytes) + require.Equal(t, uint64(maxMigrationExportScanBytes), opts.MaxScannedBytes) +} + +// Unset bounds still take the defaults, and a request under the ceiling passes +// through so a caller can still ask for smaller chunks. +func TestExportRangeVersionsOptionsKeepDefaultsAndSmallerRequests(t *testing.T) { + t.Parallel() + + i := &Internal{} + defaults := i.exportRangeVersionsOptions(&pb.ExportRangeVersionsRequest{}) + require.Equal(t, uint64(defaultMigrationExportChunkBytes), defaults.MaxBytes) + require.Equal(t, + uint64(defaultMigrationExportChunkBytes*defaultMigrationExportScanFactor), + defaults.MaxScannedBytes) + + smaller := i.exportRangeVersionsOptions(&pb.ExportRangeVersionsRequest{ + ChunkBytes: 4096, + MaxScannedBytes: 8192, + }) + require.Equal(t, uint64(4096), smaller.MaxBytes) + require.Equal(t, uint64(8192), smaller.MaxScannedBytes) +} diff --git a/adapter/redis_lua_context.go b/adapter/redis_lua_context.go index 7fddd142d..2938ec764 100644 --- a/adapter/redis_lua_context.go +++ b/adapter/redis_lua_context.go @@ -51,6 +51,11 @@ type luaScriptContext struct { // server-side probe (still correct, just not cached). negativeType map[string]bool + // rawTypeAtStart caches the TTL-unfiltered type observed at startTS. Commit + // planning needs this raw type to decide whether a logically absent key has + // expired physical rows that must be cleaned up before recreation. + rawTypeAtStart map[string]redisValueType + // keyTypeProbeCount counts how many times the server-side keyTypeAt // helper was invoked during this Eval. Only read by tests via // luaScriptContext methods; ordinary production code never reads it. @@ -283,22 +288,23 @@ func newLuaScriptContext(ctx context.Context, server *RedisServer) (*luaScriptCo } startTS := server.readTS() return &luaScriptContext{ - server: server, - startTS: startTS, - readPin: server.pinReadTS(startTS), - ctx: ctx, - touched: map[string]struct{}{}, - readKeys: map[string][]byte{}, - deleted: map[string]bool{}, - everDeleted: map[string]bool{}, - negativeType: map[string]bool{}, - strings: map[string]*luaStringState{}, - lists: map[string]*luaListState{}, - hashes: map[string]*luaHashState{}, - sets: map[string]*luaSetState{}, - zsets: map[string]*luaZSetState{}, - streams: map[string]*luaStreamState{}, - ttls: map[string]*luaTTLState{}, + server: server, + startTS: startTS, + readPin: server.pinReadTS(startTS), + ctx: ctx, + touched: map[string]struct{}{}, + readKeys: map[string][]byte{}, + deleted: map[string]bool{}, + everDeleted: map[string]bool{}, + negativeType: map[string]bool{}, + rawTypeAtStart: map[string]redisValueType{}, + strings: map[string]*luaStringState{}, + lists: map[string]*luaListState{}, + hashes: map[string]*luaHashState{}, + sets: map[string]*luaSetState{}, + zsets: map[string]*luaZSetState{}, + streams: map[string]*luaStreamState{}, + ttls: map[string]*luaTTLState{}, }, nil } @@ -544,24 +550,34 @@ func (c *luaScriptContext) keyType(key []byte) (redisValueType, error) { } c.keyTypeProbeCount++ - typ, err := c.server.keyTypeAt(c.scriptCtx(), key, c.startTS) + rawTyp, err := c.rawStartTypeForCommitPlan(c.scriptCtx(), key) if err != nil { return redisTypeNone, err } - if typ == redisTypeNone && len(c.negativeType) < maxNegativeTypeCacheEntries { - // Pin the absence result for the rest of this Eval so repeated - // BullMQ-style polling of a missing key (e.g. a "delayed" zset) - // does not re-run the ~8-seek rawKeyTypeAt probe on every - // redis.call. - // - // Bounded to keep adversarial scripts from growing the map - // unboundedly; once full, subsequent misses correctly fall - // through to the server probe without caching. - c.negativeType[string(key)] = true + typ, err := c.server.applyTTLFilter(c.scriptCtx(), key, c.startTS, rawTyp) + if err != nil { + return redisTypeNone, err + } + if typ == redisTypeNone { + c.rememberNegativeType(key) } return typ, nil } +func (c *luaScriptContext) rememberNegativeType(key []byte) { + if len(c.negativeType) >= maxNegativeTypeCacheEntries { + return + } + // Pin the absence result for the rest of this Eval so repeated + // BullMQ-style polling of a missing key (e.g. a "delayed" zset) does + // not re-run the ~8-seek rawKeyTypeAt probe on every redis.call. + // + // Bounded to keep adversarial scripts from growing the map + // unboundedly; once full, subsequent misses correctly fall through to + // the server probe without caching. + c.negativeType[string(key)] = true +} + func (c *luaScriptContext) ensureKeyNotExpired(key []byte) error { ttl, err := c.loadTTL(key) if err != nil { @@ -3749,7 +3765,7 @@ func (c *luaScriptContext) commitPlanForKey(ctx context.Context, key string, com } keyBytes := []byte(key) - rawStartType, err := c.server.rawKeyTypeAt(ctx, keyBytes, c.startTS) + rawStartType, err := c.rawStartTypeForCommitPlan(ctx, keyBytes) if err != nil { return luaKeyPlan{}, err } @@ -3783,6 +3799,29 @@ func (c *luaScriptContext) commitPlanForKey(ctx context.Context, key string, com }, nil } +func (c *luaScriptContext) rawStartTypeForCommitPlan(ctx context.Context, key []byte) (redisValueType, error) { + k := string(key) + if typ, ok := c.rawTypeAtStart[k]; ok { + return typ, nil + } + typ, err := c.server.rawKeyTypeAt(ctx, key, c.startTS) + if err != nil { + return redisTypeNone, err + } + c.rememberRawTypeAtStart(k, typ) + return typ, nil +} + +func (c *luaScriptContext) rememberRawTypeAtStart(key string, typ redisValueType) { + if c.rawTypeAtStart == nil { + c.rawTypeAtStart = map[string]redisValueType{} + } + if _, ok := c.rawTypeAtStart[key]; !ok && len(c.rawTypeAtStart) >= maxNegativeTypeCacheEntries { + return + } + c.rawTypeAtStart[key] = typ +} + func luaWideFenceReadKeysForPlan(key []byte, finalType, startType redisValueType, preserveExisting bool) [][]byte { if !preserveExisting || startType == redisTypeNone { return redisTxnWideCollectionFenceKeys(key) diff --git a/adapter/redis_lua_negative_type_cache_test.go b/adapter/redis_lua_negative_type_cache_test.go index 6253f7e30..2085543e3 100644 --- a/adapter/redis_lua_negative_type_cache_test.go +++ b/adapter/redis_lua_negative_type_cache_test.go @@ -149,44 +149,55 @@ func TestLuaNegativeTypeCache_SingleProbePerKey(t *testing.T) { // itself stays bounded. func TestLuaNegativeTypeCache_BoundedSize(t *testing.T) { t.Parallel() - nodes, _, _ := createNode(t, 3) - defer shutdown(nodes) - ctx := context.Background() - sc, err := newLuaScriptContext(ctx, nodes[0].redisServer) - require.NoError(t, err) - defer sc.Close() + sc := &luaScriptContext{negativeType: map[string]bool{}} // Probe cap+overflow unique missing keys. Each probe is a miss // (redisTypeNone); only the first `cap` should be memoized. const overflow = 50 total := maxNegativeTypeCacheEntries + overflow for i := 0; i < total; i++ { - typ, kerr := sc.keyType([]byte(fmt.Sprintf("lua:neg:cap:%d", i))) - require.NoError(t, kerr) - require.Equal(t, redisTypeNone, typ) + sc.rememberNegativeType([]byte(fmt.Sprintf("lua:neg:cap:%d", i))) } require.Equal(t, maxNegativeTypeCacheEntries, len(sc.negativeType), "negativeType map must be capped at maxNegativeTypeCacheEntries") - // Each unique key above required exactly one probe on first access. - require.Equal(t, total, sc.keyTypeProbeCount, - "each unique key must have triggered exactly one server probe") - - // Re-probing one of the first `cap` keys must hit the cache - // (no additional server probe). Re-probing an overflow key must - // miss the cache and issue another server probe. + // One of the first `cap` keys must be cached. An overflow key must + // remain uncached so keyType falls back to a server probe in real Eval + // execution while the map size stays bounded. cachedKey := []byte("lua:neg:cap:0") - _, kerr := sc.keyType(cachedKey) - require.NoError(t, kerr) - require.Equal(t, total, sc.keyTypeProbeCount, - "a key inserted before the cap must remain cached") + typ, ok := sc.cachedType(cachedKey) + require.True(t, ok, "a key inserted before the cap must remain cached") + require.Equal(t, redisTypeNone, typ) overflowKey := []byte(fmt.Sprintf("lua:neg:cap:%d", maxNegativeTypeCacheEntries+1)) - _, kerr = sc.keyType(overflowKey) - require.NoError(t, kerr) - require.Equal(t, total+1, sc.keyTypeProbeCount, - "a key probed after the cap was reached must fall back to the server probe") + _, ok = sc.cachedType(overflowKey) + require.False(t, ok, "a key probed after the cap must fall back to the server probe") + sc.rememberNegativeType(overflowKey) require.Equal(t, maxNegativeTypeCacheEntries, len(sc.negativeType), "fallback probe must NOT grow the bounded cache") } + +func TestLuaRawTypeAtStartCache_BoundedSize(t *testing.T) { + t.Parallel() + + sc := &luaScriptContext{rawTypeAtStart: map[string]redisValueType{}} + + const overflow = 50 + total := maxNegativeTypeCacheEntries + overflow + for i := 0; i < total; i++ { + sc.rememberRawTypeAtStart(fmt.Sprintf("lua:rawtype:cap:%d", i), redisTypeNone) + } + require.Equal(t, maxNegativeTypeCacheEntries, len(sc.rawTypeAtStart), + "rawTypeAtStart map must be capped at maxNegativeTypeCacheEntries") + + cachedKey := "lua:rawtype:cap:0" + require.Equal(t, redisTypeNone, sc.rawTypeAtStart[cachedKey]) + + overflowKey := fmt.Sprintf("lua:rawtype:cap:%d", maxNegativeTypeCacheEntries+1) + _, ok := sc.rawTypeAtStart[overflowKey] + require.False(t, ok, "a key probed after the cap must fall back to a later server probe") + sc.rememberRawTypeAtStart(overflowKey, redisTypeNone) + require.Equal(t, maxNegativeTypeCacheEntries, len(sc.rawTypeAtStart), + "fallback raw type probe must NOT grow the bounded cache") +} diff --git a/adapter/redis_retry_test.go b/adapter/redis_retry_test.go index a8fd428b9..fa25a175e 100644 --- a/adapter/redis_retry_test.go +++ b/adapter/redis_retry_test.go @@ -352,6 +352,34 @@ func TestRedisXAddDedupsLandedWireWriteConflict(t *testing.T) { require.Equal(t, int64(1), meta.Length, "the generated XADD entry must not be appended twice") } +func TestRedisXAddDedupRouteFenceRetryPreservesPriorProbe(t *testing.T) { + t.Parallel() + + st := store.NewMVCCStore() + coord := newDedupTestCoordinator(st, 1, true) + coord.routeFenceAtDispatch = 2 + srv := &RedisServer{ + store: st, + coordinator: coord, + scriptCache: map[string]string{}, + onePhaseTxnDedup: true, + } + conn := &recordingConn{} + + srv.xadd(conn, redcon.Command{Args: [][]byte{ + []byte(cmdXAdd), []byte("retry:stream"), []byte("*"), []byte("field"), []byte("value"), + }}) + + require.Empty(t, conn.err) + require.NotEmpty(t, conn.bulk) + require.Equal(t, 3, coord.dispatches, "attempt 1 landed, route-fenced reuse, then dedup probe retry") + require.Equal(t, 1, coord.probeNoOps, "route-fenced reuse must not replace the prior landed probe") + meta, found, err := srv.loadStreamMetaAt(context.Background(), []byte("retry:stream"), snapshotTS(coord.Clock(), st)) + require.NoError(t, err) + require.True(t, found) + require.Equal(t, int64(1), meta.Length) +} + func TestRedisXAddDedupDisabledDoesNotReplayLandedWireConflict(t *testing.T) { t.Parallel() diff --git a/adapter/redis_txn_test.go b/adapter/redis_txn_test.go index ae3cdbc33..ff52aa534 100644 --- a/adapter/redis_txn_test.go +++ b/adapter/redis_txn_test.go @@ -1746,20 +1746,21 @@ func (s *luaCleanupScanTrackingStore) ScanAt(ctx context.Context, start []byte, func newLuaCommitPlanTestContext(server *RedisServer, startTS uint64) *luaScriptContext { return &luaScriptContext{ - server: server, - startTS: startTS, - touched: map[string]struct{}{}, - readKeys: map[string][]byte{}, - deleted: map[string]bool{}, - everDeleted: map[string]bool{}, - negativeType: map[string]bool{}, - strings: map[string]*luaStringState{}, - lists: map[string]*luaListState{}, - hashes: map[string]*luaHashState{}, - sets: map[string]*luaSetState{}, - zsets: map[string]*luaZSetState{}, - streams: map[string]*luaStreamState{}, - ttls: map[string]*luaTTLState{}, + server: server, + startTS: startTS, + touched: map[string]struct{}{}, + readKeys: map[string][]byte{}, + deleted: map[string]bool{}, + everDeleted: map[string]bool{}, + negativeType: map[string]bool{}, + rawTypeAtStart: map[string]redisValueType{}, + strings: map[string]*luaStringState{}, + lists: map[string]*luaListState{}, + hashes: map[string]*luaHashState{}, + sets: map[string]*luaSetState{}, + zsets: map[string]*luaZSetState{}, + streams: map[string]*luaStreamState{}, + ttls: map[string]*luaTTLState{}, } } @@ -1773,6 +1774,8 @@ func TestLuaCommitPlanForAbsentRewriteSkipsFullLogicalCleanupScans(t *testing.T) key := "lua:absent-rewrite" scriptCtx := newLuaCommitPlanTestContext(server, 10) + scriptCtx.negativeType[key] = true + scriptCtx.rawTypeAtStart[key] = redisTypeNone scriptCtx.strings[key] = &luaStringState{loaded: true, exists: true, dirty: true, value: []byte("v")} scriptCtx.ttls[key] = &luaTTLState{loaded: true} diff --git a/adapter/s3_test.go b/adapter/s3_test.go index 5ef2897b9..f2b4ee950 100644 --- a/adapter/s3_test.go +++ b/adapter/s3_test.go @@ -799,6 +799,9 @@ func TestS3Server_ShardedStoreRoutesBucketAndObjectData(t *testing.T) { readTS := shardStore.LastCommitTS() var err error _, err = store1.GetAt(ctx, s3keys.BucketMetaKey("bucket-a"), readTS) + require.ErrorIs(t, err, store.ErrKeyNotFound) + + _, err = store2.GetAt(ctx, s3keys.BucketMetaKey("bucket-a"), readTS) require.NoError(t, err) _, err = store1.GetAt(ctx, s3keys.ObjectManifestKey("bucket-a", 1, "dir/file.txt"), readTS) diff --git a/adapter/test_util.go b/adapter/test_util.go index f4fb48041..aebb0dd31 100644 --- a/adapter/test_util.go +++ b/adapter/test_util.go @@ -645,7 +645,14 @@ func setupNodes(t *testing.T, ctx context.Context, n int, ports []portsAdress) ( } pb.RegisterRawKVServer(s, gs) pb.RegisterTransactionalKVServer(s, gs) - pb.RegisterInternalServer(s, NewInternalWithEngine(trx, result.Engine, coordinator.Clock(), relay)) + pb.RegisterInternalServer(s, NewInternalWithEngine( + trx, + result.Engine, + coordinator.Clock(), + relay, + WithInternalStore(st), + WithInternalMigrationProposer(result.Engine), + )) internalraftadmin.RegisterOperationalServices(opsCtx, s, result.Engine, []string{"Example"}) grpcAdders = append(grpcAdders, port.grpcAddress) diff --git a/distribution/engine.go b/distribution/engine.go index 67636e13b..d706b2dcb 100644 --- a/distribution/engine.go +++ b/distribution/engine.go @@ -25,11 +25,12 @@ type Route struct { GroupID uint64 // State tracks control-plane state for this route. State RouteState - // StagedVisibilityActive allows serving reads to merge staged migration rows. + // StagedVisibilityActive makes migrated versions visible through the + // staged/live merge path after cross-group CUTOVER. StagedVisibilityActive bool - // MigrationJobID identifies the active staged migration job. + // MigrationJobID identifies the migration job that owns staged visibility. MigrationJobID uint64 - // MinWriteTSExclusive rejects writes at or below the migration cutover floor. + // MinWriteTSExclusive is the post-migration write timestamp floor. MinWriteTSExclusive uint64 // Load tracks the number of accesses served by this range. Load uint64 @@ -541,17 +542,7 @@ func (e *Engine) Stats() []Route { defer e.mu.RUnlock() stats := make([]Route, len(e.routes)) for i, r := range e.routes { - stats[i] = Route{ - RouteID: r.RouteID, - Start: CloneBytes(r.Start), - End: CloneBytes(r.End), - GroupID: r.GroupID, - State: r.State, - StagedVisibilityActive: r.StagedVisibilityActive, - MigrationJobID: r.MigrationJobID, - MinWriteTSExclusive: r.MinWriteTSExclusive, - Load: r.Load, - } + stats[i] = cloneRoute(r) } return stats } @@ -587,17 +578,7 @@ func (e *Engine) intersectingRoutesLocked(start, end []byte) []Route { break } // Route intersects with scan range - result = append(result, Route{ - RouteID: r.RouteID, - Start: CloneBytes(r.Start), - End: CloneBytes(r.End), - GroupID: r.GroupID, - State: r.State, - StagedVisibilityActive: r.StagedVisibilityActive, - MigrationJobID: r.MigrationJobID, - MinWriteTSExclusive: r.MinWriteTSExclusive, - Load: r.Load, - }) + result = append(result, cloneRoute(*r)) } return result } diff --git a/distribution/engine_test.go b/distribution/engine_test.go index 90671bc98..9c8763ba0 100644 --- a/distribution/engine_test.go +++ b/distribution/engine_test.go @@ -239,7 +239,16 @@ func TestEngineApplySnapshot_ReplacesRoutesAndVersion(t *testing.T) { Version: 1, Routes: []RouteDescriptor{ {RouteID: 10, Start: []byte(""), End: []byte("m"), GroupID: 1, State: RouteStateActive}, - {RouteID: 11, Start: []byte("m"), End: nil, GroupID: 2, State: RouteStateWriteFenced}, + { + RouteID: 11, + Start: []byte("m"), + End: nil, + GroupID: 2, + State: RouteStateWriteFenced, + StagedVisibilityActive: true, + MigrationJobID: 42, + MinWriteTSExclusive: 99, + }, }, }) if err != nil { @@ -260,6 +269,19 @@ func TestEngineApplySnapshot_ReplacesRoutesAndVersion(t *testing.T) { if stats[1].RouteID != 11 || stats[1].State != RouteStateWriteFenced { t.Fatalf("unexpected second route metadata: %+v", stats[1]) } + assertStagedRouteMetadata(t, stats[1]) + route, ok := e.GetRoute([]byte("m")) + if !ok { + t.Fatalf("expected route for m") + } + assertStagedRouteMetadata(t, route) +} + +func assertStagedRouteMetadata(t *testing.T, route Route) { + t.Helper() + if !route.StagedVisibilityActive || route.MigrationJobID != 42 || route.MinWriteTSExclusive != 99 { + t.Fatalf("staged route metadata was not preserved: %+v", route) + } } func TestEngineRouteLookupsReturnMatchingCatalogVersion(t *testing.T) { diff --git a/distribution/migrator.go b/distribution/migrator.go index c22b9589c..e3041ee0b 100644 --- a/distribution/migrator.go +++ b/distribution/migrator.go @@ -2,6 +2,7 @@ package distribution import ( "bytes" + "encoding/binary" "github.com/bootjp/elastickv/internal/fskeys" "github.com/bootjp/elastickv/internal/s3keys" @@ -54,6 +55,9 @@ const ( MigrationFamilyS3GCUpload MigrationFamilyLegacyListMetaDelta MigrationFamilyS3ChunkRef + // MigrationFamilyFilesystemChunk is appended last on purpose: the values + // above are the migration wire contract, so inserting anywhere earlier + // would renumber families that peers already use. MigrationFamilyFilesystemChunk MigrationFamilyFilesystemUsage ) @@ -77,6 +81,7 @@ const ( migrationDynamoGenPrefix = "!ddb|meta|gen|" migrationDynamoItemPrefix = "!ddb|item|" migrationDynamoGSIPrefix = "!ddb|gsi|" + migrationStagedDataPrefix = "!dist|migstage|" ) const ( @@ -90,6 +95,8 @@ const ( migrationSQSMsgGroupPrefix = "!sqs|msg|group|" migrationSQSMsgByAgePrefix = "!sqs|msg|byage|" migrationSQSPartitionedSuffix = "p|" + migrationStagedDataJobIDBytes = 8 + migrationStagedDataSeparator = byte('|') ) var ( @@ -99,6 +106,70 @@ var ( ErrMigrationSourceRouteChanged = errors.New("migration source route does not match split job") ) +// IsReservedControlKey reports whether key lives in a control namespace that no +// user mutation may write. +// +// Staged migration data is included in the reserved set. It is populated by the +// typed migration import/promote paths and by FSM-internal prefix-delete +// expansion, not by externally supplied RawKV mutations. +// migrationOnlyControlPrefixes are the control namespaces that no coordinator +// dispatch ever writes: they are produced solely by the typed migration +// commands (import, promote, fence). The rest of !dist| is excluded, because +// the control plane commits route and job records through the transactional +// coordinator -- SplitRange does exactly that. +var migrationOnlyControlPrefixes = [][]byte{ + []byte(migrationStagedDataPrefix), + []byte("!migwrite|"), + []byte("!migfence|"), +} + +// IsMigrationOnlyControlKey reports whether key belongs to a namespace that +// only the typed migration commands may write. It is the transactional apply +// path's reserved-key test: the raw path can refuse every control namespace +// because nothing legitimate writes one as a raw mutation, but the catalog +// reaches the store through a transaction, so the transactional test has to be +// the narrower one. A staged-data row is the case that matters -- a client +// write landing there is promoted as user data. +func IsMigrationOnlyControlKey(key []byte) bool { + for _, prefix := range migrationOnlyControlPrefixes { + if bytes.HasPrefix(key, prefix) { + return true + } + } + return false +} + +func IsReservedControlKey(key []byte) bool { + return reservedControlPrefixIntersects(key, bytes.HasPrefix) +} + +// ReservedControlPrefixIntersects reports whether a DEL_PREFIX over prefix is +// aimed at a control namespace, either because the prefix sits inside one or +// because it is a partial spelling of one. +// +// An empty prefix is exempt. That is the whole-keyspace flush, an operation the +// caller asked for deliberately, and refusing it would break FLUSHDB. Keeping +// control keys out of a flush needs the store's prefix-exclusion argument to +// accept more than the one transaction prefix it takes today, which is a +// separate change. +func ReservedControlPrefixIntersects(prefix []byte) bool { + if len(prefix) == 0 { + return false + } + return reservedControlPrefixIntersects(prefix, func(a, b []byte) bool { + return bytes.HasPrefix(a, b) || bytes.HasPrefix(b, a) + }) +} + +func reservedControlPrefixIntersects(key []byte, match func(key, reserved []byte) bool) bool { + for _, reserved := range migrationReservedControlPrefixes { + if match(key, reserved) { + return true + } + } + return false +} + var migrationReservedControlPrefixes = [][]byte{ []byte("!dist|"), []byte("!migstage|"), @@ -145,6 +216,11 @@ var migrationInternalFamilyPrefixes = [][]byte{ []byte(s3keys.BlobPrefix), []byte(s3keys.ChunkRefPrefix), []byte(s3keys.GCUploadPrefix), + // Every family bracket's prefix belongs here so the user bracket does not + // also export those rows. A filesystem-chunk route interval can span both + // the raw !fs|chk| key and its normalized !fs|route|chk| route, and both + // route filters accept the row, so omitting this exported every chunk + // version twice under two bracket IDs. // Chunk blob payloads are peer-local: they are written straight to the // receiving node's Pebble instead of through Raft, and peers pull them by // digest over S3BlobFetch when they apply the matching chunkref. They are @@ -309,6 +385,38 @@ func (b MigrationBracket) containsFamilyShape(rawKey []byte) bool { } } +// MigrationStagedDataKey returns the target-local shadow key used while a +// cross-group split imports data before CUTOVER/promotion. +func MigrationStagedDataKey(jobID uint64, rawKey []byte) []byte { + key := make([]byte, len(migrationStagedDataPrefix)+migrationStagedDataJobIDBytes+1+len(rawKey)) + copy(key, migrationStagedDataPrefix) + binary.BigEndian.PutUint64(key[len(migrationStagedDataPrefix):], jobID) + key[len(migrationStagedDataPrefix)+migrationStagedDataJobIDBytes] = migrationStagedDataSeparator + copy(key[len(migrationStagedDataPrefix)+migrationStagedDataJobIDBytes+1:], rawKey) + return key +} + +// MigrationStagedDataKeyPrefix returns the prefix covering all staged data for +// one migration job. +func MigrationStagedDataKeyPrefix(jobID uint64) []byte { + return MigrationStagedDataKey(jobID, nil) +} + +func IsMigrationStagedDataKey(key []byte) bool { + return len(key) >= len(migrationStagedDataPrefix)+migrationStagedDataJobIDBytes+1 && + bytes.HasPrefix(key, []byte(migrationStagedDataPrefix)) && + key[len(migrationStagedDataPrefix)+migrationStagedDataJobIDBytes] == migrationStagedDataSeparator +} + +func MigrationStagedDataKeyParts(key []byte) (uint64, []byte, bool) { + if !IsMigrationStagedDataKey(key) { + return 0, nil, false + } + jobID := binary.BigEndian.Uint64(key[len(migrationStagedDataPrefix):]) + rawKey := bytes.Clone(key[len(migrationStagedDataPrefix)+migrationStagedDataJobIDBytes+1:]) + return jobID, rawKey, true +} + func (b MigrationBracket) containsDecodedS3Route(rawKey, routeStart, routeEnd []byte) bool { bucket, ok := b.decodedS3Bucket(rawKey) if !ok { @@ -470,7 +578,21 @@ func migrationFamilyBrackets() []MigrationBracket { {family: MigrationFamilyS3Blob, prefix: s3keys.BlobPrefix}, {family: MigrationFamilyS3ChunkRef, prefix: s3keys.ChunkRefPrefix}, {family: MigrationFamilyS3GCUpload, prefix: s3keys.GCUploadPrefix}, + // File chunk payloads live under !fs|chk| but route through a virtual + // !fs|route|chk| key via fskeys.ExtractRouteKey. "!fs|chk|" sorts below + // "!fs|route|chk|", so the user bracket's raw interval never reaches + // them and, without this bracket, a cross-group split completed and + // promoted while every chunk of the moved files stayed behind. The + // default route-key check in migrationBracketRouteCheck applies the + // logical route filter to the raw scan. {family: MigrationFamilyFilesystemChunk, prefix: string(fskeys.ChunkAllPrefix())}, + // Per-route usage counters are stored at !fs|usage|route| + // and normalize back to the embedded logical route key, so like chunks + // their raw key sits outside the user bracket's interval. Without this + // bracket the counter stayed on the source: after cutover the usage scan + // filtered that copy out because its logical owner had become the + // target, while target-side updates began from zero, so StatFS + // undercounted existing files and bytes. {family: MigrationFamilyFilesystemUsage, prefix: string(fskeys.UsageRouteAllPrefix())}, } diff --git a/distribution/migrator_export_plan_test.go b/distribution/migrator_export_plan_test.go index b63ecb423..0b4206427 100644 --- a/distribution/migrator_export_plan_test.go +++ b/distribution/migrator_export_plan_test.go @@ -465,6 +465,29 @@ func TestMigrationKnownInternalPrefixesAreConcreteOnly(t *testing.T) { require.False(t, bytes.Equal(prefixes[0], MigrationKnownInternalPrefixes()[0]), "prefix list must be cloned") } +func TestMigrationStagedDataKeyRoundTrip(t *testing.T) { + t.Parallel() + + raw := []byte("user|raw") + key := MigrationStagedDataKey(42, raw) + require.LessOrEqual(t, len(MigrationStagedDataKey(42, nil)), store.MaxSnapshotInternalKeyEnvelope) + require.True(t, IsMigrationStagedDataKey(key)) + require.True(t, bytes.HasPrefix(key, MigrationStagedDataKeyPrefix(42))) + require.False(t, IsMigrationStagedDataKey([]byte("!dist|migstage|short"))) + + jobID, original, ok := MigrationStagedDataKeyParts(key) + require.True(t, ok) + require.Equal(t, uint64(42), jobID) + require.Equal(t, []byte("user|raw"), original) + + raw[0] = 'X' + original[0] = 'Y' + jobID, original, ok = MigrationStagedDataKeyParts(key) + require.True(t, ok) + require.Equal(t, uint64(42), jobID) + require.Equal(t, []byte("user|raw"), original) +} + func TestValidateMigrationRouteRangeRejectsReservedControlPrefixes(t *testing.T) { t.Parallel() @@ -567,6 +590,100 @@ func legacyListMetaDeltaKey(userKey []byte, commitTS uint64, seqInTxn uint32) [] return append(key, seq[:]...) } +// File chunk payloads live under !fs|chk| but route through a virtual +// !fs|route|chk| key. "!fs|chk|" sorts below "!fs|route|chk|", so a +// filesystem-chunk route's user bracket -- whose raw interval IS the virtual +// route range -- never reaches the payloads. Without a dedicated bracket the +// export found nothing, yet the migration completed and was promoted, losing +// every chunk of the moved files on a cross-group split. +func TestPlanMigrationBracketsCoversFilesystemChunkPayloads(t *testing.T) { + t.Parallel() + + routeStart := fskeys.ChunkRouteKey(0, 1) + routeEnd := fskeys.ChunkRouteKey(0, 9) + brackets, err := PlanMigrationBrackets(routeStart, routeEnd) + require.NoError(t, err) + + var chunk *MigrationBracket + for i := range brackets { + if brackets[i].Family == MigrationFamilyFilesystemChunk { + chunk = &brackets[i] + + break + } + } + require.NotNil(t, chunk, "the plan must carry a filesystem chunk bracket") + require.Equal(t, fskeys.ChunkAllPrefix(), chunk.Start, + "the bracket must scan the raw chunk prefix, not the virtual route range") + require.True(t, chunk.RequiresRouteKeyCheck, + "raw chunk keys must still be filtered through the logical route") + + // The gap this closes: the raw payload prefix sorts below the virtual route + // interval, so the user bracket's raw range cannot reach it. + require.Negative(t, bytes.Compare(fskeys.ChunkAllPrefix(), routeStart), + "chunk payloads sort below the virtual route interval") +} + +// Invariant: every family bracket's scan prefix must also be excluded from the +// user bracket. Both filters accept the same raw row otherwise -- the family +// bracket by prefix and the user bracket by normalized route key -- so the rows +// are exported and proposed through Raft twice under separate bracket IDs. +// +// This is written as an invariant rather than a per-family case because the +// filesystem chunk family was added without its exclusion and nothing caught it. +func TestEveryFamilyBracketPrefixIsExcludedFromUserBracket(t *testing.T) { + t.Parallel() + + for _, bracket := range migrationFamilyBrackets() { + if bracket.DrainOnly { + continue + } + require.True(t, IsMigrationKnownInternalKey(bracket.Start), + "family %d prefix %q must be in migrationInternalFamilyPrefixes, "+ + "otherwise the user bracket exports the same rows a second time", + bracket.Family, bracket.Start) + } +} + +// Per-route usage counters live at !fs|usage|route| and +// normalize back to the embedded logical route key, so their raw key sits +// outside a user route's interval exactly like chunk payloads. Without a +// bracket the counter stayed on the source: after cutover the usage scan +// filtered that copy out because its logical owner had become the target, +// while target-side updates began from zero, so StatFS undercounted. +func TestPlanMigrationBracketsCoversFilesystemUsageCounters(t *testing.T) { + t.Parallel() + + routeStart := []byte("a") + routeEnd := []byte("z") + brackets, err := PlanMigrationBrackets(routeStart, routeEnd) + require.NoError(t, err) + + var usage *MigrationBracket + for i := range brackets { + if brackets[i].Family == MigrationFamilyFilesystemUsage { + usage = &brackets[i] + + break + } + } + require.NotNil(t, usage, "the plan must carry a filesystem usage bracket") + require.Equal(t, fskeys.UsageRouteAllPrefix(), usage.Start, + "the bracket must scan the raw usage prefix, not the logical route range") + require.True(t, usage.RequiresRouteKeyCheck, + "raw usage keys must still be filtered through their embedded route") + + // The gap this closes: the raw counter key sorts outside a user route + // interval, so the user bracket cannot reach it. + require.Negative(t, bytes.Compare(fskeys.UsageRouteAllPrefix(), routeStart), + "usage counters sort below a user route interval") + + // And it must round-trip: a counter for a key inside the interval + // normalizes back into that interval, so the route filter keeps it. + counter := fskeys.UsageRouteKey([]byte("customers")) + require.Equal(t, []byte("customers"), fskeys.ExtractRouteKey(counter)) +} + // !s3|chunkblob| rows never travel with a migration. They are written outside // Raft directly to the receiving node's Pebble and pulled by peers over // S3BlobFetch, and they are content-addressed, so a single row backs every diff --git a/internal/s3keys/keys_test.go b/internal/s3keys/keys_test.go index 1c4ea806f..95d8d68fa 100644 --- a/internal/s3keys/keys_test.go +++ b/internal/s3keys/keys_test.go @@ -30,6 +30,14 @@ func TestBucketGenerationKey_RoundTripsZeroByteSegments(t *testing.T) { require.Equal(t, bucket, parsed) } +func TestParseBucketGenerationKey_RejectsNonGenerationKey(t *testing.T) { + t.Parallel() + + parsed, ok := ParseBucketGenerationKey(BucketMetaKey("bucket")) + require.False(t, ok) + require.Empty(t, parsed) +} + func TestObjectManifestKey_RoundTripsZeroByteSegments(t *testing.T) { t.Parallel() diff --git a/kv/fsm.go b/kv/fsm.go index 38af3e66a..ec6faa0b0 100644 --- a/kv/fsm.go +++ b/kv/fsm.go @@ -10,7 +10,9 @@ import ( "os" "sync/atomic" + "github.com/bootjp/elastickv/distribution" "github.com/bootjp/elastickv/internal/encryption/fsmwire" + "github.com/bootjp/elastickv/internal/fskeys" "github.com/bootjp/elastickv/internal/raftengine" "github.com/bootjp/elastickv/internal/s3keys" pb "github.com/bootjp/elastickv/proto" @@ -139,6 +141,10 @@ type RouteSnapshot interface { // OwnerOf returns the Raft group ID that owned key at this // snapshot's version. (0, false) when no route covered key. OwnerOf(key []byte) (uint64, bool) + // RouteOf returns the complete route descriptor covering key. + RouteOf(key []byte) (distribution.Route, bool) + // IntersectingRoutes returns every route intersecting [start, end). + IntersectingRoutes(start, end []byte) []distribution.Route // WriteFencedForKey reports whether key is currently inside a // WriteFenced route in this snapshot. WriteFencedForKey(key []byte) bool @@ -306,6 +312,8 @@ var ErrUnknownRequestType = errors.New("unknown request type") // catches up to the promoted owner. var ErrRouteWriteFenced = errors.New("route is write-fenced; retry after route migration") +var ErrRouteWriteTimestampTooLow = errors.New("route write timestamp is below migration floor") + // ErrComposed1Violation is returned by verifyComposed1 when the // transaction's commit cannot proceed on this Raft group because the // txn's read-set or write-set keys are not owned by this group at @@ -346,12 +354,11 @@ func (f *kvFSM) Apply(data []byte) any { if f.backupFloorLoadErr != nil { return haltErr(errors.Wrap(errors.Mark(f.backupFloorLoadErr, ErrBackupApply), "kv/fsm: load backup timestamp floor")) } - if resp, handled := f.applyReservedOpcode(data); handled { + ctx := context.TODO() + if resp, handled := f.applyReservedOpcode(ctx, data); handled { return resp } - ctx := context.TODO() - reqs, err := decodeRaftRequests(data) if err != nil { return errors.WithStack(err) @@ -392,13 +399,17 @@ func (f *kvFSM) Apply(data []byte) any { // opcode with ErrEncryptionApply, which the engine's HaltApply seam // recognises as a halt — same fail-closed shape as the Stage 3 // raft-envelope unwrap path. -func (f *kvFSM) applyReservedOpcode(data []byte) (any, bool) { +func (f *kvFSM) applyReservedOpcode(ctx context.Context, data []byte) (any, bool) { if len(data) == 0 { return nil, false } switch { case data[0] == raftEncodeHLCLease: return f.applyHLCLease(data[1:]), true + case data[0] == raftEncodeMigrationImport: + return f.applyMigrationImport(ctx, data[1:]), true + case data[0] == raftEncodeMigrationPromote: + return f.applyMigrationPromote(ctx, data[1:]), true case data[0] == raftEncodeBackup: return f.applyBackup(data[1:]), true case data[0] >= fsmwire.OpEncryptionMin && data[0] <= fsmwire.OpEncryptionMax: @@ -432,6 +443,14 @@ const ( // These entries do not touch the MVCC store; they only advance the shared HLC // physicalCeiling so the logical counter can continue to increment in memory. raftEncodeHLCLease byte = 0x02 + // raftEncodeMigrationImport carries a target-group range-migration import + // batch. Every target voter applies the raw MVCC versions, import ack, and + // migration HLC floor before the RPC handler returns success. + raftEncodeMigrationImport byte = 0x09 + // raftEncodeMigrationPromote carries a target-group range-migration staged + // data promotion chunk. Every target voter atomically copies staged MVCC + // versions into the live keyspace and removes the promoted staged rows. + raftEncodeMigrationPromote byte = 0x0b ) func decodeRaftRequests(data []byte) ([]*pb.Request, error) { @@ -526,23 +545,18 @@ func (f *kvFSM) handleRawRequest(ctx context.Context, r *pb.Request, commitTS ui if err := f.verifyWriteFence(r); err != nil { return err } + floorSnap, err := f.routeFloorSnapshotForRequest(r) + if err != nil { + return err + } // DEL_PREFIX mutations are handled by the store's DeletePrefixAt which // scans and writes tombstones locally. A DEL_PREFIX request must be the // sole mutation in a request (enforced by the coordinator's toRawRequest). if hasDelPrefix, prefix := extractDelPrefix(r.Mutations); hasDelPrefix { - return f.handleDelPrefix(ctx, prefix, commitTS) + return f.handleDelPrefixWithFloorSnapshot(ctx, prefix, commitTS, floorSnap) } - for _, mut := range r.Mutations { - if mut == nil || len(mut.Key) == 0 { - return errors.WithStack(ErrInvalidRequest) - } - // Raw requests should not mutate txn-internal keys. - if isTxnInternalKey(mut.Key) { - return errors.WithStack(ErrInvalidRequest) - } - } - if err := f.assertNoConflictingTxnLocks(ctx, r.Mutations, nil, 0); err != nil { + if err := f.validateRawMutationsForApply(ctx, r, commitTS, floorSnap); err != nil { return err } @@ -559,6 +573,45 @@ func (f *kvFSM) handleRawRequest(ctx context.Context, r *pb.Request, commitTS ui return nil } +func (f *kvFSM) validateRawMutationsForApply(ctx context.Context, r *pb.Request, commitTS uint64, floorSnap RouteSnapshot) error { + bypassKeys := writeFenceBypassKeySet(r.GetWriteFenceBypassKeys()) + for _, mut := range r.GetMutations() { + if err := f.validateRawMutationForApply(ctx, mut, bypassKeys, commitTS, floorSnap); err != nil { + return err + } + } + return nil +} + +func (f *kvFSM) validateRawMutationForApply(ctx context.Context, mut *pb.Mutation, writeFenceBypassKeys map[string]struct{}, commitTS uint64, floorSnap RouteSnapshot) error { + if mut == nil || len(mut.Key) == 0 { + return errors.WithStack(ErrInvalidRequest) + } + // Raw requests should not mutate txn-internal keys. + if isTxnInternalKey(mut.Key) { + return errors.WithStack(ErrInvalidRequest) + } + // Nor the migration and catalog control namespaces. Those are written only + // by the typed internal commands (catalog applies, migration import and + // promote), never by a RawKV mutation. The check is a pure function of the + // key, so every replica reaches the same verdict for the same entry. + if distribution.IsReservedControlKey(mut.Key) { + return errors.WithStack(ErrInvalidRequest) + } + if _, bypass := writeFenceBypassKeys[string(mut.Key)]; !bypass { + if err := f.verifyRouteNotFencedForKey(mut.Key); err != nil { + return err + } + } + if err := verifyRouteWriteTimestampFloorForKeyFromSnapshot(floorSnap, mut.Key, commitTS); err != nil { + return err + } + if err := f.assertNoConflictingTxnLock(ctx, mut.Key, nil, 0); err != nil { + return err + } + return nil +} + // extractDelPrefix checks if the mutations contain a DEL_PREFIX operation. // If found, it validates that no other operation types are mixed in. func extractDelPrefix(muts []*pb.Mutation) (bool, []byte) { @@ -573,13 +626,186 @@ func extractDelPrefix(muts []*pb.Mutation) (bool, []byte) { // handleDelPrefix delegates prefix deletion to the store. Transaction-internal // keys are always excluded to preserve transactional integrity. func (f *kvFSM) handleDelPrefix(ctx context.Context, prefix []byte, commitTS uint64) error { - if err := f.store.DeletePrefixAtRaftAt(ctx, prefix, txnCommonPrefix, commitTS, f.pendingApplyIdx); err != nil { + return f.handleDelPrefixWithFloorSnapshot(ctx, prefix, commitTS, nil) +} + +func (f *kvFSM) handleDelPrefixWithFloorSnapshot(ctx context.Context, prefix []byte, commitTS uint64, floorSnap RouteSnapshot) error { + // DEL_PREFIX never reaches validateRawMutationsForApply, so the control + // namespaces are gated here as well. A prefix is refused both when it sits + // inside one and when it is broad enough to sweep one up. + if distribution.ReservedControlPrefixIntersects(prefix) { + return errors.WithStack(ErrInvalidRequest) + } + if err := f.verifyRouteNotFencedForPrefix(prefix); err != nil { + return err + } + if err := verifyRouteWriteTimestampFloorForPrefixFromSnapshot(floorSnap, prefix, commitTS); err != nil { + return err + } + deletes := []store.PrefixDelete{{ + Prefix: prefix, + ExcludePrefix: txnCommonPrefix, + }} + for _, del := range f.stagedVisibilityPrefixDeletesForApply(prefix, txnCommonPrefix) { + deletes = append(deletes, store.PrefixDelete{ + Prefix: del.prefix, + ExcludePrefix: del.excludePrefix, + }) + } + if err := f.store.DeletePrefixesAtRaftAt(ctx, deletes, commitTS, f.pendingApplyIdx); err != nil { return errors.WithStack(err) } f.notifyApplyObserver(commitTS, pb.Op_DEL_PREFIX, prefix) return nil } +func (f *kvFSM) routeFloorSnapshotForRequest(r *pb.Request) (RouteSnapshot, error) { + if f.routes == nil || f.shardGroupID == 0 { + return nil, nil + } + observedVer, pinned := DecodeObservedRouteVersion(r.GetObservedRouteVersion()) + if !pinned { + return nil, nil + } + snap, ok := f.routes.SnapshotAt(observedVer) + if !ok { + return nil, errors.WithStack(ErrComposed1VersionGCd) + } + return snap, nil +} + +type fsmStagedVisibilityPrefixDelete struct { + prefix []byte + excludePrefix []byte +} + +func (f *kvFSM) stagedVisibilityPrefixDeletesForApply(prefix []byte, excludePrefix []byte) []fsmStagedVisibilityPrefixDelete { + if f == nil || f.routes == nil || f.shardGroupID == 0 { + return nil + } + snap, ok := f.routes.Current() + if !ok { + return nil + } + start, end := routePrefixRange(prefix) + routes := snap.IntersectingRoutes(start, end) + out := make([]fsmStagedVisibilityPrefixDelete, 0, len(routes)) + seen := make(map[string]struct{}, len(routes)) + for _, route := range routes { + if route.GroupID != f.shardGroupID || !routeHasStagedVisibility(route) { + continue + } + stagedPrefix := distribution.MigrationStagedDataKey(route.MigrationJobID, prefix) + var stagedExclude []byte + if excludePrefix != nil { + stagedExclude = distribution.MigrationStagedDataKey(route.MigrationJobID, excludePrefix) + } + dedupeKey := string(stagedPrefix) + "\x00" + string(stagedExclude) + if _, ok := seen[dedupeKey]; ok { + continue + } + seen[dedupeKey] = struct{}{} + out = append(out, fsmStagedVisibilityPrefixDelete{prefix: stagedPrefix, excludePrefix: stagedExclude}) + } + return out +} + +func (f *kvFSM) verifyRouteNotFencedForKey(key []byte) error { + if f.routes == nil { + return nil + } + snap, ok := f.routes.Current() + if !ok { + return nil + } + if start, end, ok := s3BucketAuxiliaryRouteRange(key); ok { + route, found := s3BucketAuxiliaryOwnerRouteFromRange(start, end, snap.IntersectingRoutes(start, end)) + if found && route.State == distribution.RouteStateWriteFenced { + return errors.Wrapf(ErrRouteWriteFenced, "key %q route range [%q,%q)", key, start, end) + } + return nil + } + rkey := routeKey(key) + if snap.WriteFencedForKey(rkey) { + return errors.Wrapf(ErrRouteWriteFenced, "key %q routeKey %q", key, rkey) + } + return nil +} + +func (f *kvFSM) verifyRouteNotFencedForPrefix(prefix []byte) error { + if f.routes == nil { + return nil + } + snap, ok := f.routes.Current() + if !ok { + return nil + } + start, end := routePrefixRange(prefix) + if !snap.WriteFencedIntersects(start, end) { + return nil + } + return errors.Wrapf(ErrRouteWriteFenced, "prefix %q route range [%q,%q)", prefix, start, end) +} + +func verifyRouteWriteTimestampFloorForKeyFromSnapshot(snap RouteSnapshot, key []byte, commitTS uint64) error { + if snap == nil || commitTS == 0 { + return nil + } + if start, end, ok := s3BucketAuxiliaryRouteRange(key); ok { + route, found := s3BucketAuxiliaryOwnerRouteFromRange(start, end, snap.IntersectingRoutes(start, end)) + if found { + return verifyRouteWriteTimestampFloorForRange(route, key, start, end, commitTS) + } + return nil + } + rkey := routeKey(key) + if route, ok := snap.RouteOf(rkey); ok { + if err := verifyRouteWriteTimestampFloorForRoute(route, key, commitTS); err != nil { + return err + } + } + return nil +} + +func verifyRouteWriteTimestampFloorsForMutationsFromSnapshot(snap RouteSnapshot, muts []*pb.Mutation, commitTS uint64) error { + for _, mut := range muts { + if mut == nil || len(mut.Key) == 0 || isTxnInternalKey(mut.Key) { + continue + } + if err := verifyRouteWriteTimestampFloorForKeyFromSnapshot(snap, mut.Key, commitTS); err != nil { + return err + } + } + return nil +} + +func verifyRouteWriteTimestampFloorForPrefixFromSnapshot(snap RouteSnapshot, prefix []byte, commitTS uint64) error { + if snap == nil || commitTS == 0 { + return nil + } + start, end := routePrefixRange(prefix) + for _, route := range snap.IntersectingRoutes(start, end) { + if err := verifyRouteWriteTimestampFloorForRange(route, prefix, start, end, commitTS); err != nil { + return err + } + } + return nil +} + +func verifyRouteWriteTimestampFloorForRoute(route distribution.Route, key []byte, commitTS uint64) error { + if route.MinWriteTSExclusive == 0 || commitTS > route.MinWriteTSExclusive { + return nil + } + return errors.Wrapf(ErrRouteWriteTimestampTooLow, "key %q routeKey %q commit_ts=%d floor=%d", key, routeKey(key), commitTS, route.MinWriteTSExclusive) +} + +func verifyRouteWriteTimestampFloorForRange(route distribution.Route, key, start, end []byte, commitTS uint64) error { + if route.MinWriteTSExclusive == 0 || commitTS > route.MinWriteTSExclusive { + return nil + } + return errors.Wrapf(ErrRouteWriteTimestampTooLow, "key %q route range [%q,%q) commit_ts=%d floor=%d", key, start, end, commitTS, route.MinWriteTSExclusive) +} + func routePrefixRange(prefix []byte) ([]byte, []byte) { if len(prefix) == 0 { return []byte(""), nil @@ -590,6 +816,9 @@ func routePrefixRange(prefix []byte) ([]byte, []byte) { if start, ok := dynamoExactCleanupRouteKey(prefix); ok { return start, routePointRangeEnd(start) } + if start, end, ok := fskeys.ChunkScanRouteBounds(prefix, prefixScanEnd(prefix)); ok { + return start, end + } if routeKeyspaceWideRawPrefix(prefix) { return []byte(""), nil } @@ -597,6 +826,11 @@ func routePrefixRange(prefix []byte) ([]byte, []byte) { return start, prefixScanEnd(start) } +// RoutePrefixRange maps a raw key prefix to the routed key range it may touch. +func RoutePrefixRange(prefix []byte) ([]byte, []byte) { + return routePrefixRange(prefix) +} + func dynamoExactCleanupRouteKey(prefix []byte) ([]byte, bool) { switch { case bytes.HasPrefix(prefix, dynamoTableMetaPrefixBytes), @@ -666,6 +900,7 @@ var routeMappedRawPrefixes = append([][]byte{ []byte(s3keys.BlobPrefix), []byte(s3keys.GCUploadPrefix), []byte(s3keys.RoutePrefix), + fskeys.UsageRouteAllPrefix(), }, sqsConcreteInternalPrefixBytes...) var ErrNotImplemented = errors.New("not implemented") @@ -971,22 +1206,36 @@ func verifyWriteFenceFromSnapshot(mutations []*pb.Mutation, writeFenceBypassKeys if _, ok := bypassKeys[string(mut.Key)]; ok { continue } + if checked, err := verifyS3BucketAuxiliaryWriteFenceFromSnapshot(mut.Key, snap, snapVer, phase); checked { + if err != nil { + return err + } + continue + } rKey := routeKey(mut.Key) if snap.WriteFencedForKey(rKey) { return errors.Wrapf(ErrRouteWriteFenced, "%s-version v=%d: key %q routeKey %q", phase, snapVer, mut.Key, rKey) } - start, end, ok := s3BucketAuxiliaryRouteRange(mut.Key) - if ok && snap.WriteFencedIntersects(start, end) { - return errors.Wrapf(ErrRouteWriteFenced, - "%s-version v=%d: key %q route range [%q,%q)", - phase, snapVer, mut.Key, start, end) - } } return nil } +func verifyS3BucketAuxiliaryWriteFenceFromSnapshot(key []byte, snap RouteSnapshot, snapVer uint64, phase string) (bool, error) { + start, end, ok := s3BucketAuxiliaryRouteRange(key) + if !ok { + return false, nil + } + route, found := s3BucketAuxiliaryOwnerRouteFromRange(start, end, snap.IntersectingRoutes(start, end)) + if found && route.State == distribution.RouteStateWriteFenced { + return true, errors.Wrapf(ErrRouteWriteFenced, + "%s-version v=%d: key %q route range [%q,%q)", + phase, snapVer, key, start, end) + } + return true, nil +} + func writeFenceBypassKeySet(keys [][]byte) map[string]struct{} { if len(keys) == 0 { return nil @@ -1018,21 +1267,24 @@ func (f *kvFSM) verifyOwnerFromSnapshot(mutations []*pb.Mutation, bypassKeys map if _, ok := bypassKeys[string(mut.Key)]; ok { continue } - // routeKey-normalize before OwnerOf so the gate routes the - // same way as ShardRouter.ResolveGroup — raw adapter keys - // and route catalog ranges live in different lex bands - // (issue #930). - rKey := routeKey(mut.Key) - owner, found := snap.OwnerOf(rKey) + ownerKey := composed1OwnerKey(mut.Key) + owner, found := snap.OwnerOf(ownerKey) if !found || owner != f.shardGroupID { return errors.Wrapf(ErrComposed1Violation, "%s-version v=%d: key %q (routeKey %q) owned by group %d (found=%v); this FSM serves group %d", - phase, snapVer, mut.Key, rKey, owner, found, f.shardGroupID) + phase, snapVer, mut.Key, ownerKey, owner, found, f.shardGroupID) } } return nil } +func composed1OwnerKey(key []byte) []byte { + if start, _, ok := s3BucketAuxiliaryRouteRange(key); ok { + return start + } + return routeKey(key) +} + func (f *kvFSM) validateConflicts(ctx context.Context, muts []*pb.Mutation, startTS uint64) error { seen := make(map[string]struct{}, len(muts)) for _, mut := range muts { @@ -1094,7 +1346,7 @@ func (f *kvFSM) handlePrepareRequest(ctx context.Context, r *pb.Request) error { } startTS := r.Ts - uniq, err := uniqueMutations(muts) + uniq, err := f.uniqueMutationsAboveFloorForRequest(r, muts, startTS) if err != nil { return err } @@ -1145,7 +1397,7 @@ func (f *kvFSM) handleOnePhaseTxnRequest(ctx context.Context, r *pb.Request, com // duplicate (the very :duplicate-elements anomaly), so no-op the whole // apply and let the adapter reconstruct the prior result. // - // Determinism note (codex P1 round-11): the underlying CommittedVersionAt + // Determinism note: the underlying CommittedVersionAt // intentionally does NOT enforce the retention watermark — branching FSM // apply on the per-replica minRetainedTS would let replicas with stale // retention surface ErrReadTSCompacted and skip dedup while replicas that @@ -1156,7 +1408,7 @@ func (f *kvFSM) handleOnePhaseTxnRequest(ctx context.Context, r *pb.Request, com // applying this log entry. The retention-window > max-retry-latency // invariant prevents the rare case where a real never-landed retry // arrives with PrevCommitTS below pebble's compacted floor. - dedup, err := f.dedupProbeOnePhase(ctx, meta) + dedup, err := f.dedupProbeOnePhase(ctx, meta, r.ReadKeys) if err != nil { return err } @@ -1164,7 +1416,7 @@ func (f *kvFSM) handleOnePhaseTxnRequest(ctx context.Context, r *pb.Request, com return nil } - uniq, err := uniqueMutations(muts) + uniq, err := f.uniqueMutationsAboveFloorForRequest(r, muts, commitTS) if err != nil { return err } @@ -1180,6 +1432,90 @@ func (f *kvFSM) handleOnePhaseTxnRequest(ctx context.Context, r *pb.Request, com return nil } +func uniqueTxnMutations(muts []*pb.Mutation) ([]*pb.Mutation, error) { + uniq, err := uniqueMutations(muts) + if err != nil { + return nil, err + } + return uniq, nil +} + +func (f *kvFSM) uniqueMutationsAboveFloor(muts []*pb.Mutation, commitTS uint64) ([]*pb.Mutation, error) { + return f.uniqueMutationsAboveFloorWithSnapshot(muts, commitTS, nil) +} + +func (f *kvFSM) uniqueMutationsAboveFloorForRequest(r *pb.Request, muts []*pb.Mutation, commitTS uint64) ([]*pb.Mutation, error) { + floorSnap, err := f.routeFloorSnapshotForRequest(r) + if err != nil { + return nil, err + } + return f.uniqueMutationsAboveFloorWithSnapshot(muts, commitTS, floorSnap) +} + +func (f *kvFSM) uniqueMutationsAboveFloorWithSnapshot(muts []*pb.Mutation, commitTS uint64, floorSnap RouteSnapshot) ([]*pb.Mutation, error) { + uniq, err := uniqueMutations(muts) + if err != nil { + return nil, err + } + if err := rejectReservedControlMutations(uniq); err != nil { + return nil, err + } + if err := verifyRouteWriteTimestampFloorsForMutationsFromSnapshot(floorSnap, uniq, commitTS); err != nil { + return nil, err + } + return uniq, nil +} + +// rejectReservedControlMutations keeps the migration-only control namespaces +// out of transactional writes. validateRawMutationForApply refuses every +// control namespace on the raw path, but the transactional paths cannot: the +// control plane commits catalog route and job records through the coordinator, +// so SplitRange itself is a transaction writing !dist|route| and !dist|meta|. +// The narrower test still closes the case that matters -- a client write +// landing under !dist|migstage|| is promoted as user data when the same +// group is a migration target. +// +// The verdict is a pure function of the key, so every replica reaches it for +// the same entry. The typed migration commands (import, promote, fence) do not +// pass through here and keep writing these namespaces. +func rejectReservedControlMutations(muts []*pb.Mutation) error { + for _, mut := range muts { + if mut == nil { + continue + } + if distribution.IsMigrationOnlyControlKey(mut.Key) { + return errors.WithStack(ErrInvalidRequest) + } + } + return nil +} + +func (f *kvFSM) uniqueTxnMutationsAboveFloor(muts []*pb.Mutation, commitTS uint64) ([]*pb.Mutation, error) { + return f.uniqueTxnMutationsAboveFloorWithSnapshot(muts, commitTS, nil) +} + +func (f *kvFSM) uniqueTxnMutationsAboveFloorForRequest(r *pb.Request, muts []*pb.Mutation, commitTS uint64) ([]*pb.Mutation, error) { + floorSnap, err := f.routeFloorSnapshotForRequest(r) + if err != nil { + return nil, err + } + return f.uniqueTxnMutationsAboveFloorWithSnapshot(muts, commitTS, floorSnap) +} + +func (f *kvFSM) uniqueTxnMutationsAboveFloorWithSnapshot(muts []*pb.Mutation, commitTS uint64, floorSnap RouteSnapshot) ([]*pb.Mutation, error) { + uniq, err := uniqueTxnMutations(muts) + if err != nil { + return nil, err + } + if err := rejectReservedControlMutations(uniq); err != nil { + return nil, err + } + if err := verifyRouteWriteTimestampFloorsForMutationsFromSnapshot(floorSnap, uniq, commitTS); err != nil { + return nil, err + } + return uniq, nil +} + // dedupProbeOnePhase decides whether handleOnePhaseTxnRequest should no-op // because the entry is a retry whose prior attempt already landed. Extracted // to keep handleOnePhaseTxnRequest under the cyclop budget; the determinism @@ -1188,7 +1524,7 @@ func (f *kvFSM) handleOnePhaseTxnRequest(ctx context.Context, r *pb.Request, com // Returns (true, nil) → the entry must no-op (prior attempt landed). // Returns (false, nil) → fall through to normal apply. // Returns (false, err) → propagate err; apply must not proceed. -func (f *kvFSM) dedupProbeOnePhase(ctx context.Context, meta TxnMeta) (bool, error) { +func (f *kvFSM) dedupProbeOnePhase(ctx context.Context, meta TxnMeta, readKeys [][]byte) (bool, error) { if meta.PrevCommitTS == 0 { return false, nil } @@ -1196,7 +1532,51 @@ func (f *kvFSM) dedupProbeOnePhase(ctx context.Context, meta TxnMeta) (bool, err if err != nil { return false, errors.WithStack(err) } - return landed, nil + if landed { + return true, nil + } + for _, stagedKey := range stagedVisibilityKeysForReadAlias(readKeys, meta.PrimaryKey, meta.PrimaryKey) { + landed, err = f.store.CommittedVersionAt(ctx, stagedKey, meta.PrevCommitTS) + if err != nil { + return false, errors.WithStack(err) + } + if landed { + return true, nil + } + } + return false, nil +} + +func stagedVisibilityJobIDsForReadAlias(readKeys [][]byte, aliasRawKey []byte) []uint64 { + if len(readKeys) == 0 { + return nil + } + var out []uint64 + seen := make(map[uint64]struct{}) + for _, readKey := range readKeys { + jobID, rawKey, ok := distribution.MigrationStagedDataKeyParts(readKey) + if !ok || jobID == 0 || !bytes.Equal(rawKey, aliasRawKey) { + continue + } + if _, ok := seen[jobID]; ok { + continue + } + seen[jobID] = struct{}{} + out = append(out, jobID) + } + return out +} + +func stagedVisibilityKeysForReadAlias(readKeys [][]byte, aliasRawKey, stagedRawKey []byte) [][]byte { + jobIDs := stagedVisibilityJobIDsForReadAlias(readKeys, aliasRawKey) + if len(jobIDs) == 0 { + return nil + } + out := make([][]byte, 0, len(jobIDs)) + for _, jobID := range jobIDs { + out = append(out, distribution.MigrationStagedDataKey(jobID, stagedRawKey)) + } + return out } func (f *kvFSM) handleCommitRequest(ctx context.Context, r *pb.Request) error { @@ -1215,15 +1595,15 @@ func (f *kvFSM) handleCommitRequest(ctx context.Context, r *pb.Request) error { if len(meta.PrimaryKey) == 0 { return errors.WithStack(ErrTxnPrimaryKeyRequired) } - applyStartTS, err := f.commitApplyStartTS(ctx, meta.PrimaryKey, startTS, commitTS) + applyStartTS, err := f.commitApplyStartTS(ctx, meta.PrimaryKey, startTS, commitTS, r.ReadKeys) if err != nil { return err } - uniq, err := uniqueMutations(muts) + uniq, err := f.uniqueTxnMutationsAboveFloorForRequest(r, muts, commitTS) if err != nil { return err } - storeMuts, err := f.buildCommitStoreMutations(ctx, uniq, meta, startTS, commitTS) + storeMuts, err := f.buildCommitStoreMutations(ctx, uniq, meta, startTS, commitTS, r.ReadKeys) if err != nil { return err } @@ -1246,11 +1626,17 @@ func (f *kvFSM) handleCommitRequest(ctx context.Context, r *pb.Request) error { // ErrTxnAlreadyAborted. Together with the commit-record check in // appendRollbackRecord, this enforces the invariant that at most one of // {rollback marker, commit record} is present for any (primaryKey, startTS). -func (f *kvFSM) commitApplyStartTS(ctx context.Context, primaryKey []byte, startTS, commitTS uint64) (uint64, error) { +func (f *kvFSM) commitApplyStartTS(ctx context.Context, primaryKey []byte, startTS, commitTS uint64, readKeys [][]byte) (uint64, error) { recordedCommitTS, committed, err := f.txnCommitTS(ctx, primaryKey, startTS) if err != nil { return 0, err } + if !committed { + recordedCommitTS, committed, err = f.stagedTxnCommitTS(ctx, primaryKey, startTS, readKeys) + if err != nil { + return 0, err + } + } if !committed { // No commit record yet: reject if a rollback marker is present. // This catches out-of-order apply (COMMIT after ABORT), buggy @@ -1262,6 +1648,13 @@ func (f *kvFSM) commitApplyStartTS(ctx context.Context, primaryKey []byte, start if exists { return 0, errors.WithStack(ErrTxnAlreadyAborted) } + exists, rerr = f.stagedTxnRecordExists(ctx, primaryKey, txnRollbackKey(primaryKey, startTS), readKeys) + if rerr != nil { + return 0, rerr + } + if exists { + return 0, errors.WithStack(ErrTxnAlreadyAborted) + } return startTS, nil } if recordedCommitTS != commitTS { @@ -1333,19 +1726,19 @@ func (f *kvFSM) handleAbortRequest(ctx context.Context, r *pb.Request, abortTS u // abort for a secondary key (same primaryKey, same startTS) would // see the marker already present and must still clean up that // secondary's lock/intent. Idempotency is enforced per-key in - // shouldClearAbortKey (lock-missing ⇒ nothing to do) and for the + // abortCleanupMutationsForKey (lock-missing ⇒ nothing to do) and for the // rollback-marker Put in appendRollbackRecord. uniq, err := uniqueMutations(muts) if err != nil { return err } - storeMuts, abortingPrimary, err := f.buildAbortCleanupStoreMutations(ctx, uniq, meta.PrimaryKey, startTS) + storeMuts, abortingPrimary, err := f.buildAbortCleanupStoreMutations(ctx, uniq, meta.PrimaryKey, startTS, r.ReadKeys) if err != nil { return err } if abortingPrimary { - if err := f.appendRollbackRecord(ctx, meta.PrimaryKey, startTS, &storeMuts); err != nil { + if err := f.appendRollbackRecord(ctx, meta.PrimaryKey, startTS, &storeMuts, r.ReadKeys); err != nil { return err } } @@ -1384,7 +1777,7 @@ func (f *kvFSM) buildOnePhaseStoreMutations(ctx context.Context, muts []*pb.Muta return storeMuts, nil } -func (f *kvFSM) buildCommitStoreMutations(ctx context.Context, muts []*pb.Mutation, meta TxnMeta, startTS, commitTS uint64) ([]*store.KVPairMutation, error) { +func (f *kvFSM) buildCommitStoreMutations(ctx context.Context, muts []*pb.Mutation, meta TxnMeta, startTS, commitTS uint64, readKeys [][]byte) ([]*store.KVPairMutation, error) { storeMuts := make([]*store.KVPairMutation, 0, len(muts)*txnCommitStoreMutationFactor+txnCommitStoreMutationSlack) committingPrimary := false @@ -1394,7 +1787,7 @@ func (f *kvFSM) buildCommitStoreMutations(ctx context.Context, muts []*pb.Mutati committingPrimary = true } - keyMuts, err := f.commitTxnKeyMutations(ctx, key, meta.PrimaryKey, startTS) + keyMuts, err := f.commitTxnKeyMutations(ctx, key, meta.PrimaryKey, startTS, readKeys) if err != nil { return nil, err } @@ -1412,7 +1805,7 @@ func (f *kvFSM) buildCommitStoreMutations(ctx context.Context, muts []*pb.Mutati return storeMuts, nil } -func (f *kvFSM) buildAbortCleanupStoreMutations(ctx context.Context, muts []*pb.Mutation, primaryKey []byte, startTS uint64) ([]*store.KVPairMutation, bool, error) { +func (f *kvFSM) buildAbortCleanupStoreMutations(ctx context.Context, muts []*pb.Mutation, primaryKey []byte, startTS uint64, readKeys [][]byte) ([]*store.KVPairMutation, bool, error) { storeMuts := make([]*store.KVPairMutation, 0, len(muts)*txnAbortStoreMutationFactor) abortingPrimary := false for _, mut := range muts { @@ -1421,18 +1814,16 @@ func (f *kvFSM) buildAbortCleanupStoreMutations(ctx context.Context, muts []*pb. abortingPrimary = true } - shouldClear, err := f.shouldClearAbortKey(ctx, key, primaryKey, startTS) + cleanupMuts, err := f.abortCleanupMutationsForKey(ctx, key, primaryKey, startTS, readKeys) if err != nil { return nil, false, err } - if shouldClear { - storeMuts = append(storeMuts, txnCleanupMutations(key)...) - } + storeMuts = append(storeMuts, cleanupMuts...) } return storeMuts, abortingPrimary, nil } -func (f *kvFSM) appendRollbackRecord(ctx context.Context, primaryKey []byte, startTS uint64, storeMuts *[]*store.KVPairMutation) error { +func (f *kvFSM) appendRollbackRecord(ctx context.Context, primaryKey []byte, startTS uint64, storeMuts *[]*store.KVPairMutation, readKeys [][]byte) error { // Desired invariant: for any (primaryKey, startTS) pair, at most // one of {rollback marker, commit record} is present. The invariant // holds when aborts/commits flow through the symmetric guards in @@ -1451,6 +1842,12 @@ func (f *kvFSM) appendRollbackRecord(ctx context.Context, primaryKey []byte, sta if err != nil { return errors.WithStack(err) } + if !markerPresent { + markerPresent, err = f.stagedTxnRecordExists(ctx, primaryKey, txnRollbackKey(primaryKey, startTS), readKeys) + if err != nil { + return err + } + } // Verify the invariant regardless of marker presence: if a commit // record is present for this (primaryKey, startTS), refuse to @@ -1461,6 +1858,12 @@ func (f *kvFSM) appendRollbackRecord(ctx context.Context, primaryKey []byte, sta if err != nil { return errors.WithStack(err) } + if !commitExists { + commitExists, err = f.stagedTxnRecordExists(ctx, primaryKey, txnCommitKey(primaryKey, startTS), readKeys) + if err != nil { + return err + } + } if commitExists { return errors.WithStack(ErrTxnAlreadyCommitted) } @@ -1478,7 +1881,37 @@ func (f *kvFSM) appendRollbackRecord(ctx context.Context, primaryKey []byte, sta } func (f *kvFSM) txnCommitTS(ctx context.Context, primaryKey []byte, startTS uint64) (uint64, bool, error) { - b, err := f.store.GetAt(ctx, txnCommitKey(primaryKey, startTS), ^uint64(0)) + return f.txnCommitTSAtKey(ctx, txnCommitKey(primaryKey, startTS)) +} + +func (f *kvFSM) stagedTxnCommitTS(ctx context.Context, primaryKey []byte, startTS uint64, readKeys [][]byte) (uint64, bool, error) { + for _, key := range stagedVisibilityKeysForReadAlias(readKeys, primaryKey, txnCommitKey(primaryKey, startTS)) { + commitTS, committed, err := f.txnCommitTSAtKey(ctx, key) + if err != nil { + return 0, false, err + } + if committed { + return commitTS, true, nil + } + } + return 0, false, nil +} + +func (f *kvFSM) stagedTxnRecordExists(ctx context.Context, primaryKey, recordKey []byte, readKeys [][]byte) (bool, error) { + for _, key := range stagedVisibilityKeysForReadAlias(readKeys, primaryKey, recordKey) { + exists, err := f.store.ExistsAt(ctx, key, ^uint64(0)) + if err != nil { + return false, errors.WithStack(err) + } + if exists { + return true, nil + } + } + return false, nil +} + +func (f *kvFSM) txnCommitTSAtKey(ctx context.Context, key []byte) (uint64, bool, error) { + b, err := f.store.GetAt(ctx, key, ^uint64(0)) if err != nil { if errors.Is(err, store.ErrKeyNotFound) { return 0, false, nil @@ -1529,15 +1962,38 @@ func txnIntentFromPBMutation(mut *pb.Mutation, startTS uint64) (txnIntent, error } } -func txnCleanupMutations(key []byte) []*store.KVPairMutation { +func txnCleanupMutationsForKeys(lockKey, intentKey []byte) []*store.KVPairMutation { return []*store.KVPairMutation{ - {Op: store.OpTypeDelete, Key: txnLockKey(key)}, - {Op: store.OpTypeDelete, Key: txnIntentKey(key)}, + {Op: store.OpTypeDelete, Key: lockKey}, + {Op: store.OpTypeDelete, Key: intentKey}, } } -func (f *kvFSM) txnLockForCommit(ctx context.Context, key []byte) (txnLock, bool, error) { - lockBytes, err := f.store.GetAt(ctx, txnLockKey(key), ^uint64(0)) +func (f *kvFSM) txnLockForMutationKey(ctx context.Context, key []byte, readKeys [][]byte) (txnLock, []byte, []byte, bool, error) { + lockKey := txnLockKey(key) + intentKey := txnIntentKey(key) + lock, ok, err := f.txnLockAtKey(ctx, lockKey) + if err != nil || ok { + return lock, lockKey, intentKey, ok, err + } + for _, stagedLockKey := range stagedVisibilityKeysForReadAlias(readKeys, key, lockKey) { + jobID, _, ok := distribution.MigrationStagedDataKeyParts(stagedLockKey) + if !ok { + continue + } + lock, ok, err := f.txnLockAtKey(ctx, stagedLockKey) + if err != nil { + return txnLock{}, nil, nil, false, err + } + if ok { + return lock, stagedLockKey, distribution.MigrationStagedDataKey(jobID, intentKey), true, nil + } + } + return txnLock{}, nil, nil, false, nil +} + +func (f *kvFSM) txnLockAtKey(ctx context.Context, key []byte) (txnLock, bool, error) { + lockBytes, err := f.store.GetAt(ctx, key, ^uint64(0)) if err != nil { if errors.Is(err, store.ErrKeyNotFound) { return txnLock{}, false, nil @@ -1551,8 +2007,8 @@ func (f *kvFSM) txnLockForCommit(ctx context.Context, key []byte) (txnLock, bool return lock, true, nil } -func (f *kvFSM) txnIntentForCommit(ctx context.Context, key []byte) (txnIntent, bool, error) { - intentBytes, err := f.store.GetAt(ctx, txnIntentKey(key), ^uint64(0)) +func (f *kvFSM) txnIntentForCommit(ctx context.Context, intentKey []byte) (txnIntent, bool, error) { + intentBytes, err := f.store.GetAt(ctx, intentKey, ^uint64(0)) if err != nil { if errors.Is(err, store.ErrKeyNotFound) { return txnIntent{}, false, nil @@ -1577,8 +2033,8 @@ func storeMutationForIntent(key []byte, intent txnIntent) (*store.KVPairMutation } } -func (f *kvFSM) commitTxnKeyMutations(ctx context.Context, key, primaryKey []byte, startTS uint64) ([]*store.KVPairMutation, error) { - lock, ok, err := f.txnLockForCommit(ctx, key) +func (f *kvFSM) commitTxnKeyMutations(ctx context.Context, key, primaryKey []byte, startTS uint64, readKeys [][]byte) ([]*store.KVPairMutation, error) { + lock, lockKey, intentKey, ok, err := f.txnLockForMutationKey(ctx, key, readKeys) if err != nil { return nil, err } @@ -1593,7 +2049,7 @@ func (f *kvFSM) commitTxnKeyMutations(ctx context.Context, key, primaryKey []byt return nil, errors.Wrapf(ErrTxnInvalidMeta, "lock primary_key mismatch for key %s", string(key)) } - intent, ok, err := f.txnIntentForCommit(ctx, key) + intent, ok, err := f.txnIntentForCommit(ctx, intentKey) if err != nil { return nil, err } @@ -1609,37 +2065,33 @@ func (f *kvFSM) commitTxnKeyMutations(ctx context.Context, key, primaryKey []byt } out = append(out, mut) } - out = append(out, txnCleanupMutations(key)...) + out = append(out, txnCleanupMutationsForKeys(lockKey, intentKey)...) return out, nil } -// shouldClearAbortKey reports whether this abort request must emit -// cleanup (lock+intent Delete) mutations for key. It returns false +// abortCleanupMutationsForKey reports whether this abort request must emit +// cleanup (lock+intent Delete) mutations for key. It returns nil // when the lock is already missing: lock/intent are always written // and deleted together in a single ApplyMutations batch // (lock missing ⇔ intent missing), so missing lock means either // cleanup already ran for this (startTS, primaryKey) or the key was // never prepared. Emitting Deletes on already-tombstoned keys would // trigger MVCC write conflicts and has no observable effect. -func (f *kvFSM) shouldClearAbortKey(ctx context.Context, key, primaryKey []byte, startTS uint64) (bool, error) { - lockBytes, err := f.store.GetAt(ctx, txnLockKey(key), ^uint64(0)) +func (f *kvFSM) abortCleanupMutationsForKey(ctx context.Context, key, primaryKey []byte, startTS uint64, readKeys [][]byte) ([]*store.KVPairMutation, error) { + lock, lockKey, intentKey, ok, err := f.txnLockForMutationKey(ctx, key, readKeys) if err != nil { - if errors.Is(err, store.ErrKeyNotFound) { - return false, nil - } - return false, errors.WithStack(err) + return nil, err } - lock, derr := decodeTxnLock(lockBytes) - if derr != nil { - return false, errors.WithStack(derr) + if !ok { + return nil, nil } if lock.StartTS != startTS { - return false, nil + return nil, nil } if !bytes.Equal(lock.PrimaryKey, primaryKey) { - return false, errors.Wrapf(ErrTxnInvalidMeta, "abort primary_key mismatch for key %s", string(key)) + return nil, errors.Wrapf(ErrTxnInvalidMeta, "abort primary_key mismatch for key %s", string(key)) } - return true, nil + return txnCleanupMutationsForKeys(lockKey, intentKey), nil } func (f *kvFSM) assertNoConflictingTxnLock(ctx context.Context, key, primaryKey []byte, startTS uint64) error { diff --git a/kv/fsm_abort_test.go b/kv/fsm_abort_test.go index 8498c783f..c486395c2 100644 --- a/kv/fsm_abort_test.go +++ b/kv/fsm_abort_test.go @@ -4,6 +4,7 @@ import ( "context" "testing" + "github.com/bootjp/elastickv/distribution" pb "github.com/bootjp/elastickv/proto" "github.com/bootjp/elastickv/store" "github.com/stretchr/testify/require" @@ -297,7 +298,7 @@ func TestFSMAbort_AbortTSMustBeGreaterThanStartTS(t *testing.T) { // same mutation set must return nil without performing additional // writes or store mutations (reads are allowed — the idempotent path // still probes for the rollback marker and commit record via ExistsAt). -// Idempotency is enforced per-key in shouldClearAbortKey (lock +// Idempotency is enforced per-key in abortCleanupMutationsForKey (lock // already gone ⇒ skip) and by appendRollbackRecord (marker already // present ⇒ skip). The prior behaviour (write-conflict on the // rollback-marker Put) surfaced in prod as "secondary write failed" @@ -553,6 +554,96 @@ func TestFSMAbort_CommitAfterAbortIsRejected(t *testing.T) { require.ErrorIs(t, err, store.ErrKeyNotFound, "commit record must not have been written after rejection") } +func TestFSMAbort_StagedRollbackRecordIsIdempotent(t *testing.T) { + t.Parallel() + + ctx := context.Background() + st := store.NewMVCCStore() + fsm, ok := NewKvFSMWithHLC(st, NewHLC()).(*kvFSM) + require.True(t, ok) + + primary := []byte("pk") + startTS := uint64(10) + readKeys := [][]byte{distribution.MigrationStagedDataKey(9, primary)} + stagedRollbackKey := distribution.MigrationStagedDataKey(9, txnRollbackKey(primary, startTS)) + require.NoError(t, st.PutAt(ctx, stagedRollbackKey, encodeTxnRollbackRecord(), startTS, 0)) + + var storeMuts []*store.KVPairMutation + require.NoError(t, fsm.appendRollbackRecord(ctx, primary, startTS, &storeMuts, readKeys)) + require.Empty(t, storeMuts) +} + +func TestFSMAbort_RejectsStagedCommitRecord(t *testing.T) { + t.Parallel() + + ctx := context.Background() + st := store.NewMVCCStore() + fsm, ok := NewKvFSMWithHLC(st, NewHLC()).(*kvFSM) + require.True(t, ok) + + primary := []byte("pk") + startTS := uint64(10) + commitTS := uint64(20) + readKeys := [][]byte{distribution.MigrationStagedDataKey(9, primary)} + stagedCommitKey := distribution.MigrationStagedDataKey(9, txnCommitKey(primary, startTS)) + require.NoError(t, st.PutAt(ctx, stagedCommitKey, encodeTxnCommitRecord(commitTS), commitTS, 0)) + + var storeMuts []*store.KVPairMutation + err := fsm.appendRollbackRecord(ctx, primary, startTS, &storeMuts, readKeys) + require.Error(t, err) + require.ErrorIs(t, err, ErrTxnAlreadyCommitted) + require.Empty(t, storeMuts) +} + +func TestFSMAbort_CleansUpStagedPreparedArtifacts(t *testing.T) { + t.Parallel() + + ctx := context.Background() + st := store.NewMVCCStore() + fsm, ok := NewKvFSMWithHLC(st, NewHLC()).(*kvFSM) + require.True(t, ok) + + primary := []byte("pk") + startTS := uint64(10) + abortTS := uint64(20) + jobID := uint64(9) + stagedLockKey := distribution.MigrationStagedDataKey(jobID, txnLockKey(primary)) + stagedIntentKey := distribution.MigrationStagedDataKey(jobID, txnIntentKey(primary)) + require.NoError(t, st.PutAt(ctx, stagedLockKey, encodeTxnLock(txnLock{ + StartTS: startTS, + PrimaryKey: primary, + IsPrimaryKey: true, + }), startTS, 0)) + require.NoError(t, st.PutAt(ctx, stagedIntentKey, encodeTxnIntent(txnIntent{ + StartTS: startTS, + Op: txnIntentOpPut, + Value: []byte("v"), + }), startTS, 0)) + + abortReq := &pb.Request{ + IsTxn: true, + Phase: pb.Phase_ABORT, + Ts: startTS, + ReadKeys: [][]byte{ + distribution.MigrationStagedDataKey(jobID, primary), + }, + Mutations: []*pb.Mutation{ + {Op: pb.Op_PUT, Key: []byte(txnMetaPrefix), Value: EncodeTxnMeta(TxnMeta{PrimaryKey: primary, CommitTS: abortTS})}, + {Op: pb.Op_PUT, Key: primary}, + }, + } + require.NoError(t, applyFSMRequest(t, fsm, abortReq)) + + _, err := st.GetAt(ctx, stagedLockKey, ^uint64(0)) + require.ErrorIs(t, err, store.ErrKeyNotFound) + _, err = st.GetAt(ctx, stagedIntentKey, ^uint64(0)) + require.ErrorIs(t, err, store.ErrKeyNotFound) + _, err = st.GetAt(ctx, txnRollbackKey(primary, startTS), ^uint64(0)) + require.NoError(t, err) + _, err = st.GetAt(ctx, primary, ^uint64(0)) + require.ErrorIs(t, err, store.ErrKeyNotFound) +} + func TestFSMAbort_EmptyMutationsReturnsError(t *testing.T) { t.Parallel() diff --git a/kv/fsm_migration_fence_test.go b/kv/fsm_migration_fence_test.go index 807a5f551..23660780c 100644 --- a/kv/fsm_migration_fence_test.go +++ b/kv/fsm_migration_fence_test.go @@ -5,8 +5,10 @@ import ( "testing" "github.com/bootjp/elastickv/distribution" + "github.com/bootjp/elastickv/internal/fskeys" "github.com/bootjp/elastickv/internal/s3keys" pb "github.com/bootjp/elastickv/proto" + "github.com/bootjp/elastickv/store" "github.com/stretchr/testify/require" ) @@ -21,6 +23,16 @@ func newWriteFencedFSM(t *testing.T) *kvFSM { return newComposed1FSM(t, engine, 1) } +func newWriteFloorFSM(t *testing.T) *kvFSM { + t.Helper() + + engine := distribution.NewEngine() + applyComposed1Snapshot(t, engine, 1, []distribution.RouteDescriptor{ + {RouteID: 1, Start: []byte(""), End: nil, GroupID: 1, State: distribution.RouteStateActive, MinWriteTSExclusive: 100}, + }) + return newComposed1FSM(t, engine, 1) +} + func newFirstRouteWriteFencedFSM(t *testing.T) *kvFSM { t.Helper() @@ -42,6 +54,18 @@ func s3BucketAuxiliaryFenceRoutes(bucket string, rawGroupID, fencedGroupID uint6 } } +func s3BucketAuxiliarySplitRoutes(bucket string, rawGroupID, ownerGroupID, splitGroupID uint64) []distribution.RouteDescriptor { + start := s3keys.RoutePrefixForBucketAnyGeneration(bucket) + split := append(append([]byte(nil), start...), 'm') + end := prefixScanEnd(start) + return []distribution.RouteDescriptor{ + {RouteID: 1, Start: []byte(""), End: start, GroupID: rawGroupID, State: distribution.RouteStateActive}, + {RouteID: 2, Start: start, End: split, GroupID: ownerGroupID, State: distribution.RouteStateActive}, + {RouteID: 3, Start: split, End: end, GroupID: splitGroupID, State: distribution.RouteStateActive}, + {RouteID: 4, Start: end, End: nil, GroupID: rawGroupID, State: distribution.RouteStateActive}, + } +} + func newS3BucketAuxiliaryWriteFencedFSM(t *testing.T, bucket string) *kvFSM { t.Helper() @@ -97,6 +121,19 @@ func TestFSMWriteFenceBypassAllowsMarkedRawPointWrite(t *testing.T) { require.Equal(t, []byte("v"), got) } +func TestFSMWriteFenceBypassRejectsRawWriteAtBypassedRouteFloor(t *testing.T) { + t.Parallel() + + fsm := newWriteFloorFSM(t) + key := []byte("!sqs|msg|data|p|partitioned-key") + err := fsm.handleRawRequest(context.Background(), &pb.Request{ + ObservedRouteVersion: 1, + WriteFenceBypassKeys: [][]byte{key}, + Mutations: []*pb.Mutation{{Op: pb.Op_PUT, Key: key, Value: []byte("v")}}, + }, 100) + require.ErrorIs(t, err, ErrRouteWriteTimestampTooLow) +} + func TestFSMWriteFenceBypassAllowsPinnedTxnOnNonOwningGroup(t *testing.T) { t.Parallel() @@ -121,6 +158,25 @@ func TestFSMWriteFenceBypassAllowsPinnedTxnOnNonOwningGroup(t *testing.T) { require.NoError(t, err) } +func TestFSMWriteFenceBypassRejectsPinnedTxnAtBypassedRouteFloor(t *testing.T) { + t.Parallel() + + fsm := newWriteFloorFSM(t) + key := []byte("!sqs|msg|data|p|partitioned-key") + err := fsm.handleTxnRequest(context.Background(), &pb.Request{ + IsTxn: true, + Phase: pb.Phase_PREPARE, + Ts: 100, + ObservedRouteVersion: 1, + WriteFenceBypassKeys: [][]byte{key}, + Mutations: []*pb.Mutation{ + {Op: pb.Op_PUT, Key: []byte(txnMetaPrefix), Value: EncodeTxnMeta(TxnMeta{PrimaryKey: key, LockTTLms: defaultTxnLockTTLms})}, + {Op: pb.Op_DEL, Key: key}, + }, + }, 100) + require.ErrorIs(t, err, ErrRouteWriteTimestampTooLow) +} + func TestFSMWriteFenceBypassDoesNotAllowDelPrefix(t *testing.T) { t.Parallel() @@ -133,6 +189,182 @@ func TestFSMWriteFenceBypassDoesNotAllowDelPrefix(t *testing.T) { require.ErrorIs(t, err, ErrRouteWriteFenced) } +func TestFSMDelPrefixTombstonesStagedVisibilityRowsDuringApply(t *testing.T) { + t.Parallel() + + ctx := context.Background() + engine := distribution.NewEngine() + applyComposed1Snapshot(t, engine, 1, []distribution.RouteDescriptor{ + { + RouteID: 1, + Start: []byte("a"), + End: []byte("z"), + GroupID: 1, + State: distribution.RouteStateActive, + StagedVisibilityActive: true, + MigrationJobID: 9, + }, + }) + fsm := newComposed1FSM(t, engine, 1) + dropKey := []byte("b/drop") + outsideKey := []byte("c/outside") + stagedDrop := distribution.MigrationStagedDataKey(9, dropKey) + stagedOutside := distribution.MigrationStagedDataKey(9, outsideKey) + require.NoError(t, fsm.store.PutAt(ctx, stagedDrop, []byte("drop"), 20, 0)) + require.NoError(t, fsm.store.PutAt(ctx, stagedOutside, []byte("outside"), 20, 0)) + + require.NoError(t, fsm.handleDelPrefix(ctx, []byte("b/"), 101)) + + _, err := fsm.store.GetAt(ctx, stagedDrop, 150) + require.ErrorIs(t, err, store.ErrKeyNotFound) + got, err := fsm.store.GetAt(ctx, stagedOutside, 150) + require.NoError(t, err) + require.Equal(t, []byte("outside"), got) +} + +func TestFSMDelPrefixRoutesFilesystemChunkPrefixThroughVirtualRange(t *testing.T) { + t.Parallel() + + ctx := context.Background() + home := uint64(11) + inode := uint64(22) + chunkPrefix := fskeys.ChunkPrefix(home, inode) + chunkRoute := fskeys.ChunkRouteKey(home, inode) + chunkRouteEnd := prefixScanEnd(chunkRoute) + chunkRawAll := fskeys.ChunkAllPrefix() + chunkRawEnd := prefixScanEnd(chunkRawAll) + + engine := distribution.NewEngine() + applyComposed1Snapshot(t, engine, 1, []distribution.RouteDescriptor{ + {RouteID: 1, Start: []byte(""), End: chunkRawAll, GroupID: 1, State: distribution.RouteStateActive}, + {RouteID: 2, Start: chunkRawAll, End: chunkRawEnd, GroupID: 1, State: distribution.RouteStateActive}, + {RouteID: 3, Start: chunkRawEnd, End: chunkRoute, GroupID: 1, State: distribution.RouteStateActive}, + { + RouteID: 4, + Start: chunkRoute, + End: chunkRouteEnd, + GroupID: 1, + State: distribution.RouteStateActive, + StagedVisibilityActive: true, + MigrationJobID: 9, + }, + {RouteID: 5, Start: chunkRouteEnd, End: nil, GroupID: 1, State: distribution.RouteStateActive}, + }) + fsm := newComposed1FSM(t, engine, 1) + stagedChunk := distribution.MigrationStagedDataKey(9, fskeys.ChunkKey(home, inode, 0)) + require.NoError(t, fsm.store.PutAt(ctx, stagedChunk, []byte("chunk"), 20, 0)) + + require.NoError(t, fsm.handleDelPrefix(ctx, chunkPrefix, 101)) + + _, err := fsm.store.GetAt(ctx, stagedChunk, 150) + require.ErrorIs(t, err, store.ErrKeyNotFound) +} + +func TestFSMRejectsFilesystemChunkDelPrefixAgainstVirtualRouteFence(t *testing.T) { + t.Parallel() + + home := uint64(11) + inode := uint64(22) + chunkRoute := fskeys.ChunkRouteKey(home, inode) + chunkRouteEnd := prefixScanEnd(chunkRoute) + chunkRawAll := fskeys.ChunkAllPrefix() + chunkRawEnd := prefixScanEnd(chunkRawAll) + + engine := distribution.NewEngine() + applyComposed1Snapshot(t, engine, 1, []distribution.RouteDescriptor{ + {RouteID: 1, Start: []byte(""), End: chunkRawAll, GroupID: 1, State: distribution.RouteStateActive}, + {RouteID: 2, Start: chunkRawAll, End: chunkRawEnd, GroupID: 1, State: distribution.RouteStateActive}, + {RouteID: 3, Start: chunkRawEnd, End: chunkRoute, GroupID: 1, State: distribution.RouteStateActive}, + {RouteID: 4, Start: chunkRoute, End: chunkRouteEnd, GroupID: 1, State: distribution.RouteStateWriteFenced}, + {RouteID: 5, Start: chunkRouteEnd, End: nil, GroupID: 1, State: distribution.RouteStateActive}, + }) + fsm := newComposed1FSM(t, engine, 1) + + err := fsm.handleRawRequest(context.Background(), &pb.Request{ + Mutations: []*pb.Mutation{{Op: pb.Op_DEL_PREFIX, Key: fskeys.ChunkPrefix(home, inode)}}, + }, 10) + require.ErrorIs(t, err, ErrRouteWriteFenced) +} + +func TestFSMRejectsFilesystemUsageRoutePrefixWhenAnyOwnerRouteIsWriteFenced(t *testing.T) { + t.Parallel() + + fsm := newWriteFencedFSM(t) + err := fsm.handleRawRequest(context.Background(), &pb.Request{ + Mutations: []*pb.Mutation{{Op: pb.Op_DEL_PREFIX, Key: fskeys.UsageRouteAllPrefix()}}, + }, 10) + require.ErrorIs(t, err, ErrRouteWriteFenced) +} + +type recordingPrefixDeleteStore struct { + store.MVCCStore + + batchCalls int + singleCalls int + deletes []store.PrefixDelete + commitTS uint64 + appliedIndex uint64 +} + +func (s *recordingPrefixDeleteStore) DeletePrefixAtRaftAt(ctx context.Context, prefix []byte, excludePrefix []byte, commitTS, appliedIndex uint64) error { + s.singleCalls++ + return s.MVCCStore.DeletePrefixAtRaftAt(ctx, prefix, excludePrefix, commitTS, appliedIndex) +} + +func (s *recordingPrefixDeleteStore) DeletePrefixesAtRaftAt(ctx context.Context, deletes []store.PrefixDelete, commitTS, appliedIndex uint64) error { + s.batchCalls++ + s.deletes = clonePrefixDeletes(deletes) + s.commitTS = commitTS + s.appliedIndex = appliedIndex + return s.MVCCStore.DeletePrefixesAtRaftAt(ctx, deletes, commitTS, appliedIndex) +} + +func clonePrefixDeletes(deletes []store.PrefixDelete) []store.PrefixDelete { + out := make([]store.PrefixDelete, len(deletes)) + for i, del := range deletes { + out[i] = store.PrefixDelete{ + Prefix: append([]byte(nil), del.Prefix...), + ExcludePrefix: append([]byte(nil), del.ExcludePrefix...), + } + } + return out +} + +func TestFSMDelPrefixBatchesStagedAndRawTombstonesDuringApply(t *testing.T) { + t.Parallel() + + ctx := context.Background() + engine := distribution.NewEngine() + applyComposed1Snapshot(t, engine, 1, []distribution.RouteDescriptor{ + { + RouteID: 1, + Start: []byte("a"), + End: []byte("z"), + GroupID: 1, + State: distribution.RouteStateActive, + StagedVisibilityActive: true, + MigrationJobID: 9, + }, + }) + fsm := newComposed1FSM(t, engine, 1) + rec := &recordingPrefixDeleteStore{MVCCStore: fsm.store} + fsm.store = rec + fsm.pendingApplyIdx = 1234 + + require.NoError(t, fsm.handleDelPrefix(ctx, []byte("b/"), 101)) + require.Equal(t, 1, rec.batchCalls) + require.Zero(t, rec.singleCalls) + require.Equal(t, uint64(101), rec.commitTS) + require.Equal(t, uint64(1234), rec.appliedIndex) + require.Equal(t, []store.PrefixDelete{ + {Prefix: []byte("b/"), ExcludePrefix: txnCommonPrefix}, + { + Prefix: distribution.MigrationStagedDataKey(9, []byte("b/")), + ExcludePrefix: distribution.MigrationStagedDataKey(9, txnCommonPrefix), + }, + }, rec.deletes) +} + func TestFSMRejectsCurrentWriteFenceAfterObservedActiveRawPointWrite(t *testing.T) { t.Parallel() @@ -182,7 +414,7 @@ func TestFSMRejectsCurrentWriteFencedS3BucketAuxiliaryPointWrite(t *testing.T) { t.Parallel() ctx := context.Background() - const bucket = "bucket-a" + const bucket = "bucket-b" fsm := newS3BucketAuxiliaryWriteFencedFSM(t, bucket) for _, key := range [][]byte{ @@ -199,7 +431,7 @@ func TestFSMRejectsCurrentWriteFencedS3BucketAuxiliaryPointWrite(t *testing.T) { func TestFSMRejectsObservedWriteFencedS3BucketAuxiliaryPointWrite(t *testing.T) { t.Parallel() - const bucket = "bucket-a" + const bucket = "bucket-b" fsm := newS3BucketAuxiliaryWriteFencedFSM(t, bucket) err := fsm.handleRawRequest(context.Background(), &pb.Request{ @@ -211,6 +443,168 @@ func TestFSMRejectsObservedWriteFencedS3BucketAuxiliaryPointWrite(t *testing.T) require.ErrorIs(t, err, ErrRouteWriteFenced) } +func TestFSMIgnoresRawRouteFenceForS3BucketAuxiliaryWrite(t *testing.T) { + t.Parallel() + + const bucket = "bucket-a" + key := s3keys.BucketMetaKey(bucket) + engine := distribution.NewEngine() + routes := s3BucketAuxiliaryFenceRoutes(bucket, 1, 1) + routes[1].State = distribution.RouteStateActive + routes[2].State = distribution.RouteStateWriteFenced + applyComposed1Snapshot(t, engine, 1, routes) + + rawRoute, ok := engine.GetRoute(routeKey(key)) + require.True(t, ok) + require.Equal(t, distribution.RouteStateWriteFenced, rawRoute.State) + auxStart, auxEnd, ok := s3BucketAuxiliaryRouteRange(key) + require.True(t, ok) + auxRoutes := engine.GetIntersectingRoutes(auxStart, auxEnd) + require.NotEmpty(t, auxRoutes) + require.Equal(t, distribution.RouteStateActive, auxRoutes[0].State) + + fsm := newComposed1FSM(t, engine, 1) + err := fsm.handleRawRequest(context.Background(), &pb.Request{ + Mutations: []*pb.Mutation{{Op: pb.Op_PUT, Key: key, Value: []byte("meta")}}, + }, 100) + require.NoError(t, err) +} + +func TestFSMContinuesWriteFenceValidationAfterS3BucketAuxiliaryWrite(t *testing.T) { + t.Parallel() + + const bucket = "bucket-a" + auxiliaryKey := s3keys.BucketMetaKey(bucket) + fsm := newWriteFencedFSM(t) + + err := fsm.handleTxnRequest(context.Background(), &pb.Request{ + IsTxn: true, + Phase: pb.Phase_PREPARE, + Ts: 10, + Mutations: []*pb.Mutation{ + { + Op: pb.Op_PUT, + Key: []byte(txnMetaPrefix), + Value: EncodeTxnMeta(TxnMeta{PrimaryKey: auxiliaryKey, LockTTLms: defaultTxnLockTTLms}), + }, + {Op: pb.Op_PUT, Key: auxiliaryKey, Value: []byte("meta")}, + }, + }, 10) + require.NoError(t, err) + + err = fsm.handleTxnRequest(context.Background(), &pb.Request{ + IsTxn: true, + Phase: pb.Phase_PREPARE, + Ts: 11, + Mutations: []*pb.Mutation{ + { + Op: pb.Op_PUT, + Key: []byte(txnMetaPrefix), + Value: EncodeTxnMeta(TxnMeta{PrimaryKey: []byte("z"), LockTTLms: defaultTxnLockTTLms}), + }, + {Op: pb.Op_PUT, Key: []byte("z"), Value: []byte("v")}, + }, + }, 11) + require.ErrorIs(t, err, ErrRouteWriteFenced) +} + +func TestFSMIgnoresNonOwnerS3BucketAuxiliaryFenceForPointWrite(t *testing.T) { + t.Parallel() + + const bucket = "bucket-a" + key := s3keys.BucketMetaKey(bucket) + engine := distribution.NewEngine() + routes := s3BucketAuxiliarySplitRoutes(bucket, 5, 1, 1) + routes[2].State = distribution.RouteStateWriteFenced + applyComposed1Snapshot(t, engine, 1, routes) + + auxStart, auxEnd, ok := s3BucketAuxiliaryRouteRange(key) + require.True(t, ok) + auxRoutes := engine.GetIntersectingRoutes(auxStart, auxEnd) + require.Len(t, auxRoutes, 2) + require.Equal(t, distribution.RouteStateActive, auxRoutes[0].State) + require.Equal(t, distribution.RouteStateWriteFenced, auxRoutes[1].State) + + fsm := newComposed1FSM(t, engine, 1) + err := fsm.handleRawRequest(context.Background(), &pb.Request{ + Mutations: []*pb.Mutation{{Op: pb.Op_PUT, Key: key, Value: []byte("meta")}}, + }, 100) + require.NoError(t, err) +} + +func TestFSMComposed1UsesS3BucketAuxiliaryRouteOwner(t *testing.T) { + t.Parallel() + + const bucket = "bucket-b" + key := s3keys.BucketMetaKey(bucket) + engine := distribution.NewEngine() + applyComposed1Snapshot(t, engine, 1, s3BucketAuxiliaryStagedRoutes(bucket, 3, 4)) + fsm := newComposed1FSM(t, engine, 4) + + err := fsm.verifyComposed1(&pb.Request{ + IsTxn: true, + Phase: pb.Phase_PREPARE, + Ts: 10, + ObservedRouteVersion: 1, + Mutations: []*pb.Mutation{ + {Op: pb.Op_PUT, Key: []byte(txnMetaPrefix), Value: EncodeTxnMeta(TxnMeta{PrimaryKey: key, LockTTLms: defaultTxnLockTTLms})}, + {Op: pb.Op_PUT, Key: key, Value: []byte("meta")}, + }, + }) + require.NoError(t, err) +} + +func TestFSMIgnoresRawRouteFloorForS3BucketAuxiliaryWrite(t *testing.T) { + t.Parallel() + + const bucket = "bucket-a" + key := s3keys.BucketMetaKey(bucket) + engine := distribution.NewEngine() + routes := s3BucketAuxiliaryFenceRoutes(bucket, 1, 1) + routes[1].State = distribution.RouteStateActive + routes[2].MinWriteTSExclusive = ^uint64(0) + applyComposed1Snapshot(t, engine, 1, routes) + + rawRoute, ok := engine.GetRoute(routeKey(key)) + require.True(t, ok) + require.Equal(t, ^uint64(0), rawRoute.MinWriteTSExclusive) + auxStart, auxEnd, ok := s3BucketAuxiliaryRouteRange(key) + require.True(t, ok) + auxRoutes := engine.GetIntersectingRoutes(auxStart, auxEnd) + require.NotEmpty(t, auxRoutes) + require.Zero(t, auxRoutes[0].MinWriteTSExclusive) + + fsm := newComposed1FSM(t, engine, 1) + err := fsm.handleRawRequest(context.Background(), &pb.Request{ + Mutations: []*pb.Mutation{{Op: pb.Op_PUT, Key: key, Value: []byte("meta")}}, + }, 100) + require.NoError(t, err) +} + +func TestFSMIgnoresNonOwnerS3BucketAuxiliaryFloorForPointWrite(t *testing.T) { + t.Parallel() + + const bucket = "bucket-b" + key := s3keys.BucketMetaKey(bucket) + engine := distribution.NewEngine() + routes := s3BucketAuxiliarySplitRoutes(bucket, 1, 1, 1) + routes[2].MinWriteTSExclusive = ^uint64(0) + applyComposed1Snapshot(t, engine, 1, routes) + + auxStart, auxEnd, ok := s3BucketAuxiliaryRouteRange(key) + require.True(t, ok) + auxRoutes := engine.GetIntersectingRoutes(auxStart, auxEnd) + require.Len(t, auxRoutes, 2) + require.Zero(t, auxRoutes[0].MinWriteTSExclusive) + require.Equal(t, ^uint64(0), auxRoutes[1].MinWriteTSExclusive) + + fsm := newComposed1FSM(t, engine, 1) + err := fsm.handleRawRequest(context.Background(), &pb.Request{ + Mutations: []*pb.Mutation{{Op: pb.Op_PUT, Key: key, Value: []byte("meta")}}, + }, 100) + require.NoError(t, err) +} + func TestFSMRejectsCurrentWriteFencedDelPrefix(t *testing.T) { t.Parallel() @@ -319,3 +713,98 @@ func TestFSMRejectsObservedWriteFencedPrepareButAllowsAbort(t *testing.T) { } require.NotErrorIs(t, fsm.handleTxnRequest(ctx, abort, 11), ErrRouteWriteFenced) } + +func TestFSMRejectsRawPointWriteAtMigrationTimestampFloorDuringApply(t *testing.T) { + t.Parallel() + + ctx := context.Background() + fsm := newWriteFloorFSM(t) + err := fsm.handleRawRequest(ctx, &pb.Request{ + ObservedRouteVersion: 1, + Mutations: []*pb.Mutation{{Op: pb.Op_PUT, Key: []byte("z"), Value: []byte("replayed")}}, + }, 100) + require.ErrorIs(t, err, ErrRouteWriteTimestampTooLow) + _, getErr := fsm.store.GetAt(ctx, []byte("z"), ^uint64(0)) + require.ErrorIs(t, getErr, store.ErrKeyNotFound) +} + +func TestFSMRejectsDelPrefixAtMigrationTimestampFloorDuringApply(t *testing.T) { + t.Parallel() + + ctx := context.Background() + fsm := newWriteFloorFSM(t) + require.NoError(t, fsm.store.PutAt(ctx, []byte("z"), []byte("v"), 10, 0)) + + err := fsm.handleRawRequest(ctx, &pb.Request{ + ObservedRouteVersion: 1, + Mutations: []*pb.Mutation{{Op: pb.Op_DEL_PREFIX, Key: []byte("z")}}, + }, 100) + require.ErrorIs(t, err, ErrRouteWriteTimestampTooLow) + + got, getErr := fsm.store.GetAt(ctx, []byte("z"), ^uint64(0)) + require.NoError(t, getErr) + require.Equal(t, []byte("v"), got) +} + +func TestFSMRejectsOnePhaseTxnAtMigrationTimestampFloorDuringApply(t *testing.T) { + t.Parallel() + + ctx := context.Background() + fsm := newWriteFloorFSM(t) + req := &pb.Request{ + IsTxn: true, + Phase: pb.Phase_NONE, + Ts: 90, + ObservedRouteVersion: 1, + Mutations: []*pb.Mutation{ + {Op: pb.Op_PUT, Key: []byte(txnMetaPrefix), Value: EncodeTxnMeta(TxnMeta{PrimaryKey: []byte("z"), CommitTS: 100})}, + {Op: pb.Op_PUT, Key: []byte("z"), Value: []byte("low")}, + }, + } + err := fsm.handleTxnRequest(ctx, req, 100) + require.ErrorIs(t, err, ErrRouteWriteTimestampTooLow) + _, getErr := fsm.store.GetAt(ctx, []byte("z"), ^uint64(0)) + require.ErrorIs(t, getErr, store.ErrKeyNotFound) +} + +func TestFSMRejectsPrepareAtMigrationTimestampFloorDuringApply(t *testing.T) { + t.Parallel() + + ctx := context.Background() + fsm := newWriteFloorFSM(t) + prepare := &pb.Request{ + IsTxn: true, + Phase: pb.Phase_PREPARE, + Ts: 90, + ObservedRouteVersion: 1, + Mutations: []*pb.Mutation{ + {Op: pb.Op_PUT, Key: []byte(txnMetaPrefix), Value: EncodeTxnMeta(TxnMeta{PrimaryKey: []byte("z"), LockTTLms: defaultTxnLockTTLms})}, + {Op: pb.Op_PUT, Key: []byte("z"), Value: []byte("v")}, + }, + } + require.ErrorIs(t, fsm.handleTxnRequest(ctx, prepare, 90), ErrRouteWriteTimestampTooLow) +} + +func TestFSMTimestampFloorUsesObservedSnapshotDuringApply(t *testing.T) { + t.Parallel() + + ctx := context.Background() + engine := distribution.NewEngine() + applyComposed1Snapshot(t, engine, 1, []distribution.RouteDescriptor{ + {RouteID: 1, Start: []byte(""), End: nil, GroupID: 1, State: distribution.RouteStateActive}, + }) + fsm := newComposed1FSM(t, engine, 1) + applyComposed1Snapshot(t, engine, 2, []distribution.RouteDescriptor{ + {RouteID: 1, Start: []byte(""), End: nil, GroupID: 1, State: distribution.RouteStateActive, MinWriteTSExclusive: ^uint64(0)}, + }) + + err := fsm.handleRawRequest(ctx, &pb.Request{ + ObservedRouteVersion: 1, + Mutations: []*pb.Mutation{{Op: pb.Op_PUT, Key: []byte("z"), Value: []byte("proposed-before-floor")}}, + }, 100) + require.NoError(t, err) + + got, getErr := fsm.store.GetAt(ctx, []byte("z"), ^uint64(0)) + require.NoError(t, getErr) + require.Equal(t, []byte("proposed-before-floor"), got) +} diff --git a/kv/fsm_migration_import.go b/kv/fsm_migration_import.go new file mode 100644 index 000000000..43dfb285d --- /dev/null +++ b/kv/fsm_migration_import.go @@ -0,0 +1,131 @@ +package kv + +import ( + "bytes" + "context" + + "github.com/bootjp/elastickv/distribution" + pb "github.com/bootjp/elastickv/proto" + "github.com/bootjp/elastickv/store" + "github.com/cockroachdb/errors" + "google.golang.org/protobuf/proto" +) + +var ErrMigrationImportApply = errors.New("migration import: FSM apply failed; halting apply") + +// MarshalMigrationImportCommand encodes a target-group migration import batch +// as a Raft FSM command. The target Internal RPC handler uses this instead of +// mutating its local store directly so an acknowledged batch has been applied +// by the target group's voters. +func MarshalMigrationImportCommand(req *pb.ImportRangeVersionsRequest) ([]byte, error) { + if req == nil { + return nil, errors.WithStack(ErrInvalidRequest) + } + b, err := proto.Marshal(req) + if err != nil { + return nil, errors.WithStack(err) + } + if len(b) >= maxMarshaledCommandSize { + return nil, errors.New("marshaled migration import request too large") + } + return prependByte(raftEncodeMigrationImport, b), nil +} + +func (f *kvFSM) applyMigrationImport(ctx context.Context, data []byte) any { + req := &pb.ImportRangeVersionsRequest{} + if err := proto.Unmarshal(data, req); err != nil { + return haltErr(errors.Wrap(errors.Mark(err, ErrMigrationImportApply), "kv/fsm: decode migration import")) + } + staged := migrationStoreVersionsFromProto(req.GetJobId(), req.GetVersions()) + if err := validateStagedKeySizes(staged); err != nil { + // Ordinary, not halting: the batch is rejected and the migrator can + // surface it, where halting would wedge apply on every voter. + return errors.Wrap(err, "kv/fsm: apply migration import") + } + result, err := f.store.ImportVersionsRaft(ctx, store.ImportVersionsOptions{ + JobID: req.GetJobId(), + AppliedIndex: f.pendingApplyIdx, + BracketID: req.GetBracketId(), + BatchSeq: req.GetBatchSeq(), + Cursor: req.GetCursor(), + Versions: staged, + }) + if err != nil { + if isMigrationImportOrdinaryApplyError(err) { + return errors.Wrap(err, "kv/fsm: apply migration import") + } + return haltErr(errors.Wrap(errors.Mark(err, ErrMigrationImportApply), "kv/fsm: apply migration import")) + } + result.MaxImportedTS, err = f.migrationHLCFloorForApply(ctx, req, result) + if err != nil { + return haltErr(errors.Wrap(errors.Mark(err, ErrMigrationImportApply), "kv/fsm: migration import hlc floor")) + } + if f.hlc != nil && result.MaxImportedTS > 0 { + f.hlc.Observe(result.MaxImportedTS) + } + return result +} + +// isMigrationImportOrdinaryApplyError reports whether err is a verdict on the +// request bytes rather than a failure of this replica's store. Only those may +// be returned as an ordinary apply error: the engine advances setApplied past +// a response that does not implement HaltApply, and every replica applying the +// same entry decides an ordinary error identically, so the group stays in step. +// +// A store-side failure -- Pebble I/O, an encryption gate rejecting the write -- +// is per-replica. Letting it advance setApplied would leave the leader acking +// the batch while the failed voter skips the imported versions for good, which +// surfaces as missing data after failover or promotion. Those halt instead, +// matching applyMigrationPromote. +// validateStagedKeySizes rejects a batch whose staged form would not fit in the +// physical stored-key snapshot budget. +// +// MigrationStagedDataKey prepends the staged prefix, the job id, and a +// separator. store.MaxSnapshotStoredKeySize includes bounded headroom for that +// envelope so source keys at store.MaxSnapshotKeySize remain migratable, while +// still refusing batches that would make the target snapshot unrestorable. +func validateStagedKeySizes(versions []store.MVCCVersion) error { + for _, version := range versions { + if len(version.Key) > store.MaxSnapshotStoredKeySize { + return errors.Wrapf(store.ErrSnapshotKeyTooLarge, + "staged key length %d > %d", len(version.Key), store.MaxSnapshotStoredKeySize) + } + } + return nil +} + +func isMigrationImportOrdinaryApplyError(err error) bool { + return errors.Is(err, store.ErrImportBatchGap) || + errors.Is(err, store.ErrInvalidImportVersion) || + errors.Is(err, store.ErrValueTooLarge) || + errors.Is(err, store.ErrSnapshotKeyTooLarge) +} + +func (f *kvFSM) migrationHLCFloorForApply(ctx context.Context, req *pb.ImportRangeVersionsRequest, result store.ImportVersionsResult) (uint64, error) { + if result.MaxImportedTS > 0 || len(req.GetVersions()) == 0 { + return result.MaxImportedTS, nil + } + floor, err := f.store.MigrationHLCFloor(ctx, req.GetJobId()) + if err != nil { + return 0, errors.WithStack(err) + } + return floor, nil +} + +func migrationStoreVersionsFromProto(jobID uint64, in []*pb.MVCCVersion) []store.MVCCVersion { + out := make([]store.MVCCVersion, 0, len(in)) + for _, version := range in { + if version == nil { + continue + } + out = append(out, store.MVCCVersion{ + Key: distribution.MigrationStagedDataKey(jobID, version.GetKey()), + CommitTS: version.GetCommitTs(), + Tombstone: version.GetTombstone(), + Value: bytes.Clone(version.GetValue()), + KeyFamily: version.GetKeyFamily(), + ExpireAt: version.GetExpireAt(), + }) + } + return out +} diff --git a/kv/fsm_migration_import_test.go b/kv/fsm_migration_import_test.go new file mode 100644 index 000000000..5345549fc --- /dev/null +++ b/kv/fsm_migration_import_test.go @@ -0,0 +1,325 @@ +package kv + +import ( + "bytes" + "context" + "testing" + + "github.com/bootjp/elastickv/distribution" + pb "github.com/bootjp/elastickv/proto" + "github.com/bootjp/elastickv/store" + "github.com/cockroachdb/errors" + "github.com/stretchr/testify/require" + "google.golang.org/protobuf/proto" +) + +func TestMigrationStoreVersionsFromProtoStagesKeys(t *testing.T) { + t.Parallel() + + rawKey := []byte("user|k") + value := []byte("value") + got := migrationStoreVersionsFromProto(7, []*pb.MVCCVersion{ + nil, + { + Key: rawKey, + CommitTs: 11, + Value: value, + KeyFamily: distribution.MigrationFamilyUser, + ExpireAt: 123, + }, + }) + + require.Len(t, got, 1) + require.Equal(t, distribution.MigrationStagedDataKey(7, []byte("user|k")), got[0].Key) + require.Equal(t, uint64(11), got[0].CommitTS) + require.Equal(t, []byte("value"), got[0].Value) + require.Equal(t, distribution.MigrationFamilyUser, got[0].KeyFamily) + require.Equal(t, uint64(123), got[0].ExpireAt) + + rawKey[0] = 'X' + value[0] = 'X' + require.Equal(t, distribution.MigrationStagedDataKey(7, []byte("user|k")), got[0].Key) + require.Equal(t, []byte("value"), got[0].Value) +} + +func TestApplyMigrationImportWritesOnlyStagedKeys(t *testing.T) { + t.Parallel() + + ctx := context.Background() + st := store.NewMVCCStore() + hlc := NewHLC() + fsm := &kvFSM{store: st, hlc: hlc} + req := &pb.ImportRangeVersionsRequest{ + JobId: 9, + BracketId: 1, + BatchSeq: 1, + Cursor: []byte("cursor"), + Versions: []*pb.MVCCVersion{ + {Key: []byte("user|k"), CommitTs: 10, Value: []byte("v")}, + }, + } + data, err := proto.Marshal(req) + require.NoError(t, err) + + applied := fsm.applyMigrationImport(ctx, data) + result, ok := applied.(store.ImportVersionsResult) + require.True(t, ok, "got %T: %v", applied, applied) + require.Equal(t, []byte("cursor"), result.AckedCursor) + require.Equal(t, uint64(10), result.MaxImportedTS) + require.GreaterOrEqual(t, hlc.Current(), uint64(10)) + + staged := distribution.MigrationStagedDataKey(9, []byte("user|k")) + got, err := st.GetAt(ctx, staged, 10) + require.NoError(t, err) + require.Equal(t, []byte("v"), got) + _, err = st.GetAt(ctx, []byte("user|k"), 10) + require.ErrorIs(t, err, store.ErrKeyNotFound) +} + +type captureMigrationImportStore struct { + store.MVCCStore + opts store.ImportVersionsOptions +} + +func (s *captureMigrationImportStore) ImportVersionsRaft(_ context.Context, opts store.ImportVersionsOptions) (store.ImportVersionsResult, error) { + s.opts = opts + return store.ImportVersionsResult{AckedCursor: opts.Cursor, MaxImportedTS: 10}, nil +} + +func TestApplyMigrationImportThreadsPendingApplyIndex(t *testing.T) { + t.Parallel() + + capturing := &captureMigrationImportStore{} + fsm := &kvFSM{store: capturing, pendingApplyIdx: 1234} + req := &pb.ImportRangeVersionsRequest{ + JobId: 9, + BracketId: 1, + BatchSeq: 1, + Cursor: []byte("cursor"), + Versions: []*pb.MVCCVersion{ + {Key: []byte("user|k"), CommitTs: 10, Value: []byte("v")}, + }, + } + data, err := proto.Marshal(req) + require.NoError(t, err) + + applied := fsm.applyMigrationImport(context.Background(), data) + result, ok := applied.(store.ImportVersionsResult) + require.True(t, ok, "got %T: %v", applied, applied) + require.Equal(t, []byte("cursor"), result.AckedCursor) + require.Equal(t, uint64(1234), capturing.opts.AppliedIndex) + require.Equal(t, uint64(9), capturing.opts.JobID) + require.Len(t, capturing.opts.Versions, 1) + require.Equal(t, distribution.MigrationStagedDataKey(9, []byte("user|k")), capturing.opts.Versions[0].Key) +} + +// failingMigrationImportStore fails ImportVersionsRaft with a fixed error so +// the FSM's ordinary-vs-halt classification can be exercised directly. +type failingMigrationImportStore struct { + store.MVCCStore + importErr error + floorErr error + maxTS uint64 +} + +func (s *failingMigrationImportStore) ImportVersionsRaft(_ context.Context, opts store.ImportVersionsOptions) (store.ImportVersionsResult, error) { + if s.importErr != nil { + return store.ImportVersionsResult{}, s.importErr + } + return store.ImportVersionsResult{AckedCursor: opts.Cursor, MaxImportedTS: s.maxTS}, nil +} + +func (s *failingMigrationImportStore) MigrationHLCFloor(context.Context, uint64) (uint64, error) { + if s.floorErr != nil { + return 0, s.floorErr + } + return 0, nil +} + +func migrationImportCommandPayload(t *testing.T) []byte { + t.Helper() + + data, err := proto.Marshal(&pb.ImportRangeVersionsRequest{ + JobId: 9, + BracketId: 1, + BatchSeq: 1, + Cursor: []byte("cursor"), + Versions: []*pb.MVCCVersion{ + {Key: []byte("user|k"), CommitTs: 10, Value: []byte("v")}, + }, + }) + require.NoError(t, err) + + return data +} + +// A store-side import failure is per-replica: the leader can apply the batch +// and ack the RPC while this voter skips the imported versions for good. The +// response must therefore halt the apply loop instead of letting the engine +// advance setApplied past the entry. +func TestApplyMigrationImportHaltsOnStoreFailure(t *testing.T) { + t.Parallel() + + pebbleIOErr := errors.New("pebble: background error") + fsm := &kvFSM{store: &failingMigrationImportStore{importErr: pebbleIOErr}} + + applied := fsm.applyMigrationImport(context.Background(), migrationImportCommandPayload(t)) + + err := haltApplyOf(applied) + require.Error(t, err, "store failure must halt apply, got %T: %v", applied, applied) + require.True(t, errors.Is(err, ErrMigrationImportApply), "got %v", err) + require.ErrorIs(t, err, pebbleIOErr) +} + +// A failed HLC-floor read is a store read on this replica only, so it halts +// for the same reason. Replay after restart is safe: the import batch already +// committed, and the replayed entry is recognised as a duplicate. +func TestApplyMigrationImportHaltsOnHLCFloorFailure(t *testing.T) { + t.Parallel() + + floorErr := errors.New("pebble: read failed") + fsm := &kvFSM{store: &failingMigrationImportStore{maxTS: 0, floorErr: floorErr}} + + applied := fsm.applyMigrationImport(context.Background(), migrationImportCommandPayload(t)) + + err := haltApplyOf(applied) + require.Error(t, err, "hlc floor failure must halt apply, got %T: %v", applied, applied) + require.True(t, errors.Is(err, ErrMigrationImportApply), "got %v", err) + require.ErrorIs(t, err, floorErr) +} + +func TestApplyMigrationImportHaltsOnUndecodablePayload(t *testing.T) { + t.Parallel() + + fsm := &kvFSM{store: &failingMigrationImportStore{}} + + applied := fsm.applyMigrationImport(context.Background(), []byte{0xff, 0xff, 0xff, 0xff}) + + err := haltApplyOf(applied) + require.Error(t, err, "undecodable payload must halt apply, got %T: %v", applied, applied) + require.True(t, errors.Is(err, ErrMigrationImportApply), "got %v", err) +} + +// Verdicts on the request bytes are reached identically by every replica, so +// they stay ordinary errors: the RPC caller sees them and the group advances +// setApplied in step. Halting on these would turn a malformed request into a +// cluster-wide outage. +func TestApplyMigrationImportOrdinaryErrorsDoNotHalt(t *testing.T) { + t.Parallel() + + tests := map[string]error{ + "batch gap": store.ErrImportBatchGap, + "invalid version": store.ErrInvalidImportVersion, + "value too large": store.ErrValueTooLarge, + "wrapped batch gap": errors.Wrap(store.ErrImportBatchGap, "context"), + "wrapped bad version": errors.Wrap(store.ErrInvalidImportVersion, "context"), + } + for name, importErr := range tests { + t.Run(name, func(t *testing.T) { + t.Parallel() + + fsm := &kvFSM{store: &failingMigrationImportStore{importErr: importErr}} + + applied := fsm.applyMigrationImport(context.Background(), migrationImportCommandPayload(t)) + + require.NoError(t, haltApplyOf(applied), "request-shaped error must not halt apply") + err, ok := applied.(error) + require.True(t, ok, "got %T: %v", applied, applied) + require.ErrorIs(t, err, importErr) + require.False(t, errors.Is(err, ErrMigrationImportApply)) + }) + } +} + +// MigrationStagedDataKey prepends the staged prefix, the job id, and a +// separator. Source keys that already fit the ordinary snapshot key budget must +// stay migratable after that envelope is added, while truly oversized staged +// forms are still rejected before import. +func TestValidateStagedKeySizesAllowsSnapshotLimitSourceKey(t *testing.T) { + t.Parallel() + + envelope := len(distribution.MigrationStagedDataKey(1, nil)) + require.Positive(t, envelope) + require.LessOrEqual(t, envelope, store.MaxSnapshotInternalKeyEnvelope) + + atLimit := bytes.Repeat([]byte("k"), store.MaxSnapshotKeySize) + staged := migrationStoreVersionsFromProto(9, []*pb.MVCCVersion{{Key: atLimit, CommitTs: 10}}) + require.Len(t, staged, 1) + require.Greater(t, len(staged[0].Key), store.MaxSnapshotKeySize) + require.LessOrEqual(t, len(staged[0].Key), store.MaxSnapshotStoredKeySize) + require.NoError(t, validateStagedKeySizes(staged)) + + tooLong := bytes.Repeat([]byte("k"), store.MaxSnapshotStoredKeySize-envelope+1) + oversized := migrationStoreVersionsFromProto(9, []*pb.MVCCVersion{{Key: tooLong, CommitTs: 10}}) + require.Len(t, oversized, 1) + require.Equal(t, store.MaxSnapshotStoredKeySize+1, len(oversized[0].Key)) + require.ErrorIs(t, validateStagedKeySizes(oversized), store.ErrSnapshotKeyTooLarge) +} + +// The rejection must be an ordinary apply error. Halting would wedge apply on +// every voter over a batch the migrator can simply be told about. +func TestOversizedStagedKeyIsAnOrdinaryApplyError(t *testing.T) { + t.Parallel() + + require.True(t, isMigrationImportOrdinaryApplyError( + errors.Wrap(store.ErrSnapshotKeyTooLarge, "staged"))) +} + +// The size check has to be wired into the apply path, not merely available: +// source keys at the snapshot limit must be accepted, but staged forms beyond +// the stored-key snapshot budget must be refused before they are written. +func TestApplyMigrationImportAllowsSnapshotLimitSourceKey(t *testing.T) { + t.Parallel() + + ctx := context.Background() + st := store.NewMVCCStore() + t.Cleanup(func() { _ = st.Close() }) + fsm := &kvFSM{store: st, hlc: NewHLC()} + + atLimit := bytes.Repeat([]byte("k"), store.MaxSnapshotKeySize) + data, err := proto.Marshal(&pb.ImportRangeVersionsRequest{ + JobId: 9, + BracketId: 1, + BatchSeq: 1, + Versions: []*pb.MVCCVersion{{Key: atLimit, CommitTs: 10, Value: []byte("v")}}, + }) + require.NoError(t, err) + + applied := fsm.applyMigrationImport(ctx, data) + result, ok := applied.(store.ImportVersionsResult) + require.True(t, ok, "got %T: %v", applied, applied) + require.Equal(t, uint64(10), result.MaxImportedTS) + + got, getErr := st.GetAt(ctx, distribution.MigrationStagedDataKey(9, atLimit), 100) + require.NoError(t, getErr) + require.Equal(t, []byte("v"), got) +} + +func TestApplyMigrationImportRejectsOversizedStagedKey(t *testing.T) { + t.Parallel() + + ctx := context.Background() + st := store.NewMVCCStore() + t.Cleanup(func() { _ = st.Close() }) + fsm := &kvFSM{store: st, hlc: NewHLC()} + + envelope := len(distribution.MigrationStagedDataKey(9, nil)) + oversizedSource := bytes.Repeat([]byte("k"), store.MaxSnapshotStoredKeySize-envelope+1) + data, err := proto.Marshal(&pb.ImportRangeVersionsRequest{ + JobId: 9, + BracketId: 1, + BatchSeq: 1, + Versions: []*pb.MVCCVersion{{Key: oversizedSource, CommitTs: 10, Value: []byte("v")}}, + }) + require.NoError(t, err) + + applied := fsm.applyMigrationImport(ctx, data) + applyErr, ok := applied.(error) + require.True(t, ok, "got %T: %v", applied, applied) + require.ErrorIs(t, applyErr, store.ErrSnapshotKeyTooLarge) + require.NotErrorIs(t, applyErr, ErrMigrationImportApply, "must not halt apply") + + // Nothing was staged. + _, getErr := st.GetAt(ctx, distribution.MigrationStagedDataKey(9, oversizedSource), 100) + require.ErrorIs(t, getErr, store.ErrKeyNotFound) +} diff --git a/kv/fsm_migration_promote.go b/kv/fsm_migration_promote.go new file mode 100644 index 000000000..83365dde4 --- /dev/null +++ b/kv/fsm_migration_promote.go @@ -0,0 +1,123 @@ +package kv + +import ( + "context" + + "github.com/bootjp/elastickv/distribution" + pb "github.com/bootjp/elastickv/proto" + "github.com/bootjp/elastickv/store" + "github.com/cockroachdb/errors" + "google.golang.org/protobuf/proto" +) + +const ( + defaultMigrationPromoteMaxVersions = 1024 + defaultMigrationPromoteMaxBytes = 4 << 20 + defaultMigrationPromoteMaxScannedBytes = defaultMigrationPromoteMaxBytes * 4 + + // Hard server-side ceilings. The defaults above only apply to a request + // that leaves a bound unset, so without these an operator or migrator could + // ask one apply to load, re-encrypt, and commit an unbounded amount of + // staged data in a single Pebble batch -- synchronously, in every voter's + // apply loop. Clamping keeps the incremental, bounded promotion this API + // promises: the caller simply gets more rounds through the cursor, which + // PromoteVersionsResult already returns. + // + // The clamp is a pure function of the request and these constants, so every + // replica derives the same bounds from the same command and apply stays + // deterministic. + maxMigrationPromoteMaxVersions = 8192 + maxMigrationPromoteMaxBytes = 32 << 20 + maxMigrationPromoteMaxScannedBytes = maxMigrationPromoteMaxBytes * 4 +) + +var ErrMigrationPromoteApply = errors.New("migration promote: FSM apply failed; halting apply") + +// MarshalMigrationPromoteCommand encodes a target-group staged-data promotion +// chunk as a Raft FSM command. +func MarshalMigrationPromoteCommand(req *pb.PromoteStagedVersionsRequest) ([]byte, error) { + if req == nil { + return nil, errors.WithStack(ErrInvalidRequest) + } + b, err := proto.Marshal(req) + if err != nil { + return nil, errors.WithStack(err) + } + if len(b) >= maxMarshaledCommandSize { + return nil, errors.New("marshaled migration promote request too large") + } + return prependByte(raftEncodeMigrationPromote, b), nil +} + +func (f *kvFSM) applyMigrationPromote(ctx context.Context, data []byte) any { + req := &pb.PromoteStagedVersionsRequest{} + if err := proto.Unmarshal(data, req); err != nil { + return haltErr(errors.Wrap(errors.Mark(err, ErrMigrationPromoteApply), "kv/fsm: decode migration promote")) + } + promoter, ok := f.store.(store.MigrationPromoter) + if !ok { + return haltErr(errors.Wrap(errors.Mark(store.ErrNotSupported, ErrMigrationPromoteApply), "kv/fsm: migration promote store")) + } + result, err := promoter.PromoteVersions(ctx, migrationPromoteOptionsFromProto(req, f.pendingApplyIdx)) + if err != nil { + if isMigrationPromoteOrdinaryApplyError(err) { + return errors.Wrap(err, "kv/fsm: apply migration promote") + } + return haltErr(errors.Wrap(errors.Mark(err, ErrMigrationPromoteApply), "kv/fsm: apply migration promote")) + } + if f.hlc != nil && result.MaxPromotedTS > 0 { + f.hlc.Observe(result.MaxPromotedTS) + } + return result +} + +func migrationPromoteOptionsFromProto(req *pb.PromoteStagedVersionsRequest, appliedIndex uint64) store.PromoteVersionsOptions { + // Clamped in the int domain the request already decodes into, so no + // widening conversion is introduced here. + maxVersions := int(req.GetMaxVersions()) + switch { + case maxVersions <= 0: + maxVersions = defaultMigrationPromoteMaxVersions + case maxVersions > maxMigrationPromoteMaxVersions: + maxVersions = maxMigrationPromoteMaxVersions + } + maxBytes := clampMigrationPromoteBound( + req.GetMaxBytes(), defaultMigrationPromoteMaxBytes, maxMigrationPromoteMaxBytes) + maxScannedBytes := clampMigrationPromoteBound( + req.GetMaxScannedBytes(), defaultMigrationPromoteMaxScannedBytes, maxMigrationPromoteMaxScannedBytes) + prefix := distribution.MigrationStagedDataKeyPrefix(req.GetJobId()) + return store.PromoteVersionsOptions{ + JobID: req.GetJobId(), + AppliedIndex: appliedIndex, + StartKey: prefix, + EndKey: store.PrefixScanEnd(prefix), + Cursor: req.GetCursor(), + MaxVersions: maxVersions, + MaxBytes: maxBytes, + MaxScannedBytes: maxScannedBytes, + TargetKey: migrationPromoteTargetKey(req.GetJobId()), + } +} + +// clampMigrationPromoteBound resolves one promotion bound: unset takes the +// default, anything above the hard ceiling is clamped down to it. +func clampMigrationPromoteBound(requested, fallback, ceiling uint64) uint64 { + if requested == 0 { + return fallback + } + return min(requested, ceiling) +} + +func isMigrationPromoteOrdinaryApplyError(err error) bool { + return errors.Is(err, store.ErrInvalidExportCursor) +} + +func migrationPromoteTargetKey(jobID uint64) func([]byte) ([]byte, bool) { + return func(stagedKey []byte) ([]byte, bool) { + gotJobID, rawKey, ok := distribution.MigrationStagedDataKeyParts(stagedKey) + if !ok || gotJobID != jobID { + return nil, false + } + return rawKey, true + } +} diff --git a/kv/fsm_migration_promote_test.go b/kv/fsm_migration_promote_test.go new file mode 100644 index 000000000..570bf9a80 --- /dev/null +++ b/kv/fsm_migration_promote_test.go @@ -0,0 +1,155 @@ +package kv + +import ( + "context" + "testing" + + "github.com/bootjp/elastickv/distribution" + pb "github.com/bootjp/elastickv/proto" + "github.com/bootjp/elastickv/store" + "github.com/cockroachdb/errors" + "github.com/stretchr/testify/require" +) + +func TestMigrationPromoteTargetKeyRestoresRawKey(t *testing.T) { + t.Parallel() + + targetKey := migrationPromoteTargetKey(9) + raw, ok := targetKey(distribution.MigrationStagedDataKey(9, []byte("user|k"))) + require.True(t, ok) + require.Equal(t, []byte("user|k"), raw) + + _, ok = targetKey(distribution.MigrationStagedDataKey(10, []byte("user|k"))) + require.False(t, ok) +} + +func TestApplyMigrationPromoteMovesStagedVersions(t *testing.T) { + t.Parallel() + + ctx := context.Background() + st := store.NewMVCCStore() + hlc := NewHLC() + fsm := &kvFSM{store: st, hlc: hlc} + staged := distribution.MigrationStagedDataKey(9, []byte("user|k")) + require.NoError(t, st.PutAt(ctx, staged, []byte("v10"), 10, 0)) + require.NoError(t, st.DeleteAt(ctx, staged, 20)) + + cmd, err := MarshalMigrationPromoteCommand(&pb.PromoteStagedVersionsRequest{ + JobId: 9, + MaxVersions: 10, + }) + require.NoError(t, err) + applied := fsm.Apply(cmd) + result, ok := applied.(store.PromoteVersionsResult) + require.True(t, ok, "got %T: %v", applied, applied) + require.True(t, result.Done) + require.Equal(t, uint64(2), result.PromotedRows) + require.Equal(t, uint64(2), result.TotalPromotedRows) + require.Equal(t, uint64(20), result.MaxPromotedTS) + require.GreaterOrEqual(t, hlc.Current(), uint64(20)) + stateReader, ok := st.(store.MigrationPromotionStateReader) + require.True(t, ok) + state, ok, err := stateReader.MigrationPromotionState(ctx, 9) + require.NoError(t, err) + require.True(t, ok) + require.True(t, state.Done) + require.Equal(t, uint64(2), state.PromotedRows) + require.Equal(t, uint64(20), state.MaxPromotedTS) + + got, err := st.GetAt(ctx, []byte("user|k"), 10) + require.NoError(t, err) + require.Equal(t, []byte("v10"), got) + _, err = st.GetAt(ctx, []byte("user|k"), 20) + require.ErrorIs(t, err, store.ErrKeyNotFound) + left, err := st.ExportVersions(ctx, store.ExportVersionsOptions{ + StartKey: distribution.MigrationStagedDataKeyPrefix(9), + EndKey: store.PrefixScanEnd(distribution.MigrationStagedDataKeyPrefix(9)), + MaxVersions: 10, + }) + require.NoError(t, err) + require.Empty(t, left.Versions) +} + +func TestApplyMigrationPromoteMalformedPayloadHalts(t *testing.T) { + t.Parallel() + + fsm := &kvFSM{store: store.NewMVCCStore()} + err := haltApplyOf(fsm.Apply([]byte{raftEncodeMigrationPromote, 0xff, 0xff})) + require.True(t, errors.Is(err, ErrMigrationPromoteApply), "got %v", err) +} + +func TestApplyMigrationPromoteInvalidCursorReturnsOrdinaryError(t *testing.T) { + t.Parallel() + + fsm := &kvFSM{store: store.NewMVCCStore()} + cmd, err := MarshalMigrationPromoteCommand(&pb.PromoteStagedVersionsRequest{ + Cursor: []byte{0xff}, + MaxVersions: 10, + }) + require.NoError(t, err) + resp := fsm.Apply(cmd) + require.Nil(t, haltApplyOf(resp)) + err, ok := resp.(error) + require.True(t, ok, "got %T: %v", resp, resp) + require.ErrorIs(t, err, store.ErrInvalidExportCursor) + require.False(t, errors.Is(err, ErrMigrationPromoteApply)) +} + +// PromoteStagedVersions runs inside FSM apply, synchronously, on every voter. +// The per-chunk bounds arrive in the Raft command, so a request that names an +// oversized batch makes one apply load, re-encrypt, and commit that much staged +// data in a single Pebble batch on every replica at once. The defaults only fill +// in unset bounds, so they are not limits; the hard ceilings are. +func TestMigrationPromoteOptionsClampOversizedBounds(t *testing.T) { + t.Parallel() + + opts := migrationPromoteOptionsFromProto(&pb.PromoteStagedVersionsRequest{ + JobId: 9, + MaxVersions: 1 << 30, + MaxBytes: 1 << 40, + MaxScannedBytes: 1 << 42, + }, 7) + + require.Equal(t, maxMigrationPromoteMaxVersions, opts.MaxVersions) + require.Equal(t, uint64(maxMigrationPromoteMaxBytes), opts.MaxBytes) + require.Equal(t, uint64(maxMigrationPromoteMaxScannedBytes), opts.MaxScannedBytes) + require.Equal(t, uint64(7), opts.AppliedIndex) +} + +// Unset bounds still take the defaults, and a request under the ceiling is +// passed through unchanged so a caller can still ask for smaller chunks. +func TestMigrationPromoteOptionsKeepDefaultsAndSmallerRequests(t *testing.T) { + t.Parallel() + + defaults := migrationPromoteOptionsFromProto(&pb.PromoteStagedVersionsRequest{JobId: 9}, 0) + require.Equal(t, defaultMigrationPromoteMaxVersions, defaults.MaxVersions) + require.Equal(t, uint64(defaultMigrationPromoteMaxBytes), defaults.MaxBytes) + require.Equal(t, uint64(defaultMigrationPromoteMaxScannedBytes), defaults.MaxScannedBytes) + + smaller := migrationPromoteOptionsFromProto(&pb.PromoteStagedVersionsRequest{ + JobId: 9, + MaxVersions: 16, + MaxBytes: 1024, + MaxScannedBytes: 4096, + }, 0) + require.Equal(t, 16, smaller.MaxVersions) + require.Equal(t, uint64(1024), smaller.MaxBytes) + require.Equal(t, uint64(4096), smaller.MaxScannedBytes) +} + +// The clamp has to be a pure function of the command so every replica derives +// the same bounds from the same entry and apply stays deterministic. +func TestMigrationPromoteOptionsAreDeterministicPerCommand(t *testing.T) { + t.Parallel() + + req := &pb.PromoteStagedVersionsRequest{JobId: 9, MaxVersions: 1 << 30, MaxBytes: 1 << 40} + first := migrationPromoteOptionsFromProto(req, 11) + second := migrationPromoteOptionsFromProto(req, 11) + // PromoteVersionsOptions carries a closure, which never compares equal, so + // the bounds this clamp owns are compared directly. + require.Equal(t, first.MaxVersions, second.MaxVersions) + require.Equal(t, first.MaxBytes, second.MaxBytes) + require.Equal(t, first.MaxScannedBytes, second.MaxScannedBytes) + require.Equal(t, first.StartKey, second.StartKey) + require.Equal(t, first.EndKey, second.EndKey) +} diff --git a/kv/fsm_onephase_dedup_test.go b/kv/fsm_onephase_dedup_test.go index aaf027c1d..c9f06c088 100644 --- a/kv/fsm_onephase_dedup_test.go +++ b/kv/fsm_onephase_dedup_test.go @@ -4,6 +4,8 @@ import ( "context" "testing" + "github.com/bootjp/elastickv/distribution" + "github.com/bootjp/elastickv/internal/s3keys" pb "github.com/bootjp/elastickv/proto" "github.com/bootjp/elastickv/store" "github.com/stretchr/testify/require" @@ -67,6 +69,51 @@ func TestOnePhaseDedup_NoOpsWhenPriorAttemptLanded(t *testing.T) { require.Equal(t, uint64(20), latest, "newest version must remain attempt 1's at 20") } +func TestOnePhaseDedup_NoOpsWhenPriorAttemptLandedAsStagedVersion(t *testing.T) { + t.Parallel() + ctx := context.Background() + st := store.NewMVCCStore() + fsm, ok := NewKvFSMWithHLC(st, NewHLC()).(*kvFSM) + require.True(t, ok) + + key := []byte("list-item") + require.NoError(t, st.PutAt(ctx, distribution.MigrationStagedDataKey(9, key), []byte("v"), 20, 0)) + + req := onePhaseReq(30, 40, 20, key, []byte("v")) + req.ReadKeys = [][]byte{distribution.MigrationStagedDataKey(9, key)} + require.NoError(t, applyFSMRequest(t, fsm, req)) + + at40, err := st.CommittedVersionAt(ctx, key, 40) + require.NoError(t, err) + require.False(t, at40, "retry must not write a live version when the prior attempt is staged") + stagedAt20, err := st.CommittedVersionAt(ctx, distribution.MigrationStagedDataKey(9, key), 20) + require.NoError(t, err) + require.True(t, stagedAt20) +} + +func TestOnePhaseDedup_NoOpsWhenS3AuxiliaryPriorAttemptLandedAsStagedVersion(t *testing.T) { + t.Parallel() + ctx := context.Background() + st := store.NewMVCCStore() + const bucket = "bucket-a" + fsm, ok := NewKvFSMWithHLC(st, NewHLC()).(*kvFSM) + require.True(t, ok) + + key := s3keys.BucketMetaKey(bucket) + require.NoError(t, st.PutAt(ctx, distribution.MigrationStagedDataKey(9, key), []byte("v"), 20, 0)) + + req := onePhaseReq(30, 40, 20, key, []byte("v")) + req.ReadKeys = [][]byte{distribution.MigrationStagedDataKey(9, key)} + require.NoError(t, applyFSMRequest(t, fsm, req)) + + at40, err := st.CommittedVersionAt(ctx, key, 40) + require.NoError(t, err) + require.False(t, at40, "retry must not write a live S3 auxiliary version when the prior attempt is staged") + stagedAt20, err := st.CommittedVersionAt(ctx, distribution.MigrationStagedDataKey(9, key), 20) + require.NoError(t, err) + require.True(t, stagedAt20) +} + // TestOnePhaseDedup_AppliesWhenPriorAttemptDidNotLand covers the truncated / // never-applied case: prev_commit_ts is set but no version exists at exactly // that timestamp (attempt 1's entry lost the log race). The probe misses and diff --git a/kv/fsm_reserved_control_test.go b/kv/fsm_reserved_control_test.go new file mode 100644 index 000000000..6a5077f64 --- /dev/null +++ b/kv/fsm_reserved_control_test.go @@ -0,0 +1,137 @@ +package kv + +import ( + "context" + "testing" + + "github.com/bootjp/elastickv/distribution" + pb "github.com/bootjp/elastickv/proto" + "github.com/bootjp/elastickv/store" + "github.com/stretchr/testify/require" +) + +// A RawKV mutation must not reach the migration and catalog control +// namespaces. Those are written only by the typed internal commands, and a user +// write that lands in one is later promoted as ordinary data. +func TestValidateRawMutationRejectsReservedControlKeys(t *testing.T) { + t.Parallel() + + ctx := context.Background() + st := store.NewMVCCStore() + t.Cleanup(func() { _ = st.Close() }) + f, ok := NewKvFSMWithHLC(st, NewHLC()).(*kvFSM) + require.True(t, ok) + + for _, key := range [][]byte{ + []byte("!dist|meta|version"), + []byte("!dist|route|0001"), + []byte("!dist|job|7"), + []byte("!dist|jobhist|7"), + distribution.MigrationStagedDataKey(7, []byte("victim")), + []byte("!migstage|7|victim"), + []byte("!migwrite|7"), + []byte("!migfence|7"), + } { + err := f.validateRawMutationForApply(ctx, &pb.Mutation{Op: pb.Op_PUT, Key: key, Value: []byte("v")}, nil, 10, nil) + require.ErrorIs(t, err, ErrInvalidRequest, "key %q must be refused", key) + } + + // Ordinary user keys are unaffected. + require.NoError(t, f.validateRawMutationForApply(ctx, + &pb.Mutation{Op: pb.Op_PUT, Key: []byte("user-key"), Value: []byte("v")}, nil, 10, nil)) +} + +// DEL_PREFIX never reaches validateRawMutationsForApply, so it is gated in +// handleDelPrefix. A partial spelling of a control prefix is refused too, +// because the delete would sweep the namespace up. +func TestHandleDelPrefixRejectsReservedControlPrefixes(t *testing.T) { + t.Parallel() + + ctx := context.Background() + st := store.NewMVCCStore() + t.Cleanup(func() { _ = st.Close() }) + f, ok := NewKvFSMWithHLC(st, NewHLC()).(*kvFSM) + require.True(t, ok) + + for _, prefix := range [][]byte{ + []byte("!dist|"), + []byte("!dist|route|"), + []byte("!dist"), + distribution.MigrationStagedDataKey(7, []byte("user:")), + []byte("!migwrite|"), + []byte("!migfence"), + } { + require.ErrorIs(t, f.handleDelPrefix(ctx, prefix, 10), ErrInvalidRequest, + "prefix %q must be refused", prefix) + } + + // The whole-keyspace flush is a deliberate operation and stays allowed. + require.NoError(t, f.handleDelPrefix(ctx, nil, 11)) + // So does an ordinary user prefix. + require.NoError(t, f.handleDelPrefix(ctx, []byte("user:"), 12)) +} + +// Staged migration data is private to typed migration paths. A user-supplied +// RawKV request must not be able to forge a staged key that promotion later +// treats as migrated data. +func TestValidateRawMutationRejectsStagedDataKeys(t *testing.T) { + t.Parallel() + + ctx := context.Background() + st := store.NewMVCCStore() + t.Cleanup(func() { _ = st.Close() }) + f, ok := NewKvFSMWithHLC(st, NewHLC()).(*kvFSM) + require.True(t, ok) + + staged := distribution.MigrationStagedDataKey(7, []byte("user-key")) + require.ErrorIs(t, f.validateRawMutationForApply(ctx, + &pb.Mutation{Op: pb.Op_PUT, Key: staged, Value: []byte("v")}, nil, 10, nil), ErrInvalidRequest) + require.ErrorIs(t, f.handleDelPrefix(ctx, distribution.MigrationStagedDataKey(7, []byte("user:")), 11), ErrInvalidRequest) +} + +// The transactional paths reach the store through their own helpers, which did +// not carry the reserved-key check the raw path has. A TransactionalKV request +// could therefore write catalog or staged-migration state directly, and a +// forged !migstage|| row is promoted as user data when the same group is +// the migration target. +func TestTxnMutationHelpersRejectReservedControlKeys(t *testing.T) { + t.Parallel() + + st := store.NewMVCCStore() + t.Cleanup(func() { _ = st.Close() }) + f, ok := NewKvFSMWithHLC(st, NewHLC()).(*kvFSM) + require.True(t, ok) + + for _, key := range [][]byte{ + distribution.MigrationStagedDataKey(7, []byte("victim")), + []byte("!migwrite|7"), + []byte("!migfence|7"), + } { + muts := []*pb.Mutation{{Op: pb.Op_PUT, Key: key, Value: []byte("v")}} + _, err := f.uniqueMutationsAboveFloor(muts, 10) + require.ErrorIs(t, err, ErrInvalidRequest, + "prepare/one-phase must refuse %q", key) + _, err = f.uniqueTxnMutationsAboveFloor(muts, 10) + require.ErrorIs(t, err, ErrInvalidRequest, + "commit must refuse %q", key) + } + + // Ordinary user keys, the transaction-internal keys the txn paths + // legitimately write, and the catalog records the control plane commits + // through the coordinator are unaffected. SplitRange is a transaction that + // writes !dist|route| and !dist|meta|, so refusing those here would break + // the control plane itself. + for _, key := range [][]byte{ + []byte("user-key"), + txnLockKey([]byte("user-key")), + txnIntentKey([]byte("user-key")), + []byte("!dist|route|0001"), + []byte("!dist|meta|version"), + } { + muts := []*pb.Mutation{{Op: pb.Op_PUT, Key: key, Value: []byte("v")}} + _, err := f.uniqueMutationsAboveFloor(muts, 10) + require.NoError(t, err, "key %q must be accepted", key) + _, err = f.uniqueTxnMutationsAboveFloor(muts, 10) + require.NoError(t, err, "key %q must be accepted", key) + } +} diff --git a/kv/fsm_txn_test.go b/kv/fsm_txn_test.go index db35beab5..97fb7d17f 100644 --- a/kv/fsm_txn_test.go +++ b/kv/fsm_txn_test.go @@ -5,6 +5,7 @@ import ( "fmt" "testing" + "github.com/bootjp/elastickv/distribution" pb "github.com/bootjp/elastickv/proto" "github.com/bootjp/elastickv/store" "github.com/stretchr/testify/require" @@ -213,6 +214,94 @@ func TestCommitIsIdempotentAfterCommitRecordExists(t *testing.T) { require.Equal(t, commitTS, gotCommitTS) } +func TestCommitApplyStartTSUsesStagedCommitRecord(t *testing.T) { + t.Parallel() + + ctx := context.Background() + st := store.NewMVCCStore() + fsm, ok := NewKvFSMWithHLC(st, NewHLC()).(*kvFSM) + require.True(t, ok) + + primary := []byte("pk") + startTS := uint64(10) + commitTS := uint64(20) + readKeys := [][]byte{distribution.MigrationStagedDataKey(9, primary)} + stagedCommitKey := distribution.MigrationStagedDataKey(9, txnCommitKey(primary, startTS)) + require.NoError(t, st.PutAt(ctx, stagedCommitKey, encodeTxnCommitRecord(commitTS), commitTS, 0)) + + applyStartTS, err := fsm.commitApplyStartTS(ctx, primary, startTS, commitTS, readKeys) + require.NoError(t, err) + require.Equal(t, commitTS, applyStartTS) +} + +func TestCommitRejectsStagedRollbackRecord(t *testing.T) { + t.Parallel() + + ctx := context.Background() + st := store.NewMVCCStore() + fsm, ok := NewKvFSMWithHLC(st, NewHLC()).(*kvFSM) + require.True(t, ok) + + primary := []byte("pk") + startTS := uint64(10) + commitTS := uint64(20) + readKeys := [][]byte{distribution.MigrationStagedDataKey(9, primary)} + stagedRollbackKey := distribution.MigrationStagedDataKey(9, txnRollbackKey(primary, startTS)) + require.NoError(t, st.PutAt(ctx, stagedRollbackKey, encodeTxnRollbackRecord(), startTS, 0)) + + _, err := fsm.commitApplyStartTS(ctx, primary, startTS, commitTS, readKeys) + require.Error(t, err) + require.ErrorIs(t, err, ErrTxnAlreadyAborted) +} + +func TestCommitCleansUpStagedPreparedArtifacts(t *testing.T) { + t.Parallel() + + ctx := context.Background() + st := store.NewMVCCStore() + fsm, ok := NewKvFSMWithHLC(st, NewHLC()).(*kvFSM) + require.True(t, ok) + + primary := []byte("pk") + startTS := uint64(10) + commitTS := uint64(20) + jobID := uint64(9) + stagedLockKey := distribution.MigrationStagedDataKey(jobID, txnLockKey(primary)) + stagedIntentKey := distribution.MigrationStagedDataKey(jobID, txnIntentKey(primary)) + require.NoError(t, st.PutAt(ctx, stagedLockKey, encodeTxnLock(txnLock{ + StartTS: startTS, + PrimaryKey: primary, + IsPrimaryKey: true, + }), startTS, 0)) + require.NoError(t, st.PutAt(ctx, stagedIntentKey, encodeTxnIntent(txnIntent{ + StartTS: startTS, + Op: txnIntentOpPut, + Value: []byte("v"), + }), startTS, 0)) + + commit := &pb.Request{ + IsTxn: true, + Phase: pb.Phase_COMMIT, + Ts: startTS, + ReadKeys: [][]byte{ + distribution.MigrationStagedDataKey(jobID, primary), + }, + Mutations: []*pb.Mutation{ + {Op: pb.Op_PUT, Key: []byte(txnMetaPrefix), Value: EncodeTxnMeta(TxnMeta{PrimaryKey: primary, CommitTS: commitTS})}, + {Op: pb.Op_PUT, Key: primary}, + }, + } + require.NoError(t, applyFSMRequest(t, fsm, commit)) + + value, err := st.GetAt(ctx, primary, ^uint64(0)) + require.NoError(t, err) + require.Equal(t, []byte("v"), value) + _, err = st.GetAt(ctx, stagedLockKey, ^uint64(0)) + require.ErrorIs(t, err, store.ErrKeyNotFound) + _, err = st.GetAt(ctx, stagedIntentKey, ^uint64(0)) + require.ErrorIs(t, err, store.ErrKeyNotFound) +} + func TestCommitIsIdempotentOnSecondaryShardWhenKeyAlreadyCommitted(t *testing.T) { t.Parallel() diff --git a/kv/leader_routed_store.go b/kv/leader_routed_store.go index c9aa87b7d..2a1ef53f8 100644 --- a/kv/leader_routed_store.go +++ b/kv/leader_routed_store.go @@ -594,6 +594,13 @@ func (s *LeaderRoutedStore) DeletePrefixAtRaftAt(ctx context.Context, prefix []b return errors.WithStack(s.local.DeletePrefixAtRaftAt(ctx, prefix, excludePrefix, commitTS, appliedIndex)) } +func (s *LeaderRoutedStore) DeletePrefixesAtRaftAt(ctx context.Context, deletes []store.PrefixDelete, commitTS, appliedIndex uint64) error { + if s == nil || s.local == nil { + return errors.WithStack(store.ErrNotSupported) + } + return errors.WithStack(s.local.DeletePrefixesAtRaftAt(ctx, deletes, commitTS, appliedIndex)) +} + // LastAppliedIndex forwards to the local store when it implements // raftengine.AppliedIndexReader. Defensive: in production today the // kvFSM holds a *pebbleStore directly (not a LeaderRoutedStore — that @@ -721,6 +728,14 @@ func (s *LeaderRoutedStore) ImportVersions(ctx context.Context, opts store.Impor return result, errors.WithStack(err) } +func (s *LeaderRoutedStore) ImportVersionsRaft(ctx context.Context, opts store.ImportVersionsOptions) (store.ImportVersionsResult, error) { + if s == nil || s.local == nil { + return store.ImportVersionsResult{}, errors.WithStack(store.ErrNotSupported) + } + result, err := s.local.ImportVersionsRaft(ctx, opts) + return result, errors.WithStack(err) +} + func (s *LeaderRoutedStore) MigrationHLCFloor(ctx context.Context, jobID uint64) (uint64, error) { if s == nil || s.local == nil { return 0, errors.WithStack(store.ErrNotSupported) diff --git a/kv/migrator_filter.go b/kv/migrator_filter.go index acbf1a796..7036e64bd 100644 --- a/kv/migrator_filter.go +++ b/kv/migrator_filter.go @@ -3,6 +3,7 @@ package kv import ( "bytes" + "github.com/bootjp/elastickv/distribution" "github.com/bootjp/elastickv/internal/s3keys" ) @@ -21,10 +22,18 @@ func RouteKeyFilterForGroup(rangeStart, rangeEnd []byte, sourceGroupID uint64, r end := bytes.Clone(rangeEnd) return func(rawKey []byte) bool { if resolver != nil { - if gid, ok := resolver.ResolveGroup(rawKey); ok { + // Transaction-family brackets carry the user key wrapped in + // !txn|int| / !txn|cmt| / ... . The resolver only recognises bare + // partitioned-family prefixes, so probing the wrapper would answer + // "not mine" and drop the key onto routeKey's !sqs|route|global + // collapse -- a different verdict from the one the embedded row + // gets in its own data bracket, which is how a migration ends up + // copying SQS rows while leaving their intents behind. + probe := partitionResolverProbeKey(rawKey) + if gid, ok := resolver.ResolveGroup(probe); ok { return gid == sourceGroupID } - if resolver.RecognisesPartitionedKey(rawKey) { + if resolver.RecognisesPartitionedKey(probe) { return false } } @@ -36,6 +45,17 @@ func RouteKeyFilterForGroup(rangeStart, rangeEnd []byte, sourceGroupID uint64, r } } +// partitionResolverProbeKey returns the key a PartitionResolver should be asked +// about: the embedded user key for a transaction-internal wrapper, the key +// itself otherwise. Route ownership of a txn key is the ownership of the user +// key it locks, so both must be resolved the same way. +func partitionResolverProbeKey(rawKey []byte) []byte { + if embedded, ok := txnRouteKey(rawKey); ok { + return embedded + } + return rawKey +} + func s3BucketAuxiliaryRouteInRange(rawKey, routeStart, routeEnd []byte) bool { bucketRouteStart, bucketRouteEnd, ok := s3BucketAuxiliaryRouteRange(rawKey) if !ok { @@ -59,6 +79,23 @@ func s3BucketAuxiliaryRouteRange(rawKey []byte) ([]byte, []byte, bool) { return bucketRouteStart, prefixScanEnd(bucketRouteStart), true } +func s3BucketAuxiliaryOwnerRoute(rawKey []byte, routes []distribution.Route) (distribution.Route, bool) { + start, end, ok := s3BucketAuxiliaryRouteRange(rawKey) + if !ok { + return distribution.Route{}, false + } + return s3BucketAuxiliaryOwnerRouteFromRange(start, end, routes) +} + +func s3BucketAuxiliaryOwnerRouteFromRange(start []byte, end []byte, routes []distribution.Route) (distribution.Route, bool) { + for _, route := range routes { + if migrationRouteRangesIntersect(route.Start, route.End, start, end) { + return route, true + } + } + return distribution.Route{}, false +} + func keyInMigrationRouteRange(key, routeStart, routeEnd []byte) bool { if key == nil { return false diff --git a/kv/migrator_filter_test.go b/kv/migrator_filter_test.go new file mode 100644 index 000000000..ba99d689d --- /dev/null +++ b/kv/migrator_filter_test.go @@ -0,0 +1,71 @@ +package kv + +import ( + "testing" + + "github.com/stretchr/testify/require" +) + +// A transaction-family bracket holds the user key wrapped in !txn|int| and +// friends. The group filter must resolve those through the same partition +// resolver that decides the embedded row's own bracket: a partitioned HT-FIFO +// SQS row and its intent belong to one group, and answering differently for +// the two is how a migration copies the data and strands the transaction +// state. +func TestRouteKeyFilterForGroupResolvesTxnWrappedPartitionedKeys(t *testing.T) { + t.Parallel() + + const ( + sourceGroup = uint64(7) + otherGroup = uint64(9) + ) + partitioned := []byte("!sqs|msg|data|p|q|0|m1") + elsewhere := []byte("!sqs|msg|data|p|q|1|m1") + resolver := &fakePartitionResolver{ + routes: map[string]uint64{ + string(partitioned): sourceGroup, + string(elsewhere): otherGroup, + }, + recognisedPrefix: []byte("!sqs|msg|data|p|"), + } + + // A route range that excludes the !sqs|route|global collapse target, so a + // filter that falls through to the byte-range path answers false. + filter := RouteKeyFilterForGroup([]byte("a"), []byte("b"), sourceGroup, resolver) + + for _, tc := range []struct { + name string + key []byte + want bool + }{ + {name: "bare partitioned row", key: partitioned, want: true}, + {name: "intent", key: txnIntentKey(partitioned), want: true}, + {name: "lock", key: txnLockKey(partitioned), want: true}, + {name: "meta", key: append(append([]byte{}, txnMetaPrefixBytes...), partitioned...), want: true}, + {name: "intent for another group", key: txnIntentKey(elsewhere), want: false}, + {name: "bare row for another group", key: elsewhere, want: false}, + } { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + require.Equal(t, tc.want, filter(tc.key), "key %q", tc.key) + }) + } +} + +// A recognised-but-unresolved partitioned key fails closed whether or not it +// arrives wrapped: the resolver cannot say which group owns it, and guessing +// through the byte-range route would export it from the wrong source. +func TestRouteKeyFilterForGroupFailsClosedOnWrappedUnresolvedKey(t *testing.T) { + t.Parallel() + + unresolved := []byte("!sqs|msg|data|p|unknown|0|m1") + resolver := &fakePartitionResolver{ + routes: map[string]uint64{}, + recognisedPrefix: []byte("!sqs|msg|data|p|"), + } + filter := RouteKeyFilterForGroup(nil, nil, 7, resolver) + + require.False(t, filter(unresolved)) + require.False(t, filter(txnIntentKey(unresolved))) + require.False(t, filter(txnCommitKey(unresolved, 42))) +} diff --git a/kv/route_history.go b/kv/route_history.go index c6ac85608..8bc50c58e 100644 --- a/kv/route_history.go +++ b/kv/route_history.go @@ -61,6 +61,14 @@ func (s distributionRouteSnapshot) OwnerOf(key []byte) (uint64, bool) { return s.snap.OwnerOf(key) } +func (s distributionRouteSnapshot) RouteOf(key []byte) (distribution.Route, bool) { + return s.snap.RouteOf(key) +} + +func (s distributionRouteSnapshot) IntersectingRoutes(start, end []byte) []distribution.Route { + return s.snap.IntersectingRoutes(start, end) +} + func (s distributionRouteSnapshot) WriteFencedForKey(key []byte) bool { route, ok := s.snap.RouteOf(key) return ok && route.State == distribution.RouteStateWriteFenced diff --git a/kv/shard_key.go b/kv/shard_key.go index 1d2b609a8..156bf8b78 100644 --- a/kv/shard_key.go +++ b/kv/shard_key.go @@ -115,6 +115,12 @@ func RouteKey(key []byte) []byte { return routeKey(key) } +// RouteOwnershipKey normalizes a stored key to the catalog key used when +// answering route ownership queries. +func RouteOwnershipKey(key []byte) []byte { + return routeOwnershipKey(key) +} + func routeKey(key []byte) []byte { if key == nil { return nil @@ -125,6 +131,19 @@ func routeKey(key []byte) []byte { return normalizeRouteKey(key) } +func routeOwnershipKey(key []byte) []byte { + if embedded, ok := txnRouteKey(key); ok { + key = embedded + } + if bucket, ok := s3keys.ParseBucketMetaKey(key); ok { + return s3keys.RoutePrefixForBucketAnyGeneration(bucket) + } + if bucket, ok := s3keys.ParseBucketGenerationKey(key); ok { + return s3keys.RoutePrefixForBucketAnyGeneration(bucket) + } + return routeKey(key) +} + func routeFilterKey(key []byte) []byte { if key == nil { return nil diff --git a/kv/shard_key_test.go b/kv/shard_key_test.go index 05030bd46..13eb99bb7 100644 --- a/kv/shard_key_test.go +++ b/kv/shard_key_test.go @@ -2,6 +2,7 @@ package kv import ( "encoding/base64" + "encoding/binary" "testing" "github.com/bootjp/elastickv/internal/fskeys" @@ -31,6 +32,16 @@ func TestRouteKey_NormalizesTxnWrappedS3Key(t *testing.T) { require.Equal(t, s3keys.RouteKey("bucket-a", 7, "path/to/object"), routeKey(txnLockKey(embedded))) } +func TestRouteOwnershipKey_NormalizesS3BucketAuxiliaryKeys(t *testing.T) { + t.Parallel() + + want := s3keys.RoutePrefixForBucketAnyGeneration("bucket-a") + require.Equal(t, want, RouteOwnershipKey(s3keys.BucketMetaKey("bucket-a"))) + require.Equal(t, want, RouteOwnershipKey(s3keys.BucketGenerationKey("bucket-a"))) + require.Equal(t, want, RouteOwnershipKey(txnLockKey(s3keys.BucketMetaKey("bucket-a")))) + require.Equal(t, want, RouteOwnershipKey(txnLockKey(s3keys.BucketGenerationKey("bucket-a")))) +} + func TestRouteKey_NormalizesFilesystemChunkKey(t *testing.T) { t.Parallel() @@ -302,3 +313,14 @@ func TestStreamScansAreNotWideColumnCanonicalizable(t *testing.T) { require.True(t, redisWideColumnCanonicalizableScan(start)) } } + +func legacyListMetaDeltaKey(userKey []byte, commitTS uint64) []byte { + const seqInTxn = uint32(1) + key := store.LegacyListMetaDeltaScanPrefix(userKey) + var ts [8]byte + binary.BigEndian.PutUint64(ts[:], commitTS) + key = append(key, ts[:]...) + var seq [4]byte + binary.BigEndian.PutUint32(seq[:], seqInTxn) + return append(key, seq[:]...) +} diff --git a/kv/shard_router.go b/kv/shard_router.go index 4f6cd094c..c8b97a867 100644 --- a/kv/shard_router.go +++ b/kv/shard_router.go @@ -134,6 +134,9 @@ func (s *ShardRouter) ResolveGroup(rawKey []byte) (uint64, bool) { return 0, false } } + if route, ok := s.s3BucketAuxiliaryOwnerRouteForKey(rawKey); ok { + return route.GroupID, true + } // Engine routes against the user-key view of the byte-range // space; routeKey may rewrite SQS / DynamoDB / Redis-internal // keys to a stable per-table or per-namespace route key so the @@ -145,6 +148,17 @@ func (s *ShardRouter) ResolveGroup(rawKey []byte) (uint64, bool) { return route.GroupID, true } +func (s *ShardRouter) s3BucketAuxiliaryOwnerRouteForKey(rawKey []byte) (distribution.Route, bool) { + if s == nil || s.engine == nil { + return distribution.Route{}, false + } + start, end, ok := s3BucketAuxiliaryRouteRange(rawKey) + if !ok { + return distribution.Route{}, false + } + return s3BucketAuxiliaryOwnerRouteFromRange(start, end, s.engine.GetIntersectingRoutes(start, end)) +} + // Register associates a raft group ID with its transactional manager and store. func (s *ShardRouter) Register(group uint64, tm Transactional, st store.MVCCStore) { s.mu.Lock() diff --git a/kv/shard_store.go b/kv/shard_store.go index 8ae1f9fbf..af112d6b8 100644 --- a/kv/shard_store.go +++ b/kv/shard_store.go @@ -35,9 +35,11 @@ type ShardStore struct { } var ( - ErrCrossShardMutationBatchNotSupported = errors.New("cross-shard mutation batches are not supported") - ErrReadRouteVersionUnavailable = errors.New("read route version is not locally available") - ErrFilesystemPlacementTargetNotFound = errors.New("filesystem placement target group has no routable home slot") + ErrCrossShardMutationBatchNotSupported = errors.New("cross-shard mutation batches are not supported") + ErrExplicitGroupStagedVisibilityUnresolved = errors.New("explicit group read cannot resolve staged visibility route") + ErrExplicitGroupRouteOwnerMismatch = errors.New("explicit group read does not own the requested key range") + ErrReadRouteVersionUnavailable = errors.New("read route version is not locally available") + ErrFilesystemPlacementTargetNotFound = errors.New("filesystem placement target group has no routable home slot") ) // NewShardStore creates a sharded MVCC store wrapper. @@ -114,6 +116,40 @@ func (s *ShardStore) FilesystemGroupIDs() []uint64 { return groupIDs } +// LocalStoreForKey resolves the local store that owns a raw key in this +// process. S3 chunkblob backfill uses this for peer-local auxiliary rows, which +// bypass Raft but must still follow the same route catalog as the public paths. +func (s *ShardStore) LocalStoreForKey(key []byte) (store.MVCCStore, bool) { + _, group, _, ok := s.routeAndGroupForKeyWithVersion(key) + if !ok || group == nil || group.Store == nil { + return nil, false + } + return group.Store, true +} + +// LocalStores returns every local physical store once. It is intentionally not +// filtered by current route ownership because startup backfill must recover +// chunkrefs restored from snapshots or retained after a route move. +func (s *ShardStore) LocalStores() []store.MVCCStore { + if s == nil { + return nil + } + groupIDs := make([]uint64, 0, len(s.groups)) + for groupID := range s.groups { + groupIDs = append(groupIDs, groupID) + } + slices.Sort(groupIDs) + stores := make([]store.MVCCStore, 0, len(groupIDs)) + for _, groupID := range groupIDs { + group := s.groups[groupID] + if group == nil || group.Store == nil { + continue + } + stores = append(stores, group.Store) + } + return stores +} + // ResolveFilesystemHomeSlot finds a home token whose file route belongs to // targetGroup. It derives candidates from current route boundaries and verifies // each candidate against the live catalog before returning it. @@ -231,8 +267,12 @@ func (s *ShardStore) getGroupAtWithReadFence(ctx context.Context, groupID uint64 if !ok || g.Store == nil { return nil, store.ErrKeyNotFound } + route, err := s.routeForExplicitGroupKey(groupID, key) + if err != nil { + return nil, err + } - return s.getGroupAt(ctx, g, key, ts, groupID, readRouteVersion) + return s.getGroupAt(ctx, g, route, key, ts, groupID, readRouteVersion) } func (s *ShardStore) getRouteAt(ctx context.Context, route distribution.Route, key []byte, ts uint64, readRouteVersion uint64) ([]byte, error) { @@ -240,15 +280,24 @@ func (s *ShardStore) getRouteAt(ctx context.Context, route distribution.Route, k if !ok || g.Store == nil { return nil, store.ErrKeyNotFound } - return s.getGroupAt(ctx, g, key, ts, route.GroupID, readRouteVersion) + return s.getGroupAt(ctx, g, route, key, ts, route.GroupID, readRouteVersion) +} + +func (s *ShardStore) getGroupAt(ctx context.Context, g *ShardGroup, route distribution.Route, key []byte, ts uint64, groupID uint64, readRouteVersion uint64) ([]byte, error) { + if engineForGroup(g) == nil { + return s.localGetAt(ctx, g, route, key, ts) + } + if isLinearizableRaftLeader(ctx, engineForGroup(g)) { + return s.leaderGetAt(ctx, g, route, key, ts) + } + return s.proxyRawGet(ctx, g, key, ts, groupID, readRouteVersion) } // routeHasLatestVersionVisibleAt reports whether this route can answer the read -// itself, which is what stops the point-read fallback from trying the legacy -// wide-column route. If the exact local/remote presence probe is unavailable, -// stop the fallback anyway: the old latest-commit heuristic cannot distinguish -// "no version" from "a tombstone at or before ts plus a newer future version", -// which can resurrect legacy values during rolling upgrades. +// itself, which is what stops the point-read fallback from trying a legacy +// route. If the exact local/remote presence probe is unavailable, stop the +// fallback anyway: the old latest-commit heuristic cannot distinguish "no +// version" from "a tombstone at or before ts plus a newer future version". func (s *ShardStore) routeHasLatestVersionVisibleAt(ctx context.Context, route distribution.Route, key []byte, ts uint64, readRouteVersion uint64) (bool, error) { if exists, ok, err := s.routeHasVersionAtOrBefore(ctx, route, key, ts); ok || err != nil { return exists, err @@ -259,11 +308,6 @@ func (s *ShardStore) routeHasLatestVersionVisibleAt(ctx context.Context, route d return true, nil } -// routeHasVersionAtOrBeforeRemote asks the route's leader the exact question a -// non-leader replica cannot answer locally. The second bool reports whether the -// peer answered at all: one that predates version_visible_at_ts leaves the -// response fields unset, and the caller stops legacy fallback rather than -// silently treating "unknown" as "no version". func (s *ShardStore) routeHasVersionAtOrBeforeRemote(ctx context.Context, route distribution.Route, key []byte, ts uint64, readRouteVersion uint64) (bool, bool, error) { if ts == 0 { return false, false, nil @@ -343,32 +387,48 @@ func (s *ShardStore) routeMissStopsPointFallback(ctx context.Context, routes []d return s.routeHasLatestVersionVisibleAt(ctx, route, key, ts, readRouteVersion) } -func (s *ShardStore) getGroupAt(ctx context.Context, g *ShardGroup, key []byte, ts uint64, groupID uint64, readRouteVersion uint64) ([]byte, error) { - if engineForGroup(g) == nil { - return s.localGetAt(ctx, g, key, ts) - } - if isLinearizableRaftLeader(ctx, engineForGroup(g)) { - return s.leaderGetAt(ctx, g, key, ts) - } - return s.proxyRawGet(ctx, g, key, ts, groupID, readRouteVersion) -} +const pointReadRouteCandidateCapacity = 2 func (s *ShardStore) pointReadRoutesWithVersion(key []byte) ([]distribution.Route, uint64) { + if routes, version, ok := s.s3BucketAuxiliaryPointReadRoutesWithVersion(key); ok { + return routes, version + } primaryKey := routeKey(key) legacyKey := legacyPointRouteKey(key) if legacyKey == nil || bytes.Equal(legacyKey, primaryKey) { - primary, version, ok := s.engine.GetRouteWithVersion(primaryKey) - if !ok { - return nil, version + return s.singlePointReadRouteWithVersion(primaryKey) + } + return s.primaryAndLegacyPointReadRoutesWithVersion(primaryKey, legacyKey) +} + +func (s *ShardStore) s3BucketAuxiliaryPointReadRoutesWithVersion(key []byte) ([]distribution.Route, uint64, bool) { + start, end, auxiliary := s3BucketAuxiliaryRouteRange(key) + if !auxiliary { + return nil, 0, false + } + catalogRoutes, version := s.engine.GetIntersectingRoutesWithVersion(nil, nil) + routes := make([]distribution.Route, 0, pointReadRouteCandidateCapacity) + if route, ok := s3BucketAuxiliaryOwnerRouteFromRange(start, end, catalogRoutes); ok { + routes = append(routes, route) + } + normalizedKey := routeKey(key) + for _, route := range catalogRoutes { + if routeContainsKey(route, normalizedKey) { + routes = append(routes, route) } - return []distribution.Route{primary}, version } + return appendDistinctRoutesByGroup(nil, routes), version, true +} + +func (s *ShardStore) singlePointReadRouteWithVersion(key []byte) ([]distribution.Route, uint64) { + primary, version, ok := s.engine.GetRouteWithVersion(key) + if !ok { + return nil, version + } + return []distribution.Route{primary}, version +} - // Both candidates must come from one catalog snapshot. Resolving them with - // two GetRouteWithVersion calls let the primary route move in between: the - // request then carried the newer read_route_version while still naming the - // stale primary group, so the leader's fence passed after waiting for a - // version the route it actually read was never valid under. +func (s *ShardStore) primaryAndLegacyPointReadRoutesWithVersion(primaryKey []byte, legacyKey []byte) ([]distribution.Route, uint64) { resolved, version := s.engine.ResolveRoutesWithVersion( distribution.RouteQuery{Start: primaryKey, Exact: true}, distribution.RouteQuery{Start: legacyKey, Exact: true}, @@ -401,16 +461,19 @@ func isLinearizableRaftLeader(ctx context.Context, engine raftengine.LeaderView) return err == nil } -func (s *ShardStore) leaderGetAt(ctx context.Context, g *ShardGroup, key []byte, ts uint64) ([]byte, error) { +func (s *ShardStore) leaderGetAt(ctx context.Context, g *ShardGroup, route distribution.Route, key []byte, ts uint64) ([]byte, error) { if !isTxnInternalKey(key) { if err := s.maybeResolveTxnLock(ctx, g, key, ts); err != nil { return nil, err } } - return s.localGetAt(ctx, g, key, ts) + return s.localGetAt(ctx, g, route, key, ts) } -func (s *ShardStore) localGetAt(ctx context.Context, g *ShardGroup, key []byte, ts uint64) ([]byte, error) { +func (s *ShardStore) localGetAt(ctx context.Context, g *ShardGroup, route distribution.Route, key []byte, ts uint64) ([]byte, error) { + if routeHasStagedVisibility(route) { + return s.getAtWithStagedVisibility(ctx, g, route, key, ts) + } val, err := g.Store.GetAt(ctx, key, ts) if err != nil { return nil, errors.WithStack(err) @@ -418,6 +481,179 @@ func (s *ShardStore) localGetAt(ctx context.Context, g *ShardGroup, key []byte, return val, nil } +func routeHasStagedVisibility(route distribution.Route) bool { + return route.StagedVisibilityActive && route.MigrationJobID != 0 +} + +func (s *ShardStore) routeForExplicitGroupKey(groupID uint64, key []byte) (distribution.Route, error) { + fallback := distribution.Route{GroupID: groupID} + if s == nil || s.engine == nil { + return fallback, nil + } + if route, ok := s.s3BucketAuxiliaryOwnerRouteForKey(key); ok { + if route.GroupID == groupID { + return route, nil + } + if routeHasStagedVisibility(route) { + return distribution.Route{}, errors.Wrapf(ErrExplicitGroupStagedVisibilityUnresolved, "group_id=%d key=%q", groupID, key) + } + return distribution.Route{}, errors.Wrapf( + ErrExplicitGroupRouteOwnerMismatch, + "group_id=%d catalog_group_id=%d key=%q", groupID, route.GroupID, key) + } + if route, ok := s.engine.GetRoute(routeKey(key)); ok { + if route.GroupID == groupID { + return route, nil + } + if routeHasStagedVisibility(route) { + return distribution.Route{}, errors.Wrapf(ErrExplicitGroupStagedVisibilityUnresolved, "group_id=%d key=%q", groupID, key) + } + if !explicitGroupResolverOwnedKey(key) { + return distribution.Route{}, errors.Wrapf( + ErrExplicitGroupRouteOwnerMismatch, + "group_id=%d catalog_group_id=%d key=%q", groupID, route.GroupID, key) + } + } + return fallback, nil +} + +// explicitGroupResolverOwnedKey reports whether key belongs to a keyspace whose +// owning group is chosen by a resolver rather than by the byte-range catalog -- +// SQS HT-FIFO's (queue, partition) resolver being the one such keyspace today. +// Only those keys may be served from a group the catalog does not name. +// +// Everything else must fail closed on a mismatch. Once a migration's promotion +// clears StagedVisibilityActive, the source's former range belongs to the +// target, and a coordinator that has not yet applied that catalog version keeps +// forwarding the source group. Falling back to the requested group then serves +// the source's own MVCC -- the value from before the cutover -- for as long as +// source cleanup takes. The staged-visibility rejection above stops covering +// that window the moment the flag is cleared, which is exactly when the +// pre-cutover data is still sitting there. +// fsChunkAllPrefix is the raw filesystem chunk keyspace prefix, hoisted so the +// explicit-group gate does not allocate it per read. +var fsChunkAllPrefix = fskeys.ChunkAllPrefix() + +func explicitGroupResolverOwnedKey(key []byte) bool { + if sqsRouteKey(key) != nil { + return true + } + // Filesystem chunks are placed per group by the filesystem's own home-slot + // placement, not by the byte-range catalog: ListFilePlacementStats scans the + // whole chunk keyspace once per FilesystemGroupIDs entry + // (internal/filesystem/placement.go scanPlacementChunks), so most of those + // groups are not the catalog owner of the range and never will be. + return bytes.HasPrefix(key, fsChunkAllPrefix) +} + +func (s *ShardStore) getAtWithStagedVisibility(ctx context.Context, g *ShardGroup, route distribution.Route, key []byte, ts uint64) ([]byte, error) { + if err := ensureReadTSRetained(g.Store, ts); err != nil { + return nil, err + } + // Staged first, then live. The two probes are separate store reads, so a + // promotion batch can land between them, and the order decides whether that + // is observable. Promotion only ever moves a row staged -> live: the live + // set only gains keys, the staged set only loses them. Reading the shrinking + // side first and the growing side second means a key present in either at + // any instant during the pair is seen by at least one probe. + // + // The reverse order has two holes. A staged-only key read live-then-staged + // misses both sides and returns not-found for a key that existed + // throughout; and a key whose staged version is newer than its live one can + // return the stale live value when promotion removes the alias in between. + stagedKey := distribution.MigrationStagedDataKey(route.MigrationJobID, key) + staged, stagedOK, err := latestMVCCVersionAt(ctx, g.Store, stagedKey, ts) + if err != nil { + return nil, err + } + live, liveOK, err := latestMVCCVersionAt(ctx, g.Store, key, ts) + if err != nil { + return nil, err + } + if stagedOK { + staged.Key = bytes.Clone(key) + } + winner, ok := newerMigrationVersion(live, liveOK, staged, stagedOK) + if !ok || !migrationVersionVisible(winner, ts) { + return nil, store.ErrKeyNotFound + } + return bytes.Clone(winner.Value), nil +} + +// exactKeyScanEnd returns the exclusive upper bound that selects only key. +// +// prefixScanEnd(key) selects every key that has key as a prefix, which turns an +// exact-key probe into a prefix scan: probing an absent "a" walks "ab", "az" and +// all of their versions before concluding "a" is not there. A byte-string has no +// value strictly between key and key+0x00, so this bound covers key alone. +func exactKeyScanEnd(key []byte) []byte { + // The immediate successor of key in byte-string order. + const successorByte = byte(0) + out := make([]byte, 0, len(key)+1) + out = append(out, key...) + return append(out, successorByte) +} + +func latestMVCCVersionAt(ctx context.Context, st store.MVCCStore, key []byte, ts uint64) (store.MVCCVersion, bool, error) { + opts := store.ExportVersionsOptions{ + StartKey: key, + EndKey: exactKeyScanEnd(key), + MaxCommitTSInclusive: ts, + MaxVersions: 1, + MaxScannedBytes: 0, + MinCommitTSExclusive: 0, + MaxBytes: 0, + KeyFamily: 0, + AcceptKey: func(rawKey []byte) bool { + return bytes.Equal(rawKey, key) + }, + } + for { + result, err := st.ExportVersions(ctx, opts) + if err != nil { + return store.MVCCVersion{}, false, errors.WithStack(err) + } + for _, version := range result.Versions { + if bytes.Equal(version.Key, key) { + return version, true, nil + } + } + if result.Done || len(result.NextCursor) == 0 { + return store.MVCCVersion{}, false, nil + } + opts.Cursor = result.NextCursor + } +} + +func newerMigrationVersion(a store.MVCCVersion, aOK bool, b store.MVCCVersion, bOK bool) (store.MVCCVersion, bool) { + switch { + case !aOK: + return b, bOK + case !bOK: + return a, true + case b.CommitTS >= a.CommitTS: + return b, true + default: + return a, true + } +} + +func migrationVersionVisible(version store.MVCCVersion, ts uint64) bool { + return !version.Tombstone && (version.ExpireAt == 0 || version.ExpireAt > ts) +} + +func ensureReadTSRetained(st store.MVCCStore, ts uint64) error { + retention, ok := st.(store.RetentionController) + if !ok { + return nil + } + minRetainedTS := retention.MinRetainedTS() + if minRetainedTS != 0 && ts != 0 && ts != ^uint64(0) && ts < minRetainedTS { + return errors.WithStack(store.ErrReadTSCompacted) + } + return nil +} + func (s *ShardStore) ExistsAt(ctx context.Context, key []byte, ts uint64) (bool, error) { v, err := s.GetAt(ctx, key, ts) if err != nil { @@ -510,11 +746,7 @@ func (s *ShardStore) ScanAtWithReadFence(ctx context.Context, start []byte, end } if reverse { if groupID != 0 { - // Unbounded group-scoped reverse scans are served here too, not - // short-circuited by the caller: routing them through the fenced - // path is what applies awaitReadRouteVersion above and the - // migration read fence to an explicitly selected group. - return s.scanRouteAtDirectionWithReadFence(ctx, distribution.Route{GroupID: groupID}, start, end, limit, ts, true, readRouteVersion, routeStart, routeEnd) + return s.scanExplicitGroupAtWithReadFence(ctx, groupID, start, end, limit, ts, true, readRouteVersion, routeStart, routeEnd) } return s.reverseScanAtWithReadFence(ctx, start, end, limit, ts, readRouteVersion, routeStart, routeEnd) } @@ -527,7 +759,7 @@ func (s *ShardStore) scanAtWithReadFence(ctx context.Context, start []byte, end } if groupID != 0 { - return s.scanRouteAtDirectionWithReadFence(ctx, distribution.Route{GroupID: groupID}, start, end, limit, ts, false, readRouteVersion, routeStart, routeEnd) + return s.scanExplicitGroupAtWithReadFence(ctx, groupID, start, end, limit, ts, false, readRouteVersion, routeStart, routeEnd) } routes, clampToRoutes, routeVersion := s.routesForFencedScanWithVersion(start, end, routeStart, routeEnd) @@ -540,11 +772,6 @@ func (s *ShardStore) scanAtWithReadFence(ctx context.Context, start []byte, end return bytes.Compare(out[i].Key, out[j].Key) < 0 }) out = dedupeSortedScanResults(out) - // Not canonicalized again here. Every route page this merged already went - // through canonicalizeRedisWideColumnScanResults in its local, leader, or - // proxy page path, and the canonical row keeps the physical key, so a second - // pass repeats the point read for every surviving row -- 2N reads for a page - // of N, and on a remote group that is a fenced RPC each. if len(out) > limit { out = out[:limit] } @@ -570,7 +797,7 @@ func (s *ShardStore) ScanKeysAtWithReadFence(ctx context.Context, start []byte, return keysFromKVs(kvs), nil } if groupID != 0 { - return s.scanKeyRouteAtWithReadFence(ctx, distribution.Route{GroupID: groupID}, start, end, limit, ts, readRouteVersion) + return s.scanExplicitGroupKeysAtWithReadFence(ctx, groupID, start, end, limit, ts, readRouteVersion, nil, nil) } routes, clampToRoutes, routeVersion := s.routesForScanWithVersion(start, end) @@ -588,11 +815,58 @@ func (s *ShardStore) ScanKeysAtWithReadFence(ctx context.Context, start []byte, return out, nil } +func (s *ShardStore) canonicalizeRedisWideColumnScanResults(ctx context.Context, kvs []*store.KVPair, start []byte, ts uint64, readRouteVersion uint64) ([]*store.KVPair, error) { + if len(kvs) == 0 || !redisWideColumnCanonicalizableScan(start) { + return kvs, nil + } + out := make([]*store.KVPair, 0, len(kvs)) + for _, kvp := range kvs { + if kvp == nil || redisWideColumnLegacyPointRouteKey(kvp.Key) == nil { + out = append(out, kvp) + continue + } + value, err := s.GetAtWithReadFence(ctx, kvp.Key, ts, 0, readRouteVersion) + if err != nil { + if errors.Is(err, store.ErrKeyNotFound) { + continue + } + return nil, err + } + out = append(out, &store.KVPair{Key: bytes.Clone(kvp.Key), Value: value}) + } + return out, nil +} + +func (s *ShardStore) canonicalizeRedisWideColumnScanKeys(ctx context.Context, keys [][]byte, start []byte, ts uint64, readRouteVersion uint64) ([][]byte, error) { + if len(keys) == 0 || !redisWideColumnCanonicalizableScan(start) { + return keys, nil + } + out := make([][]byte, 0, len(keys)) + for _, key := range keys { + if key == nil || redisWideColumnLegacyPointRouteKey(key) == nil { + out = append(out, key) + continue + } + if _, err := s.GetAtWithReadFence(ctx, key, ts, 0, readRouteVersion); err != nil { + if errors.Is(err, store.ErrKeyNotFound) { + continue + } + return nil, err + } + out = append(out, bytes.Clone(key)) + } + return out, nil +} + func (s *ShardStore) ScanAtPhysicalLimit(ctx context.Context, start []byte, end []byte, visibleLimit, physicalLimit int, ts uint64) ([]*store.KVPair, bool, error) { if visibleLimit <= 0 || physicalLimit <= 0 { return []*store.KVPair{}, false, nil } routes, clampToRoutes := s.routesForForwardScan(start, end) + if routesContainStagedVisibility(routes) { + kvs, err := s.ScanAt(ctx, start, end, visibleLimit, ts) + return kvs, false, err + } if len(routes) != 1 || clampToRoutes { kvs, err := s.ScanAt(ctx, start, end, visibleLimit, ts) return kvs, false, err @@ -606,18 +880,69 @@ func (s *ShardStore) ScanAtPhysicalLimit(ctx context.Context, start []byte, end // Normal callers should use ScanAt so range scans keep following the // distribution engine's route table. func (s *ShardStore) ScanGroupAt(ctx context.Context, groupID uint64, start []byte, end []byte, limit int, ts uint64) ([]*store.KVPair, error) { + return s.scanExplicitGroupAtWithReadFence(ctx, groupID, start, end, limit, ts, false, 0, nil, nil) +} + +func (s *ShardStore) scanExplicitGroupAtWithReadFence(ctx context.Context, groupID uint64, start []byte, end []byte, limit int, ts uint64, reverse bool, readRouteVersion uint64, routeStart []byte, routeEnd []byte) ([]*store.KVPair, error) { if limit <= 0 { return []*store.KVPair{}, nil } - return s.scanRouteAtDirection(ctx, distribution.Route{GroupID: groupID}, start, end, limit, ts, false) + routes, clampToRoutes, err := s.routesForExplicitGroupScanWithRouteBounds(groupID, start, end, routeStart, routeEnd) + if err != nil { + return nil, err + } + routeFilterPresent := routeScanBoundsPresent(routeStart, routeEnd) + dedupeByKey := s3BucketAuxiliaryScanBounds(start, end) + if !clampToRoutes && !routeFilterPresent { + routes, dedupeByKey = prepareUnclampedRawScanRoutes(routes, dedupeByKey) + } + ownerFilterRoutes := s.s3BucketAuxiliaryOwnerFilterRoutes(start, end, routes, dedupeByKey) + return s.scanExplicitGroupRoutesAtWithReadFence(ctx, routes, ownerFilterRoutes, start, end, limit, ts, reverse, readRouteVersion, routeStart, routeEnd, clampToRoutes, dedupeByKey) +} + +func (s *ShardStore) s3BucketAuxiliaryOwnerFilterRoutes(start []byte, end []byte, routes []distribution.Route, dedupeByKey bool) []distribution.Route { + if !dedupeByKey { + return routes + } + if candidates, _, ok := s.routesForS3BucketAuxiliaryScan(start, end); ok && len(candidates) > 0 { + return candidates + } + return routes +} + +func (s *ShardStore) scanExplicitGroupRoutesAtWithReadFence(ctx context.Context, routes []distribution.Route, ownerFilterRoutes []distribution.Route, start []byte, end []byte, limit int, ts uint64, reverse bool, readRouteVersion uint64, routeStart []byte, routeEnd []byte, clampToRoutes bool, dedupeByKey bool) ([]*store.KVPair, error) { + out := make([]*store.KVPair, 0) + for i := 0; i < len(routes); i++ { + route := routes[i] + if reverse && clampToRoutes { + route = routes[len(routes)-1-i] + } + scanStart := start + scanEnd := end + if clampToRoutes { + scanStart = clampScanStart(start, route.Start) + scanEnd = clampScanEnd(end, route.End) + } + kvs, err := s.scanRouteAtDirectionWithS3AuxiliaryOwnerFilter(ctx, ownerFilterRoutes, route, scanStart, scanEnd, limit, ts, reverse, true, readRouteVersion, routeStart, routeEnd, dedupeByKey) + if err != nil { + return nil, err + } + if clampToRoutes { + out = append(out, kvs...) + if len(out) >= limit { + clear(out[limit:]) + return out[:limit], nil + } + continue + } + out = mergeAndTrimScanResultsWithOptions(out, kvs, limit, reverse, dedupeByKey) + } + return out, nil } // ReverseScanGroupAt reverse-scans a range on the explicitly selected Raft group. func (s *ShardStore) ReverseScanGroupAt(ctx context.Context, groupID uint64, start []byte, end []byte, limit int, ts uint64) ([]*store.KVPair, error) { - if limit <= 0 { - return []*store.KVPair{}, nil - } - return s.scanRouteAtDirection(ctx, distribution.Route{GroupID: groupID}, start, end, limit, ts, true) + return s.scanExplicitGroupAtWithReadFence(ctx, groupID, start, end, limit, ts, true, 0, nil, nil) } // ScanGroupKeysAt scans keys on the explicitly selected Raft group without @@ -644,67 +969,21 @@ func (s *ShardStore) reverseScanAtWithReadFence(ctx context.Context, start []byt if err != nil { return nil, err } - // Already canonicalized per page; see scanAtWithReadFence. if len(out) > limit { out = out[:limit] } return out, nil } -func (s *ShardStore) canonicalizeRedisWideColumnScanResults(ctx context.Context, kvs []*store.KVPair, start []byte, ts uint64, readRouteVersion uint64) ([]*store.KVPair, error) { - if len(kvs) == 0 { - return kvs, nil - } - if !redisWideColumnCanonicalizableScan(start) { - return kvs, nil - } - out := make([]*store.KVPair, 0, len(kvs)) - for _, kvp := range kvs { - if kvp == nil || redisWideColumnLegacyPointRouteKey(kvp.Key) == nil { - out = append(out, kvp) - continue - } - value, err := s.GetAtWithReadFence(ctx, kvp.Key, ts, 0, readRouteVersion) - if err != nil { - if errors.Is(err, store.ErrKeyNotFound) { - continue - } - return nil, err - } - out = append(out, &store.KVPair{Key: bytes.Clone(kvp.Key), Value: value}) - } - return out, nil -} - -func (s *ShardStore) canonicalizeRedisWideColumnScanKeys(ctx context.Context, keys [][]byte, start []byte, ts uint64, readRouteVersion uint64) ([][]byte, error) { - if len(keys) == 0 { - return keys, nil - } - if !redisWideColumnCanonicalizableScan(start) { - return keys, nil - } - out := make([][]byte, 0, len(keys)) - for _, key := range keys { - if key == nil || redisWideColumnLegacyPointRouteKey(key) == nil { - out = append(out, key) - continue - } - if _, err := s.GetAtWithReadFence(ctx, key, ts, 0, readRouteVersion); err != nil { - if errors.Is(err, store.ErrKeyNotFound) { - continue - } - return nil, err - } - out = append(out, bytes.Clone(key)) - } - return out, nil -} - func (s *ShardStore) ReverseScanAtPhysicalLimit(ctx context.Context, start []byte, end []byte, visibleLimit, physicalLimit int, ts uint64) ([]*store.KVPair, bool, error) { if visibleLimit <= 0 || physicalLimit <= 0 { return []*store.KVPair{}, false, nil } routes, clampToRoutes := s.routesForReverseScan(start, end) + if routesContainStagedVisibility(routes) { + kvs, err := s.ReverseScanAt(ctx, start, end, visibleLimit, ts) + return kvs, false, err + } if len(routes) != 1 || clampToRoutes { kvs, err := s.ReverseScanAt(ctx, start, end, visibleLimit, ts) return kvs, false, err @@ -712,38 +991,8 @@ func (s *ShardStore) ReverseScanAtPhysicalLimit(ctx context.Context, start []byt return s.scanRouteAtDirectionPhysicalLimit(ctx, routes[0], start, end, visibleLimit, physicalLimit, ts, true) } -func (s *ShardStore) AllowExactScanFallbackAfterPhysicalLimit(ctx context.Context, start []byte, end []byte, visibleLimit, physicalLimit int, _ uint64, _ bool) bool { - if visibleLimit <= 0 || physicalLimit <= 0 { - return false - } - g := s.exactFallbackPhysicalLimitGroup(start, end) - if g == nil { - return false - } - if _, ok := g.Store.(physicalLimitedStore); !ok { - return false - } - engine := engineForGroup(g) - return engine == nil || isLinearizableRaftLeader(ctx, engine) -} - -func (s *ShardStore) exactFallbackPhysicalLimitGroup(start []byte, end []byte) *ShardGroup { - if s == nil || s.engine == nil { - return nil - } - routes, clampToRoutes := s.routesForScan(start, end) - if len(routes) != 1 || clampToRoutes { - return nil - } - g, ok := s.groupForID(routes[0].GroupID) - if !ok || g == nil || g.Store == nil { - return nil - } - return g -} - -func (s *ShardStore) routesForForwardScan(start []byte, end []byte) ([]distribution.Route, bool) { - return s.routesForScan(start, end) +func (s *ShardStore) routesForForwardScan(start []byte, end []byte) ([]distribution.Route, bool) { + return s.routesForScan(start, end) } func (s *ShardStore) routesForReverseScan(start []byte, end []byte) ([]distribution.Route, bool) { @@ -829,6 +1078,36 @@ func scanReadFenceRouteKey(route distribution.Route, start []byte, clampToRoutes return bytes.Clone(start) } +func (s *ShardStore) AllowExactScanFallbackAfterPhysicalLimit(ctx context.Context, start []byte, end []byte, visibleLimit, physicalLimit int, _ uint64, _ bool) bool { + if visibleLimit <= 0 || physicalLimit <= 0 { + return false + } + g := s.exactFallbackPhysicalLimitGroup(start, end) + if g == nil { + return false + } + if _, ok := g.Store.(physicalLimitedStore); !ok { + return false + } + engine := engineForGroup(g) + return engine == nil || isLinearizableRaftLeader(ctx, engine) +} + +func (s *ShardStore) exactFallbackPhysicalLimitGroup(start []byte, end []byte) *ShardGroup { + if s == nil || s.engine == nil { + return nil + } + routes, clampToRoutes := s.routesForScan(start, end) + if len(routes) != 1 || clampToRoutes { + return nil + } + g, ok := s.groupForID(routes[0].GroupID) + if !ok || g == nil || g.Store == nil { + return nil + } + return g +} + func (s *ShardStore) routesForScan(start []byte, end []byte) ([]distribution.Route, bool) { routes, clampToRoutes, _ := s.routesForScanWithVersion(start, end) return routes, clampToRoutes @@ -839,13 +1118,21 @@ func (s *ShardStore) routesForScanWithVersion(start []byte, end []byte) ([]distr routes, version := s.engine.GetIntersectingRoutesWithVersion(routeStart, routeEnd) return routes, false, version } - if routes, version, ok := s.routesForEncodedScanWithVersion(start, end); ok { + if routes, version, ok := s.routesForFilesystemUsageScanWithVersion(start, end); ok { + return routes, false, version + } + if routes, version, ok := s.routesForS3BucketAuxiliaryScan(start, end); ok { return routes, false, version } if routes, version, ok := s.routesForRedisWideColumnScanWithVersion(start, end); ok { return routes, false, version } - + if routes, version, ok := s.routesForFilesystemChunkScanWithVersion(start, end); ok { + return routes, false, version + } + if selected, ok := s.routesForInternalScanWithVersion(start, end); ok { + return selected.routes, false, selected.version + } routes, version := s.engine.GetIntersectingRoutesWithVersion(start, end) // If the scan can include internal list keys (which use a fixed prefix), // avoid clamping to shard range bounds because those keys may be ordered @@ -857,14 +1144,47 @@ func (s *ShardStore) routesForScanWithVersion(start []byte, end []byte) ([]distr return routes, true, version } +type internalScanRouteSelection struct { + routes []distribution.Route + version uint64 +} + +func (s *ShardStore) routesForInternalScanWithVersion(start []byte, end []byte) (internalScanRouteSelection, bool) { + if isBroadLegacyListDeltaScan(start) { + routes, version := s.engine.GetIntersectingRoutesWithVersion(nil, nil) + return internalScanRouteSelection{routes: routes, version: version}, true + } + if store.ExtractLegacyListUserKeyFromDeltaScanPrefix(start) != nil { + catalogRoutes, version := s.engine.GetIntersectingRoutesWithVersion(nil, nil) + return internalScanRouteSelection{routes: routesForLegacyListDeltaScan(catalogRoutes, start, end), version: version}, true + } + if routeStart, routeEnd, exact, ok := redisWideColumnScanRouteRange(start, end); ok { + if !exact { + routes, version := s.engine.GetIntersectingRoutesWithVersion(routeStart, routeEnd) + return internalScanRouteSelection{routes: routes, version: version}, true + } + route, version, ok := s.engine.GetRouteWithVersion(routeStart) + if !ok { + return internalScanRouteSelection{routes: []distribution.Route{}, version: version}, true + } + return internalScanRouteSelection{routes: []distribution.Route{route}, version: version}, true + } + // Remaining internal collection scans route by their logical user key. + if userKey := scanRouteUserKey(start); userKey != nil { + route, version, ok := s.engine.GetRouteWithVersion(userKey) + if !ok { + return internalScanRouteSelection{routes: []distribution.Route{}, version: version}, true + } + return internalScanRouteSelection{routes: []distribution.Route{route}, version: version}, true + } + return internalScanRouteSelection{}, false +} + func (s *ShardStore) routesForRedisWideColumnScanWithVersion(start []byte, end []byte) ([]distribution.Route, uint64, bool) { routeStart, routeEnd, exact, ok := redisWideColumnScanRouteRange(start, end) if !ok { return nil, 0, false } - // Same single-snapshot requirement as pointReadRoutesWithVersion: the - // primary and legacy candidates are resolved together so a catalog update - // between them cannot pair a stale route with a newer fence version. queries := []distribution.RouteQuery{{Start: routeStart, End: routeEnd, Exact: exact}} legacyStart, legacyEnd, hasLegacy := redisWideColumnLegacyScanRouteRange(start, end) if hasLegacy { @@ -896,41 +1216,199 @@ func appendDistinctRoutesByGroup(routes []distribution.Route, candidates []distr return routes } -func (s *ShardStore) routesForEncodedScanWithVersion(start []byte, end []byte) ([]distribution.Route, uint64, bool) { - if routes, version, ok := s.routesForFilesystemUsageScanWithVersion(start, end); ok { - return routes, version, true +func (s *ShardStore) routesForS3BucketAuxiliaryScan(start []byte, end []byte) ([]distribution.Route, uint64, bool) { + if s == nil || s.engine == nil || !s3BucketAuxiliaryScanBounds(start, end) { + return nil, 0, false } - if routes, version, ok := s.routesForFilesystemChunkScanWithVersion(start, end); ok { - return routes, version, true + catalogRoutes, version := s.engine.GetIntersectingRoutesWithVersion(nil, nil) + routes := make([]distribution.Route, 0) + for _, route := range catalogRoutes { + if migrationRouteRangesIntersect(route.Start, route.End, start, end) { + routes = append(routes, route) + } } - if isBroadLegacyListDeltaScan(start) { - routes, version := s.engine.GetIntersectingRoutesWithVersion(nil, nil) - return routes, version, true + routeStart, routeEnd := s3BucketAuxiliaryScanRouteRange(start, end) + for _, route := range catalogRoutes { + if migrationRouteRangesIntersect(route.Start, route.End, routeStart, routeEnd) { + routes = append(routes, route) + } } - if store.ExtractLegacyListUserKeyFromDeltaScanPrefix(start) != nil { - catalogRoutes, version := s.engine.GetIntersectingRoutesWithVersion(nil, nil) - return routesForLegacyListDeltaScan(catalogRoutes, start, end), version, true + return routes, version, true +} + +func s3BucketAuxiliaryScanBounds(start []byte, end []byte) bool { + if !bytes.HasPrefix(start, []byte(s3keys.BucketMetaPrefix)) && + !bytes.HasPrefix(start, []byte(s3keys.BucketGenerationPrefix)) { + return false } - if routeStart, exact, ok := listAuxiliaryScanRouteRange(start, end); ok { - if !exact { - routes, version := s.engine.GetIntersectingRoutesWithVersion(routeStart, nil) - return routes, version, true + if end == nil { + return true + } + return bytes.Compare(start, end) < 0 +} + +func s3BucketAuxiliaryScanRouteRange(start []byte, end []byte) ([]byte, []byte) { + if routeStart, routeEnd, ok := s3BucketAuxiliaryRouteRange(start); ok && end != nil && bytes.Compare(end, prefixScanEnd(start)) <= 0 { + return routeStart, routeEnd + } + routeStart := []byte(s3keys.RoutePrefix) + return routeStart, prefixScanEnd(routeStart) +} + +type repeatedRawScanRouteKey struct { + groupID uint64 + staged bool + migrationJobID uint64 + routeStart string + routeEnd string +} + +func dedupeRepeatedRawScanRoutes(routes []distribution.Route) []distribution.Route { + if len(routes) <= 1 { + return routes + } + out := make([]distribution.Route, 0, len(routes)) + seen := make(map[repeatedRawScanRouteKey]struct{}, len(routes)) + for _, route := range routes { + key := repeatedRawScanRouteDedupeKey(route) + if _, ok := seen[key]; ok { + continue } - route, version, ok := s.engine.GetRouteWithVersion(routeStart) - if !ok { - return []distribution.Route{}, version, true + seen[key] = struct{}{} + out = append(out, route) + } + return out +} + +func repeatedRawScanRouteDedupeKey(route distribution.Route) repeatedRawScanRouteKey { + key := repeatedRawScanRouteKey{groupID: route.GroupID} + if routeHasStagedVisibility(route) { + key.staged = true + key.migrationJobID = route.MigrationJobID + key.routeStart = string(route.Start) + key.routeEnd = string(route.End) + } + return key +} + +func prepareUnclampedRawScanRoutes(routes []distribution.Route, dedupeByKey bool) ([]distribution.Route, bool) { + if routesContainStagedVisibility(routes) { + routes = dedupeRepeatedRawScanRoutes(routes) + return orderRawScanRoutesForStagedVisibility(routes), true + } + if !dedupeByKey { + routes = dedupeRepeatedRawScanRoutes(routes) + } + return routes, dedupeByKey +} + +func orderRawScanRoutesForStagedVisibility(routes []distribution.Route) []distribution.Route { + if !routesContainStagedVisibility(routes) { + return routes + } + out := make([]distribution.Route, 0, len(routes)) + staged := make([]distribution.Route, 0) + for _, route := range routes { + if routeHasStagedVisibility(route) { + staged = append(staged, route) + continue } - return []distribution.Route{route}, version, true + out = append(out, route) } - userKey := listScanUserKey(start) - if userKey == nil { - return nil, 0, false + return append(out, staged...) +} + +func routesContainStagedVisibility(routes []distribution.Route) bool { + for _, route := range routes { + if routeHasStagedVisibility(route) { + return true + } } - route, version, ok := s.engine.GetRouteWithVersion(userKey) - if !ok { - return []distribution.Route{}, version, true + return false +} + +func (s *ShardStore) routesForExplicitGroupScanWithRouteBounds(groupID uint64, start []byte, end []byte, routeStart []byte, routeEnd []byte) ([]distribution.Route, bool, error) { + if routeScanBoundsPresent(routeStart, routeEnd) { + return s.routesForExplicitGroupRouteBounds(groupID, start, end, routeStart, routeEnd) } - return []distribution.Route{route}, version, true + return s.routesForExplicitGroupScan(groupID, start, end) +} + +func (s *ShardStore) routesForExplicitGroupRouteBounds(groupID uint64, start []byte, end []byte, routeStart []byte, routeEnd []byte) ([]distribution.Route, bool, error) { + fallback := []distribution.Route{{GroupID: groupID}} + if s == nil || s.engine == nil { + return fallback, false, nil + } + routes := s.engine.GetIntersectingRoutes(routeStart, normalizedRouteScanEnd(routeEnd)) + matched := make([]distribution.Route, 0, len(routes)) + for _, route := range routes { + if route.GroupID == groupID { + matched = append(matched, route) + continue + } + if routeHasStagedVisibility(route) { + return nil, false, errors.Wrapf(ErrExplicitGroupStagedVisibilityUnresolved, "group_id=%d range=[%q,%q)", groupID, start, end) + } + } + if len(matched) == 0 { + if len(routes) > 0 && !explicitGroupResolverOwnedKey(start) { + return nil, false, errors.Wrapf( + ErrExplicitGroupRouteOwnerMismatch, + "group_id=%d range=[%q,%q)", groupID, start, end) + } + return fallback, false, nil + } + return matched, false, nil +} + +func (s *ShardStore) routesForExplicitGroupScan(groupID uint64, start []byte, end []byte) ([]distribution.Route, bool, error) { + fallback := []distribution.Route{{GroupID: groupID}} + if s == nil || s.engine == nil { + return fallback, false, nil + } + routeStart, routeEnd, routeMapped := explicitGroupScanRouteBounds(start, end) + routes := s.engine.GetIntersectingRoutes(routeStart, routeEnd) + matched := make([]distribution.Route, 0, len(routes)) + for _, route := range routes { + if route.GroupID == groupID { + matched = append(matched, route) + continue + } + if routeHasStagedVisibility(route) { + return nil, false, errors.Wrapf(ErrExplicitGroupStagedVisibilityUnresolved, "group_id=%d range=[%q,%q)", groupID, start, end) + } + } + if len(matched) > 0 { + if routeMapped { + matched = dedupeRepeatedRawScanRoutes(matched) + } + return matched, !routeMapped, nil + } + if len(routes) > 0 && !explicitGroupResolverOwnedKey(start) { + return nil, false, errors.Wrapf( + ErrExplicitGroupRouteOwnerMismatch, + "group_id=%d range=[%q,%q)", groupID, start, end) + } + return fallback, false, nil +} + +func explicitGroupScanRouteBounds(start []byte, end []byte) ([]byte, []byte, bool) { + routeStart := routeKey(start) + if len(start) == 0 { + routeStart = []byte("") + } + routeEnd := end + routeMapped := !bytes.Equal(routeStart, start) + if end != nil { + normalizedEnd := routeKey(end) + if !bytes.Equal(normalizedEnd, end) { + routeMapped = true + } + } + if routeMapped && len(routeStart) != 0 { + routeEnd = prefixScanEnd(routeStart) + } + return routeStart, routeEnd, routeMapped } const legacyListDeltaRouteCandidateCapacity = 2 @@ -970,19 +1448,47 @@ func isBroadLegacyListDeltaScan(start []byte) bool { return logicalUserKey == nil || !bytes.Equal(start, store.LegacyListMetaDeltaScanPrefix(logicalUserKey)) } -func listScanUserKey(start []byte) []byte { - // Internal list keys route by their logical user key rather than their raw - // storage prefix. Well-formed full keys and exact scan prefixes are matched - // first, under the strict length checks listRouteKey applies. A resumed scan - // starts from a cursor inside the prefix, so only that case falls through to - // the permissive form that tolerates a trailing cursor tail. - if userKey := listRouteKey(start); userKey != nil { - return userKey - } - if userKey := store.ExtractListUserKeyFromDeltaScanKey(start); userKey != nil { - return userKey +// isLegacyListDeltaScan reports whether start addresses the legacy list-delta +// family at all, whether broad or scoped to a single user key. +// +// Route-group marking must cover both shapes. Redis cleanup and compaction build +// their delete elems as {Op: Del, Key: pair.Key, GroupID: pair.RouteGroupID}; a +// zero GroupID falls back to routing by the raw "!lst|meta|d|..." key rather +// than the logical list key, so after a split the delete lands on the wrong +// shard and the stale delta survives. Route *selection* still distinguishes the +// two shapes -- see routesForInternalScanWithVersion, where an exact scan is +// deliberately narrowed to the logical key's routes. +func isLegacyListDeltaScan(start []byte) bool { + return bytes.HasPrefix(start, []byte(store.LegacyListMetaDeltaPrefix)) +} + +func shouldMarkRouteGroupOnScan(start []byte, explicitGroup bool, routeStart []byte, routeEnd []byte) bool { + return !explicitGroup && !routeScanBoundsPresent(routeStart, routeEnd) && isLegacyListDeltaScan(start) +} + +func scanRouteUserKey(start []byte) []byte { + for _, extract := range scanRouteUserKeyExtractors { + if userKey := extract(start); userKey != nil { + return userKey + } } - return store.ExtractListUserKeyFromClaimScanKey(start) + return nil +} + +var scanRouteUserKeyExtractors = []func([]byte) []byte{ + store.ExtractListUserKeyFromDeltaScanPrefix, + store.ExtractListUserKey, + store.ExtractListUserKeyFromClaimScanPrefix, + store.ExtractHashUserKeyFromField, + store.ExtractHashUserKeyFromDeltaScanPrefix, + store.ExtractSetUserKeyFromMember, + store.ExtractSetUserKeyFromDeltaScanPrefix, + store.ExtractZSetUserKeyFromMember, + store.ExtractZSetUserKeyFromScore, + store.ExtractZSetUserKeyFromScoreScanPrefix, + store.ExtractZSetUserKeyFromDeltaScanPrefix, + store.ExtractStreamUserKeyFromMeta, + store.ExtractStreamUserKeyFromEntryScanPrefix, } func (s *ShardStore) routesForFencedScanWithVersion(start []byte, end []byte, routeStart []byte, routeEnd []byte) ([]distribution.Route, bool, uint64) { @@ -1004,31 +1510,50 @@ func normalizedRouteScanEnd(routeEnd []byte) []byte { return routeEnd } -func (s *ShardStore) scanRoutesAtWithReadFence(ctx context.Context, routes []distribution.Route, start []byte, end []byte, limit int, ts uint64, clampToRoutes bool, readRouteVersion uint64, routeStart []byte, routeEnd []byte) ([]*store.KVPair, error) { - out := make([]*store.KVPair, 0) - seenGroups := make(map[uint64]struct{}) +type scanRouteOwnerFilterPlan struct { + routes []distribution.Route + routeFilterPresent bool + filterUsageOwners bool + dedupeByKey bool +} + +func prepareScanRouteOwnerFilters(routes []distribution.Route, start []byte, end []byte, clampToRoutes bool, routeStart []byte, routeEnd []byte) scanRouteOwnerFilterPlan { routeFilterPresent := routeScanBoundsPresent(routeStart, routeEnd) filterUsageOwners := !clampToRoutes && !routeFilterPresent && filesystemUsageScanOverlap(start, end) + dedupeByKey := s3BucketAuxiliaryScanBounds(start, end) + if !clampToRoutes && !routeFilterPresent { + routes, dedupeByKey = prepareUnclampedRawScanRoutes(routes, dedupeByKey) + } + return scanRouteOwnerFilterPlan{ + routes: routes, + routeFilterPresent: routeFilterPresent, + filterUsageOwners: filterUsageOwners, + dedupeByKey: dedupeByKey, + } +} + +func (s *ShardStore) scanRoutesAtWithReadFence(ctx context.Context, routes []distribution.Route, start []byte, end []byte, limit int, ts uint64, clampToRoutes bool, readRouteVersion uint64, routeStart []byte, routeEnd []byte) ([]*store.KVPair, error) { + out := make([]*store.KVPair, 0) + plan := prepareScanRouteOwnerFilters(routes, start, end, clampToRoutes, routeStart, routeEnd) + routes = plan.routes for _, route := range routes { scanStart := start scanEnd := end if clampToRoutes { scanStart = clampScanStart(start, route.Start) scanEnd = clampScanEnd(end, route.End) - } else if !routeFilterPresent { - if _, seen := seenGroups[route.GroupID]; seen { - continue - } - seenGroups[route.GroupID] = struct{}{} } - kvs, err := s.scanRouteAtWithOptionalFilesystemUsageOwnerFilter( - ctx, route, scanStart, scanEnd, limit, ts, false, - readRouteVersion, routeStart, routeEnd, filterUsageOwners, + kvs, err := s.scanRouteAtWithMigrationOwnerFilters( + ctx, routes, route, scanStart, scanEnd, limit, ts, false, !clampToRoutes, + readRouteVersion, routeStart, routeEnd, plan.filterUsageOwners, plan.dedupeByKey, ) if err != nil { return nil, err } + if isLegacyListDeltaScan(start) && !plan.routeFilterPresent { + kvs = markScanRouteGroup(kvs, route.GroupID, true) + } if clampToRoutes { out = append(out, kvs...) if len(out) >= limit { @@ -1037,7 +1562,7 @@ func (s *ShardStore) scanRoutesAtWithReadFence(ctx context.Context, routes []dis } continue } - out = mergeAndTrimScanResults(out, kvs, limit) + out = mergeAndTrimScanResultsWithOptions(out, kvs, limit, false, plan.dedupeByKey) } return out, nil } @@ -1050,6 +1575,7 @@ func (s *ShardStore) scanRouteAtWithOptionalFilesystemUsageOwnerFilter( limit int, ts uint64, reverse bool, + explicitGroup bool, readRouteVersion uint64, routeStart []byte, routeEnd []byte, @@ -1057,16 +1583,44 @@ func (s *ShardStore) scanRouteAtWithOptionalFilesystemUsageOwnerFilter( ) ([]*store.KVPair, error) { if filterUsageOwners { return s.scanRouteAtWithFilesystemUsageOwnerFilter( - ctx, route, start, end, limit, ts, reverse, + ctx, route, start, end, limit, ts, reverse, explicitGroup, readRouteVersion, routeStart, routeEnd, ) } return s.scanRouteAtDirectionWithReadFence( - ctx, route, start, end, limit, ts, reverse, + ctx, route, start, end, limit, ts, reverse, explicitGroup, readRouteVersion, routeStart, routeEnd, ) } +func (s *ShardStore) scanRouteAtWithMigrationOwnerFilters( + ctx context.Context, + routes []distribution.Route, + route distribution.Route, + start []byte, + end []byte, + limit int, + ts uint64, + reverse bool, + explicitGroup bool, + readRouteVersion uint64, + routeStart []byte, + routeEnd []byte, + filterUsageOwners bool, + dedupeByKey bool, +) ([]*store.KVPair, error) { + if filterUsageOwners { + return s.scanRouteAtWithOptionalFilesystemUsageOwnerFilter( + ctx, route, start, end, limit, ts, reverse, explicitGroup, + readRouteVersion, routeStart, routeEnd, true, + ) + } + return s.scanRouteAtDirectionWithS3AuxiliaryOwnerFilter( + ctx, routes, route, start, end, limit, ts, reverse, explicitGroup, + readRouteVersion, routeStart, routeEnd, dedupeByKey, + ) +} + func (s *ShardStore) routesForFilesystemUsageScanWithVersion(start []byte, end []byte) ([]distribution.Route, uint64, bool) { if !filesystemUsageScanOverlap(start, end) { return nil, 0, false @@ -1078,19 +1632,18 @@ func (s *ShardStore) routesForFilesystemUsageScanWithVersion(start []byte, end [ } func (s *ShardStore) routesForFilesystemChunkScanWithVersion(start []byte, end []byte) ([]distribution.Route, uint64, bool) { + allRoutes, version := s.engine.GetIntersectingRoutesWithVersion(nil, nil) if routeStart, routeEnd, ok := fskeys.ChunkScanRouteBounds(start, end); ok { - allRoutes, version := s.engine.GetIntersectingRoutesWithVersion(nil, nil) return intersectingRoutes(allRoutes, routeStart, routeEnd), version, true } chunkStart, chunkEnd, ok := filesystemChunkScanOverlap(start, end) if !ok { - return nil, 0, false + return nil, version, false } routeStart, routeEnd, ok := fskeys.ChunkScanRouteBounds(chunkStart, chunkEnd) if !ok { - return nil, 0, false + return nil, version, false } - allRoutes, version := s.engine.GetIntersectingRoutesWithVersion(nil, nil) // Raw scans can continue from the chunk keyspace into later filesystem // families, so include both raw and virtual chunk route groups rather than // narrowing the scan to chunks only. @@ -1150,7 +1703,7 @@ func filesystemChunkScanOverlap(start []byte, end []byte) ([]byte, []byte, bool) func (s *ShardStore) scanKeyRoutesAtWithReadFence(ctx context.Context, routes []distribution.Route, start []byte, end []byte, limit int, ts uint64, clampToRoutes bool, readRouteVersion uint64) ([][]byte, error) { out := make([][]byte, 0) - seenGroups := make(map[uint64]struct{}) + seenRoutes := make(map[repeatedRawScanRouteKey]struct{}) filterUsageOwners := !clampToRoutes && filesystemUsageScanOverlap(start, end) for _, route := range routes { scanStart := start @@ -1159,10 +1712,11 @@ func (s *ShardStore) scanKeyRoutesAtWithReadFence(ctx context.Context, routes [] scanStart = clampScanStart(start, route.Start) scanEnd = clampScanEnd(end, route.End) } else { - if _, seen := seenGroups[route.GroupID]; seen { + dedupeKey := repeatedRawScanRouteDedupeKey(route) + if _, seen := seenRoutes[dedupeKey]; seen { continue } - seenGroups[route.GroupID] = struct{}{} + seenRoutes[dedupeKey] = struct{}{} } var keys [][]byte @@ -1172,7 +1726,7 @@ func (s *ShardStore) scanKeyRoutesAtWithReadFence(ctx context.Context, routes [] ctx, route, scanStart, scanEnd, limit, ts, readRouteVersion, ) } else { - keys, err = s.scanKeyRouteAtWithReadFence(ctx, route, scanStart, scanEnd, limit, ts, readRouteVersion) + keys, err = s.scanKeyRouteAtWithReadFence(ctx, route, scanStart, scanEnd, limit, ts, !clampToRoutes, readRouteVersion) } if err != nil { return nil, err @@ -1189,6 +1743,20 @@ func (s *ShardStore) scanKeyRoutesAtWithReadFence(ctx context.Context, routes [] return out, nil } +func (s *ShardStore) scanExplicitGroupKeysAtWithReadFence(ctx context.Context, groupID uint64, start []byte, end []byte, limit int, ts uint64, readRouteVersion uint64, routeStart []byte, routeEnd []byte) ([][]byte, error) { + if limit <= 0 { + return [][]byte{}, nil + } + routes, clampToRoutes, err := s.routesForExplicitGroupScanWithRouteBounds(groupID, start, end, routeStart, routeEnd) + if err != nil { + return nil, err + } + if !clampToRoutes && !routeScanBoundsPresent(routeStart, routeEnd) { + routes, _ = prepareUnclampedRawScanRoutes(routes, false) + } + return s.scanKeyRoutesAtWithReadFence(ctx, routes, start, end, limit, ts, clampToRoutes, readRouteVersion) +} + func (s *ShardStore) filterFilesystemUsageKVsForGroup(kvs []*store.KVPair, groupID uint64) []*store.KVPair { write := 0 for _, pair := range kvs { @@ -1227,6 +1795,7 @@ func (s *ShardStore) scanRouteAtWithFilesystemUsageOwnerFilter( limit int, ts uint64, reverse bool, + explicitGroup bool, readRouteVersion uint64, routeStart []byte, routeEnd []byte, @@ -1240,7 +1809,7 @@ func (s *ShardStore) scanRouteAtWithFilesystemUsageOwnerFilter( cursorEnd := end for len(out) < limit { page, err := s.scanRouteAtDirectionWithReadFence( - ctx, route, cursorStart, cursorEnd, limit, ts, reverse, + ctx, route, cursorStart, cursorEnd, limit, ts, reverse, explicitGroup, readRouteVersion, routeStart, routeEnd, ) if err != nil { @@ -1252,7 +1821,7 @@ func (s *ShardStore) scanRouteAtWithFilesystemUsageOwnerFilter( if reverse { out = mergeAndTrimReverseScanResults(out, page, limit) } else { - out = mergeAndTrimScanResults(out, page, limit) + out = mergeAndTrimScanResultsWithOptions(out, page, limit, false, false) } if len(out) >= limit || pageLen < limit || advanceKey == nil { break @@ -1288,7 +1857,7 @@ func (s *ShardStore) scanKeyRouteAtWithFilesystemUsageOwnerFilter( out := make([][]byte, 0, limit) cursor := start for len(out) < limit { - page, err := s.scanKeyRouteAtWithReadFence(ctx, route, cursor, end, limit, ts, readRouteVersion) + page, err := s.scanKeyRouteAtWithReadFence(ctx, route, cursor, end, limit, ts, true, readRouteVersion) if err != nil { return nil, err } @@ -1320,10 +1889,9 @@ func (s *ShardStore) reverseScanRoutesAtWithReadFence( routeEnd []byte, ) ([]*store.KVPair, error) { out := make([]*store.KVPair, 0) - seenGroups := make(map[uint64]struct{}) - routeFilterPresent := routeScanBoundsPresent(routeStart, routeEnd) - filterUsageOwners := !clampToRoutes && !routeFilterPresent && filesystemUsageScanOverlap(start, end) - for i := range routes { + plan := prepareScanRouteOwnerFilters(routes, start, end, clampToRoutes, routeStart, routeEnd) + routes = plan.routes + for i := 0; i < len(routes); i++ { route := routes[i] if clampToRoutes { route = routes[len(routes)-1-i] @@ -1342,23 +1910,17 @@ func (s *ShardStore) reverseScanRoutesAtWithReadFence( // shards), keys from different routes may interleave in descending order. // Fetch up to limit from every route and merge+sort descending so the // result honours the ReverseScanAt contract. - // De-duplicate by GroupID: after a range split both halves share the same - // GroupID (same backing shard store), so only scan each group once unless - // route filters make each descriptor's logical interval distinct. - if !routeFilterPresent { - if _, seen := seenGroups[route.GroupID]; seen { - continue - } - seenGroups[route.GroupID] = struct{}{} - } - kvs, err := s.scanRouteAtWithOptionalFilesystemUsageOwnerFilter( - ctx, route, start, end, limit, ts, true, - readRouteVersion, routeStart, routeEnd, filterUsageOwners, + kvs, err := s.scanRouteAtWithMigrationOwnerFilters( + ctx, routes, route, start, end, limit, ts, true, true, + readRouteVersion, routeStart, routeEnd, plan.filterUsageOwners, plan.dedupeByKey, ) if err != nil { return nil, err } - out = mergeAndTrimReverseScanResults(out, kvs, limit) + if isLegacyListDeltaScan(start) && !plan.routeFilterPresent { + kvs = markScanRouteGroup(kvs, route.GroupID, true) + } + out = mergeAndTrimScanResultsWithOptions(out, kvs, limit, true, plan.dedupeByKey) } return out, nil } @@ -1371,7 +1933,7 @@ func (s *ShardStore) scanKeyRouteAt( limit int, ts uint64, ) ([][]byte, error) { - return s.scanKeyRouteAtWithReadFence(ctx, route, start, end, limit, ts, 0) + return s.scanKeyRouteAtWithReadFence(ctx, route, start, end, limit, ts, false, 0) } func (s *ShardStore) scanKeyRouteAtWithReadFence( @@ -1381,6 +1943,7 @@ func (s *ShardStore) scanKeyRouteAtWithReadFence( end []byte, limit int, ts uint64, + explicitGroup bool, readRouteVersion uint64, ) ([][]byte, error) { g, ok := s.groupForID(route.GroupID) @@ -1389,14 +1952,32 @@ func (s *ShardStore) scanKeyRouteAtWithReadFence( } if engineForGroup(g) == nil { + if routeHasStagedVisibility(route) { + return scanKeysWithRefill(start, end, limit, func(cursor []byte, pageLimit int) ([][]byte, error) { + kvs, err := s.scanRouteWithStagedVisibility(ctx, g, route, cursor, end, pageLimit, ts, false) + return keysFromKVs(kvs), err + }, func(keys [][]byte) ([][]byte, error) { + return s.canonicalizeRedisWideColumnScanKeys(ctx, keys, start, ts, readRouteVersion) + }) + } return s.scanKeysRouteLocal(ctx, g, start, end, limit, ts, readRouteVersion) } if isLinearizableRaftLeader(ctx, engineForGroup(g)) { - return s.scanKeysRouteAtLeader(ctx, g, start, end, limit, ts, readRouteVersion) + return s.scanKeysRouteAtLeader(ctx, g, route, start, end, limit, ts, readRouteVersion) } - return s.proxyScanKeysAt(ctx, g, start, end, limit, ts, route.GroupID, readRouteVersion) + routeBoundsPresent := routeHasStagedVisibility(route) + var routeStart, routeEnd []byte + if routeBoundsPresent { + routeStart = route.Start + routeEnd = route.End + } + groupID := proxyScanGroupID(route, explicitGroup, readRouteVersion, routeStart, routeEnd) + return s.proxyScanKeysAt( + ctx, g, start, end, limit, ts, groupID, readRouteVersion, + routeStart, routeEnd, routeBoundsPresent, + ) } func (s *ShardStore) scanKeysRouteLocal( @@ -1422,6 +2003,7 @@ func (s *ShardStore) scanKeysRouteLocal( func (s *ShardStore) scanKeysRouteAtLeader( ctx context.Context, g *ShardGroup, + route distribution.Route, start []byte, end []byte, limit int, @@ -1431,7 +2013,40 @@ func (s *ShardStore) scanKeysRouteAtLeader( if limit <= 0 { return [][]byte{}, nil } + if routeHasStagedVisibility(route) { + return s.scanStagedKeysRouteAtLeader(ctx, g, route, start, end, limit, ts, readRouteVersion) + } + return s.scanUnstagedKeysRouteAtLeader(ctx, g, route, start, end, limit, ts, readRouteVersion) +} + +func (s *ShardStore) scanStagedKeysRouteAtLeader( + ctx context.Context, + g *ShardGroup, + route distribution.Route, + start []byte, + end []byte, + limit int, + ts uint64, + readRouteVersion uint64, +) ([][]byte, error) { + return scanKeysWithRefill(start, end, limit, func(cursor []byte, pageLimit int) ([][]byte, error) { + kvs, err := s.scanRouteAtLeader(ctx, g, route, cursor, end, pageLimit, ts, false) + return keysFromKVs(kvs), err + }, func(keys [][]byte) ([][]byte, error) { + return s.canonicalizeRedisWideColumnScanKeys(ctx, keys, start, ts, readRouteVersion) + }) +} +func (s *ShardStore) scanUnstagedKeysRouteAtLeader( + ctx context.Context, + g *ShardGroup, + route distribution.Route, + start []byte, + end []byte, + limit int, + ts uint64, + readRouteVersion uint64, +) ([][]byte, error) { out := make([][]byte, 0, limit) cursor := start for len(out) < limit { @@ -1440,7 +2055,7 @@ func (s *ShardStore) scanKeysRouteAtLeader( return nil, errors.WithStack(err) } if len(keys) == 0 { - keys, err := s.scanLockOnlyVisibleKeysAtLeader(ctx, g, cursor, end, start, ts, limit, readRouteVersion) + keys, err := s.scanLockOnlyVisibleKeysAtLeader(ctx, g, route, cursor, end, start, ts, limit, readRouteVersion) if err != nil { return nil, err } @@ -1454,7 +2069,7 @@ func (s *ShardStore) scanKeysRouteAtLeader( if err != nil { return nil, err } - kvs, err := s.resolveScanLocks(ctx, g, keyKVs, lockKVs, ts) + kvs, err := s.resolveScanLocks(ctx, g, route, keyKVs, lockKVs, ts) if err != nil { return nil, err } @@ -1476,6 +2091,7 @@ func (s *ShardStore) scanKeysRouteAtLeader( func (s *ShardStore) scanLockOnlyVisibleKeysAtLeader( ctx context.Context, g *ShardGroup, + route distribution.Route, cursor []byte, end []byte, start []byte, @@ -1483,7 +2099,7 @@ func (s *ShardStore) scanLockOnlyVisibleKeysAtLeader( limit int, readRouteVersion uint64, ) ([][]byte, error) { - keys, err := s.scanLockOnlyKeysAtLeader(ctx, g, cursor, end, ts, limit) + keys, err := s.scanLockOnlyKeysAtLeader(ctx, g, route, cursor, end, ts, limit) if err != nil { return nil, err } @@ -1491,12 +2107,13 @@ func (s *ShardStore) scanLockOnlyVisibleKeysAtLeader( } func (s *ShardStore) visibleScanKeysForReadFence(ctx context.Context, keys [][]byte, start []byte, ts uint64, readRouteVersion uint64) ([][]byte, error) { - return s.canonicalizeRedisWideColumnScanKeys(ctx, filterTxnInternalKeys(keys), start, ts, readRouteVersion) + return s.canonicalizeRedisWideColumnScanKeys(ctx, filterScanInternalKeys(keys), start, ts, readRouteVersion) } func (s *ShardStore) scanLockOnlyKeysAtLeader( ctx context.Context, g *ShardGroup, + route distribution.Route, start []byte, end []byte, ts uint64, @@ -1509,11 +2126,11 @@ func (s *ShardStore) scanLockOnlyKeysAtLeader( if len(lockKVs) == 0 { return nil, nil } - kvs, err := s.resolveScanLocks(ctx, g, nil, lockKVs, ts) + kvs, err := s.resolveScanLocks(ctx, g, route, nil, lockKVs, ts) if err != nil { return nil, err } - return filterTxnInternalKeys(keysFromKVs(kvs)), nil + return filterScanInternalKeys(keysFromKVs(kvs)), nil } func (s *ShardStore) proxyScanKeysAt( @@ -1525,9 +2142,15 @@ func (s *ShardStore) proxyScanKeysAt( ts uint64, groupID uint64, readRouteVersion uint64, + routeStart []byte, + routeEnd []byte, + routeBoundsPresent bool, ) ([][]byte, error) { return scanKeysWithRefill(start, end, limit, func(cursor []byte, pageLimit int) ([][]byte, error) { - return s.proxyRawScanKeysAt(ctx, g, cursor, end, pageLimit, ts, groupID, readRouteVersion) + return s.proxyRawScanKeysAt( + ctx, g, cursor, end, pageLimit, ts, groupID, readRouteVersion, + routeStart, routeEnd, routeBoundsPresent, + ) }, nil) } @@ -1553,7 +2176,7 @@ func scanKeysWithRefill( break } - visibleKeys := filterTxnInternalKeys(keys) + visibleKeys := filterScanInternalKeys(keys) if filter != nil { visibleKeys, err = filter(visibleKeys) if err != nil { @@ -1610,14 +2233,14 @@ func (s *ShardStore) clampedReverseScanRouteAtWithReadFence( scanStart := clampScanStart(start, route.Start) scanEnd := clampScanEnd(end, route.End) - kvs, err := s.scanRouteAtDirectionWithReadFence(ctx, route, scanStart, scanEnd, limit-currentLen, ts, true, readRouteVersion, routeStart, routeEnd) + kvs, err := s.scanRouteAtDirectionWithReadFence(ctx, route, scanStart, scanEnd, limit-currentLen, ts, true, false, readRouteVersion, routeStart, routeEnd) if err != nil { return nil, false, err } return kvs, false, nil } -func (s *ShardStore) scanRouteAtDirection( +func (s *ShardStore) scanRouteAtDirectionWithReadFence( ctx context.Context, route distribution.Route, start []byte, @@ -1625,26 +2248,412 @@ func (s *ShardStore) scanRouteAtDirection( limit int, ts uint64, reverse bool, + explicitGroup bool, + readRouteVersion uint64, + routeStart []byte, + routeEnd []byte, ) ([]*store.KVPair, error) { - return s.scanRouteAtDirectionWithReadFence(ctx, route, start, end, limit, ts, reverse, 0, nil, nil) + if routeScanBoundsPresent(routeStart, routeEnd) { + return s.scanRouteAtDirectionWithReadFenceRouteFilter(ctx, route, start, end, limit, ts, reverse, explicitGroup, readRouteVersion, routeStart, routeEnd) + } + return s.scanRouteAtDirectionWithReadFenceOnce(ctx, route, start, end, limit, ts, reverse, explicitGroup, readRouteVersion, routeStart, routeEnd) } -func (s *ShardStore) scanRouteAtDirectionWithReadFence( +func (s *ShardStore) scanRouteAtDirectionWithS3AuxiliaryOwnerFilter( ctx context.Context, + routes []distribution.Route, route distribution.Route, start []byte, end []byte, limit int, ts uint64, reverse bool, + explicitGroup bool, readRouteVersion uint64, routeStart []byte, routeEnd []byte, + dedupeByKey bool, ) ([]*store.KVPair, error) { - if routeScanBoundsPresent(routeStart, routeEnd) { - return s.scanRouteAtDirectionWithReadFenceRouteFilter(ctx, route, start, end, limit, ts, reverse, readRouteVersion, routeStart, routeEnd) + if !dedupeByKey { + return s.scanRouteAtDirectionWithReadFence(ctx, route, start, end, limit, ts, reverse, explicitGroup, readRouteVersion, routeStart, routeEnd) + } + out := make([]*store.KVPair, 0, limit) + scanStart := start + scanEnd := end + for len(out) < limit { + page, err := s.scanRouteAtDirectionWithReadFence(ctx, route, scanStart, scanEnd, limit, ts, reverse, explicitGroup, readRouteVersion, routeStart, routeEnd) + if err != nil { + return nil, err + } + filtered, err := s.filterS3AuxiliaryKVsOwnedByRoute(ctx, page, routes, route, ts, readRouteVersion) + if err != nil { + return nil, err + } + out = append(out, filtered...) + if len(out) >= limit { + clear(out[limit:]) + return out[:limit], nil + } + if len(page) < limit { + return out, nil + } + lastKey := lastKVKey(page) + if lastKey == nil { + return out, nil + } + if reverse { + scanEnd = lastKey + } else { + scanStart = nextScanCursor(lastKey) + } + } + return out, nil +} + +func (s *ShardStore) filterS3AuxiliaryKVsOwnedByRoute( + ctx context.Context, + kvs []*store.KVPair, + routes []distribution.Route, + route distribution.Route, + ts uint64, + readRouteVersion uint64, +) ([]*store.KVPair, error) { + include, batches := s3BucketAuxiliaryOwnerProbePlan(kvs, routes, route) + if err := s.applyS3BucketAuxiliaryOwnerProbeBatches(ctx, include, batches, ts, readRouteVersion); err != nil { + return nil, err + } + return includedKVs(kvs, include), nil +} + +func s3BucketAuxiliaryOwnerProbePlan( + kvs []*store.KVPair, + routes []distribution.Route, + route distribution.Route, +) ([]bool, map[s3BucketAuxiliaryOwnerProbeRouteKey]*s3BucketAuxiliaryOwnerProbeBatch) { + include := make([]bool, len(kvs)) + batches := make(map[s3BucketAuxiliaryOwnerProbeRouteKey]*s3BucketAuxiliaryOwnerProbeBatch) + for i, kvp := range kvs { + if kvp == nil { + continue + } + owner, auxiliary := s3BucketAuxiliaryOwnerRoute(kvp.Key, routes) + if auxiliary && !routeMatchesS3BucketAuxiliaryOwner(route, owner) { + key := s3BucketAuxiliaryOwnerProbeRouteKeyFor(owner) + batch := batches[key] + if batch == nil { + batch = &s3BucketAuxiliaryOwnerProbeBatch{owner: owner} + batches[key] = batch + } + batch.probes = append(batch.probes, s3BucketAuxiliaryOwnerProbe{index: i, key: kvp.Key}) + continue + } + include[i] = true + } + return include, batches +} + +func (s *ShardStore) applyS3BucketAuxiliaryOwnerProbeBatches( + ctx context.Context, + include []bool, + batches map[s3BucketAuxiliaryOwnerProbeRouteKey]*s3BucketAuxiliaryOwnerProbeBatch, + ts uint64, + readRouteVersion uint64, +) error { + for _, batch := range batches { + keys := batch.probeKeys() + covered, err := s.s3BucketAuxiliaryOwnerHasVersionsAt(ctx, batch.owner, keys, ts, readRouteVersion) + if err != nil { + return err + } + if len(covered) != len(batch.probes) { + return errors.WithStack(errors.Newf("s3 auxiliary owner version probe returned %d results for %d keys", len(covered), len(batch.probes))) + } + for i, found := range covered { + if found { + continue + } + include[batch.probes[i].index] = true + } + } + return nil +} + +func includedKVs(kvs []*store.KVPair, include []bool) []*store.KVPair { + out := make([]*store.KVPair, 0, len(kvs)) + for i, kvp := range kvs { + if include[i] { + out = append(out, kvp) + } + } + return out +} + +type s3BucketAuxiliaryOwnerProbe struct { + index int + key []byte +} + +type s3BucketAuxiliaryOwnerProbeBatch struct { + owner distribution.Route + probes []s3BucketAuxiliaryOwnerProbe +} + +func (b *s3BucketAuxiliaryOwnerProbeBatch) probeKeys() [][]byte { + keys := make([][]byte, 0, len(b.probes)) + for _, probe := range b.probes { + keys = append(keys, probe.key) + } + return keys +} + +type s3BucketAuxiliaryOwnerProbeRouteKey struct { + groupID uint64 + routeID uint64 + staged bool + migrationJobID uint64 + routeStart string + routeEnd string +} + +func s3BucketAuxiliaryOwnerProbeRouteKeyFor(route distribution.Route) s3BucketAuxiliaryOwnerProbeRouteKey { + return s3BucketAuxiliaryOwnerProbeRouteKey{ + groupID: route.GroupID, + routeID: route.RouteID, + staged: routeHasStagedVisibility(route), + migrationJobID: route.MigrationJobID, + routeStart: string(route.Start), + routeEnd: string(route.End), + } +} + +func (s *ShardStore) s3BucketAuxiliaryOwnerHasVersionsAt( + ctx context.Context, + owner distribution.Route, + keys [][]byte, + ts uint64, + readRouteVersion uint64, +) ([]bool, error) { + if !routeHasStagedVisibility(owner) { + return s.s3BucketAuxiliaryLiveOwnerHasVersionsAt(ctx, owner, keys, ts, readRouteVersion) + } + return s.s3BucketAuxiliaryStagedOwnerHasVersionsAt(ctx, owner, keys, ts, readRouteVersion) +} + +func (s *ShardStore) s3BucketAuxiliaryLiveOwnerHasVersionsAt( + ctx context.Context, + owner distribution.Route, + keys [][]byte, + ts uint64, + readRouteVersion uint64, +) ([]bool, error) { + visible, ok, err := s.ownerRouteHasVersionsAtOrBefore(ctx, owner, keys, ts, readRouteVersion) + if err != nil { + return nil, err + } + return existenceProbeResults(visible, ok, owner, keys) +} + +func (s *ShardStore) s3BucketAuxiliaryStagedOwnerHasVersionsAt( + ctx context.Context, + owner distribution.Route, + keys [][]byte, + ts uint64, + readRouteVersion uint64, +) ([]bool, error) { + stagedKeys := s3BucketAuxiliaryStagedDataKeys(owner, keys) + staged, stagedOK, err := s.ownerRouteHasVersionsAtOrBefore(ctx, owner, stagedKeys, ts, readRouteVersion) + if err != nil { + return nil, err + } + + out, liveKeys, liveIndex := s3BucketAuxiliaryLiveProbePlan(keys, staged) + if len(liveKeys) == 0 { + return out, nil + } + + live, liveOK, err := s.ownerRouteHasVersionsAtOrBefore(ctx, owner, liveKeys, ts, readRouteVersion) + if err != nil { + return nil, err + } + if err := applyS3BucketAuxiliaryLiveProbeResults(out, live, liveIndex, liveKeys, owner, stagedOK && liveOK); err != nil { + return nil, err + } + return out, nil +} + +func s3BucketAuxiliaryStagedDataKeys(owner distribution.Route, keys [][]byte) [][]byte { + stagedKeys := make([][]byte, 0, len(keys)) + for _, key := range keys { + stagedKeys = append(stagedKeys, distribution.MigrationStagedDataKey(owner.MigrationJobID, key)) + } + return stagedKeys +} + +func s3BucketAuxiliaryLiveProbePlan(keys [][]byte, staged []bool) ([]bool, [][]byte, []int) { + out := make([]bool, len(keys)) + liveKeys := make([][]byte, 0, len(keys)) + liveIndex := make([]int, 0, len(keys)) + for i, found := range staged { + if found { + out[i] = true + continue + } + liveKeys = append(liveKeys, keys[i]) + liveIndex = append(liveIndex, i) + } + return out, liveKeys, liveIndex +} + +func applyS3BucketAuxiliaryLiveProbeResults( + out []bool, + live []bool, + liveIndex []int, + liveKeys [][]byte, + owner distribution.Route, + answered bool, +) error { + for i, found := range live { + if found { + out[liveIndex[i]] = true + continue + } + if !answered { + return ownerVersionProbeUnavailable(owner, liveKeys[i]) + } + } + return nil +} + +func existenceProbeResults(found []bool, answered bool, owner distribution.Route, keys [][]byte) ([]bool, error) { + if answered { + return found, nil + } + for i, visible := range found { + if visible { + continue + } + return nil, ownerVersionProbeUnavailable(owner, keys[i]) + } + return found, nil +} + +func ownerVersionProbeUnavailable(route distribution.Route, key []byte) error { + return errors.Wrapf(ErrLeaderNotFound, "s3 auxiliary owner version probe unavailable group_id=%d key=%q", route.GroupID, key) +} + +func (s *ShardStore) ownerRouteHasVersionsAtOrBefore( + ctx context.Context, + owner distribution.Route, + keys [][]byte, + ts uint64, + readRouteVersion uint64, +) ([]bool, bool, error) { + if exists, ok, err := s.routeHasVersionsAtOrBefore(ctx, owner, keys, ts); ok || err != nil { + return exists, ok, err + } + return s.routeHasVersionsAtOrBeforeRemote(ctx, owner, keys, ts, readRouteVersion) +} + +func (s *ShardStore) routeHasVersionsAtOrBefore( + ctx context.Context, + route distribution.Route, + keys [][]byte, + ts uint64, +) ([]bool, bool, error) { + out := make([]bool, len(keys)) + g, ok := s.groupForID(route.GroupID) + if !ok || g.Store == nil { + return out, true, nil + } + if engine := engineForGroup(g); engine != nil && !isLinearizableRaftLeader(ctx, engine) { + return out, false, nil + } + exists, err := versionsExistAtOrBefore(ctx, g.Store, keys, ts) + return exists, true, errors.WithStack(err) +} + +func versionsExistAtOrBefore(ctx context.Context, st store.MVCCStore, keys [][]byte, ts uint64) ([]bool, error) { + out := make([]bool, len(keys)) + versions, err := latestCandidateVersionsAt(ctx, st, keys, ts) + if err != nil { + return nil, err + } + for i, key := range keys { + _, out[i] = versions[string(key)] + } + return out, nil +} + +func (s *ShardStore) routeHasVersionsAtOrBeforeRemote( + ctx context.Context, + route distribution.Route, + keys [][]byte, + ts uint64, + readRouteVersion uint64, +) ([]bool, bool, error) { + out := make([]bool, len(keys)) + if len(keys) == 0 { + return out, true, nil + } + if ts == 0 { + return out, false, nil + } + cli, ok, err := s.routeVersionPresenceClient(route) + if err != nil { + return nil, false, err + } + if !ok { + return out, false, nil + } + rpcCtx, cancel := context.WithTimeout(ctx, proxyForwardTimeout) + defer cancel() + resp, err := cli.RawLatestCommitTS(rpcCtx, &pb.RawLatestCommitTSRequest{ + KeyBatch: pb.EncodeRawLatestCommitTSKeyBatch(keys), + ReadRouteVersion: readRouteVersion, + GroupId: route.GroupID, + VersionVisibleAtTs: ts, + }) + if err != nil { + return nil, false, errors.WithStack(err) } - return s.scanRouteAtDirectionWithReadFenceOnce(ctx, route, start, end, limit, ts, reverse, readRouteVersion, routeStart, routeEnd) + if !resp.GetVersionVisibleSupported() { + return out, false, nil + } + visible := resp.GetVersionVisibleResults() + if len(visible) != len(keys) { + return nil, false, errors.WithStack(errors.Newf("s3 auxiliary owner version probe returned %d results for %d keys", len(visible), len(keys))) + } + return append([]bool(nil), visible...), true, nil +} + +func (s *ShardStore) routeVersionPresenceClient(route distribution.Route) (pb.RawKVClient, bool, error) { + g, ok := s.groupForID(route.GroupID) + if !ok || g == nil { + return nil, false, nil + } + engine := engineForGroup(g) + if engine == nil { + return nil, false, nil + } + addr := leaderAddrFromEngine(engine) + if addr == "" { + return nil, false, nil + } + conn, err := s.connCache.ConnFor(addr) + if err != nil { + return nil, false, err + } + return pb.NewRawKVClient(conn), true, nil +} + +func routeMatchesS3BucketAuxiliaryOwner(route distribution.Route, owner distribution.Route) bool { + if route.GroupID != owner.GroupID { + return false + } + if route.RouteID == 0 || owner.RouteID == 0 { + return true + } + return route.RouteID == owner.RouteID } func (s *ShardStore) scanRouteAtDirectionWithReadFenceRouteFilter( @@ -1655,6 +2664,7 @@ func (s *ShardStore) scanRouteAtDirectionWithReadFenceRouteFilter( limit int, ts uint64, reverse bool, + explicitGroup bool, readRouteVersion uint64, routeStart []byte, routeEnd []byte, @@ -1669,7 +2679,7 @@ func (s *ShardStore) scanRouteAtDirectionWithReadFenceRouteFilter( for len(out) < limit { remaining := limit - len(out) batchLimit := routeFilteredScanBatchLimit(remaining) - kvs, cursorKVs, err := s.scanRouteAtDirectionWithReadFenceRouteFilterPage(ctx, route, scanStart, scanEnd, batchLimit, remaining, ts, reverse, readRouteVersion, filterStart, filterEnd) + kvs, cursorKVs, err := s.scanRouteAtDirectionWithReadFenceRouteFilterPage(ctx, route, scanStart, scanEnd, batchLimit, remaining, ts, reverse, explicitGroup, readRouteVersion, filterStart, filterEnd) if err != nil { return nil, err } @@ -1710,6 +2720,7 @@ func (s *ShardStore) scanRouteAtDirectionWithReadFenceRouteFilterPage( visibleLimit int, ts uint64, reverse bool, + explicitGroup bool, readRouteVersion uint64, routeStart []byte, routeEnd []byte, @@ -1718,21 +2729,25 @@ func (s *ShardStore) scanRouteAtDirectionWithReadFenceRouteFilterPage( if !ok || g == nil || g.Store == nil { return nil, nil, nil } + markRouteGroup := shouldMarkRouteGroupOnScan(start, explicitGroup, routeStart, routeEnd) + engine := engineForGroup(g) - if engineForGroup(g) == nil { - return s.scanRouteAtDirectionWithReadFenceRouteFilterLocalPage(ctx, g, start, end, limit, ts, reverse, readRouteVersion) + if engine == nil { + kvs, cursorKVs, err := s.scanRouteAtDirectionWithReadFenceRouteFilterLocalPage(ctx, g, route, start, end, limit, ts, reverse, readRouteVersion) + return markRouteFilterPage(kvs, cursorKVs, err, route.GroupID, markRouteGroup) } - - if isLinearizableRaftLeader(ctx, engineForGroup(g)) { - return s.scanRouteAtDirectionWithReadFenceRouteFilterLeaderPage(ctx, g, start, end, limit, visibleLimit, ts, reverse, readRouteVersion, routeStart, routeEnd) + if isLinearizableRaftLeader(ctx, engine) { + kvs, cursorKVs, err := s.scanRouteAtDirectionWithReadFenceRouteFilterLeaderPage(ctx, g, route, start, end, limit, visibleLimit, ts, reverse, readRouteVersion, routeStart, routeEnd) + return markRouteFilterPage(kvs, cursorKVs, err, route.GroupID, markRouteGroup) } - - return s.scanRouteAtDirectionWithReadFenceRouteFilterProxyPage(ctx, route, g, start, end, limit, ts, reverse, readRouteVersion, routeStart, routeEnd) + kvs, cursorKVs, err := s.scanRouteAtDirectionWithReadFenceRouteFilterProxyPage(ctx, route, g, start, end, limit, ts, reverse, explicitGroup, readRouteVersion, routeStart, routeEnd) + return markRouteFilterPage(kvs, cursorKVs, err, route.GroupID, markRouteGroup) } func (s *ShardStore) scanRouteAtDirectionWithReadFenceRouteFilterLocalPage( ctx context.Context, g *ShardGroup, + route distribution.Route, start []byte, end []byte, limit int, @@ -1740,11 +2755,11 @@ func (s *ShardStore) scanRouteAtDirectionWithReadFenceRouteFilterLocalPage( reverse bool, readRouteVersion uint64, ) ([]*store.KVPair, []*store.KVPair, error) { - kvs, err := s.scanRouteLocal(ctx, g, start, end, limit, ts, reverse) + kvs, err := s.scanRouteLocal(ctx, g, route, start, end, limit, ts, reverse) if err != nil { - return nil, nil, errors.WithStack(err) + return nil, nil, err } - filtered, err := s.canonicalizeRedisWideColumnScanResults(ctx, filterTxnInternalKVs(kvs), start, ts, readRouteVersion) + filtered, err := s.canonicalizeRedisWideColumnScanResults(ctx, filterScanInternalKVs(kvs), start, ts, readRouteVersion) if err != nil { return nil, nil, err } @@ -1754,6 +2769,7 @@ func (s *ShardStore) scanRouteAtDirectionWithReadFenceRouteFilterLocalPage( func (s *ShardStore) scanRouteAtDirectionWithReadFenceRouteFilterLeaderPage( ctx context.Context, g *ShardGroup, + route distribution.Route, start []byte, end []byte, limit int, @@ -1764,7 +2780,7 @@ func (s *ShardStore) scanRouteAtDirectionWithReadFenceRouteFilterLeaderPage( routeStart []byte, routeEnd []byte, ) ([]*store.KVPair, []*store.KVPair, error) { - kvs, cursorKVs, err := s.scanRouteAtLeaderRouteFilter(ctx, g, start, end, limit, visibleLimit, ts, reverse, routeStart, routeEnd) + kvs, cursorKVs, err := s.scanRouteAtLeaderRouteFilter(ctx, g, route, start, end, limit, visibleLimit, ts, reverse, routeStart, routeEnd) if err != nil { return nil, nil, err } @@ -1781,29 +2797,37 @@ func (s *ShardStore) scanRouteAtDirectionWithReadFenceRouteFilterProxyPage( limit int, ts uint64, reverse bool, + explicitGroup bool, readRouteVersion uint64, routeStart []byte, routeEnd []byte, ) ([]*store.KVPair, []*store.KVPair, error) { - kvs, err := s.proxyRawScanAt(ctx, g, start, end, limit, ts, reverse, route.GroupID, readRouteVersion, routeStart, routeEnd) + groupID := proxyScanGroupID(route, explicitGroup, readRouteVersion, routeStart, routeEnd) + kvs, err := s.proxyRawScanAt(ctx, g, start, end, limit, ts, reverse, groupID, readRouteVersion, routeStart, routeEnd) if err != nil { return nil, nil, err } - // Canonicalized here even though a peer running this branch already did it - // on its side. The peer may be running the parent binary instead: that one - // accepts group_id on RawScanAt and serves rawScanAtExplicitGroup while - // ignoring the fence and route-bounds fields this request adds, and it has - // no canonicalization at all, so it answers with physical rows. Dropping - // this pass would let a legacy hash/set/zset row through that a logical - // tombstone should have suppressed. It can go once the response can say it - // is already canonical. - filtered, err := s.canonicalizeRedisWideColumnScanResults(ctx, filterTxnInternalKVs(kvs), start, ts, readRouteVersion) + filtered, err := s.canonicalizeRedisWideColumnScanResults(ctx, filterScanInternalKVs(kvs), start, ts, readRouteVersion) if err != nil { return nil, nil, err } return filtered, kvs, nil } +func markRouteFilterPage(kvs []*store.KVPair, cursorKVs []*store.KVPair, err error, groupID uint64, markRouteGroup bool) ([]*store.KVPair, []*store.KVPair, error) { + if err != nil { + return nil, nil, err + } + return markScanRouteGroup(kvs, groupID, markRouteGroup), markScanRouteGroup(cursorKVs, groupID, markRouteGroup), nil +} + +func proxyScanGroupID(route distribution.Route, explicitGroup bool, readRouteVersion uint64, routeStart []byte, routeEnd []byte) uint64 { + if explicitGroup || readRouteVersion == 0 || routeScanBoundsPresent(routeStart, routeEnd) { + return route.GroupID + } + return 0 +} + func (s *ShardStore) scanRouteAtDirectionWithReadFenceOnce( ctx context.Context, route distribution.Route, @@ -1812,6 +2836,7 @@ func (s *ShardStore) scanRouteAtDirectionWithReadFenceOnce( limit int, ts uint64, reverse bool, + explicitGroup bool, readRouteVersion uint64, routeStart []byte, routeEnd []byte, @@ -1822,30 +2847,80 @@ func (s *ShardStore) scanRouteAtDirectionWithReadFenceOnce( } if !reverse { - return s.scanRouteAtForward(ctx, route, g, start, end, limit, ts, readRouteVersion, routeStart, routeEnd) - } - return s.scanRouteAtReverse(ctx, route, g, start, end, limit, ts, readRouteVersion, routeStart, routeEnd) -} - -const routeFilteredScanBatchMin = 128 - -func routeFilteredScanBatchLimit(remaining int) int { - if remaining <= 0 { - return 0 - } - limit := remaining - if limit < routeFilteredScanBatchMin { - limit = routeFilteredScanBatchMin - } - if maxLimit := store.MaxDeltaScanLimit + 1; limit > maxLimit { - limit = maxLimit + return s.scanRouteAtForward(ctx, route, g, start, end, limit, ts, explicitGroup, readRouteVersion, routeStart, routeEnd) } - return limit + return s.scanRouteAtReverse(ctx, route, g, start, end, limit, ts, explicitGroup, readRouteVersion, routeStart, routeEnd) } -func routeFilteredScanDone(kvs []*store.KVPair, batchLimit int, outLen int, limit int) bool { - return len(kvs) == 0 || outLen >= limit || len(kvs) < batchLimit -} +func (s *ShardStore) scanRouteAtReverse( + ctx context.Context, + route distribution.Route, + g *ShardGroup, + start []byte, + end []byte, + limit int, + ts uint64, + explicitGroup bool, + readRouteVersion uint64, + routeStart []byte, + routeEnd []byte, +) ([]*store.KVPair, error) { + markRouteGroup := shouldMarkRouteGroupOnScan(start, explicitGroup, routeStart, routeEnd) + engine := engineForGroup(g) + if engine == nil { + kvs, err := s.scanRouteLocal(ctx, g, route, start, end, limit, ts, true) + if err != nil { + return nil, errors.WithStack(err) + } + filtered, err := s.canonicalizeRedisWideColumnScanResults(ctx, filterScanInternalKVs(kvs), start, ts, readRouteVersion) + if err != nil { + return nil, err + } + return markScanRouteGroup(filtered, route.GroupID, markRouteGroup), nil + } + + if isLinearizableRaftLeader(ctx, engine) { + kvs, err := s.scanRouteAtLeader(ctx, g, route, start, end, limit, ts, true) + if err != nil { + return nil, err + } + kvs, err = s.canonicalizeRedisWideColumnScanResults(ctx, kvs, start, ts, readRouteVersion) + return markScanRouteGroup(kvs, route.GroupID, markRouteGroup), err + } + + groupID := proxyScanGroupID(route, explicitGroup, readRouteVersion, routeStart, routeEnd) + kvs, err := s.proxyRawScanAt(ctx, g, start, end, limit, ts, true, groupID, readRouteVersion, routeStart, routeEnd) + if err != nil { + return nil, err + } + // The leader's RawScanAt is expected to perform lock resolution and filtering + // via ShardStore.ScanAt, so avoid N+1 proxy gets here. + filtered, err := s.canonicalizeRedisWideColumnScanResults(ctx, filterScanInternalKVs(kvs), start, ts, readRouteVersion) + if err != nil { + return nil, err + } + return markScanRouteGroup(filtered, route.GroupID, markRouteGroup), nil +} + +const routeFilteredScanBatchMin = 128 + +func routeFilteredScanBatchLimit(remaining int) int { + if remaining <= 0 { + return 0 + } + limit := remaining + if limit < routeFilteredScanBatchMin { + limit = routeFilteredScanBatchMin + } + if maxLimit := store.MaxDeltaScanLimit + 1; limit > maxLimit { + limit = maxLimit + } + return limit +} + +func routeFilteredScanDone(kvs []*store.KVPair, batchLimit int, outLen int, limit int) bool { + return len(kvs) == 0 || outLen >= limit || len(kvs) < batchLimit +} func nextRouteFilteredScanWindow(kvs []*store.KVPair, scanStart []byte, scanEnd []byte, reverse bool) ([]byte, []byte, bool) { lastKey := kvs[len(kvs)-1].Key @@ -1909,7 +2984,10 @@ func minScanEnd(a []byte, b []byte) []byte { } func routeKeyInScanBounds(key []byte, routeStart []byte, routeEnd []byte) bool { - key = routeFilterKey(key) + if s3BucketAuxiliaryRouteInRange(key, routeStart, routeEnd) { + return true + } + key = routeKey(key) if len(routeStart) > 0 && bytes.Compare(key, routeStart) < 0 { return false } @@ -1933,6 +3011,7 @@ func (s *ShardStore) scanRouteAtForward( end []byte, limit int, ts uint64, + explicitGroup bool, readRouteVersion uint64, routeStart []byte, routeEnd []byte, @@ -1944,11 +3023,11 @@ func (s *ShardStore) scanRouteAtForward( out := make([]*store.KVPair, 0, limit) cursor := start for len(out) < limit { - page, err := s.scanRouteAtForwardPage(ctx, route, g, cursor, end, limit, ts, readRouteVersion, routeStart, routeEnd) + page, err := s.scanRouteAtForwardPage(ctx, route, g, cursor, end, limit, ts, explicitGroup, readRouteVersion, routeStart, routeEnd) if err != nil { return nil, err } - out = mergeAndTrimScanResults(out, page.kvs, limit) + out = mergeAndTrimScanResultsWithOptions(out, page.kvs, limit, false, false) if len(out) >= limit { break } @@ -1968,7 +3047,8 @@ func (s *ShardStore) scanRouteAtForward( if len(out) > limit { out = out[:limit] } - return out, nil + markRouteGroup := shouldMarkRouteGroupOnScan(start, explicitGroup, routeStart, routeEnd) + return markScanRouteGroup(out, route.GroupID, markRouteGroup), nil } func (s *ShardStore) scanRouteAtForwardPage( @@ -1979,36 +3059,36 @@ func (s *ShardStore) scanRouteAtForwardPage( end []byte, limit int, ts uint64, + explicitGroup bool, readRouteVersion uint64, routeStart []byte, routeEnd []byte, ) (scanRoutePage, error) { engine := engineForGroup(g) if engine == nil { - return s.scanRouteAtForwardLocalPage(ctx, g, start, end, limit, ts, readRouteVersion) + return s.scanRouteAtForwardLocalPage(ctx, g, route, start, end, limit, ts, readRouteVersion) } - if isLinearizableRaftLeader(ctx, engine) { - return s.scanRouteAtForwardLeaderPage(ctx, g, start, end, limit, ts, readRouteVersion) + return s.scanRouteAtForwardLeaderPage(ctx, g, route, start, end, limit, ts, readRouteVersion) } - - return s.scanRouteAtForwardProxyPage(ctx, route, g, start, end, limit, ts, readRouteVersion, routeStart, routeEnd) + return s.scanRouteAtForwardProxyPage(ctx, route, g, start, end, limit, ts, explicitGroup, readRouteVersion, routeStart, routeEnd) } func (s *ShardStore) scanRouteAtForwardLocalPage( ctx context.Context, g *ShardGroup, + route distribution.Route, start []byte, end []byte, limit int, ts uint64, readRouteVersion uint64, ) (scanRoutePage, error) { - raw, err := s.scanRouteLocal(ctx, g, start, end, limit, ts, false) + raw, err := s.scanRouteLocal(ctx, g, route, start, end, limit, ts, false) if err != nil { return scanRoutePage{}, errors.WithStack(err) } - kvs, err := s.canonicalizeRedisWideColumnScanResults(ctx, filterTxnInternalKVs(raw), start, ts, readRouteVersion) + kvs, err := s.canonicalizeRedisWideColumnScanResults(ctx, filterScanInternalKVs(raw), start, ts, readRouteVersion) if err != nil { return scanRoutePage{}, err } @@ -2022,13 +3102,14 @@ func (s *ShardStore) scanRouteAtForwardLocalPage( func (s *ShardStore) scanRouteAtForwardLeaderPage( ctx context.Context, g *ShardGroup, + route distribution.Route, start []byte, end []byte, limit int, ts uint64, readRouteVersion uint64, ) (scanRoutePage, error) { - raw, err := g.Store.ScanAt(ctx, start, end, limit, ts) + raw, err := s.scanRouteLocal(ctx, g, route, start, end, limit, ts, false) if err != nil { return scanRoutePage{}, errors.WithStack(err) } @@ -2037,11 +3118,11 @@ func (s *ShardStore) scanRouteAtForwardLeaderPage( if err != nil { return scanRoutePage{}, err } - kvs, err := s.resolveScanLocks(ctx, g, raw, lockKVs, ts) + kvs, err := s.resolveScanLocks(ctx, g, route, raw, lockKVs, ts) if err != nil { return scanRoutePage{}, err } - kvs, err = s.canonicalizeRedisWideColumnScanResults(ctx, filterTxnInternalKVs(kvs), start, ts, readRouteVersion) + kvs, err = s.canonicalizeRedisWideColumnScanResults(ctx, filterScanInternalKVs(kvs), start, ts, readRouteVersion) if err != nil { return scanRoutePage{}, err } @@ -2060,18 +3141,17 @@ func (s *ShardStore) scanRouteAtForwardProxyPage( end []byte, limit int, ts uint64, + explicitGroup bool, readRouteVersion uint64, routeStart []byte, routeEnd []byte, ) (scanRoutePage, error) { - raw, err := s.proxyRawScanAt(ctx, g, start, end, limit, ts, false, route.GroupID, readRouteVersion, routeStart, routeEnd) + groupID := proxyScanGroupID(route, explicitGroup, readRouteVersion, routeStart, routeEnd) + raw, err := s.proxyRawScanAt(ctx, g, start, end, limit, ts, false, groupID, readRouteVersion, routeStart, routeEnd) if err != nil { return scanRoutePage{}, err } - // Canonicalized here; see - // scanRouteAtDirectionWithReadFenceRouteFilterProxyPage for why a peer's - // page cannot be assumed canonical. - kvs, err := s.canonicalizeRedisWideColumnScanResults(ctx, filterTxnInternalKVs(raw), start, ts, readRouteVersion) + kvs, err := s.canonicalizeRedisWideColumnScanResults(ctx, filterScanInternalKVs(raw), start, ts, readRouteVersion) if err != nil { return scanRoutePage{}, err } @@ -2082,282 +3162,617 @@ func (s *ShardStore) scanRouteAtForwardProxyPage( }, nil } -func (s *ShardStore) scanRouteAtReverse( +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) +} + +func (s *ShardStore) scanRouteAtDirectionPhysicalLimit( ctx context.Context, route distribution.Route, - g *ShardGroup, start []byte, end []byte, - limit int, + visibleLimit int, + physicalLimit int, ts uint64, - readRouteVersion uint64, - routeStart []byte, - routeEnd []byte, -) ([]*store.KVPair, error) { - if limit <= 0 { - return []*store.KVPair{}, nil + reverse bool, +) ([]*store.KVPair, bool, error) { + g, ok := s.groupForID(route.GroupID) + if !ok || g == nil || g.Store == nil { + return nil, false, nil } + markRouteGroup := shouldMarkRouteGroupOnScan(start, false, nil, nil) - out := make([]*store.KVPair, 0, limit) - cursorEnd := end - for len(out) < limit { - page, err := s.scanRouteAtReversePage(ctx, route, g, start, cursorEnd, limit, ts, readRouteVersion, routeStart, routeEnd) - if err != nil { - return nil, err + if engineForGroup(g) == nil { + if routeHasStagedVisibility(route) { + kvs, err := s.scanRouteLocal(ctx, g, route, start, end, visibleLimit, ts, reverse) + return markScanRouteGroup(kvs, route.GroupID, markRouteGroup), false, err } - out = mergeAndTrimReverseScanResults(out, page.kvs, limit) - if len(out) >= limit { - break + kvs, limitReached, err := scanLocalPhysicalLimit(ctx, g.Store, start, end, visibleLimit, physicalLimit, ts, reverse) + if err != nil { + return nil, limitReached, errors.WithStack(err) } - if !page.full || page.advanceKey == nil { - break + return markScanRouteGroup(filterScanInternalKVs(kvs), route.GroupID, markRouteGroup), limitReached, nil + } + + if isLinearizableRaftLeader(ctx, engineForGroup(g)) { + if routeHasStagedVisibility(route) { + kvs, err := s.scanRouteAtLeader(ctx, g, route, start, end, visibleLimit, ts, reverse) + return markScanRouteGroup(kvs, route.GroupID, markRouteGroup), false, err } - if len(start) > 0 && bytes.Compare(page.advanceKey, start) <= 0 { - break + kvs, limitReached, err := s.scanRouteAtLeaderPhysicalLimit(ctx, g, route, start, end, visibleLimit, physicalLimit, ts, reverse) + return markScanRouteGroup(kvs, route.GroupID, markRouteGroup), limitReached, err + } + + // RawScanAt cannot enforce physicalLimit, so report truncation and let + // callers fail closed instead of proxying an unbounded physical scan. + return nil, true, nil +} + +func scanLocalPhysicalLimit( + ctx context.Context, + st store.MVCCStore, + start []byte, + end []byte, + visibleLimit int, + physicalLimit int, + ts uint64, + reverse bool, +) ([]*store.KVPair, bool, error) { + scanner, ok := st.(physicalLimitedStore) + if !ok { + if reverse { + kvs, err := st.ReverseScanAt(ctx, start, end, visibleLimit, ts) + return kvs, false, errors.WithStack(err) } - cursorEnd = page.advanceKey + kvs, err := st.ScanAt(ctx, start, end, visibleLimit, ts) + return kvs, false, errors.WithStack(err) } + return scanPhysicalLimitLocal(ctx, scanner, start, end, visibleLimit, physicalLimit, ts, reverse) +} - 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] +func scanPhysicalLimitLocal( + ctx context.Context, + scanner physicalLimitedStore, + start []byte, + end []byte, + visibleLimit int, + physicalLimit int, + ts uint64, + reverse bool, +) ([]*store.KVPair, bool, error) { + if reverse { + kvs, limitReached, err := scanner.ReverseScanAtPhysicalLimit(ctx, start, end, visibleLimit, physicalLimit, ts) + return kvs, limitReached, errors.WithStack(err) } - return out, nil + kvs, limitReached, err := scanner.ScanAtPhysicalLimit(ctx, start, end, visibleLimit, physicalLimit, ts) + return kvs, limitReached, errors.WithStack(err) } -func (s *ShardStore) scanRouteAtReversePage( +func (s *ShardStore) scanRouteLocal( ctx context.Context, - route distribution.Route, g *ShardGroup, + route distribution.Route, start []byte, end []byte, limit int, ts uint64, - readRouteVersion uint64, - routeStart []byte, - routeEnd []byte, -) (scanRoutePage, error) { - engine := engineForGroup(g) - if engine == nil { - return s.scanRouteAtReverseLocalPage(ctx, g, start, end, limit, ts, readRouteVersion) + reverse bool, +) ([]*store.KVPair, error) { + if routeHasStagedVisibility(route) { + return s.scanRouteWithStagedVisibility(ctx, g, route, start, end, limit, ts, reverse) } - - if isLinearizableRaftLeader(ctx, engine) { - return s.scanRouteAtReverseLeaderPage(ctx, g, start, end, limit, ts, readRouteVersion) + if reverse { + kvs, err := g.Store.ReverseScanAt(ctx, start, end, limit, ts) + return kvs, errors.WithStack(err) } - - return s.scanRouteAtReverseProxyPage(ctx, route, g, start, end, limit, ts, readRouteVersion, routeStart, routeEnd) + kvs, err := g.Store.ScanAt(ctx, start, end, limit, ts) + return kvs, errors.WithStack(err) } -func (s *ShardStore) scanRouteAtReverseLocalPage( +func (s *ShardStore) scanRouteAtLeaderPhysicalLimit( ctx context.Context, g *ShardGroup, + route distribution.Route, start []byte, end []byte, - limit int, + visibleLimit int, + physicalLimit int, ts uint64, - readRouteVersion uint64, -) (scanRoutePage, error) { - raw, err := s.scanRouteLocal(ctx, g, start, end, limit, ts, true) + reverse bool, +) ([]*store.KVPair, bool, error) { + kvs, limitReached, err := scanLocalPhysicalLimit(ctx, g.Store, start, end, visibleLimit, physicalLimit, ts, reverse) if err != nil { - return scanRoutePage{}, errors.WithStack(err) + return nil, limitReached, errors.WithStack(err) } - kvs, err := s.canonicalizeRedisWideColumnScanResults(ctx, filterTxnInternalKVs(raw), start, ts, readRouteVersion) + lockStart, lockEnd := scanLockBoundsForKVsDirection(kvs, start, end, visibleLimit, reverse) + lockKVs, err := scanTxnLockRangeAt(ctx, g, lockStart, lockEnd, ts, visibleLimit) if err != nil { - return scanRoutePage{}, err + return nil, limitReached, err } - return scanRoutePage{ - kvs: kvs, - advanceKey: lastKVKey(raw), - full: len(raw) >= limit, - }, nil + resolved, err := s.resolveScanLocks(ctx, g, route, kvs, lockKVs, ts) + return resolved, limitReached, err } -func (s *ShardStore) scanRouteAtReverseLeaderPage( +func (s *ShardStore) scanRouteAtLeader( ctx context.Context, g *ShardGroup, + route distribution.Route, start []byte, end []byte, limit int, ts uint64, - readRouteVersion uint64, -) (scanRoutePage, error) { - raw, err := g.Store.ReverseScanAt(ctx, start, end, limit, ts) + reverse bool, +) ([]*store.KVPair, error) { + var ( + kvs []*store.KVPair + err error + ) + switch { + case routeHasStagedVisibility(route): + kvs, err = s.scanRouteWithStagedVisibility(ctx, g, route, start, end, limit, ts, reverse) + case reverse: + kvs, err = g.Store.ReverseScanAt(ctx, start, end, limit, ts) + default: + kvs, err = g.Store.ScanAt(ctx, start, end, limit, ts) + } if err != nil { - return scanRoutePage{}, errors.WithStack(err) + return nil, errors.WithStack(err) } - lockStart, lockEnd := scanLockBoundsForKVsDirection(raw, start, end, limit, true) + lockStart, lockEnd := scanLockBoundsForKVsDirection(kvs, start, end, limit, reverse) lockKVs, err := scanTxnLockRangeAt(ctx, g, lockStart, lockEnd, ts, limit) if err != nil { - return scanRoutePage{}, err + return nil, err } - kvs, err := s.resolveScanLocks(ctx, g, raw, lockKVs, ts) - if err != nil { - return scanRoutePage{}, err + return s.resolveScanLocks(ctx, g, route, kvs, lockKVs, ts) +} + +const ( + stagedVisibilityMaxCandidateWindow = 8192 + stagedVisibilityWindowGrowthFactor = 2 +) + +func (s *ShardStore) scanRouteWithStagedVisibility( + ctx context.Context, + g *ShardGroup, + route distribution.Route, + start []byte, + end []byte, + limit int, + ts uint64, + reverse bool, +) ([]*store.KVPair, error) { + if err := ensureReadTSRetained(g.Store, ts); err != nil { + return nil, err } - kvs, err = s.canonicalizeRedisWideColumnScanResults(ctx, filterTxnInternalKVs(kvs), start, ts, readRouteVersion) - if err != nil { - return scanRoutePage{}, err + out := make([]*store.KVPair, 0, limit) + scanStart := bytes.Clone(start) + scanEnd := bytes.Clone(end) + for len(out) < limit { + remaining := limit - len(out) + kvs, boundary, hasMore, err := s.scanRouteWithStagedVisibilityPage(ctx, g, route, scanStart, scanEnd, remaining, ts, reverse) + if err != nil { + return nil, err + } + out = append(out, kvs...) + if len(out) >= limit { + clear(out[limit:]) + return out[:limit], nil + } + if !hasMore { + return out, nil + } + if reverse { + scanEnd = boundary + } else { + scanStart = exclusiveScanStartAfter(boundary) + } + } + return out, nil +} + +func (s *ShardStore) scanRouteWithStagedVisibilityPage( + ctx context.Context, + g *ShardGroup, + route distribution.Route, + start []byte, + end []byte, + limit int, + ts uint64, + reverse bool, +) ([]*store.KVPair, []byte, bool, error) { + stagedStart, stagedEnd := stagedVisibilityScanBounds(route.MigrationJobID, start, end) + window := stagedVisibilityCandidateWindow(limit) + for { + // Staged before live, for the reason on getAtWithStagedVisibility: + // promotion only moves rows staged -> live, so reading the shrinking + // side first keeps a key that is mid-promotion visible to one of the + // two scans. + stagedKVs, err := scanVisibleCandidates(ctx, g.Store, stagedStart, stagedEnd, window, ts, reverse) + if err != nil { + return nil, nil, false, err + } + liveKVs, err := scanVisibleCandidates(ctx, g.Store, start, end, window, ts, reverse) + if err != nil { + return nil, nil, false, err + } + versions, err := s.latestStagedVisibilityCandidates(ctx, g.Store, route, liveKVs, stagedKVs, ts) + if err != nil { + return nil, nil, false, err + } + out := visibleLogicalKVs(versions, ts, reverse) + out = filterRouteScanKVs(out, route.Start, route.End) + liveExhausted := len(liveKVs) < window + stagedExhausted := len(stagedKVs) < window + boundary, hasBoundary := stagedVisibilityCandidateBoundary(liveKVs, stagedKVs, liveExhausted, stagedExhausted, reverse) + exhausted := liveExhausted && stagedExhausted + out = stagedVisibilityKVsWithinPageBoundary(out, boundary, hasBoundary, exhausted, reverse) + if len(out) >= limit { + clear(out[limit:]) + return out[:limit], boundary, !exhausted && hasBoundary, nil + } + if exhausted { + return out, nil, false, nil + } + nextWindow := nextStagedVisibilityCandidateWindow(window) + if nextWindow == window { + return out, boundary, hasBoundary, nil + } + window = nextWindow + } +} + +func stagedVisibilityKVsWithinPageBoundary(kvs []*store.KVPair, boundary []byte, hasBoundary bool, exhausted bool, reverse bool) []*store.KVPair { + if exhausted || !hasBoundary { + return kvs + } + return stagedVisibilityKVsWithinBoundary(kvs, boundary, reverse) +} + +func stagedVisibilityKVsWithinBoundary(kvs []*store.KVPair, boundary []byte, reverse bool) []*store.KVPair { + if len(boundary) == 0 { + return kvs + } + n := 0 + for _, kvp := range kvs { + if kvp == nil { + continue + } + cmp := bytes.Compare(kvp.Key, boundary) + if (!reverse && cmp <= 0) || (reverse && cmp >= 0) { + kvs[n] = kvp + n++ + } + } + clear(kvs[n:]) + return kvs[:n] +} + +func stagedVisibilityCandidateBoundary(liveKVs []*store.KVPair, stagedKVs []*store.KVPair, liveExhausted bool, stagedExhausted bool, reverse bool) ([]byte, bool) { + liveBoundary := stagedVisibilityBoundary{reverse: reverse} + for _, kvp := range liveKVs { + if kvp == nil { + continue + } + if isMigrationStagedDataKey(kvp.Key) { + continue + } + liveBoundary.visit(kvp.Key) + } + stagedBoundary := stagedVisibilityBoundary{reverse: reverse} + for _, kvp := range stagedKVs { + rawKey, stagedOK := stagedVisibilityRawCandidateKey(kvp) + if !stagedOK { + continue + } + stagedBoundary.visit(rawKey) + } + return mergeStagedVisibilityBoundaries(liveBoundary, stagedBoundary, liveExhausted, stagedExhausted, reverse) +} + +func mergeStagedVisibilityBoundaries(live stagedVisibilityBoundary, staged stagedVisibilityBoundary, liveExhausted bool, stagedExhausted bool, reverse bool) ([]byte, bool) { + if !live.ok { + return staged.key, staged.ok + } + if !staged.ok { + return live.key, live.ok + } + if !liveExhausted && !stagedExhausted { + return nearerStagedVisibilityBoundary(live.key, staged.key, reverse), true + } + if liveExhausted && stagedExhausted { + return fartherStagedVisibilityBoundary(live.key, staged.key, reverse), true + } + if liveExhausted { + return staged.key, true + } + return live.key, true +} + +func nearerStagedVisibilityBoundary(a []byte, b []byte, reverse bool) []byte { + cmp := bytes.Compare(a, b) + if (!reverse && cmp <= 0) || (reverse && cmp >= 0) { + return a + } + return b +} + +func fartherStagedVisibilityBoundary(a []byte, b []byte, reverse bool) []byte { + cmp := bytes.Compare(a, b) + if (!reverse && cmp >= 0) || (reverse && cmp <= 0) { + return a + } + return b +} + +type stagedVisibilityBoundary struct { + key []byte + ok bool + reverse bool +} + +func (b *stagedVisibilityBoundary) visit(key []byte) { + if !b.ok { + b.key = bytes.Clone(key) + b.ok = true + return + } + cmp := bytes.Compare(key, b.key) + if (!b.reverse && cmp > 0) || (b.reverse && cmp < 0) { + b.key = bytes.Clone(key) + } +} + +func stagedVisibilityRawCandidateKey(kvp *store.KVPair) ([]byte, bool) { + if kvp == nil { + return nil, false + } + _, rawKey, ok := distribution.MigrationStagedDataKeyParts(kvp.Key) + return rawKey, ok +} + +func exclusiveScanStartAfter(key []byte) []byte { + if key == nil { + return nil + } + out := bytes.Clone(key) + return append(out, 0) +} + +func stagedVisibilityCandidateWindow(limit int) int { + if limit <= 0 { + return 0 + } + if limit > stagedVisibilityMaxCandidateWindow { + return stagedVisibilityMaxCandidateWindow + } + return limit +} + +func nextStagedVisibilityCandidateWindow(window int) int { + if window >= stagedVisibilityMaxCandidateWindow { + return window + } + next := window * stagedVisibilityWindowGrowthFactor + if next < window || next > stagedVisibilityMaxCandidateWindow { + return stagedVisibilityMaxCandidateWindow + } + return next +} + +// stagedVisibilityCandidateScanBudget bounds the range export that resolves a +// staged-visibility page's candidate keys in one pass. +// +// The export spans from the smallest candidate through the largest and filters +// for the candidate set, so with no budget it scans and decodes every version +// in between -- tombstoned keys and dense MVCC history included -- even though +// at most stagedVisibilityMaxCandidateWindow exact keys are wanted. A route +// whose visible rows are sparse could therefore make an ordinary scan page +// consume unbounded I/O on the serving leader. +// +// With the budget the range pass stays the fast path for the dense case, where +// the candidates sit close together, and a page that would have run away instead +// falls back to one exact-key probe per unresolved candidate -- bounded by the +// candidate count, which the caller already bounds. +const stagedVisibilityCandidateScanBudget = 1 << 20 + +func scanVisibleCandidates(ctx context.Context, st store.MVCCStore, start, end []byte, limit int, ts uint64, reverse bool) ([]*store.KVPair, error) { + if limit <= 0 { + return []*store.KVPair{}, nil + } + if reverse { + kvs, err := st.ReverseScanAt(ctx, start, end, limit, ts) + return kvs, errors.WithStack(err) } - return scanRoutePage{ - kvs: kvs, - advanceKey: lastKVKey(raw), - full: len(raw) >= limit, - }, nil + kvs, err := st.ScanAt(ctx, start, end, limit, ts) + return kvs, errors.WithStack(err) } -func (s *ShardStore) scanRouteAtReverseProxyPage( +func (s *ShardStore) latestStagedVisibilityCandidates( ctx context.Context, + st store.MVCCStore, route distribution.Route, - g *ShardGroup, - start []byte, - end []byte, - limit int, + liveKVs []*store.KVPair, + stagedKVs []*store.KVPair, ts uint64, - readRouteVersion uint64, - routeStart []byte, - routeEnd []byte, -) (scanRoutePage, error) { - raw, err := s.proxyRawScanAt(ctx, g, start, end, limit, ts, true, route.GroupID, readRouteVersion, routeStart, routeEnd) +) (map[string]store.MVCCVersion, error) { + keys := stagedVisibilityCandidateKeys(liveKVs, stagedKVs) + stagedKeys := make([][]byte, 0, len(keys)) + for _, key := range keys { + stagedKeys = append(stagedKeys, distribution.MigrationStagedDataKey(route.MigrationJobID, key)) + } + // Staged before live: same promotion ordering as the point read. + stagedVersions, err := latestCandidateVersionsAt(ctx, st, stagedKeys, ts) if err != nil { - return scanRoutePage{}, err + return nil, err } - // Canonicalized here; see - // scanRouteAtDirectionWithReadFenceRouteFilterProxyPage for why a peer's - // page cannot be assumed canonical. - kvs, err := s.canonicalizeRedisWideColumnScanResults(ctx, filterTxnInternalKVs(raw), start, ts, readRouteVersion) + liveVersions, err := latestCandidateVersionsAt(ctx, st, keys, ts) if err != nil { - return scanRoutePage{}, err + return nil, err } - return scanRoutePage{ - kvs: kvs, - 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) + out := make(map[string]store.MVCCVersion, len(keys)) + for _, key := range keys { + live, liveOK := liveVersions[string(key)] + stagedKey := distribution.MigrationStagedDataKey(route.MigrationJobID, key) + staged, stagedOK := stagedVersions[string(stagedKey)] + if stagedOK { + staged.Key = bytes.Clone(key) + } + if winner, ok := newerMigrationVersion(live, liveOK, staged, stagedOK); ok { + out[string(key)] = winner + } + } + return out, nil } -func (s *ShardStore) scanRouteAtDirectionPhysicalLimit( - ctx context.Context, - route distribution.Route, - start []byte, - end []byte, - visibleLimit int, - physicalLimit int, - ts uint64, - reverse bool, -) ([]*store.KVPair, bool, error) { - g, ok := s.groupForID(route.GroupID) - if !ok || g == nil || g.Store == nil { - return nil, false, nil +func latestCandidateVersionsAt(ctx context.Context, st store.MVCCStore, keys [][]byte, ts uint64) (map[string]store.MVCCVersion, error) { + if len(keys) == 0 { + return map[string]store.MVCCVersion{}, nil } - - if engineForGroup(g) == nil { - kvs, limitReached, err := scanLocalPhysicalLimit(ctx, g.Store, start, end, visibleLimit, physicalLimit, ts, reverse) - if err != nil { - return nil, limitReached, errors.WithStack(err) + candidates := make(map[string]struct{}, len(keys)) + sortedKeys := make([][]byte, 0, len(keys)) + for _, key := range keys { + id := string(key) + if _, exists := candidates[id]; exists { + continue } - return filterTxnInternalKVs(kvs), limitReached, nil + candidates[id] = struct{}{} + sortedKeys = append(sortedKeys, key) } - - if isLinearizableRaftLeader(ctx, engineForGroup(g)) { - return s.scanRouteAtLeaderPhysicalLimit(ctx, g, start, end, visibleLimit, physicalLimit, ts, reverse) + sort.Slice(sortedKeys, func(i, j int) bool { + return bytes.Compare(sortedKeys[i], sortedKeys[j]) < 0 + }) + accepted := make(map[string]struct{}, len(sortedKeys)) + result, err := st.ExportVersions(ctx, store.ExportVersionsOptions{ + StartKey: sortedKeys[0], + EndKey: prefixScanEnd(sortedKeys[len(sortedKeys)-1]), + MaxCommitTSInclusive: ts, + MaxVersions: len(sortedKeys), + MaxBytes: ^uint64(0), + MaxScannedBytes: stagedVisibilityCandidateScanBudget, + AcceptKey: func(key []byte) bool { + _, ok := candidates[string(key)] + return ok + }, + AcceptVersion: func(key []byte, _ []byte) bool { + id := string(key) + if _, ok := accepted[id]; ok { + return false + } + accepted[id] = struct{}{} + return true + }, + }) + if err != nil { + return nil, errors.WithStack(err) } - - // RawScanAt cannot enforce physicalLimit, so report truncation and let - // callers fail closed instead of proxying an unbounded physical scan. - return nil, true, nil + out := make(map[string]store.MVCCVersion, len(sortedKeys)) + for _, version := range result.Versions { + out[string(version.Key)] = version + } + if result.Done { + // The export walked the whole enclosing range, so a candidate missing + // from the result has no visible version. + return out, nil + } + return probeRemainingCandidateVersionsAt(ctx, st, sortedKeys, out, ts) } -func scanLocalPhysicalLimit( +// probeRemainingCandidateVersionsAt resolves the candidates the bounded range +// export did not reach, one exact-key probe each. +func probeRemainingCandidateVersionsAt( ctx context.Context, st store.MVCCStore, - start []byte, - end []byte, - visibleLimit int, - physicalLimit int, + sortedKeys [][]byte, + out map[string]store.MVCCVersion, ts uint64, - reverse bool, -) ([]*store.KVPair, bool, error) { - scanner, ok := st.(physicalLimitedStore) - if !ok { - if reverse { - kvs, err := st.ReverseScanAt(ctx, start, end, visibleLimit, ts) - return kvs, false, errors.WithStack(err) +) (map[string]store.MVCCVersion, error) { + for _, key := range sortedKeys { + if _, ok := out[string(key)]; ok { + continue + } + version, found, err := latestMVCCVersionAt(ctx, st, key, ts) + if err != nil { + return nil, err + } + if found { + out[string(key)] = version } - kvs, err := st.ScanAt(ctx, start, end, visibleLimit, ts) - return kvs, false, errors.WithStack(err) } - return scanPhysicalLimitLocal(ctx, scanner, start, end, visibleLimit, physicalLimit, ts, reverse) + return out, nil } -func scanPhysicalLimitLocal( - ctx context.Context, - scanner physicalLimitedStore, - start []byte, - end []byte, - visibleLimit int, - physicalLimit int, - ts uint64, - reverse bool, -) ([]*store.KVPair, bool, error) { - if reverse { - kvs, limitReached, err := scanner.ReverseScanAtPhysicalLimit(ctx, start, end, visibleLimit, physicalLimit, ts) - return kvs, limitReached, errors.WithStack(err) +func stagedVisibilityCandidateKeys(liveKVs []*store.KVPair, stagedKVs []*store.KVPair) [][]byte { + seen := make(map[string][]byte, len(liveKVs)+len(stagedKVs)) + for _, kvp := range liveKVs { + if kvp == nil { + continue + } + if isMigrationStagedDataKey(kvp.Key) { + continue + } + seen[string(kvp.Key)] = bytes.Clone(kvp.Key) } - kvs, limitReached, err := scanner.ScanAtPhysicalLimit(ctx, start, end, visibleLimit, physicalLimit, ts) - return kvs, limitReached, errors.WithStack(err) + for _, kvp := range stagedKVs { + if kvp == nil { + continue + } + _, rawKey, ok := distribution.MigrationStagedDataKeyParts(kvp.Key) + if !ok { + continue + } + seen[string(rawKey)] = bytes.Clone(rawKey) + } + out := make([][]byte, 0, len(seen)) + for _, key := range seen { + out = append(out, key) + } + return out } -func (s *ShardStore) scanRouteLocal( - ctx context.Context, - g *ShardGroup, - start []byte, - end []byte, - limit int, - ts uint64, - reverse bool, -) ([]*store.KVPair, error) { - if reverse { - kvs, err := g.Store.ReverseScanAt(ctx, start, end, limit, ts) - return kvs, errors.WithStack(err) - } - kvs, err := g.Store.ScanAt(ctx, start, end, limit, ts) - return kvs, errors.WithStack(err) +func isMigrationStagedDataKey(key []byte) bool { + _, _, ok := distribution.MigrationStagedDataKeyParts(key) + return ok } -func (s *ShardStore) scanRouteAtLeaderPhysicalLimit( - ctx context.Context, - g *ShardGroup, - start []byte, - end []byte, - visibleLimit int, - physicalLimit int, - ts uint64, - reverse bool, -) ([]*store.KVPair, bool, error) { - kvs, limitReached, err := scanLocalPhysicalLimit(ctx, g.Store, start, end, visibleLimit, physicalLimit, ts, reverse) - if err != nil { - return nil, limitReached, errors.WithStack(err) +func stagedVisibilityScanBounds(jobID uint64, start []byte, end []byte) ([]byte, []byte) { + prefix := distribution.MigrationStagedDataKeyPrefix(jobID) + scanStart := prefix + if len(start) > 0 { + scanStart = distribution.MigrationStagedDataKey(jobID, start) } - lockStart, lockEnd := scanLockBoundsForKVsDirection(kvs, start, end, visibleLimit, reverse) - lockKVs, err := scanTxnLockRangeAt(ctx, g, lockStart, lockEnd, ts, visibleLimit) - if err != nil { - return nil, limitReached, err + scanEnd := prefixScanEnd(prefix) + if len(end) > 0 { + scanEnd = distribution.MigrationStagedDataKey(jobID, end) } - resolved, err := s.resolveScanLocks(ctx, g, kvs, lockKVs, ts) - return resolved, limitReached, err + return scanStart, scanEnd +} + +func visibleLogicalKVs(versions map[string]store.MVCCVersion, ts uint64, reverse bool) []*store.KVPair { + out := make([]*store.KVPair, 0, len(versions)) + for _, version := range versions { + if !migrationVersionVisible(version, ts) { + continue + } + out = append(out, &store.KVPair{ + Key: bytes.Clone(version.Key), + Value: bytes.Clone(version.Value), + }) + } + sort.Slice(out, func(i, j int) bool { + cmp := bytes.Compare(out[i].Key, out[j].Key) + if reverse { + return cmp > 0 + } + return cmp < 0 + }) + return out } func (s *ShardStore) scanRouteAtLeaderRouteFilter( ctx context.Context, g *ShardGroup, + route distribution.Route, start []byte, end []byte, limit int, @@ -2371,9 +3786,12 @@ func (s *ShardStore) scanRouteAtLeaderRouteFilter( kvs []*store.KVPair err error ) - if reverse { + switch { + case routeHasStagedVisibility(route): + kvs, err = s.scanRouteWithStagedVisibility(ctx, g, route, start, end, limit, ts, reverse) + case reverse: kvs, err = g.Store.ReverseScanAt(ctx, start, end, limit, ts) - } else { + default: kvs, err = g.Store.ScanAt(ctx, start, end, limit, ts) } if err != nil { @@ -2385,7 +3803,7 @@ func (s *ShardStore) scanRouteAtLeaderRouteFilter( if err != nil { return nil, nil, err } - resolved, err := s.resolveScanLocks(ctx, g, filteredKVs, lockKVs, ts) + resolved, err := s.resolveScanLocks(ctx, g, route, filteredKVs, lockKVs, ts) if err == nil { sort.Slice(resolved, func(i, j int) bool { if reverse { @@ -2472,21 +3890,31 @@ func scanUserKey(kvp *store.KVPair) ([]byte, bool) { if kvp == nil || kvp.Key == nil { return nil, false } + if isMigrationStagedDataKey(kvp.Key) { + return nil, false + } if !isTxnInternalKey(kvp.Key) { return kvp.Key, true } return txnUserKeyFromLockKey(kvp.Key) } -func mergeAndTrimScanResults(out []*store.KVPair, kvs []*store.KVPair, limit int) []*store.KVPair { +func mergeAndTrimScanResultsWithOptions(out []*store.KVPair, kvs []*store.KVPair, limit int, reverse bool, dedupeByKey bool) []*store.KVPair { if len(kvs) == 0 { return out } - out = append(out, kvs...) - sort.SliceStable(out, func(i, j int) bool { - return bytes.Compare(out[i].Key, out[j].Key) < 0 + if dedupeByKey { + out = appendReplacingKVsByKey(out, kvs) + } else { + out = append(out, kvs...) + } + sort.Slice(out, func(i, j int) bool { + cmp := bytes.Compare(out[i].Key, out[j].Key) + if reverse { + return cmp > 0 + } + return cmp < 0 }) - out = dedupeSortedScanResults(out) if len(out) <= limit { return out } @@ -2539,19 +3967,30 @@ func mergeAndTrimScanKeys(out [][]byte, keys [][]byte, limit int) [][]byte { } func mergeAndTrimReverseScanResults(out []*store.KVPair, kvs []*store.KVPair, limit int) []*store.KVPair { - if len(kvs) == 0 { - return out + return mergeAndTrimScanResultsWithOptions(out, kvs, limit, true, false) +} + +func appendReplacingKVsByKey(out []*store.KVPair, kvs []*store.KVPair) []*store.KVPair { + indexByKey := make(map[string]int, len(out)+len(kvs)) + for i, kvp := range out { + if kvp == nil { + continue + } + indexByKey[string(kvp.Key)] = i } - out = append(out, kvs...) - sort.SliceStable(out, func(i, j int) bool { - return bytes.Compare(out[i].Key, out[j].Key) > 0 - }) - out = dedupeSortedScanResults(out) - if len(out) <= limit { - return out + for _, kvp := range kvs { + if kvp == nil { + continue + } + key := string(kvp.Key) + if idx, ok := indexByKey[key]; ok { + out[idx] = kvp + continue + } + indexByKey[key] = len(out) + out = append(out, kvp) } - clear(out[limit:]) - return out[:limit] + return out } func kvPairsFromKeys(keys [][]byte) []*store.KVPair { @@ -2586,13 +4025,13 @@ func lastKVKey(kvs []*store.KVPair) []byte { return nil } -func filterTxnInternalKeys(keys [][]byte) [][]byte { +func filterScanInternalKeys(keys [][]byte) [][]byte { if len(keys) == 0 { return keys } out := make([][]byte, 0, len(keys)) for _, key := range keys { - if key == nil || isTxnInternalKey(key) { + if key == nil || isScanInternalKey(key) { continue } out = append(out, key) @@ -2629,80 +4068,119 @@ func clampScanEnd(end []byte, routeEnd []byte) []byte { } func (s *ShardStore) PutAt(ctx context.Context, key []byte, value []byte, commitTS uint64, expireAt uint64) error { - g, err := s.writeGroupForKey(key, commitTS) - if err != nil { + route, g, ok := s.routeAndGroupForKey(key) + if !ok || g.Store == nil { + return store.ErrNotSupported + } + if err := ensureRouteWriteTimestampFloor(route, key, commitTS); err != nil { + return err + } + if err := s.ensureS3BucketAuxiliaryWriteTimestampFloor(key, commitTS); err != nil { return err } return errors.WithStack(g.Store.PutAt(ctx, key, value, commitTS, expireAt)) } func (s *ShardStore) DeleteAt(ctx context.Context, key []byte, commitTS uint64) error { - g, err := s.writeGroupForKey(key, commitTS) - if err != nil { + route, g, ok := s.routeAndGroupForKey(key) + if !ok || g.Store == nil { + return store.ErrNotSupported + } + if err := ensureRouteWriteTimestampFloor(route, key, commitTS); err != nil { + return err + } + if err := s.ensureS3BucketAuxiliaryWriteTimestampFloor(key, commitTS); err != nil { return err } return errors.WithStack(g.Store.DeleteAt(ctx, key, commitTS)) } func (s *ShardStore) PutWithTTLAt(ctx context.Context, key []byte, value []byte, commitTS uint64, expireAt uint64) error { - g, err := s.writeGroupForKey(key, commitTS) - if err != nil { + route, g, ok := s.routeAndGroupForKey(key) + if !ok || g.Store == nil { + return store.ErrNotSupported + } + if err := ensureRouteWriteTimestampFloor(route, key, commitTS); err != nil { + return err + } + if err := s.ensureS3BucketAuxiliaryWriteTimestampFloor(key, commitTS); err != nil { return err } return errors.WithStack(g.Store.PutWithTTLAt(ctx, key, value, commitTS, expireAt)) } func (s *ShardStore) ExpireAt(ctx context.Context, key []byte, expireAt uint64, commitTS uint64) error { - g, err := s.writeGroupForKey(key, commitTS) - if err != nil { + route, g, ok := s.routeAndGroupForKey(key) + if !ok || g.Store == nil { + return store.ErrNotSupported + } + if err := ensureRouteWriteTimestampFloor(route, key, commitTS); err != nil { + return err + } + if err := s.ensureS3BucketAuxiliaryWriteTimestampFloor(key, commitTS); err != nil { return err } + if routeHasStagedVisibility(route) { + return s.expireStagedVisibleAt(ctx, g, route, key, expireAt, commitTS) + } return errors.WithStack(g.Store.ExpireAt(ctx, key, expireAt, commitTS)) } -func (s *ShardStore) LatestCommitTS(ctx context.Context, key []byte) (uint64, bool, error) { - return s.LatestCommitTSWithReadFence(ctx, key, 0) +// expireStagedVisibleAt applies an expiration to the version visible through +// the staged/live merge. +// +// Between cutover and promotion a key can be visible through its staged alias +// while the live key is missing, older, or tombstoned by a newer staged delete. +// Resolve the staged/live winner before writing the expiration, then record the +// result as a live MVCC version, which is where every other post-cutover write +// goes. +func (s *ShardStore) expireStagedVisibleAt( + ctx context.Context, + g *ShardGroup, + route distribution.Route, + key []byte, + expireAt uint64, + commitTS uint64, +) error { + // Staged before live. Reading live first lets a promotion land in between + // and leave the staged probe empty, so the older live payload would be + // copied into a fresh version at commitTS -- a permanent rollback of the + // value rather than a TTL on the current one. + stagedKey := distribution.MigrationStagedDataKey(route.MigrationJobID, key) + staged, stagedOK, err := latestMVCCVersionAt(ctx, g.Store, stagedKey, commitTS) + if err != nil { + return err + } + live, liveOK, err := latestMVCCVersionAt(ctx, g.Store, key, commitTS) + if err != nil { + return err + } + if stagedOK { + staged.Key = bytes.Clone(key) + } + winner, ok := newerMigrationVersion(live, liveOK, staged, stagedOK) + if !ok || !migrationVersionVisible(winner, commitTS) { + return store.ErrKeyNotFound + } + return errors.WithStack(g.Store.PutWithTTLAt(ctx, key, winner.Value, commitTS, expireAt)) } -func (s *ShardStore) LatestCommitTSGroupWithReadFence(ctx context.Context, key []byte, groupID uint64, readRouteVersion uint64) (uint64, bool, error) { - if err := s.awaitReadRouteVersion(ctx, readRouteVersion); err != nil { - return 0, false, err - } - return s.latestCommitTSForRoute(ctx, distribution.Route{GroupID: groupID}, key, readRouteVersion) +func (s *ShardStore) LatestCommitTS(ctx context.Context, key []byte) (uint64, bool, error) { + return s.LatestCommitTSWithReadFence(ctx, key, 0) } func (s *ShardStore) LatestCommitTSWithReadFence(ctx context.Context, key []byte, readRouteVersion uint64) (uint64, bool, error) { if err := s.awaitReadRouteVersion(ctx, readRouteVersion); err != nil { return 0, false, err } - routes, routeVersion := s.pointReadRoutesWithVersion(key) + route, g, routeVersion, ok := s.routeAndGroupForKeyWithVersion(key) readRouteVersion = max(readRouteVersion, routeVersion) - if len(routes) == 0 { - return 0, false, nil - } - var latest uint64 - found := false - for _, route := range routes { - ts, exists, err := s.latestCommitTSForRoute(ctx, route, key, readRouteVersion) - if err != nil { - return 0, false, err - } - if exists && (!found || ts > latest) { - latest = ts - found = true - } - } - return latest, found, nil -} - -func (s *ShardStore) latestCommitTSForRoute(ctx context.Context, route distribution.Route, key []byte, readRouteVersion uint64) (uint64, bool, error) { - g, ok := s.groupForID(route.GroupID) if !ok || g.Store == nil { return 0, false, nil } if engineForGroup(g) == nil { - ts, exists, err := g.Store.LatestCommitTS(ctx, key) + ts, exists, err := s.localLatestCommitTS(ctx, g, route, key) if err != nil { return 0, false, errors.WithStack(err) } @@ -2714,7 +4192,7 @@ func (s *ShardStore) latestCommitTSForRoute(ctx context.Context, route distribut // round-trip (same rationale as isLinearizableRaftLeader). if engine := engineForGroup(g); isLeaderEngine(engine) { if _, err := leaseReadEngineCtx(ctx, engine); err == nil { - ts, exists, err := g.Store.LatestCommitTS(ctx, key) + ts, exists, err := s.localLatestCommitTS(ctx, g, route, key) if err != nil { return 0, false, errors.WithStack(err) } @@ -2722,10 +4200,118 @@ func (s *ShardStore) latestCommitTSForRoute(ctx context.Context, route distribut } } - return s.proxyLatestCommitTS(ctx, g, route.GroupID, key, readRouteVersion) + return s.proxyLatestCommitTS(ctx, g, key, readRouteVersion) +} + +func (s *ShardStore) LatestCommitTSGroupWithReadFence(ctx context.Context, key []byte, groupID uint64, readRouteVersion uint64) (uint64, bool, error) { + if groupID == 0 { + return s.LatestCommitTSWithReadFence(ctx, key, readRouteVersion) + } + if err := s.awaitReadRouteVersion(ctx, readRouteVersion); err != nil { + return 0, false, err + } + g, ok := s.groupForID(groupID) + if !ok || g.Store == nil { + return 0, false, nil + } + + if engineForGroup(g) == nil { + ts, exists, err := g.Store.LatestCommitTS(ctx, key) + return ts, exists, errors.WithStack(err) + } + + if engine := engineForGroup(g); isLeaderEngine(engine) { + if _, err := leaseReadEngineCtx(ctx, engine); err == nil { + ts, exists, err := g.Store.LatestCommitTS(ctx, key) + return ts, exists, errors.WithStack(err) + } + } + + return s.proxyLatestCommitTSGroup(ctx, g, key, groupID, readRouteVersion) +} + +func (s *ShardStore) VersionsExistAtOrBeforeGroupWithReadFence(ctx context.Context, keys [][]byte, groupID uint64, ts uint64, readRouteVersion uint64) ([]bool, bool, error) { + if groupID == 0 { + return s.versionsExistAtOrBeforeWithReadFence(ctx, keys, ts, readRouteVersion) + } + if err := s.awaitReadRouteVersion(ctx, readRouteVersion); err != nil { + return nil, false, err + } + return s.routeHasVersionsAtOrBefore(ctx, distribution.Route{GroupID: groupID}, keys, ts) +} + +func (s *ShardStore) versionsExistAtOrBeforeWithReadFence(ctx context.Context, keys [][]byte, ts uint64, readRouteVersion uint64) ([]bool, bool, error) { + out := make([]bool, len(keys)) + if err := s.awaitReadRouteVersion(ctx, readRouteVersion); err != nil { + return nil, false, err + } + for i, key := range keys { + route, _, _, ok := s.routeAndGroupForKeyWithVersion(key) + if !ok { + continue + } + exists, answered, err := s.routeHasVersionsAtOrBefore(ctx, route, [][]byte{key}, ts) + if err != nil || !answered { + return out, answered, err + } + out[i] = exists[0] + } + return out, true, nil +} + +func (s *ShardStore) localLatestCommitTS(ctx context.Context, g *ShardGroup, route distribution.Route, key []byte) (uint64, bool, error) { + if !routeHasStagedVisibility(route) { + liveTS, liveExists, err := g.Store.LatestCommitTS(ctx, key) + return liveTS, liveExists, errors.WithStack(err) + } + // Staged before live: same promotion ordering as every other pair of + // probes across the two namespaces. + stagedTS, stagedExists, err := g.Store.LatestCommitTS(ctx, distribution.MigrationStagedDataKey(route.MigrationJobID, key)) + if err != nil { + return 0, false, errors.WithStack(err) + } + liveTS, liveExists, err := g.Store.LatestCommitTS(ctx, key) + if err != nil { + return 0, false, errors.WithStack(err) + } + switch { + case !liveExists: + return stagedTS, stagedExists, nil + case !stagedExists: + return liveTS, true, nil + case stagedTS >= liveTS: + return stagedTS, true, nil + default: + return liveTS, true, nil + } +} + +func (s *ShardStore) proxyLatestCommitTS(ctx context.Context, g *ShardGroup, key []byte, readRouteVersion uint64) (uint64, bool, error) { + engine := engineForGroup(g) + if engine == nil { + return 0, false, nil + } + addr := leaderAddrFromEngine(engine) + if addr == "" { + return 0, false, errors.WithStack(ErrLeaderNotFound) + } + + conn, err := s.connCache.ConnFor(addr) + if err != nil { + return 0, false, err + } + + ctx, cancel := context.WithTimeout(ctx, proxyForwardTimeout) + defer cancel() + cli := pb.NewRawKVClient(conn) + resp, err := cli.RawLatestCommitTS(ctx, &pb.RawLatestCommitTSRequest{Key: key, ReadRouteVersion: readRouteVersion}) + if err != nil { + return 0, false, errors.WithStack(err) + } + return resp.Ts, resp.Exists, nil } -func (s *ShardStore) proxyLatestCommitTS(ctx context.Context, g *ShardGroup, groupID uint64, key []byte, readRouteVersion uint64) (uint64, bool, error) { +func (s *ShardStore) proxyLatestCommitTSGroup(ctx context.Context, g *ShardGroup, key []byte, groupID uint64, readRouteVersion uint64) (uint64, bool, error) { engine := engineForGroup(g) if engine == nil { return 0, false, nil @@ -2743,7 +4329,11 @@ func (s *ShardStore) proxyLatestCommitTS(ctx context.Context, g *ShardGroup, gro ctx, cancel := context.WithTimeout(ctx, proxyForwardTimeout) defer cancel() cli := pb.NewRawKVClient(conn) - resp, err := cli.RawLatestCommitTS(ctx, &pb.RawLatestCommitTSRequest{Key: key, ReadRouteVersion: readRouteVersion, GroupId: groupID}) + resp, err := cli.RawLatestCommitTS(ctx, &pb.RawLatestCommitTSRequest{ + Key: key, + GroupId: groupID, + ReadRouteVersion: readRouteVersion, + }) if err != nil { return 0, false, errors.WithStack(err) } @@ -2850,7 +4440,7 @@ func newScanLockPlan(size int) *scanLockPlan { } } -func (s *ShardStore) resolveScanLocks(ctx context.Context, g *ShardGroup, kvs []*store.KVPair, lockKVs []*store.KVPair, ts uint64) ([]*store.KVPair, error) { +func (s *ShardStore) resolveScanLocks(ctx context.Context, g *ShardGroup, route distribution.Route, kvs []*store.KVPair, lockKVs []*store.KVPair, ts uint64) ([]*store.KVPair, error) { if len(kvs) == 0 && len(lockKVs) == 0 { return kvs, nil } @@ -2865,7 +4455,7 @@ func (s *ShardStore) resolveScanLocks(ctx context.Context, g *ShardGroup, kvs [] if err := applyScanLockResolutions(ctx, g, plan); err != nil { return nil, err } - return s.materializeScanLockResults(ctx, g, ts, plan.items) + return s.materializeScanLockResults(ctx, g, route, ts, plan.items) } func (s *ShardStore) planScanLockResolutions(ctx context.Context, g *ShardGroup, kvs []*store.KVPair, lockKVs []*store.KVPair, ts uint64) (*scanLockPlan, error) { @@ -2904,7 +4494,7 @@ func (s *ShardStore) planScanLockFromLockKVP(ctx context.Context, plan *scanLock } func (s *ShardStore) planScanLockItem(ctx context.Context, g *ShardGroup, ts uint64, plan *scanLockPlan, kvp *store.KVPair) error { - if kvp == nil || isTxnInternalKey(kvp.Key) { + if kvp == nil || isScanInternalKey(kvp.Key) { plan.items = append(plan.items, scanItem{skip: true}) return nil } @@ -3064,12 +4654,6 @@ func scanTxnLockPagesAtWithRouteFilter(ctx context.Context, st store.MVCCStore, out := make([]*store.KVPair, 0, min(limit, lockPageLimit)) cursor := start scanned := 0 - // The raw budget is deliberately larger than the accepted-lock budget. A - // route-bound scan walks a physical lock range that also holds locks for - // other logical routes, and counting those against `limit` let a page of - // unrelated locks fail the scan with ErrTxnLocked before a single relevant - // lock was seen. Only accepted locks signal a real conflict; this cap just - // keeps the paging bounded. rawLimit := routeFilteredRawLockScanLimit(limit) for { lockKVs, nextCursor, done, err := scanTxnLockPageAt(ctx, st, cursor, end, ts) @@ -3096,8 +4680,6 @@ func scanTxnLockPagesAtWithRouteFilter(ctx context.Context, st store.MVCCStore, } } -// routeFilteredRawLockScanLimit caps how many physical locks a route-bound scan -// may page through while looking for locks that belong to its own route. func routeFilteredRawLockScanLimit(limit int) int { if limit <= 0 { return lockPageLimit @@ -3184,7 +4766,7 @@ func applyScanLockResolutions(ctx context.Context, g *ShardGroup, plan *scanLock return nil } -func (s *ShardStore) materializeScanLockResults(ctx context.Context, g *ShardGroup, ts uint64, items []scanItem) ([]*store.KVPair, error) { +func (s *ShardStore) materializeScanLockResults(ctx context.Context, g *ShardGroup, route distribution.Route, ts uint64, items []scanItem) ([]*store.KVPair, error) { out := make([]*store.KVPair, 0, len(items)) for _, item := range items { if item.skip { @@ -3194,7 +4776,7 @@ func (s *ShardStore) materializeScanLockResults(ctx context.Context, g *ShardGro out = append(out, item.kvp) continue } - v, err := s.localGetAt(ctx, g, item.kvp.Key, ts) + v, err := s.localGetAt(ctx, g, route, item.kvp.Key, ts) if err != nil { if errors.Is(err, store.ErrKeyNotFound) { continue @@ -3206,7 +4788,7 @@ func (s *ShardStore) materializeScanLockResults(ctx context.Context, g *ShardGro return out, nil } -func filterTxnInternalKVs(kvs []*store.KVPair) []*store.KVPair { +func filterScanInternalKVs(kvs []*store.KVPair) []*store.KVPair { if len(kvs) == 0 { return kvs } @@ -3215,7 +4797,7 @@ func filterTxnInternalKVs(kvs []*store.KVPair) []*store.KVPair { if kvp == nil { continue } - if isTxnInternalKey(kvp.Key) { + if isScanInternalKey(kvp.Key) { continue } out = append(out, kvp) @@ -3223,6 +4805,22 @@ func filterTxnInternalKVs(kvs []*store.KVPair) []*store.KVPair { return out } +func isScanInternalKey(key []byte) bool { + return isTxnInternalKey(key) || isMigrationStagedDataKey(key) +} + +func markScanRouteGroup(kvs []*store.KVPair, groupID uint64, mark bool) []*store.KVPair { + if !mark || groupID == 0 { + return kvs + } + for _, kvp := range kvs { + if kvp != nil { + kvp.RouteGroupID = groupID + } + } + return kvs +} + type txnStatus int const ( @@ -3399,20 +4997,30 @@ func cleanupTSWithNow(startTS, now uint64) uint64 { // All mutations must belong to the same shard. Cross-shard mutation batches are // not supported. func (s *ShardStore) ApplyMutations(ctx context.Context, mutations []*store.KVPairMutation, readKeys [][]byte, startTS, commitTS uint64) error { - group, err := s.resolveSingleShardGroup(mutations, commitTS) + group, err := s.resolveSingleShardGroup(mutations) if err != nil || group == nil { return err } + if err := s.ensureMutationWriteTimestampFloors(mutations, commitTS); err != nil { + return err + } + readKeys = s.readKeysWithStagedVisibilityAliases(group, readKeys) + readKeys = s.readKeysWithStagedVisibilityMutationAliases(group, readKeys, mutations) return errors.WithStack(group.Store.ApplyMutations(ctx, mutations, readKeys, startTS, commitTS)) } // ApplyMutationsRaft is the raft-apply variant; see store.MVCCStore for the // durability contract. Only the FSM may call this method. func (s *ShardStore) ApplyMutationsRaft(ctx context.Context, mutations []*store.KVPairMutation, readKeys [][]byte, startTS, commitTS uint64) error { - group, err := s.resolveSingleShardGroup(mutations, commitTS) + group, err := s.resolveSingleShardGroup(mutations) if err != nil || group == nil { return err } + if err := s.ensureMutationWriteTimestampFloors(mutations, commitTS); err != nil { + return err + } + readKeys = s.readKeysWithStagedVisibilityAliases(group, readKeys) + readKeys = s.readKeysWithStagedVisibilityMutationAliases(group, readKeys, mutations) return errors.WithStack(group.Store.ApplyMutationsRaft(ctx, mutations, readKeys, startTS, commitTS)) } @@ -3420,29 +5028,23 @@ func (s *ShardStore) ApplyMutationsRaft(ctx context.Context, mutations []*store. // appliedIndex through to the single owning shard so the leaf can // bundle metaAppliedIndex with the mutation. See PR #910 design §2. func (s *ShardStore) ApplyMutationsRaftAt(ctx context.Context, mutations []*store.KVPairMutation, readKeys [][]byte, startTS, commitTS, appliedIndex uint64) error { - group, err := s.resolveSingleShardGroup(mutations, commitTS) + group, err := s.resolveSingleShardGroup(mutations) if err != nil || group == nil { return err } + if err := s.ensureMutationWriteTimestampFloors(mutations, commitTS); err != nil { + return err + } + readKeys = s.readKeysWithStagedVisibilityAliases(group, readKeys) + readKeys = s.readKeysWithStagedVisibilityMutationAliases(group, readKeys, mutations) return errors.WithStack(group.Store.ApplyMutationsRaftAt(ctx, mutations, readKeys, startTS, commitTS, appliedIndex)) } -func (s *ShardStore) writeGroupForKey(key []byte, commitTS uint64) (*ShardGroup, error) { - route, ok := s.engine.GetRoute(routeKey(key)) - if !ok { - return nil, store.ErrNotSupported - } - if err := ensureRouteWriteAllowed(route, key, commitTS); err != nil { - return nil, err - } - if err := ensureLogicalRouteWriteAllowed(s.engine, key, commitTS); err != nil { - return nil, err - } - g, ok := s.groupForID(route.GroupID) - if !ok || g.Store == nil { - return nil, store.ErrNotSupported +func ensureRouteWriteTimestampFloor(route distribution.Route, key []byte, commitTS uint64) error { + if route.MinWriteTSExclusive == 0 || commitTS == 0 || commitTS > route.MinWriteTSExclusive { + return nil } - return g, nil + return errors.Wrapf(ErrRouteWriteTimestampTooLow, "key %q routeKey %q commit_ts=%d floor=%d", key, routeKey(key), commitTS, route.MinWriteTSExclusive) } // logicalRouteFloorKey returns the logical Redis user key whose route floor @@ -3460,12 +5062,7 @@ func logicalRouteFloorKey(key []byte) []byte { // raw key, but every other part of the system treats the user key as their // owner: route-bound scans classify them with routeFilterKey, and prefix-write // floors go through routesForRedisListPrefixWrite. Checking only the raw-key -// route therefore let a fenced user key keep accepting its auxiliary writes, -// which is the MinWriteTSExclusive bypass this closes. -// -// This deliberately does not change placement. Making routeKey apply the same -// decoders would relocate existing rows to the user key's group, which is a -// data-migration decision that needs its own design proposal. +// route therefore lets a fenced user key keep accepting auxiliary writes. func ensureLogicalRouteWriteAllowed(engine *distribution.Engine, key []byte, commitTS uint64) error { if engine == nil { return nil @@ -3485,22 +5082,14 @@ func ensureRouteWriteAllowed(route distribution.Route, key []byte, commitTS uint if route.MinWriteTSExclusive == 0 || commitTS > route.MinWriteTSExclusive { return nil } - return errors.Wrapf( - store.NewWriteConflictError(key), - "route min_write_ts_exclusive=%d rejects commit_ts=%d", - route.MinWriteTSExclusive, - commitTS, - ) + return routeWriteTimestampFloorError(route, key, commitTS) } -func (s *ShardStore) ensurePrefixWriteAllowed(prefix []byte, commitTS uint64) error { - routes := s.routesForPrefixWrite(prefix) - for _, route := range routes { - if err := ensureRouteWriteAllowed(route, prefix, commitTS); err != nil { - return err - } - } - return nil +func routeWriteTimestampFloorError(route distribution.Route, key []byte, commitTS uint64) error { + return errors.Join( + errors.Wrapf(ErrRouteWriteTimestampTooLow, "key %q routeKey %q commit_ts=%d floor=%d", key, routeKey(key), commitTS, route.MinWriteTSExclusive), + store.NewWriteConflictError(key), + ) } func (s *ShardStore) routesForPrefixWrite(prefix []byte) []distribution.Route { @@ -3522,16 +5111,6 @@ func (s *ShardStore) routesForPrefixWrite(prefix []byte) []distribution.Route { return s.engine.GetIntersectingRoutes(prefix, prefixScanEnd(prefix)) } -// S3 object keys route through !s3route|, never -// through the raw !s3|obj|head| / !s3|upload| / !s3|blob| / !s3|chunkref| / -// !s3|gc|upload| intervals. bucketDeleteSafetyNetElems DEL_PREFIXes six of those -// raw families per bucket delete, and checked against the raw interval those six -// miss the object routes' own MinWriteTSExclusive floors: a cleanup timestamp at -// or below a floor is admitted, and its tombstones then sit hidden behind -// migrated object versions -- data retained for a bucket reported as deleted. -// -// Project the prefix onto the object route range for the same bucket generation. -// This is a range, not one route: a bucket's objects can span several routes. func (s *ShardStore) routesForS3PrefixWrite(prefix []byte) ([]distribution.Route, bool) { routeStart, ok := s3keys.BucketScopedRoutePrefix(prefix) if !ok { @@ -3540,23 +5119,9 @@ func (s *ShardStore) routesForS3PrefixWrite(prefix []byte) ([]distribution.Route return s.engine.GetIntersectingRoutes(routeStart, prefixScanEnd(routeStart)), true } -// DynamoDB rows route through their logical table key -// (!ddb|route|table|), never through the raw !ddb|item| / !ddb|gsi| -// interval, so a DEL_PREFIX checked against that raw interval can miss the -// table's own migration write floor and install tombstones below already -// migrated versions. Project the prefix onto the same table route the rows -// underneath it use. -// -// Only the item and GSI families are projected. Their table segment is -// terminated by '|', so a prefix that reaches the terminator covers exactly one -// table. The !ddb|meta| families have no terminator in the prefix, so -// "!ddb|meta|table|foo" would also cover table "foobar" and cannot be collapsed -// to a single route. func (s *ShardStore) routesForDynamoPrefixWrite(prefix []byte) ([]distribution.Route, bool) { for _, family := range dynamoTablePrefixWriteFamilies { if bytes.HasPrefix(family, prefix) { - // The prefix stops at or inside the family marker itself, so it - // spans every table and has no single logical key to project onto. return s.engine.GetIntersectingRoutes(nil, nil), true } if !bytes.HasPrefix(prefix, family) { @@ -3564,8 +5129,6 @@ func (s *ShardStore) routesForDynamoPrefixWrite(prefix []byte) ([]distribution.R } route := dynamoRouteKey(prefix) if route == nil { - // A partial table segment: the prefix can still match more than one - // table, so every table route stays a candidate. return s.engine.GetIntersectingRoutes(nil, nil), true } return s.routesForLogicalPrefixWriteRange(route, nil, true), true @@ -3638,44 +5201,125 @@ func (s *ShardStore) routesForLogicalPrefixWriteRange(routeStart []byte, routeEn return s.engine.GetIntersectingRoutes(routeStart, routeEnd) } -// resolveSingleShardGroup returns the shard group that owns every mutation in -// the batch, or an error if the batch is cross-shard, below a route write -// timestamp floor, or references an unknown group. A nil group with nil error -// means "empty batch; caller should no-op". -func (s *ShardStore) resolveSingleShardGroup(mutations []*store.KVPairMutation, commitTS uint64) (*ShardGroup, error) { - if len(mutations) == 0 { - return nil, nil +func (s *ShardStore) ensureMutationWriteTimestampFloors(mutations []*store.KVPairMutation, commitTS uint64) error { + if commitTS == 0 { + return nil } - var firstGroup *ShardGroup - for _, mutation := range mutations { - g, err := s.writeGroupForMutation(mutation, commitTS) - if err != nil { - return nil, err + for _, mut := range mutations { + if mut == nil || len(mut.Key) == 0 || isTxnInternalKey(mut.Key) { + continue + } + route, _, ok := s.routeAndGroupForKey(mut.Key) + if !ok { + return store.ErrNotSupported + } + if err := ensureRouteWriteTimestampFloor(route, mut.Key, commitTS); err != nil { + return err } - if g == nil { + if err := s.ensureS3BucketAuxiliaryWriteTimestampFloor(mut.Key, commitTS); err != nil { + return err + } + } + return nil +} + +func (s *ShardStore) ensureS3BucketAuxiliaryWriteTimestampFloor(key []byte, commitTS uint64) error { + if s == nil || s.engine == nil || commitTS == 0 { + return nil + } + start, end, ok := s3BucketAuxiliaryRouteRange(key) + if !ok { + return nil + } + for _, route := range s.engine.GetIntersectingRoutes(start, end) { + if route.MinWriteTSExclusive != 0 && commitTS <= route.MinWriteTSExclusive { + return errors.Wrapf(ErrRouteWriteTimestampTooLow, "key %q route range [%q,%q) commit_ts=%d floor=%d", key, start, end, commitTS, route.MinWriteTSExclusive) + } + } + return nil +} + +func (s *ShardStore) readKeysWithStagedVisibilityAliases(group *ShardGroup, readKeys [][]byte) [][]byte { + if len(readKeys) == 0 { + return readKeys + } + out := readKeys + copied := false + for _, key := range readKeys { + alias, ok := s.stagedVisibilityReadKeyAlias(group, key) + if !ok { continue } - if firstGroup == nil { - firstGroup = g + if !copied { + out = append([][]byte(nil), readKeys...) + copied = true + } + out = append(out, alias) + } + return out +} + +func (s *ShardStore) readKeysWithStagedVisibilityMutationAliases(group *ShardGroup, readKeys [][]byte, mutations []*store.KVPairMutation) [][]byte { + out := readKeys + copied := false + for _, mut := range mutations { + if mut == nil { continue } - if g != firstGroup { - return nil, errors.WithStack(ErrCrossShardMutationBatchNotSupported) + alias, ok := s.stagedVisibilityReadKeyAlias(group, mut.Key) + if !ok { + continue + } + if !copied { + out = append([][]byte(nil), readKeys...) + copied = true } + out = append(out, alias) } - return firstGroup, nil + return out +} + +func (s *ShardStore) stagedVisibilityReadKeyAlias(group *ShardGroup, key []byte) ([]byte, bool) { + if s == nil || s.engine == nil || group == nil || len(key) == 0 { + return nil, false + } + if _, _, ok := distribution.MigrationStagedDataKeyParts(key); ok { + return nil, false + } + route, g, ok := s.routeAndGroupForKey(key) + if !ok || g != group || !routeHasStagedVisibility(route) { + return nil, false + } + return distribution.MigrationStagedDataKey(route.MigrationJobID, key), true } -func (s *ShardStore) writeGroupForMutation(mutation *store.KVPairMutation, commitTS uint64) (*ShardGroup, error) { - if mutation == nil { +// resolveSingleShardGroup returns the shard group that owns every +// mutation in the batch, or an error if the batch is cross-shard or +// references an unknown group. A nil group with nil error means "empty +// batch — caller should no-op". +func (s *ShardStore) resolveSingleShardGroup(mutations []*store.KVPairMutation) (*ShardGroup, error) { + if len(mutations) == 0 { return nil, nil } - return s.writeGroupForKey(mutation.Key, commitTS) + firstGroup, ok := s.groupForKey(mutations[0].Key) + if !ok || firstGroup == nil || firstGroup.Store == nil { + return nil, store.ErrNotSupported + } + for i := 1; i < len(mutations); i++ { + g, ok := s.groupForKey(mutations[i].Key) + if !ok || g == nil || g.Store == nil { + return nil, store.ErrNotSupported + } + if g != firstGroup { + return nil, errors.WithStack(ErrCrossShardMutationBatchNotSupported) + } + } + return firstGroup, nil } // DeletePrefixAt applies a prefix delete to every shard in the store. func (s *ShardStore) DeletePrefixAt(ctx context.Context, prefix []byte, excludePrefix []byte, commitTS uint64) error { - if err := s.ensurePrefixWriteAllowed(prefix, commitTS); err != nil { + if err := s.ensurePrefixWriteTimestampFloors(prefix, commitTS); err != nil { return err } for _, g := range s.groups { @@ -3686,12 +5330,67 @@ func (s *ShardStore) DeletePrefixAt(ctx context.Context, prefix []byte, excludeP return errors.WithStack(err) } } + for _, del := range s.stagedVisibilityPrefixDeletes(prefix, excludePrefix) { + if err := del.group.Store.DeletePrefixAt(ctx, del.prefix, del.excludePrefix, commitTS); err != nil { + return errors.WithStack(err) + } + } + return nil +} + +type stagedVisibilityPrefixDelete struct { + group *ShardGroup + prefix []byte + excludePrefix []byte +} + +func (s *ShardStore) stagedVisibilityPrefixDeletes(prefix []byte, excludePrefix []byte) []stagedVisibilityPrefixDelete { + if s == nil || s.engine == nil { + return nil + } + start, end := routePrefixRange(prefix) + routes := s.engine.GetIntersectingRoutes(start, end) + out := make([]stagedVisibilityPrefixDelete, 0, len(routes)) + seen := make(map[string]struct{}, len(routes)) + for _, route := range routes { + if !routeHasStagedVisibility(route) { + continue + } + g := s.groups[route.GroupID] + if g == nil || g.Store == nil { + continue + } + stagedPrefix := distribution.MigrationStagedDataKey(route.MigrationJobID, prefix) + var stagedExclude []byte + if excludePrefix != nil { + stagedExclude = distribution.MigrationStagedDataKey(route.MigrationJobID, excludePrefix) + } + dedupeKey := string(stagedPrefix) + "\x00" + string(stagedExclude) + if _, ok := seen[dedupeKey]; ok { + continue + } + seen[dedupeKey] = struct{}{} + out = append(out, stagedVisibilityPrefixDelete{group: g, prefix: stagedPrefix, excludePrefix: stagedExclude}) + } + return out +} + +func (s *ShardStore) ensurePrefixWriteTimestampFloors(prefix []byte, commitTS uint64) error { + if s == nil || s.engine == nil || commitTS == 0 { + return nil + } + start, end := routePrefixRange(prefix) + for _, route := range s.engine.GetIntersectingRoutes(start, end) { + if route.MinWriteTSExclusive != 0 && commitTS <= route.MinWriteTSExclusive { + return errors.Wrapf(ErrRouteWriteTimestampTooLow, "prefix %q route range [%q,%q) commit_ts=%d floor=%d", prefix, start, end, commitTS, route.MinWriteTSExclusive) + } + } return nil } // DeletePrefixAtRaft is the raft-apply variant of DeletePrefixAt. func (s *ShardStore) DeletePrefixAtRaft(ctx context.Context, prefix []byte, excludePrefix []byte, commitTS uint64) error { - if err := s.ensurePrefixWriteAllowed(prefix, commitTS); err != nil { + if err := s.ensurePrefixWriteTimestampFloors(prefix, commitTS); err != nil { return err } for _, g := range s.groups { @@ -3702,6 +5401,11 @@ func (s *ShardStore) DeletePrefixAtRaft(ctx context.Context, prefix []byte, excl return errors.WithStack(err) } } + for _, del := range s.stagedVisibilityPrefixDeletes(prefix, excludePrefix) { + if err := del.group.Store.DeletePrefixAtRaft(ctx, del.prefix, del.excludePrefix, commitTS); err != nil { + return errors.WithStack(err) + } + } return nil } @@ -3720,13 +5424,32 @@ func (s *ShardStore) DeletePrefixAtRaft(ctx context.Context, prefix []byte, excl // is the receiver only when an aggregate (admin / coordinator) path // is replaying a global FLUSHALL, which is not raft-applied. func (s *ShardStore) DeletePrefixAtRaftAt(ctx context.Context, prefix []byte, excludePrefix []byte, commitTS, appliedIndex uint64) error { - if err := s.ensurePrefixWriteAllowed(prefix, commitTS); err != nil { - return err + return s.DeletePrefixesAtRaftAt(ctx, []store.PrefixDelete{{Prefix: prefix, ExcludePrefix: excludePrefix}}, commitTS, appliedIndex) +} + +func (s *ShardStore) DeletePrefixesAtRaftAt(ctx context.Context, deletes []store.PrefixDelete, commitTS, appliedIndex uint64) error { + if len(deletes) == 0 { + return nil + } + stagedByGroup := make(map[*ShardGroup][]store.PrefixDelete) + for _, del := range deletes { + if err := s.ensurePrefixWriteTimestampFloors(del.Prefix, commitTS); err != nil { + return err + } + for _, staged := range s.stagedVisibilityPrefixDeletes(del.Prefix, del.ExcludePrefix) { + stagedByGroup[staged.group] = append(stagedByGroup[staged.group], store.PrefixDelete{ + Prefix: staged.prefix, + ExcludePrefix: staged.excludePrefix, + }) + } } for _, g := range s.groups { if g == nil || g.Store == nil { continue } + groupDeletes := make([]store.PrefixDelete, 0, len(deletes)+len(stagedByGroup[g])) + groupDeletes = append(groupDeletes, deletes...) + groupDeletes = append(groupDeletes, stagedByGroup[g]...) // Pass appliedIndex through to every group. In the // single-group call-path (the production raft-apply case) // this is correct: appliedIndex IS that group's raft entry @@ -3737,7 +5460,7 @@ func (s *ShardStore) DeletePrefixAtRaftAt(ctx context.Context, prefix []byte, ex // case impossible to reach in production. Tests that // exercise ShardStore.DeletePrefixAtRaftAt across multiple // groups MUST pass appliedIndex=0 to opt out. - if err := g.Store.DeletePrefixAtRaftAt(ctx, prefix, excludePrefix, commitTS, appliedIndex); err != nil { + if err := g.Store.DeletePrefixesAtRaftAt(ctx, groupDeletes, commitTS, appliedIndex); err != nil { return errors.WithStack(err) } } @@ -3880,6 +5603,10 @@ func (s *ShardStore) ImportVersions(context.Context, store.ImportVersionsOptions return store.ImportVersionsResult{}, store.ErrNotSupported } +func (s *ShardStore) ImportVersionsRaft(context.Context, store.ImportVersionsOptions) (store.ImportVersionsResult, error) { + return store.ImportVersionsResult{}, store.ErrNotSupported +} + func (s *ShardStore) MigrationHLCFloor(context.Context, uint64) (uint64, error) { return 0, store.ErrNotSupported } @@ -3925,43 +5652,44 @@ func (s *ShardStore) closeGroup(g *ShardGroup) error { } func (s *ShardStore) groupForKey(key []byte) (*ShardGroup, bool) { - route, ok := s.engine.GetRoute(routeKey(key)) - if !ok { - return nil, false - } - g, ok := s.groups[route.GroupID] + _, g, ok := s.routeAndGroupForKey(key) return g, ok } -// LocalStoreForKey returns this process's store for the key's owning group -// without a leader fence or network proxy. It is reserved for node-local -// auxiliary state such as content-addressed S3 chunk blobs; replicated state -// must continue through the normal ShardStore or Coordinator paths. -func (s *ShardStore) LocalStoreForKey(key []byte) (store.MVCCStore, bool) { - g, ok := s.groupForKey(key) - if !ok || g == nil || g.Store == nil { - return nil, false - } - return g.Store, true +func (s *ShardStore) routeAndGroupForKey(key []byte) (distribution.Route, *ShardGroup, bool) { + route, g, _, ok := s.routeAndGroupForKeyWithVersion(key) + return route, g, ok } -// LocalStores returns every process-local shard store in stable group order. -// It is used by node-local auxiliary maintenance that must recover state after -// snapshot restore without leader routing. -func (s *ShardStore) LocalStores() []store.MVCCStore { - groupIDs := make([]uint64, 0, len(s.groups)) - for groupID := range s.groups { - groupIDs = append(groupIDs, groupID) +func (s *ShardStore) routeAndGroupForKeyWithVersion(key []byte) (distribution.Route, *ShardGroup, uint64, bool) { + if s == nil || s.engine == nil { + return distribution.Route{}, nil, 0, false } - sort.Slice(groupIDs, func(i, j int) bool { return groupIDs[i] < groupIDs[j] }) - stores := make([]store.MVCCStore, 0, len(groupIDs)) - for _, groupID := range groupIDs { - group := s.groups[groupID] - if group != nil && group.Store != nil { - stores = append(stores, group.Store) + if start, end, auxiliary := s3BucketAuxiliaryRouteRange(key); auxiliary { + routes, version := s.engine.GetIntersectingRoutesWithVersion(start, end) + if route, ok := s3BucketAuxiliaryOwnerRouteFromRange(start, end, routes); ok { + g, groupOK := s.groups[route.GroupID] + return route, g, version, groupOK } + return distribution.Route{}, nil, version, false } - return stores + route, version, ok := s.engine.GetRouteWithVersion(routeKey(key)) + if !ok { + return distribution.Route{}, nil, version, false + } + g, ok := s.groups[route.GroupID] + return route, g, version, ok +} + +func (s *ShardStore) s3BucketAuxiliaryOwnerRouteForKey(key []byte) (distribution.Route, bool) { + if s == nil || s.engine == nil { + return distribution.Route{}, false + } + start, end, ok := s3BucketAuxiliaryRouteRange(key) + if !ok { + return distribution.Route{}, false + } + return s3BucketAuxiliaryOwnerRouteFromRange(start, end, s.engine.GetIntersectingRoutes(start, end)) } func (s *ShardStore) proxyRawGet(ctx context.Context, g *ShardGroup, key []byte, ts uint64, groupID uint64, readRouteVersion uint64) ([]byte, error) { @@ -4060,6 +5788,9 @@ func (s *ShardStore) proxyRawScanKeysAt( ts uint64, groupID uint64, readRouteVersion uint64, + routeStart []byte, + routeEnd []byte, + routeBoundsPresent bool, ) ([][]byte, error) { engine := engineForGroup(g) if engine == nil { @@ -4079,13 +5810,16 @@ func (s *ShardStore) proxyRawScanKeysAt( defer cancel() cli := pb.NewRawKVClient(conn) resp, err := cli.RawScanAt(ctx, &pb.RawScanAtRequest{ - StartKey: start, - EndKey: end, - Limit: int64(limit), - Ts: ts, - GroupId: groupID, - ReadRouteVersion: readRouteVersion, - KeysOnly: true, + StartKey: start, + EndKey: end, + Limit: int64(limit), + Ts: ts, + GroupId: groupID, + ReadRouteVersion: readRouteVersion, + KeysOnly: true, + RouteStart: bytes.Clone(routeStart), + RouteEnd: bytes.Clone(routeEnd), + RouteBoundsPresent: routeBoundsPresent, }) if err != nil { return nil, errors.WithStack(err) diff --git a/kv/shard_store_test.go b/kv/shard_store_test.go index 2dc325870..b0e7fa0ff 100644 --- a/kv/shard_store_test.go +++ b/kv/shard_store_test.go @@ -3,8 +3,8 @@ package kv import ( "bytes" "context" + "fmt" "sync" - "sync/atomic" "testing" "time" @@ -17,6 +17,1231 @@ import ( "github.com/stretchr/testify/require" ) +type exportCountingStore struct { + store.MVCCStore + exportCalls int +} + +func (s *exportCountingStore) ExportVersions(ctx context.Context, opts store.ExportVersionsOptions) (store.ExportVersionsResult, error) { + s.exportCalls++ + return s.MVCCStore.ExportVersions(ctx, opts) +} + +func newStagedVisibilityShardStore(t *testing.T) (*ShardStore, *ShardGroup) { + t.Helper() + + engine := distribution.NewEngine() + require.NoError(t, engine.ApplySnapshot(distribution.CatalogSnapshot{ + Version: 1, + Routes: []distribution.RouteDescriptor{ + { + RouteID: 1, + Start: []byte("a"), + End: []byte("z"), + GroupID: 1, + State: distribution.RouteStateActive, + StagedVisibilityActive: true, + MigrationJobID: 9, + MinWriteTSExclusive: 100, + }, + }, + })) + group := &ShardGroup{Store: store.NewMVCCStore()} + return NewShardStore(engine, map[uint64]*ShardGroup{1: group}), group +} + +func newStagedVisibilityPebbleShardStore(t *testing.T) (*ShardStore, *ShardGroup) { + t.Helper() + + engine := distribution.NewEngine() + require.NoError(t, engine.ApplySnapshot(distribution.CatalogSnapshot{ + Version: 1, + Routes: []distribution.RouteDescriptor{ + { + RouteID: 1, + Start: []byte("a"), + End: []byte("z"), + GroupID: 1, + State: distribution.RouteStateActive, + StagedVisibilityActive: true, + MigrationJobID: 9, + MinWriteTSExclusive: 100, + }, + }, + })) + st, err := store.NewPebbleStore(t.TempDir()) + require.NoError(t, err) + t.Cleanup(func() { require.NoError(t, st.Close()) }) + group := &ShardGroup{Store: st} + return NewShardStore(engine, map[uint64]*ShardGroup{1: group}), group +} + +func TestShardStoreGetAt_MergesStagedVisibility(t *testing.T) { + t.Parallel() + + ctx := context.Background() + st, group := newStagedVisibilityShardStore(t) + rawKey := []byte("k") + stagedKey := distribution.MigrationStagedDataKey(9, rawKey) + + require.NoError(t, group.Store.PutAt(ctx, rawKey, []byte("live-old"), 10, 0)) + require.NoError(t, group.Store.PutAt(ctx, stagedKey, []byte("staged-new"), 20, 0)) + got, err := st.GetAt(ctx, rawKey, 25) + require.NoError(t, err) + require.Equal(t, []byte("staged-new"), got) + + require.NoError(t, group.Store.PutAt(ctx, rawKey, []byte("live-new"), 30, 0)) + got, err = st.GetAt(ctx, rawKey, 35) + require.NoError(t, err) + require.Equal(t, []byte("live-new"), got) + + require.NoError(t, group.Store.DeleteAt(ctx, stagedKey, 40)) + _, err = st.GetAt(ctx, rawKey, 45) + require.ErrorIs(t, err, store.ErrKeyNotFound) +} + +func TestShardStoreGetAt_MergesStagedVisibilityPebbleExactKey(t *testing.T) { + t.Parallel() + + ctx := context.Background() + st, group := newStagedVisibilityPebbleShardStore(t) + rawKey := []byte("k") + stagedKey := distribution.MigrationStagedDataKey(9, rawKey) + + require.NoError(t, group.Store.PutAt(ctx, rawKey, []byte("live-old"), 10, 0)) + require.NoError(t, group.Store.PutAt(ctx, stagedKey, []byte("staged-new"), 20, 0)) + + got, err := st.GetAt(ctx, rawKey, 25) + require.NoError(t, err) + require.Equal(t, []byte("staged-new"), got) +} + +func TestShardStoreGetAt_MergesStagedVisibilityForS3BucketAuxiliary(t *testing.T) { + t.Parallel() + + ctx := context.Background() + const bucket = "bucket-a" + routeStart := s3keys.RoutePrefixForBucketAnyGeneration(bucket) + routeEnd := prefixScanEnd(routeStart) + engine := distribution.NewEngine() + require.NoError(t, engine.ApplySnapshot(distribution.CatalogSnapshot{ + Version: 1, + Routes: []distribution.RouteDescriptor{ + { + RouteID: 1, + Start: routeStart, + End: routeEnd, + GroupID: 2, + State: distribution.RouteStateActive, + StagedVisibilityActive: true, + MigrationJobID: 9, + }, + }, + })) + group := &ShardGroup{Store: store.NewMVCCStore()} + st := NewShardStore(engine, map[uint64]*ShardGroup{2: group}) + + for _, tc := range []struct { + name string + key []byte + value []byte + }{ + {name: "bucket meta", key: s3keys.BucketMetaKey(bucket), value: []byte("meta")}, + {name: "bucket generation", key: s3keys.BucketGenerationKey(bucket), value: []byte("generation")}, + } { + t.Run(tc.name, func(t *testing.T) { + require.NoError(t, group.Store.PutAt(ctx, distribution.MigrationStagedDataKey(9, tc.key), tc.value, 20, 0)) + + got, err := st.GetAt(ctx, tc.key, 25) + require.NoError(t, err) + require.Equal(t, tc.value, got) + }) + } +} + +func TestShardStoreGetAt_RoutesS3BucketAuxiliaryToPromotedOwner(t *testing.T) { + t.Parallel() + + ctx := context.Background() + const bucket = "bucket-a" + engine := distribution.NewEngine() + require.NoError(t, engine.ApplySnapshot(distribution.CatalogSnapshot{ + Version: 1, + Routes: s3BucketAuxiliaryPromotedRoutes(), + })) + groups := map[uint64]*ShardGroup{ + 1: {Store: store.NewMVCCStore()}, + 2: {Store: store.NewMVCCStore()}, + } + st := NewShardStore(engine, groups) + key := s3keys.BucketMetaKey(bucket) + require.NoError(t, groups[1].Store.PutAt(ctx, key, []byte("stale-source"), 10, 0)) + require.NoError(t, groups[2].Store.PutAt(ctx, key, []byte("promoted-target"), 20, 0)) + + got, err := st.GetAt(ctx, key, 30) + require.NoError(t, err) + require.Equal(t, []byte("promoted-target"), got) +} + +func TestShardStoreS3BucketAuxiliaryScanFiltersStagedRoutesToBucketRange(t *testing.T) { + t.Parallel() + + ctx := context.Background() + const ( + bucketA = "bucket-a" + bucketB = "bucket-b" + bucketC = "bucket-c" + ) + routeStartA := s3keys.RoutePrefixForBucketAnyGeneration(bucketA) + routeEndA := prefixScanEnd(routeStartA) + routeStartB := s3keys.RoutePrefixForBucketAnyGeneration(bucketB) + routeEndB := prefixScanEnd(routeStartB) + require.Less(t, bytes.Compare(routeStartA, routeStartB), 0) + + engine := distribution.NewEngine() + require.NoError(t, engine.ApplySnapshot(distribution.CatalogSnapshot{ + Version: 1, + Routes: []distribution.RouteDescriptor{ + {RouteID: 1, Start: []byte(""), End: routeStartA, GroupID: 1, State: distribution.RouteStateActive}, + {RouteID: 2, Start: routeStartA, End: routeEndA, GroupID: 1, State: distribution.RouteStateActive, StagedVisibilityActive: true, MigrationJobID: 9}, + {RouteID: 3, Start: routeEndA, End: routeStartB, GroupID: 1, State: distribution.RouteStateActive}, + {RouteID: 4, Start: routeStartB, End: routeEndB, GroupID: 1, State: distribution.RouteStateActive, StagedVisibilityActive: true, MigrationJobID: 10}, + {RouteID: 5, Start: routeEndB, End: nil, GroupID: 1, State: distribution.RouteStateActive}, + }, + })) + group := &ShardGroup{Store: store.NewMVCCStore()} + st := NewShardStore(engine, map[uint64]*ShardGroup{1: group}) + + keyA := s3keys.BucketMetaKey(bucketA) + keyB := s3keys.BucketMetaKey(bucketB) + keyC := s3keys.BucketMetaKey(bucketC) + require.NoError(t, group.Store.PutAt(ctx, keyA, []byte("live-a"), 10, 0)) + require.NoError(t, group.Store.PutAt(ctx, keyB, []byte("live-b"), 10, 0)) + require.NoError(t, group.Store.PutAt(ctx, keyC, []byte("live-c"), 10, 0)) + + exactA, err := st.ScanAt(ctx, keyA, prefixScanEnd(keyA), 10, 20) + require.NoError(t, err) + require.Equal(t, []*store.KVPair{{Key: keyA, Value: []byte("live-a")}}, exactA) + + all, err := st.ScanAt(ctx, []byte(s3keys.BucketMetaPrefix), prefixScanEnd([]byte(s3keys.BucketMetaPrefix)), 10, 20) + require.NoError(t, err) + require.Equal(t, []*store.KVPair{ + {Key: keyA, Value: []byte("live-a")}, + {Key: keyB, Value: []byte("live-b")}, + {Key: keyC, Value: []byte("live-c")}, + }, all) + + reverseAll, err := st.ReverseScanAt(ctx, []byte(s3keys.BucketMetaPrefix), prefixScanEnd([]byte(s3keys.BucketMetaPrefix)), 10, 20) + require.NoError(t, err) + require.Equal(t, []*store.KVPair{ + {Key: keyC, Value: []byte("live-c")}, + {Key: keyB, Value: []byte("live-b")}, + {Key: keyA, Value: []byte("live-a")}, + }, reverseAll) +} + +func TestShardStoreS3BucketAuxiliaryScanUsesPromotedOwner(t *testing.T) { + t.Parallel() + + ctx := context.Background() + const ( + migratedBucket = "bucket-a" + otherBucket = "bucket-z" + ) + engine := distribution.NewEngine() + require.NoError(t, engine.ApplySnapshot(distribution.CatalogSnapshot{ + Version: 1, + Routes: s3BucketAuxiliaryPromotedRoutes(), + })) + groups := map[uint64]*ShardGroup{ + 1: {Store: store.NewMVCCStore()}, + 2: {Store: store.NewMVCCStore()}, + } + st := NewShardStore(engine, groups) + migratedKey := s3keys.BucketMetaKey(migratedBucket) + otherKey := s3keys.BucketMetaKey(otherBucket) + require.NoError(t, groups[1].Store.PutAt(ctx, migratedKey, []byte("stale-source"), 10, 0)) + require.NoError(t, groups[2].Store.PutAt(ctx, migratedKey, []byte("promoted-target"), 20, 0)) + require.NoError(t, groups[1].Store.PutAt(ctx, otherKey, []byte("raw-owner"), 15, 0)) + + start := []byte(s3keys.BucketMetaPrefix) + end := prefixScanEnd(start) + kvs, err := st.ScanAt(ctx, start, end, 10, 30) + require.NoError(t, err) + require.Equal(t, []*store.KVPair{ + {Key: migratedKey, Value: []byte("promoted-target")}, + {Key: otherKey, Value: []byte("raw-owner")}, + }, kvs) + + reverse, err := st.ReverseScanAt(ctx, start, end, 10, 30) + require.NoError(t, err) + require.Equal(t, []*store.KVPair{ + {Key: otherKey, Value: []byte("raw-owner")}, + {Key: migratedKey, Value: []byte("promoted-target")}, + }, reverse) +} + +func TestShardStoreS3BucketAuxiliaryScanPreservesLegacyRawOnlyRows(t *testing.T) { + t.Parallel() + + ctx := context.Background() + for _, tc := range []struct { + name string + prefix string + keyFor func(string) []byte + value []byte + rawValue []byte + }{ + {name: "bucket meta", prefix: s3keys.BucketMetaPrefix, keyFor: s3keys.BucketMetaKey, value: []byte("legacy-meta"), rawValue: []byte("raw-meta")}, + {name: "bucket generation", prefix: s3keys.BucketGenerationPrefix, keyFor: s3keys.BucketGenerationKey, value: []byte("legacy-generation"), rawValue: []byte("raw-generation")}, + } { + t.Run(tc.name, func(t *testing.T) { + const ( + migratedBucket = "bucket-a" + otherBucket = "bucket-z" + ) + engine := distribution.NewEngine() + require.NoError(t, engine.ApplySnapshot(distribution.CatalogSnapshot{ + Version: 1, + Routes: s3BucketAuxiliaryPromotedRoutes(), + })) + groups := map[uint64]*ShardGroup{ + 1: {Store: store.NewMVCCStore()}, + 2: {Store: store.NewMVCCStore()}, + } + st := NewShardStore(engine, groups) + migratedKey := tc.keyFor(migratedBucket) + otherKey := tc.keyFor(otherBucket) + require.NoError(t, groups[1].Store.PutAt(ctx, migratedKey, tc.value, 10, 0)) + require.NoError(t, groups[1].Store.PutAt(ctx, otherKey, tc.rawValue, 15, 0)) + + start := []byte(tc.prefix) + end := prefixScanEnd(start) + kvs, err := st.ScanAt(ctx, start, end, 10, 30) + require.NoError(t, err) + require.Equal(t, []*store.KVPair{ + {Key: migratedKey, Value: tc.value}, + {Key: otherKey, Value: tc.rawValue}, + }, kvs) + + reverse, err := st.ReverseScanAt(ctx, start, end, 10, 30) + require.NoError(t, err) + require.Equal(t, []*store.KVPair{ + {Key: otherKey, Value: tc.rawValue}, + {Key: migratedKey, Value: tc.value}, + }, reverse) + + exact, err := st.ScanAt(ctx, migratedKey, prefixScanEnd(migratedKey), 10, 30) + require.NoError(t, err) + require.Equal(t, []*store.KVPair{{Key: migratedKey, Value: tc.value}}, exact) + }) + } +} + +func TestShardStoreRouteBoundedS3BucketAuxiliaryScanKeepsStagedRows(t *testing.T) { + t.Parallel() + + ctx := context.Background() + const bucket = "bucket-a" + routeStart := s3keys.RoutePrefixForBucketAnyGeneration(bucket) + routeEnd := prefixScanEnd(routeStart) + engine := distribution.NewEngine() + require.NoError(t, engine.ApplySnapshot(distribution.CatalogSnapshot{ + Version: 1, + Routes: []distribution.RouteDescriptor{ + {RouteID: 1, Start: []byte(""), End: routeStart, GroupID: 1, State: distribution.RouteStateActive}, + {RouteID: 2, Start: routeStart, End: routeEnd, GroupID: 1, State: distribution.RouteStateActive, StagedVisibilityActive: true, MigrationJobID: 9}, + {RouteID: 3, Start: routeEnd, End: nil, GroupID: 1, State: distribution.RouteStateActive}, + }, + })) + group := &ShardGroup{Store: store.NewMVCCStore()} + st := NewShardStore(engine, map[uint64]*ShardGroup{1: group}) + key := s3keys.BucketMetaKey(bucket) + require.NoError(t, group.Store.PutAt(ctx, distribution.MigrationStagedDataKey(9, key), []byte("staged"), 20, 0)) + + kvs, err := st.ScanAtWithReadFence(ctx, []byte(s3keys.BucketMetaPrefix), prefixScanEnd([]byte(s3keys.BucketMetaPrefix)), 10, 25, false, 0, 1, routeStart, routeEnd) + require.NoError(t, err) + require.Equal(t, []*store.KVPair{{Key: key, Value: []byte("staged")}}, kvs) +} + +func TestShardStoreRejectsS3BucketAuxiliaryWriteAtMigrationTimestampFloor(t *testing.T) { + t.Parallel() + + ctx := context.Background() + const bucket = "bucket-a" + routeStart := s3keys.RoutePrefixForBucketAnyGeneration(bucket) + routeEnd := prefixScanEnd(routeStart) + engine := distribution.NewEngine() + require.NoError(t, engine.ApplySnapshot(distribution.CatalogSnapshot{ + Version: 1, + Routes: []distribution.RouteDescriptor{ + { + RouteID: 1, + Start: []byte(""), + End: routeStart, + GroupID: 1, + State: distribution.RouteStateActive, + MinWriteTSExclusive: 100, + }, + { + RouteID: 2, + Start: routeStart, + End: routeEnd, + GroupID: 2, + State: distribution.RouteStateActive, + MinWriteTSExclusive: 100, + }, + { + RouteID: 3, + Start: routeEnd, + End: nil, + GroupID: 1, + State: distribution.RouteStateActive, + MinWriteTSExclusive: 100, + }, + }, + })) + st := NewShardStore(engine, map[uint64]*ShardGroup{ + 1: {Store: store.NewMVCCStore()}, + 2: {Store: store.NewMVCCStore()}, + }) + + for _, key := range [][]byte{ + s3keys.BucketMetaKey(bucket), + s3keys.BucketGenerationKey(bucket), + } { + require.ErrorIs(t, st.PutAt(ctx, key, []byte("v"), 100, 0), ErrRouteWriteTimestampTooLow) + require.ErrorIs(t, st.ApplyMutations(ctx, []*store.KVPairMutation{{Op: store.OpTypePut, Key: key, Value: []byte("v")}}, nil, 90, 100), ErrRouteWriteTimestampTooLow) + } +} + +func TestShardStoreStagedVisibilityReadTSCompacted(t *testing.T) { + t.Parallel() + + ctx := context.Background() + st, group := newStagedVisibilityShardStore(t) + retention, ok := group.Store.(store.RetentionController) + require.True(t, ok) + retention.SetMinRetainedTS(15) + + _, err := st.GetAt(ctx, []byte("k"), 10) + require.ErrorIs(t, err, store.ErrReadTSCompacted) + _, err = st.ScanAt(ctx, []byte("a"), []byte("z"), 10, 10) + require.ErrorIs(t, err, store.ErrReadTSCompacted) +} + +func TestShardStoreScanAndLatestCommitTS_MergeStagedVisibility(t *testing.T) { + t.Parallel() + + ctx := context.Background() + st, group := newStagedVisibilityShardStore(t) + + require.NoError(t, group.Store.PutAt(ctx, []byte("b"), []byte("live-b"), 10, 0)) + require.NoError(t, group.Store.PutAt(ctx, []byte("c"), []byte("live-c"), 30, 0)) + require.NoError(t, group.Store.PutAt(ctx, distribution.MigrationStagedDataKey(9, []byte("b")), []byte("staged-b"), 20, 0)) + require.NoError(t, group.Store.PutAt(ctx, distribution.MigrationStagedDataKey(9, []byte("d")), []byte("staged-d"), 15, 0)) + require.NoError(t, group.Store.PutAt(ctx, distribution.MigrationStagedDataKey(9, []byte("e")), []byte("staged-e"), 40, 0)) + require.NoError(t, group.Store.DeleteAt(ctx, []byte("d"), 25)) + + kvs, err := st.ScanAt(ctx, []byte("a"), []byte("z"), 10, 50) + require.NoError(t, err) + require.Equal(t, []*store.KVPair{ + {Key: []byte("b"), Value: []byte("staged-b")}, + {Key: []byte("c"), Value: []byte("live-c")}, + {Key: []byte("e"), Value: []byte("staged-e")}, + }, kvs) + + kvs, err = st.ReverseScanAt(ctx, []byte("a"), []byte("z"), 10, 50) + require.NoError(t, err) + require.Equal(t, []*store.KVPair{ + {Key: []byte("e"), Value: []byte("staged-e")}, + {Key: []byte("c"), Value: []byte("live-c")}, + {Key: []byte("b"), Value: []byte("staged-b")}, + }, kvs) + + ts, exists, err := st.LatestCommitTS(ctx, []byte("b")) + require.NoError(t, err) + require.True(t, exists) + require.Equal(t, uint64(20), ts) + + ts, exists, err = st.LatestCommitTS(ctx, []byte("d")) + require.NoError(t, err) + require.True(t, exists) + require.Equal(t, uint64(25), ts) + + ts, exists, err = st.LatestCommitTS(ctx, []byte("e")) + require.NoError(t, err) + require.True(t, exists) + require.Equal(t, uint64(40), ts) +} + +func TestShardStoreStagedVisibilityScanUsesTwoRangeExports(t *testing.T) { + t.Parallel() + + ctx := context.Background() + st, group := newStagedVisibilityShardStore(t) + counting := &exportCountingStore{MVCCStore: group.Store} + group.Store = counting + require.NoError(t, group.Store.PutAt(ctx, []byte("b"), []byte("live"), 10, 0)) + require.NoError(t, group.Store.PutAt(ctx, distribution.MigrationStagedDataKey(9, []byte("c")), []byte("staged"), 20, 0)) + + kvs, err := st.ScanAt(ctx, []byte("a"), []byte("z"), 10, 30) + require.NoError(t, err) + require.Equal(t, []*store.KVPair{ + {Key: []byte("b"), Value: []byte("live")}, + {Key: []byte("c"), Value: []byte("staged")}, + }, kvs) + require.Equal(t, 2, counting.exportCalls) +} + +func TestStagedVisibilityScanBoundsTreatsEmptyEndAsUnbounded(t *testing.T) { + t.Parallel() + + prefix := distribution.MigrationStagedDataKeyPrefix(9) + start, end := stagedVisibilityScanBounds(9, []byte{}, []byte{}) + require.Equal(t, prefix, start) + require.Equal(t, prefixScanEnd(prefix), end) +} + +func TestShardStoreScanAt_FiltersStagedShadowRowsFromLiveCandidates(t *testing.T) { + t.Parallel() + + ctx := context.Background() + st, group := newStagedVisibilityShardStore(t) + rawKey := []byte("b") + require.NoError(t, group.Store.PutAt(ctx, distribution.MigrationStagedDataKey(9, rawKey), []byte("staged"), 20, 0)) + + kvs, err := st.ScanAt(ctx, []byte(""), nil, 10, 30) + require.NoError(t, err) + require.Equal(t, []*store.KVPair{{Key: rawKey, Value: []byte("staged")}}, kvs) +} + +func TestShardStoreScanAt_PreservesNonStagedRoutesDuringBroadStagedVisibilityScan(t *testing.T) { + t.Parallel() + + ctx := context.Background() + engine := distribution.NewEngine() + require.NoError(t, engine.ApplySnapshot(distribution.CatalogSnapshot{ + Version: 1, + Routes: []distribution.RouteDescriptor{ + {RouteID: 1, Start: []byte("a"), End: []byte("m"), GroupID: 1, State: distribution.RouteStateActive}, + { + RouteID: 2, + Start: []byte("m"), + End: []byte("t"), + GroupID: 1, + State: distribution.RouteStateActive, + StagedVisibilityActive: true, + MigrationJobID: 9, + }, + {RouteID: 3, Start: []byte("t"), End: []byte("z"), GroupID: 1, State: distribution.RouteStateActive}, + }, + })) + group := &ShardGroup{Store: store.NewMVCCStore()} + st := NewShardStore(engine, map[uint64]*ShardGroup{1: group}) + + require.NoError(t, group.Store.PutAt(ctx, []byte("b"), []byte("live-b"), 10, 0)) + require.NoError(t, group.Store.PutAt(ctx, []byte("n"), []byte("live-n"), 10, 0)) + require.NoError(t, group.Store.PutAt(ctx, []byte("u"), []byte("live-u"), 10, 0)) + require.NoError(t, group.Store.PutAt(ctx, distribution.MigrationStagedDataKey(9, []byte("n")), []byte("staged-n"), 20, 0)) + require.NoError(t, group.Store.PutAt(ctx, distribution.MigrationStagedDataKey(9, []byte("o")), []byte("staged-o"), 20, 0)) + + kvs, err := st.ScanAt(ctx, nil, nil, 10, 30) + require.NoError(t, err) + require.Equal(t, []*store.KVPair{ + {Key: []byte("b"), Value: []byte("live-b")}, + {Key: []byte("n"), Value: []byte("staged-n")}, + {Key: []byte("o"), Value: []byte("staged-o")}, + {Key: []byte("u"), Value: []byte("live-u")}, + }, kvs) + + kvs, err = st.ReverseScanAt(ctx, nil, nil, 10, 30) + require.NoError(t, err) + require.Equal(t, []*store.KVPair{ + {Key: []byte("u"), Value: []byte("live-u")}, + {Key: []byte("o"), Value: []byte("staged-o")}, + {Key: []byte("n"), Value: []byte("staged-n")}, + {Key: []byte("b"), Value: []byte("live-b")}, + }, kvs) +} + +func TestShardStoreScanAt_RoutesS3BucketAuxiliaryStagedVisibility(t *testing.T) { + t.Parallel() + + ctx := context.Background() + const bucket = "bucket-a" + routeStart := s3keys.RoutePrefixForBucketAnyGeneration(bucket) + routeEnd := prefixScanEnd(routeStart) + engine := distribution.NewEngine() + require.NoError(t, engine.ApplySnapshot(distribution.CatalogSnapshot{ + Version: 1, + Routes: []distribution.RouteDescriptor{ + {RouteID: 1, Start: []byte(""), End: routeStart, GroupID: 1, State: distribution.RouteStateActive}, + { + RouteID: 2, + Start: routeStart, + End: routeEnd, + GroupID: 2, + State: distribution.RouteStateActive, + StagedVisibilityActive: true, + MigrationJobID: 9, + }, + {RouteID: 3, Start: routeEnd, End: nil, GroupID: 1, State: distribution.RouteStateActive}, + }, + })) + groups := map[uint64]*ShardGroup{ + 1: {Store: store.NewMVCCStore()}, + 2: {Store: store.NewMVCCStore()}, + } + st := NewShardStore(engine, groups) + + for _, tc := range []struct { + name string + prefix string + key []byte + value []byte + }{ + {name: "bucket meta", prefix: s3keys.BucketMetaPrefix, key: s3keys.BucketMetaKey(bucket), value: []byte("meta")}, + {name: "bucket generation", prefix: s3keys.BucketGenerationPrefix, key: s3keys.BucketGenerationKey(bucket), value: []byte("generation")}, + } { + t.Run(tc.name, func(t *testing.T) { + require.NoError(t, groups[2].Store.PutAt(ctx, distribution.MigrationStagedDataKey(9, tc.key), tc.value, 20, 0)) + + kvs, err := st.ScanAt(ctx, []byte(tc.prefix), prefixScanEnd([]byte(tc.prefix)), 10, 30) + require.NoError(t, err) + require.Contains(t, kvs, &store.KVPair{Key: tc.key, Value: tc.value}) + }) + } +} + +func TestShardStoreS3BucketAuxiliaryScanHonorsStagedTombstone(t *testing.T) { + t.Parallel() + + ctx := context.Background() + const migratedBucket = "bucket-a" + routeStart := s3keys.RoutePrefixForBucketAnyGeneration(migratedBucket) + routeEnd := prefixScanEnd(routeStart) + engine := distribution.NewEngine() + require.NoError(t, engine.ApplySnapshot(distribution.CatalogSnapshot{ + Version: 1, + Routes: []distribution.RouteDescriptor{ + {RouteID: 1, Start: []byte(""), End: routeStart, GroupID: 1, State: distribution.RouteStateActive}, + {RouteID: 2, Start: routeStart, End: routeEnd, GroupID: 2, State: distribution.RouteStateActive, StagedVisibilityActive: true, MigrationJobID: 9}, + {RouteID: 3, Start: routeEnd, End: nil, GroupID: 1, State: distribution.RouteStateActive}, + }, + })) + groups := map[uint64]*ShardGroup{ + 1: {Store: store.NewMVCCStore()}, + 2: {Store: store.NewMVCCStore()}, + } + st := NewShardStore(engine, groups) + deletedKey := s3keys.BucketMetaKey(migratedBucket) + visibleKey := s3keys.BucketMetaKey("bucket-z") + require.NoError(t, groups[1].Store.PutAt(ctx, deletedKey, []byte("stale"), 10, 0)) + require.NoError(t, groups[1].Store.PutAt(ctx, visibleKey, []byte("visible"), 10, 0)) + require.NoError(t, groups[2].Store.DeleteAt(ctx, distribution.MigrationStagedDataKey(9, deletedKey), 20)) + + start := []byte(s3keys.BucketMetaPrefix) + end := prefixScanEnd(start) + kvs, err := st.ScanAt(ctx, start, end, 1, 30) + require.NoError(t, err) + require.Equal(t, []*store.KVPair{{Key: visibleKey, Value: []byte("visible")}}, kvs) + + kvs, err = st.ScanAt(ctx, deletedKey, prefixScanEnd(deletedKey), 1, 30) + require.NoError(t, err) + require.Empty(t, kvs) +} + +type versionVisibleRawKVServer struct { + pb.UnimplementedRawKVServer + + mu sync.Mutex + visible map[string]bool + latestReqs []*pb.RawLatestCommitTSRequest +} + +func (s *versionVisibleRawKVServer) RawGet(context.Context, *pb.RawGetRequest) (*pb.RawGetResponse, error) { + return &pb.RawGetResponse{}, nil +} + +func (s *versionVisibleRawKVServer) RawScanAt(context.Context, *pb.RawScanAtRequest) (*pb.RawScanAtResponse, error) { + return &pb.RawScanAtResponse{}, nil +} + +func (s *versionVisibleRawKVServer) RawLatestCommitTS(_ context.Context, req *pb.RawLatestCommitTSRequest) (*pb.RawLatestCommitTSResponse, error) { + s.mu.Lock() + defer s.mu.Unlock() + s.latestReqs = append(s.latestReqs, &pb.RawLatestCommitTSRequest{ + Key: bytes.Clone(req.GetKey()), + KeyBatch: bytes.Clone(req.GetKeyBatch()), + GroupId: req.GetGroupId(), + ReadRouteVersion: req.GetReadRouteVersion(), + VersionVisibleAtTs: req.GetVersionVisibleAtTs(), + }) + keys, err := pb.DecodeRawLatestCommitTSKeyBatch(req.GetKeyBatch(), store.MaxDeltaScanLimit+1) + if err != nil { + return nil, err + } + if len(keys) > 0 { + results := make([]bool, len(keys)) + for i, key := range keys { + results[i] = s.visible[string(key)] + } + return &pb.RawLatestCommitTSResponse{ + VersionVisibleResults: results, + VersionVisibleSupported: true, + }, nil + } + return &pb.RawLatestCommitTSResponse{ + VersionVisible: s.visible[string(req.GetKey())], + VersionVisibleSupported: true, + }, nil +} + +func TestShardStoreS3BucketAuxiliaryOwnerProbeUsesLeaderRoutedReadFence(t *testing.T) { + t.Parallel() + + ctx := context.Background() + const bucket = "bucket-a" + migratedKey := s3keys.BucketMetaKey(bucket) + stagedKey := distribution.MigrationStagedDataKey(9, migratedKey) + probe := &versionVisibleRawKVServer{ + visible: map[string]bool{string(stagedKey): true}, + } + addr, stop := startRawKVServer(t, probe) + t.Cleanup(stop) + + engine := distribution.NewEngine() + require.NoError(t, engine.ApplySnapshot(distribution.CatalogSnapshot{ + Version: 77, + Routes: s3BucketAuxiliaryStagedRoutes(bucket, 1, 2), + })) + groups := map[uint64]*ShardGroup{ + 1: {Store: store.NewMVCCStore()}, + 2: { + Store: store.NewMVCCStore(), + Engine: &followerProxyEngine{leader: addr}, + }, + } + st := NewShardStore(engine, groups) + visibleKey := s3keys.BucketMetaKey("bucket-z") + require.NoError(t, groups[1].Store.PutAt(ctx, migratedKey, []byte("stale-source"), 10, 0)) + require.NoError(t, groups[1].Store.PutAt(ctx, visibleKey, []byte("visible"), 10, 0)) + + start := []byte(s3keys.BucketMetaPrefix) + kvs, err := st.ScanAtWithReadFence(ctx, start, prefixScanEnd(start), 10, 30, false, 0, 77, nil, nil) + require.NoError(t, err) + require.Equal(t, []*store.KVPair{{Key: visibleKey, Value: []byte("visible")}}, kvs) + + probe.mu.Lock() + defer probe.mu.Unlock() + // Staged is probed first, and a staged hit answers the question, so the + // live probe is never issued. Probing live first would have to fall + // through to staged anyway, and that pair is what a concurrent promotion + // slips between. + require.Len(t, probe.latestReqs, 1) + keys, err := pb.DecodeRawLatestCommitTSKeyBatch(probe.latestReqs[0].GetKeyBatch(), store.MaxDeltaScanLimit+1) + require.NoError(t, err) + require.Equal(t, [][]byte{stagedKey}, keys) + for _, req := range probe.latestReqs { + require.Equal(t, uint64(2), req.GetGroupId()) + require.Equal(t, uint64(77), req.GetReadRouteVersion()) + require.Equal(t, uint64(30), req.GetVersionVisibleAtTs()) + } +} + +func TestShardStoreS3BucketAuxiliaryOwnerProbeBatchesFollowerChecks(t *testing.T) { + t.Parallel() + + ctx := context.Background() + keys := [][]byte{ + s3keys.BucketMetaKey("bucket-a"), + s3keys.BucketMetaKey("bucket-b"), + s3keys.BucketMetaKey("bucket-c"), + } + probe := &versionVisibleRawKVServer{visible: map[string]bool{}} + for _, key := range keys { + probe.visible[string(distribution.MigrationStagedDataKey(9, key))] = true + } + addr, stop := startRawKVServer(t, probe) + t.Cleanup(stop) + + engine := distribution.NewEngine() + require.NoError(t, engine.ApplySnapshot(distribution.CatalogSnapshot{ + Version: 77, + Routes: []distribution.RouteDescriptor{ + {RouteID: 1, Start: []byte(""), End: []byte(s3keys.RoutePrefix), GroupID: 1, State: distribution.RouteStateActive}, + { + RouteID: 2, Start: []byte(s3keys.RoutePrefix), End: prefixScanEnd([]byte(s3keys.RoutePrefix)), + GroupID: 2, State: distribution.RouteStateActive, StagedVisibilityActive: true, MigrationJobID: 9, + }, + {RouteID: 3, Start: prefixScanEnd([]byte(s3keys.RoutePrefix)), End: nil, GroupID: 1, State: distribution.RouteStateActive}, + }, + })) + groups := map[uint64]*ShardGroup{ + 1: {Store: store.NewMVCCStore()}, + 2: {Store: store.NewMVCCStore(), Engine: &followerProxyEngine{leader: addr}}, + } + st := NewShardStore(engine, groups) + for _, key := range keys { + require.NoError(t, groups[1].Store.PutAt(ctx, key, []byte("stale-source"), 10, 0)) + } + + start := []byte(s3keys.BucketMetaPrefix) + kvs, err := st.ScanAtWithReadFence(ctx, start, prefixScanEnd(start), 10, 30, false, 0, 77, nil, nil) + require.NoError(t, err) + require.Empty(t, kvs) + + probe.mu.Lock() + defer probe.mu.Unlock() + require.Len(t, probe.latestReqs, 1) + require.Equal(t, uint64(2), probe.latestReqs[0].GetGroupId()) + require.Equal(t, uint64(77), probe.latestReqs[0].GetReadRouteVersion()) + require.Equal(t, uint64(30), probe.latestReqs[0].GetVersionVisibleAtTs()) + decodedKeys, err := pb.DecodeRawLatestCommitTSKeyBatch(probe.latestReqs[0].GetKeyBatch(), store.MaxDeltaScanLimit+1) + require.NoError(t, err) + require.Equal(t, [][]byte{ + distribution.MigrationStagedDataKey(9, keys[0]), + distribution.MigrationStagedDataKey(9, keys[1]), + distribution.MigrationStagedDataKey(9, keys[2]), + }, decodedKeys) +} + +func TestShardStoreS3BucketAuxiliaryOwnerProbeFailsWhenLeaderUnavailable(t *testing.T) { + t.Parallel() + + ctx := context.Background() + const bucket = "bucket-a" + migratedKey := s3keys.BucketMetaKey(bucket) + engine := distribution.NewEngine() + require.NoError(t, engine.ApplySnapshot(distribution.CatalogSnapshot{ + Version: 77, + Routes: s3BucketAuxiliaryStagedRoutes(bucket, 1, 2), + })) + groups := map[uint64]*ShardGroup{ + 1: {Store: store.NewMVCCStore()}, + 2: { + Store: store.NewMVCCStore(), + Engine: &followerProxyEngine{}, + }, + } + st := NewShardStore(engine, groups) + require.NoError(t, groups[1].Store.PutAt(ctx, migratedKey, []byte("stale-source"), 10, 0)) + + start := []byte(s3keys.BucketMetaPrefix) + _, err := st.ScanAtWithReadFence(ctx, start, prefixScanEnd(start), 10, 30, false, 0, 77, nil, nil) + require.ErrorIs(t, err, ErrLeaderNotFound) +} + +func TestShardStoreExplicitGroupS3BucketAuxiliaryScanKeepsOwnerRoutes(t *testing.T) { + t.Parallel() + + ctx := context.Background() + const bucket = "bucket-a" + migratedKey := s3keys.BucketMetaKey(bucket) + engine := distribution.NewEngine() + require.NoError(t, engine.ApplySnapshot(distribution.CatalogSnapshot{ + Version: 77, + Routes: s3BucketAuxiliaryStagedRoutes(bucket, 1, 2), + })) + groups := map[uint64]*ShardGroup{ + 1: {Store: store.NewMVCCStore()}, + 2: {Store: store.NewMVCCStore()}, + } + st := NewShardStore(engine, groups) + require.NoError(t, groups[1].Store.PutAt(ctx, migratedKey, []byte("stale-source"), 10, 0)) + require.NoError(t, groups[2].Store.PutAt(ctx, distribution.MigrationStagedDataKey(9, migratedKey), []byte("owner"), 20, 0)) + + start := []byte(s3keys.BucketMetaPrefix) + kvs, err := st.ScanAtWithReadFence(ctx, start, prefixScanEnd(start), 10, 30, false, 1, 77, nil, nil) + require.NoError(t, err) + require.Empty(t, kvs) +} + +func TestShardStoreGetAt_ContinuesLatestVersionExportPages(t *testing.T) { + t.Parallel() + + ctx := context.Background() + st, group := newStagedVisibilityPebbleShardStore(t) + rawKey := []byte("k") + large := bytes.Repeat([]byte("x"), 1<<20) + require.NoError(t, group.Store.PutAt(ctx, rawKey, []byte("old"), 20, 0)) + require.NoError(t, group.Store.PutAt(ctx, rawKey, large, 30, 0)) + + got, err := st.GetAt(ctx, rawKey, 25) + require.NoError(t, err) + require.Equal(t, []byte("old"), got) +} + +func TestShardStoreDeletePrefixAtDeletesStagedVisibilityRows(t *testing.T) { + t.Parallel() + + ctx := context.Background() + st, group := newStagedVisibilityShardStore(t) + dropKey := []byte("b/drop") + keepKey := []byte("b/keep") + outsideKey := []byte("c/outside") + + require.NoError(t, group.Store.PutAt(ctx, distribution.MigrationStagedDataKey(9, dropKey), []byte("drop"), 20, 0)) + require.NoError(t, group.Store.PutAt(ctx, distribution.MigrationStagedDataKey(9, keepKey), []byte("keep"), 20, 0)) + require.NoError(t, group.Store.PutAt(ctx, distribution.MigrationStagedDataKey(9, outsideKey), []byte("outside"), 20, 0)) + + require.NoError(t, st.DeletePrefixAt(ctx, []byte("b/"), []byte("b/keep"), 101)) + + _, err := st.GetAt(ctx, dropKey, 150) + require.ErrorIs(t, err, store.ErrKeyNotFound) + got, err := st.GetAt(ctx, keepKey, 150) + require.NoError(t, err) + require.Equal(t, []byte("keep"), got) + got, err = st.GetAt(ctx, outsideKey, 150) + require.NoError(t, err) + require.Equal(t, []byte("outside"), got) +} + +func TestShardStoreRouteFilteredLeaderScanUsesStagedVisibility(t *testing.T) { + t.Parallel() + + ctx := context.Background() + st, group := newStagedVisibilityShardStore(t) + require.NoError(t, group.Store.PutAt(ctx, distribution.MigrationStagedDataKey(9, []byte("b")), []byte("staged-b"), 20, 0)) + route, _, ok := st.routeAndGroupForKey([]byte("b")) + require.True(t, ok) + + filtered, cursorKVs, err := st.scanRouteAtLeaderRouteFilter( + ctx, + group, + route, + []byte("a"), + []byte("z"), + 10, + 10, + 25, + false, + []byte("b"), + []byte("c"), + ) + require.NoError(t, err) + require.Equal(t, []*store.KVPair{{Key: []byte("b"), Value: []byte("staged-b")}}, filtered) + require.Equal(t, []*store.KVPair{{Key: []byte("b"), Value: []byte("staged-b")}}, cursorKVs) +} + +func TestShardStoreExplicitGroupReads_MergeStagedVisibility(t *testing.T) { + t.Parallel() + + ctx := context.Background() + st, group := newStagedVisibilityShardStore(t) + + require.NoError(t, group.Store.PutAt(ctx, []byte("b"), []byte("live-b"), 10, 0)) + require.NoError(t, group.Store.PutAt(ctx, distribution.MigrationStagedDataKey(9, []byte("b")), []byte("staged-b"), 20, 0)) + require.NoError(t, group.Store.PutAt(ctx, distribution.MigrationStagedDataKey(9, []byte("c")), []byte("staged-c"), 30, 0)) + + got, err := st.GetGroupAt(ctx, 1, []byte("b"), 25) + require.NoError(t, err) + require.Equal(t, []byte("staged-b"), got) + + kvs, err := st.ScanGroupAt(ctx, 1, []byte("a"), []byte("z"), 10, 35) + require.NoError(t, err) + require.Equal(t, []*store.KVPair{ + {Key: []byte("b"), Value: []byte("staged-b")}, + {Key: []byte("c"), Value: []byte("staged-c")}, + }, kvs) + + keys, err := st.ScanKeysAtWithReadFence(ctx, []byte("a"), []byte("z"), 10, 35, 1, 0) + require.NoError(t, err) + require.Equal(t, [][]byte{[]byte("b"), []byte("c")}, keys) + + kvs, err = st.ReverseScanGroupAt(ctx, 1, []byte("a"), []byte("z"), 10, 35) + require.NoError(t, err) + require.Equal(t, []*store.KVPair{ + {Key: []byte("c"), Value: []byte("staged-c")}, + {Key: []byte("b"), Value: []byte("staged-b")}, + }, kvs) + + kvs, err = st.ScanAtWithReadFence(ctx, []byte("a"), []byte("z"), 10, 35, false, 1, 0, []byte("a"), []byte("z")) + require.NoError(t, err) + require.Equal(t, []*store.KVPair{ + {Key: []byte("b"), Value: []byte("staged-b")}, + {Key: []byte("c"), Value: []byte("staged-c")}, + }, kvs) + + kvs, err = st.ScanAtWithReadFence(ctx, []byte("a"), []byte("z"), 10, 35, true, 1, 0, nil, nil) + require.NoError(t, err) + require.Equal(t, []*store.KVPair{ + {Key: []byte("c"), Value: []byte("staged-c")}, + {Key: []byte("b"), Value: []byte("staged-b")}, + }, kvs) +} + +func TestShardStoreExplicitGroupReads_FailClosedWhenRouteMovedToStagedGroup(t *testing.T) { + t.Parallel() + + ctx := context.Background() + engine := distribution.NewEngine() + require.NoError(t, engine.ApplySnapshot(distribution.CatalogSnapshot{ + Version: 1, + Routes: []distribution.RouteDescriptor{ + { + RouteID: 1, + Start: []byte("a"), + End: []byte("z"), + GroupID: 2, + State: distribution.RouteStateActive, + StagedVisibilityActive: true, + MigrationJobID: 9, + }, + }, + })) + 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("b"), []byte("old-source"), 10, 0)) + require.NoError(t, groups[2].Store.PutAt(ctx, distribution.MigrationStagedDataKey(9, []byte("b")), []byte("staged-target"), 20, 0)) + + _, err := st.GetGroupAt(ctx, 1, []byte("b"), 25) + require.ErrorIs(t, err, ErrExplicitGroupStagedVisibilityUnresolved) + + _, err = st.ScanGroupAt(ctx, 1, []byte("a"), []byte("z"), 10, 25) + require.ErrorIs(t, err, ErrExplicitGroupStagedVisibilityUnresolved) +} + +func TestShardStoreExplicitGroupScan_NormalizesRouteMappedBoundsForStagedRoutes(t *testing.T) { + t.Parallel() + + ctx := context.Background() + engine := distribution.NewEngine() + require.NoError(t, engine.ApplySnapshot(distribution.CatalogSnapshot{ + Version: 1, + Routes: []distribution.RouteDescriptor{ + { + RouteID: 1, + Start: sqsGlobalRouteKey, + End: prefixScanEnd(sqsGlobalRouteKey), + GroupID: 2, + State: distribution.RouteStateActive, + StagedVisibilityActive: true, + MigrationJobID: 9, + }, + }, + })) + groups := map[uint64]*ShardGroup{ + 1: {Store: store.NewMVCCStore()}, + 2: {Store: store.NewMVCCStore()}, + } + st := NewShardStore(engine, groups) + start := []byte("!sqs|msg|vis|p|") + end := prefixScanEnd(start) + require.NoError(t, groups[1].Store.PutAt(ctx, []byte("!sqs|msg|vis|p|orders|1"), []byte("old-source"), 10, 0)) + + _, err := st.ScanGroupAt(ctx, 1, start, end, 10, 25) + require.ErrorIs(t, err, ErrExplicitGroupStagedVisibilityUnresolved) +} + +func TestShardStoreExplicitGroupRead_IgnoresUnrelatedStagedRoutes(t *testing.T) { + t.Parallel() + + ctx := context.Background() + engine := distribution.NewEngine() + require.NoError(t, engine.ApplySnapshot(distribution.CatalogSnapshot{ + Version: 1, + Routes: []distribution.RouteDescriptor{ + // The read's own range is owned by the group it names, so the + // staged route below is genuinely unrelated to it. + {RouteID: 1, Start: []byte("a"), End: []byte("m"), GroupID: 1, State: distribution.RouteStateActive}, + { + RouteID: 2, + Start: []byte("m"), + End: []byte("z"), + GroupID: 2, + State: distribution.RouteStateActive, + StagedVisibilityActive: true, + MigrationJobID: 9, + }, + }, + })) + 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("b"), []byte("explicit-group"), 10, 0)) + + got, err := st.GetGroupAt(ctx, 1, []byte("b"), 25) + require.NoError(t, err) + require.Equal(t, []byte("explicit-group"), got) + + kvs, err := st.ScanGroupAt(ctx, 1, []byte("b"), []byte("c"), 10, 25) + require.NoError(t, err) + require.Equal(t, []*store.KVPair{{Key: []byte("b"), Value: []byte("explicit-group")}}, kvs) +} + +func TestShardStoreScanAt_ContinuesStagedVisibilityAfterCandidateWindow(t *testing.T) { + t.Parallel() + + ctx := context.Background() + st, group := newStagedVisibilityShardStore(t) + limit := stagedVisibilityMaxCandidateWindow + 3 + for i := range limit { + key := []byte(fmt.Sprintf("k%05d", i)) + require.NoError(t, group.Store.PutAt(ctx, distribution.MigrationStagedDataKey(9, key), []byte(fmt.Sprintf("v%05d", i)), 10, 0)) + } + + kvs, err := st.ScanAt(ctx, []byte("a"), []byte("z"), limit, 20) + require.NoError(t, err) + require.Len(t, kvs, limit) + require.Equal(t, []byte("k00000"), kvs[0].Key) + require.Equal(t, []byte(fmt.Sprintf("k%05d", limit-1)), kvs[limit-1].Key) + + kvs, err = st.ReverseScanAt(ctx, []byte("a"), []byte("z"), limit, 20) + require.NoError(t, err) + require.Len(t, kvs, limit) + require.Equal(t, []byte(fmt.Sprintf("k%05d", limit-1)), kvs[0].Key) + require.Equal(t, []byte("k00000"), kvs[limit-1].Key) +} + +func TestShardStoreScanAtRestrictsStagedVisibilityToSafeFrontier(t *testing.T) { + t.Parallel() + + ctx := context.Background() + st, group := newStagedVisibilityShardStore(t) + limit := stagedVisibilityMaxCandidateWindow + 1 + for i := range limit { + key := []byte(fmt.Sprintf("k%05d", i)) + require.NoError(t, group.Store.PutAt(ctx, key, []byte(fmt.Sprintf("live%05d", i)), 10, 0)) + } + require.NoError(t, group.Store.PutAt(ctx, distribution.MigrationStagedDataKey(9, []byte("x-staged")), []byte("staged"), 10, 0)) + + kvs, err := st.ScanAt(ctx, []byte("a"), []byte("z"), limit, 20) + require.NoError(t, err) + require.Len(t, kvs, limit) + require.Equal(t, []byte("k00000"), kvs[0].Key) + require.Equal(t, []byte(fmt.Sprintf("k%05d", limit-1)), kvs[limit-1].Key) + for _, kvp := range kvs { + require.NotEqual(t, []byte("x-staged"), kvp.Key) + } +} + +func TestStagedVisibilityCandidateBoundary_UsesSafeFrontier(t *testing.T) { + t.Parallel() + + live := []*store.KVPair{{Key: []byte("a")}, {Key: []byte("c")}} + staged := []*store.KVPair{ + {Key: distribution.MigrationStagedDataKey(9, []byte("b"))}, + {Key: distribution.MigrationStagedDataKey(9, []byte("z"))}, + } + boundary, ok := stagedVisibilityCandidateBoundary(live, staged, false, false, false) + require.True(t, ok) + require.Equal(t, []byte("c"), boundary) + + boundary, ok = stagedVisibilityCandidateBoundary(live, staged, false, false, true) + require.True(t, ok) + require.Equal(t, []byte("b"), boundary) +} + +func TestShardStoreApplyMutations_ValidatesStagedReadKeys(t *testing.T) { + t.Parallel() + + ctx := context.Background() + st, group := newStagedVisibilityShardStore(t) + readKey := []byte("k") + require.NoError(t, group.Store.PutAt(ctx, distribution.MigrationStagedDataKey(9, readKey), []byte("staged"), 20, 0)) + + err := st.ApplyMutations(ctx, []*store.KVPairMutation{ + {Op: store.OpTypePut, Key: []byte("m"), Value: []byte("write")}, + }, [][]byte{readKey}, 10, 101) + require.ErrorIs(t, err, store.ErrWriteConflict) +} + +func TestShardStoreApplyMutations_ValidatesStagedWriteKeys(t *testing.T) { + t.Parallel() + + ctx := context.Background() + st, group := newStagedVisibilityShardStore(t) + writeKey := []byte("k") + require.NoError(t, group.Store.PutAt(ctx, distribution.MigrationStagedDataKey(9, writeKey), []byte("staged"), 20, 0)) + + apply := []struct { + name string + fn func(context.Context, []*store.KVPairMutation, [][]byte, uint64, uint64) error + }{ + { + name: "direct", + fn: st.ApplyMutations, + }, + { + name: "raft", + fn: st.ApplyMutationsRaft, + }, + { + name: "raft_at", + fn: func(ctx context.Context, muts []*store.KVPairMutation, readKeys [][]byte, startTS, commitTS uint64) error { + return st.ApplyMutationsRaftAt(ctx, muts, readKeys, startTS, commitTS, 1) + }, + }, + } + for _, tc := range apply { + t.Run(tc.name, func(t *testing.T) { + err := tc.fn(ctx, []*store.KVPairMutation{ + {Op: store.OpTypePut, Key: writeKey, Value: []byte("write")}, + }, nil, 10, 101) + require.ErrorIs(t, err, store.ErrWriteConflict) + }) + } +} + +func TestShardStorePhysicalLimitFallsBackToStagedVisibilityScan(t *testing.T) { + t.Parallel() + + ctx := context.Background() + st, group := newStagedVisibilityShardStore(t) + require.NoError(t, group.Store.PutAt(ctx, []byte("b/live"), []byte("live"), 10, 0)) + require.NoError(t, group.Store.PutAt(ctx, distribution.MigrationStagedDataKey(9, []byte("b/staged")), []byte("staged"), 20, 0)) + + kvs, limitReached, err := st.ScanAtPhysicalLimit(ctx, []byte("b"), []byte("c"), 10, 10, 50) + require.NoError(t, err) + require.False(t, limitReached) + require.Equal(t, []*store.KVPair{ + {Key: []byte("b/live"), Value: []byte("live")}, + {Key: []byte("b/staged"), Value: []byte("staged")}, + }, kvs) + + kvs, limitReached, err = st.ReverseScanAtPhysicalLimit(ctx, []byte("b"), []byte("c"), 10, 10, 50) + require.NoError(t, err) + require.False(t, limitReached) + require.Equal(t, []*store.KVPair{ + {Key: []byte("b/staged"), Value: []byte("staged")}, + {Key: []byte("b/live"), Value: []byte("live")}, + }, kvs) +} + +func TestShardStoreRejectsWritesAtMigrationTimestampFloor(t *testing.T) { + t.Parallel() + + ctx := context.Background() + st, _ := newStagedVisibilityShardStore(t) + + err := st.PutAt(ctx, []byte("k"), []byte("low"), 100, 0) + require.ErrorIs(t, err, ErrRouteWriteTimestampTooLow) + require.NoError(t, st.PutAt(ctx, []byte("k"), []byte("ok"), 101, 0)) +} + +func TestShardStoreRaftApplyRejectsMigrationTimestampFloor(t *testing.T) { + t.Parallel() + + ctx := context.Background() + st, _ := newStagedVisibilityShardStore(t) + + require.ErrorIs(t, st.ApplyMutationsRaft(ctx, []*store.KVPairMutation{ + {Op: store.OpTypePut, Key: []byte("k-raft"), Value: []byte("v")}, + }, nil, 90, 100), ErrRouteWriteTimestampTooLow) + require.ErrorIs(t, st.ApplyMutationsRaftAt(ctx, []*store.KVPairMutation{ + {Op: store.OpTypePut, Key: []byte("k-raft-at"), Value: []byte("v")}, + }, nil, 90, 100, 1), ErrRouteWriteTimestampTooLow) + require.ErrorIs(t, st.DeletePrefixAtRaft(ctx, []byte("k-raft"), nil, 100), ErrRouteWriteTimestampTooLow) + require.ErrorIs(t, st.DeletePrefixAtRaftAt(ctx, []byte("k-raft-at"), nil, 100, 2), ErrRouteWriteTimestampTooLow) +} + type followerProxyEngine struct { leader string } @@ -89,88 +1314,6 @@ func TestShardStoreScanAt_IncludesListKeysAcrossShards(t *testing.T) { require.Equal(t, itemKey, kvs[0].Key) } -// SplitRange normalizes every boundary through kv.RouteKey before storing it -// (adapter/distribution_server.go), so a boundary governing a Redis wide-column -// family is always a decoded user key -- never a raw !hs|fld|-prefixed byte -// string. The fence must therefore cover the group owning the logical user key, -// plus the legacy raw-prefix group when rows may still be routed that way -// (redisWideColumnLegacyScanRouteRange contributes that second query). -// -// The earlier form of this test split at prefix+'m', a raw-prefixed boundary -// SplitRange cannot produce, and asserted a raw range intersection. That -// comparison puts raw storage bytes against decoded user-key boundaries and -// selects the wrong owning route, so it did not describe a reachable state. -func TestShardStoreReadFenceGroupKeysForRangeCoversOwningAndLegacyGroups(t *testing.T) { - t.Parallel() - - for _, tc := range []struct { - name string - userKey string - wantOwner uint64 - wantGroups []uint64 - }{ - // "alpha" sorts below the boundary, so the owning group and the legacy - // raw-prefix group are the same one and collapse to a single fence key. - {name: "owner below boundary", userKey: "alpha", wantOwner: 1, wantGroups: []uint64{1}}, - // "zebra" sorts above it while the raw !hs|fld| prefix sorts below, so - // the two queries land on different groups and both must be fenced. - {name: "owner above boundary", userKey: "zebra", wantOwner: 2, wantGroups: []uint64{2, 1}}, - } { - t.Run(tc.name, func(t *testing.T) { - t.Parallel() - - prefix := store.HashFieldScanPrefix([]byte(tc.userKey)) - engine := distribution.NewEngine() - engine.UpdateRoute([]byte(""), []byte("m"), 1) - engine.UpdateRoute([]byte("m"), nil, 2) - st := NewShardStore(engine, map[uint64]*ShardGroup{ - 1: {}, - 2: {}, - }) - - routes, _ := st.routesForForwardScan(prefix, store.PrefixScanEnd(prefix)) - gotGroups := make([]uint64, 0, len(routes)) - for _, route := range routes { - gotGroups = append(gotGroups, route.GroupID) - } - require.Equal(t, tc.wantGroups, gotGroups) - require.Equal(t, tc.wantOwner, routes[0].GroupID, - "the logical user key's owner must be resolved first") - - got := st.ReadFenceGroupKeysForRange(prefix, store.PrefixScanEnd(prefix)) - require.Len(t, got, len(tc.wantGroups)) - require.Contains(t, got, prefix, - "the queried prefix must be fenced so groupForKey re-derives the owner") - }) - } -} - -func TestShardStoreReadFenceGroupKeysForListRangeUsesStorageRepresentative(t *testing.T) { - t.Parallel() - - for _, userKey := range [][]byte{ - []byte("!sqs|foo"), - []byte("!redis|str|foo"), - } { - t.Run(string(userKey), func(t *testing.T) { - t.Parallel() - prefix := store.ListMetaDeltaScanPrefix(userKey) - engine := distribution.NewEngine() - engine.UpdateRoute([]byte(""), userKey, 1) - engine.UpdateRoute(userKey, nil, 2) - st := NewShardStore(engine, map[uint64]*ShardGroup{ - 1: {}, - 2: {}, - }) - - got := st.ReadFenceGroupKeysForRange(prefix, store.PrefixScanEnd(prefix)) - - require.Equal(t, [][]byte{prefix}, got) - require.Equal(t, userKey, routeKey(got[0])) - }) - } -} - func TestShardStoreScanAt_RoutesListItemScansByUserKey(t *testing.T) { t.Parallel() @@ -203,40 +1346,31 @@ func TestShardStoreScanAt_RoutesListItemScansByUserKey(t *testing.T) { require.Equal(t, k2, kvs[2].Key) } -func TestShardStoreScanAtWithReadFence_RoutesListAuxiliaryScansByUserKey(t *testing.T) { +func TestShardStoreScanAt_RoutesListDeltaScansByUserKey(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()}, - } - t.Cleanup(func() { - _ = groups[1].Store.Close() - _ = groups[2].Store.Close() - }) - shardStore := NewShardStore(engine, groups) - userKey := []byte("x") - deltaKey := store.ListMetaDeltaKey(userKey, 10, 0) - claimKey := store.ListClaimKey(userKey, 1) - require.NoError(t, groups[2].Store.PutAt(ctx, deltaKey, []byte("delta"), 10, 0)) - require.NoError(t, groups[2].Store.PutAt(ctx, claimKey, []byte("claim"), 11, 0)) - + userKey := []byte("x") // routes to group 2; raw !lst|* prefixes route to group 1. for _, tc := range []struct { - name string - prefix []byte - key []byte + name string + key []byte + scanStart []byte + legacyRouting bool }{ - {name: "delta", prefix: store.ListMetaDeltaScanPrefix(userKey), key: deltaKey}, - {name: "claim", prefix: store.ListClaimScanPrefix(userKey), key: claimKey}, + {name: "current", key: store.ListMetaDeltaKey(userKey, 10, 1), scanStart: store.ListMetaDeltaScanPrefix(userKey)}, + {name: "legacy", key: legacyListMetaDeltaKey(userKey, 10), scanStart: store.LegacyListMetaDeltaScanPrefix(userKey), legacyRouting: true}, } { t.Run(tc.name, func(t *testing.T) { - kvs, err := shardStore.ScanAtWithReadFence( - ctx, tc.prefix, prefixScanEnd(tc.prefix), 10, ^uint64(0), false, 0, engine.Version(), nil, nil, - ) + t.Parallel() + st := newTwoRouteShardStoreForScanTest() + deltaValue := store.MarshalListMetaDelta(store.ListMetaDelta{LenDelta: 1}) + if tc.legacyRouting { + require.NoError(t, st.groups[1].Store.PutAt(ctx, tc.key, deltaValue, 1, 0)) + } else { + require.NoError(t, st.PutAt(ctx, tc.key, deltaValue, 1, 0)) + } + + kvs, err := st.ScanAt(ctx, tc.scanStart, store.PrefixScanEnd(tc.scanStart), 10, ^uint64(0)) require.NoError(t, err) require.Len(t, kvs, 1) require.Equal(t, tc.key, kvs[0].Key) @@ -244,33 +1378,68 @@ func TestShardStoreScanAtWithReadFence_RoutesListAuxiliaryScansByUserKey(t *test } } -func TestShardStoreScanAt_RoutesBareListAuxiliaryScansAcrossShards(t *testing.T) { +func TestShardStoreScanAt_BroadLegacyListDeltaScansAllRoutes(t *testing.T) { t.Parallel() ctx := context.Background() + st := newTwoRouteShardStoreForScanTest() + deltaValue := store.MarshalListMetaDelta(store.ListMetaDelta{LenDelta: 1}) + leftKey := legacyListMetaDeltaKey([]byte("left-list"), 10) + rightKey := legacyListMetaDeltaKey([]byte("right-list"), 11) + require.NoError(t, st.groups[1].Store.PutAt(ctx, leftKey, deltaValue, 1, 0)) + require.NoError(t, st.groups[2].Store.PutAt(ctx, rightKey, deltaValue, 1, 0)) + + kvs, err := st.ScanAt(ctx, []byte(store.LegacyListMetaDeltaPrefix), store.PrefixScanEnd([]byte(store.LegacyListMetaDeltaPrefix)), 10, ^uint64(0)) + require.NoError(t, err) + require.Len(t, kvs, 2) + require.Equal(t, leftKey, kvs[0].Key) + require.Equal(t, uint64(1), kvs[0].RouteGroupID) + require.Equal(t, rightKey, kvs[1].Key) + require.Equal(t, uint64(2), kvs[1].RouteGroupID) +} + +func TestShardStoreScanAt_RoutesWideColumnScansByUserKey(t *testing.T) { + t.Parallel() + + ctx := context.Background() + for _, tc := range []struct { + name string + key []byte + scanStart []byte + }{ + {name: "hash field", key: store.HashFieldKey([]byte("x"), []byte("f")), scanStart: store.HashFieldScanPrefix([]byte("x"))}, + {name: "hash delta", key: store.HashMetaDeltaKey([]byte("x"), 10, 0), scanStart: store.HashMetaDeltaScanPrefix([]byte("x"))}, + {name: "set member", key: store.SetMemberKey([]byte("x"), []byte("m")), scanStart: store.SetMemberScanPrefix([]byte("x"))}, + {name: "set delta", key: store.SetMetaDeltaKey([]byte("x"), 10, 0), scanStart: store.SetMetaDeltaScanPrefix([]byte("x"))}, + {name: "zset member", key: store.ZSetMemberKey([]byte("x"), []byte("m")), scanStart: store.ZSetMemberScanPrefix([]byte("x"))}, + {name: "zset score", key: store.ZSetScoreKey([]byte("x"), 1.5, []byte("m")), scanStart: store.ZSetScoreScanPrefix([]byte("x"))}, + {name: "zset delta", key: store.ZSetMetaDeltaKey([]byte("x"), 10, 0), scanStart: store.ZSetMetaDeltaScanPrefix([]byte("x"))}, + {name: "stream meta", key: store.StreamMetaKey([]byte("x")), scanStart: store.StreamMetaKey([]byte("x"))}, + {name: "stream entry", key: store.StreamEntryKey([]byte("x"), 10, 0), scanStart: store.StreamEntryScanPrefix([]byte("x"))}, + } { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + st := newTwoRouteShardStoreForScanTest() + require.NoError(t, st.PutAt(ctx, tc.key, []byte("v"), 1, 0)) + + kvs, err := st.ScanAt(ctx, tc.scanStart, store.PrefixScanEnd(tc.scanStart), 10, ^uint64(0)) + require.NoError(t, err) + require.Len(t, kvs, 1) + require.Equal(t, tc.key, kvs[0].Key) + }) + } +} + +func newTwoRouteShardStoreForScanTest() *ShardStore { 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()}, } - t.Cleanup(func() { - _ = groups[1].Store.Close() - _ = groups[2].Store.Close() - }) - st := NewShardStore(engine, groups) - - left := store.ListMetaDeltaKey([]byte("anna"), 10, 0) - right := store.ListMetaDeltaKey([]byte("zoey"), 11, 0) - require.NoError(t, groups[1].Store.PutAt(ctx, left, []byte("left"), 10, 0)) - require.NoError(t, groups[2].Store.PutAt(ctx, right, []byte("right"), 11, 0)) - - prefix := []byte(store.ListMetaDeltaPrefix) - kvs, err := st.ScanAt(ctx, prefix, prefixScanEnd(prefix), 10, ^uint64(0)) - require.NoError(t, err) - require.Len(t, kvs, 2) - require.Equal(t, [][]byte{left, right}, [][]byte{kvs[0].Key, kvs[1].Key}) + return NewShardStore(engine, groups) } func TestShardStoreScanGroupAt_UsesExplicitGroup(t *testing.T) { @@ -295,127 +1464,75 @@ func TestShardStoreScanGroupAt_UsesExplicitGroup(t *testing.T) { require.Equal(t, []byte("msg-2"), kvs[0].Value) } -func TestShardStoreGetGroupAt_UsesExplicitGroup(t *testing.T) { +func TestShardStoreScanGroupAt_DoesNotClampRouteMappedRawBounds(t *testing.T) { t.Parallel() ctx := context.Background() engine := distribution.NewEngine() - engine.UpdateRoute([]byte(""), nil, 1) + engine.UpdateRoute([]byte(""), nil, 42) groups := map[uint64]*ShardGroup{ - 1: {Store: store.NewMVCCStore()}, 42: {Store: store.NewMVCCStore()}, } st := NewShardStore(engine, groups) - key := []byte("!sqs|msg|data|p|orders|partition-2|msg-2") - require.NoError(t, groups[42].Store.PutAt(ctx, key, []byte("payload"), 7, 0)) + start := []byte("!sqs|msg|vis|p|") + key := []byte("!sqs|msg|vis|p|orders|partition-2") + require.NoError(t, groups[42].Store.PutAt(ctx, key, []byte("msg-2"), 7, 0)) - val, err := st.GetGroupAt(ctx, 42, key, 7) + kvs, err := st.ScanGroupAt(ctx, 42, start, prefixScanEnd(start), 10, 7) require.NoError(t, err) - require.Equal(t, []byte("payload"), val) - - _, err = st.GetAt(ctx, key, 7) - require.ErrorIs(t, err, store.ErrKeyNotFound) + require.Equal(t, []*store.KVPair{{Key: key, Value: []byte("msg-2")}}, kvs) } -func TestShardStoreWritePathsRejectRouteWriteTimestampFloor(t *testing.T) { +func TestShardStoreScanGroupAt_DeduplicatesRouteMappedSameGroupSplits(t *testing.T) { t.Parallel() ctx := context.Background() engine := distribution.NewEngine() + routeEnd := prefixScanEnd(sqsGlobalRouteKey) + split := append(bytes.Clone(sqsGlobalRouteKey), 'm') require.NoError(t, engine.ApplySnapshot(distribution.CatalogSnapshot{ Version: 1, Routes: []distribution.RouteDescriptor{ - { - RouteID: 1, - Start: []byte(""), - End: nil, - GroupID: 1, - State: distribution.RouteStateActive, - MinWriteTSExclusive: 10, - }, + {RouteID: 1, Start: sqsGlobalRouteKey, End: split, GroupID: 42, State: distribution.RouteStateActive}, + {RouteID: 2, Start: split, End: routeEnd, GroupID: 42, State: distribution.RouteStateActive}, }, })) groups := map[uint64]*ShardGroup{ - 1: {Store: store.NewMVCCStore()}, + 42: {Store: store.NewMVCCStore()}, } - t.Cleanup(func() { - _ = groups[1].Store.Close() - }) st := NewShardStore(engine, groups) - require.ErrorIs(t, st.PutAt(ctx, []byte("put-stale"), []byte("v"), 10, 0), store.ErrWriteConflict) - require.NoError(t, st.PutAt(ctx, []byte("put-fresh"), []byte("v"), 11, 0)) - - require.ErrorIs(t, st.DeleteAt(ctx, []byte("delete-stale"), 10), store.ErrWriteConflict) - require.NoError(t, st.DeleteAt(ctx, []byte("delete-fresh"), 11)) - - require.ErrorIs(t, st.PutWithTTLAt(ctx, []byte("ttl-stale"), []byte("v"), 10, 99), store.ErrWriteConflict) - require.NoError(t, st.PutWithTTLAt(ctx, []byte("ttl-fresh"), []byte("v"), 11, 99)) - - require.ErrorIs(t, st.ExpireAt(ctx, []byte("expire-stale"), 99, 10), store.ErrWriteConflict) - require.NoError(t, st.PutAt(ctx, []byte("expire-fresh"), []byte("v"), 11, 0)) - require.NoError(t, st.ExpireAt(ctx, []byte("expire-fresh"), 99, 12)) - - require.ErrorIs(t, st.ApplyMutations(ctx, []*store.KVPairMutation{ - {Op: store.OpTypePut, Key: []byte("apply-stale"), Value: []byte("v")}, - }, nil, 0, 10), store.ErrWriteConflict) - require.NoError(t, st.ApplyMutations(ctx, []*store.KVPairMutation{ - {Op: store.OpTypePut, Key: []byte("apply-fresh"), Value: []byte("v")}, - }, nil, 0, 11)) - - require.ErrorIs(t, st.ApplyMutationsRaft(ctx, []*store.KVPairMutation{ - {Op: store.OpTypePut, Key: []byte("raft-stale"), Value: []byte("v")}, - }, nil, 0, 10), store.ErrWriteConflict) - require.NoError(t, st.ApplyMutationsRaft(ctx, []*store.KVPairMutation{ - {Op: store.OpTypePut, Key: []byte("raft-fresh"), Value: []byte("v")}, - }, nil, 0, 11)) - - require.ErrorIs(t, st.ApplyMutationsRaftAt(ctx, []*store.KVPairMutation{ - {Op: store.OpTypePut, Key: []byte("raft-at-stale"), Value: []byte("v")}, - }, nil, 0, 10, 1), store.ErrWriteConflict) - require.NoError(t, st.ApplyMutationsRaftAt(ctx, []*store.KVPairMutation{ - {Op: store.OpTypePut, Key: []byte("raft-at-fresh"), Value: []byte("v")}, - }, nil, 0, 11, 2)) - - require.ErrorIs(t, st.DeletePrefixAt(ctx, []byte("prefix-stale"), nil, 10), store.ErrWriteConflict) - require.NoError(t, st.DeletePrefixAt(ctx, []byte("prefix-fresh"), nil, 11)) - - require.ErrorIs(t, st.DeletePrefixAtRaft(ctx, []byte("raft-prefix-stale"), nil, 10), store.ErrWriteConflict) - require.NoError(t, st.DeletePrefixAtRaft(ctx, []byte("raft-prefix-fresh"), nil, 11)) + start := []byte("!sqs|msg|vis|p|") + key := []byte("!sqs|msg|vis|p|orders|partition-2") + require.NoError(t, groups[42].Store.PutAt(ctx, key, []byte("msg-2"), 7, 0)) - require.ErrorIs(t, st.DeletePrefixAtRaftAt(ctx, []byte("raft-at-prefix-stale"), nil, 10, 3), store.ErrWriteConflict) - require.NoError(t, st.DeletePrefixAtRaftAt(ctx, []byte("raft-at-prefix-fresh"), nil, 11, 4)) + kvs, err := st.ScanGroupAt(ctx, 42, start, prefixScanEnd(start), 10, 7) + require.NoError(t, err) + require.Equal(t, []*store.KVPair{{Key: key, Value: []byte("msg-2")}}, kvs) } -func TestShardStoreDeletePrefixChecksRedisLogicalRouteFloors(t *testing.T) { +func TestShardStoreGetGroupAt_UsesExplicitGroup(t *testing.T) { t.Parallel() ctx := context.Background() engine := distribution.NewEngine() - require.NoError(t, engine.ApplySnapshot(distribution.CatalogSnapshot{ - Version: 1, - Routes: []distribution.RouteDescriptor{ - {RouteID: 1, Start: []byte(""), End: []byte("m"), GroupID: 1, State: distribution.RouteStateActive}, - {RouteID: 2, Start: []byte("m"), GroupID: 2, State: distribution.RouteStateActive, MinWriteTSExclusive: 100}, - }, - })) + engine.UpdateRoute([]byte(""), nil, 1) groups := map[uint64]*ShardGroup{ - 1: {Store: store.NewMVCCStore()}, - 2: {Store: store.NewMVCCStore()}, + 1: {Store: store.NewMVCCStore()}, + 42: {Store: store.NewMVCCStore()}, } - t.Cleanup(func() { - _ = groups[1].Store.Close() - _ = groups[2].Store.Close() - }) st := NewShardStore(engine, groups) - require.ErrorIs(t, st.DeletePrefixAt(ctx, []byte(store.HashFieldPrefix), nil, 100), store.ErrWriteConflict) - require.ErrorIs(t, st.DeletePrefixAt(ctx, []byte("!lst|"), nil, 100), store.ErrWriteConflict) - require.ErrorIs(t, st.DeletePrefixAt(ctx, []byte("!redis|hash|"), nil, 100), store.ErrWriteConflict) - require.NoError(t, st.DeletePrefixAt(ctx, store.HashFieldScanPrefix([]byte("alpha")), nil, 100)) - require.ErrorIs(t, st.DeletePrefixAt(ctx, store.HashFieldScanPrefix([]byte("zulu")), nil, 100), store.ErrWriteConflict) - require.NoError(t, st.DeletePrefixAt(ctx, store.HashFieldScanPrefix([]byte("zulu")), nil, 101)) + key := []byte("!sqs|msg|data|p|orders|partition-2|msg-2") + require.NoError(t, groups[42].Store.PutAt(ctx, key, []byte("payload"), 7, 0)) + + val, err := st.GetGroupAt(ctx, 42, key, 7) + require.NoError(t, err) + require.Equal(t, []byte("payload"), val) + + _, err = st.GetAt(ctx, key, 7) + require.ErrorIs(t, err, store.ErrKeyNotFound) } func TestShardStore_ForwardsReadFenceStamps(t *testing.T) { @@ -461,7 +1578,6 @@ func TestShardStore_ForwardsReadFenceStamps(t *testing.T) { fake.mu.Lock() require.Equal(t, uint64(100), fake.lastGetReq.GetReadRouteVersion()) require.Equal(t, uint64(100), fake.lastLatestReq.GetReadRouteVersion()) - require.Equal(t, uint64(1), fake.lastLatestReq.GetGroupId()) require.Equal(t, uint64(100), fake.lastScanReq.GetReadRouteVersion()) require.Equal(t, uint64(1), fake.lastScanReq.GetGroupId()) require.Equal(t, []byte("a"), fake.lastScanReq.GetRouteStart()) @@ -473,7 +1589,7 @@ func TestShardStore_ForwardsReadFenceStamps(t *testing.T) { require.NoError(t, err) fake.mu.Lock() - require.Equal(t, uint64(1), fake.lastScanReq.GetGroupId()) + require.Equal(t, uint64(0), fake.lastScanReq.GetGroupId()) require.Equal(t, uint64(100), fake.lastScanReq.GetReadRouteVersion()) fake.mu.Unlock() @@ -481,7 +1597,7 @@ func TestShardStore_ForwardsReadFenceStamps(t *testing.T) { require.NoError(t, err) fake.mu.Lock() - require.Equal(t, uint64(1), fake.lastScanReq.GetGroupId()) + require.Equal(t, uint64(0), fake.lastScanReq.GetGroupId()) require.Equal(t, uint64(100), fake.lastScanReq.GetReadRouteVersion()) require.True(t, fake.lastScanReq.GetKeysOnly()) fake.mu.Unlock() @@ -499,7 +1615,7 @@ func TestShardStore_ForwardsReadFenceStamps(t *testing.T) { require.NoError(t, err) fake.mu.Lock() - require.Equal(t, uint64(1), fake.lastScanReq.GetGroupId()) + require.Equal(t, uint64(0), fake.lastScanReq.GetGroupId()) require.Equal(t, uint64(100), fake.lastScanReq.GetReadRouteVersion()) require.False(t, fake.lastScanReq.GetRouteBoundsPresent()) fake.mu.Unlock() @@ -822,64 +1938,6 @@ func TestShardStoreScanAtWithReadFence_FiltersSuppliedBoundsByRouteKey(t *testin require.Equal(t, left, kvs[0].Key) } -func TestShardStoreScanAtWithReadFence_FiltersRedisAuxiliaryBoundsByRouteKey(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()}, - } - t.Cleanup(func() { _ = groups[1].Store.Close() }) - st := NewShardStore(engine, groups) - - for _, tc := range []struct { - name string - prefix []byte - left []byte - right []byte - }{ - { - name: "list delta", - prefix: []byte(store.ListMetaDeltaPrefix), - left: store.ListMetaDeltaKey([]byte("alpha"), 10, 0), - right: store.ListMetaDeltaKey([]byte("zulu"), 11, 0), - }, - { - name: "list claim", - prefix: []byte(store.ListClaimPrefix), - left: store.ListClaimKey([]byte("alpha"), 1), - right: store.ListClaimKey([]byte("zulu"), 1), - }, - { - name: "stream meta", - prefix: []byte(store.StreamMetaPrefix), - left: store.StreamMetaKey([]byte("alpha")), - right: store.StreamMetaKey([]byte("zulu")), - }, - { - name: "stream entry", - prefix: []byte(store.StreamEntryPrefix), - left: store.StreamEntryKey([]byte("alpha"), 1, 0), - right: store.StreamEntryKey([]byte("zulu"), 1, 0), - }, - } { - t.Run(tc.name, func(t *testing.T) { - require.NoError(t, groups[1].Store.PutAt(ctx, tc.left, []byte("left"), 1, 0)) - require.NoError(t, groups[1].Store.PutAt(ctx, tc.right, []byte("right"), 2, 0)) - - kvs, err := st.ScanAtWithReadFence(ctx, tc.prefix, prefixScanEnd(tc.prefix), 1, 2, false, 0, st.ReadRouteVersion(), []byte("m"), nil) - require.NoError(t, err) - require.Len(t, kvs, 1) - require.Equal(t, tc.right, kvs[0].Key) - }) - } -} - func TestShardStoreScanAtWithReadFence_FiltersByEachRouteBounds(t *testing.T) { t.Parallel() @@ -910,7 +1968,7 @@ func TestShardStoreScanAtWithReadFence_FiltersByEachRouteBounds(t *testing.T) { require.Equal(t, right, kvs[1].Key) } -func TestShardStoreScanAtWithReadFence_ServesExplicitGroupReverse(t *testing.T) { +func TestShardStoreScanAtWithReadFence_AllowsExplicitGroupRouteBoundReverse(t *testing.T) { t.Parallel() ctx := context.Background() @@ -929,15 +1987,13 @@ func TestShardStoreScanAtWithReadFence_ServesExplicitGroupReverse(t *testing.T) require.NoError(t, groups[1].Store.PutAt(ctx, left, []byte("left"), 1, 0)) require.NoError(t, groups[1].Store.PutAt(ctx, right, []byte("right"), 2, 0)) - // An unbounded explicit-group reverse scan is served through the fenced - // route path rather than rejected. Rejecting it here only pushed callers - // back onto the unfenced ReverseScanGroupAt shortcut in the gRPC server. - unbounded, err := st.ScanAtWithReadFence(ctx, rawPrefix, prefixScanEnd(rawPrefix), 1, 2, true, 1, st.ReadRouteVersion(), nil, nil) + kvs, err := st.ScanAtWithReadFence(ctx, rawPrefix, prefixScanEnd(rawPrefix), 1, 2, true, 1, st.ReadRouteVersion(), nil, nil) require.NoError(t, err) - require.Len(t, unbounded, 1) - require.Equal(t, right, unbounded[0].Key) + require.Len(t, kvs, 1) + require.Equal(t, right, kvs[0].Key) + require.Equal(t, []byte("right"), kvs[0].Value) - kvs, err := st.ScanAtWithReadFence(ctx, rawPrefix, prefixScanEnd(rawPrefix), -1, 2, true, 1, st.ReadRouteVersion(), []byte("m"), nil) + kvs, err = st.ScanAtWithReadFence(ctx, rawPrefix, prefixScanEnd(rawPrefix), -1, 2, true, 1, st.ReadRouteVersion(), []byte("m"), nil) require.NoError(t, err) require.Empty(t, kvs) @@ -1208,7 +2264,24 @@ func TestShardStoreScanKeysRouteAtLeaderRefillsAfterTxnInternalKeys(t *testing.T 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), 0) + keys, err := st.scanKeysRouteAtLeader(ctx, g, distribution.Route{GroupID: 1}, []byte(""), nil, 1, ^uint64(0), 0) + require.NoError(t, err) + require.Equal(t, [][]byte{[]byte("a")}, keys) +} + +func TestShardStoreScanKeysRouteAtLeaderRefillsAfterStagedControlKeys(t *testing.T) { + t.Parallel() + + ctx := context.Background() + g := &ShardGroup{Store: store.NewMVCCStore()} + st := NewShardStore(distribution.NewEngine(), map[uint64]*ShardGroup{1: g}) + t.Cleanup(func() { _ = st.Close() }) + + stagedKey := distribution.MigrationStagedDataKey(9, []byte("shadow")) + require.NoError(t, g.Store.PutAt(ctx, stagedKey, []byte("internal"), 1, 0)) + require.NoError(t, g.Store.PutAt(ctx, []byte("a"), []byte("visible"), 2, 0)) + + keys, err := st.scanKeysRouteAtLeader(ctx, g, distribution.Route{GroupID: 1}, []byte(""), nil, 1, ^uint64(0), 0) require.NoError(t, err) require.Equal(t, [][]byte{[]byte("a")}, keys) } @@ -1223,11 +2296,57 @@ func TestShardStoreScanKeysRouteAtLeaderPreservesEmptyKey(t *testing.T) { 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), 0) + keys, err := st.scanKeysRouteAtLeader(ctx, g, distribution.Route{GroupID: 1}, nil, nil, 2, ^uint64(0), 0) require.NoError(t, err) require.Equal(t, [][]byte{[]byte(""), []byte("a")}, keys) } +func TestShardStoreScanKeysRouteAtLeaderIncludesStagedOnlyKeys(t *testing.T) { + t.Parallel() + + ctx := context.Background() + g := &ShardGroup{Store: store.NewMVCCStore()} + st := NewShardStore(distribution.NewEngine(), map[uint64]*ShardGroup{1: g}) + route := distribution.Route{ + GroupID: 1, + StagedVisibilityActive: true, + MigrationJobID: 9, + } + key := []byte("staged-key") + require.NoError(t, g.Store.PutAt(ctx, distribution.MigrationStagedDataKey(route.MigrationJobID, key), []byte("value"), 1, 0)) + + keys, err := st.scanKeysRouteAtLeader(ctx, g, route, []byte(""), nil, 10, ^uint64(0), 0) + require.NoError(t, err) + require.Equal(t, [][]byte{key}, keys) +} + +func TestShardStoreScanKeysAtIncludesStagedOnlyKeys(t *testing.T) { + t.Parallel() + + ctx := context.Background() + engine := distribution.NewEngine() + require.NoError(t, engine.ApplySnapshot(distribution.CatalogSnapshot{ + Version: 1, + Routes: []distribution.RouteDescriptor{{ + RouteID: 1, + Start: []byte(""), + End: nil, + GroupID: 1, + State: distribution.RouteStateActive, + StagedVisibilityActive: true, + MigrationJobID: 9, + }}, + })) + g := &ShardGroup{Store: store.NewMVCCStore()} + st := NewShardStore(engine, map[uint64]*ShardGroup{1: g}) + key := []byte("staged-key") + require.NoError(t, g.Store.PutAt(ctx, distribution.MigrationStagedDataKey(9, key), []byte("value"), 1, 0)) + + keys, err := st.ScanKeysAt(ctx, []byte(""), nil, 10, ^uint64(0)) + require.NoError(t, err) + require.Equal(t, [][]byte{key}, keys) +} + func TestShardStoreProxyScanKeysAtUsesSelectedGroup(t *testing.T) { t.Parallel() @@ -1259,6 +2378,46 @@ func TestShardStoreProxyScanKeysAtUsesSelectedGroup(t *testing.T) { require.True(t, fake.lastScanKeysOnly) } +func TestShardStoreProxyScanKeysAtCarriesStagedRouteBounds(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() + g := &ShardGroup{ + Engine: &followerProxyEngine{leader: addr}, + Store: store.NewMVCCStore(), + } + st := NewShardStore(distribution.NewEngine(), map[uint64]*ShardGroup{42: g}) + t.Cleanup(func() { _ = st.Close() }) + route := distribution.Route{ + Start: []byte("a"), + End: []byte("m"), + GroupID: 42, + StagedVisibilityActive: true, + MigrationJobID: 9, + } + + keys, err := st.scanKeyRouteAtWithReadFence(ctx, route, []byte("a"), []byte("m"), 10, ^uint64(0), false, 7) + require.NoError(t, err) + require.Equal(t, [][]byte{[]byte("k")}, keys) + + fake.mu.Lock() + defer fake.mu.Unlock() + require.Equal(t, uint64(42), fake.lastScanReq.GetGroupId()) + require.Equal(t, uint64(7), fake.lastScanReq.GetReadRouteVersion()) + require.Equal(t, []byte("a"), fake.lastScanReq.GetRouteStart()) + require.Equal(t, []byte("m"), fake.lastScanReq.GetRouteEnd()) + require.True(t, fake.lastScanReq.GetRouteBoundsPresent()) + require.True(t, fake.lastScanReq.GetKeysOnly()) +} + func TestShardStoreProxyScanAtUsesSelectedGroup(t *testing.T) { t.Parallel() @@ -1312,7 +2471,7 @@ func TestShardStoreProxyForwardPageAdvancesFromRawPage(t *testing.T) { } st := NewShardStore(distribution.NewEngine(), map[uint64]*ShardGroup{42: g}) - page, err := st.scanRouteAtForwardPage(ctx, distribution.Route{GroupID: 42}, g, []byte(""), nil, 2, ^uint64(0), 0, nil, nil) + page, err := st.scanRouteAtForwardPage(ctx, distribution.Route{GroupID: 42}, g, []byte(""), nil, 2, ^uint64(0), true, 0, nil, nil) require.NoError(t, err) require.True(t, page.full) require.Equal(t, internalKey, page.advanceKey) @@ -1756,231 +2915,7 @@ func TestShardStoreScanAt_RoutesExactRedisWideColumnScanToOneShard(t *testing.T) engine := distribution.NewEngine() engine.UpdateRoute([]byte(""), []byte("am"), 1) - engine.UpdateRoute([]byte("am"), nil, 2) - groups := map[uint64]*ShardGroup{ - 1: {Store: store.NewMVCCStore()}, - 2: {Store: store.NewMVCCStore()}, - } - t.Cleanup(func() { - _ = groups[1].Store.Close() - _ = groups[2].Store.Close() - }) - st := NewShardStore(engine, groups) - - start := store.HashFieldScanPrefix([]byte("alice")) - routes, clamp, _ := st.routesForScanWithVersion(start, prefixScanEnd(start)) - require.False(t, clamp) - require.Len(t, routes, 1) - require.Equal(t, uint64(1), routes[0].GroupID) -} - -func TestShardStoreRoutesForWideColumnBoundedPatternIncludesLegacyRawRoute(t *testing.T) { - t.Parallel() - - 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()}, - } - t.Cleanup(func() { - _ = groups[1].Store.Close() - _ = groups[2].Store.Close() - }) - st := NewShardStore(engine, groups) - - start := store.HashFieldScanPrefix([]byte("m")) - routes, clamp, _ := st.routesForScanWithVersion(start, prefixScanEnd([]byte(store.HashFieldPrefix))) - require.False(t, clamp) - require.Len(t, routes, 2) - require.Equal(t, uint64(2), routes[0].GroupID) - require.Equal(t, uint64(1), routes[1].GroupID) -} - -func TestShardStoreRedisWideColumnReadsLegacyRawRoute(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()}, - } - t.Cleanup(func() { - _ = groups[1].Store.Close() - _ = groups[2].Store.Close() - }) - st := NewShardStore(engine, groups) - - key := store.HashFieldKey([]byte("zulu"), []byte("field")) - require.NoError(t, groups[1].Store.PutAt(ctx, key, []byte("legacy"), 5, 0)) - - value, err := st.GetAt(ctx, key, 5) - require.NoError(t, err) - require.Equal(t, []byte("legacy"), value) - - ts, exists, err := st.LatestCommitTS(ctx, key) - require.NoError(t, err) - require.True(t, exists) - require.Equal(t, uint64(5), ts) - - prefix := store.HashFieldScanPrefix([]byte("zulu")) - kvs, err := st.ScanAt(ctx, prefix, prefixScanEnd(prefix), 10, 5) - require.NoError(t, err) - require.Len(t, kvs, 1) - require.Equal(t, []byte("legacy"), kvs[0].Value) - - require.NoError(t, st.PutAt(ctx, key, []byte("current"), 6, 0)) - value, err = st.GetAt(ctx, key, 6) - require.NoError(t, err) - require.Equal(t, []byte("current"), value) - - ts, exists, err = st.LatestCommitTS(ctx, key) - require.NoError(t, err) - require.True(t, exists) - require.Equal(t, uint64(6), ts) - - kvs, err = st.ScanAt(ctx, prefix, prefixScanEnd(prefix), 10, 6) - require.NoError(t, err) - require.Len(t, kvs, 1) - require.Equal(t, []byte("current"), kvs[0].Value) - - kvs, err = st.ReverseScanAt(ctx, prefix, prefixScanEnd(prefix), 10, 6) - require.NoError(t, err) - require.Len(t, kvs, 1) - require.Equal(t, []byte("current"), kvs[0].Value) - - require.NoError(t, st.DeleteAt(ctx, key, 7)) - _, err = st.GetAt(ctx, key, 7) - require.ErrorIs(t, err, store.ErrKeyNotFound) - - kvs, err = st.ScanAt(ctx, prefix, prefixScanEnd(prefix), 10, 7) - require.NoError(t, err) - require.Empty(t, kvs) - - kvs, err = st.ReverseScanAt(ctx, prefix, prefixScanEnd(prefix), 10, 7) - require.NoError(t, err) - require.Empty(t, kvs) - - require.NoError(t, st.PutAt(ctx, key, []byte("future"), 9, 0)) - _, err = st.GetAt(ctx, key, 8) - require.ErrorIs(t, err, store.ErrKeyNotFound) - - kvs, err = st.ScanAt(ctx, prefix, prefixScanEnd(prefix), 10, 8) - require.NoError(t, err) - require.Empty(t, kvs) - - value, err = st.GetAt(ctx, key, 9) - require.NoError(t, err) - require.Equal(t, []byte("future"), value) -} - -func TestShardStoreRedisWideColumnScanRefillsAfterLogicalTombstones(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()}, - } - t.Cleanup(func() { - _ = groups[1].Store.Close() - _ = groups[2].Store.Close() - }) - st := NewShardStore(engine, groups) - - userKey := []byte("zulu") - a := store.HashFieldKey(userKey, []byte("a")) - b := store.HashFieldKey(userKey, []byte("b")) - c := store.HashFieldKey(userKey, []byte("c")) - d := store.HashFieldKey(userKey, []byte("d")) - for _, item := range []struct { - key []byte - value []byte - }{ - {key: a, value: []byte("legacy-a")}, - {key: b, value: []byte("legacy-b")}, - {key: c, value: []byte("legacy-c")}, - {key: d, value: []byte("legacy-d")}, - } { - require.NoError(t, groups[1].Store.PutAt(ctx, item.key, item.value, 5, 0)) - } - require.NoError(t, st.DeleteAt(ctx, a, 7)) - require.NoError(t, st.DeleteAt(ctx, b, 7)) - - prefix := store.HashFieldScanPrefix(userKey) - kvs, err := st.ScanAt(ctx, prefix, prefixScanEnd(prefix), 2, 7) - require.NoError(t, err) - require.Len(t, kvs, 2) - require.Equal(t, c, kvs[0].Key) - require.Equal(t, []byte("legacy-c"), kvs[0].Value) - require.Equal(t, d, kvs[1].Key) - require.Equal(t, []byte("legacy-d"), kvs[1].Value) - - keys, err := st.ScanKeysAt(ctx, prefix, prefixScanEnd(prefix), 2, 7) - require.NoError(t, err) - require.Equal(t, [][]byte{c, d}, keys) -} - -func TestShardStoreReverseRedisWideColumnScanRefillsAfterLogicalTombstones(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()}, - } - t.Cleanup(func() { - _ = groups[1].Store.Close() - _ = groups[2].Store.Close() - }) - st := NewShardStore(engine, groups) - - userKey := []byte("zulu") - a := store.HashFieldKey(userKey, []byte("a")) - b := store.HashFieldKey(userKey, []byte("b")) - c := store.HashFieldKey(userKey, []byte("c")) - d := store.HashFieldKey(userKey, []byte("d")) - for _, item := range []struct { - key []byte - value []byte - }{ - {key: a, value: []byte("legacy-a")}, - {key: b, value: []byte("legacy-b")}, - {key: c, value: []byte("legacy-c")}, - {key: d, value: []byte("legacy-d")}, - } { - require.NoError(t, groups[1].Store.PutAt(ctx, item.key, item.value, 5, 0)) - } - require.NoError(t, st.DeleteAt(ctx, d, 7)) - require.NoError(t, st.DeleteAt(ctx, c, 7)) - - prefix := store.HashFieldScanPrefix(userKey) - kvs, err := st.ReverseScanAt(ctx, prefix, prefixScanEnd(prefix), 2, 7) - require.NoError(t, err) - require.Len(t, kvs, 2) - require.Equal(t, b, kvs[0].Key) - require.Equal(t, []byte("legacy-b"), kvs[0].Value) - require.Equal(t, a, kvs[1].Key) - require.Equal(t, []byte("legacy-a"), kvs[1].Value) -} - -func TestShardStoreReverseRedisWideColumnScanPrefersLogicalRoute(t *testing.T) { - t.Parallel() - - ctx := context.Background() - engine := distribution.NewEngine() - engine.UpdateRoute([]byte(""), []byte("m"), 1) - engine.UpdateRoute([]byte("m"), nil, 2) + engine.UpdateRoute([]byte("am"), nil, 2) groups := map[uint64]*ShardGroup{ 1: {Store: store.NewMVCCStore()}, 2: {Store: store.NewMVCCStore()}, @@ -1991,15 +2926,11 @@ func TestShardStoreReverseRedisWideColumnScanPrefersLogicalRoute(t *testing.T) { }) st := NewShardStore(engine, groups) - key := store.HashFieldKey([]byte("zulu"), []byte("field")) - require.NoError(t, groups[1].Store.PutAt(ctx, key, []byte("legacy"), 5, 0)) - require.NoError(t, st.PutAt(ctx, key, []byte("current"), 6, 0)) - - prefix := store.HashFieldScanPrefix([]byte("zulu")) - kvs, err := st.ReverseScanAt(ctx, prefix, prefixScanEnd(prefix), 10, 6) - require.NoError(t, err) - require.Len(t, kvs, 1) - require.Equal(t, []byte("current"), kvs[0].Value) + start := store.HashFieldScanPrefix([]byte("alice")) + routes, clamp, _ := st.routesForScanWithVersion(start, prefixScanEnd(start)) + require.False(t, clamp) + require.Len(t, routes, 1) + require.Equal(t, uint64(1), routes[0].GroupID) } func TestShardStoreScanAt_RoutesFilesystemChunkScansByChunkRouteKey(t *testing.T) { @@ -2810,693 +3741,417 @@ func TestScanLockBoundsForKVs_ReverseInternalOnlyPageUsesOriginalRange(t *testin require.Equal(t, []byte("z"), lockEnd) } -// Redis list-delta/claim and stream rows are placed by their raw key but are -// treated as owned by the logical user key by route-bound scans and by -// prefix-write floors. Point writes checked only the raw-key route, so a fenced -// user key still accepted its auxiliary rows. -func TestShardStorePointWriteChecksRedisAuxiliaryLogicalRouteFloor(t *testing.T) { +// Reverse-scan counterpart of TestShardStoreS3BucketAuxiliaryScanHonorsStagedTombstone. +// A staged tombstone must hide the stale live row from the old raw route in both +// scan directions, not just forward. +func TestShardStoreS3BucketAuxiliaryReverseScanHonorsStagedTombstone(t *testing.T) { t.Parallel() ctx := context.Background() + const migratedBucket = "bucket-a" + routeStart := s3keys.RoutePrefixForBucketAnyGeneration(migratedBucket) + routeEnd := prefixScanEnd(routeStart) engine := distribution.NewEngine() require.NoError(t, engine.ApplySnapshot(distribution.CatalogSnapshot{ Version: 1, Routes: []distribution.RouteDescriptor{ - // Raw "!..." keys sort below "m" and land on the unfenced route. - {RouteID: 1, Start: []byte(""), End: []byte("m"), GroupID: 1, State: distribution.RouteStateActive}, - // The logical user key "zulu" lands on the fenced route. - {RouteID: 2, Start: []byte("m"), GroupID: 2, State: distribution.RouteStateActive, MinWriteTSExclusive: 100}, + {RouteID: 1, Start: []byte(""), End: routeStart, GroupID: 1, State: distribution.RouteStateActive}, + {RouteID: 2, Start: routeStart, End: routeEnd, GroupID: 2, State: distribution.RouteStateActive, StagedVisibilityActive: true, MigrationJobID: 9}, + {RouteID: 3, Start: routeEnd, End: nil, GroupID: 1, State: distribution.RouteStateActive}, }, })) groups := map[uint64]*ShardGroup{ 1: {Store: store.NewMVCCStore()}, 2: {Store: store.NewMVCCStore()}, } - t.Cleanup(func() { - _ = groups[1].Store.Close() - _ = groups[2].Store.Close() - }) st := NewShardStore(engine, groups) + deletedKey := s3keys.BucketMetaKey(migratedBucket) + visibleKey := s3keys.BucketMetaKey("bucket-z") + require.NoError(t, groups[1].Store.PutAt(ctx, deletedKey, []byte("stale"), 10, 0)) + require.NoError(t, groups[1].Store.PutAt(ctx, visibleKey, []byte("visible"), 10, 0)) + require.NoError(t, groups[2].Store.DeleteAt(ctx, distribution.MigrationStagedDataKey(9, deletedKey), 20)) - fenced := []byte("zulu") - unfenced := []byte("alpha") - - tests := []struct { - name string - key []byte - commitTS uint64 - wantReject bool - }{ - { - name: "list delta under a fenced user key", - key: store.ListMetaDeltaKey(fenced, 10, 0), - commitTS: 100, - wantReject: true, - }, - { - name: "list claim under a fenced user key", - key: store.ListClaimKey(fenced, 1), - commitTS: 100, - wantReject: true, - }, - { - name: "stream meta under a fenced user key", - key: store.StreamMetaKey(fenced), - commitTS: 100, - wantReject: true, - }, - { - name: "stream entry under a fenced user key", - key: store.StreamEntryKey(fenced, 123, 4), - commitTS: 100, - wantReject: true, - }, - { - name: "stream entry above the fenced floor is admitted", - key: store.StreamEntryKey(fenced, 123, 5), - commitTS: 101, - wantReject: false, - }, - { - name: "auxiliary row under an unfenced user key is admitted", - key: store.ListMetaDeltaKey(unfenced, 10, 0), - commitTS: 100, - wantReject: false, - }, - } - - for _, tt := range tests { - t.Run(tt.name, func(t *testing.T) { - err := st.PutAt(ctx, tt.key, []byte("v"), tt.commitTS, 0) - if tt.wantReject { - require.ErrorIs(t, err, store.ErrWriteConflict) - return - } - require.NoError(t, err) - }) - } -} - -// The coordinator's admission check must reject the same rows, so a fenced user -// key's auxiliary writes never reach Raft in the first place. -func TestShardedCoordinatorRejectsRedisAuxiliaryWriteUnderLogicalFloor(t *testing.T) { - t.Parallel() + start := []byte(s3keys.BucketMetaPrefix) + end := prefixScanEnd(start) - engine := distribution.NewEngine() - require.NoError(t, engine.ApplySnapshot(distribution.CatalogSnapshot{ - Version: 1, - Routes: []distribution.RouteDescriptor{ - {RouteID: 1, Start: []byte(""), End: []byte("m"), GroupID: 1, State: distribution.RouteStateActive}, - {RouteID: 2, Start: []byte("m"), GroupID: 2, State: distribution.RouteStateActive, MinWriteTSExclusive: 100}, - }, - })) - c := &ShardedCoordinator{engine: engine} + // Reverse over the whole family: the tombstoned bucket must not appear. + kvs, err := st.ReverseScanAt(ctx, start, end, 10, 30) + require.NoError(t, err) + require.Equal(t, []*store.KVPair{{Key: visibleKey, Value: []byte("visible")}}, kvs) - fenced := store.StreamEntryKey([]byte("zulu"), 123, 4) - require.ErrorIs(t, - c.ensureMutationsWriteAllowed([]*pb.Mutation{{Op: pb.Op_PUT, Key: fenced, Value: []byte("v")}}, 100), - store.ErrWriteConflict) - require.NoError(t, - c.ensureMutationsWriteAllowed([]*pb.Mutation{{Op: pb.Op_PUT, Key: fenced, Value: []byte("v")}}, 101)) + // Reverse scoped to the tombstoned bucket alone. + kvs, err = st.ReverseScanAt(ctx, deletedKey, prefixScanEnd(deletedKey), 10, 30) + require.NoError(t, err) + require.Empty(t, kvs) } -// A tombstone on the primary route must hide the legacy wide-column value even -// when a newer version sits above the read timestamp. The remote fallback probe -// used to compare only the newest commit timestamp, so latest > ts read as "not -// visible here", the point read fell through to the legacy route, and the -// snapshot read between the tombstone and the newer write resurrected the old -// value. -func TestShardStorePointReadStopsLegacyFallbackOnRemoteTombstone(t *testing.T) { +// An exact per-user-key legacy delta scan must carry RouteGroupID too. Redis +// cleanup and compaction build their deletes as {Del, pair.Key, GroupID: +// pair.RouteGroupID}; a zero GroupID routes the delete by the raw +// "!lst|meta|d|..." key instead of the logical list key, so after a split it +// lands on the wrong shard and the stale delta survives. +func TestShardStoreScanAt_ExactLegacyListDeltaScanMarksRouteGroup(t *testing.T) { t.Parallel() - const readTS = uint64(100) - - tests := []struct { - name string - versionVisible bool - versionSupported bool - wantLegacyValue bool - }{ - { - name: "leader reports a version visible at the read ts", - versionVisible: true, - versionSupported: true, - wantLegacyValue: false, - }, - { - name: "leader reports no version at or before the read ts", - versionVisible: false, - versionSupported: true, - wantLegacyValue: true, - }, - { - // Pre-upgrade peer: the probe is unanswered, so the caller stops - // legacy fallback rather than using the unsafe latest-commit - // heuristic. - name: "peer predating the probe stops legacy fallback", - versionVisible: false, - versionSupported: false, - wantLegacyValue: false, - }, - } - - for _, tt := range tests { - t.Run(tt.name, func(t *testing.T) { - t.Parallel() - - fake := &fakeRawKVServer{ - getResp: &pb.RawGetResponse{Exists: false}, - latestResp: &pb.RawLatestCommitTSResponse{ - // Newer than readTS, which is what defeats the heuristic. - Ts: readTS + 100, - Exists: true, - VersionVisible: tt.versionVisible, - VersionVisibleSupported: tt.versionSupported, - }, - } - addr, stop := startRawKVServer(t, fake) - t.Cleanup(stop) - - engine := distribution.NewEngine() - require.NoError(t, engine.ApplySnapshot(distribution.CatalogSnapshot{ - Version: 5, - Routes: []distribution.RouteDescriptor{ - // Raw "!hs|fld|..." keys sort below "m" and stay on group 2. - {RouteID: 1, Start: []byte(""), End: []byte("m"), GroupID: 2, State: distribution.RouteStateActive}, - // The logical user key "zulu" lives on group 1. - {RouteID: 2, Start: []byte("m"), GroupID: 1, State: distribution.RouteStateActive}, - }, - })) + ctx := context.Background() + st := newTwoRouteShardStoreForScanTest() + deltaValue := store.MarshalListMetaDelta(store.ListMetaDelta{LenDelta: 1}) - legacyStore := store.NewMVCCStore() - groups := map[uint64]*ShardGroup{ - 1: {Store: store.NewMVCCStore(), Engine: &stubFollowerEngine{leaderAddr: addr}}, - 2: {Store: legacyStore}, - } - st := NewShardStore(engine, groups) - t.Cleanup(func() { _ = st.Close() }) + // "right-list" sorts into the second route's group. + userKey := []byte("right-list") + key := legacyListMetaDeltaKey(userKey, 11) + require.NoError(t, st.groups[2].Store.PutAt(ctx, key, deltaValue, 1, 0)) - ctx := context.Background() - fieldKey := store.HashFieldKey([]byte("zulu"), []byte("f")) - require.NoError(t, legacyStore.PutAt(ctx, fieldKey, []byte("legacy"), 1, 0)) + scanStart := store.LegacyListMetaDeltaScanPrefix(userKey) + require.False(t, isBroadLegacyListDeltaScan(scanStart), + "this test is only meaningful for the exact-scan shape") - got, err := st.GetAt(ctx, fieldKey, readTS) - if tt.wantLegacyValue { - require.NoError(t, err) - require.Equal(t, []byte("legacy"), got) - return - } - require.ErrorIs(t, err, store.ErrKeyNotFound) - }) - } + kvs, err := st.ScanAt(ctx, scanStart, store.PrefixScanEnd(scanStart), 10, ^uint64(0)) + require.NoError(t, err) + require.Len(t, kvs, 1) + require.Equal(t, key, kvs[0].Key) + require.Equal(t, uint64(2), kvs[0].RouteGroupID, + "an exact legacy delta scan must still report the owning route group") } -// The server half: a group-scoped presence probe is answered from the group's -// own store, and requests that do not ask leave both response fields unset. -func TestShardStoreVersionExistsAtOrBeforeGroupWithReadFence(t *testing.T) { +// A coordinator that has not yet applied a promotion keeps forwarding the +// pre-cutover source group. Once StagedVisibilityActive is cleared the source's +// former range belongs to the target, and the staged-visibility rejection stops +// covering the request -- exactly while the source's pre-cutover MVCC is still +// sitting there waiting for cleanup. Serving that is a stale read, so the +// mismatch must fail closed instead. +func TestShardStoreExplicitGroupRead_FailsClosedAfterPromotionClearsStaging(t *testing.T) { t.Parallel() ctx := context.Background() engine := distribution.NewEngine() require.NoError(t, engine.ApplySnapshot(distribution.CatalogSnapshot{ - Version: 1, + Version: 2, Routes: []distribution.RouteDescriptor{ - {RouteID: 1, Start: []byte(""), GroupID: 1, State: distribution.RouteStateActive}, + // Promotion completed: the range now belongs to group 2 and the + // staged-visibility flag is gone. + {RouteID: 1, Start: []byte("a"), End: []byte("z"), GroupID: 2, State: distribution.RouteStateActive}, }, })) - st := NewShardStore(engine, map[uint64]*ShardGroup{1: {Store: store.NewMVCCStore()}}) - t.Cleanup(func() { _ = st.Close() }) + 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("b"), []byte("pre-cutover"), 10, 0)) + require.NoError(t, groups[2].Store.PutAt(ctx, []byte("b"), []byte("post-cutover"), 20, 0)) - key := []byte("k") - require.NoError(t, st.PutAt(ctx, key, []byte("v"), 50, 0)) + _, err := st.GetGroupAt(ctx, 1, []byte("b"), 25) + require.ErrorIs(t, err, ErrExplicitGroupRouteOwnerMismatch) - visible, ok, err := st.VersionExistsAtOrBeforeGroupWithReadFence(ctx, key, 1, 100, 0) - require.NoError(t, err) - require.True(t, ok) - require.True(t, visible) + _, err = st.ScanGroupAt(ctx, 1, []byte("a"), []byte("z"), 10, 25) + require.ErrorIs(t, err, ErrExplicitGroupRouteOwnerMismatch) - visible, ok, err = st.VersionExistsAtOrBeforeGroupWithReadFence(ctx, key, 1, 10, 0) - require.NoError(t, err) - require.True(t, ok) - require.False(t, visible) + _, err = st.ScanAtWithReadFence(ctx, []byte("a"), []byte("z"), 10, 25, false, 1, 0, []byte("a"), []byte("z")) + require.ErrorIs(t, err, ErrExplicitGroupRouteOwnerMismatch) - // Unknown group: authoritative "no version" rather than an unanswered probe. - visible, ok, err = st.VersionExistsAtOrBeforeGroupWithReadFence(ctx, key, 99, 100, 0) + // The group the catalog does name still serves the post-cutover value. + got, err := st.GetGroupAt(ctx, 2, []byte("b"), 25) require.NoError(t, err) - require.True(t, ok) - require.False(t, visible) + require.Equal(t, []byte("post-cutover"), got) } -// fenceGroupRouter is a Coordinator that only implements the group routing the -// fence dedup path uses, resolving exactly the way production does: normalize -// the key first, then look the route up. -type fenceGroupRouter struct { - Coordinator - engine *distribution.Engine -} - -func (f *fenceGroupRouter) EngineGroupIDForKey(key []byte) uint64 { - route, ok := f.engine.GetRoute(routeKey(key)) - if !ok { - return 0 - } - return route.GroupID -} - -// The fence's representative keys are re-normalized by every downstream -// consumer (LeaseReadGroupKeys -> EngineGroupIDForKey -> ResolveGroup -> -// routeKey). For a Redis wide-column range the owner key and the legacy -// raw-prefix key both normalize to the same user key, so resolving groups from -// bytes collapses the two into one and leaves the legacy group unfenced. The -// group id has to survive on the target instead. -func TestReadFenceTargetsSurviveGroupDedup(t *testing.T) { +// SQS resolves its owning group through the (queue, partition) resolver rather +// than the byte-range catalog, so a catalog route naming another group must not +// reject those reads. +func TestShardStoreExplicitGroupRead_AllowsResolverOwnedKeysOnMismatch(t *testing.T) { t.Parallel() + ctx := context.Background() engine := distribution.NewEngine() - engine.UpdateRoute([]byte(""), []byte("m"), 1) - engine.UpdateRoute([]byte("m"), nil, 2) - st := NewShardStore(engine, map[uint64]*ShardGroup{1: {}, 2: {}}) + engine.UpdateRoute([]byte(""), nil, 1) + groups := map[uint64]*ShardGroup{ + 1: {Store: store.NewMVCCStore()}, + 42: {Store: store.NewMVCCStore()}, + } + st := NewShardStore(engine, groups) - prefix := store.HashFieldScanPrefix([]byte("zebra")) - targets := st.ReadFenceTargetsForRange(prefix, store.PrefixScanEnd(prefix)) - require.Len(t, targets, 2, "owner group plus legacy raw-prefix group") + key := []byte("!sqs|msg|data|p|orders|partition-2|msg-2") + require.NoError(t, groups[42].Store.PutAt(ctx, key, []byte("payload"), 7, 0)) - carried := make([]uint64, 0, len(targets)) - reResolved := make([]uint64, 0, len(targets)) - for _, target := range targets { - carried = append(carried, target.GroupID) - route, ok := engine.GetRoute(routeKey(target.Key)) - require.True(t, ok) - reResolved = append(reResolved, route.GroupID) - } + got, err := st.GetGroupAt(ctx, 42, key, 7) + require.NoError(t, err) + require.Equal(t, []byte("payload"), got) - require.ElementsMatch(t, []uint64{2, 1}, carried, - "the fence must name both groups") - // Re-deriving from bytes is exactly what loses the legacy group; asserting - // it here pins why GroupID is carried rather than recomputed. - require.Equal(t, []uint64{2, 2}, reResolved, - "both representative keys normalize to the owner, so bytes alone are lossy") + start := []byte("!sqs|msg|data|p|orders|partition-2|") + kvs, err := st.ScanGroupAt(ctx, 42, start, prefixScanEnd(start), 10, 7) + require.NoError(t, err) + require.Equal(t, []*store.KVPair{{Key: key, Value: []byte("payload")}}, kvs) } -// LeaseReadGroupTargets must keep both groups, where the key-only -// LeaseReadGroupKeys collapses them. -func TestLeaseReadGroupTargetsKeepsLegacyGroup(t *testing.T) { +// Filesystem placement stats scan the whole chunk keyspace once per filesystem +// group, so most of those groups are not the catalog owner of the range. The +// explicit-group gate must let them through the way it lets SQS through. +func TestShardStoreExplicitGroupScan_AllowsFilesystemChunkKeys(t *testing.T) { t.Parallel() + ctx := context.Background() engine := distribution.NewEngine() - engine.UpdateRoute([]byte(""), []byte("m"), 1) - engine.UpdateRoute([]byte("m"), nil, 2) - st := NewShardStore(engine, map[uint64]*ShardGroup{1: {}, 2: {}}) - router := &fenceGroupRouter{engine: engine} - - prefix := store.HashFieldScanPrefix([]byte("zebra")) - targets := st.ReadFenceTargetsForRange(prefix, store.PrefixScanEnd(prefix)) - - kept := LeaseReadGroupTargets(router, targets) - keptGroups := make([]uint64, 0, len(kept)) - for _, target := range kept { - keptGroups = append(keptGroups, target.GroupID) + engine.UpdateRoute([]byte(""), nil, 1) + groups := map[uint64]*ShardGroup{ + 1: {Store: store.NewMVCCStore()}, + 2: {Store: store.NewMVCCStore()}, } - require.ElementsMatch(t, []uint64{2, 1}, keptGroups, - "both fenced groups must survive dedup") + st := NewShardStore(engine, groups) - collapsed := LeaseReadGroupKeys(router, st.ReadFenceGroupKeysForRange(prefix, store.PrefixScanEnd(prefix))) - require.Len(t, collapsed, 1, - "the key-only path collapses to one group; this is the gap targets close") -} + chunkKey := fskeys.ChunkKey(3, 7, 0) + require.NoError(t, groups[2].Store.PutAt(ctx, chunkKey, []byte("chunk"), 7, 0)) -// syntheticGroupRouter mimics Coordinate: EngineGroupIDForKey returns a constant -// that exists only to collapse single-group deployments to one lease. It is not -// a real group id and no group map contains it. -type syntheticGroupRouter struct { - Coordinator + start := fskeys.ChunkAllPrefix() + kvs, err := st.ScanGroupAt(ctx, 2, start, prefixScanEnd(start), 10, 7) + require.NoError(t, err) + require.Equal(t, []*store.KVPair{{Key: chunkKey, Value: []byte("chunk")}}, kvs) } -func (syntheticGroupRouter) EngineGroupIDForKey([]byte) uint64 { return 1 } - -// A point key carries GroupID 0, meaning "resolve from Key". Dedup must not -// stamp the resolved id onto it: on a single-group Coordinate that id is -// synthetic, and a later group-routed lease read or leader check would look it -// up in a group map that has never heard of it and fail closed with -// ErrLeaderNotFound. -func TestLeaseReadGroupTargetsKeepsKeyResolutionForPointKeys(t *testing.T) { +// Between cutover and promotion a key can be visible through its staged alias +// while the live key holds nothing. Both store implementations read the live key +// first and return ErrKeyNotFound when it is absent, so an expiration issued in +// that window failed for a value the same route serves happily through GetAt. +func TestExpireAtAppliesToStagedOnlyValues(t *testing.T) { t.Parallel() - targets := []ReadFenceTarget{{Key: []byte("!redis|str|k")}} - got := LeaseReadGroupTargets(syntheticGroupRouter{}, targets) - - require.Len(t, got, 1) - require.Zero(t, got[0].GroupID, - "a key-resolved target must stay key-resolved through dedup") - require.Equal(t, []byte("!redis|str|k"), got[0].Key) -} + ctx := context.Background() + st, group := newStagedVisibilityShardStore(t) -// A DynamoDB DEL_PREFIX cleanup must be checked against the same route its rows -// use. Rows route through !ddb|route|table|
; if the prefix keeps -// resolving through the raw !ddb|item| interval, a table whose route carries a -// migration write floor is skipped entirely and the cleanup installs tombstones -// below already-migrated versions, leaving the deleted generation's rows behind. -func TestEnsurePrefixWriteAllowedRoutesDynamoPrefixThroughTableRoute(t *testing.T) { - t.Parallel() + key := []byte("b") + require.NoError(t, group.Store.PutAt(ctx, + distribution.MigrationStagedDataKey(9, key), []byte("staged-b"), 20, 0)) - const ( - fencedTable = "orders" - freeTable = "carts" - ) - tableRoute := func(name string) []byte { - return dynamoRouteKey([]byte(DynamoItemPrefix + name + "|1|pk")) - } - fencedRoute := tableRoute(fencedTable) - require.NotNil(t, fencedRoute) + // The value is visible even though the live key has nothing. + got, err := st.GetAt(ctx, key, 25) + require.NoError(t, err) + require.Equal(t, []byte("staged-b"), got) - engine := distribution.NewEngine() - // Route 1 owns the raw !ddb|item| interval and carries no floor. Route 2 - // owns the fenced table's logical route key and does. - require.NoError(t, engine.ApplySnapshot(distribution.CatalogSnapshot{ - Version: 1, - Routes: []distribution.RouteDescriptor{ - {RouteID: 1, Start: []byte(""), End: fencedRoute, GroupID: 1, State: distribution.RouteStateActive}, - {RouteID: 2, Start: fencedRoute, End: nil, GroupID: 2, State: distribution.RouteStateActive, MinWriteTSExclusive: 100}, - }, - })) - st := &ShardStore{engine: engine} + // The expiry is in the future relative to the commit timestamp, so the value + // survives; the point is that ExpireAt no longer fails outright. + require.NoError(t, st.ExpireAt(ctx, key, 5_000, 300)) - itemPrefix := []byte(DynamoItemPrefix + fencedTable + "|1|") - require.ErrorIs(t, - st.ensurePrefixWriteAllowed(itemPrefix, 50), - store.ErrWriteConflict, - "a cleanup at or below the table route's floor must be rejected", - ) - require.NoError(t, st.ensurePrefixWriteAllowed(itemPrefix, 101)) + got, err = st.GetAt(ctx, key, 300) + require.NoError(t, err) + require.Equal(t, []byte("staged-b"), got) - gsiPrefix := []byte(DynamoGSIPrefix + fencedTable + "|1|") - require.ErrorIs(t, st.ensurePrefixWriteAllowed(gsiPrefix, 50), store.ErrWriteConflict) + // The expiration was recorded as a live MVCC version, which is where every + // other post-cutover write goes. + live, err := group.Store.GetAt(ctx, key, 300) + require.NoError(t, err) + require.Equal(t, []byte("staged-b"), live) - // A table whose logical route has no floor is unaffected. - require.NotNil(t, tableRoute(freeTable)) - require.NoError(t, st.ensurePrefixWriteAllowed([]byte(DynamoItemPrefix+freeTable+"|1|"), 50)) + // And it takes effect once the read passes the expiry. + _, err = st.GetAt(ctx, key, 6_000) + require.ErrorIs(t, err, store.ErrKeyNotFound) } -// A prefix that stops before the table terminator can still match several -// tables, so it must keep fanning out rather than collapsing onto one route. -func TestEnsurePrefixWriteAllowedFansOutPartialDynamoPrefixes(t *testing.T) { +func TestExpireAtUsesNewerStagedValueOverLiveValue(t *testing.T) { t.Parallel() - fencedRoute := dynamoRouteKey([]byte(DynamoItemPrefix + "orders|1|pk")) - require.NotNil(t, fencedRoute) + ctx := context.Background() + st, group := newStagedVisibilityShardStore(t) - engine := distribution.NewEngine() - require.NoError(t, engine.ApplySnapshot(distribution.CatalogSnapshot{ - Version: 1, - Routes: []distribution.RouteDescriptor{ - {RouteID: 1, Start: []byte(""), End: fencedRoute, GroupID: 1, State: distribution.RouteStateActive}, - {RouteID: 2, Start: fencedRoute, End: nil, GroupID: 2, State: distribution.RouteStateActive, MinWriteTSExclusive: 100}, - }, - })) - st := &ShardStore{engine: engine} + key := []byte("b") + require.NoError(t, group.Store.PutAt(ctx, key, []byte("live-b"), 10, 0)) + require.NoError(t, group.Store.PutAt(ctx, + distribution.MigrationStagedDataKey(9, key), []byte("staged-b"), 20, 0)) - for _, prefix := range [][]byte{ - []byte(DynamoItemPrefix), - []byte(DynamoItemPrefix + "order"), - []byte("!ddb|"), - } { - require.ErrorIs(t, st.ensurePrefixWriteAllowed(prefix, 50), store.ErrWriteConflict, - "prefix %q spans the fenced table and must still see its floor", prefix) - } + require.NoError(t, st.ExpireAt(ctx, key, 5_000, 300)) + + got, err := st.GetAt(ctx, key, 300) + require.NoError(t, err) + require.Equal(t, []byte("staged-b"), got) + + live, err := group.Store.GetAt(ctx, key, 300) + require.NoError(t, err) + require.Equal(t, []byte("staged-b"), live) } -// bucketDeleteSafetyNetElems DEL_PREFIXes six raw S3 families per bucket delete, -// but every key underneath them routes through !s3route|. -// Checked against the raw interval those six miss the object routes' own write -// floors, so a cleanup at or below a floor is admitted and its tombstones sit -// hidden behind migrated object versions -- data retained for a bucket the -// operator was told is gone. -func TestEnsurePrefixWriteAllowedRoutesS3BucketPrefixesThroughObjectRoutes(t *testing.T) { +func TestExpireAtHonorsNewerStagedTombstone(t *testing.T) { t.Parallel() - const ( - bucket = "orders" - generation = uint64(3) - ) - bucketRoute := s3keys.RoutePrefixForBucket(bucket, generation) - // Every object of this bucket generation sorts under bucketRoute. - require.True(t, bytes.HasPrefix( - s3keys.ExtractRouteKey(s3keys.ObjectManifestKey(bucket, generation, "a/b")), bucketRoute)) + ctx := context.Background() + st, group := newStagedVisibilityShardStore(t) - // !s3route| sorts before !s3|, so the bucket's object route range and the - // raw !s3| families are three distinct routes. Only the middle one, which - // holds the objects, carries a floor. - engine := distribution.NewEngine() - require.NoError(t, engine.ApplySnapshot(distribution.CatalogSnapshot{ - Version: 1, - Routes: []distribution.RouteDescriptor{ - {RouteID: 1, Start: []byte(""), End: bucketRoute, GroupID: 1, State: distribution.RouteStateActive}, - {RouteID: 2, Start: bucketRoute, End: prefixScanEnd(bucketRoute), GroupID: 2, State: distribution.RouteStateActive, MinWriteTSExclusive: 100}, - {RouteID: 3, Start: prefixScanEnd(bucketRoute), End: nil, GroupID: 3, State: distribution.RouteStateActive}, - }, - })) - st := &ShardStore{engine: engine} - - for name, prefix := range map[string][]byte{ - "manifest": s3keys.ObjectManifestPrefixForBucket(bucket, generation), - "uploadMeta": s3keys.UploadMetaPrefixForBucket(bucket, generation), - "uploadPart": s3keys.UploadPartPrefixForBucket(bucket, generation), - "blob": s3keys.BlobPrefixForBucket(bucket, generation), - "chunkRef": s3keys.ChunkRefPrefixForBucket(bucket, generation), - "gcUpload": s3keys.GCUploadPrefixForBucket(bucket, generation), - } { - require.ErrorIs(t, st.ensurePrefixWriteAllowed(prefix, 50), store.ErrWriteConflict, - "%s prefix must see the object route floor", name) - require.NoError(t, st.ensurePrefixWriteAllowed(prefix, 101), "%s prefix above the floor", name) - } + key := []byte("b") + require.NoError(t, group.Store.PutAt(ctx, key, []byte("live-b"), 10, 0)) + require.NoError(t, group.Store.DeleteAt(ctx, distribution.MigrationStagedDataKey(9, key), 20)) - // The route-space prefix in the same batch already resolved correctly and - // must keep doing so. - require.ErrorIs(t, - st.ensurePrefixWriteAllowed(s3keys.RoutePrefixForBucket(bucket, generation), 50), - store.ErrWriteConflict, - ) + require.ErrorIs(t, st.ExpireAt(ctx, key, 5_000, 300), store.ErrKeyNotFound) + + _, err := st.GetAt(ctx, key, 300) + require.ErrorIs(t, err, store.ErrKeyNotFound) + live, err := group.Store.GetAt(ctx, key, 300) + require.NoError(t, err) + require.Equal(t, []byte("live-b"), live) } -// A bucket whose object routes carry no floor is unaffected, and a prefix that -// cannot be projected still fans out over the raw interval. -func TestEnsurePrefixWriteAllowedKeepsS3FanoutForUnprojectablePrefixes(t *testing.T) { +// A key with nothing on either side still reports ErrKeyNotFound. +func TestExpireAtStillFailsWhenNothingIsVisible(t *testing.T) { t.Parallel() - bucketRoute := s3keys.RoutePrefixForBucket("orders", 3) - - // The floor is on the raw !s3| interval this time, and the bucket's object - // route range is clean. - engine := distribution.NewEngine() - require.NoError(t, engine.ApplySnapshot(distribution.CatalogSnapshot{ - Version: 1, - Routes: []distribution.RouteDescriptor{ - {RouteID: 1, Start: []byte(""), End: bucketRoute, GroupID: 1, State: distribution.RouteStateActive}, - {RouteID: 2, Start: bucketRoute, End: prefixScanEnd(bucketRoute), GroupID: 2, State: distribution.RouteStateActive}, - {RouteID: 3, Start: prefixScanEnd(bucketRoute), End: nil, GroupID: 3, State: distribution.RouteStateActive, MinWriteTSExclusive: 100}, - }, - })) - st := &ShardStore{engine: engine} - - // Projectable: only the object route range is consulted, and it has no floor. - require.NoError(t, st.ensurePrefixWriteAllowed(s3keys.ObjectManifestPrefixForBucket("orders", 3), 50)) + ctx := context.Background() + st, _ := newStagedVisibilityShardStore(t) - // Not projectable: the bare family still fans out over the raw interval, - // which does carry a floor. - require.ErrorIs(t, - st.ensurePrefixWriteAllowed([]byte(s3keys.ObjectManifestPrefix), 50), - store.ErrWriteConflict, - ) + require.ErrorIs(t, st.ExpireAt(ctx, []byte("absent"), 40, 300), store.ErrKeyNotFound) } -// countingGetStore counts the point reads canonicalization performs. -type countingGetStore struct { +// promotingExportStore runs a hook after the first ExportVersions call, which +// is how a promotion batch is landed exactly between the two probes +// getAtWithStagedVisibility makes. +type promotingExportStore struct { store.MVCCStore - gets atomic.Int64 + afterFirst func() + calls int } -func (s *countingGetStore) GetAt(ctx context.Context, key []byte, ts uint64) ([]byte, error) { - s.gets.Add(1) - return s.MVCCStore.GetAt(ctx, key, ts) +func (s *promotingExportStore) ExportVersions( + ctx context.Context, + opts store.ExportVersionsOptions, +) (store.ExportVersionsResult, error) { + res, err := s.MVCCStore.ExportVersions(ctx, opts) + s.fireAfterFirst() + return res, err } -// Every route page a scan merges has already been through -// canonicalizeRedisWideColumnScanResults in its local, leader, or proxy page -// path, and the canonical row keeps the physical key. A second pass over the -// merged result therefore repeats the point read for every surviving row: 2N -// reads for a page of N, and on a remote group a fenced RPC each. -func TestScanAtCanonicalizesWideColumnRowsOnce(t *testing.T) { - t.Parallel() - - ctx := context.Background() - userKey := []byte("user:key") - fields := [][]byte{[]byte("f1"), []byte("f2"), []byte("f3")} - - newStore := func() (*ShardStore, *countingGetStore) { - engine := distribution.NewEngine() - engine.UpdateRoute([]byte(""), nil, 1) - counting := &countingGetStore{MVCCStore: store.NewMVCCStore()} - for _, field := range fields { - require.NoError(t, counting.PutAt(ctx, store.HashFieldKey(userKey, field), []byte("v"), 10, 0)) - } - counting.gets.Store(0) - return NewShardStore(engine, map[uint64]*ShardGroup{1: {Store: counting}}), counting +// ScanAt fires the promotion when the *live* range is scanned. That is the +// only interleaving that distinguishes the two orderings: live-first means the +// live scan misses the key and the staged scan that follows misses it too, +// while staged-first has already captured it before promotion runs. +func (s *promotingExportStore) ScanAt( + ctx context.Context, + start, end []byte, + limit int, + ts uint64, +) ([]*store.KVPair, error) { + kvs, err := s.MVCCStore.ScanAt(ctx, start, end, limit, ts) + if !isMigrationStagedDataKey(start) && s.afterFirst != nil { + s.fireAfterFirst() } + return kvs, err +} - start := store.HashFieldScanPrefix(userKey) - end := prefixScanEnd(start) - - st, counting := newStore() - kvs, err := st.ScanAt(ctx, start, end, 10, 20) - require.NoError(t, err) - require.Len(t, kvs, len(fields)) - require.Equal(t, int64(len(fields)), counting.gets.Load(), - "one canonicalizing point read per returned row, not two") +// LatestCommitTS fires the promotion on the live-key probe, which is the +// interleaving that separates the two orderings for the watermark reads. +func (s *promotingExportStore) LatestCommitTS(ctx context.Context, key []byte) (uint64, bool, error) { + ts, exists, err := s.MVCCStore.LatestCommitTS(ctx, key) + if !isMigrationStagedDataKey(key) && s.afterFirst != nil { + s.fireAfterFirst() + } + return ts, exists, err +} - st, counting = newStore() - kvs, err = st.ReverseScanAt(ctx, start, end, 10, 20) - require.NoError(t, err) - require.Len(t, kvs, len(fields)) - require.Equal(t, int64(len(fields)), counting.gets.Load(), - "the reverse path had the same redundant outer pass") +func (s *promotingExportStore) fireAfterFirst() { + s.calls++ + if s.calls == 1 && s.afterFirst != nil { + s.afterFirst() + } } -// normalizeRouteKey places stream writes on the logical user-key route, so -// metadata written before that normalization sits on the physical !stream|meta| -// route. A split can put the two on different groups, and a point read that -// consults only the logical route reports an existing stream as missing -- after -// which the caller never scans, so the legacy candidate the scan path carries is -// never reached and commands re-initialize fresh metadata over live entries. -func TestGetAtReadsLegacyStreamRouteDuringRollingUpgrade(t *testing.T) { +// A promotion batch landing between the staged and live probes must not make a +// key disappear. Promotion writes the live version and drops the staged alias, +// so a live-then-staged order misses both sides of a staged-only key: live has +// not been written yet at the first probe, and the alias is gone by the second. +func TestShardStoreGetAt_StagedVisibilitySurvivesPromotionBetweenProbes(t *testing.T) { t.Parallel() ctx := context.Background() - userKey := []byte("z-stream") - metaKey := store.StreamMetaKey(userKey) - // The physical key and the logical user key land on different routes. - require.NotEqual(t, RouteKey(metaKey), metaKey) - engine := distribution.NewEngine() require.NoError(t, engine.ApplySnapshot(distribution.CatalogSnapshot{ Version: 1, - Routes: []distribution.RouteDescriptor{ - // The raw !stream| prefix sorts below "z", so the physical key is on - // the lower route and the user key on the upper one. - {RouteID: 1, Start: []byte(""), End: userKey, GroupID: 1, State: distribution.RouteStateActive}, - {RouteID: 2, Start: userKey, End: nil, GroupID: 2, State: distribution.RouteStateActive}, - }, + Routes: []distribution.RouteDescriptor{{ + RouteID: 1, Start: []byte("a"), End: []byte("z"), GroupID: 1, + State: distribution.RouteStateActive, StagedVisibilityActive: true, + MigrationJobID: 9, MinWriteTSExclusive: 100, + }}, })) - groups := map[uint64]*ShardGroup{ - 1: {Store: store.NewMVCCStore()}, - 2: {Store: store.NewMVCCStore()}, - } - t.Cleanup(func() { - for _, g := range groups { - _ = g.Store.Close() - } - }) - st := NewShardStore(engine, groups) - // Written by a node that predates the stream route normalization. - require.NoError(t, groups[1].Store.PutAt(ctx, metaKey, []byte("legacy-meta"), 10, 0)) + inner := store.NewMVCCStore() + t.Cleanup(func() { _ = inner.Close() }) + rawKey := []byte("k") + stagedKey := distribution.MigrationStagedDataKey(9, rawKey) + require.NoError(t, inner.PutAt(ctx, stagedKey, []byte("staged-only"), 20, 0)) - got, err := st.GetAt(ctx, metaKey, 20) - require.NoError(t, err, "the stream must not read as missing") - require.Equal(t, []byte("legacy-meta"), got) + promoting := &promotingExportStore{MVCCStore: inner} + promoting.afterFirst = func() { + // The promotion batch: the row becomes live at its original commit ts + // and the staged alias goes away. + require.NoError(t, inner.PutAt(ctx, rawKey, []byte("staged-only"), 20, 0)) + require.NoError(t, inner.DeleteAt(ctx, stagedKey, 21)) + } + st := NewShardStore(engine, map[uint64]*ShardGroup{1: {Store: promoting}}) + + got, err := st.GetAt(ctx, rawKey, 25) + require.NoError(t, err, "a promotion between the probes must not hide the key") + require.Equal(t, []byte("staged-only"), got) } -// A stream entry key gets the same treatment, and the logical route still wins -// once a value is there. -func TestGetAtPrefersLogicalStreamRouteWhenBothExist(t *testing.T) { +// Every place that reads the live and staged namespaces as two separate store +// calls has to read staged first, for the reason getAtWithStagedVisibility +// documents. Fixing only the point read left the scan, the TTL winner, and the +// watermark on the old order, each with its own way of losing the key. +func TestStagedVisibilityProbesReadStagedFirst(t *testing.T) { t.Parallel() ctx := context.Background() - userKey := []byte("z-stream") - entryKey := store.StreamEntryKey(userKey, 123, 4) - engine := distribution.NewEngine() require.NoError(t, engine.ApplySnapshot(distribution.CatalogSnapshot{ Version: 1, - Routes: []distribution.RouteDescriptor{ - {RouteID: 1, Start: []byte(""), End: userKey, GroupID: 1, State: distribution.RouteStateActive}, - {RouteID: 2, Start: userKey, End: nil, GroupID: 2, State: distribution.RouteStateActive}, - }, + Routes: []distribution.RouteDescriptor{{ + RouteID: 1, Start: []byte("a"), End: []byte("z"), GroupID: 1, + State: distribution.RouteStateActive, StagedVisibilityActive: true, + MigrationJobID: 9, MinWriteTSExclusive: 100, + }}, })) - groups := map[uint64]*ShardGroup{ - 1: {Store: store.NewMVCCStore()}, - 2: {Store: store.NewMVCCStore()}, - } - t.Cleanup(func() { - for _, g := range groups { - _ = g.Store.Close() - } - }) - st := NewShardStore(engine, groups) - - require.NoError(t, groups[1].Store.PutAt(ctx, entryKey, []byte("legacy"), 10, 0)) - require.NoError(t, groups[2].Store.PutAt(ctx, entryKey, []byte("current"), 11, 0)) - - got, err := st.GetAt(ctx, entryKey, 20) - require.NoError(t, err) - require.Equal(t, []byte("current"), got) -} - -// The canonicalization predicate answers a different question -- whether a scan -// row is a physical form needing a point read -- and streams have no such form, -// so it must stay wide-column only even though the route candidate now covers -// both. -func TestLegacyPointRouteKeyIsNotTheCanonicalizationPredicate(t *testing.T) { - t.Parallel() - - userKey := []byte("user:key") - streamKey := store.StreamMetaKey(userKey) - require.NotNil(t, legacyPointRouteKey(streamKey)) - require.Nil(t, redisWideColumnLegacyPointRouteKey(streamKey)) - - hashKey := store.HashFieldKey(userKey, []byte("f")) - require.NotNil(t, legacyPointRouteKey(hashKey)) - require.NotNil(t, redisWideColumnLegacyPointRouteKey(hashKey)) - - require.Nil(t, legacyPointRouteKey([]byte("plain"))) -} - -// A proxied page cannot be assumed canonical. The peer may be running the parent -// binary, which accepts group_id on RawScanAt and serves rawScanAtExplicitGroup -// while ignoring the fence and route-bounds fields, and has no canonicalization -// at all -- so it answers with physical rows. Dropping the local pass would let a -// legacy hash/set/zset row through that a logical tombstone should suppress. -func TestScanAtCanonicalizesProxiedPages(t *testing.T) { - t.Parallel() - userKey := []byte("user:key") - rows := make([]*pb.RawKVPair, 0, 3) - for _, field := range [][]byte{[]byte("f1"), []byte("f2"), []byte("f3")} { - rows = append(rows, &pb.RawKVPair{Key: store.HashFieldKey(userKey, field), Value: []byte("v")}) + rawKey := []byte("k") + stagedKey := distribution.MigrationStagedDataKey(9, rawKey) + + newStore := func(t *testing.T) (*ShardStore, *promotingExportStore, store.MVCCStore) { + t.Helper() + inner := store.NewMVCCStore() + t.Cleanup(func() { _ = inner.Close() }) + require.NoError(t, inner.PutAt(ctx, stagedKey, []byte("staged-only"), 20, 0)) + promoting := &promotingExportStore{MVCCStore: inner} + promoting.afterFirst = func() { + // The real promotion batch: PromoteVersions moves the staged row to + // its live key and removes the staged version physically. Modelling + // it with a Delete would write a tombstone that legitimately hides + // the key, which is a different scenario. + promoter, ok := inner.(store.MigrationPromoter) + require.True(t, ok) + _, err := promoter.PromoteVersions(ctx, store.PromoteVersionsOptions{ + JobID: 9, + StartKey: distribution.MigrationStagedDataKeyPrefix(9), + EndKey: prefixScanEnd(distribution.MigrationStagedDataKeyPrefix(9)), + MaxVersions: 16, + TargetKey: func(staged []byte) ([]byte, bool) { + _, raw, ok := distribution.MigrationStagedDataKeyParts(staged) + return raw, ok + }, + }) + require.NoError(t, err) + } + return NewShardStore(engine, map[uint64]*ShardGroup{1: {Store: promoting}}), promoting, inner } - fake := &fakeRawKVServer{scanResp: &pb.RawScanAtResponse{Kv: rows}} - addr, stop := startRawKVServer(t, fake) - t.Cleanup(stop) - engine := distribution.NewEngine() - require.NoError(t, engine.ApplySnapshot(distribution.CatalogSnapshot{ - Version: 1, - Routes: []distribution.RouteDescriptor{ - {RouteID: 1, Start: []byte(""), GroupID: 1, State: distribution.RouteStateActive}, - }, - })) - st := NewShardStore(engine, map[uint64]*ShardGroup{ - 1: {Store: store.NewMVCCStore(), Engine: &stubFollowerEngine{leaderAddr: addr}}, + t.Run("scan", func(t *testing.T) { + t.Parallel() + st, _, _ := newStore(t) + kvs, err := st.ScanAt(ctx, []byte("a"), []byte("z"), 10, 25) + require.NoError(t, err) + require.Len(t, kvs, 1, "a promotion between the scans must not drop the key") + require.Equal(t, rawKey, kvs[0].Key) }) - t.Cleanup(func() { _ = st.Close() }) - - start := store.HashFieldScanPrefix(userKey) - _, err := st.ScanAt(context.Background(), start, prefixScanEnd(start), 10, 20) - require.NoError(t, err) - fake.mu.Lock() - defer fake.mu.Unlock() - require.Equal(t, len(rows), fake.getCalls, - "every proxied row must be canonicalized locally until the peer can say it already did") + t.Run("point read", func(t *testing.T) { + t.Parallel() + st, _, _ := newStore(t) + got, err := st.GetAt(ctx, rawKey, 25) + require.NoError(t, err) + require.Equal(t, []byte("staged-only"), got) + }) } func (e *followerProxyEngine) SnapshotEvery() uint64 { return 0 } diff --git a/kv/sharded_coordinator.go b/kv/sharded_coordinator.go index 500b78b3b..a15a48eff 100644 --- a/kv/sharded_coordinator.go +++ b/kv/sharded_coordinator.go @@ -799,13 +799,6 @@ func (c *ShardedCoordinator) Dispatch(ctx context.Context, reqs *OperationGroup[ return nil, err } - // DEL_PREFIX cannot be routed to a single shard because the prefix may - // span multiple shards (or be nil, meaning "all keys"). Broadcast the - // operation to every shard group so each FSM scans locally. - if hasDelPrefixElem(reqs.Elems) { - return c.dispatchDelPrefixBroadcast(ctx, reqs.IsTxn, reqs.Elems) - } - // Capture whether the caller supplied a non-zero StartTS BEFORE // the coordinator-allocates-on-zero branch below mutates the // field. A caller-supplied StartTS names a specific snapshot @@ -837,10 +830,27 @@ func (c *ShardedCoordinator) Dispatch(ctx context.Context, reqs *OperationGroup[ } if reqs.IsTxn { + if resp, handled, err := c.dispatchBeforeShardRouting(ctx, reqs); handled { + return resp, err + } return c.dispatchTxnWithComposed1Retry(ctx, reqs, callerSuppliedStartTS) } - return c.dispatchNonTxn(ctx, reqs) + return c.dispatchRawWithComposed1Retry(ctx, reqs) +} + +func (c *ShardedCoordinator) dispatchBeforeShardRouting(ctx context.Context, reqs *OperationGroup[OP]) (*CoordinateResponse, bool, error) { + // DEL_PREFIX cannot be routed to a single shard because the prefix may + // span multiple shards (or be nil, meaning "all keys"). Broadcast the + // operation to every shard group so each FSM scans locally. + if hasDelPrefixElem(reqs.Elems) { + resp, err := c.dispatchDelPrefixBroadcast(ctx, reqs.IsTxn, reqs.Elems, reqs.ObservedRouteVersion) + return resp, true, err + } + if err := c.rejectWriteFencedPointElems(reqs.Elems); err != nil { + return nil, true, err + } + return nil, false, nil } // dispatchTxnWithComposed1Retry runs the M4 Composed-1 retry loop @@ -946,47 +956,6 @@ func (c *ShardedCoordinator) anyResolverClaimedKey(elems []*Elem[OP]) bool { return false } -func (c *ShardedCoordinator) partitionResolverClaimsPointKey(key []byte) bool { - if c == nil || c.router == nil || c.router.partitionResolver == nil || len(key) == 0 { - return false - } - _, ok := c.router.partitionResolver.ResolveGroup(key) - return ok -} - -func (c *ShardedCoordinator) writeFenceBypassKeysForElems(elems []*Elem[OP]) [][]byte { - if len(elems) == 0 { - return nil - } - out := make([][]byte, 0, len(elems)) - for _, elem := range elems { - if elem == nil || elem.Op == DelPrefix || (elem.GroupID == 0 && !c.partitionResolverClaimsPointKey(elem.Key)) { - continue - } - out = append(out, bytes.Clone(elem.Key)) - } - return out -} - -func (c *ShardedCoordinator) writeFenceBypassKeysByGroup(elems []*Elem[OP]) map[uint64][][]byte { - out := make(map[uint64][][]byte) - for _, elem := range elems { - if elem == nil || elem.Op == DelPrefix { - continue - } - gid := elem.GroupID - if gid == 0 { - var ok bool - gid, ok = c.router.ResolveGroup(elem.Key) - if !ok || !c.partitionResolverClaimsPointKey(elem.Key) { - continue - } - } - out[gid] = append(out[gid], bytes.Clone(elem.Key)) - } - return out -} - func (c *ShardedCoordinator) dispatchTxnWithComposed1Retry(ctx context.Context, reqs *OperationGroup[OP], callerSuppliedStartTS bool) (*CoordinateResponse, error) { c.maybeAutoPinObservedRouteVersion(reqs, callerSuppliedStartTS) @@ -1113,6 +1082,9 @@ func isComposed1RetryableError(err error) bool { // shard router. Extracted from Dispatch to keep that method's branch // count within the cyclop budget after the 7a registration gate landed. func (c *ShardedCoordinator) dispatchNonTxn(ctx context.Context, reqs *OperationGroup[OP]) (*CoordinateResponse, error) { + if err := c.rejectWriteFencedPointElems(reqs.Elems); err != nil { + return nil, err + } if hasExplicitGroupElem(reqs.Elems) { return c.dispatchExplicitGroupNonTxn(ctx, reqs) } @@ -1158,6 +1130,61 @@ func (c *ShardedCoordinator) dispatchExplicitGroupNonTxn(ctx context.Context, re return &CoordinateResponse{CommitIndex: maxIndex}, nil } +func (c *ShardedCoordinator) dispatchRawWithComposed1Retry(ctx context.Context, reqs *OperationGroup[OP]) (*CoordinateResponse, error) { + for attempt := 0; attempt <= composed1RetryAttempts; attempt++ { + resp, handled, err := c.dispatchBeforeShardRouting(ctx, reqs) + if !handled { + resp, err = c.dispatchNonTxn(ctx, reqs) + } + if err == nil { + return resp, nil + } + if !errors.Is(err, ErrComposed1VersionGCd) || attempt == composed1RetryAttempts || c.engine == nil { + return resp, err + } + if !c.canRetryRawVersionGC(reqs) { + return resp, err + } + reqs.ObservedRouteVersion = c.engine.Version() + } + return nil, errors.WithStack(ErrInvalidRequest) +} + +func (c *ShardedCoordinator) canRetryRawVersionGC(reqs *OperationGroup[OP]) bool { + if c == nil || c.router == nil || reqs == nil || hasDelPrefixElem(reqs.Elems) { + return false + } + var ( + firstGID uint64 + seen bool + ) + for _, elem := range reqs.Elems { + gid, ok := c.rawElemGroupID(elem) + if !ok { + return false + } + if !seen { + firstGID = gid + seen = true + continue + } + if gid != firstGID { + return false + } + } + return seen +} + +func (c *ShardedCoordinator) rawElemGroupID(elem *Elem[OP]) (uint64, bool) { + if elem == nil { + return 0, false + } + if elem.GroupID != 0 { + return elem.GroupID, true + } + return c.router.ResolveGroup(elem.Key) +} + // hasDelPrefixElem returns true if any element is a DelPrefix operation. func hasDelPrefixElem(elems []*Elem[OP]) bool { for _, e := range elems { @@ -1185,18 +1212,24 @@ func validateDelPrefixOnly(elems []*Elem[OP]) error { // pb.Request (the FSM's extractDelPrefix processes only the first DEL_PREFIX // mutation per request). All requests are batched into a single Commit call // per shard group. -func (c *ShardedCoordinator) dispatchDelPrefixBroadcast(ctx context.Context, isTxn bool, elems []*Elem[OP]) (*CoordinateResponse, error) { +func (c *ShardedCoordinator) dispatchDelPrefixBroadcast(ctx context.Context, isTxn bool, elems []*Elem[OP], observedRouteVersion uint64) (*CoordinateResponse, error) { if isTxn { return nil, errors.Wrap(ErrInvalidRequest, "DEL_PREFIX not supported in transactions") } if err := validateDelPrefixOnly(elems); err != nil { return nil, err } + if err := c.rejectWriteFencedDelPrefixes(elems); err != nil { + return nil, err + } ts, err := c.allocateTimestamp(ctx, "allocate DEL_PREFIX broadcast ts") if err != nil { return nil, err } + if err := c.rejectWriteTimestampFloorDelPrefixes(elems, ts); err != nil { + return nil, err + } requests := make([]*pb.Request, 0, len(elems)) for _, elem := range elems { mut := elemToMutation(elem) @@ -1204,16 +1237,197 @@ func (c *ShardedCoordinator) dispatchDelPrefixBroadcast(ctx context.Context, isT return nil, err } requests = append(requests, &pb.Request{ - IsTxn: false, - Phase: pb.Phase_NONE, - Ts: ts, - Mutations: []*pb.Mutation{mut}, + IsTxn: false, + Phase: pb.Phase_NONE, + Ts: ts, + Mutations: []*pb.Mutation{mut}, + ObservedRouteVersion: observedRouteVersion, }) } return c.broadcastToAllGroups(ctx, requests) } +func (c *ShardedCoordinator) rejectWriteFencedPointElems(elems []*Elem[OP]) error { + if c == nil || c.engine == nil { + return nil + } + for _, elem := range elems { + if elem == nil || elem.GroupID != 0 { + continue + } + if err := c.rejectWriteFencedPointKey(elem.Key); err != nil { + return err + } + } + return nil +} + +func (c *ShardedCoordinator) rejectWriteFencedPointKey(key []byte) error { + if c.partitionResolverRecognisesPointKey(key) { + return nil + } + start, end, ok := s3BucketAuxiliaryRouteRange(key) + if ok { + route, found := s3BucketAuxiliaryOwnerRouteFromRange(start, end, c.engine.GetIntersectingRoutes(start, end)) + if found && route.State == distribution.RouteStateWriteFenced { + return errors.Wrapf(ErrRouteWriteFenced, "key %q route range [%q,%q)", key, start, end) + } + return nil + } + rkey := routeKey(key) + if route, ok := c.engine.GetRoute(rkey); ok && route.State == distribution.RouteStateWriteFenced { + return errors.Wrapf(ErrRouteWriteFenced, "key %q routeKey %q", key, rkey) + } + return nil +} + +func (c *ShardedCoordinator) partitionResolverRecognisesPointKey(key []byte) bool { + if c == nil || c.router == nil || c.router.partitionResolver == nil || len(key) == 0 { + return false + } + if _, ok := c.router.partitionResolver.ResolveGroup(key); ok { + return true + } + return c.router.partitionResolver.RecognisesPartitionedKey(key) +} + +func (c *ShardedCoordinator) writeFenceBypassKeysForElems(elems []*Elem[OP]) [][]byte { + if len(elems) == 0 { + return nil + } + out := make([][]byte, 0, len(elems)) + for _, elem := range elems { + if elem == nil || elem.Op == DelPrefix || (elem.GroupID == 0 && !c.partitionResolverClaimsPointKey(elem.Key)) { + continue + } + out = append(out, bytes.Clone(elem.Key)) + } + return out +} + +func (c *ShardedCoordinator) writeFenceBypassKeysByGroup(elems []*Elem[OP]) map[uint64][][]byte { + out := make(map[uint64][][]byte) + for _, elem := range elems { + if elem == nil || elem.Op == DelPrefix { + continue + } + gid := elem.GroupID + if gid == 0 { + var ok bool + gid, ok = c.router.ResolveGroup(elem.Key) + if !ok || !c.partitionResolverClaimsPointKey(elem.Key) { + continue + } + } + out[gid] = append(out[gid], bytes.Clone(elem.Key)) + } + return out +} + +func (c *ShardedCoordinator) partitionResolverClaimsPointKey(key []byte) bool { + if c == nil || c.router == nil || c.router.partitionResolver == nil || len(key) == 0 { + return false + } + _, ok := c.router.partitionResolver.ResolveGroup(key) + return ok +} + +func (c *ShardedCoordinator) rejectWriteFencedDelPrefixes(elems []*Elem[OP]) error { + if c == nil || c.engine == nil { + return nil + } + for _, elem := range elems { + if elem == nil { + continue + } + start, end := routePrefixRange(elem.Key) + for _, route := range c.engine.GetIntersectingRoutes(start, end) { + if route.State == distribution.RouteStateWriteFenced { + return errors.Wrapf(ErrRouteWriteFenced, "prefix %q route range [%q,%q)", elem.Key, start, end) + } + } + } + return nil +} + +func (c *ShardedCoordinator) rejectWriteTimestampFloorPointKey(key []byte, commitTS uint64) error { + if c == nil || c.engine == nil || commitTS == 0 { + return nil + } + if c.skipPointWriteTimestampFloor(key) { + return nil + } + if err := ensureLogicalRouteWriteAllowed(c.engine, key, commitTS); err != nil { + return err + } + if checked, err := c.rejectS3BucketAuxiliaryWriteTimestampFloor(key, commitTS); checked || err != nil { + return err + } + return c.rejectRawRouteWriteTimestampFloor(key, commitTS) +} + +func (c *ShardedCoordinator) skipPointWriteTimestampFloor(key []byte) bool { + // Same exemption rejectWriteFencedPointKey applies, and for the same + // reason: in partition-resolved keyspaces such as HT-FIFO SQS, routeKey + // collapses a concrete partition key onto the global SQS route, so that + // route's floor is not this key's floor. Applying it here rejected writes + // the write-fence precheck deliberately lets through, and the Raft-side + // gate already covers these keys through the resolver's own routing. + return c.partitionResolverRecognisesPointKey(key) || isTxnInternalKey(key) +} + +func (c *ShardedCoordinator) rejectS3BucketAuxiliaryWriteTimestampFloor(key []byte, commitTS uint64) (bool, error) { + start, end, ok := s3BucketAuxiliaryRouteRange(key) + if !ok { + return false, nil + } + route, found := s3BucketAuxiliaryOwnerRouteFromRange(start, end, c.engine.GetIntersectingRoutes(start, end)) + if found && route.MinWriteTSExclusive != 0 && commitTS <= route.MinWriteTSExclusive { + return true, errors.Join( + errors.Wrapf(ErrRouteWriteTimestampTooLow, "key %q route range [%q,%q) commit_ts=%d floor=%d", key, start, end, commitTS, route.MinWriteTSExclusive), + store.NewWriteConflictError(key), + ) + } + return true, nil +} + +func (c *ShardedCoordinator) rejectRawRouteWriteTimestampFloor(key []byte, commitTS uint64) error { + rkey := routeKey(key) + if route, ok := c.engine.GetRoute(rkey); ok && route.MinWriteTSExclusive != 0 && commitTS <= route.MinWriteTSExclusive { + return routeWriteTimestampFloorError(route, key, commitTS) + } + return nil +} + +func (c *ShardedCoordinator) rejectWriteTimestampFloorDelPrefixes(elems []*Elem[OP], commitTS uint64) error { + if c == nil || c.engine == nil || commitTS == 0 { + return nil + } + for _, elem := range elems { + if elem == nil { + continue + } + if err := c.rejectWriteTimestampFloorDelPrefix(elem.Key, commitTS); err != nil { + return err + } + } + return nil +} + +func (c *ShardedCoordinator) rejectWriteTimestampFloorDelPrefix(prefix []byte, commitTS uint64) error { + start, end := routePrefixRange(prefix) + for _, route := range c.engine.GetIntersectingRoutes(start, end) { + if route.MinWriteTSExclusive != 0 && commitTS <= route.MinWriteTSExclusive { + return errors.Join( + errors.Wrapf(ErrRouteWriteTimestampTooLow, "prefix %q route range [%q,%q) commit_ts=%d floor=%d", prefix, start, end, commitTS, route.MinWriteTSExclusive), + store.NewWriteConflictError(prefix), + ) + } + } + return nil +} + // broadcastToAllGroups sends the same set of requests to every configured // all-shard data group in parallel and returns the maximum commit index. func (c *ShardedCoordinator) broadcastToAllGroups(ctx context.Context, requests []*pb.Request) (*CoordinateResponse, error) { @@ -1325,6 +1539,10 @@ func (c *ShardedCoordinator) dispatchMultiShardTxn(ctx context.Context, startTS, if err != nil { return nil, err } + groupedReadKeys = c.groupedReadKeysWithStagedVisibilityMutationAliases(groupedReadKeys, grouped) + if groupedReadKeyCount(groupedReadKeys) > maxReadKeys { + return nil, errors.WithStack(ErrInvalidRequest) + } prepared, err := c.prewriteTxn(ctx, startTS, commitTS, primaryKey, grouped, gids, groupedReadKeys, observedRouteVersion, bypassKeysByGroup) if err != nil { return nil, err @@ -1399,6 +1617,11 @@ func (c *ShardedCoordinator) dispatchSingleShardTxn(ctx context.Context, startTS if err != nil { return nil, err } + readKeys = c.readKeysWithStagedVisibilityAliasesForGroup(gid, readKeys) + readKeys = c.readKeysWithStagedVisibilityMutationAliasesForGroup(gid, readKeys, elems) + if len(readKeys) > maxReadKeys { + return nil, errors.WithStack(ErrInvalidRequest) + } // ReadKeys are included in the Raft log entry so the FSM validates // read-write conflicts atomically under applyMu. prevCommitTS, when set, // carries the one-phase dedup probe key for a retry that reuses a failed @@ -1415,6 +1638,48 @@ func (c *ShardedCoordinator) dispatchSingleShardTxn(ctx context.Context, startTS return &CoordinateResponse{CommitIndex: resp.CommitIndex, CommitTS: commitTS}, nil } +func (c *ShardedCoordinator) readKeysWithStagedVisibilityAliasesForGroup(gid uint64, readKeys [][]byte) [][]byte { + if len(readKeys) == 0 { + return readKeys + } + var out [][]byte + for _, key := range readKeys { + alias, ok := c.stagedVisibilityReadKeyAlias(gid, key) + if !ok { + continue + } + if out == nil { + out = append([][]byte(nil), readKeys...) + } + out = append(out, alias) + } + if out == nil { + return readKeys + } + return out +} + +func (c *ShardedCoordinator) readKeysWithStagedVisibilityMutationAliasesForGroup(gid uint64, readKeys [][]byte, elems []*Elem[OP]) [][]byte { + var out [][]byte + for _, elem := range elems { + if elem == nil { + continue + } + alias, ok := c.stagedVisibilityReadKeyAlias(gid, elem.Key) + if !ok { + continue + } + if out == nil { + out = append([][]byte(nil), readKeys...) + } + out = append(out, alias) + } + if out == nil { + return readKeys + } + return out +} + type preparedGroup struct { gid uint64 keys []*pb.Mutation @@ -2225,7 +2490,7 @@ func (c *ShardedCoordinator) groupForKey(key []byte) (*ShardGroup, bool) { // catalog RouteID for !sqs|route|global. Partition-aware keyviz // is a Phase 3.D follow-up. func (c *ShardedCoordinator) routeAndGroupForKey(key []byte) (uint64, *ShardGroup, bool) { - gid, ok := c.router.ResolveGroup(key) + gid, routeID, ok := c.resolveGroupAndRouteForKey(key) if !ok { return 0, nil, false } @@ -2233,21 +2498,43 @@ func (c *ShardedCoordinator) routeAndGroupForKey(key []byte) (uint64, *ShardGrou if !ok { return 0, nil, false } - var routeID uint64 - if route, found := c.engine.GetRoute(routeKey(key)); found { - routeID = route.RouteID - } return routeID, g, true } func (c *ShardedCoordinator) engineGroupIDForKey(key []byte) uint64 { - gid, ok := c.router.ResolveGroup(key) + gid, _, ok := c.resolveGroupAndRouteForKey(key) if !ok { return 0 } return gid } +func (c *ShardedCoordinator) resolveGroupAndRouteForKey(key []byte) (uint64, uint64, bool) { + if route, ok := c.s3BucketAuxiliaryOwnerRouteForKey(key); ok { + return route.GroupID, route.RouteID, true + } + gid, ok := c.router.ResolveGroup(key) + if !ok { + return 0, 0, false + } + var routeID uint64 + if route, found := c.engine.GetRoute(routeKey(key)); found { + routeID = route.RouteID + } + return gid, routeID, true +} + +func (c *ShardedCoordinator) s3BucketAuxiliaryOwnerRouteForKey(key []byte) (distribution.Route, bool) { + if c == nil || c.engine == nil { + return distribution.Route{}, false + } + start, end, ok := s3BucketAuxiliaryRouteRange(key) + if !ok { + return distribution.Route{}, false + } + return s3BucketAuxiliaryOwnerRouteFromRange(start, end, c.engine.GetIntersectingRoutes(start, end)) +} + // EngineGroupIDForKey reports the Raft group ID that owns key, or 0 when // the key cannot be routed. Callers that batch lease checks across many // keys use it to collapse keys sharing a group into a single lease read @@ -2279,7 +2566,7 @@ func (c *ShardedCoordinator) groupReadKeysByShardID(readKeys [][]byte) (map[uint } grouped := make(map[uint64][][]byte) for _, key := range readKeys { - gid, ok := c.router.ResolveGroup(key) + gid, _, ok := c.resolveGroupAndRouteForKey(key) if !ok || gid == 0 { return nil, errors.Wrapf(ErrInvalidRequest, "no route for txn read key %q — recognised-but-"+ @@ -2287,10 +2574,65 @@ func (c *ShardedCoordinator) groupReadKeysByShardID(readKeys [][]byte) (map[uint "preserve OCC read-set integrity", key) } grouped[gid] = append(grouped[gid], key) + if alias, ok := c.stagedVisibilityReadKeyAlias(gid, key); ok { + grouped[gid] = append(grouped[gid], alias) + } } return grouped, nil } +func (c *ShardedCoordinator) groupedReadKeysWithStagedVisibilityMutationAliases(groupedReadKeys map[uint64][][]byte, groupedMutations map[uint64][]*pb.Mutation) map[uint64][][]byte { + out := groupedReadKeys + for gid, muts := range groupedMutations { + for _, mut := range muts { + if mut == nil { + continue + } + alias, ok := c.stagedVisibilityReadKeyAlias(gid, mut.Key) + if !ok { + continue + } + if out == nil { + out = make(map[uint64][][]byte) + } + out[gid] = append(out[gid], alias) + } + } + return out +} + +func groupedReadKeyCount(grouped map[uint64][][]byte) int { + var count int + for _, keys := range grouped { + count += len(keys) + } + return count +} + +func (c *ShardedCoordinator) stagedVisibilityReadKeyAlias(gid uint64, key []byte) ([]byte, bool) { + if c == nil || c.engine == nil || len(key) == 0 { + return nil, false + } + if _, _, ok := distribution.MigrationStagedDataKeyParts(key); ok { + return nil, false + } + if route, ok := c.s3BucketAuxiliaryOwnerRouteForKey(key); ok { + return stagedVisibilityReadKeyAliasForRoute(gid, key, route) + } + route, ok := c.engine.GetRoute(routeKey(key)) + if !ok { + return nil, false + } + return stagedVisibilityReadKeyAliasForRoute(gid, key, route) +} + +func stagedVisibilityReadKeyAliasForRoute(gid uint64, key []byte, route distribution.Route) ([]byte, bool) { + if route.GroupID != gid || !routeHasStagedVisibility(route) { + return nil, false + } + return distribution.MigrationStagedDataKey(route.MigrationJobID, key), true +} + // validateReadOnlyShards checks read-write conflicts on shards that have // read keys but no mutations in this transaction. writeGIDs is the set of // shards that already received a PREPARE with their readKeys attached. @@ -2340,7 +2682,7 @@ func (c *ShardedCoordinator) validateReadKeysOnShard(ctx context.Context, gid ui return errors.WithStack(err) } for _, key := range keys { - ts, exists, err := g.Store.LatestCommitTS(ctx, key) + ts, exists, err := c.latestCommitTSForReadKeyOnShard(ctx, gid, g, key) if err != nil { return errors.WithStack(err) } @@ -2351,6 +2693,51 @@ func (c *ShardedCoordinator) validateReadKeysOnShard(ctx context.Context, gid ui return nil } +func (c *ShardedCoordinator) latestCommitTSForReadKeyOnShard(ctx context.Context, gid uint64, g *ShardGroup, key []byte) (uint64, bool, error) { + route, ok := c.stagedVisibilityRouteForReadKey(gid, key) + if !ok { + liveTS, liveExists, err := g.Store.LatestCommitTS(ctx, key) + return liveTS, liveExists, errors.WithStack(err) + } + // Staged before live, for the reason on getAtWithStagedVisibility. + // Missing the version here is worse than a stale read: OCC would see no + // commit above startTS and let the transaction commit on a read it should + // have conflicted with. + stagedTS, stagedExists, err := g.Store.LatestCommitTS(ctx, distribution.MigrationStagedDataKey(route.MigrationJobID, key)) + if err != nil { + return 0, false, errors.WithStack(err) + } + liveTS, liveExists, err := g.Store.LatestCommitTS(ctx, key) + if err != nil { + return 0, false, errors.WithStack(err) + } + return maxStagedVisibilityLatestCommitTS(liveTS, liveExists, stagedTS, stagedExists), liveExists || stagedExists, nil +} + +func (c *ShardedCoordinator) stagedVisibilityRouteForReadKey(gid uint64, key []byte) (distribution.Route, bool) { + if c == nil || c.engine == nil { + return distribution.Route{}, false + } + if _, _, ok := distribution.MigrationStagedDataKeyParts(key); ok { + return distribution.Route{}, false + } + if route, ok := c.s3BucketAuxiliaryOwnerRouteForKey(key); ok { + return route, route.GroupID == gid && routeHasStagedVisibility(route) + } + route, ok := c.engine.GetRoute(routeKey(key)) + return route, ok && route.GroupID == gid && routeHasStagedVisibility(route) +} + +func maxStagedVisibilityLatestCommitTS(liveTS uint64, liveExists bool, stagedTS uint64, stagedExists bool) uint64 { + if !liveExists { + return stagedTS + } + if !stagedExists || liveTS > stagedTS { + return liveTS + } + return stagedTS +} + var _ Coordinator = (*ShardedCoordinator)(nil) func validateOperationGroup(reqs *OperationGroup[OP]) error { @@ -2401,12 +2788,38 @@ func (c *ShardedCoordinator) rawLogsWithGroups(ctx context.Context, reqs *Operat Phase: pb.Phase_NONE, Ts: ts, Mutations: muts, + ObservedRouteVersion: reqs.ObservedRouteVersion, WriteFenceBypassKeys: bypassKeysByGroup[gid], }) } return logs, gids, nil } +func (c *ShardedCoordinator) rejectWriteTimestampFloorMutations(muts []*pb.Mutation, commitTS uint64) error { + if c == nil || c.engine == nil || commitTS == 0 { + return nil + } + for _, mut := range muts { + if err := c.rejectWriteTimestampFloorMutation(mut, commitTS); err != nil { + return err + } + } + return nil +} + +func (c *ShardedCoordinator) rejectWriteTimestampFloorMutation(mut *pb.Mutation, commitTS uint64) error { + if mut == nil { + return nil + } + if mut.GetOp() == pb.Op_DEL_PREFIX { + return c.rejectWriteTimestampFloorDelPrefix(mut.Key, commitTS) + } + if len(mut.Key) == 0 { + return nil + } + return c.rejectWriteTimestampFloorPointKey(mut.Key, commitTS) +} + func (c *ShardedCoordinator) rawLogTimestamp(ctx context.Context) (uint64, error) { if c.tsAllocator != nil { return 0, nil @@ -2610,38 +3023,12 @@ func (c *ShardedCoordinator) ensureMutationWriteAllowed(mut *pb.Mutation, commit } return nil case pb.Op_PUT, pb.Op_DEL: - return c.ensurePointMutationWriteAllowed(mut.Key, commitTS) + return c.rejectWriteTimestampFloorPointKey(mut.Key, commitTS) default: return errors.WithStack(ErrInvalidRequest) } } -func (c *ShardedCoordinator) ensurePointMutationWriteAllowed(key []byte, commitTS uint64) error { - if c.engine == nil { - return nil - } - // Applied before the raw-route lookup below so a fenced user key rejects - // its list-delta/claim and stream rows even when the raw key's own route - // has no floor. - if err := ensureLogicalRouteWriteAllowed(c.engine, key, commitTS); err != nil { - return err - } - route, ok := c.engine.GetRoute(routeKey(key)) - if !ok { - return nil - } - if c.router != nil { - gid, ok := c.router.ResolveGroup(key) - if !ok { - return errors.Wrapf(ErrInvalidRequest, "no route for key %q", key) - } - if gid != route.GroupID { - return nil - } - } - return ensureRouteWriteAllowed(route, key, commitTS) -} - func (c *ShardedCoordinator) groupMutations(reqs []*Elem[OP], label keyviz.Label) (map[uint64][]*pb.Mutation, []uint64, error) { grouped := make(map[uint64][]*pb.Mutation) for _, req := range reqs { @@ -2650,21 +3037,18 @@ func (c *ShardedCoordinator) groupMutations(reqs []*Elem[OP], label keyviz.Label } mut := elemToMutation(req) gid := req.GroupID + _, routeID, routeOK := c.resolveGroupAndRouteForKey(mut.Key) if gid == 0 { - var ok bool - gid, ok = c.router.ResolveGroup(mut.Key) + resolvedGID, resolvedRouteID, ok := c.resolveGroupAndRouteForKey(mut.Key) if !ok { return nil, nil, errors.Wrapf(ErrInvalidRequest, "no route for key %q", mut.Key) } + gid = resolvedGID + routeID = resolvedRouteID } else if _, ok := c.groups[gid]; !ok { return nil, nil, errors.Wrapf(ErrInvalidRequest, "no shard group %d for key %q", gid, mut.Key) - } - // Engine RouteID for keyviz observation; partition-resolved - // keys observe under the !sqs|route|global RouteID until - // partition-aware keyviz lands. - var routeID uint64 - if route, found := c.engine.GetRoute(routeKey(mut.Key)); found { - routeID = route.RouteID + } else if !routeOK { + return nil, nil, errors.Wrapf(ErrInvalidRequest, "no route for key %q", mut.Key) } c.observeMutation(routeID, mut, label) grouped[gid] = append(grouped[gid], mut) diff --git a/kv/sharded_coordinator_del_prefix_test.go b/kv/sharded_coordinator_del_prefix_test.go index 544f91c5f..0cfda39b5 100644 --- a/kv/sharded_coordinator_del_prefix_test.go +++ b/kv/sharded_coordinator_del_prefix_test.go @@ -6,6 +6,7 @@ import ( "testing" "github.com/bootjp/elastickv/distribution" + "github.com/bootjp/elastickv/internal/s3keys" pb "github.com/bootjp/elastickv/proto" "github.com/bootjp/elastickv/store" "github.com/stretchr/testify/require" @@ -121,7 +122,250 @@ func TestShardedCoordinator_DelPrefixBroadcastsToAllGroups(t *testing.T) { "same DEL_PREFIX element must use the same timestamp across shards") } -func TestShardedCoordinator_DelPrefixRejectsRouteWriteTimestampFloor(t *testing.T) { +func newMigrationFloorEngine(t *testing.T, floor uint64) *distribution.Engine { + t.Helper() + + engine := distribution.NewEngine() + require.NoError(t, engine.ApplySnapshot(distribution.CatalogSnapshot{ + Version: 1, + Routes: []distribution.RouteDescriptor{ + { + RouteID: 1, + Start: []byte(""), + End: nil, + GroupID: 1, + State: distribution.RouteStateActive, + MinWriteTSExclusive: floor, + }, + }, + })) + return engine +} + +func TestShardedCoordinatorRejectsPointWriteAtMigrationTimestampFloor(t *testing.T) { + t.Parallel() + + g1Txn := &recordingTransactional{} + coord := NewShardedCoordinator(newMigrationFloorEngine(t, ^uint64(0)), map[uint64]*ShardGroup{ + 1: {Txn: g1Txn}, + }, 1, NewHLC(), nil) + + _, err := coord.Dispatch(context.Background(), &OperationGroup[OP]{ + Elems: []*Elem[OP]{{Op: Put, Key: []byte("z"), Value: []byte("v")}}, + }) + require.ErrorIs(t, err, ErrRouteWriteTimestampTooLow) + require.Empty(t, g1Txn.requests, "coordinator must reject before proposing a floor-violating point write") +} + +func TestShardedCoordinatorRejectsDelPrefixAtMigrationTimestampFloor(t *testing.T) { + t.Parallel() + + g1Txn := &recordingTransactional{} + coord := NewShardedCoordinator(newMigrationFloorEngine(t, ^uint64(0)), map[uint64]*ShardGroup{ + 1: {Txn: g1Txn}, + }, 1, NewHLC(), nil) + + _, err := coord.Dispatch(context.Background(), &OperationGroup[OP]{ + Elems: []*Elem[OP]{{Op: DelPrefix, Key: []byte("z")}}, + }) + require.ErrorIs(t, err, ErrRouteWriteTimestampTooLow) + require.Empty(t, g1Txn.requests, "coordinator must reject before broadcasting a floor-violating prefix delete") +} + +func TestShardedCoordinatorRejectsRawDelPrefixMutationAtMigrationTimestampFloor(t *testing.T) { + t.Parallel() + + coord := NewShardedCoordinator(newMigrationFloorEngine(t, ^uint64(0)), map[uint64]*ShardGroup{ + 1: {Txn: &recordingTransactional{}}, + }, 1, NewHLC(), nil) + + for _, mut := range []*pb.Mutation{ + {Op: pb.Op_DEL_PREFIX, Key: []byte("z")}, + {Op: pb.Op_DEL_PREFIX, Key: nil}, + } { + err := coord.rejectWriteTimestampFloorMutations([]*pb.Mutation{mut}, 100) + require.ErrorIs(t, err, ErrRouteWriteTimestampTooLow) + } +} + +func TestShardedCoordinator_DelPrefixDoesNotAutoPinObservedRouteVersion(t *testing.T) { + t.Parallel() + + engine := distribution.NewEngine() + require.NoError(t, engine.ApplySnapshot(distribution.CatalogSnapshot{ + Version: 7, + Routes: []distribution.RouteDescriptor{ + {RouteID: 1, Start: nil, End: nil, GroupID: 1, State: distribution.RouteStateActive}, + }, + })) + txn := &recordingTransactional{} + coord := NewShardedCoordinator(engine, map[uint64]*ShardGroup{ + 1: {Txn: txn}, + }, 1, NewHLC(), nil) + + _, err := coord.Dispatch(context.Background(), &OperationGroup[OP]{ + Elems: []*Elem[OP]{{Op: DelPrefix, Key: []byte("user:")}}, + }) + require.NoError(t, err) + require.Len(t, txn.requests, 1) + require.Zero(t, txn.requests[0].GetObservedRouteVersion()) +} + +func TestShardedCoordinator_RawWriteDoesNotAutoPinObservedRouteVersion(t *testing.T) { + t.Parallel() + + engine := distribution.NewEngine() + require.NoError(t, engine.ApplySnapshot(distribution.CatalogSnapshot{ + Version: 9, + Routes: []distribution.RouteDescriptor{ + {RouteID: 1, Start: nil, End: nil, GroupID: 1, State: distribution.RouteStateActive}, + }, + })) + txn := &recordingTransactional{} + coord := NewShardedCoordinator(engine, map[uint64]*ShardGroup{ + 1: {Txn: txn}, + }, 1, NewHLC(), nil) + + _, err := coord.Dispatch(context.Background(), &OperationGroup[OP]{ + Elems: []*Elem[OP]{{Op: Put, Key: []byte("k"), Value: []byte("v")}}, + }) + require.NoError(t, err) + require.Len(t, txn.requests, 1) + require.Zero(t, txn.requests[0].GetObservedRouteVersion()) +} + +func TestShardedCoordinator_RetriesRawWriteWhenObservedRouteVersionGCd(t *testing.T) { + t.Parallel() + + engine := distribution.NewEngine() + require.NoError(t, engine.ApplySnapshot(distribution.CatalogSnapshot{ + Version: 9, + Routes: []distribution.RouteDescriptor{ + {RouteID: 1, Start: nil, End: nil, GroupID: 1, State: distribution.RouteStateActive}, + }, + })) + txn := &recordingTransactional{ + errs: []error{ErrComposed1VersionGCd}, + } + coord := NewShardedCoordinator(engine, map[uint64]*ShardGroup{ + 1: {Txn: txn}, + }, 1, NewHLC(), nil) + + _, err := coord.Dispatch(context.Background(), &OperationGroup[OP]{ + ObservedRouteVersion: 3, + Elems: []*Elem[OP]{{Op: Put, Key: []byte("k"), Value: []byte("v")}}, + }) + require.NoError(t, err) + require.Len(t, txn.requests, 2) + require.Equal(t, uint64(3), txn.requests[0].GetObservedRouteVersion()) + require.Equal(t, uint64(9), txn.requests[1].GetObservedRouteVersion()) +} + +func TestShardedCoordinator_DoesNotRetryDelPrefixWhenObservedRouteVersionGCd(t *testing.T) { + t.Parallel() + + engine := distribution.NewEngine() + require.NoError(t, engine.ApplySnapshot(distribution.CatalogSnapshot{ + Version: 7, + Routes: []distribution.RouteDescriptor{ + {RouteID: 1, Start: nil, End: nil, GroupID: 1, State: distribution.RouteStateActive}, + }, + })) + txn := &recordingTransactional{ + errs: []error{ErrComposed1VersionGCd}, + } + coord := NewShardedCoordinator(engine, map[uint64]*ShardGroup{ + 1: {Txn: txn}, + }, 1, NewHLC(), nil) + + _, err := coord.Dispatch(context.Background(), &OperationGroup[OP]{ + ObservedRouteVersion: 2, + Elems: []*Elem[OP]{{Op: DelPrefix, Key: []byte("user:")}}, + }) + require.ErrorIs(t, err, ErrComposed1VersionGCd) + require.Len(t, txn.requests, 1) + require.Equal(t, uint64(2), txn.requests[0].GetObservedRouteVersion()) +} + +func TestShardedCoordinator_DoesNotRetryMultiShardRawWriteWhenObservedRouteVersionGCd(t *testing.T) { + t.Parallel() + + engine := distribution.NewEngine() + require.NoError(t, engine.ApplySnapshot(distribution.CatalogSnapshot{ + Version: 11, + Routes: []distribution.RouteDescriptor{ + {RouteID: 1, Start: []byte(""), End: []byte("m"), GroupID: 1, State: distribution.RouteStateActive}, + {RouteID: 2, Start: []byte("m"), End: nil, GroupID: 2, State: distribution.RouteStateActive}, + }, + })) + g1Txn := &recordingTransactional{} + g2Txn := &recordingTransactional{errs: []error{ErrComposed1VersionGCd}} + coord := NewShardedCoordinator(engine, map[uint64]*ShardGroup{ + 1: {Txn: g1Txn}, + 2: {Txn: g2Txn}, + }, 1, NewHLC(), nil) + + _, err := coord.Dispatch(context.Background(), &OperationGroup[OP]{ + ObservedRouteVersion: 3, + Elems: []*Elem[OP]{ + {Op: Put, Key: []byte("a"), Value: []byte("v1")}, + {Op: Put, Key: []byte("z"), Value: []byte("v2")}, + }, + }) + + require.ErrorIs(t, err, ErrComposed1VersionGCd) + require.Len(t, g1Txn.requests, 1) + require.Len(t, g2Txn.requests, 1) + require.Equal(t, uint64(3), g1Txn.requests[0].GetObservedRouteVersion()) + require.Equal(t, uint64(3), g2Txn.requests[0].GetObservedRouteVersion()) +} + +func TestShardedCoordinator_DoesNotRetryRawWriteWhenRetryRouteBecomesMultiShard(t *testing.T) { + t.Parallel() + + engine := distribution.NewEngine() + require.NoError(t, engine.ApplySnapshot(distribution.CatalogSnapshot{ + Version: 11, + Routes: []distribution.RouteDescriptor{ + {RouteID: 1, Start: []byte(""), End: nil, GroupID: 1, State: distribution.RouteStateActive}, + }, + })) + g1Txn := &recordingTransactional{ + errs: []error{ErrComposed1VersionGCd}, + onCommit: func(call int, _ *pb.Request) { + if call != 0 { + return + } + require.NoError(t, engine.ApplySnapshot(distribution.CatalogSnapshot{ + Version: 12, + Routes: []distribution.RouteDescriptor{ + {RouteID: 1, Start: []byte(""), End: []byte("m"), GroupID: 1, State: distribution.RouteStateActive}, + {RouteID: 2, Start: []byte("m"), End: nil, GroupID: 2, State: distribution.RouteStateActive}, + }, + })) + }, + } + g2Txn := &recordingTransactional{} + coord := NewShardedCoordinator(engine, map[uint64]*ShardGroup{ + 1: {Txn: g1Txn}, + 2: {Txn: g2Txn}, + }, 1, NewHLC(), nil) + + _, err := coord.Dispatch(context.Background(), &OperationGroup[OP]{ + ObservedRouteVersion: 3, + Elems: []*Elem[OP]{ + {Op: Put, Key: []byte("a"), Value: []byte("v1")}, + {Op: Put, Key: []byte("z"), Value: []byte("v2")}, + }, + }) + + require.ErrorIs(t, err, ErrComposed1VersionGCd) + require.Len(t, g1Txn.requests, 1) + require.Len(t, g2Txn.requests, 0) + require.Equal(t, uint64(3), g1Txn.requests[0].GetObservedRouteVersion()) +} + +func TestShardedCoordinatorRejectsPointWriteOnWriteFencedRoute(t *testing.T) { t.Parallel() engine := distribution.NewEngine() @@ -129,10 +373,58 @@ func TestShardedCoordinator_DelPrefixRejectsRouteWriteTimestampFloor(t *testing. Version: 1, Routes: []distribution.RouteDescriptor{ {RouteID: 1, Start: []byte(""), End: []byte("m"), GroupID: 1, State: distribution.RouteStateActive}, - {RouteID: 2, Start: []byte("m"), GroupID: 2, State: distribution.RouteStateActive, MinWriteTSExclusive: ^uint64(0)}, + {RouteID: 2, Start: []byte("m"), End: nil, GroupID: 2, State: distribution.RouteStateWriteFenced}, + }, + })) + + g2Txn := &recordingTransactional{} + coord := NewShardedCoordinator(engine, map[uint64]*ShardGroup{ + 1: {Txn: &recordingTransactional{}}, + 2: {Txn: g2Txn}, + }, 1, NewHLC(), nil) + + _, err := coord.Dispatch(context.Background(), &OperationGroup[OP]{ + Elems: []*Elem[OP]{{Op: Put, Key: []byte("z"), Value: []byte("v")}}, + }) + require.ErrorIs(t, err, ErrRouteWriteFenced) + require.Empty(t, g2Txn.requests, "coordinator must reject before proposing to the fenced shard") +} + +func TestShardedCoordinatorRejectsEmptyKeyWriteOnLeadingWriteFencedRoute(t *testing.T) { + t.Parallel() + + engine := distribution.NewEngine() + require.NoError(t, engine.ApplySnapshot(distribution.CatalogSnapshot{ + Version: 1, + Routes: []distribution.RouteDescriptor{ + {RouteID: 1, Start: []byte(""), End: []byte("m"), GroupID: 1, State: distribution.RouteStateWriteFenced}, + {RouteID: 2, Start: []byte("m"), End: nil, GroupID: 2, State: distribution.RouteStateActive}, }, })) + g1Txn := &recordingTransactional{} + coord := NewShardedCoordinator(engine, map[uint64]*ShardGroup{ + 1: {Txn: g1Txn}, + 2: {Txn: &recordingTransactional{}}, + }, 1, NewHLC(), nil) + + _, err := coord.Dispatch(context.Background(), &OperationGroup[OP]{ + Elems: []*Elem[OP]{{Op: Put, Key: []byte{}, Value: []byte("v")}}, + }) + require.ErrorIs(t, err, ErrRouteWriteFenced) + require.Empty(t, g1Txn.requests, "coordinator must reject the empty key before proposing to the fenced shard") +} + +func TestShardedCoordinatorRejectsS3BucketAuxiliaryPointWriteOnWriteFencedRoute(t *testing.T) { + t.Parallel() + + const bucket = "bucket-a" + engine := distribution.NewEngine() + require.NoError(t, engine.ApplySnapshot(distribution.CatalogSnapshot{ + Version: 1, + Routes: s3BucketAuxiliaryFenceRoutes(bucket, 1, 2), + })) + g1Txn := &recordingTransactional{} g2Txn := &recordingTransactional{} coord := NewShardedCoordinator(engine, map[uint64]*ShardGroup{ @@ -140,16 +432,426 @@ func TestShardedCoordinator_DelPrefixRejectsRouteWriteTimestampFloor(t *testing. 2: {Txn: g2Txn}, }, 1, NewHLC(), nil) + for _, key := range [][]byte{ + s3keys.BucketMetaKey(bucket), + s3keys.BucketGenerationKey(bucket), + } { + _, err := coord.Dispatch(context.Background(), &OperationGroup[OP]{ + Elems: []*Elem[OP]{{Op: Put, Key: key, Value: []byte("v")}}, + }) + require.ErrorIs(t, err, ErrRouteWriteFenced) + require.Empty(t, g1Txn.requests, "coordinator must reject before proposing to the raw-key shard") + require.Empty(t, g2Txn.requests, "coordinator must reject before proposing to the fenced shard") + } +} + +func s3BucketAuxiliaryStagedRoutes(bucket string, rawGroupID, stagedGroupID uint64) []distribution.RouteDescriptor { + routes := s3BucketAuxiliaryFenceRoutes(bucket, rawGroupID, stagedGroupID) + routes[1].State = distribution.RouteStateActive + routes[1].StagedVisibilityActive = true + routes[1].MigrationJobID = 9 + return routes +} + +func s3BucketAuxiliaryPromotedRoutes() []distribution.RouteDescriptor { + const bucket = "bucket-a" + routes := s3BucketAuxiliaryFenceRoutes(bucket, 1, 2) + routes[1].State = distribution.RouteStateActive + return routes +} + +func TestShardedCoordinatorRoutesS3BucketAuxiliaryWriteToStagedOwner(t *testing.T) { + t.Parallel() + + const bucket = "bucket-a" + engine := distribution.NewEngine() + require.NoError(t, engine.ApplySnapshot(distribution.CatalogSnapshot{ + Version: 1, + Routes: s3BucketAuxiliaryStagedRoutes(bucket, 1, 2), + })) + + g1Txn := &recordingTransactional{} + g2Txn := &recordingTransactional{responses: []*TransactionResponse{{CommitIndex: 22}}} + coord := NewShardedCoordinator(engine, map[uint64]*ShardGroup{ + 1: {Txn: g1Txn}, + 2: {Txn: g2Txn}, + }, 1, NewHLC(), nil) + + key := s3keys.BucketMetaKey(bucket) + route, ok := coord.s3BucketAuxiliaryOwnerRouteForKey(key) + require.True(t, ok) + require.Equal(t, uint64(2), route.GroupID) + _, err := coord.Dispatch(context.Background(), &OperationGroup[OP]{ - Elems: []*Elem[OP]{ - {Op: DelPrefix, Key: nil}, + Elems: []*Elem[OP]{{Op: Put, Key: key, Value: []byte("meta")}}, + }) + require.NoError(t, err) + require.Empty(t, g1Txn.requests) + require.Len(t, g2Txn.requests, 1) + require.Equal(t, key, g2Txn.requests[0].Mutations[0].Key) +} + +func TestShardedCoordinatorRoutesS3BucketAuxiliaryWriteToPromotedOwner(t *testing.T) { + t.Parallel() + + const bucket = "bucket-a" + engine := distribution.NewEngine() + require.NoError(t, engine.ApplySnapshot(distribution.CatalogSnapshot{ + Version: 1, + Routes: s3BucketAuxiliaryPromotedRoutes(), + })) + + g1Txn := &recordingTransactional{} + g2Txn := &recordingTransactional{responses: []*TransactionResponse{{CommitIndex: 22}}} + coord := NewShardedCoordinator(engine, map[uint64]*ShardGroup{ + 1: {Txn: g1Txn}, + 2: {Txn: g2Txn}, + }, 1, NewHLC(), nil) + + key := s3keys.BucketMetaKey(bucket) + route, ok := coord.s3BucketAuxiliaryOwnerRouteForKey(key) + require.True(t, ok) + require.Equal(t, uint64(2), route.GroupID) + require.False(t, routeHasStagedVisibility(route)) + + _, err := coord.Dispatch(context.Background(), &OperationGroup[OP]{ + Elems: []*Elem[OP]{{Op: Put, Key: key, Value: []byte("meta")}}, + }) + require.NoError(t, err) + require.Empty(t, g1Txn.requests) + require.Len(t, g2Txn.requests, 1) + require.Equal(t, key, g2Txn.requests[0].Mutations[0].Key) +} + +func TestShardedCoordinatorIgnoresRawRouteFloorForS3BucketAuxiliaryWrite(t *testing.T) { + t.Parallel() + + const bucket = "bucket-a" + key := s3keys.BucketMetaKey(bucket) + engine := distribution.NewEngine() + routes := s3BucketAuxiliaryStagedRoutes(bucket, 1, 2) + routes[2].MinWriteTSExclusive = ^uint64(0) + require.NoError(t, engine.ApplySnapshot(distribution.CatalogSnapshot{ + Version: 1, + Routes: routes, + })) + + rawRoute, ok := engine.GetRoute(routeKey(key)) + require.True(t, ok) + require.Equal(t, ^uint64(0), rawRoute.MinWriteTSExclusive) + auxStart, auxEnd, ok := s3BucketAuxiliaryRouteRange(key) + require.True(t, ok) + auxRoutes := engine.GetIntersectingRoutes(auxStart, auxEnd) + require.NotEmpty(t, auxRoutes) + require.Zero(t, auxRoutes[0].MinWriteTSExclusive) + + g1Txn := &recordingTransactional{} + g2Txn := &recordingTransactional{responses: []*TransactionResponse{{CommitIndex: 22}}} + coord := NewShardedCoordinator(engine, map[uint64]*ShardGroup{ + 1: {Txn: g1Txn}, + 2: {Txn: g2Txn}, + }, 1, NewHLC(), nil) + + _, err := coord.Dispatch(context.Background(), &OperationGroup[OP]{ + Elems: []*Elem[OP]{{Op: Put, Key: key, Value: []byte("meta")}}, + }) + require.NoError(t, err) + require.Empty(t, g1Txn.requests) + require.Len(t, g2Txn.requests, 1) +} + +func TestShardedCoordinatorIgnoresRawRouteFenceForS3BucketAuxiliaryWrite(t *testing.T) { + t.Parallel() + + const bucket = "bucket-a" + key := s3keys.BucketMetaKey(bucket) + engine := distribution.NewEngine() + routes := s3BucketAuxiliaryStagedRoutes(bucket, 1, 2) + routes[2].State = distribution.RouteStateWriteFenced + require.NoError(t, engine.ApplySnapshot(distribution.CatalogSnapshot{ + Version: 1, + Routes: routes, + })) + + rawRoute, ok := engine.GetRoute(routeKey(key)) + require.True(t, ok) + require.Equal(t, distribution.RouteStateWriteFenced, rawRoute.State) + auxStart, auxEnd, ok := s3BucketAuxiliaryRouteRange(key) + require.True(t, ok) + auxRoutes := engine.GetIntersectingRoutes(auxStart, auxEnd) + require.NotEmpty(t, auxRoutes) + require.Equal(t, distribution.RouteStateActive, auxRoutes[0].State) + + g1Txn := &recordingTransactional{} + g2Txn := &recordingTransactional{responses: []*TransactionResponse{{CommitIndex: 22}}} + coord := NewShardedCoordinator(engine, map[uint64]*ShardGroup{ + 1: {Txn: g1Txn}, + 2: {Txn: g2Txn}, + }, 1, NewHLC(), nil) + + _, err := coord.Dispatch(context.Background(), &OperationGroup[OP]{ + Elems: []*Elem[OP]{{Op: Put, Key: key, Value: []byte("meta")}}, + }) + require.NoError(t, err) + require.Empty(t, g1Txn.requests) + require.Len(t, g2Txn.requests, 1) +} + +func TestShardedCoordinatorIgnoresNonOwnerS3BucketAuxiliaryFenceForPointWrite(t *testing.T) { + t.Parallel() + + const bucket = "bucket-a" + key := s3keys.BucketMetaKey(bucket) + engine := distribution.NewEngine() + routes := s3BucketAuxiliarySplitRoutes(bucket, 1, 2, 3) + routes[2].State = distribution.RouteStateWriteFenced + require.NoError(t, engine.ApplySnapshot(distribution.CatalogSnapshot{ + Version: 1, + Routes: routes, + })) + + g1Txn := &recordingTransactional{} + g2Txn := &recordingTransactional{responses: []*TransactionResponse{{CommitIndex: 22}}} + g3Txn := &recordingTransactional{} + coord := NewShardedCoordinator(engine, map[uint64]*ShardGroup{ + 1: {Txn: g1Txn}, + 2: {Txn: g2Txn}, + 3: {Txn: g3Txn}, + }, 1, NewHLC(), nil) + + _, err := coord.Dispatch(context.Background(), &OperationGroup[OP]{ + Elems: []*Elem[OP]{{Op: Put, Key: key, Value: []byte("meta")}}, + }) + require.NoError(t, err) + require.Empty(t, g1Txn.requests) + require.Len(t, g2Txn.requests, 1) + require.Empty(t, g3Txn.requests) +} + +func TestShardedCoordinatorRejectsS3BucketAuxiliaryPointWriteAtMigrationTimestampFloor(t *testing.T) { + t.Parallel() + + const bucket = "bucket-a" + start := s3keys.RoutePrefixForBucketAnyGeneration(bucket) + end := prefixScanEnd(start) + engine := distribution.NewEngine() + require.NoError(t, engine.ApplySnapshot(distribution.CatalogSnapshot{ + Version: 1, + Routes: []distribution.RouteDescriptor{ + {RouteID: 1, Start: []byte(""), End: start, GroupID: 1, State: distribution.RouteStateActive}, + {RouteID: 2, Start: start, End: end, GroupID: 2, State: distribution.RouteStateActive, MinWriteTSExclusive: ^uint64(0)}, + {RouteID: 3, Start: end, End: nil, GroupID: 1, State: distribution.RouteStateActive}, }, + })) + + g1Txn := &recordingTransactional{} + g2Txn := &recordingTransactional{} + coord := NewShardedCoordinator(engine, map[uint64]*ShardGroup{ + 1: {Txn: g1Txn}, + 2: {Txn: g2Txn}, + }, 1, NewHLC(), nil) + + for _, key := range [][]byte{ + s3keys.BucketMetaKey(bucket), + s3keys.BucketGenerationKey(bucket), + } { + _, err := coord.Dispatch(context.Background(), &OperationGroup[OP]{ + Elems: []*Elem[OP]{{Op: Put, Key: key, Value: []byte("v")}}, + }) + require.ErrorIs(t, err, ErrRouteWriteTimestampTooLow) + require.Empty(t, g1Txn.requests, "coordinator must reject before proposing to the raw-key shard") + require.Empty(t, g2Txn.requests, "coordinator must reject before proposing to the floor-fenced shard") + } +} + +func TestShardedCoordinatorIgnoresNonOwnerS3BucketAuxiliaryFloorForPointWrite(t *testing.T) { + t.Parallel() + + const bucket = "bucket-b" + key := s3keys.BucketMetaKey(bucket) + engine := distribution.NewEngine() + routes := s3BucketAuxiliarySplitRoutes(bucket, 1, 2, 3) + routes[2].MinWriteTSExclusive = ^uint64(0) + require.NoError(t, engine.ApplySnapshot(distribution.CatalogSnapshot{ + Version: 1, + Routes: routes, + })) + + g1Txn := &recordingTransactional{} + g2Txn := &recordingTransactional{responses: []*TransactionResponse{{CommitIndex: 22}}} + g3Txn := &recordingTransactional{} + coord := NewShardedCoordinator(engine, map[uint64]*ShardGroup{ + 1: {Txn: g1Txn}, + 2: {Txn: g2Txn}, + 3: {Txn: g3Txn}, + }, 1, NewHLC(), nil) + + _, err := coord.Dispatch(context.Background(), &OperationGroup[OP]{ + Elems: []*Elem[OP]{{Op: Put, Key: key, Value: []byte("meta")}}, + }) + require.NoError(t, err) + require.Empty(t, g1Txn.requests) + require.Len(t, g2Txn.requests, 1) + require.Empty(t, g3Txn.requests) +} + +func TestShardedCoordinatorRejectsDelPrefixIntersectingWriteFencedRoute(t *testing.T) { + t.Parallel() + + engine := distribution.NewEngine() + require.NoError(t, engine.ApplySnapshot(distribution.CatalogSnapshot{ + Version: 1, + Routes: []distribution.RouteDescriptor{ + {RouteID: 1, Start: []byte(""), End: []byte("m"), GroupID: 1, State: distribution.RouteStateActive}, + {RouteID: 2, Start: []byte("m"), End: nil, GroupID: 2, State: distribution.RouteStateWriteFenced}, + }, + })) + + g1Txn := &recordingTransactional{} + g2Txn := &recordingTransactional{} + coord := NewShardedCoordinator(engine, map[uint64]*ShardGroup{ + 1: {Txn: g1Txn}, + 2: {Txn: g2Txn}, + }, 1, NewHLC(), nil) + + _, err := coord.Dispatch(context.Background(), &OperationGroup[OP]{ + Elems: []*Elem[OP]{{Op: DelPrefix, Key: []byte("z")}}, + }) + require.ErrorIs(t, err, ErrRouteWriteFenced) + require.Empty(t, g1Txn.requests) + require.Empty(t, g2Txn.requests) +} + +func TestShardedCoordinatorRejectsFullRangeDelPrefixWhenRouteIsWriteFenced(t *testing.T) { + t.Parallel() + + engine := distribution.NewEngine() + require.NoError(t, engine.ApplySnapshot(distribution.CatalogSnapshot{ + Version: 1, + Routes: []distribution.RouteDescriptor{ + {RouteID: 1, Start: []byte(""), End: []byte("m"), GroupID: 1, State: distribution.RouteStateActive}, + {RouteID: 2, Start: []byte("m"), End: nil, GroupID: 2, State: distribution.RouteStateWriteFenced}, + }, + })) + + g1Txn := &recordingTransactional{} + g2Txn := &recordingTransactional{} + coord := NewShardedCoordinator(engine, map[uint64]*ShardGroup{ + 1: {Txn: g1Txn}, + 2: {Txn: g2Txn}, + }, 1, NewHLC(), nil) + + _, err := coord.Dispatch(context.Background(), &OperationGroup[OP]{ + Elems: []*Elem[OP]{{Op: DelPrefix, Key: nil}}, }) - require.ErrorIs(t, err, store.ErrWriteConflict) + require.ErrorIs(t, err, ErrRouteWriteFenced) require.Empty(t, g1Txn.requests) require.Empty(t, g2Txn.requests) } +func TestShardedCoordinatorRejectsBroadInternalDelPrefixWhenRouteIsWriteFenced(t *testing.T) { + t.Parallel() + + engine := distribution.NewEngine() + require.NoError(t, engine.ApplySnapshot(distribution.CatalogSnapshot{ + Version: 1, + Routes: []distribution.RouteDescriptor{ + {RouteID: 1, Start: []byte(""), End: []byte("m"), GroupID: 1, State: distribution.RouteStateActive}, + {RouteID: 2, Start: []byte("m"), End: nil, GroupID: 2, State: distribution.RouteStateWriteFenced}, + }, + })) + + for _, prefix := range [][]byte{ + []byte("!redis|"), + []byte("!lst|"), + } { + t.Run(string(prefix), func(t *testing.T) { + t.Parallel() + + g1Txn := &recordingTransactional{} + g2Txn := &recordingTransactional{} + coord := NewShardedCoordinator(engine, map[uint64]*ShardGroup{ + 1: {Txn: g1Txn}, + 2: {Txn: g2Txn}, + }, 1, NewHLC(), nil) + + _, err := coord.Dispatch(context.Background(), &OperationGroup[OP]{ + Elems: []*Elem[OP]{{Op: DelPrefix, Key: prefix}}, + }) + require.ErrorIs(t, err, ErrRouteWriteFenced) + require.Empty(t, g1Txn.requests) + require.Empty(t, g2Txn.requests) + }) + } +} + +func TestShardedCoordinatorAllowsRawSQSLookingDelPrefixWhenUnrelatedRouteIsWriteFenced(t *testing.T) { + t.Parallel() + + engine := distribution.NewEngine() + require.NoError(t, engine.ApplySnapshot(distribution.CatalogSnapshot{ + Version: 1, + Routes: []distribution.RouteDescriptor{ + {RouteID: 1, Start: []byte(""), End: []byte("m"), GroupID: 1, State: distribution.RouteStateActive}, + {RouteID: 2, Start: []byte("m"), End: nil, GroupID: 2, State: distribution.RouteStateWriteFenced}, + }, + })) + + g1Txn := &recordingTransactional{} + g2Txn := &recordingTransactional{} + coord := NewShardedCoordinator(engine, map[uint64]*ShardGroup{ + 1: {Txn: g1Txn}, + 2: {Txn: g2Txn}, + }, 1, NewHLC(), nil) + + _, err := coord.Dispatch(context.Background(), &OperationGroup[OP]{ + Elems: []*Elem[OP]{{Op: DelPrefix, Key: []byte("!sqs|foo")}}, + }) + require.NoError(t, err) + require.NotEmpty(t, g1Txn.requests) + require.NotEmpty(t, g2Txn.requests) +} + +func TestShardedCoordinatorAllowsS3BucketDelPrefixWhenUnrelatedRouteIsWriteFenced(t *testing.T) { + t.Parallel() + + const ( + activeBucket = "bucket-a" + fencedBucket = "bucket-b" + generation = uint64(7) + ) + activeStart := s3keys.RoutePrefixForBucketAnyGeneration(activeBucket) + activeEnd := prefixScanEnd(activeStart) + fencedStart := s3keys.RoutePrefixForBucketAnyGeneration(fencedBucket) + fencedEnd := prefixScanEnd(fencedStart) + + engine := distribution.NewEngine() + require.NoError(t, engine.ApplySnapshot(distribution.CatalogSnapshot{ + Version: 1, + Routes: []distribution.RouteDescriptor{ + {RouteID: 1, Start: []byte(""), End: activeStart, GroupID: 1, State: distribution.RouteStateActive}, + {RouteID: 2, Start: activeStart, End: activeEnd, GroupID: 1, State: distribution.RouteStateActive}, + {RouteID: 3, Start: activeEnd, End: fencedStart, GroupID: 1, State: distribution.RouteStateActive}, + {RouteID: 4, Start: fencedStart, End: fencedEnd, GroupID: 2, State: distribution.RouteStateWriteFenced}, + {RouteID: 5, Start: fencedEnd, End: nil, GroupID: 1, State: distribution.RouteStateActive}, + }, + })) + + g1Txn := &recordingTransactional{} + g2Txn := &recordingTransactional{} + coord := NewShardedCoordinator(engine, map[uint64]*ShardGroup{ + 1: {Txn: g1Txn}, + 2: {Txn: g2Txn}, + }, 1, NewHLC(), nil) + + _, err := coord.Dispatch(context.Background(), &OperationGroup[OP]{ + Elems: []*Elem[OP]{{Op: DelPrefix, Key: s3keys.ObjectManifestPrefixForBucket(activeBucket, generation)}}, + }) + require.NoError(t, err) + require.NotEmpty(t, g1Txn.requests) + require.NotEmpty(t, g2Txn.requests, "DEL_PREFIX still broadcasts to every group after the narrow fence precheck passes") +} + // TestShardedCoordinator_DelPrefixRejectsTxn verifies that DEL_PREFIX inside // a transactional group is rejected. func TestShardedCoordinator_DelPrefixRejectsTxn(t *testing.T) { diff --git a/kv/sharded_coordinator_test.go b/kv/sharded_coordinator_test.go new file mode 100644 index 000000000..f2f1dd280 --- /dev/null +++ b/kv/sharded_coordinator_test.go @@ -0,0 +1,101 @@ +package kv + +import ( + "context" + "testing" + + "github.com/bootjp/elastickv/distribution" + "github.com/bootjp/elastickv/internal/s3keys" + "github.com/bootjp/elastickv/store" + "github.com/stretchr/testify/require" +) + +// The OCC read-key validation reads the same two namespaces as the store, and +// missing a version here is worse than a stale read: with no commit above +// startTS the transaction commits on a read it should have conflicted with. +// Promotion landing between a live-first pair hides exactly that version. +func TestValidateReadKeysOnShardReadsStagedBeforeLive(t *testing.T) { + t.Parallel() + + ctx := context.Background() + engine := distribution.NewEngine() + require.NoError(t, engine.ApplySnapshot(distribution.CatalogSnapshot{ + Version: 1, + Routes: []distribution.RouteDescriptor{{ + RouteID: 1, Start: []byte("a"), End: []byte("z"), GroupID: 1, + State: distribution.RouteStateActive, StagedVisibilityActive: true, + MigrationJobID: 9, MinWriteTSExclusive: 100, + }}, + })) + + rawKey := []byte("k") + stagedKey := distribution.MigrationStagedDataKey(9, rawKey) + inner := store.NewMVCCStore() + t.Cleanup(func() { _ = inner.Close() }) + // A staged commit above the transaction's startTS: the read must conflict. + require.NoError(t, inner.PutAt(ctx, stagedKey, []byte("staged"), 60, 0)) + + promoting := &promotingExportStore{MVCCStore: inner} + promoting.afterFirst = func() { + promoter, ok := inner.(store.MigrationPromoter) + require.True(t, ok) + _, err := promoter.PromoteVersions(ctx, store.PromoteVersionsOptions{ + JobID: 9, + StartKey: distribution.MigrationStagedDataKeyPrefix(9), + EndKey: prefixScanEnd(distribution.MigrationStagedDataKeyPrefix(9)), + MaxVersions: 16, + TargetKey: func(staged []byte) ([]byte, bool) { + _, raw, ok := distribution.MigrationStagedDataKeyParts(staged) + return raw, ok + }, + }) + require.NoError(t, err) + } + + c := &ShardedCoordinator{ + engine: engine, + groups: map[uint64]*ShardGroup{1: {Store: promoting}}, + } + + // latestCommitTSForReadKeyOnShard is the unit under test; + // validateReadKeysOnShard wraps it in a Raft read barrier that needs a + // live engine, which this case does not depend on. A missed version here + // is what makes validateReadKeysOnShard let the transaction through. + ts, exists, err := c.latestCommitTSForReadKeyOnShard(ctx, 1, c.groups[1], rawKey) + require.NoError(t, err) + require.True(t, exists, + "a promotion between the probes must not hide the conflicting commit") + require.Equal(t, uint64(60), ts) + require.Greater(t, ts, uint64(50), "the read must still conflict with startTS=50") +} + +func TestValidateReadKeysOnShardUsesS3BucketAuxiliaryOwnerForStagedProbe(t *testing.T) { + t.Parallel() + + ctx := context.Background() + const bucket = "bucket-a" + engine := distribution.NewEngine() + require.NoError(t, engine.ApplySnapshot(distribution.CatalogSnapshot{ + Version: 1, + Routes: s3BucketAuxiliaryStagedRoutes(bucket, 1, 2), + })) + + rawKey := s3keys.BucketMetaKey(bucket) + stagedKey := distribution.MigrationStagedDataKey(9, rawKey) + ownerStore := store.NewMVCCStore() + t.Cleanup(func() { _ = ownerStore.Close() }) + require.NoError(t, ownerStore.PutAt(ctx, stagedKey, []byte("staged"), 60, 0)) + + c := &ShardedCoordinator{ + engine: engine, + groups: map[uint64]*ShardGroup{ + 1: {Store: store.NewMVCCStore()}, + 2: {Store: ownerStore}, + }, + } + + ts, exists, err := c.latestCommitTSForReadKeyOnShard(ctx, 2, c.groups[2], rawKey) + require.NoError(t, err) + require.True(t, exists) + require.Equal(t, uint64(60), ts) +} diff --git a/kv/sharded_coordinator_txn_test.go b/kv/sharded_coordinator_txn_test.go index 67e5de8e6..5f39e6fb4 100644 --- a/kv/sharded_coordinator_txn_test.go +++ b/kv/sharded_coordinator_txn_test.go @@ -9,6 +9,7 @@ import ( "github.com/bootjp/elastickv/distribution" "github.com/bootjp/elastickv/internal/raftengine" + "github.com/bootjp/elastickv/internal/s3keys" "github.com/bootjp/elastickv/keyviz" pb "github.com/bootjp/elastickv/proto" "github.com/bootjp/elastickv/store" @@ -49,6 +50,158 @@ func (s *recordingTransactional) Abort(_ context.Context, _ []*pb.Request) (*Tra return &TransactionResponse{}, nil } +func TestShardedCoordinatorValidateReadKeysOnShard_UsesStagedVisibility(t *testing.T) { + t.Parallel() + + ctx := context.Background() + engine := distribution.NewEngine() + require.NoError(t, engine.ApplySnapshot(distribution.CatalogSnapshot{ + Version: 1, + Routes: []distribution.RouteDescriptor{ + { + RouteID: 1, + Start: []byte("a"), + End: []byte("z"), + GroupID: 1, + State: distribution.RouteStateActive, + StagedVisibilityActive: true, + MigrationJobID: 9, + }, + }, + })) + st := store.NewMVCCStore() + readKey := []byte("k") + require.NoError(t, st.PutAt(ctx, distribution.MigrationStagedDataKey(9, readKey), []byte("staged"), 20, 0)) + coord := NewShardedCoordinator(engine, map[uint64]*ShardGroup{ + 1: {Engine: stubLeaderEngine{}, Store: st}, + }, 1, NewHLC(), nil) + + err := coord.validateReadKeysOnShard(ctx, 1, [][]byte{readKey}, 10) + require.ErrorIs(t, err, store.ErrWriteConflict) +} + +func TestShardedCoordinatorDispatchTxn_AddsStagedReadKeyAlias(t *testing.T) { + t.Parallel() + + engine := distribution.NewEngine() + require.NoError(t, engine.ApplySnapshot(distribution.CatalogSnapshot{ + Version: 1, + Routes: []distribution.RouteDescriptor{ + { + RouteID: 1, + Start: []byte("a"), + End: []byte("z"), + GroupID: 1, + State: distribution.RouteStateActive, + StagedVisibilityActive: true, + MigrationJobID: 9, + }, + }, + })) + txn := &recordingTransactional{} + coord := NewShardedCoordinator(engine, map[uint64]*ShardGroup{ + 1: {Txn: txn}, + }, 1, NewHLC(), nil) + + readKey := []byte("k") + _, err := coord.Dispatch(context.Background(), &OperationGroup[OP]{ + IsTxn: true, + StartTS: 10, + CommitTS: 101, + Elems: []*Elem[OP]{{Op: Put, Key: []byte("m"), Value: []byte("write")}}, + ReadKeys: [][]byte{readKey}, + }) + require.NoError(t, err) + require.Len(t, txn.requests, 1) + require.Equal(t, [][]byte{ + readKey, + distribution.MigrationStagedDataKey(9, readKey), + distribution.MigrationStagedDataKey(9, []byte("m")), + }, txn.requests[0].ReadKeys) +} + +func TestShardedCoordinatorDispatchTxn_AddsStagedWriteKeyAlias(t *testing.T) { + t.Parallel() + + engine := distribution.NewEngine() + require.NoError(t, engine.ApplySnapshot(distribution.CatalogSnapshot{ + Version: 1, + Routes: []distribution.RouteDescriptor{ + { + RouteID: 1, + Start: []byte("a"), + End: []byte("z"), + GroupID: 1, + State: distribution.RouteStateActive, + StagedVisibilityActive: true, + MigrationJobID: 9, + }, + }, + })) + txn := &recordingTransactional{} + coord := NewShardedCoordinator(engine, map[uint64]*ShardGroup{ + 1: {Txn: txn}, + }, 1, NewHLC(), nil) + + writeKey := []byte("k") + _, err := coord.Dispatch(context.Background(), &OperationGroup[OP]{ + IsTxn: true, + StartTS: 10, + CommitTS: 101, + Elems: []*Elem[OP]{{Op: Put, Key: writeKey, Value: []byte("write")}}, + }) + require.NoError(t, err) + require.Len(t, txn.requests, 1) + require.Equal(t, [][]byte{ + distribution.MigrationStagedDataKey(9, writeKey), + }, txn.requests[0].ReadKeys) +} + +func TestShardedCoordinatorPrewrite_AddsStagedWriteKeyAlias(t *testing.T) { + t.Parallel() + + engine := distribution.NewEngine() + require.NoError(t, engine.ApplySnapshot(distribution.CatalogSnapshot{ + Version: 1, + Routes: []distribution.RouteDescriptor{ + { + RouteID: 1, + Start: []byte("a"), + End: []byte("m"), + GroupID: 1, + State: distribution.RouteStateActive, + StagedVisibilityActive: true, + MigrationJobID: 9, + }, + {RouteID: 2, Start: []byte("m"), End: []byte("z"), GroupID: 2, State: distribution.RouteStateActive}, + }, + })) + g1Txn := &recordingTransactional{} + g2Txn := &recordingTransactional{} + coord := NewShardedCoordinator(engine, map[uint64]*ShardGroup{ + 1: {Txn: g1Txn}, + 2: {Txn: g2Txn}, + }, 1, NewHLC(), nil) + + writeKey := []byte("b") + _, err := coord.Dispatch(context.Background(), &OperationGroup[OP]{ + IsTxn: true, + StartTS: 10, + CommitTS: 101, + Elems: []*Elem[OP]{ + {Op: Put, Key: writeKey, Value: []byte("write-b")}, + {Op: Put, Key: []byte("x"), Value: []byte("write-x")}, + }, + }) + require.NoError(t, err) + require.NotEmpty(t, g1Txn.requests) + require.NotEmpty(t, g2Txn.requests) + require.Equal(t, [][]byte{ + distribution.MigrationStagedDataKey(9, writeKey), + }, g1Txn.requests[0].ReadKeys) + require.Empty(t, g2Txn.requests[0].ReadKeys) +} + func cloneTxnRequest(req *pb.Request) *pb.Request { if req == nil { return nil @@ -441,43 +594,24 @@ func TestShardedCoordinatorDispatchTxn_UsesProvidedCommitTS(t *testing.T) { require.Equal(t, commitTS, commitMeta2.CommitTS) } -func TestShardedCoordinatorDispatchTxn_RejectsRouteWriteTimestampFloor(t *testing.T) { +func TestShardedCoordinatorDispatchTxn_RejectsMigrationTimestampFloor(t *testing.T) { t.Parallel() - engine := distribution.NewEngine() - require.NoError(t, engine.ApplySnapshot(distribution.CatalogSnapshot{ - Version: 1, - Routes: []distribution.RouteDescriptor{ - {RouteID: 1, Start: []byte(""), GroupID: 1, State: distribution.RouteStateActive, MinWriteTSExclusive: 20}, - }, - })) - g1Txn := &recordingTransactional{} - coord := NewShardedCoordinator(engine, map[uint64]*ShardGroup{ + coord := NewShardedCoordinator(newMigrationFloorEngine(t, 100), map[uint64]*ShardGroup{ 1: {Txn: g1Txn}, }, 1, NewHLC(), nil) _, err := coord.Dispatch(context.Background(), &OperationGroup[OP]{ IsTxn: true, - StartTS: 10, - CommitTS: 20, + StartTS: 90, + CommitTS: 100, Elems: []*Elem[OP]{ - {Op: Put, Key: []byte("b"), Value: []byte("v")}, + {Op: Put, Key: []byte("z"), Value: []byte("v")}, }, }) - require.ErrorIs(t, err, store.ErrWriteConflict) - require.Empty(t, g1Txn.requests) - - _, err = coord.Dispatch(context.Background(), &OperationGroup[OP]{ - IsTxn: true, - StartTS: 10, - CommitTS: 21, - Elems: []*Elem[OP]{ - {Op: Put, Key: []byte("b"), Value: []byte("v")}, - }, - }) - require.NoError(t, err) - require.Len(t, g1Txn.requests, 1) + require.ErrorIs(t, err, ErrRouteWriteTimestampTooLow) + require.Empty(t, g1Txn.requests, "coordinator must reject before preparing a floor-violating txn") } func TestCommitSecondaryWithRetry_RetriesAndSucceeds(t *testing.T) { @@ -609,6 +743,51 @@ func TestGroupReadKeysByShardID_FailsClosedOnUnroutable(t *testing.T) { require.ErrorIs(t, err, ErrInvalidRequest) } +func TestGroupReadKeysByShardID_RoutesS3BucketAuxiliaryToStagedOwner(t *testing.T) { + t.Parallel() + + const bucket = "bucket-a" + engine := distribution.NewEngine() + require.NoError(t, engine.ApplySnapshot(distribution.CatalogSnapshot{ + Version: 1, + Routes: s3BucketAuxiliaryStagedRoutes(bucket, 1, 2), + })) + coord := NewShardedCoordinator(engine, map[uint64]*ShardGroup{ + 1: {}, + 2: {}, + }, 1, NewHLC(), nil) + + key := s3keys.BucketMetaKey(bucket) + grouped, err := coord.groupReadKeysByShardID([][]byte{key}) + require.NoError(t, err) + require.Empty(t, grouped[1]) + require.Equal(t, [][]byte{ + key, + distribution.MigrationStagedDataKey(9, key), + }, grouped[2]) +} + +func TestGroupReadKeysByShardID_RoutesS3BucketAuxiliaryToPromotedOwner(t *testing.T) { + t.Parallel() + + const bucket = "bucket-a" + engine := distribution.NewEngine() + require.NoError(t, engine.ApplySnapshot(distribution.CatalogSnapshot{ + Version: 1, + Routes: s3BucketAuxiliaryPromotedRoutes(), + })) + coord := NewShardedCoordinator(engine, map[uint64]*ShardGroup{ + 1: {}, + 2: {}, + }, 1, NewHLC(), nil) + + key := s3keys.BucketMetaKey(bucket) + grouped, err := coord.groupReadKeysByShardID([][]byte{key}) + require.NoError(t, err) + require.Empty(t, grouped[1]) + require.Equal(t, [][]byte{key}, grouped[2]) +} + // --------------------------------------------------------------------------- // validateReadOnlyShards // --------------------------------------------------------------------------- @@ -915,3 +1094,44 @@ func TestShardedCoordinatorDispatchTxn_CrossShardPropagatesObservedRouteVersion( req.Phase) } } + +// In partition-resolved keyspaces such as HT-FIFO SQS, routeKey collapses a +// concrete partition key onto the global SQS route, so that route's write floor +// is not the key's floor. rejectWriteFencedPointKey already exempts these keys; +// the timestamp-floor precheck must match, or it rejects writes the fence +// precheck deliberately lets through. +func TestShardedCoordinatorFloorPrecheckSkipsResolverOwnedKeys(t *testing.T) { + t.Parallel() + + const partitionKey = "!sqs|msg|data|p|queue|7" + + engine := distribution.NewEngine() + require.NoError(t, engine.ApplySnapshot(distribution.CatalogSnapshot{ + Version: 1, + Routes: []distribution.RouteDescriptor{ + // The route routeKey() collapses partition keys onto, carrying a floor. + {RouteID: 1, Start: []byte(""), End: nil, GroupID: 1, State: distribution.RouteStateActive, MinWriteTSExclusive: 100}, + }, + })) + c := NewShardedCoordinator(engine, map[uint64]*ShardGroup{1: {}}, 1, NewHLC(), nil) + + // Without a resolver the floor applies, which is what makes the exemption + // below meaningful rather than vacuous. + require.ErrorIs(t, + c.rejectWriteTimestampFloorPointKey([]byte(partitionKey), 100), + ErrRouteWriteTimestampTooLow) + + c.WithPartitionResolver(&fakePartitionResolver{ + routes: map[string]uint64{partitionKey: 1}, + recognisedPrefix: []byte("!sqs|msg|data|p|"), + }) + + require.NoError(t, + c.rejectWriteTimestampFloorPointKey([]byte(partitionKey), 100), + "a resolver-owned key must not be judged by the route routeKey collapses it onto") + + // A key the resolver does not own still gets the floor. + require.ErrorIs(t, + c.rejectWriteTimestampFloorPointKey([]byte("ordinary-key"), 100), + ErrRouteWriteTimestampTooLow) +} diff --git a/kv/staged_visibility_candidates_benchmark_test.go b/kv/staged_visibility_candidates_benchmark_test.go new file mode 100644 index 000000000..6290ce1d7 --- /dev/null +++ b/kv/staged_visibility_candidates_benchmark_test.go @@ -0,0 +1,187 @@ +package kv + +import ( + "bytes" + "context" + "fmt" + "testing" + + "github.com/bootjp/elastickv/store" + "github.com/stretchr/testify/require" +) + +// seedSparseCandidateRange builds the shape latestCandidateVersionsAt is worst +// at: a handful of candidate keys spread across a range that is otherwise dense +// with versions nobody asked for. +func seedSparseCandidateRange(tb testing.TB, st store.MVCCStore, candidates, fillerPerGap, versionsPerFiller int) [][]byte { + tb.Helper() + + ctx := context.Background() + keys := make([][]byte, 0, candidates) + commitTS := uint64(1) + for c := range candidates { + key := []byte(fmt.Sprintf("k%06d", c*(fillerPerGap+1))) + require.NoError(tb, st.PutAt(ctx, key, []byte("candidate"), commitTS, 0)) + commitTS++ + keys = append(keys, key) + for f := 1; f <= fillerPerGap; f++ { + filler := []byte(fmt.Sprintf("k%06d", c*(fillerPerGap+1)+f)) + for v := range versionsPerFiller { + _ = v + require.NoError(tb, st.PutAt(ctx, filler, []byte("filler-value-padding"), commitTS, 0)) + commitTS++ + } + } + } + return keys +} + +// BenchmarkLatestCandidateVersionsAt measures the candidate resolution a staged +// visibility scan page performs. The sparse case is the one the scan budget +// exists for: the enclosing range holds far more versions than the candidate +// set, and an unbounded export decodes all of them. +func BenchmarkLatestCandidateVersionsAt(b *testing.B) { + for _, tc := range []struct { + name string + candidates int + fillerPerGap int + versionsPerFiller int + }{ + {name: "dense", candidates: 64, fillerPerGap: 0, versionsPerFiller: 0}, + {name: "sparse", candidates: 64, fillerPerGap: 64, versionsPerFiller: 16}, + } { + b.Run(tc.name, func(b *testing.B) { + st := store.NewMVCCStore() + b.Cleanup(func() { _ = st.Close() }) + keys := seedSparseCandidateRange(b, st, tc.candidates, tc.fillerPerGap, tc.versionsPerFiller) + ctx := context.Background() + readTS := ^uint64(0) >> 1 + + b.ReportAllocs() + b.ResetTimer() + for range b.N { + got, err := latestCandidateVersionsAt(ctx, st, keys, readTS) + if err != nil { + b.Fatal(err) + } + if len(got) != len(keys) { + b.Fatalf("resolved %d of %d candidates", len(got), len(keys)) + } + } + }) + } +} + +// The bounded export must still resolve every candidate: whatever it does not +// reach is probed by exact key, so the result is the same set either way. +func TestLatestCandidateVersionsAtResolvesSparseCandidates(t *testing.T) { + t.Parallel() + + st := store.NewMVCCStore() + t.Cleanup(func() { _ = st.Close() }) + keys := seedSparseCandidateRange(t, st, 32, 64, 16) + + got, err := latestCandidateVersionsAt(context.Background(), st, keys, ^uint64(0)>>1) + require.NoError(t, err) + require.Len(t, got, len(keys)) + for _, key := range keys { + version, ok := got[string(key)] + require.True(t, ok, "candidate %q must resolve", key) + require.Equal(t, []byte("candidate"), version.Value) + } +} + +// A candidate with no visible version stays absent rather than being invented +// by the probe fallback. +func TestLatestCandidateVersionsAtOmitsAbsentCandidates(t *testing.T) { + t.Parallel() + + ctx := context.Background() + st := store.NewMVCCStore() + t.Cleanup(func() { _ = st.Close() }) + keys := seedSparseCandidateRange(t, st, 4, 64, 8) + missing := []byte("k999999") + keys = append(keys, missing) + + got, err := latestCandidateVersionsAt(ctx, st, keys, ^uint64(0)>>1) + require.NoError(t, err) + require.Len(t, got, len(keys)-1) + _, ok := got[string(missing)] + require.False(t, ok) +} + +// exportRecordingStore captures the options every ExportVersions call carries. +type exportRecordingStore struct { + store.MVCCStore + scannedBudgets []uint64 + endKeys [][]byte + exports int +} + +func (s *exportRecordingStore) ExportVersions( + ctx context.Context, + opts store.ExportVersionsOptions, +) (store.ExportVersionsResult, error) { + s.exports++ + s.scannedBudgets = append(s.scannedBudgets, opts.MaxScannedBytes) + s.endKeys = append(s.endKeys, bytes.Clone(opts.EndKey)) + return s.MVCCStore.ExportVersions(ctx, opts) +} + +// The candidate export spans from the smallest candidate through the largest +// and filters for the candidate set, so an unbounded scan budget decodes every +// intervening version -- tombstones and dense MVCC history included -- for a +// page that wants at most a bounded number of exact keys. A sparse route could +// make one ordinary scan page consume unbounded I/O on the serving leader. +func TestLatestCandidateVersionsAtBoundsTheRangeScan(t *testing.T) { + t.Parallel() + + base := store.NewMVCCStore() + t.Cleanup(func() { _ = base.Close() }) + keys := seedSparseCandidateRange(t, base, 32, 64, 16) + recording := &exportRecordingStore{MVCCStore: base} + + got, err := latestCandidateVersionsAt(context.Background(), recording, keys, ^uint64(0)>>1) + require.NoError(t, err) + require.Len(t, got, len(keys), "every candidate still resolves") + + require.NotEmpty(t, recording.scannedBudgets) + require.Equal(t, uint64(stagedVisibilityCandidateScanBudget), recording.scannedBudgets[0], + "the range pass must carry a finite scan budget") + require.Greater(t, recording.exports, 1, + "a range pass that hit the budget must fall back to exact-key probes") +} + +// The probe fallback must cover exactly the candidate key. prefixScanEnd(key) +// covers every key that has it as a prefix, so probing an absent "a" would walk +// "ab", "az" and all of their versions before concluding "a" is not there -- +// unbounded work that the scan budget only splits into repeated exports, once +// per unresolved candidate. +func TestLatestCandidateVersionsAtProbesExactKeysOnly(t *testing.T) { + t.Parallel() + + base := store.NewMVCCStore() + t.Cleanup(func() { _ = base.Close() }) + keys := seedSparseCandidateRange(t, base, 32, 64, 16) + // A candidate with no version of its own, but with neighbours that share it + // as a prefix: a prefix-bounded probe would scan all of them. + absent := []byte("k000000-absent") + ctx := context.Background() + for i := range 64 { + require.NoError(t, base.PutAt(ctx, []byte(fmt.Sprintf("%s-%03d", absent, i)), []byte("neighbour"), 1, 0)) + } + keys = append(keys, absent) + + recording := &exportRecordingStore{MVCCStore: base} + got, err := latestCandidateVersionsAt(ctx, recording, keys, ^uint64(0)>>1) + require.NoError(t, err) + _, ok := got[string(absent)] + require.False(t, ok, "the absent candidate has no visible version") + + require.Contains(t, recording.endKeys, exactKeyScanEnd(absent), + "the probe must be bounded to the candidate key alone") + for _, end := range recording.endKeys[1:] { + require.NotEqual(t, prefixScanEnd(absent), end, + "no probe may use a prefix bound") + } +} diff --git a/kv/tso_test.go b/kv/tso_test.go index 348c9242d..50f6f4844 100644 --- a/kv/tso_test.go +++ b/kv/tso_test.go @@ -483,6 +483,21 @@ func TestShardedCoordinatorRawFollowerDefersTSOAllocationToLeaderPath(t *testing require.EqualValues(t, 0, txn.requests[0].Ts) } +func TestShardedCoordinatorRejectsTSORawPointWriteAfterStamping(t *testing.T) { + t.Parallel() + + txn := &recordingTransactional{} + coord := NewShardedCoordinator(newMigrationFloorEngine(t, testTSOInitialBase), map[uint64]*ShardGroup{ + 1: {Txn: txn}, + }, 1, NewHLC(), nil).WithTSOAllocator(&fakeTSOAllocator{nextBase: testTSOInitialBase, leader: true}) + + _, err := coord.Dispatch(context.Background(), &OperationGroup[OP]{ + Elems: []*Elem[OP]{{Op: Put, Key: []byte("z"), Value: []byte("v")}}, + }) + require.ErrorIs(t, err, ErrRouteWriteTimestampTooLow) + require.Empty(t, txn.requests, "coordinator must reject after TSO stamping before proposing") +} + func TestShardedCoordinatorUsesTSOAllocatorForRawTxnAndDelPrefix(t *testing.T) { t.Parallel() diff --git a/main.go b/main.go index 1c9847665..c38871fec 100644 --- a/main.go +++ b/main.go @@ -2004,6 +2004,7 @@ func startServersAfterStartupRotation(waitRotateOnStartup startupRotationWaiter, shardStore: in.shardStore, coordinate: adapterCoordinate, distServer: in.distServer, + routeEngine: in.cfg.engine, adminServer: adminServer, adminGRPCOpts: adminGRPCOpts, redisAddress: *redisAddr, @@ -3063,6 +3064,7 @@ func startRaftServers( forwardDeps adminForwardServerDeps, confChangeInterceptor internalraftadmin.MembershipChangeInterceptor, encWiring encryptionWriteWiring, + sqsPartitionResolver kv.PartitionResolver, defaultGroupID uint64, s3BlobObserver adapter.S3BlobOffloadObserver, s3BlobPushBlocked func() bool, @@ -3109,7 +3111,11 @@ func startRaftServers( rt.engine, coordinate.Clock(), relay, - internalOpts..., + append(internalOpts, + adapter.WithInternalStore(rt.store), + adapter.WithInternalMigrationProposer(proposerForGroup(rt, shardGroups)), + adapter.WithInternalMigrationExportRouting(rt.spec.id, sqsPartitionResolver), + )..., )) pb.RegisterDistributionServer(gs, distServer) registerAdminServerIfPresent(gs, adminServer) @@ -3585,6 +3591,7 @@ type runtimeServerRunner struct { shardStore *kv.ShardStore coordinate kv.Coordinator distServer *adapter.DistributionServer + routeEngine *distribution.Engine adminServer *adapter.AdminServer adminGRPCOpts adminGRPCInterceptors redisAddress string @@ -3696,6 +3703,10 @@ func (r *runtimeServerRunner) startRaftTransport() error { buckets: newBucketsSource(r.s3Server), roles: r.roleStore, } + var sqsPartitionResolver kv.PartitionResolver + if r.sqsPartitionResolver != nil { + sqsPartitionResolver = r.sqsPartitionResolver + } if err := startRaftServers( r.ctx, r.lc, @@ -3714,6 +3725,7 @@ func (r *runtimeServerRunner) startRaftTransport() error { forwardDeps, r.encryptionConfChangeInterceptor, r.encWiring, + sqsPartitionResolver, r.defaultGroupID, r.metricsRegistry.S3BlobOffloadObserver(), s3BlobPushBlocked, diff --git a/main_bootstrap_e2e_test.go b/main_bootstrap_e2e_test.go index 0298e8fc1..e912f031c 100644 --- a/main_bootstrap_e2e_test.go +++ b/main_bootstrap_e2e_test.go @@ -728,7 +728,7 @@ func startRuntimeServersWithBoundListeners( if err := startBoundRedisServer(ctx, eg, listeners.redis, shardStore, coordinate, leaderRedis, redisAddr, relay); err != nil { return waitErrgroupAfterStartupFailure(cancel, eg, err) } - if err := startBoundGRPCServer(ctx, eg, rt, shardStore, coordinate, distServer, relay, listeners.grpc); err != nil { + if err := startBoundGRPCServer(ctx, eg, rt, shardStore, coordinate, distServer, relay, nil, listeners.grpc); err != nil { return waitErrgroupAfterStartupFailure(cancel, eg, err) } if err := startBoundDynamoDBServer(ctx, eg, listeners.dynamo, shardStore, coordinate); err != nil { @@ -757,7 +757,7 @@ func startRuntimeServersWithBoundMultiGroupListeners( return waitErrgroupAfterStartupFailure(cancel, eg, err) } for _, rt := range runtimes { - if err := startBoundGRPCServer(ctx, eg, rt, shardStore, coordinate, distServer, relay, listeners.grpc[rt.spec.id]); err != nil { + if err := startBoundGRPCServer(ctx, eg, rt, shardStore, coordinate, distServer, relay, nil, listeners.grpc[rt.spec.id]); err != nil { return waitErrgroupAfterStartupFailure(cancel, eg, err) } } @@ -775,6 +775,7 @@ func startBoundGRPCServer( coordinate kv.Coordinator, distServer *adapter.DistributionServer, relay *adapter.RedisPubSubRelay, + sqsPartitionResolver kv.PartitionResolver, listener net.Listener, ) error { if rt == nil || rt.engine == nil { @@ -789,7 +790,15 @@ func startBoundGRPCServer( grpcSvc := adapter.NewGRPCServer(shardStore, coordinate) pb.RegisterRawKVServer(gs, grpcSvc) pb.RegisterTransactionalKVServer(gs, grpcSvc) - pb.RegisterInternalServer(gs, adapter.NewInternalWithEngine(trx, rt.engine, coordinate.Clock(), relay)) + pb.RegisterInternalServer(gs, adapter.NewInternalWithEngine( + trx, + rt.engine, + coordinate.Clock(), + relay, + adapter.WithInternalStore(rt.store), + adapter.WithInternalMigrationProposer(rt.engine), + adapter.WithInternalMigrationExportRouting(rt.spec.id, sqsPartitionResolver), + )) pb.RegisterDistributionServer(gs, distServer) rt.registerGRPC(gs) internalraftadmin.RegisterOperationalServices(ctx, gs, rt.engine, []string{"RawKV"}) diff --git a/proto/internal.pb.go b/proto/internal.pb.go index 94cdb3d50..7dec69e86 100644 --- a/proto/internal.pb.go +++ b/proto/internal.pb.go @@ -750,6 +750,15 @@ type ExportRangeVersionsRequest struct { RouteStart []byte `protobuf:"bytes,7,opt,name=route_start,json=routeStart,proto3" json:"route_start,omitempty"` RouteEnd []byte `protobuf:"bytes,8,opt,name=route_end,json=routeEnd,proto3" json:"route_end,omitempty"` MaxScannedBytes uint64 `protobuf:"varint,9,opt,name=max_scanned_bytes,json=maxScannedBytes,proto3" json:"max_scanned_bytes,omitempty"` + // Migration bracket family tag copied into exported MVCCVersion.key_family. + // Zero is invalid on the RPC path: callers must pass the bracket family they + // are exporting so target promotion can keep family-specific metadata. + KeyFamily uint32 `protobuf:"varint,10,opt,name=key_family,json=keyFamily,proto3" json:"key_family,omitempty"` + // Applies the user-bracket exclusion list for known internal families. + ExcludeKnownInternal bool `protobuf:"varint,11,opt,name=exclude_known_internal,json=excludeKnownInternal,proto3" json:"exclude_known_internal,omitempty"` + // Bracket-local raw-prefix exclusions, e.g. non-partitioned SQS brackets + // excluding their partitioned subprefixes. + ExcludePrefixes [][]byte `protobuf:"bytes,12,rep,name=exclude_prefixes,json=excludePrefixes,proto3" json:"exclude_prefixes,omitempty"` unknownFields protoimpl.UnknownFields sizeCache protoimpl.SizeCache } @@ -847,6 +856,27 @@ func (x *ExportRangeVersionsRequest) GetMaxScannedBytes() uint64 { return 0 } +func (x *ExportRangeVersionsRequest) GetKeyFamily() uint32 { + if x != nil { + return x.KeyFamily + } + return 0 +} + +func (x *ExportRangeVersionsRequest) GetExcludeKnownInternal() bool { + if x != nil { + return x.ExcludeKnownInternal + } + return false +} + +func (x *ExportRangeVersionsRequest) GetExcludePrefixes() [][]byte { + if x != nil { + return x.ExcludePrefixes + } + return nil +} + type ExportRangeVersionsResponse struct { state protoimpl.MessageState `protogen:"open.v1"` Versions []*MVCCVersion `protobuf:"bytes,1,rep,name=versions,proto3" json:"versions,omitempty"` @@ -1111,6 +1141,150 @@ func (x *ImportRangeVersionsResponse) GetAckedCursor() []byte { return nil } +type PromoteStagedVersionsRequest struct { + state protoimpl.MessageState `protogen:"open.v1"` + JobId uint64 `protobuf:"varint,1,opt,name=job_id,json=jobId,proto3" json:"job_id,omitempty"` + Cursor []byte `protobuf:"bytes,2,opt,name=cursor,proto3" json:"cursor,omitempty"` + MaxVersions uint32 `protobuf:"varint,3,opt,name=max_versions,json=maxVersions,proto3" json:"max_versions,omitempty"` + MaxBytes uint64 `protobuf:"varint,4,opt,name=max_bytes,json=maxBytes,proto3" json:"max_bytes,omitempty"` + MaxScannedBytes uint64 `protobuf:"varint,5,opt,name=max_scanned_bytes,json=maxScannedBytes,proto3" json:"max_scanned_bytes,omitempty"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *PromoteStagedVersionsRequest) Reset() { + *x = PromoteStagedVersionsRequest{} + mi := &file_internal_proto_msgTypes[16] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *PromoteStagedVersionsRequest) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*PromoteStagedVersionsRequest) ProtoMessage() {} + +func (x *PromoteStagedVersionsRequest) ProtoReflect() protoreflect.Message { + mi := &file_internal_proto_msgTypes[16] + if x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use PromoteStagedVersionsRequest.ProtoReflect.Descriptor instead. +func (*PromoteStagedVersionsRequest) Descriptor() ([]byte, []int) { + return file_internal_proto_rawDescGZIP(), []int{16} +} + +func (x *PromoteStagedVersionsRequest) GetJobId() uint64 { + if x != nil { + return x.JobId + } + return 0 +} + +func (x *PromoteStagedVersionsRequest) GetCursor() []byte { + if x != nil { + return x.Cursor + } + return nil +} + +func (x *PromoteStagedVersionsRequest) GetMaxVersions() uint32 { + if x != nil { + return x.MaxVersions + } + return 0 +} + +func (x *PromoteStagedVersionsRequest) GetMaxBytes() uint64 { + if x != nil { + return x.MaxBytes + } + return 0 +} + +func (x *PromoteStagedVersionsRequest) GetMaxScannedBytes() uint64 { + if x != nil { + return x.MaxScannedBytes + } + return 0 +} + +type PromoteStagedVersionsResponse struct { + state protoimpl.MessageState `protogen:"open.v1"` + NextCursor []byte `protobuf:"bytes,1,opt,name=next_cursor,json=nextCursor,proto3" json:"next_cursor,omitempty"` + Done bool `protobuf:"varint,2,opt,name=done,proto3" json:"done,omitempty"` + PromotedRows uint64 `protobuf:"varint,3,opt,name=promoted_rows,json=promotedRows,proto3" json:"promoted_rows,omitempty"` + MaxPromotedTs uint64 `protobuf:"varint,4,opt,name=max_promoted_ts,json=maxPromotedTs,proto3" json:"max_promoted_ts,omitempty"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *PromoteStagedVersionsResponse) Reset() { + *x = PromoteStagedVersionsResponse{} + mi := &file_internal_proto_msgTypes[17] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *PromoteStagedVersionsResponse) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*PromoteStagedVersionsResponse) ProtoMessage() {} + +func (x *PromoteStagedVersionsResponse) ProtoReflect() protoreflect.Message { + mi := &file_internal_proto_msgTypes[17] + if x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use PromoteStagedVersionsResponse.ProtoReflect.Descriptor instead. +func (*PromoteStagedVersionsResponse) Descriptor() ([]byte, []int) { + return file_internal_proto_rawDescGZIP(), []int{17} +} + +func (x *PromoteStagedVersionsResponse) GetNextCursor() []byte { + if x != nil { + return x.NextCursor + } + return nil +} + +func (x *PromoteStagedVersionsResponse) GetDone() bool { + if x != nil { + return x.Done + } + return false +} + +func (x *PromoteStagedVersionsResponse) GetPromotedRows() uint64 { + if x != nil { + return x.PromotedRows + } + return 0 +} + +func (x *PromoteStagedVersionsResponse) GetMaxPromotedTs() uint64 { + if x != nil { + return x.MaxPromotedTs + } + return 0 +} + var File_internal_proto protoreflect.FileDescriptor const file_internal_proto_rawDesc = "" + @@ -1150,7 +1324,7 @@ const file_internal_proto_rawDesc = "" + "\achannel\x18\x01 \x01(\fR\achannel\x12\x18\n" + "\amessage\x18\x02 \x01(\fR\amessage\"8\n" + "\x14RelayPublishResponse\x12 \n" + - "\vsubscribers\x18\x01 \x01(\x03R\vsubscribers\"\xc5\x02\n" + + "\vsubscribers\x18\x01 \x01(\x03R\vsubscribers\"\xc5\x03\n" + "\x1aExportRangeVersionsRequest\x12\x1f\n" + "\vrange_start\x18\x01 \x01(\fR\n" + "rangeStart\x12\x1b\n" + @@ -1163,7 +1337,12 @@ const file_internal_proto_rawDesc = "" + "\vroute_start\x18\a \x01(\fR\n" + "routeStart\x12\x1b\n" + "\troute_end\x18\b \x01(\fR\brouteEnd\x12*\n" + - "\x11max_scanned_bytes\x18\t \x01(\x04R\x0fmaxScannedBytes\"|\n" + + "\x11max_scanned_bytes\x18\t \x01(\x04R\x0fmaxScannedBytes\x12\x1d\n" + + "\n" + + "key_family\x18\n" + + " \x01(\rR\tkeyFamily\x124\n" + + "\x16exclude_known_internal\x18\v \x01(\bR\x14excludeKnownInternal\x12)\n" + + "\x10exclude_prefixes\x18\f \x03(\fR\x0fexcludePrefixes\"|\n" + "\x1bExportRangeVersionsResponse\x12(\n" + "\bversions\x18\x01 \x03(\v2\f.MVCCVersionR\bversions\x12\x1f\n" + "\vnext_cursor\x18\x02 \x01(\fR\n" + @@ -1185,7 +1364,19 @@ const file_internal_proto_rawDesc = "" + "bracket_id\x18\x04 \x01(\x04R\tbracketId\x12\x1b\n" + "\tbatch_seq\x18\x05 \x01(\x04R\bbatchSeq\"@\n" + "\x1bImportRangeVersionsResponse\x12!\n" + - "\facked_cursor\x18\x01 \x01(\fR\vackedCursor*&\n" + + "\facked_cursor\x18\x01 \x01(\fR\vackedCursor\"\xb9\x01\n" + + "\x1cPromoteStagedVersionsRequest\x12\x15\n" + + "\x06job_id\x18\x01 \x01(\x04R\x05jobId\x12\x16\n" + + "\x06cursor\x18\x02 \x01(\fR\x06cursor\x12!\n" + + "\fmax_versions\x18\x03 \x01(\rR\vmaxVersions\x12\x1b\n" + + "\tmax_bytes\x18\x04 \x01(\x04R\bmaxBytes\x12*\n" + + "\x11max_scanned_bytes\x18\x05 \x01(\x04R\x0fmaxScannedBytes\"\xa1\x01\n" + + "\x1dPromoteStagedVersionsResponse\x12\x1f\n" + + "\vnext_cursor\x18\x01 \x01(\fR\n" + + "nextCursor\x12\x12\n" + + "\x04done\x18\x02 \x01(\bR\x04done\x12#\n" + + "\rpromoted_rows\x18\x03 \x01(\x04R\fpromotedRows\x12&\n" + + "\x0fmax_promoted_ts\x18\x04 \x01(\x04R\rmaxPromotedTs*&\n" + "\x02Op\x12\a\n" + "\x03PUT\x10\x00\x12\a\n" + "\x03DEL\x10\x01\x12\x0e\n" + @@ -1196,14 +1387,15 @@ const file_internal_proto_rawDesc = "" + "\aPREPARE\x10\x01\x12\n" + "\n" + "\x06COMMIT\x10\x02\x12\t\n" + - "\x05ABORT\x10\x032\xc5\x03\n" + + "\x05ABORT\x10\x032\x9f\x04\n" + "\bInternal\x12.\n" + "\aForward\x12\x0f.ForwardRequest\x1a\x10.ForwardResponse\"\x00\x12U\n" + "\x14ForwardAdminProposal\x12\x1c.ForwardAdminProposalRequest\x1a\x1d.ForwardAdminProposalResponse\"\x00\x12I\n" + "\x10ForwardLeaseRead\x12\x18.ForwardLeaseReadRequest\x1a\x19.ForwardLeaseReadResponse\"\x00\x12=\n" + "\fRelayPublish\x12\x14.RelayPublishRequest\x1a\x15.RelayPublishResponse\"\x00\x12T\n" + "\x13ExportRangeVersions\x12\x1b.ExportRangeVersionsRequest\x1a\x1c.ExportRangeVersionsResponse\"\x000\x01\x12R\n" + - "\x13ImportRangeVersions\x12\x1b.ImportRangeVersionsRequest\x1a\x1c.ImportRangeVersionsResponse\"\x00B#Z!github.com/bootjp/elastickv/protob\x06proto3" + "\x13ImportRangeVersions\x12\x1b.ImportRangeVersionsRequest\x1a\x1c.ImportRangeVersionsResponse\"\x00\x12X\n" + + "\x15PromoteStagedVersions\x12\x1d.PromoteStagedVersionsRequest\x1a\x1e.PromoteStagedVersionsResponse\"\x00B#Z!github.com/bootjp/elastickv/protob\x06proto3" var ( file_internal_proto_rawDescOnce sync.Once @@ -1218,26 +1410,28 @@ func file_internal_proto_rawDescGZIP() []byte { } var file_internal_proto_enumTypes = make([]protoimpl.EnumInfo, 2) -var file_internal_proto_msgTypes = make([]protoimpl.MessageInfo, 16) +var file_internal_proto_msgTypes = make([]protoimpl.MessageInfo, 18) var file_internal_proto_goTypes = []any{ - (Op)(0), // 0: Op - (Phase)(0), // 1: Phase - (*Mutation)(nil), // 2: Mutation - (*Request)(nil), // 3: Request - (*RaftCommand)(nil), // 4: RaftCommand - (*ForwardRequest)(nil), // 5: ForwardRequest - (*ForwardResponse)(nil), // 6: ForwardResponse - (*ForwardAdminProposalRequest)(nil), // 7: ForwardAdminProposalRequest - (*ForwardAdminProposalResponse)(nil), // 8: ForwardAdminProposalResponse - (*ForwardLeaseReadRequest)(nil), // 9: ForwardLeaseReadRequest - (*ForwardLeaseReadResponse)(nil), // 10: ForwardLeaseReadResponse - (*RelayPublishRequest)(nil), // 11: RelayPublishRequest - (*RelayPublishResponse)(nil), // 12: RelayPublishResponse - (*ExportRangeVersionsRequest)(nil), // 13: ExportRangeVersionsRequest - (*ExportRangeVersionsResponse)(nil), // 14: ExportRangeVersionsResponse - (*MVCCVersion)(nil), // 15: MVCCVersion - (*ImportRangeVersionsRequest)(nil), // 16: ImportRangeVersionsRequest - (*ImportRangeVersionsResponse)(nil), // 17: ImportRangeVersionsResponse + (Op)(0), // 0: Op + (Phase)(0), // 1: Phase + (*Mutation)(nil), // 2: Mutation + (*Request)(nil), // 3: Request + (*RaftCommand)(nil), // 4: RaftCommand + (*ForwardRequest)(nil), // 5: ForwardRequest + (*ForwardResponse)(nil), // 6: ForwardResponse + (*ForwardAdminProposalRequest)(nil), // 7: ForwardAdminProposalRequest + (*ForwardAdminProposalResponse)(nil), // 8: ForwardAdminProposalResponse + (*ForwardLeaseReadRequest)(nil), // 9: ForwardLeaseReadRequest + (*ForwardLeaseReadResponse)(nil), // 10: ForwardLeaseReadResponse + (*RelayPublishRequest)(nil), // 11: RelayPublishRequest + (*RelayPublishResponse)(nil), // 12: RelayPublishResponse + (*ExportRangeVersionsRequest)(nil), // 13: ExportRangeVersionsRequest + (*ExportRangeVersionsResponse)(nil), // 14: ExportRangeVersionsResponse + (*MVCCVersion)(nil), // 15: MVCCVersion + (*ImportRangeVersionsRequest)(nil), // 16: ImportRangeVersionsRequest + (*ImportRangeVersionsResponse)(nil), // 17: ImportRangeVersionsResponse + (*PromoteStagedVersionsRequest)(nil), // 18: PromoteStagedVersionsRequest + (*PromoteStagedVersionsResponse)(nil), // 19: PromoteStagedVersionsResponse } var file_internal_proto_depIdxs = []int32{ 0, // 0: Mutation.op:type_name -> Op @@ -1253,14 +1447,16 @@ var file_internal_proto_depIdxs = []int32{ 11, // 10: Internal.RelayPublish:input_type -> RelayPublishRequest 13, // 11: Internal.ExportRangeVersions:input_type -> ExportRangeVersionsRequest 16, // 12: Internal.ImportRangeVersions:input_type -> ImportRangeVersionsRequest - 6, // 13: Internal.Forward:output_type -> ForwardResponse - 8, // 14: Internal.ForwardAdminProposal:output_type -> ForwardAdminProposalResponse - 10, // 15: Internal.ForwardLeaseRead:output_type -> ForwardLeaseReadResponse - 12, // 16: Internal.RelayPublish:output_type -> RelayPublishResponse - 14, // 17: Internal.ExportRangeVersions:output_type -> ExportRangeVersionsResponse - 17, // 18: Internal.ImportRangeVersions:output_type -> ImportRangeVersionsResponse - 13, // [13:19] is the sub-list for method output_type - 7, // [7:13] is the sub-list for method input_type + 18, // 13: Internal.PromoteStagedVersions:input_type -> PromoteStagedVersionsRequest + 6, // 14: Internal.Forward:output_type -> ForwardResponse + 8, // 15: Internal.ForwardAdminProposal:output_type -> ForwardAdminProposalResponse + 10, // 16: Internal.ForwardLeaseRead:output_type -> ForwardLeaseReadResponse + 12, // 17: Internal.RelayPublish:output_type -> RelayPublishResponse + 14, // 18: Internal.ExportRangeVersions:output_type -> ExportRangeVersionsResponse + 17, // 19: Internal.ImportRangeVersions:output_type -> ImportRangeVersionsResponse + 19, // 20: Internal.PromoteStagedVersions:output_type -> PromoteStagedVersionsResponse + 14, // [14:21] is the sub-list for method output_type + 7, // [7:14] is the sub-list for method input_type 7, // [7:7] is the sub-list for extension type_name 7, // [7:7] is the sub-list for extension extendee 0, // [0:7] is the sub-list for field type_name @@ -1277,7 +1473,7 @@ func file_internal_proto_init() { GoPackagePath: reflect.TypeOf(x{}).PkgPath(), RawDescriptor: unsafe.Slice(unsafe.StringData(file_internal_proto_rawDesc), len(file_internal_proto_rawDesc)), NumEnums: 2, - NumMessages: 16, + NumMessages: 18, NumExtensions: 0, NumServices: 1, }, diff --git a/proto/internal.proto b/proto/internal.proto index de11c805e..4713cf430 100644 --- a/proto/internal.proto +++ b/proto/internal.proto @@ -11,6 +11,7 @@ service Internal { rpc RelayPublish(RelayPublishRequest) returns (RelayPublishResponse) {} rpc ExportRangeVersions(ExportRangeVersionsRequest) returns (stream ExportRangeVersionsResponse) {} rpc ImportRangeVersions(ImportRangeVersionsRequest) returns (ImportRangeVersionsResponse) {} + rpc PromoteStagedVersions(PromoteStagedVersionsRequest) returns (PromoteStagedVersionsResponse) {} } // internal.proto is node to node communication message in raft replication. @@ -121,6 +122,15 @@ message ExportRangeVersionsRequest { bytes route_start = 7; bytes route_end = 8; uint64 max_scanned_bytes = 9; + // Migration bracket family tag copied into exported MVCCVersion.key_family. + // Zero is invalid on the RPC path: callers must pass the bracket family they + // are exporting so target promotion can keep family-specific metadata. + uint32 key_family = 10; + // Applies the user-bracket exclusion list for known internal families. + bool exclude_known_internal = 11; + // Bracket-local raw-prefix exclusions, e.g. non-partitioned SQS brackets + // excluding their partitioned subprefixes. + repeated bytes exclude_prefixes = 12; } message ExportRangeVersionsResponse { @@ -149,3 +159,18 @@ message ImportRangeVersionsRequest { message ImportRangeVersionsResponse { bytes acked_cursor = 1; } + +message PromoteStagedVersionsRequest { + uint64 job_id = 1; + bytes cursor = 2; + uint32 max_versions = 3; + uint64 max_bytes = 4; + uint64 max_scanned_bytes = 5; +} + +message PromoteStagedVersionsResponse { + bytes next_cursor = 1; + bool done = 2; + uint64 promoted_rows = 3; + uint64 max_promoted_ts = 4; +} diff --git a/proto/internal_grpc.pb.go b/proto/internal_grpc.pb.go index cd5733d80..51f41b71b 100644 --- a/proto/internal_grpc.pb.go +++ b/proto/internal_grpc.pb.go @@ -19,12 +19,13 @@ import ( const _ = grpc.SupportPackageIsVersion9 const ( - Internal_Forward_FullMethodName = "/Internal/Forward" - Internal_ForwardAdminProposal_FullMethodName = "/Internal/ForwardAdminProposal" - Internal_ForwardLeaseRead_FullMethodName = "/Internal/ForwardLeaseRead" - Internal_RelayPublish_FullMethodName = "/Internal/RelayPublish" - Internal_ExportRangeVersions_FullMethodName = "/Internal/ExportRangeVersions" - Internal_ImportRangeVersions_FullMethodName = "/Internal/ImportRangeVersions" + Internal_Forward_FullMethodName = "/Internal/Forward" + Internal_ForwardAdminProposal_FullMethodName = "/Internal/ForwardAdminProposal" + Internal_ForwardLeaseRead_FullMethodName = "/Internal/ForwardLeaseRead" + Internal_RelayPublish_FullMethodName = "/Internal/RelayPublish" + Internal_ExportRangeVersions_FullMethodName = "/Internal/ExportRangeVersions" + Internal_ImportRangeVersions_FullMethodName = "/Internal/ImportRangeVersions" + Internal_PromoteStagedVersions_FullMethodName = "/Internal/PromoteStagedVersions" ) // InternalClient is the client API for Internal service. @@ -38,6 +39,7 @@ type InternalClient interface { RelayPublish(ctx context.Context, in *RelayPublishRequest, opts ...grpc.CallOption) (*RelayPublishResponse, error) ExportRangeVersions(ctx context.Context, in *ExportRangeVersionsRequest, opts ...grpc.CallOption) (grpc.ServerStreamingClient[ExportRangeVersionsResponse], error) ImportRangeVersions(ctx context.Context, in *ImportRangeVersionsRequest, opts ...grpc.CallOption) (*ImportRangeVersionsResponse, error) + PromoteStagedVersions(ctx context.Context, in *PromoteStagedVersionsRequest, opts ...grpc.CallOption) (*PromoteStagedVersionsResponse, error) } type internalClient struct { @@ -117,6 +119,16 @@ func (c *internalClient) ImportRangeVersions(ctx context.Context, in *ImportRang return out, nil } +func (c *internalClient) PromoteStagedVersions(ctx context.Context, in *PromoteStagedVersionsRequest, opts ...grpc.CallOption) (*PromoteStagedVersionsResponse, error) { + cOpts := append([]grpc.CallOption{grpc.StaticMethod()}, opts...) + out := new(PromoteStagedVersionsResponse) + err := c.cc.Invoke(ctx, Internal_PromoteStagedVersions_FullMethodName, in, out, cOpts...) + if err != nil { + return nil, err + } + return out, nil +} + // InternalServer is the server API for Internal service. // All implementations must embed UnimplementedInternalServer // for forward compatibility. @@ -128,6 +140,7 @@ type InternalServer interface { RelayPublish(context.Context, *RelayPublishRequest) (*RelayPublishResponse, error) ExportRangeVersions(*ExportRangeVersionsRequest, grpc.ServerStreamingServer[ExportRangeVersionsResponse]) error ImportRangeVersions(context.Context, *ImportRangeVersionsRequest) (*ImportRangeVersionsResponse, error) + PromoteStagedVersions(context.Context, *PromoteStagedVersionsRequest) (*PromoteStagedVersionsResponse, error) mustEmbedUnimplementedInternalServer() } @@ -156,6 +169,9 @@ func (UnimplementedInternalServer) ExportRangeVersions(*ExportRangeVersionsReque func (UnimplementedInternalServer) ImportRangeVersions(context.Context, *ImportRangeVersionsRequest) (*ImportRangeVersionsResponse, error) { return nil, status.Error(codes.Unimplemented, "method ImportRangeVersions not implemented") } +func (UnimplementedInternalServer) PromoteStagedVersions(context.Context, *PromoteStagedVersionsRequest) (*PromoteStagedVersionsResponse, error) { + return nil, status.Error(codes.Unimplemented, "method PromoteStagedVersions not implemented") +} func (UnimplementedInternalServer) mustEmbedUnimplementedInternalServer() {} func (UnimplementedInternalServer) testEmbeddedByValue() {} @@ -278,6 +294,24 @@ func _Internal_ImportRangeVersions_Handler(srv interface{}, ctx context.Context, return interceptor(ctx, in, info, handler) } +func _Internal_PromoteStagedVersions_Handler(srv interface{}, ctx context.Context, dec func(interface{}) error, interceptor grpc.UnaryServerInterceptor) (interface{}, error) { + in := new(PromoteStagedVersionsRequest) + if err := dec(in); err != nil { + return nil, err + } + if interceptor == nil { + return srv.(InternalServer).PromoteStagedVersions(ctx, in) + } + info := &grpc.UnaryServerInfo{ + Server: srv, + FullMethod: Internal_PromoteStagedVersions_FullMethodName, + } + handler := func(ctx context.Context, req interface{}) (interface{}, error) { + return srv.(InternalServer).PromoteStagedVersions(ctx, req.(*PromoteStagedVersionsRequest)) + } + return interceptor(ctx, in, info, handler) +} + // Internal_ServiceDesc is the grpc.ServiceDesc for Internal service. // It's only intended for direct use with grpc.RegisterService, // and not to be introspected or modified (even as a copy) @@ -305,6 +339,10 @@ var Internal_ServiceDesc = grpc.ServiceDesc{ MethodName: "ImportRangeVersions", Handler: _Internal_ImportRangeVersions_Handler, }, + { + MethodName: "PromoteStagedVersions", + Handler: _Internal_PromoteStagedVersions_Handler, + }, }, Streams: []grpc.StreamDesc{ { diff --git a/proto/raw_latest_commit_ts_batch.go b/proto/raw_latest_commit_ts_batch.go new file mode 100644 index 000000000..f3782b9f2 --- /dev/null +++ b/proto/raw_latest_commit_ts_batch.go @@ -0,0 +1,58 @@ +package proto + +import ( + "encoding/binary" + "fmt" + "io" +) + +// EncodeRawLatestCommitTSKeyBatch packs exact-key version-presence probes into +// a single protobuf bytes field. The service-side decoder can reject an +// oversized count before allocating per-key slices. +func EncodeRawLatestCommitTSKeyBatch(keys [][]byte) []byte { + size := binary.MaxVarintLen64 + for _, key := range keys { + size += binary.MaxVarintLen64 + len(key) + } + out := make([]byte, 0, size) + out = binary.AppendUvarint(out, uint64(len(keys))) + for _, key := range keys { + out = binary.AppendUvarint(out, uint64(len(key))) + out = append(out, key...) + } + return out +} + +func DecodeRawLatestCommitTSKeyBatch(data []byte, maxKeys int) ([][]byte, error) { + if len(data) == 0 { + return nil, nil + } + if maxKeys < 0 { + return nil, fmt.Errorf("raw latest commit timestamp key batch max_keys must be non-negative") + } + count, n := binary.Uvarint(data) + if n <= 0 { + return nil, fmt.Errorf("raw latest commit timestamp key batch has invalid count") + } + if count > uint64(maxKeys) { + return nil, fmt.Errorf("raw latest commit timestamp key batch has %d keys, max %d", count, maxKeys) + } + data = data[n:] + keys := make([][]byte, 0, count) + for range count { + keyLen, n := binary.Uvarint(data) + if n <= 0 { + return nil, fmt.Errorf("raw latest commit timestamp key batch has invalid key length") + } + data = data[n:] + if keyLen > uint64(len(data)) { + return nil, io.ErrUnexpectedEOF + } + keys = append(keys, data[:keyLen]) + data = data[keyLen:] + } + if len(data) != 0 { + return nil, fmt.Errorf("raw latest commit timestamp key batch has trailing bytes") + } + return keys, nil +} diff --git a/proto/service.pb.go b/proto/service.pb.go index 5390269b5..e26eed3f7 100644 --- a/proto/service.pb.go +++ b/proto/service.pb.go @@ -413,8 +413,12 @@ type RawLatestCommitTSRequest struct { // below this timestamp. Comparing only `ts` cannot tell a tombstone at or // before the read timestamp apart from a newer version above it. VersionVisibleAtTs uint64 `protobuf:"varint,4,opt,name=version_visible_at_ts,json=versionVisibleAtTs,proto3" json:"version_visible_at_ts,omitempty"` - unknownFields protoimpl.UnknownFields - sizeCache protoimpl.SizeCache + // Optional batch form for version_visible_at_ts probes. Encoded as a single + // bounded key_batch blob so untrusted clients cannot expand one request into + // an unbounded repeated-bytes slice before the service enforces page limits. + KeyBatch []byte `protobuf:"bytes,5,opt,name=key_batch,json=keyBatch,proto3" json:"key_batch,omitempty"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache } func (x *RawLatestCommitTSRequest) Reset() { @@ -475,6 +479,13 @@ func (x *RawLatestCommitTSRequest) GetVersionVisibleAtTs() uint64 { return 0 } +func (x *RawLatestCommitTSRequest) GetKeyBatch() []byte { + if x != nil { + return x.KeyBatch + } + return nil +} + type RawLatestCommitTSResponse struct { state protoimpl.MessageState `protogen:"open.v1"` Ts uint64 `protobuf:"varint,1,opt,name=ts,proto3" json:"ts,omitempty"` @@ -482,8 +493,9 @@ type RawLatestCommitTSResponse struct { // Answer to version_visible_at_ts. Only meaningful when // version_visible_supported is set; a server that predates the probe leaves // both unset and the caller falls back to comparing `ts`. - VersionVisible bool `protobuf:"varint,3,opt,name=version_visible,json=versionVisible,proto3" json:"version_visible,omitempty"` - VersionVisibleSupported bool `protobuf:"varint,4,opt,name=version_visible_supported,json=versionVisibleSupported,proto3" json:"version_visible_supported,omitempty"` + VersionVisible bool `protobuf:"varint,3,opt,name=version_visible,json=versionVisible,proto3" json:"version_visible,omitempty"` + VersionVisibleSupported bool `protobuf:"varint,4,opt,name=version_visible_supported,json=versionVisibleSupported,proto3" json:"version_visible_supported,omitempty"` + VersionVisibleResults []bool `protobuf:"varint,5,rep,packed,name=version_visible_results,json=versionVisibleResults,proto3" json:"version_visible_results,omitempty"` unknownFields protoimpl.UnknownFields sizeCache protoimpl.SizeCache } @@ -546,6 +558,13 @@ func (x *RawLatestCommitTSResponse) GetVersionVisibleSupported() bool { return false } +func (x *RawLatestCommitTSResponse) GetVersionVisibleResults() []bool { + if x != nil { + return x.VersionVisibleResults + } + return nil +} + type RawScanAtRequest struct { state protoimpl.MessageState `protogen:"open.v1"` StartKey []byte `protobuf:"bytes,1,opt,name=start_key,json=startKey,proto3" json:"start_key,omitempty"` @@ -2593,17 +2612,19 @@ const file_service_proto_rawDesc = "" + "\x03key\x18\x01 \x01(\fR\x03key\"P\n" + "\x11RawDeleteResponse\x12!\n" + "\fcommit_index\x18\x01 \x01(\x04R\vcommitIndex\x12\x18\n" + - "\asuccess\x18\x02 \x01(\bR\asuccess\"\xa8\x01\n" + + "\asuccess\x18\x02 \x01(\bR\asuccess\"\xc5\x01\n" + "\x18RawLatestCommitTSRequest\x12\x10\n" + "\x03key\x18\x01 \x01(\fR\x03key\x12,\n" + "\x12read_route_version\x18\x02 \x01(\x04R\x10readRouteVersion\x12\x19\n" + "\bgroup_id\x18\x03 \x01(\x04R\agroupId\x121\n" + - "\x15version_visible_at_ts\x18\x04 \x01(\x04R\x12versionVisibleAtTs\"\xa8\x01\n" + + "\x15version_visible_at_ts\x18\x04 \x01(\x04R\x12versionVisibleAtTs\x12\x1b\n" + + "\tkey_batch\x18\x05 \x01(\fR\bkeyBatch\"\xe0\x01\n" + "\x19RawLatestCommitTSResponse\x12\x0e\n" + "\x02ts\x18\x01 \x01(\x04R\x02ts\x12\x16\n" + "\x06exists\x18\x02 \x01(\bR\x06exists\x12'\n" + "\x0fversion_visible\x18\x03 \x01(\bR\x0eversionVisible\x12:\n" + - "\x19version_visible_supported\x18\x04 \x01(\bR\x17versionVisibleSupported\"\xde\x02\n" + + "\x19version_visible_supported\x18\x04 \x01(\bR\x17versionVisibleSupported\x126\n" + + "\x17version_visible_results\x18\x05 \x03(\bR\x15versionVisibleResults\"\xde\x02\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" + diff --git a/proto/service.proto b/proto/service.proto index 6a0ccf570..b4f13b929 100644 --- a/proto/service.proto +++ b/proto/service.proto @@ -76,6 +76,10 @@ message RawLatestCommitTSRequest { // below this timestamp. Comparing only `ts` cannot tell a tombstone at or // before the read timestamp apart from a newer version above it. uint64 version_visible_at_ts = 4; + // Optional batch form for version_visible_at_ts probes. Encoded as a single + // bounded key_batch blob so untrusted clients cannot expand one request into + // an unbounded repeated-bytes slice before the service enforces page limits. + bytes key_batch = 5; } message RawLatestCommitTSResponse { @@ -86,6 +90,7 @@ message RawLatestCommitTSResponse { // both unset and the caller falls back to comparing `ts`. bool version_visible = 3; bool version_visible_supported = 4; + repeated bool version_visible_results = 5; } message RawScanAtRequest { diff --git a/store/lsm_migration.go b/store/lsm_migration.go index cd9716ef4..de2f6a1aa 100644 --- a/store/lsm_migration.go +++ b/store/lsm_migration.go @@ -10,6 +10,13 @@ import ( ) func (s *pebbleStore) ExportVersions(ctx context.Context, opts ExportVersionsOptions) (ExportVersionsResult, error) { + s.dbMu.RLock() + defer s.dbMu.RUnlock() + + return s.exportVersionsLocked(ctx, opts) +} + +func (s *pebbleStore) exportVersionsLocked(ctx context.Context, opts ExportVersionsOptions) (ExportVersionsResult, error) { opts = normalizeExportVersionsOptions(opts) pos, err := decodeExportCursorForOptions(opts) if err != nil { @@ -19,9 +26,6 @@ func (s *pebbleStore) ExportVersions(ctx context.Context, opts ExportVersionsOpt return ExportVersionsResult{}, errors.WithStack(ErrInvalidExportBudget) } - s.dbMu.RLock() - defer s.dbMu.RUnlock() - iter, err := s.db.NewIter(pebbleExportIterOptions(opts)) if err != nil { return ExportVersionsResult{}, errors.WithStack(err) @@ -248,26 +252,17 @@ func (s *pebbleStore) exportPebbleVersion( commitTS uint64, result *ExportVersionsResult, ) (bool, error) { - tag := exportCursorTagScanned rawValue := iter.Value() result.ScannedBytes += versionExportSize(userKey, len(rawValue)) + tag := exportCursorTagScanned if shouldExportPebbleVersion(opts, userKey, commitTS) { - version, err := s.decodeExportedPebbleVersion(iter, userKey, commitTS, opts.KeyFamily) - if err != nil { + emitted, cont, err := s.appendExportedPebbleVersion(iter, opts, userKey, commitTS, result) + if err != nil || !cont { return false, err } - if opts.AcceptVersion != nil && !opts.AcceptVersion(version.Key, version.Value) { - result.NextCursor = encodeExportCursor(userKey, commitTS, exportCursorTagScanned) - if finishExportIfLimited(opts, result) { - result.Done = false - return false, nil - } - return true, nil + if emitted { + tag = exportCursorTagEmitted } - result.Versions = append(result.Versions, version) - result.ExportedBytes += versionExportSize(userKey, len(version.Value)) - result.AcceptedRows++ - tag = exportCursorTagEmitted } result.NextCursor = encodeExportCursor(userKey, commitTS, tag) if finishExportIfLimited(opts, result) { @@ -277,6 +272,41 @@ func (s *pebbleStore) exportPebbleVersion( return true, nil } +// appendExportedPebbleVersion decodes the version at the iterator and puts it +// on the page. It returns emitted=false when a filter rejected the version, +// and cont=false when the page had to stop: either because the row would push +// a page that already holds rows past its byte budget -- the cursor stays on +// the previous row so this one starts the next page -- or because a rejected +// version finished the chunk. +func (s *pebbleStore) appendExportedPebbleVersion( + iter *pebble.Iterator, + opts ExportVersionsOptions, + userKey []byte, + commitTS uint64, + result *ExportVersionsResult, +) (emitted bool, cont bool, err error) { + version, err := s.decodeExportedPebbleVersion(iter, userKey, commitTS, opts.KeyFamily) + if err != nil { + return false, false, err + } + if opts.AcceptVersion != nil && !opts.AcceptVersion(version.Key, version.Value) { + result.NextCursor = encodeExportCursor(userKey, commitTS, exportCursorTagScanned) + if finishExportIfLimited(opts, result) { + result.Done = false + return false, false, nil + } + return false, true, nil + } + if exportPageWouldOverflow(opts, result, userKey, len(version.Value)) { + result.Done = false + return false, false, nil + } + result.Versions = append(result.Versions, version) + result.ExportedBytes += versionExportSize(userKey, len(version.Value)) + result.AcceptedRows++ + return true, true, nil +} + func shouldExportPebbleVersion(opts ExportVersionsOptions, userKey []byte, commitTS uint64) bool { if shouldSkipMigrationExportKey(userKey) { return false @@ -310,6 +340,14 @@ func (s *pebbleStore) decodeExportedPebbleVersion(iter *pebble.Iterator, userKey } func (s *pebbleStore) ImportVersions(ctx context.Context, opts ImportVersionsOptions) (ImportVersionsResult, error) { + return s.importVersionsWithOpts(ctx, opts, s.directApplyWriteOpts(), true) +} + +func (s *pebbleStore) ImportVersionsRaft(ctx context.Context, opts ImportVersionsOptions) (ImportVersionsResult, error) { + return s.importVersionsWithOpts(ctx, opts, s.raftApplyWriteOpts(), false) +} + +func (s *pebbleStore) importVersionsWithOpts(ctx context.Context, opts ImportVersionsOptions, writeOpts *pebble.WriteOptions, gateRegistration bool) (ImportVersionsResult, error) { s.dbMu.RLock() defer s.dbMu.RUnlock() @@ -321,11 +359,14 @@ func (s *pebbleStore) ImportVersions(ctx context.Context, opts ImportVersionsOpt return ImportVersionsResult{}, err } if duplicate { + if err := s.commitPebbleImportAppliedIndex(opts.AppliedIndex, writeOpts); err != nil { + return ImportVersionsResult{}, err + } return ImportVersionsResult{AckedCursor: ackedCursor, Duplicate: true}, nil } batchMax := importBatchMaxTS(opts.Versions) - if err := s.commitPebbleImportBatch(opts, batchMax); err != nil { + if err := s.commitPebbleImportBatch(opts, batchMax, writeOpts, gateRegistration); err != nil { return ImportVersionsResult{}, errors.WithStack(err) } s.log.InfoContext(ctx, "import_versions", @@ -338,6 +379,18 @@ func (s *pebbleStore) ImportVersions(ctx context.Context, opts ImportVersionsOpt return ImportVersionsResult{AckedCursor: bytes.Clone(opts.Cursor), MaxImportedTS: batchMax}, nil } +func (s *pebbleStore) commitPebbleImportAppliedIndex(appliedIndex uint64, writeOpts *pebble.WriteOptions) error { + if appliedIndex == 0 { + return nil + } + batch := s.db.NewBatch() + defer batch.Close() + if err := setPebbleUint64InBatch(batch, metaAppliedIndexBytes, appliedIndex); err != nil { + return err + } + return errors.WithStack(batch.Commit(writeOpts)) +} + func (s *pebbleStore) validatePebbleImportBatch(opts ImportVersionsOptions) (bool, []byte, error) { existing, hasExisting, err := s.readMigrationImportAck(opts.JobID, opts.BracketID) if err != nil { @@ -358,10 +411,10 @@ func (s *pebbleStore) validatePebbleImportBatch(opts ImportVersionsOptions) (boo return false, nil, nil } -func (s *pebbleStore) commitPebbleImportBatch(opts ImportVersionsOptions, batchMax uint64) error { +func (s *pebbleStore) commitPebbleImportBatch(opts ImportVersionsOptions, batchMax uint64, writeOpts *pebble.WriteOptions, gateRegistration bool) error { batch := s.db.NewBatch() defer batch.Close() - if err := s.applyImportVersionsBatch(batch, opts.Versions); err != nil { + if err := s.applyImportVersionsBatch(batch, opts.Versions, gateRegistration); err != nil { return err } if err := s.stageMigrationImportAck(batch, opts.JobID, opts.BracketID, migrationImportAck{ @@ -375,7 +428,10 @@ func (s *pebbleStore) commitPebbleImportBatch(opts ImportVersionsOptions, batchM return err } defer unlock() - if err := batch.Commit(s.directApplyWriteOpts()); err != nil { + if err := stagePebbleAppliedIndex(batch, opts.AppliedIndex); err != nil { + return err + } + if err := batch.Commit(writeOpts); err != nil { return errors.WithStack(err) } if batchMax > 0 { @@ -384,6 +440,13 @@ func (s *pebbleStore) commitPebbleImportBatch(opts ImportVersionsOptions, batchM return nil } +func stagePebbleAppliedIndex(batch *pebble.Batch, appliedIndex uint64) error { + if appliedIndex == 0 { + return nil + } + return setPebbleUint64InBatch(batch, metaAppliedIndexBytes, appliedIndex) +} + func (s *pebbleStore) stageMigrationImportAck(batch *pebble.Batch, jobID, bracketID uint64, ack migrationImportAck) error { acks, err := s.readMigrationImportAcks() if err != nil { @@ -403,7 +466,7 @@ func (s *pebbleStore) stageMigrationClockMetadataIfNeeded(batch *pebble.Batch, j return s.stageMigrationClockMetadata(batch, jobID, batchMax) } -func (s *pebbleStore) applyImportVersionsBatch(batch *pebble.Batch, versions []MVCCVersion) error { +func (s *pebbleStore) applyImportVersionsBatch(batch *pebble.Batch, versions []MVCCVersion, gateRegistration bool) error { for _, version := range versions { k, err := encodePebbleUserVersionKey(version.Key, version.CommitTS) if err != nil { @@ -413,7 +476,7 @@ func (s *pebbleStore) applyImportVersionsBatch(batch *pebble.Batch, versions []M if version.Tombstone { encoded = encodeValue(nil, true, 0, encStateCleartext) } else { - body, encState, err := s.encryptForKey(k, version.Value, version.ExpireAt, true) + body, encState, err := s.encryptForKey(k, version.Value, version.ExpireAt, gateRegistration) if err != nil { return err } @@ -531,6 +594,9 @@ func (s *pebbleStore) RetireMigration(ctx context.Context, jobID uint64) error { if err := s.stageRetireMigrationHLCFloor(batch, jobID); err != nil { return err } + if err := s.stageRetireMigrationPromotionState(batch, jobID); err != nil { + return err + } return errors.WithStack(batch.Commit(s.directApplyWriteOpts())) } @@ -556,6 +622,15 @@ func (s *pebbleStore) stageRetireMigrationHLCFloor(batch *pebble.Batch, jobID ui return stageMigrationMetadataMap(batch, migrationHLCFloorMetaKeyBytes, len(floors), encodeMigrationHLCFloors(floors)) } +func (s *pebbleStore) stageRetireMigrationPromotionState(batch *pebble.Batch, jobID uint64) error { + states, err := s.readPebblePromotionStates() + if err != nil { + return err + } + delete(states, jobID) + return stageMigrationMetadataMap(batch, migrationPromoteMetaKeyBytes, len(states), encodeMigrationPromotionStates(states)) +} + func stageMigrationMetadataMap(batch *pebble.Batch, key []byte, entries int, encoded []byte) error { if entries == 0 { return errors.WithStack(batch.Delete(key, nil)) diff --git a/store/lsm_store.go b/store/lsm_store.go index 628e536bd..6cb2152e0 100644 --- a/store/lsm_store.go +++ b/store/lsm_store.go @@ -55,10 +55,10 @@ const ( // maxPebbleEncodedKeySize is the limit for encoded Pebble on-disk keys, // which are the user key concatenated with the 8-byte inverted timestamp. - // Using maxSnapshotKeySize+timestampSize (instead of just maxSnapshotKeySize) - // avoids rejecting keys that are valid at the user-key level but slightly - // exceed maxSnapshotKeySize once the timestamp suffix is appended. - maxPebbleEncodedKeySize = maxSnapshotKeySize + timestampSize + // Using maxSnapshotStoredKeySize+timestampSize avoids rejecting logical keys + // that fit maxSnapshotKeySize but gain a bounded internal envelope before the + // timestamp suffix is appended. + maxPebbleEncodedKeySize = maxSnapshotStoredKeySize + timestampSize // defaultPebbleCacheBytes is the fallback process-wide Pebble block-cache // capacity when the node's effective memory budget cannot be discovered. @@ -689,6 +689,9 @@ func writeTempDBMetadata(db *pebble.DB, meta streamingMVCCRestoreMetadata) error if err := batch.Set(migrationHLCFloorMetaKeyBytes, encodeMigrationHLCFloors(meta.migrationHLCFloors), nil); err != nil { return errors.WithStack(err) } + if err := batch.Set(migrationPromoteMetaKeyBytes, encodeMigrationPromotionStates(meta.migrationPromotions), nil); err != nil { + return errors.WithStack(err) + } return errors.WithStack(batch.Commit(pebble.Sync)) } @@ -2538,7 +2541,7 @@ func (s *pebbleStore) stageLastCommitTSInBatch(b *pebble.Batch, commitTS uint64, // ELASTICKV_FSM_SYNC_MODE=nosync. Raft-apply callers must use // DeletePrefixAtRaft instead. func (s *pebbleStore) DeletePrefixAt(ctx context.Context, prefix []byte, excludePrefix []byte, commitTS uint64) error { - return s.deletePrefixAtWithOpts(ctx, prefix, excludePrefix, commitTS, s.directApplyWriteOpts(), 0) + return s.deletePrefixesAtWithOpts(ctx, []PrefixDelete{{Prefix: prefix, ExcludePrefix: excludePrefix}}, commitTS, s.directApplyWriteOpts(), 0) } // DeletePrefixAtRaft is the raft-apply variant of DeletePrefixAt. Durability @@ -2549,7 +2552,7 @@ func (s *pebbleStore) DeletePrefixAt(ctx context.Context, prefix []byte, exclude // DeletePrefixAtRaftAt to bundle metaAppliedIndex atomically — see // PR #910 design §2 "why both leaves". func (s *pebbleStore) DeletePrefixAtRaft(ctx context.Context, prefix []byte, excludePrefix []byte, commitTS uint64) error { - return s.deletePrefixAtWithOpts(ctx, prefix, excludePrefix, commitTS, s.raftApplyWriteOpts(), 0) + return s.deletePrefixesAtWithOpts(ctx, []PrefixDelete{{Prefix: prefix, ExcludePrefix: excludePrefix}}, commitTS, s.raftApplyWriteOpts(), 0) } // DeletePrefixAtRaftAt is DeletePrefixAtRaft with the raft entry @@ -2560,24 +2563,25 @@ func (s *pebbleStore) DeletePrefixAtRaft(ctx context.Context, prefix []byte, exc // LastAppliedIndex behind the true applied count for any workload // that uses DEL_PREFIX. PR #910 design §2. func (s *pebbleStore) DeletePrefixAtRaftAt(ctx context.Context, prefix []byte, excludePrefix []byte, commitTS, appliedIndex uint64) error { - return s.deletePrefixAtWithOpts(ctx, prefix, excludePrefix, commitTS, s.raftApplyWriteOpts(), appliedIndex) + return s.DeletePrefixesAtRaftAt(ctx, []PrefixDelete{{Prefix: prefix, ExcludePrefix: excludePrefix}}, commitTS, appliedIndex) +} + +func (s *pebbleStore) DeletePrefixesAtRaftAt(ctx context.Context, deletes []PrefixDelete, commitTS, appliedIndex uint64) error { + return s.deletePrefixesAtWithOpts(ctx, deletes, commitTS, s.raftApplyWriteOpts(), appliedIndex) } -func (s *pebbleStore) deletePrefixAtWithOpts(_ context.Context, prefix []byte, excludePrefix []byte, commitTS uint64, writeOpts *pebble.WriteOptions, appliedIndex uint64) error { +func (s *pebbleStore) deletePrefixesAtWithOpts(_ context.Context, deletes []PrefixDelete, commitTS uint64, writeOpts *pebble.WriteOptions, appliedIndex uint64) error { + if len(deletes) == 0 { + return nil + } + s.dbMu.RLock() defer s.dbMu.RUnlock() s.applyMu.Lock() defer s.applyMu.Unlock() - var lowerBound []byte - if len(prefix) > 0 { - lowerBound = encodeKey(prefix, math.MaxUint64) - } - - iter, err := s.db.NewIter(&pebble.IterOptions{ - LowerBound: lowerBound, - }) + iter, err := s.newDeletePrefixesIterator(deletes) if err != nil { return errors.WithStack(err) } @@ -2586,10 +2590,34 @@ func (s *pebbleStore) deletePrefixAtWithOpts(_ context.Context, prefix []byte, e batch := s.db.NewBatch() defer batch.Close() - if err := s.scanDeletePrefix(iter, batch, prefix, excludePrefix, commitTS); err != nil { + if err := s.stageDeletePrefixes(iter, batch, deletes, commitTS); err != nil { return err } + return s.commitDeletePrefixesBatch(batch, commitTS, writeOpts, appliedIndex) +} +func (s *pebbleStore) newDeletePrefixesIterator(deletes []PrefixDelete) (*pebble.Iterator, error) { + var lowerBound []byte + if len(deletes) == 1 && len(deletes[0].Prefix) > 0 { + lowerBound = encodeKey(deletes[0].Prefix, math.MaxUint64) + } + iter, err := s.db.NewIter(&pebble.IterOptions{LowerBound: lowerBound}) + if err != nil { + return nil, errors.WithStack(err) + } + return iter, nil +} + +func (s *pebbleStore) stageDeletePrefixes(iter *pebble.Iterator, batch *pebble.Batch, deletes []PrefixDelete, commitTS uint64) error { + for _, del := range deletes { + if err := s.scanDeletePrefix(iter, batch, del.Prefix, del.ExcludePrefix, commitTS); err != nil { + return err + } + } + return nil +} + +func (s *pebbleStore) commitDeletePrefixesBatch(batch *pebble.Batch, commitTS uint64, writeOpts *pebble.WriteOptions, appliedIndex uint64) error { // Persist lastCommitTS update atomically with the tombstones. s.mtx.Lock() defer s.mtx.Unlock() @@ -2614,7 +2642,6 @@ func (s *pebbleStore) deletePrefixAtWithOpts(_ context.Context, prefix []byte, e return errors.WithStack(err) } s.updateLastCommitTS(newLastTS) - return nil } @@ -3361,10 +3388,11 @@ func writeNativeSnapshotToTempDir(r io.Reader, tmpDir string, ts uint64) error { // place after the CRC32 checksum is verified, preserving the existing store // on failure. type streamingMVCCRestoreMetadata struct { - lastCommitTS uint64 - minRetainedTS uint64 - migrationAcks map[migrationAckID]migrationImportAck - migrationHLCFloors map[uint64]uint64 + lastCommitTS uint64 + minRetainedTS uint64 + migrationAcks map[migrationAckID]migrationImportAck + migrationHLCFloors map[uint64]uint64 + migrationPromotions map[uint64]PromotionState } func readStreamingMVCCRestoreHeader(r io.Reader) (io.Reader, hash.Hash32, uint32, streamingMVCCRestoreMetadata, error) { @@ -3375,15 +3403,16 @@ func readStreamingMVCCRestoreHeader(r io.Reader) (io.Reader, hash.Hash32, uint32 hash := crc32.NewIEEE() body := io.TeeReader(r, hash) - lastCommitTS, minRetainedTS, migrationAcks, migrationHLCFloors, err := readMVCCSnapshotMetadata(body, version) + lastCommitTS, minRetainedTS, migrationAcks, migrationHLCFloors, migrationPromotions, err := readMVCCSnapshotMetadata(body, version) if err != nil { return nil, nil, 0, streamingMVCCRestoreMetadata{}, err } meta := streamingMVCCRestoreMetadata{ - lastCommitTS: lastCommitTS, - minRetainedTS: minRetainedTS, - migrationAcks: migrationAcks, - migrationHLCFloors: migrationHLCFloors, + lastCommitTS: lastCommitTS, + minRetainedTS: minRetainedTS, + migrationAcks: migrationAcks, + migrationHLCFloors: migrationHLCFloors, + migrationPromotions: migrationPromotions, } return body, hash, expectedChecksum, meta, nil } diff --git a/store/lsm_store_applied_index_test.go b/store/lsm_store_applied_index_test.go index 8a30710d3..bad8be62a 100644 --- a/store/lsm_store_applied_index_test.go +++ b/store/lsm_store_applied_index_test.go @@ -1,6 +1,7 @@ package store import ( + "bytes" "context" "encoding/binary" "os" @@ -97,6 +98,118 @@ func TestApplyMutationsRaftAt_BundlesMetaAppliedIndex(t *testing.T) { require.Equal(t, []byte("v1"), val) } +func TestPromoteVersions_BundlesMetaAppliedIndex(t *testing.T) { + ctx := context.Background() + st := newApplyIndexPebbleStore(t) + ps := pebbleStoreApplied(t, st) + + stage := func(raw string) []byte { + return append([]byte("stage|"), []byte(raw)...) + } + targetKey := func(staged []byte) ([]byte, bool) { + return bytes.TrimPrefix(staged, []byte("stage|")), bytes.HasPrefix(staged, []byte("stage|")) + } + prefix := []byte("stage|") + + require.NoError(t, ps.PutAt(ctx, stage("k"), []byte("v10"), 10, 0)) + + const entryIdx uint64 = 77 + result, err := ps.PromoteVersions(ctx, PromoteVersionsOptions{ + JobID: 9, + AppliedIndex: entryIdx, + StartKey: prefix, + EndKey: PrefixScanEnd(prefix), + MaxVersions: 10, + TargetKey: targetKey, + }) + require.NoError(t, err) + require.True(t, result.Done) + require.Equal(t, uint64(1), result.PromotedRows) + + got, present, err := ps.LastAppliedIndex() + require.NoError(t, err) + require.True(t, present, "PromoteVersions must persist metaAppliedIndex") + require.Equal(t, entryIdx, got) + + val, err := ps.GetAt(ctx, []byte("k"), 10) + require.NoError(t, err) + require.Equal(t, []byte("v10"), val) + _, err = ps.GetAt(ctx, stage("k"), 10) + require.ErrorIs(t, err, ErrKeyNotFound) + + const retryEntryIdx uint64 = 78 + retry, err := ps.PromoteVersions(ctx, PromoteVersionsOptions{ + JobID: 9, + AppliedIndex: retryEntryIdx, + StartKey: prefix, + EndKey: PrefixScanEnd(prefix), + MaxVersions: 10, + TargetKey: targetKey, + }) + require.NoError(t, err) + require.True(t, retry.Done) + require.Equal(t, uint64(1), retry.TotalPromotedRows) + + got, present, err = ps.LastAppliedIndex() + require.NoError(t, err) + require.True(t, present, "completed PromoteVersions retry must persist metaAppliedIndex") + require.Equal(t, retryEntryIdx, got) +} + +func TestImportVersionsRaft_BundlesMetaAppliedIndex(t *testing.T) { + ctx := context.Background() + st := newApplyIndexPebbleStore(t) + ps := pebbleStoreApplied(t, st) + + const entryIdx uint64 = 123 + result, err := ps.ImportVersionsRaft(ctx, ImportVersionsOptions{ + JobID: 9, + AppliedIndex: entryIdx, + BracketID: 1, + BatchSeq: 1, + Cursor: []byte("c1"), + Versions: []MVCCVersion{ + {Key: []byte("stage|k"), CommitTS: 100, Value: []byte("v100")}, + }, + }) + require.NoError(t, err) + require.Equal(t, []byte("c1"), result.AckedCursor) + require.Equal(t, uint64(100), result.MaxImportedTS) + + got, present, err := ps.LastAppliedIndex() + require.NoError(t, err) + require.True(t, present, "ImportVersionsRaft must persist metaAppliedIndex") + require.Equal(t, entryIdx, got) + + val, err := ps.GetAt(ctx, []byte("stage|k"), 100) + require.NoError(t, err) + require.Equal(t, []byte("v100"), val) + + const retryEntryIdx uint64 = 124 + duplicate, err := ps.ImportVersionsRaft(ctx, ImportVersionsOptions{ + JobID: 9, + AppliedIndex: retryEntryIdx, + BracketID: 1, + BatchSeq: 1, + Cursor: []byte("ignored"), + Versions: []MVCCVersion{ + {Key: []byte("stage|k"), CommitTS: 100, Value: []byte("changed")}, + }, + }) + require.NoError(t, err) + require.True(t, duplicate.Duplicate) + require.Equal(t, []byte("c1"), duplicate.AckedCursor) + + got, present, err = ps.LastAppliedIndex() + require.NoError(t, err) + require.True(t, present, "duplicate ImportVersionsRaft retry must still advance metaAppliedIndex") + require.Equal(t, retryEntryIdx, got) + + val, err = ps.GetAt(ctx, []byte("stage|k"), 100) + require.NoError(t, err) + require.Equal(t, []byte("v100"), val, "duplicate import must not rewrite the acknowledged batch") +} + func TestApplyMutationsRaftAt_AlreadyLandedAdvancesStaleAppliedIndex(t *testing.T) { ctx := context.Background() st := newApplyIndexPebbleStore(t) @@ -228,6 +341,38 @@ func TestDeletePrefixAtRaftAt_BundlesMetaAppliedIndex(t *testing.T) { require.Equal(t, entryIdx, got) } +func TestDeletePrefixesAtRaftAt_BundlesMetaAppliedIndex(t *testing.T) { + ctx := context.Background() + st := newApplyIndexPebbleStore(t) + ps := pebbleStoreApplied(t, st) + + const seedTS uint64 = 50 + require.NoError(t, ps.ApplyMutations(ctx, []*KVPairMutation{ + {Op: OpTypePut, Key: []byte("p/k1"), Value: []byte("v")}, + {Op: OpTypePut, Key: []byte("q/drop"), Value: []byte("v")}, + {Op: OpTypePut, Key: []byte("q/keep"), Value: []byte("v")}, + }, nil, seedTS, seedTS)) + + const entryIdx uint64 = 100 + require.NoError(t, ps.DeletePrefixesAtRaftAt(ctx, []PrefixDelete{ + {Prefix: []byte("p/")}, + {Prefix: []byte("q/"), ExcludePrefix: []byte("q/keep")}, + }, 200, entryIdx)) + + _, err := ps.GetAt(ctx, []byte("p/k1"), 250) + require.ErrorIs(t, err, ErrKeyNotFound) + _, err = ps.GetAt(ctx, []byte("q/drop"), 250) + require.ErrorIs(t, err, ErrKeyNotFound) + gotVal, err := ps.GetAt(ctx, []byte("q/keep"), 250) + require.NoError(t, err) + require.Equal(t, []byte("v"), gotVal) + + got, present, err := ps.LastAppliedIndex() + require.NoError(t, err) + require.True(t, present, "DeletePrefixesAtRaftAt must persist metaAppliedIndex") + require.Equal(t, entryIdx, got) +} + // TestSetDurableAppliedIndex_UsesPebbleSync exercises the // nosync-mode independence claim — even when ELASTICKV_FSM_SYNC_MODE // is nosync, the checkpoint must use pebble.Sync. We can't directly diff --git a/store/lsm_store_registration_gate_test.go b/store/lsm_store_registration_gate_test.go index 4dbfeb6a5..f8b6dd713 100644 --- a/store/lsm_store_registration_gate_test.go +++ b/store/lsm_store_registration_gate_test.go @@ -1,6 +1,7 @@ package store import ( + "bytes" "context" "path/filepath" "testing" @@ -96,6 +97,30 @@ func TestRegistrationGate_DirectPathFailsClosedBeforeRegistration(t *testing.T) mustGet(t, f.mvcc, []byte("a"), 250, "1") }) + t.Run("ImportVersions", func(t *testing.T) { + t.Parallel() + registered := false + f := newRegGateStore(t, ®istered) + opts := ImportVersionsOptions{ + JobID: 1, + BracketID: 1, + BatchSeq: 1, + Cursor: []byte("cursor"), + Versions: []MVCCVersion{ + {Key: []byte("import"), CommitTS: 100, Value: []byte("v")}, + }, + } + _, err := f.mvcc.ImportVersions(ctx, opts) + if !errors.Is(err, ErrWriterNotRegistered) { + t.Fatalf("ImportVersions pre-registration: got %v, want ErrWriterNotRegistered", err) + } + registered = true + if _, err := f.mvcc.ImportVersions(ctx, opts); err != nil { + t.Fatalf("ImportVersions post-registration: %v", err) + } + mustGet(t, f.mvcc, []byte("import"), 150, "v") + }) + t.Run("ExpireAt", func(t *testing.T) { t.Parallel() // ExpireAt re-encrypts the latest value, so seed a value first @@ -133,6 +158,63 @@ func TestRegistrationGate_FSMApplyPathNeverGated(t *testing.T) { t.Fatalf("ApplyMutationsRaft must not be gated, got: %v", err) } mustGet(t, f.mvcc, []byte("raft"), 150, "applied") + + _, err := f.mvcc.ImportVersionsRaft(ctx, ImportVersionsOptions{ + JobID: 2, + BracketID: 1, + BatchSeq: 1, + Cursor: []byte("cursor"), + Versions: []MVCCVersion{ + {Key: []byte("raft-import"), CommitTS: 200, Value: []byte("imported")}, + }, + }) + if err != nil { + t.Fatalf("ImportVersionsRaft must not be gated, got: %v", err) + } + mustGet(t, f.mvcc, []byte("raft-import"), 250, "imported") +} + +func TestRegistrationGate_PromoteVersionsNeverGated(t *testing.T) { + t.Parallel() + ctx := context.Background() + registered := true + f := newRegGateStore(t, ®istered) + stage := func(raw string) []byte { + return append([]byte("stage|"), []byte(raw)...) + } + targetKey := func(staged []byte) ([]byte, bool) { + return bytes.TrimPrefix(staged, []byte("stage|")), bytes.HasPrefix(staged, []byte("stage|")) + } + + if err := f.mvcc.PutAt(ctx, stage("promote"), []byte("value"), 100, 0); err != nil { + t.Fatalf("seed staged PutAt: %v", err) + } + registered = false + if err := f.mvcc.PutAt(ctx, []byte("direct"), []byte("blocked"), 110, 0); !errors.Is(err, ErrWriterNotRegistered) { + t.Fatalf("direct PutAt pre-registration: got %v, want ErrWriterNotRegistered", err) + } + promoter, ok := f.mvcc.(MigrationPromoter) + if !ok { + t.Fatalf("expected MigrationPromoter, got %T", f.mvcc) + } + + result, err := promoter.PromoteVersions(ctx, PromoteVersionsOptions{ + JobID: 11, + StartKey: []byte("stage|"), + EndKey: PrefixScanEnd([]byte("stage|")), + MaxVersions: 10, + TargetKey: targetKey, + }) + if err != nil { + t.Fatalf("PromoteVersions pre-registration: %v", err) + } + if !result.Done || result.PromotedRows != 1 { + t.Fatalf("PromoteVersions result = %+v, want done with one promoted row", result) + } + mustGet(t, f.mvcc, []byte("promote"), 150, "value") + if _, err := f.mvcc.GetAt(ctx, stage("promote"), 150); !errors.Is(err, ErrKeyNotFound) { + t.Fatalf("staged version after promotion: got %v, want ErrKeyNotFound", err) + } } // TestRegistrationGate_NotEncryptingIsUngated confirms the gate is only diff --git a/store/lsm_store_sync_mode_test.go b/store/lsm_store_sync_mode_test.go index e5e065bfd..29349b997 100644 --- a/store/lsm_store_sync_mode_test.go +++ b/store/lsm_store_sync_mode_test.go @@ -1,6 +1,7 @@ package store import ( + "bytes" "context" "testing" @@ -116,6 +117,65 @@ func TestDirectApplyWriteOpts_AlwaysSync(t *testing.T) { }) } +func TestPromoteVersionsWriteOptsFollowApplyContext(t *testing.T) { + t.Run("raft-applied promotion observes nosync", func(t *testing.T) { + ps := newPebbleStoreWithFSMApplyWriteOptsForTest(t, t.TempDir(), pebble.NoSync, fsmSyncModeNoSync) + defer ps.Close() + + require.Same(t, pebble.NoSync, ps.promotionWriteOpts(123), + "promotion with an applied index must use raft apply write options") + }) + + t.Run("direct promotion stays sync", func(t *testing.T) { + ps := newPebbleStoreWithFSMApplyWriteOptsForTest(t, t.TempDir(), pebble.NoSync, fsmSyncModeNoSync) + defer ps.Close() + + require.Same(t, pebble.Sync, ps.promotionWriteOpts(0), + "promotion without an applied index has no raft durability backstop") + }) +} + +func TestPromoteVersionsRaftNoSyncFunctionalEquivalence(t *testing.T) { + dir := t.TempDir() + ps := newPebbleStoreWithFSMApplyWriteOptsForTest(t, dir, pebble.NoSync, fsmSyncModeNoSync) + defer ps.Close() + + ctx := context.Background() + stage := func(raw string) []byte { + return append([]byte("stage|"), []byte(raw)...) + } + targetKey := func(staged []byte) ([]byte, bool) { + return bytes.TrimPrefix(staged, []byte("stage|")), bytes.HasPrefix(staged, []byte("stage|")) + } + prefix := []byte("stage|") + + require.NoError(t, ps.PutAt(ctx, stage("k"), []byte("v10"), 10, 0)) + + const entryIdx uint64 = 88 + result, err := ps.PromoteVersions(ctx, PromoteVersionsOptions{ + JobID: 12, + AppliedIndex: entryIdx, + StartKey: prefix, + EndKey: PrefixScanEnd(prefix), + MaxVersions: 10, + TargetKey: targetKey, + }) + require.NoError(t, err) + require.True(t, result.Done) + require.Equal(t, uint64(1), result.PromotedRows) + + val, err := ps.GetAt(ctx, []byte("k"), 10) + require.NoError(t, err) + require.Equal(t, []byte("v10"), val) + _, err = ps.GetAt(ctx, stage("k"), 10) + require.ErrorIs(t, err, ErrKeyNotFound) + + got, present, err := ps.LastAppliedIndex() + require.NoError(t, err) + require.True(t, present) + require.Equal(t, entryIdx, got) +} + // TestDirectApplyMutations_NoSyncConfigured_StillWritesDurably is the // functional twin of TestDirectApplyWriteOpts_AlwaysSync: it exercises // the public ApplyMutations and DeletePrefixAt entry points with a diff --git a/store/lsm_store_test.go b/store/lsm_store_test.go index dbd6605e1..fb4c9a3ef 100644 --- a/store/lsm_store_test.go +++ b/store/lsm_store_test.go @@ -1212,3 +1212,39 @@ func TestPebbleStore_SnapshotRestore_MaxSizeKey(t *testing.T) { require.NoError(t, err) assert.Equal(t, []byte("val"), got) } + +func TestPebbleStore_SnapshotRestore_MaxStoredKey(t *testing.T) { + dir, err := os.MkdirTemp("", "pebble-maxstoredkey-snap-*") + require.NoError(t, err) + defer os.RemoveAll(dir) + + s, err := NewPebbleStore(dir) + require.NoError(t, err) + defer func() { assert.NoError(t, s.Close()) }() + + ctx := context.Background() + bigKey := bytes.Repeat([]byte("k"), maxSnapshotStoredKeySize) + require.NoError(t, s.PutAt(ctx, bigKey, []byte("val"), 1, 0)) + + snap, err := s.Snapshot() + require.NoError(t, err) + defer func() { assert.NoError(t, snap.Close()) }() + + var buf bytes.Buffer + _, err = snap.WriteTo(&buf) + require.NoError(t, err) + + dir2, err := os.MkdirTemp("", "pebble-maxstoredkey-restore-*") + require.NoError(t, err) + defer os.RemoveAll(dir2) + + s2, err := NewPebbleStore(dir2) + require.NoError(t, err) + defer func() { assert.NoError(t, s2.Close()) }() + + require.NoError(t, s2.Restore(bytes.NewReader(buf.Bytes()))) + + got, err := s2.GetAt(ctx, bigKey, 1) + require.NoError(t, err) + assert.Equal(t, []byte("val"), got) +} diff --git a/store/migration_promote.go b/store/migration_promote.go new file mode 100644 index 000000000..d8a8bc27b --- /dev/null +++ b/store/migration_promote.go @@ -0,0 +1,496 @@ +package store + +import ( + "bytes" + "context" + "encoding/binary" + + "github.com/cockroachdb/errors" + "github.com/cockroachdb/pebble/v2" +) + +const ( + migrationPromotionDoneFlag byte = 1 + migrationPromotionStateVersion2Flag byte = 1 << 7 +) + +type promotedVersion struct { + staged MVCCVersion + target MVCCVersion +} + +func validatePromoteVersionsOptions(opts PromoteVersionsOptions) error { + if opts.TargetKey == nil { + return errors.New("migration promote target key mapper is required") + } + return nil +} + +func promotedVersionsFromStaged(opts PromoteVersionsOptions, versions []MVCCVersion) ([]promotedVersion, PromoteVersionsResult, error) { + out := make([]promotedVersion, 0, len(versions)) + result := PromoteVersionsResult{PromotedRows: uint64(len(versions))} //nolint:gosec // len is bounded by MaxVersions. + for _, staged := range versions { + targetKey, ok := opts.TargetKey(staged.Key) + if !ok { + return nil, PromoteVersionsResult{}, errors.WithStack(errors.Newf("migration promote target key rejected staged key %q", string(staged.Key))) + } + target := MVCCVersion{ + Key: bytes.Clone(targetKey), + CommitTS: staged.CommitTS, + Tombstone: staged.Tombstone, + Value: bytes.Clone(staged.Value), + KeyFamily: staged.KeyFamily, + ExpireAt: staged.ExpireAt, + } + if err := validateImportVersion(target); err != nil { + return nil, PromoteVersionsResult{}, err + } + result.PromotedBytes += versionExportSize(target.Key, len(target.Value)) + if target.CommitTS > result.MaxPromotedTS { + result.MaxPromotedTS = target.CommitTS + } + out = append(out, promotedVersion{ + staged: staged, + target: target, + }) + } + return out, result, nil +} + +func (s *mvccStore) PromoteVersions(ctx context.Context, opts PromoteVersionsOptions) (PromoteVersionsResult, error) { + if err := validatePromoteVersionsOptions(opts); err != nil { + return PromoteVersionsResult{}, err + } + if opts.MaxVersions <= 0 { + return PromoteVersionsResult{Done: true}, nil + } + + s.mtx.Lock() + defer s.mtx.Unlock() + + state, cursor := s.promotionStateAndCursorLocked(opts) + if opts.JobID != 0 && state.Done { + return PromoteVersionsResult{Done: true, TotalPromotedRows: state.PromotedRows, MaxPromotedTS: state.MaxPromotedTS}, nil + } + exported, toPromote, promoted, err := s.planMemoryPromotionLocked(ctx, opts, cursor) + if err != nil { + return PromoteVersionsResult{}, err + } + s.applyMemoryPromotionLocked(toPromote) + result, updatedState := finishPromotionResult(opts, state, exported, promoted) + if result.MaxPromotedTS > s.lastCommitTS { + s.lastCommitTS = result.MaxPromotedTS + } + if updatedState != nil { + s.migrationPromotions[opts.JobID] = clonePromotionState(*updatedState) + } + return result, nil +} + +func (s *mvccStore) planMemoryPromotionLocked( + ctx context.Context, + opts PromoteVersionsOptions, + cursor []byte, +) (ExportVersionsResult, []promotedVersion, PromoteVersionsResult, error) { + exportOpts := normalizeExportVersionsOptions(ExportVersionsOptions{ + StartKey: opts.StartKey, + EndKey: opts.EndKey, + Cursor: cursor, + MaxVersions: opts.MaxVersions, + MaxBytes: opts.MaxBytes, + MaxScannedBytes: opts.MaxScannedBytes, + }) + pos, err := decodeExportCursorForOptions(exportOpts) + if err != nil { + return ExportVersionsResult{}, nil, PromoteVersionsResult{}, err + } + exported, err := s.exportMemoryVersionsLocked(ctx, exportOpts, pos) + if err != nil { + return ExportVersionsResult{}, nil, PromoteVersionsResult{}, err + } + toPromote, promoted, err := promotedVersionsFromStaged(opts, exported.Versions) + return exported, toPromote, promoted, err +} + +func (s *mvccStore) applyMemoryPromotionLocked(toPromote []promotedVersion) { + for _, version := range toPromote { + if version.target.Tombstone { + s.deleteVersionLocked(version.target.Key, version.target.CommitTS) + } else { + s.putVersionLocked(version.target.Key, version.target.Value, version.target.CommitTS, version.target.ExpireAt) + } + s.removeVersionLocked(version.staged.Key, version.staged.CommitTS) + } +} + +func (s *mvccStore) promotionStateAndCursorLocked(opts PromoteVersionsOptions) (PromotionState, []byte) { + if opts.JobID == 0 { + return PromotionState{}, opts.Cursor + } + state, ok := s.migrationPromotions[opts.JobID] + if !ok { + return PromotionState{}, nil + } + state = clonePromotionState(state) + return state, state.Cursor +} + +func (s *mvccStore) MigrationPromotionState(_ context.Context, jobID uint64) (PromotionState, bool, error) { + s.mtx.RLock() + defer s.mtx.RUnlock() + state, ok := s.migrationPromotions[jobID] + return clonePromotionState(state), ok, nil +} + +func (s *mvccStore) exportMemoryVersionsLocked(ctx context.Context, opts ExportVersionsOptions, pos exportCursorPosition) (ExportVersionsResult, error) { + result := newExportVersionsResult(opts.MaxVersions) + it := s.tree.Iterator() + if !s.seekMemoryExportStart(&it, opts.StartKey, pos) { + result.Done = true + return result, nil + } + for ok := true; ok; ok = it.Next() { + key, keyOK := it.Key().([]byte) + if err := checkExportKey(ctx, key, keyOK, opts.EndKey); err != nil { + if errors.Is(err, errExportReachedEnd) { + result.Done = true + result.NextCursor = nil + return result, nil + } + return ExportVersionsResult{}, err + } + if !keyOK { + continue + } + done, err := exportMemoryIteratorKey(ctx, opts, pos, key, it.Value(), &result) + if err != nil || !done { + return result, err + } + } + result.Done = true + result.NextCursor = nil + return result, nil +} + +func (s *mvccStore) removeVersionLocked(key []byte, commitTS uint64) bool { + existing, ok := s.tree.Get(key) + if !ok { + return false + } + versions, _ := existing.([]VersionedValue) + idx := findVersionIndex(versions, commitTS) + if idx < 0 { + return false + } + next := make([]VersionedValue, len(versions)-1) + copy(next, versions[:idx]) + copy(next[idx:], versions[idx+1:]) + if len(next) == 0 { + s.tree.Remove(key) + return true + } + s.tree.Put(bytes.Clone(key), next) + return true +} + +func findVersionIndex(versions []VersionedValue, commitTS uint64) int { + for i := range versions { + if versions[i].TS == commitTS { + return i + } + } + return -1 +} + +func (s *pebbleStore) PromoteVersions(ctx context.Context, opts PromoteVersionsOptions) (PromoteVersionsResult, error) { + if err := validatePromoteVersionsOptions(opts); err != nil { + return PromoteVersionsResult{}, err + } + if opts.MaxVersions <= 0 { + return PromoteVersionsResult{Done: true}, nil + } + + s.dbMu.RLock() + defer s.dbMu.RUnlock() + + s.applyMu.Lock() + defer s.applyMu.Unlock() + + state, cursor, err := s.pebblePromotionStateAndCursor(opts) + if err != nil { + return PromoteVersionsResult{}, err + } + writeOpts := s.promotionWriteOpts(opts.AppliedIndex) + if opts.JobID != 0 && state.Done { + result := PromoteVersionsResult{Done: true, TotalPromotedRows: state.PromotedRows, MaxPromotedTS: state.MaxPromotedTS} + return s.finishPebblePromotion(nil, opts.JobID, nil, result, opts.AppliedIndex, state.MaxPromotedTS, writeOpts) + } + opts.Cursor = cursor + exported, toPromote, promoted, err := s.planPebblePromotionLocked(ctx, opts) + if err != nil { + return PromoteVersionsResult{}, err + } + result, stateToWrite := finishPromotionResult(opts, state, exported, promoted) + return s.finishPebblePromotion( + toPromote, + opts.JobID, + stateToWrite, + result, + opts.AppliedIndex, + result.MaxPromotedTS, + writeOpts, + ) +} + +func (s *pebbleStore) promotionWriteOpts(appliedIndex uint64) *pebble.WriteOptions { + if appliedIndex > 0 { + return s.raftApplyWriteOpts() + } + return s.directApplyWriteOpts() +} + +func (s *pebbleStore) finishPebblePromotion( + toPromote []promotedVersion, + jobID uint64, + stateToWrite *PromotionState, + result PromoteVersionsResult, + appliedIndex uint64, + maxPromotedTS uint64, + writeOpts *pebble.WriteOptions, +) (PromoteVersionsResult, error) { + if len(toPromote) == 0 && stateToWrite == nil && appliedIndex == 0 && maxPromotedTS == 0 { + return result, nil + } + if err := s.commitPebblePromoteVersions( + toPromote, + jobID, + stateToWrite, + appliedIndex, + maxPromotedTS, + writeOpts, + ); err != nil { + return PromoteVersionsResult{}, err + } + return result, nil +} + +func (s *pebbleStore) planPebblePromotionLocked( + ctx context.Context, + opts PromoteVersionsOptions, +) (ExportVersionsResult, []promotedVersion, PromoteVersionsResult, error) { + exported, err := s.exportVersionsLocked(ctx, ExportVersionsOptions{ + StartKey: opts.StartKey, + EndKey: opts.EndKey, + Cursor: opts.Cursor, + MaxVersions: opts.MaxVersions, + MaxBytes: opts.MaxBytes, + MaxScannedBytes: opts.MaxScannedBytes, + }) + if err != nil { + return ExportVersionsResult{}, nil, PromoteVersionsResult{}, err + } + toPromote, promoted, err := promotedVersionsFromStaged(opts, exported.Versions) + return exported, toPromote, promoted, err +} + +func finishPromotionResult( + opts PromoteVersionsOptions, + state PromotionState, + exported ExportVersionsResult, + promoted PromoteVersionsResult, +) (PromoteVersionsResult, *PromotionState) { + promoted.NextCursor = exported.NextCursor + promoted.Done = exported.Done + promoted.ScannedBytes = exported.ScannedBytes + promoted.TotalPromotedRows = promoted.PromotedRows + if opts.JobID == 0 { + return promoted, nil + } + state.Cursor = bytes.Clone(exported.NextCursor) + state.Done = exported.Done + state.PromotedRows += promoted.PromotedRows + if promoted.MaxPromotedTS > state.MaxPromotedTS { + state.MaxPromotedTS = promoted.MaxPromotedTS + } + state.LastError = "" + promoted.TotalPromotedRows = state.PromotedRows + promoted.MaxPromotedTS = state.MaxPromotedTS + return promoted, &state +} + +func (s *pebbleStore) pebblePromotionStateAndCursor(opts PromoteVersionsOptions) (PromotionState, []byte, error) { + if opts.JobID == 0 { + return PromotionState{}, opts.Cursor, nil + } + state, ok, err := s.readPebblePromotionState(opts.JobID) + if err != nil { + return PromotionState{}, nil, err + } + if !ok { + return PromotionState{}, nil, nil + } + return state, state.Cursor, nil +} + +func (s *pebbleStore) MigrationPromotionState(_ context.Context, jobID uint64) (PromotionState, bool, error) { + s.dbMu.RLock() + defer s.dbMu.RUnlock() + return s.readPebblePromotionState(jobID) +} + +func (s *pebbleStore) readPebblePromotionState(jobID uint64) (PromotionState, bool, error) { + states, err := s.readPebblePromotionStates() + if err != nil { + return PromotionState{}, false, err + } + state, ok := states[jobID] + return clonePromotionState(state), ok, nil +} + +func (s *pebbleStore) readPebblePromotionStates() (map[uint64]PromotionState, error) { + val, closer, err := s.db.Get(migrationPromoteMetaKeyBytes) + if err != nil { + if errors.Is(err, pebble.ErrNotFound) { + return make(map[uint64]PromotionState), nil + } + return nil, errors.WithStack(err) + } + defer func() { _ = closer.Close() }() + states, ok := decodeMigrationPromotionStates(val) + if !ok { + return nil, errors.New("corrupt migration promotion state metadata") + } + return states, nil +} + +func (s *pebbleStore) stagePebblePromotionState(batch *pebble.Batch, jobID uint64, state PromotionState) error { + states, err := s.readPebblePromotionStates() + if err != nil { + return err + } + states[jobID] = clonePromotionState(state) + return errors.WithStack(batch.Set(migrationPromoteMetaKeyBytes, encodeMigrationPromotionStates(states), nil)) +} + +func (s *pebbleStore) commitPebblePromoteVersions( + versions []promotedVersion, + jobID uint64, + state *PromotionState, + appliedIndex uint64, + maxPromotedTS uint64, + writeOpts *pebble.WriteOptions, +) error { + batch := s.db.NewBatch() + defer batch.Close() + targets := make([]MVCCVersion, 0, len(versions)) + for _, version := range versions { + targets = append(targets, version.target) + } + // Promotion is replayed from the Raft FSM, so it must not fail closed on + // this node's local writer-registration state. + if err := s.applyImportVersionsBatch(batch, targets, false); err != nil { + return err + } + for _, version := range versions { + if err := batch.Delete(encodeKey(version.staged.Key, version.staged.CommitTS), nil); err != nil { + return errors.WithStack(err) + } + } + if state != nil { + if err := s.stagePebblePromotionState(batch, jobID, *state); err != nil { + return err + } + } + if appliedIndex > 0 { + if err := setPebbleUint64InBatch(batch, metaAppliedIndexBytes, appliedIndex); err != nil { + return err + } + } + return s.commitPebblePromotionBatch(batch, maxPromotedTS, writeOpts) +} + +func (s *pebbleStore) commitPebblePromotionBatch( + batch *pebble.Batch, + maxPromotedTS uint64, + writeOpts *pebble.WriteOptions, +) error { + if maxPromotedTS > 0 { + s.mtx.Lock() + defer s.mtx.Unlock() + newLastTS := s.lastCommitTS + if maxPromotedTS > newLastTS { + newLastTS = maxPromotedTS + } + if err := setPebbleUint64InBatch(batch, metaLastCommitTSBytes, newLastTS); err != nil { + return err + } + if err := batch.Commit(writeOpts); err != nil { + return errors.WithStack(err) + } + s.updateLastCommitTS(newLastTS) + return nil + } + if err := batch.Commit(writeOpts); err != nil { + return errors.WithStack(err) + } + return nil +} + +func encodePromotionState(state PromotionState) []byte { + buf := make([]byte, 0, 1+2*migrationUint64Bytes+binary.MaxVarintLen64*2+len(state.Cursor)+len(state.LastError)) + flags := migrationPromotionStateVersion2Flag + if state.Done { + flags |= migrationPromotionDoneFlag + } + buf = append(buf, flags) + buf = binary.BigEndian.AppendUint64(buf, state.PromotedRows) + buf = binary.BigEndian.AppendUint64(buf, state.MaxPromotedTS) + buf = binary.AppendUvarint(buf, lenAsUint64(len(state.Cursor))) + buf = append(buf, state.Cursor...) + buf = binary.AppendUvarint(buf, lenAsUint64(len(state.LastError))) + buf = append(buf, state.LastError...) + return buf +} + +func decodePromotionState(data []byte) (PromotionState, bool) { + if len(data) < 1+migrationUint64Bytes { + return PromotionState{}, false + } + flags := data[0] + state := PromotionState{ + Done: flags&migrationPromotionDoneFlag != 0, + PromotedRows: binary.BigEndian.Uint64(data[1 : 1+migrationUint64Bytes]), + } + rest := data[1+migrationUint64Bytes:] + if flags&migrationPromotionStateVersion2Flag != 0 { + if len(rest) < migrationUint64Bytes { + return PromotionState{}, false + } + state.MaxPromotedTS = binary.BigEndian.Uint64(rest[:migrationUint64Bytes]) + rest = rest[migrationUint64Bytes:] + } + cursorLen, n := binary.Uvarint(rest) + if n <= 0 || cursorLen > lenAsUint64(len(rest[n:])) { + return PromotionState{}, false + } + rest = rest[n:] + cursorEnd := int(cursorLen) //nolint:gosec // bounded by len(rest) above. + state.Cursor = bytes.Clone(rest[:cursorEnd]) + rest = rest[cursorEnd:] + errLen, n := binary.Uvarint(rest) + if n <= 0 || errLen != lenAsUint64(len(rest[n:])) { + return PromotionState{}, false + } + state.LastError = string(rest[n:]) + return state, true +} + +func clonePromotionState(state PromotionState) PromotionState { + state.Cursor = bytes.Clone(state.Cursor) + return state +} + +var _ MigrationPromoter = (*mvccStore)(nil) +var _ MigrationPromoter = (*pebbleStore)(nil) +var _ MigrationPromotionStateReader = (*mvccStore)(nil) +var _ MigrationPromotionStateReader = (*pebbleStore)(nil) diff --git a/store/migration_versions.go b/store/migration_versions.go index 82d885ffb..e68eb273f 100644 --- a/store/migration_versions.go +++ b/store/migration_versions.go @@ -18,6 +18,7 @@ const ( migrationAckMetaKey = "_migack" migrationHLCFloorMetaKey = "_mighlc" + migrationPromoteMetaKey = "_migpromote" migrationMetadataVersion = 1 migrationAckPrefix = "!migstage|ack|" @@ -29,6 +30,7 @@ const ( var ( migrationAckMetaKeyBytes = []byte(migrationAckMetaKey) migrationHLCFloorMetaKeyBytes = []byte(migrationHLCFloorMetaKey) + migrationPromoteMetaKeyBytes = []byte(migrationPromoteMetaKey) ) type exportCursorPosition struct { @@ -89,6 +91,54 @@ func decodeExportCursor(cursor []byte) (exportCursorPosition, error) { return exportCursorPosition{key: key, commitTS: commitTS, tag: tag, hasKey: true}, nil } +// ValidateExportCursorForRange verifies that an export cursor decodes and +// resumes inside the supplied key interval. Skipped-key cursors are accepted +// only when they describe a key outside the interval. +func ValidateExportCursorForRange(cursor, startKey, endKey []byte) error { + pos, err := decodeExportCursor(cursor) + if err != nil { + return err + } + return validateExportCursorPositionForRange(pos, startKey, endKey) +} + +// ValidatePromotionCursorForRange verifies a promotion cursor before it is +// proposed to Raft. Promotion scans emit only accepted positions, so callers +// must not resume from sparse-scan-only cursor tags. +func ValidatePromotionCursorForRange(cursor, startKey, endKey []byte) error { + pos, err := decodeExportCursor(cursor) + if err != nil { + return err + } + if !pos.hasKey { + return nil + } + if pos.tag != exportCursorTagEmitted { + return errors.WithStack(ErrInvalidExportCursor) + } + return validateExportCursorPositionForRange(pos, startKey, endKey) +} + +func validateExportCursorPositionForRange(pos exportCursorPosition, startKey, endKey []byte) error { + if !pos.hasKey { + return nil + } + if pos.tag == exportCursorTagSkippedKey { + opts := ExportVersionsOptions{StartKey: startKey, EndKey: endKey} + if !exportSkippedCursorOutsideRange(opts, pos.key) { + return errors.WithStack(ErrInvalidExportCursor) + } + return nil + } + if startKey != nil && bytes.Compare(pos.key, startKey) < 0 { + return errors.WithStack(ErrInvalidExportCursor) + } + if endKey != nil && bytes.Compare(pos.key, endKey) >= 0 { + return errors.WithStack(ErrInvalidExportCursor) + } + return nil +} + func decodeExportCursorForOptions(opts ExportVersionsOptions) (exportCursorPosition, error) { pos, err := decodeExportCursor(opts.Cursor) if err != nil { @@ -97,7 +147,7 @@ func decodeExportCursorForOptions(opts ExportVersionsOptions) (exportCursorPosit if err := validateExportCursorRange(opts, pos); err != nil { return exportCursorPosition{}, err } - return pos, nil + return normalizeExportCursorPositionForRange(opts, pos), nil } func validateExportCursorRange(opts ExportVersionsOptions, pos exportCursorPosition) error { @@ -124,6 +174,16 @@ func exportSkippedCursorOutsideRange(opts ExportVersionsOptions, key []byte) boo (opts.EndKey != nil && bytes.Compare(key, opts.EndKey) >= 0) } +func normalizeExportCursorPositionForRange(opts ExportVersionsOptions, pos exportCursorPosition) exportCursorPosition { + if !pos.hasKey || pos.tag != exportCursorTagSkippedKey || opts.StartKey == nil { + return pos + } + if bytes.Compare(pos.key, opts.StartKey) >= 0 { + return pos + } + return exportCursorPosition{} +} + func normalizeExportVersionsOptions(opts ExportVersionsOptions) ExportVersionsOptions { if opts.EndKey != nil && len(opts.EndKey) == 0 { opts.EndKey = nil @@ -145,7 +205,8 @@ func exportUsesSparseScanBudget(opts ExportVersionsOptions) bool { func isMigrationMetadataKey(rawKey []byte) bool { return bytes.Equal(rawKey, migrationAckMetaKeyBytes) || - bytes.Equal(rawKey, migrationHLCFloorMetaKeyBytes) + bytes.Equal(rawKey, migrationHLCFloorMetaKeyBytes) || + bytes.Equal(rawKey, migrationPromoteMetaKeyBytes) } func encodeMigrationImportAcks(acks map[migrationAckID]migrationImportAck) []byte { @@ -250,16 +311,71 @@ func decodeMigrationHLCFloors(data []byte) (map[uint64]uint64, bool) { return floors, len(rest) == 0 } +func encodeMigrationPromotionStates(states map[uint64]PromotionState) []byte { + jobIDs := make([]uint64, 0, len(states)) + for jobID := range states { + jobIDs = append(jobIDs, jobID) + } + sort.Slice(jobIDs, func(i, j int) bool { return jobIDs[i] < jobIDs[j] }) + + buf := make([]byte, 0, 1+binary.MaxVarintLen64+len(jobIDs)*(migrationUint64Bytes+binary.MaxVarintLen64)) + buf = append(buf, migrationMetadataVersion) + buf = binary.AppendUvarint(buf, lenAsUint64(len(jobIDs))) + for _, jobID := range jobIDs { + encoded := encodePromotionState(states[jobID]) + buf = binary.BigEndian.AppendUint64(buf, jobID) + buf = binary.AppendUvarint(buf, lenAsUint64(len(encoded))) + buf = append(buf, encoded...) + } + return buf +} + +func decodeMigrationPromotionStates(data []byte) (map[uint64]PromotionState, bool) { + if len(data) == 0 || data[0] != migrationMetadataVersion { + return nil, false + } + rest := data[1:] + count, n := binary.Uvarint(rest) + if n <= 0 { + return nil, false + } + rest = rest[n:] + states := make(map[uint64]PromotionState) + for i := uint64(0); i < count; i++ { + if len(rest) < migrationUint64Bytes { + return nil, false + } + jobID := binary.BigEndian.Uint64(rest[:migrationUint64Bytes]) + rest = rest[migrationUint64Bytes:] + stateLen, n := binary.Uvarint(rest) + if n <= 0 { + return nil, false + } + rest = rest[n:] + if stateLen > lenAsUint64(len(rest)) { + return nil, false + } + stateEnd := int(stateLen) //nolint:gosec // bounded by len(rest) above. + state, ok := decodePromotionState(rest[:stateEnd]) + if !ok { + return nil, false + } + states[jobID] = state + rest = rest[stateEnd:] + } + return states, len(rest) == 0 +} + func validateImportVersion(version MVCCVersion) error { if version.CommitTS == 0 { - return errors.New("migration import version has zero commit_ts") + return errors.Wrap(ErrInvalidImportVersion, "migration import version has zero commit_ts") } if version.Tombstone { if version.ExpireAt != 0 { - return errors.New("migration import tombstone carries expire_at") + return errors.Wrap(ErrInvalidImportVersion, "migration import tombstone carries expire_at") } if len(version.Value) != 0 { - return errors.New("migration import tombstone carries value") + return errors.Wrap(ErrInvalidImportVersion, "migration import tombstone carries value") } return nil } @@ -394,24 +510,53 @@ func exportMemoryIteratorKey( return exportMemoryVersionsForKey(ctx, opts, cursorCommitTS, key, versions, result) } +// exportPageWouldOverflow reports whether adding one more row of valueLen +// bytes would push a page that already holds rows past its byte budget. +// MaxBytes is otherwise only consulted after a row has been appended, so a +// page can overshoot it by the whole size of its last row. That overshoot is +// what pushes an otherwise ordinary page past the migration transport limit -- +// a 3 MiB page followed by a 61 MiB row -- and the same cursor rebuilds the +// same page on every retry. Stopping first leaves the oversized row to start +// the next page, where it is alone and within the transport limit; a row too +// large even alone still goes through, because a page that holds nothing yet +// has to make progress. +func exportPageWouldOverflow(opts ExportVersionsOptions, result *ExportVersionsResult, key []byte, valueLen int) bool { + if opts.MaxBytes == 0 || len(result.Versions) == 0 { + return false + } + return result.ExportedBytes+versionExportSize(key, valueLen) > opts.MaxBytes +} + func finishExportIfLimited(opts ExportVersionsOptions, result *ExportVersionsResult) bool { return len(result.Versions) >= opts.MaxVersions || (opts.MaxBytes > 0 && result.ExportedBytes >= opts.MaxBytes) || (opts.MaxScannedBytes > 0 && result.ScannedBytes >= opts.MaxScannedBytes) } -func appendMemoryExportVersion(opts ExportVersionsOptions, key []byte, version VersionedValue, result *ExportVersionsResult) byte { +// appendMemoryExportVersion returns the cursor tag for this version and, as +// deferred, whether the page had to stop before it. AcceptVersion is called at +// most once per version here: it is not required to be a pure predicate, so +// the page-size guard sits after the filters rather than re-running them. +func appendMemoryExportVersion( + opts ExportVersionsOptions, + key []byte, + version VersionedValue, + result *ExportVersionsResult, +) (byte, bool) { if shouldSkipMigrationExportKey(key) { - return exportCursorTagScanned + return exportCursorTagScanned, false } if opts.AcceptKey != nil && !opts.AcceptKey(key) { - return exportCursorTagScanned + return exportCursorTagScanned, false + } + if opts.MaxCommitTSInclusive != 0 && version.TS > opts.MaxCommitTSInclusive { + return exportCursorTagScanned, false } if opts.AcceptVersion != nil && !opts.AcceptVersion(key, version.Value) { - return exportCursorTagScanned + return exportCursorTagScanned, false } - if opts.MaxCommitTSInclusive != 0 && version.TS > opts.MaxCommitTSInclusive { - return exportCursorTagScanned + if exportPageWouldOverflow(opts, result, key, len(version.Value)) { + return exportCursorTagScanned, true } result.Versions = append(result.Versions, MVCCVersion{ Key: bytes.Clone(key), @@ -423,7 +568,7 @@ func appendMemoryExportVersion(opts ExportVersionsOptions, key []byte, version V }) result.ExportedBytes += versionExportSize(key, len(version.Value)) result.AcceptedRows++ - return exportCursorTagEmitted + return exportCursorTagEmitted, false } func shouldSkipMigrationExportKey(key []byte) bool { @@ -444,12 +589,23 @@ func shouldSkipMemoryVersion(cursorCommitTS uint64, version VersionedValue) bool return cursorCommitTS != 0 && version.TS >= cursorCommitTS } -func exportMemoryVersion(opts ExportVersionsOptions, cursorCommitTS uint64, key []byte, version VersionedValue, result *ExportVersionsResult) bool { +// exportMemoryVersion returns (continue, deferred): deferred means the page is +// full for this version and it must start the next one. +func exportMemoryVersion( + opts ExportVersionsOptions, + cursorCommitTS uint64, + key []byte, + version VersionedValue, + result *ExportVersionsResult, +) (bool, bool) { if shouldSkipMemoryVersion(cursorCommitTS, version) { - return true + return true, false } - tag := appendMemoryExportVersion(opts, key, version, result) - return finishMemoryExportPosition(opts, key, version, tag, result) + tag, deferred := appendMemoryExportVersion(opts, key, version, result) + if deferred { + return false, true + } + return finishMemoryExportPosition(opts, key, version, tag, result), false } func exportMemoryVersionsForKey( @@ -473,7 +629,14 @@ func exportMemoryVersionsForKey( } return true, nil } - if !exportMemoryVersion(opts, cursorCommitTS, key, versions[i], result) { + cont, deferred := exportMemoryVersion(opts, cursorCommitTS, key, versions[i], result) + if deferred { + // The page is full for this row. The cursor still points at the row + // before it, so the next page starts here and carries it alone. + result.Done = false + return false, nil + } + if !cont { return !finishExportIfLimited(opts, result), nil } if !result.Done && finishExportIfLimited(opts, result) { @@ -521,6 +684,10 @@ func (s *mvccStore) ImportVersions(_ context.Context, opts ImportVersionsOptions return ImportVersionsResult{AckedCursor: bytes.Clone(opts.Cursor), MaxImportedTS: batchMax}, nil } +func (s *mvccStore) ImportVersionsRaft(ctx context.Context, opts ImportVersionsOptions) (ImportVersionsResult, error) { + return s.ImportVersions(ctx, opts) +} + func (s *mvccStore) MigrationHLCFloor(_ context.Context, jobID uint64) (uint64, error) { s.mtx.RLock() defer s.mtx.RUnlock() @@ -539,5 +706,6 @@ func (s *mvccStore) RetireMigration(ctx context.Context, jobID uint64) error { } } delete(s.migrationHLCFloors, jobID) + delete(s.migrationPromotions, jobID) return nil } diff --git a/store/migration_versions_test.go b/store/migration_versions_test.go index 1dee297f3..59ca5a206 100644 --- a/store/migration_versions_test.go +++ b/store/migration_versions_test.go @@ -9,6 +9,7 @@ import ( "testing" "github.com/bootjp/elastickv/internal/encryption" + "github.com/cockroachdb/pebble/v2" "github.com/stretchr/testify/require" ) @@ -834,6 +835,9 @@ func TestImportVersionsIdempotencyAndMetadata(t *testing.T) { func TestRetireMigrationRemovesOnlySelectedJobMetadata(t *testing.T) { runMigrationStoreSuite(t, func(t *testing.T, st MVCCStore) { ctx := context.Background() + seedPromotionState(t, ctx, st, 1, []byte("stage|job1|"), []byte("job1-promoted")) + seedPromotionState(t, ctx, st, 2, []byte("stage|job2|"), []byte("job2-promoted")) + _, err := st.ImportVersions(ctx, ImportVersionsOptions{ JobID: 1, BracketID: 2, @@ -874,6 +878,14 @@ func TestRetireMigrationRemovesOnlySelectedJobMetadata(t *testing.T) { floor, err = st.MigrationHLCFloor(ctx, 2) require.NoError(t, err) require.Equal(t, uint64(30), floor) + stateReader, ok := st.(MigrationPromotionStateReader) + require.True(t, ok) + _, ok = migrationPromotionState(t, ctx, stateReader, 1) + require.False(t, ok) + state, ok := migrationPromotionState(t, ctx, stateReader, 2) + require.True(t, ok) + require.True(t, state.Done) + require.Equal(t, uint64(1), state.PromotedRows) res, err := st.ImportVersions(ctx, ImportVersionsOptions{ JobID: 1, @@ -905,6 +917,31 @@ func TestRetireMigrationRemovesOnlySelectedJobMetadata(t *testing.T) { }) } +func seedPromotionState(t *testing.T, ctx context.Context, st MVCCStore, jobID uint64, prefix []byte, target []byte) { + t.Helper() + promoter, ok := st.(MigrationPromoter) + require.True(t, ok) + require.NoError(t, st.PutAt(ctx, append(bytes.Clone(prefix), 'k'), []byte("v"), 10+jobID, 0)) + result, err := promoter.PromoteVersions(ctx, PromoteVersionsOptions{ + JobID: jobID, + StartKey: prefix, + EndKey: PrefixScanEnd(prefix), + MaxVersions: 10, + TargetKey: func(staged []byte) ([]byte, bool) { + return target, bytes.HasPrefix(staged, prefix) + }, + }) + require.NoError(t, err) + require.True(t, result.Done) +} + +func migrationPromotionState(t *testing.T, ctx context.Context, reader MigrationPromotionStateReader, jobID uint64) (PromotionState, bool) { + t.Helper() + state, ok, err := reader.MigrationPromotionState(ctx, jobID) + require.NoError(t, err) + return state, ok +} + func TestPebbleImportMetadataPersistsAcrossReopen(t *testing.T) { ctx := context.Background() dir, err := os.MkdirTemp("", "migration-import-persist-*") @@ -964,6 +1001,8 @@ func TestPebbleRetireMigrationPersistsAcrossReopen(t *testing.T) { Versions: []MVCCVersion{{Key: []byte("kept-k"), CommitTS: 109, Value: []byte("v109")}}, }) require.NoError(t, err) + seedPromotionState(t, ctx, st, 9, []byte("stage|job9|"), []byte("job9-promoted")) + seedPromotionState(t, ctx, st, 10, []byte("stage|job10|"), []byte("job10-promoted")) require.NoError(t, st.RetireMigration(ctx, 9)) require.NoError(t, st.Close()) @@ -976,6 +1015,14 @@ func TestPebbleRetireMigrationPersistsAcrossReopen(t *testing.T) { floor, err = reopened.MigrationHLCFloor(ctx, 10) require.NoError(t, err) require.Equal(t, uint64(109), floor) + stateReader, ok := reopened.(MigrationPromotionStateReader) + require.True(t, ok) + _, ok = migrationPromotionState(t, ctx, stateReader, 9) + require.False(t, ok) + state, ok := migrationPromotionState(t, ctx, stateReader, 10) + require.True(t, ok) + require.True(t, state.Done) + require.Equal(t, uint64(1), state.PromotedRows) res, err := reopened.ImportVersions(ctx, ImportVersionsOptions{ JobID: 9, @@ -1070,3 +1117,478 @@ func TestExportVersionsRejectsZeroVersionBudget(t *testing.T) { } }) } + +func TestExportVersionsAppliesTimestampBoundBeforeAcceptVersion(t *testing.T) { + runMigrationStoreSuite(t, func(t *testing.T, st MVCCStore) { + ctx := context.Background() + require.NoError(t, st.PutAt(ctx, []byte("k"), []byte("eligible"), 20, 0)) + require.NoError(t, st.PutAt(ctx, []byte("k"), []byte("too-new"), 30, 0)) + accepted := false + + result, err := st.ExportVersions(ctx, ExportVersionsOptions{ + MaxCommitTSInclusive: 25, + MaxVersions: 1, + AcceptVersion: func(_ []byte, _ []byte) bool { + if accepted { + return false + } + accepted = true + return true + }, + }) + require.NoError(t, err) + require.Equal(t, []MVCCVersion{{Key: []byte("k"), CommitTS: 20, Value: []byte("eligible")}}, result.Versions) + }) +} + +func TestExportVersionsSkippedCursorBeforeStartResumesAtStartKey(t *testing.T) { + runMigrationStoreSuite(t, func(t *testing.T, st MVCCStore) { + ctx := context.Background() + require.NoError(t, st.PutAt(ctx, []byte("a"), []byte("a10"), 10, 0)) + require.NoError(t, st.PutAt(ctx, []byte("m"), []byte("m20"), 20, 0)) + + res, err := st.ExportVersions(ctx, ExportVersionsOptions{ + StartKey: []byte("m"), + EndKey: []byte("z"), + Cursor: encodeExportCursor([]byte("a"), 10, exportCursorTagSkippedKey), + MaxVersions: 10, + }) + require.NoError(t, err) + require.True(t, res.Done) + require.Equal(t, []MVCCVersion{{Key: []byte("m"), CommitTS: 20, Value: []byte("m20")}}, res.Versions) + }) +} + +func TestValidateExportCursorForRangeRejectsSkippedCursorInsideRange(t *testing.T) { + t.Parallel() + + err := ValidateExportCursorForRange( + encodeExportCursor([]byte("stage|k"), 10, exportCursorTagSkippedKey), + []byte("stage|"), + PrefixScanEnd([]byte("stage|")), + ) + require.ErrorIs(t, err, ErrInvalidExportCursor) + + err = ValidateExportCursorForRange( + encodeExportCursor([]byte("outside|k"), 10, exportCursorTagSkippedKey), + []byte("stage|"), + PrefixScanEnd([]byte("stage|")), + ) + require.NoError(t, err) +} + +func TestValidatePromotionCursorForRangeAcceptsOnlyEmittedPositions(t *testing.T) { + t.Parallel() + + prefix := []byte("stage|") + key := []byte("stage|k") + for _, tc := range []struct { + name string + cursor []byte + wantErr bool + }{ + {name: "empty cursor"}, + {name: "emitted cursor", cursor: encodeExportCursor(key, 10, exportCursorTagEmitted)}, + {name: "scanned cursor", cursor: encodeExportCursor(key, 10, exportCursorTagScanned), wantErr: true}, + {name: "pruned-key cursor", cursor: encodeExportCursor(key, 10, exportCursorTagPrunedKey), wantErr: true}, + {name: "skipped-key cursor", cursor: encodeExportCursor(key, 10, exportCursorTagSkippedKey), wantErr: true}, + {name: "emitted cursor outside range", cursor: encodeExportCursor([]byte("other|k"), 10, exportCursorTagEmitted), wantErr: true}, + } { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + + err := ValidatePromotionCursorForRange(tc.cursor, prefix, PrefixScanEnd(prefix)) + if tc.wantErr { + require.ErrorIs(t, err, ErrInvalidExportCursor) + return + } + require.NoError(t, err) + }) + } +} + +func TestPromoteVersionsMovesStagedVersionsAndDeletesStagedRows(t *testing.T) { + runMigrationStoreSuite(t, func(t *testing.T, st MVCCStore) { + ctx := context.Background() + promoter, ok := st.(MigrationPromoter) + require.True(t, ok) + stateReader, ok := st.(MigrationPromotionStateReader) + require.True(t, ok) + + stage := func(raw string) []byte { + return append([]byte("stage|"), []byte(raw)...) + } + targetKey := func(staged []byte) ([]byte, bool) { + return bytes.TrimPrefix(staged, []byte("stage|")), bytes.HasPrefix(staged, []byte("stage|")) + } + prefix := []byte("stage|") + + require.NoError(t, st.PutAt(ctx, []byte("k"), []byte("old"), 5, 0)) + require.NoError(t, st.PutAt(ctx, stage("k"), []byte("v10"), 10, 0)) + require.NoError(t, st.PutWithTTLAt(ctx, stage("k"), []byte("v20"), 20, 55)) + require.NoError(t, st.DeleteAt(ctx, stage("k"), 30)) + require.NoError(t, st.PutAt(ctx, stage("z"), []byte("z15"), 15, 0)) + + first, err := promoter.PromoteVersions(ctx, PromoteVersionsOptions{ + JobID: 99, + StartKey: prefix, + EndKey: PrefixScanEnd(prefix), + MaxVersions: 2, + TargetKey: targetKey, + }) + require.NoError(t, err) + require.False(t, first.Done) + require.Equal(t, uint64(2), first.PromotedRows) + require.Equal(t, uint64(2), first.TotalPromotedRows) + require.Equal(t, uint64(30), first.MaxPromotedTS) + require.NotEmpty(t, first.NextCursor) + state, ok, err := stateReader.MigrationPromotionState(ctx, 99) + require.NoError(t, err) + require.True(t, ok) + require.False(t, state.Done) + require.Equal(t, first.NextCursor, state.Cursor) + require.Equal(t, uint64(2), state.PromotedRows) + require.Equal(t, uint64(30), state.MaxPromotedTS) + + got, err := st.GetAt(ctx, []byte("k"), 25) + require.NoError(t, err) + require.Equal(t, []byte("v20"), got) + _, err = st.GetAt(ctx, []byte("k"), 35) + require.ErrorIs(t, err, ErrKeyNotFound) + + stagedLeft, err := st.ExportVersions(ctx, ExportVersionsOptions{ + StartKey: prefix, + EndKey: PrefixScanEnd(prefix), + MaxVersions: 10, + }) + require.NoError(t, err) + require.Equal(t, []MVCCVersion{ + {Key: stage("k"), CommitTS: 10, Value: []byte("v10")}, + {Key: stage("z"), CommitTS: 15, Value: []byte("z15")}, + }, stagedLeft.Versions) + + second, err := promoter.PromoteVersions(ctx, PromoteVersionsOptions{ + JobID: 99, + StartKey: prefix, + EndKey: PrefixScanEnd(prefix), + MaxVersions: 10, + TargetKey: targetKey, + }) + require.NoError(t, err) + require.True(t, second.Done) + require.Empty(t, second.NextCursor) + require.Equal(t, uint64(2), second.PromotedRows) + require.Equal(t, uint64(4), second.TotalPromotedRows) + require.Equal(t, uint64(30), second.MaxPromotedTS) + state, ok, err = stateReader.MigrationPromotionState(ctx, 99) + require.NoError(t, err) + require.True(t, ok) + require.True(t, state.Done) + require.Empty(t, state.Cursor) + require.Equal(t, uint64(4), state.PromotedRows) + require.Equal(t, uint64(30), state.MaxPromotedTS) + + retry, err := promoter.PromoteVersions(ctx, PromoteVersionsOptions{ + JobID: 99, + StartKey: prefix, + EndKey: PrefixScanEnd(prefix), + MaxVersions: 10, + TargetKey: targetKey, + }) + require.NoError(t, err) + require.True(t, retry.Done) + require.Zero(t, retry.PromotedRows) + require.Equal(t, uint64(4), retry.TotalPromotedRows) + require.Equal(t, uint64(30), retry.MaxPromotedTS) + + got, err = st.GetAt(ctx, []byte("k"), 10) + require.NoError(t, err) + require.Equal(t, []byte("v10"), got) + got, err = st.GetAt(ctx, []byte("z"), 15) + require.NoError(t, err) + require.Equal(t, []byte("z15"), got) + + stagedLeft, err = st.ExportVersions(ctx, ExportVersionsOptions{ + StartKey: prefix, + EndKey: PrefixScanEnd(prefix), + MaxVersions: 10, + }) + require.NoError(t, err) + require.True(t, stagedLeft.Done) + require.Empty(t, stagedLeft.Versions) + }) +} + +func TestPromoteVersionsIgnoresClientCursorWhenStateMissing(t *testing.T) { + runMigrationStoreSuite(t, func(t *testing.T, st MVCCStore) { + ctx := context.Background() + promoter, ok := st.(MigrationPromoter) + require.True(t, ok) + stateReader, ok := st.(MigrationPromotionStateReader) + require.True(t, ok) + + stage := func(raw string) []byte { + return append([]byte("stage|"), []byte(raw)...) + } + targetKey := func(staged []byte) ([]byte, bool) { + return bytes.TrimPrefix(staged, []byte("stage|")), bytes.HasPrefix(staged, []byte("stage|")) + } + prefix := []byte("stage|") + + require.NoError(t, st.PutAt(ctx, stage("a"), []byte("a10"), 10, 0)) + require.NoError(t, st.PutAt(ctx, stage("z"), []byte("z20"), 20, 0)) + staleCursor := encodeExportCursor(stage("m"), 1, exportCursorTagEmitted) + + result, err := promoter.PromoteVersions(ctx, PromoteVersionsOptions{ + JobID: 202, + StartKey: prefix, + EndKey: PrefixScanEnd(prefix), + Cursor: staleCursor, + MaxVersions: 10, + TargetKey: targetKey, + }) + require.NoError(t, err) + require.True(t, result.Done) + require.Equal(t, uint64(2), result.PromotedRows) + require.Equal(t, uint64(2), result.TotalPromotedRows) + + state, ok, err := stateReader.MigrationPromotionState(ctx, 202) + require.NoError(t, err) + require.True(t, ok) + require.True(t, state.Done) + require.Equal(t, uint64(2), state.PromotedRows) + require.Equal(t, uint64(20), state.MaxPromotedTS) + + got, err := st.GetAt(ctx, []byte("a"), 10) + require.NoError(t, err) + require.Equal(t, []byte("a10"), got) + got, err = st.GetAt(ctx, []byte("z"), 20) + require.NoError(t, err) + require.Equal(t, []byte("z20"), got) + stagedLeft, err := st.ExportVersions(ctx, ExportVersionsOptions{ + StartKey: prefix, + EndKey: PrefixScanEnd(prefix), + MaxVersions: 10, + }) + require.NoError(t, err) + require.True(t, stagedLeft.Done) + require.Empty(t, stagedLeft.Versions) + }) +} + +func TestPebblePromoteVersionsAdvancesLastCommitTS(t *testing.T) { + ctx := context.Background() + dir := t.TempDir() + st, err := NewPebbleStore(dir) + require.NoError(t, err) + closed := false + t.Cleanup(func() { + if !closed { + require.NoError(t, st.Close()) + } + }) + ps, ok := st.(*pebbleStore) + require.True(t, ok) + + stage := func(raw string) []byte { + return append([]byte("stage|"), []byte(raw)...) + } + targetKey := func(staged []byte) ([]byte, bool) { + return bytes.TrimPrefix(staged, []byte("stage|")), bytes.HasPrefix(staged, []byte("stage|")) + } + prefix := []byte("stage|") + + const promotedTS uint64 = 100 + require.NoError(t, ps.db.Set(encodeKey(stage("k"), promotedTS), encodeValue([]byte("v100"), false, 0, encStateCleartext), pebble.NoSync)) + require.Zero(t, ps.LastCommitTS()) + + result, err := ps.PromoteVersions(ctx, PromoteVersionsOptions{ + JobID: 101, + StartKey: prefix, + EndKey: PrefixScanEnd(prefix), + MaxVersions: 10, + TargetKey: targetKey, + }) + require.NoError(t, err) + require.True(t, result.Done) + require.Equal(t, uint64(1), result.PromotedRows) + require.Equal(t, promotedTS, result.MaxPromotedTS) + require.Equal(t, promotedTS, ps.LastCommitTS()) + state, ok, err := ps.MigrationPromotionState(ctx, 101) + require.NoError(t, err) + require.True(t, ok) + require.Equal(t, promotedTS, state.MaxPromotedTS) + + metaTS, err := readPebbleUint64(ps.db, metaLastCommitTSBytes) + require.NoError(t, err) + require.Equal(t, promotedTS, metaTS) + val, err := ps.GetAt(ctx, []byte("k"), ps.LastCommitTS()) + require.NoError(t, err) + require.Equal(t, []byte("v100"), val) + + require.NoError(t, st.Close()) + closed = true + reopened, err := NewPebbleStore(dir) + require.NoError(t, err) + defer func() { require.NoError(t, reopened.Close()) }() + require.Equal(t, promotedTS, reopened.LastCommitTS()) + val, err = reopened.GetAt(ctx, []byte("k"), reopened.LastCommitTS()) + require.NoError(t, err) + require.Equal(t, []byte("v100"), val) + reopenedPromoter, ok := reopened.(MigrationPromoter) + require.True(t, ok) + retry, err := reopenedPromoter.PromoteVersions(ctx, PromoteVersionsOptions{ + JobID: 101, + StartKey: prefix, + EndKey: PrefixScanEnd(prefix), + MaxVersions: 10, + TargetKey: targetKey, + }) + require.NoError(t, err) + require.True(t, retry.Done) + require.Zero(t, retry.PromotedRows) + require.Equal(t, uint64(1), retry.TotalPromotedRows) + require.Equal(t, promotedTS, retry.MaxPromotedTS) +} + +func TestPromotionStateCodecPreservesMaxPromotedTS(t *testing.T) { + t.Parallel() + + state := PromotionState{ + Cursor: []byte("cursor"), + Done: true, + PromotedRows: 7, + MaxPromotedTS: 42, + LastError: "boom", + } + decoded, ok := decodePromotionState(encodePromotionState(state)) + require.True(t, ok) + require.Equal(t, state, decoded) + + old := []byte{migrationPromotionDoneFlag} + old = binary.BigEndian.AppendUint64(old, 3) + old = binary.AppendUvarint(old, lenAsUint64(len("old-cursor"))) + old = append(old, "old-cursor"...) + old = binary.AppendUvarint(old, lenAsUint64(len("old-error"))) + old = append(old, "old-error"...) + decoded, ok = decodePromotionState(old) + require.True(t, ok) + require.True(t, decoded.Done) + require.Equal(t, uint64(3), decoded.PromotedRows) + require.Zero(t, decoded.MaxPromotedTS) + require.Equal(t, []byte("old-cursor"), decoded.Cursor) + require.Equal(t, "old-error", decoded.LastError) +} + +func TestPebbleRestoreStreamingSnapshotPreservesMigrationPromotionState(t *testing.T) { + ctx := context.Background() + src := NewMVCCStore() + promoter, ok := src.(MigrationPromoter) + require.True(t, ok) + + prefix := []byte("stage|") + targetKey := func(staged []byte) ([]byte, bool) { + return bytes.TrimPrefix(staged, prefix), bytes.HasPrefix(staged, prefix) + } + require.NoError(t, src.PutAt(ctx, []byte("stage|a"), []byte("va"), 100, 0)) + require.NoError(t, src.PutAt(ctx, []byte("stage|b"), []byte("vb"), 110, 0)) + + first, err := promoter.PromoteVersions(ctx, PromoteVersionsOptions{ + JobID: 12, + StartKey: prefix, + EndKey: PrefixScanEnd(prefix), + MaxVersions: 1, + TargetKey: targetKey, + }) + require.NoError(t, err) + require.False(t, first.Done) + require.Equal(t, uint64(1), first.TotalPromotedRows) + + snap, err := src.Snapshot() + require.NoError(t, err) + raw := snapshotBytes(t, snap) + require.NoError(t, snap.Close()) + + dstDir, err := os.MkdirTemp("", "migration-streaming-snapshot-dst-*") + require.NoError(t, err) + t.Cleanup(func() { require.NoError(t, os.RemoveAll(dstDir)) }) + dst, err := NewPebbleStore(dstDir) + require.NoError(t, err) + t.Cleanup(func() { require.NoError(t, dst.Close()) }) + require.NoError(t, dst.Restore(bytes.NewReader(raw))) + + dstPromoter, ok := any(dst).(MigrationPromoter) + require.True(t, ok) + stateReader, ok := any(dst).(MigrationPromotionStateReader) + require.True(t, ok) + state, ok, err := stateReader.MigrationPromotionState(ctx, 12) + require.NoError(t, err) + require.True(t, ok) + require.False(t, state.Done) + require.Equal(t, first.NextCursor, state.Cursor) + require.Equal(t, uint64(1), state.PromotedRows) + require.Equal(t, uint64(100), state.MaxPromotedTS) + + restored, err := dstPromoter.PromoteVersions(ctx, PromoteVersionsOptions{ + JobID: 12, + StartKey: prefix, + EndKey: PrefixScanEnd(prefix), + MaxVersions: 10, + TargetKey: targetKey, + }) + require.NoError(t, err) + require.True(t, restored.Done) + require.Equal(t, uint64(1), restored.PromotedRows) + require.Equal(t, uint64(2), restored.TotalPromotedRows) + require.Equal(t, uint64(110), restored.MaxPromotedTS) +} + +func TestExportVersionsSplitsBeforeOverflowingTheByteBudget(t *testing.T) { + runMigrationStoreSuite(t, func(t *testing.T, st MVCCStore) { + ctx := context.Background() + + small := bytes.Repeat([]byte("s"), 1<<10) + large := bytes.Repeat([]byte("l"), 8<<10) + require.NoError(t, st.PutAt(ctx, []byte("a"), small, 10, 0)) + require.NoError(t, st.PutAt(ctx, []byte("b"), large, 11, 0)) + + opts := ExportVersionsOptions{ + MaxVersions: 16, + MaxBytes: 4 << 10, + MaxCommitTSInclusive: 100, + EndKey: []byte("z"), + } + first, err := st.ExportVersions(ctx, opts) + require.NoError(t, err) + require.False(t, first.Done) + require.Len(t, first.Versions, 1, "the oversized row must not join this page") + require.Equal(t, []byte("a"), first.Versions[0].Key) + require.LessOrEqual(t, first.ExportedBytes, opts.MaxBytes, "the page stays inside its budget") + + opts.Cursor = first.NextCursor + second, err := st.ExportVersions(ctx, opts) + require.NoError(t, err) + require.Len(t, second.Versions, 1, "the oversized row goes out alone") + require.Equal(t, []byte("b"), second.Versions[0].Key) + require.Equal(t, large, second.Versions[0].Value) + }) +} + +func TestExportVersionsEmitsSingleOversizedRow(t *testing.T) { + runMigrationStoreSuite(t, func(t *testing.T, st MVCCStore) { + ctx := context.Background() + + huge := bytes.Repeat([]byte("h"), 8<<10) + require.NoError(t, st.PutAt(ctx, []byte("a"), huge, 10, 0)) + + got, err := st.ExportVersions(ctx, ExportVersionsOptions{ + MaxVersions: 16, + MaxBytes: 1 << 10, + MaxCommitTSInclusive: 100, + EndKey: []byte("z"), + }) + require.NoError(t, err) + require.Len(t, got.Versions, 1) + require.Equal(t, huge, got.Versions[0].Value) + }) +} diff --git a/store/mvcc_store.go b/store/mvcc_store.go index cc173c9c0..073db8183 100644 --- a/store/mvcc_store.go +++ b/store/mvcc_store.go @@ -25,10 +25,27 @@ type VersionedValue struct { } const ( - mvccSnapshotVersionV1 = uint32(1) - mvccSnapshotVersion = uint32(2) - maxSnapshotKeySize = 1 << 20 // 1 MiB per key - maxSnapshotVersionCount = 1 << 20 // 1M versions per key + mvccSnapshotVersionV1 = uint32(1) + mvccSnapshotVersionV2 = uint32(2) + mvccSnapshotVersionV3 = uint32(3) + maxSnapshotKeySize = 1 << 20 // 1 MiB per logical key + // maxSnapshotInternalKeyEnvelope reserves snapshot headroom for internal + // stored-key envelopes such as migration staging. Logical source keys keep + // the full maxSnapshotKeySize budget so existing snapshot-valid rows remain + // migratable. + maxSnapshotInternalKeyEnvelope = 32 + maxSnapshotStoredKeySize = maxSnapshotKeySize + maxSnapshotInternalKeyEnvelope + + // MaxSnapshotKeySize is the largest logical/source key size guaranteed to + // round-trip through snapshots. + MaxSnapshotKeySize = maxSnapshotKeySize + // MaxSnapshotInternalKeyEnvelope is the stored-key envelope headroom + // snapshots reserve for internal namespaces. + MaxSnapshotInternalKeyEnvelope = maxSnapshotInternalKeyEnvelope + // MaxSnapshotStoredKeySize is the largest physical stored key a snapshot can + // carry after adding a bounded internal envelope to a logical key. + MaxSnapshotStoredKeySize = maxSnapshotStoredKeySize + maxSnapshotVersionCount = 1 << 20 // 1M versions per key ) // maxSnapshotValueSize caps the allowed size of a single value during streaming @@ -61,13 +78,14 @@ func byteSliceComparator(a, b any) int { // mvccStore is an in-memory MVCC implementation backed by a treemap for // deterministic iteration order and range scans. type mvccStore struct { - tree *treemap.Map // key []byte -> []VersionedValue - mtx sync.RWMutex - log *slog.Logger - lastCommitTS uint64 - minRetainedTS uint64 - migrationAcks map[migrationAckID]migrationImportAck - migrationHLCFloors map[uint64]uint64 + tree *treemap.Map // key []byte -> []VersionedValue + mtx sync.RWMutex + log *slog.Logger + lastCommitTS uint64 + minRetainedTS uint64 + migrationAcks map[migrationAckID]migrationImportAck + migrationHLCFloors map[uint64]uint64 + migrationPromotions map[uint64]PromotionState // writeConflicts mirrors the per-(kind, key_prefix) counter from // the pebble-backed store so the in-memory implementation shows up // in the same Prometheus series (even if the counts are usually @@ -114,9 +132,10 @@ func NewMVCCStore(opts ...MVCCStoreOption) MVCCStore { log: slog.New(slog.NewJSONHandler(os.Stdout, &slog.HandlerOptions{ Level: slog.LevelWarn, })), - migrationAcks: make(map[migrationAckID]migrationImportAck), - migrationHLCFloors: make(map[uint64]uint64), - writeConflicts: newWriteConflictCounter(), + migrationAcks: make(map[migrationAckID]migrationImportAck), + migrationHLCFloors: make(map[uint64]uint64), + migrationPromotions: make(map[uint64]PromotionState), + writeConflicts: newWriteConflictCounter(), } for _, opt := range opts { opt(s) @@ -787,10 +806,22 @@ func (s *mvccStore) DeletePrefixAtRaft(ctx context.Context, prefix []byte, exclu // DeletePrefixAtRaftAt satisfies the MVCCStore interface — see // ApplyMutationsRaftAt for the appliedIndex disposition rationale. func (s *mvccStore) DeletePrefixAtRaftAt(ctx context.Context, prefix []byte, excludePrefix []byte, commitTS, _ uint64) error { - return s.DeletePrefixAt(ctx, prefix, excludePrefix, commitTS) + return s.DeletePrefixesAtRaftAt(ctx, []PrefixDelete{{Prefix: prefix, ExcludePrefix: excludePrefix}}, commitTS, 0) +} + +func (s *mvccStore) DeletePrefixesAtRaftAt(_ context.Context, deletes []PrefixDelete, commitTS, _ uint64) error { + return s.deletePrefixesAt(deletes, commitTS) } func (s *mvccStore) DeletePrefixAt(_ context.Context, prefix []byte, excludePrefix []byte, commitTS uint64) error { + return s.deletePrefixesAt([]PrefixDelete{{Prefix: prefix, ExcludePrefix: excludePrefix}}, commitTS) +} + +func (s *mvccStore) deletePrefixesAt(deletes []PrefixDelete, commitTS uint64) error { + if len(deletes) == 0 { + return nil + } + s.mtx.Lock() defer s.mtx.Unlock() @@ -798,29 +829,9 @@ func (s *mvccStore) DeletePrefixAt(_ context.Context, prefix []byte, excludePref // Collect matching keys first since we cannot modify the tree while iterating. var toDelete [][]byte - it := s.tree.Iterator() - var started bool - if len(prefix) > 0 { - started = seekForwardIteratorStart(s.tree, &it, prefix) - } else { - started = it.First() - } - for ok := started; ok; ok = it.Next() { - k, keyOK := it.Key().([]byte) - if !keyOK { - continue - } - if len(prefix) > 0 && !bytes.HasPrefix(k, prefix) { - break - } - if len(excludePrefix) > 0 && bytes.HasPrefix(k, excludePrefix) { - continue - } - versions, _ := it.Value().([]VersionedValue) - if _, visible := visibleValue(versions, commitTS); !visible { - continue - } - toDelete = append(toDelete, k) + seen := make(map[string]struct{}) + for _, del := range deletes { + toDelete = s.collectDeletePrefixKeysLocked(del, commitTS, seen, toDelete) } for _, k := range toDelete { @@ -903,12 +914,20 @@ func isStreamingMVCCSnapshot(r *bufio.Reader) (bool, error) { } func (s *mvccStore) writeSnapshotFile(f *os.File) error { - checksumOffset, err := writeMVCCSnapshotHeader(f) + // The version and the body are chosen under one lock: the version depends on + // which metadata maps are populated, so releasing between the header and the + // body would let a concurrent migration write a body the header does not + // describe. + s.mtx.RLock() + defer s.mtx.RUnlock() + + version := mvccSnapshotVersionForState(s.migrationAcks, s.migrationHLCFloors, s.migrationPromotions) + checksumOffset, err := writeMVCCSnapshotHeader(f, version) if err != nil { return err } - sum, err := s.writeSnapshotBody(f) + sum, err := s.writeSnapshotBodyLocked(f, version) if err != nil { return err } @@ -916,11 +935,33 @@ func (s *mvccStore) writeSnapshotFile(f *os.File) error { return finalizeMVCCSnapshotFile(f, checksumOffset, sum) } -func writeMVCCSnapshotHeader(f *os.File) (int64, error) { +// mvccSnapshotVersionForState returns the oldest snapshot layout that can carry +// this state. Every version is its predecessor plus one trailing metadata +// field, so a store holding none of that metadata still emits the layout older +// binaries already parse. Emitting the newest layout unconditionally would +// strand a not-yet-upgraded follower mid-rolling-upgrade on a snapshot it +// rejects as unsupported, even on a cluster where the migration opcodes have +// never been enabled and there is nothing extra to carry. +func mvccSnapshotVersionForState( + acks map[migrationAckID]migrationImportAck, + floors map[uint64]uint64, + promotions map[uint64]PromotionState, +) uint32 { + switch { + case len(promotions) > 0: + return mvccSnapshotVersionV3 + case len(acks) > 0 || len(floors) > 0: + return mvccSnapshotVersionV2 + default: + return mvccSnapshotVersionV1 + } +} + +func writeMVCCSnapshotHeader(f *os.File, version uint32) (int64, error) { if _, err := f.Write(mvccSnapshotMagic[:]); err != nil { return 0, errors.WithStack(err) } - if err := binary.Write(f, binary.LittleEndian, mvccSnapshotVersion); err != nil { + if err := binary.Write(f, binary.LittleEndian, version); err != nil { return 0, errors.WithStack(err) } checksumOffset, err := f.Seek(0, io.SeekCurrent) @@ -933,24 +974,20 @@ func writeMVCCSnapshotHeader(f *os.File) (int64, error) { return checksumOffset, nil } -func (s *mvccStore) writeSnapshotBody(f *os.File) (uint32, error) { +func (s *mvccStore) writeSnapshotBodyLocked(f *os.File, version uint32) (uint32, error) { hash := crc32.NewIEEE() bw := bufio.NewWriter(f) w := io.MultiWriter(bw, hash) - s.mtx.RLock() - defer s.mtx.RUnlock() - - if err := binary.Write(w, binary.LittleEndian, s.lastCommitTS); err != nil { - return 0, errors.WithStack(err) - } - if err := binary.Write(w, binary.LittleEndian, s.minRetainedTS); err != nil { - return 0, errors.WithStack(err) - } - if err := writeMVCCSnapshotBytes(w, encodeMigrationImportAcks(s.migrationAcks)); err != nil { - return 0, err - } - if err := writeMVCCSnapshotBytes(w, encodeMigrationHLCFloors(s.migrationHLCFloors)); err != nil { + if err := writeMVCCSnapshotMetadata( + w, + version, + s.lastCommitTS, + s.minRetainedTS, + s.migrationAcks, + s.migrationHLCFloors, + s.migrationPromotions, + ); err != nil { return 0, err } iter := s.tree.Iterator() @@ -973,6 +1010,39 @@ func (s *mvccStore) writeSnapshotBody(f *os.File) (uint32, error) { return hash.Sum32(), nil } +func writeMVCCSnapshotMetadata( + w io.Writer, + version uint32, + lastCommitTS uint64, + minRetainedTS uint64, + migrationAcks map[migrationAckID]migrationImportAck, + migrationHLCFloors map[uint64]uint64, + migrationPromotions map[uint64]PromotionState, +) error { + if err := binary.Write(w, binary.LittleEndian, lastCommitTS); err != nil { + return errors.WithStack(err) + } + if err := binary.Write(w, binary.LittleEndian, minRetainedTS); err != nil { + return errors.WithStack(err) + } + if version == mvccSnapshotVersionV1 { + return nil + } + if err := writeMVCCSnapshotBytes(w, encodeMigrationImportAcks(migrationAcks)); err != nil { + return err + } + if err := writeMVCCSnapshotBytes(w, encodeMigrationHLCFloors(migrationHLCFloors)); err != nil { + return err + } + if version < mvccSnapshotVersionV3 { + return nil + } + if err := writeMVCCSnapshotBytes(w, encodeMigrationPromotionStates(migrationPromotions)); err != nil { + return err + } + return nil +} + func finalizeMVCCSnapshotFile(f *os.File, checksumOffset int64, sum uint32) error { if _, err := f.Seek(checksumOffset, io.SeekStart); err != nil { return errors.WithStack(err) @@ -1033,6 +1103,39 @@ func writeMVCCSnapshotVersion(w io.Writer, version VersionedValue) error { return nil } +func (s *mvccStore) collectDeletePrefixKeysLocked(del PrefixDelete, commitTS uint64, seen map[string]struct{}, toDelete [][]byte) [][]byte { + it := s.tree.Iterator() + var started bool + if len(del.Prefix) > 0 { + started = seekForwardIteratorStart(s.tree, &it, del.Prefix) + } else { + started = it.First() + } + for ok := started; ok; ok = it.Next() { + k, keyOK := it.Key().([]byte) + if !keyOK { + continue + } + if len(del.Prefix) > 0 && !bytes.HasPrefix(k, del.Prefix) { + break + } + if len(del.ExcludePrefix) > 0 && bytes.HasPrefix(k, del.ExcludePrefix) { + continue + } + versions, _ := it.Value().([]VersionedValue) + if _, visible := visibleValue(versions, commitTS); !visible { + continue + } + dedupeKey := string(k) + if _, ok := seen[dedupeKey]; ok { + continue + } + seen[dedupeKey] = struct{}{} + toDelete = append(toDelete, bytes.Clone(k)) + } + return toDelete +} + func mvccSnapshotTombstoneByte(tombstone bool) byte { if tombstone { return 1 @@ -1046,7 +1149,7 @@ func (s *mvccStore) restoreStreamingSnapshot(r io.Reader) error { return err } - tree, lastCommitTS, minRetainedTS, migrationAcks, migrationHLCFloors, actual, err := restoreStreamingMVCCSnapshotBody(r, version) + tree, lastCommitTS, minRetainedTS, migrationAcks, migrationHLCFloors, migrationPromotions, actual, err := restoreStreamingMVCCSnapshotBody(r, version) if err != nil { return err } @@ -1061,6 +1164,7 @@ func (s *mvccStore) restoreStreamingSnapshot(r io.Reader) error { s.minRetainedTS = minRetainedTS s.migrationAcks = migrationAcks s.migrationHLCFloors = migrationHLCFloors + s.migrationPromotions = migrationPromotions return nil } @@ -1077,7 +1181,7 @@ func readMVCCSnapshotHeader(r io.Reader) (uint32, uint32, error) { if err := binary.Read(r, binary.LittleEndian, &version); err != nil { return 0, 0, errors.WithStack(err) } - if version != mvccSnapshotVersionV1 && version != mvccSnapshotVersion { + if version != mvccSnapshotVersionV1 && version != mvccSnapshotVersionV2 && version != mvccSnapshotVersionV3 { return 0, 0, errors.WithStack(errors.Newf("unsupported mvcc snapshot version %d", version)) } @@ -1088,55 +1192,86 @@ func readMVCCSnapshotHeader(r io.Reader) (uint32, uint32, error) { return version, expected, nil } -func restoreStreamingMVCCSnapshotBody(r io.Reader, version uint32) (*treemap.Map, uint64, uint64, map[migrationAckID]migrationImportAck, map[uint64]uint64, uint32, error) { +func restoreStreamingMVCCSnapshotBody( + r io.Reader, + version uint32, +) (*treemap.Map, uint64, uint64, map[migrationAckID]migrationImportAck, map[uint64]uint64, map[uint64]PromotionState, uint32, error) { hash := crc32.NewIEEE() body := io.TeeReader(r, hash) - lastCommitTS, minRetainedTS, migrationAcks, migrationHLCFloors, err := readMVCCSnapshotMetadata(body, version) + lastCommitTS, minRetainedTS, migrationAcks, migrationHLCFloors, migrationPromotions, err := readMVCCSnapshotMetadata(body, version) if err != nil { - return nil, 0, 0, nil, nil, 0, err + return nil, 0, 0, nil, nil, nil, 0, err } tree, err := readMVCCSnapshotTree(body) if err != nil { - return nil, 0, 0, nil, nil, 0, err + return nil, 0, 0, nil, nil, nil, 0, err } - return tree, lastCommitTS, minRetainedTS, migrationAcks, migrationHLCFloors, hash.Sum32(), nil + return tree, lastCommitTS, minRetainedTS, migrationAcks, migrationHLCFloors, migrationPromotions, hash.Sum32(), nil } -func readMVCCSnapshotMetadata(r io.Reader, version uint32) (uint64, uint64, map[migrationAckID]migrationImportAck, map[uint64]uint64, error) { +func readMVCCSnapshotMetadata( + r io.Reader, + version uint32, +) (uint64, uint64, map[migrationAckID]migrationImportAck, map[uint64]uint64, map[uint64]PromotionState, error) { var lastCommitTS uint64 if err := binary.Read(r, binary.LittleEndian, &lastCommitTS); err != nil { - return 0, 0, nil, nil, errors.WithStack(err) + return 0, 0, nil, nil, nil, errors.WithStack(err) } var minRetainedTS uint64 if err := binary.Read(r, binary.LittleEndian, &minRetainedTS); err != nil { - return 0, 0, nil, nil, errors.WithStack(err) + return 0, 0, nil, nil, nil, errors.WithStack(err) } if version == mvccSnapshotVersionV1 { - return lastCommitTS, minRetainedTS, make(map[migrationAckID]migrationImportAck), make(map[uint64]uint64), nil + return lastCommitTS, + minRetainedTS, + make(map[migrationAckID]migrationImportAck), + make(map[uint64]uint64), + make(map[uint64]PromotionState), + nil } ackData, err := readMVCCSnapshotBytes(r, "snapshot migration acks") if err != nil { - return 0, 0, nil, nil, err + return 0, 0, nil, nil, nil, err } migrationAcks, ok := decodeMigrationImportAcks(ackData) if !ok { - return 0, 0, nil, nil, errors.New("invalid snapshot migration acks") + return 0, 0, nil, nil, nil, errors.New("invalid snapshot migration acks") } floorData, err := readMVCCSnapshotBytes(r, "snapshot migration hlc floors") if err != nil { - return 0, 0, nil, nil, err + return 0, 0, nil, nil, nil, err } migrationHLCFloors, ok := decodeMigrationHLCFloors(floorData) if !ok { - return 0, 0, nil, nil, errors.New("invalid snapshot migration hlc floors") + return 0, 0, nil, nil, nil, errors.New("invalid snapshot migration hlc floors") } - return lastCommitTS, minRetainedTS, migrationAcks, migrationHLCFloors, nil + migrationPromotions, err := readMVCCSnapshotPromotionMetadata(r, version) + if err != nil { + return 0, 0, nil, nil, nil, err + } + + return lastCommitTS, minRetainedTS, migrationAcks, migrationHLCFloors, migrationPromotions, nil +} + +func readMVCCSnapshotPromotionMetadata(r io.Reader, version uint32) (map[uint64]PromotionState, error) { + if version < mvccSnapshotVersionV3 { + return make(map[uint64]PromotionState), nil + } + promotionData, err := readMVCCSnapshotBytes(r, "snapshot migration promotions") + if err != nil { + return nil, err + } + migrationPromotions, ok := decodeMigrationPromotionStates(promotionData) + if !ok { + return nil, errors.New("invalid snapshot migration promotions") + } + return migrationPromotions, nil } func readMVCCSnapshotBytes(r io.Reader, field string) ([]byte, error) { @@ -1177,8 +1312,8 @@ func readMVCCSnapshotEntry(r io.Reader) ([]byte, []VersionedValue, bool, error) } return nil, nil, false, errors.WithStack(err) } - if keyLen > maxSnapshotKeySize { - return nil, nil, false, errors.Wrapf(ErrSnapshotKeyTooLarge, "%d > %d", keyLen, maxSnapshotKeySize) + if keyLen > maxSnapshotStoredKeySize { + return nil, nil, false, errors.Wrapf(ErrSnapshotKeyTooLarge, "%d > %d", keyLen, maxSnapshotStoredKeySize) } key := make([]byte, keyLen) diff --git a/store/mvcc_store_snapshot_test.go b/store/mvcc_store_snapshot_test.go index 7b0bd39a3..c52c050de 100644 --- a/store/mvcc_store_snapshot_test.go +++ b/store/mvcc_store_snapshot_test.go @@ -38,6 +38,27 @@ func TestMVCCStore_SnapshotRestoreRoundTrip(t *testing.T) { require.Equal(t, []byte("v2"), v) } +func TestMVCCStore_SnapshotRestoreMaxStoredKey(t *testing.T) { + t.Parallel() + + ctx := context.Background() + src := newTestMVCCStore(t) + key := bytes.Repeat([]byte("k"), MaxSnapshotStoredKeySize) + require.NoError(t, src.PutAt(ctx, key, []byte("v"), 10, 0)) + + snap, err := src.Snapshot() + require.NoError(t, err) + defer snap.Close() + raw := snapshotBytes(t, snap) + + dst := newTestMVCCStore(t) + require.NoError(t, dst.Restore(bytes.NewReader(raw))) + + got, err := dst.GetAt(ctx, key, 10) + require.NoError(t, err) + require.Equal(t, []byte("v"), got) +} + func TestMVCCStore_RestoreRejectsInvalidChecksum(t *testing.T) { t.Parallel() @@ -61,6 +82,19 @@ func TestMVCCStore_RestoreClearsMigrationMetadata(t *testing.T) { ctx := context.Background() st := newTestMVCCStore(t) + promoter, ok := any(st).(MigrationPromoter) + require.True(t, ok) + stateReader, ok := any(st).(MigrationPromotionStateReader) + require.True(t, ok) + + prefix := []byte("stage|") + stage := func(raw string) []byte { + return append([]byte("stage|"), []byte(raw)...) + } + targetKey := func(staged []byte) ([]byte, bool) { + return bytes.TrimPrefix(staged, prefix), bytes.HasPrefix(staged, prefix) + } + require.NoError(t, st.PutAt(ctx, []byte("base"), []byte("v1"), 10, 0)) snap, err := st.Snapshot() @@ -80,13 +114,46 @@ func TestMVCCStore_RestoreClearsMigrationMetadata(t *testing.T) { require.NoError(t, err) require.Equal(t, uint64(50), floor) + require.NoError(t, st.PutAt(ctx, stage("stale"), []byte("old"), 70, 0)) + promoted, err := promoter.PromoteVersions(ctx, PromoteVersionsOptions{ + JobID: 7, + StartKey: prefix, + EndKey: PrefixScanEnd(prefix), + MaxVersions: 10, + TargetKey: targetKey, + }) + require.NoError(t, err) + require.True(t, promoted.Done) + state, ok, err := stateReader.MigrationPromotionState(ctx, 7) + require.NoError(t, err) + require.True(t, ok) + require.True(t, state.Done) + require.NoError(t, st.Restore(bytes.NewReader(raw))) floor, err = st.MigrationHLCFloor(ctx, 7) require.NoError(t, err) require.Zero(t, floor) + _, ok, err = stateReader.MigrationPromotionState(ctx, 7) + require.NoError(t, err) + require.False(t, ok) _, err = st.GetAt(ctx, []byte("imported"), 50) require.ErrorIs(t, err, ErrKeyNotFound) + require.NoError(t, st.PutAt(ctx, stage("fresh"), []byte("new"), 80, 0)) + promoted, err = promoter.PromoteVersions(ctx, PromoteVersionsOptions{ + JobID: 7, + StartKey: prefix, + EndKey: PrefixScanEnd(prefix), + MaxVersions: 10, + TargetKey: targetKey, + }) + require.NoError(t, err) + require.True(t, promoted.Done) + require.Equal(t, uint64(1), promoted.PromotedRows) + got, err := st.GetAt(ctx, []byte("fresh"), 80) + require.NoError(t, err) + require.Equal(t, []byte("new"), got) + res, err := st.ImportVersions(ctx, ImportVersionsOptions{ JobID: 7, BracketID: 3, @@ -151,6 +218,74 @@ func TestMVCCStore_SnapshotRestorePreservesMigrationMetadata(t *testing.T) { require.ErrorIs(t, err, ErrKeyNotFound) } +func TestMVCCStore_SnapshotRestorePreservesMigrationPromotionState(t *testing.T) { + t.Parallel() + + ctx := context.Background() + st := newTestMVCCStore(t) + promoter, ok := any(st).(MigrationPromoter) + require.True(t, ok) + stateReader, ok := any(st).(MigrationPromotionStateReader) + require.True(t, ok) + + prefix := []byte("stage|") + targetKey := func(staged []byte) ([]byte, bool) { + return bytes.TrimPrefix(staged, prefix), bytes.HasPrefix(staged, prefix) + } + require.NoError(t, st.PutAt(ctx, []byte("stage|a"), []byte("va"), 100, 0)) + require.NoError(t, st.PutAt(ctx, []byte("stage|b"), []byte("vb"), 110, 0)) + + first, err := promoter.PromoteVersions(ctx, PromoteVersionsOptions{ + JobID: 11, + StartKey: prefix, + EndKey: PrefixScanEnd(prefix), + MaxVersions: 1, + TargetKey: targetKey, + }) + require.NoError(t, err) + require.False(t, first.Done) + require.Equal(t, uint64(1), first.TotalPromotedRows) + require.Equal(t, uint64(100), first.MaxPromotedTS) + + snap, err := st.Snapshot() + require.NoError(t, err) + defer snap.Close() + raw := snapshotBytes(t, snap) + + rest, err := promoter.PromoteVersions(ctx, PromoteVersionsOptions{ + JobID: 11, + StartKey: prefix, + EndKey: PrefixScanEnd(prefix), + MaxVersions: 10, + TargetKey: targetKey, + }) + require.NoError(t, err) + require.True(t, rest.Done) + require.Equal(t, uint64(2), rest.TotalPromotedRows) + + require.NoError(t, st.Restore(bytes.NewReader(raw))) + state, ok, err := stateReader.MigrationPromotionState(ctx, 11) + require.NoError(t, err) + require.True(t, ok) + require.False(t, state.Done) + require.Equal(t, first.NextCursor, state.Cursor) + require.Equal(t, uint64(1), state.PromotedRows) + require.Equal(t, uint64(100), state.MaxPromotedTS) + + restored, err := promoter.PromoteVersions(ctx, PromoteVersionsOptions{ + JobID: 11, + StartKey: prefix, + EndKey: PrefixScanEnd(prefix), + MaxVersions: 10, + TargetKey: targetKey, + }) + require.NoError(t, err) + require.True(t, restored.Done) + require.Equal(t, uint64(1), restored.PromotedRows) + require.Equal(t, uint64(2), restored.TotalPromotedRows) + require.Equal(t, uint64(110), restored.MaxPromotedTS) +} + func TestMVCCStore_ApplyMutations_WriteConflict(t *testing.T) { t.Parallel() diff --git a/store/mvcc_store_snapshot_version_test.go b/store/mvcc_store_snapshot_version_test.go new file mode 100644 index 000000000..a461550a0 --- /dev/null +++ b/store/mvcc_store_snapshot_version_test.go @@ -0,0 +1,150 @@ +package store + +import ( + "bytes" + "context" + "encoding/binary" + "testing" + + "github.com/stretchr/testify/require" +) + +// snapshotHeaderVersion reads the layout version a snapshot advertises, which +// is what an older binary checks against its own accept list before it will +// restore from it. +func snapshotHeaderVersion(t *testing.T, raw []byte) uint32 { + t.Helper() + + require.Greater(t, len(raw), len(mvccSnapshotMagic)+4) + require.Equal(t, mvccSnapshotMagic[:], raw[:len(mvccSnapshotMagic)]) + return binary.LittleEndian.Uint32(raw[len(mvccSnapshotMagic) : len(mvccSnapshotMagic)+4]) +} + +// Each snapshot layout is its predecessor plus one trailing metadata field, so +// a replica must advertise the oldest layout that can carry the state it +// actually holds. Advertising the newest one unconditionally strands a +// not-yet-upgraded follower during a rolling upgrade: it rejects the snapshot +// as an unsupported version and cannot catch up, even on a cluster where no +// migration has ever run and the extra fields would all be empty. +func TestMVCCStoreSnapshotEmitsOldestSufficientVersion(t *testing.T) { + t.Parallel() + + prefix := []byte("stage|") + seedPromotion := func(t *testing.T, st MVCCStore) { + t.Helper() + promoter, ok := any(st).(MigrationPromoter) + require.True(t, ok) + require.NoError(t, st.PutAt(context.Background(), append(bytes.Clone(prefix), []byte("k")...), []byte("v"), 70, 0)) + promoted, err := promoter.PromoteVersions(context.Background(), PromoteVersionsOptions{ + JobID: 7, + StartKey: prefix, + EndKey: PrefixScanEnd(prefix), + MaxVersions: 10, + TargetKey: func(staged []byte) ([]byte, bool) { + return bytes.TrimPrefix(staged, prefix), bytes.HasPrefix(staged, prefix) + }, + }) + require.NoError(t, err) + require.True(t, promoted.Done) + } + seedImport := func(t *testing.T, st MVCCStore) { + t.Helper() + _, err := st.ImportVersions(context.Background(), ImportVersionsOptions{ + JobID: 7, + BracketID: 3, + BatchSeq: 1, + Cursor: []byte("c"), + Versions: []MVCCVersion{{Key: []byte("imported"), CommitTS: 50, Value: []byte("v50")}}, + }) + require.NoError(t, err) + } + + for _, tc := range []struct { + name string + seed func(*testing.T, MVCCStore) + want uint32 + }{ + { + name: "no migration metadata stays on the original layout", + want: mvccSnapshotVersionV1, + }, + { + name: "import acks and hlc floors need v2", + seed: seedImport, + want: mvccSnapshotVersionV2, + }, + { + name: "promotion state needs v3", + seed: seedPromotion, + want: mvccSnapshotVersionV3, + }, + } { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + + ctx := context.Background() + st := newTestMVCCStore(t) + require.NoError(t, st.PutAt(ctx, []byte("base"), []byte("v1"), 10, 0)) + if tc.seed != nil { + tc.seed(t, st) + } + + snap, err := st.Snapshot() + require.NoError(t, err) + defer snap.Close() + raw := snapshotBytes(t, snap) + + require.Equal(t, tc.want, snapshotHeaderVersion(t, raw)) + + // Whatever version was chosen, the body must match it: restoring + // through the real reader reproduces the data and the metadata. + dst := newTestMVCCStore(t) + require.NoError(t, dst.Restore(bytes.NewReader(raw))) + got, err := dst.GetAt(ctx, []byte("base"), 10) + require.NoError(t, err) + require.Equal(t, []byte("v1"), got) + }) + } +} + +// The metadata a newer layout exists to carry must survive the round trip, so +// the version choice cannot be made by simply dropping fields. +func TestMVCCStoreSnapshotRoundTripsMigrationMetadataAtItsVersion(t *testing.T) { + t.Parallel() + + ctx := context.Background() + st := newTestMVCCStore(t) + _, err := st.ImportVersions(ctx, ImportVersionsOptions{ + JobID: 7, + BracketID: 3, + BatchSeq: 1, + Cursor: []byte("c"), + Versions: []MVCCVersion{{Key: []byte("imported"), CommitTS: 50, Value: []byte("v50")}}, + }) + require.NoError(t, err) + + snap, err := st.Snapshot() + require.NoError(t, err) + defer snap.Close() + raw := snapshotBytes(t, snap) + require.Equal(t, mvccSnapshotVersionV2, snapshotHeaderVersion(t, raw)) + + dst := newTestMVCCStore(t) + require.NoError(t, dst.Restore(bytes.NewReader(raw))) + floor, err := dst.MigrationHLCFloor(ctx, 7) + require.NoError(t, err) + require.Equal(t, uint64(50), floor, "the hlc floor is why v2 exists; it must survive") + + // The recorded ack is what makes a replayed batch a duplicate rather than a + // second application, so it has to survive the restore as well. + result, err := dst.ImportVersions(ctx, ImportVersionsOptions{ + JobID: 7, + BracketID: 3, + BatchSeq: 1, + Cursor: []byte("replayed"), + Versions: []MVCCVersion{{Key: []byte("imported"), CommitTS: 50, Value: []byte("v50")}}, + }) + require.NoError(t, err) + require.True(t, result.Duplicate) + require.Equal(t, []byte("c"), result.AckedCursor) +} diff --git a/store/store.go b/store/store.go index 0fae79552..853d70b69 100644 --- a/store/store.go +++ b/store/store.go @@ -38,6 +38,13 @@ var ErrInvalidExportCursor = errors.New("invalid export cursor") var ErrInvalidExportBudget = errors.New("migration export requires a positive version budget") var ErrImportBatchGap = errors.New("migration import batch gap") +// ErrInvalidImportVersion marks a migration import version that is malformed +// on its face (zero commit_ts, a tombstone carrying a value or expire_at). +// It is a property of the request bytes, so every replica applying the same +// Raft entry reaches the same verdict -- which is what lets kv/fsm classify +// it as an ordinary apply error instead of halting the apply loop. +var ErrInvalidImportVersion = errors.New("invalid migration import version") + // validateValueSize returns ErrValueTooLarge when the value exceeds maxSnapshotValueSize. func validateValueSize(value []byte) error { if len(value) > maxSnapshotValueSize { @@ -76,6 +83,12 @@ type KVPair struct { RouteGroupID uint64 } +// PrefixDelete describes one prefix tombstone operation in a batched apply. +type PrefixDelete struct { + Prefix []byte + ExcludePrefix []byte +} + // MVCCVersion is a raw committed MVCC version for range migration. // Unlike scan results, it preserves tombstones and TTL expiry metadata. type MVCCVersion struct { @@ -114,11 +127,14 @@ type ExportVersionsResult struct { // ImportVersionsOptions applies one idempotent migration-import batch. type ImportVersionsOptions struct { - JobID uint64 - BracketID uint64 - BatchSeq uint64 - Versions []MVCCVersion - Cursor []byte + JobID uint64 + // AppliedIndex is the optional Raft entry index to bundle with Pebble + // import batches as metaAppliedIndex. Zero leaves the meta key unchanged. + AppliedIndex uint64 + BracketID uint64 + BatchSeq uint64 + Versions []MVCCVersion + Cursor []byte } // ImportVersionsResult reports the cursor durably acknowledged by the target. @@ -128,6 +144,53 @@ type ImportVersionsResult struct { Duplicate bool } +// PromoteVersionsOptions atomically copies staged MVCC versions to their +// target keys and physically removes the staged versions. +type PromoteVersionsOptions struct { + JobID uint64 + // AppliedIndex is the optional Raft entry index to bundle with Pebble + // promotion batches as metaAppliedIndex. Zero leaves the meta key unchanged. + AppliedIndex uint64 + StartKey []byte + EndKey []byte + Cursor []byte + MaxVersions int + MaxBytes uint64 + MaxScannedBytes uint64 + TargetKey func(stagedKey []byte) ([]byte, bool) +} + +// PromoteVersionsResult reports one resumable staged-version promotion chunk. +type PromoteVersionsResult struct { + NextCursor []byte + Done bool + PromotedRows uint64 + TotalPromotedRows uint64 + PromotedBytes uint64 + MaxPromotedTS uint64 + ScannedBytes uint64 +} + +// PromotionState is the target-local durable cursor for staged data promotion. +type PromotionState struct { + Cursor []byte + Done bool + PromotedRows uint64 + MaxPromotedTS uint64 + LastError string +} + +// MigrationPromoter is implemented by stores that can promote staged range +// migration data into the live keyspace. +type MigrationPromoter interface { + PromoteVersions(ctx context.Context, opts PromoteVersionsOptions) (PromoteVersionsResult, error) +} + +// MigrationPromotionStateReader reads target-local staged promotion state. +type MigrationPromotionStateReader interface { + MigrationPromotionState(ctx context.Context, jobID uint64) (PromotionState, bool, error) +} + // OpType describes a mutation kind. type OpType int @@ -269,6 +332,11 @@ type MVCCStore interface { // bundles metaAppliedIndex in that batch so DEL_PREFIX entries // also advance the meta key. PR #910 design §2 "why both leaves". DeletePrefixAtRaftAt(ctx context.Context, prefix []byte, excludePrefix []byte, commitTS, appliedIndex uint64) error + // DeletePrefixesAtRaftAt applies several prefix deletes in one + // raft-apply batch. It is used when one logical raft command must + // tombstone multiple physical namespaces without exposing a partial + // apply or advancing metaAppliedIndex separately from any tombstone. + DeletePrefixesAtRaftAt(ctx context.Context, deletes []PrefixDelete, commitTS, appliedIndex uint64) error // LastCommitTS returns the highest commit timestamp applied on this node. LastCommitTS() uint64 // WriteConflictCountsByPrefix returns a snapshot of the MVCC @@ -285,6 +353,11 @@ type MVCCStore interface { // ImportVersions applies a migration import batch idempotently by // (jobID, bracketID, batchSeq), preserving tombstones and expireAt. ImportVersions(ctx context.Context, opts ImportVersionsOptions) (ImportVersionsResult, error) + // ImportVersionsRaft is the raft-apply variant of ImportVersions. It + // preserves the same idempotency contract while using the FSM write path. + // When opts.AppliedIndex is non-zero, the implementation must durably + // bundle metaAppliedIndex with the import batch. + ImportVersionsRaft(ctx context.Context, opts ImportVersionsOptions) (ImportVersionsResult, error) // MigrationHLCFloor returns the full-HLC target-local migration floor // persisted by ImportVersions for jobID. MigrationHLCFloor(ctx context.Context, jobID uint64) (uint64, error)