diff --git a/adapter/distribution_server.go b/adapter/distribution_server.go index f8b973f3d..96742050b 100644 --- a/adapter/distribution_server.go +++ b/adapter/distribution_server.go @@ -76,15 +76,16 @@ var ( defaultCatalogReloadRetryAttempts = 20 defaultCatalogReloadRetryInterval = 10 * time.Millisecond - errDistributionCatalogNotConfigured = errors.New("route catalog is not configured") - errDistributionUnknownRoute = errors.New("unknown route") - errDistributionInvalidSplitKey = errors.New("invalid split key") - errDistributionSplitKeyAtBoundary = errors.New("split key at route boundary") - errDistributionCatalogConflict = errors.New("catalog version conflict") - errDistributionRouteIDOverflow = errors.New("route id overflow") - errDistributionNotLeader = errors.New("not leader for distribution catalog") - errDistributionCoordinatorRequired = errors.New("distribution coordinator is not configured") - errDistributionEngineNotConfigured = errors.New("distribution engine is not configured") + errDistributionCatalogNotConfigured = errors.New("route catalog is not configured") + errDistributionUnknownRoute = errors.New("unknown route") + errDistributionInvalidSplitKey = errors.New("invalid split key") + errDistributionSplitKeyAtBoundary = errors.New("split key at route boundary") + errDistributionCatalogConflict = errors.New("catalog version conflict") + errDistributionRouteIDOverflow = errors.New("route id overflow") + 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") ) // NewDistributionServer creates a new server. @@ -160,6 +161,62 @@ 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 + } + route, ok := snapshot.RouteOf(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 +} + // SplitRange splits a route into two child routes in the same raft group. func (s *DistributionServer) SplitRange(ctx context.Context, req *pb.SplitRangeRequest) (*pb.SplitRangeResponse, error) { // SplitRange performs a read-modify-write cycle across catalog and engine. @@ -578,6 +635,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 876ef40a5..a9d8da293 100644 --- a/adapter/distribution_server_test.go +++ b/adapter/distribution_server_test.go @@ -42,7 +42,12 @@ 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}, @@ -60,11 +65,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) { @@ -159,6 +190,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() diff --git a/adapter/internal.go b/adapter/internal.go index aa98f16b9..dd466144d 100644 --- a/adapter/internal.go +++ b/adapter/internal.go @@ -3,11 +3,18 @@ package adapter import ( "bytes" "context" + "os" + "strings" + "github.com/bootjp/elastickv/distribution" "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" ) type InternalOption func(*Internal) @@ -18,12 +25,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 WithInternalMigrationPromoteGate(gate func(context.Context) error) InternalOption { + return func(i *Internal) { + i.migrationPromoteGate = gate + } +} + +func WithInternalRouteEngine(engine *distribution.Engine) InternalOption { + return func(i *Internal) { + i.routeEngine = engine + } +} + 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, + migrationPromoteGate: defaultMigrationPromoteGate, } for _, opt := range opts { opt(i) @@ -32,11 +64,15 @@ func NewInternalWithEngine(txm kv.Transactional, leader raftengine.LeaderView, c } type Internal struct { - leader raftengine.LeaderView - transactionManager kv.Transactional - clock *kv.HLC - tsAllocator kv.TimestampAllocator - relay *RedisPubSubRelay + leader raftengine.LeaderView + transactionManager kv.Transactional + clock *kv.HLC + tsAllocator kv.TimestampAllocator + relay *RedisPubSubRelay + store store.MVCCStore + migrationProposer raftengine.Proposer + migrationPromoteGate func(context.Context) error + routeEngine *distribution.Engine pb.UnimplementedInternalServer } @@ -47,6 +83,12 @@ 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 +) + 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) @@ -85,6 +127,371 @@ 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 exportRangeVersionsRequestFullyUnbounded(req) { + return errors.WithStack(status.Error(codes.InvalidArgument, "migration export requires a raw or route bound")) + } + return nil +} + +func exportRangeVersionsRequestFullyUnbounded(req *pb.ExportRangeVersionsRequest) bool { + return len(req.GetRangeStart()) == 0 && + len(req.GetRangeEnd()) == 0 && + len(req.GetRouteStart()) == 0 && + len(req.GetRouteEnd()) == 0 +} + +func (i *Internal) streamExportRangeVersions(req *pb.ExportRangeVersionsRequest, stream pb.Internal_ExportRangeVersionsServer) error { + opts := 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")) + } + if err := stream.Send(&pb.ExportRangeVersionsResponse{ + Versions: protoMVCCVersionsFromStore(result.Versions), + NextCursor: result.NextCursor, + Done: result.Done, + }); err != nil { + return errors.WithStack(err) + } + if result.Done { + return nil + } + opts.Cursor = result.NextCursor + } +} + +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 + } + 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) 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 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 { + switch strings.ToLower(strings.TrimSpace(os.Getenv("ELASTICKV_ENABLE_MIGRATION_PROMOTE_OPCODE"))) { + 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 exportRangeVersionsOptions(req *pb.ExportRangeVersionsRequest) store.ExportVersionsOptions { + chunkBytes := uint64(req.GetChunkBytes()) + if chunkBytes == 0 { + chunkBytes = defaultMigrationExportChunkBytes + } + maxScannedBytes := req.GetMaxScannedBytes() + if maxScannedBytes == 0 { + maxScannedBytes = chunkBytes * defaultMigrationExportScanFactor + } + opts := store.ExportVersionsOptions{ + StartKey: req.GetRangeStart(), + EndKey: req.GetRangeEnd(), + MinCommitTSExclusive: req.GetMinCommitTs(), + MaxCommitTSInclusive: req.GetMaxCommitTs(), + Cursor: req.GetCursor(), + MaxVersions: defaultMigrationExportMaxVersions, + MaxBytes: chunkBytes, + MaxScannedBytes: maxScannedBytes, + KeyFamily: req.GetKeyFamily(), + AcceptKey: migrationExportFilter(req), + } + return opts +} + +func migrationExportFilter(req *pb.ExportRangeVersionsRequest) func([]byte) bool { + routeFilter := kv.RouteKeyFilter(req.GetRouteStart(), req.GetRouteEnd()) + if migrationFamilyRequiresDecodedS3(req.GetKeyFamily()) { + routeFilter = decodedS3BucketRouteFilter(req.GetKeyFamily(), req.GetRouteStart(), req.GetRouteEnd()) + } + excludeKnownInternal := req.GetExcludeKnownInternal() || req.GetKeyFamily() == distribution.MigrationFamilyUser + bracket := distribution.MigrationBracket{ + Family: req.GetKeyFamily(), + Start: bytes.Clone(req.GetRangeStart()), + End: bytes.Clone(req.GetRangeEnd()), + ExcludeKnownInternal: excludeKnownInternal, + ExcludePrefixes: cloneByteSlices(req.GetExcludePrefixes()), + } + return func(rawKey []byte) bool { + return bracket.ContainsRawKey(rawKey) && routeFilter(rawKey) + } +} + +func migrationFamilyRequiresDecodedS3(family uint32) bool { + return family == distribution.MigrationFamilyS3BucketMeta || + family == distribution.MigrationFamilyS3BucketGeneration +} + +func decodedS3BucketRouteFilter(family uint32, routeStart, routeEnd []byte) func([]byte) bool { + allowRawRouteMatch := !s3BucketRouteBounds(routeStart, routeEnd) + return func(rawKey []byte) bool { + bucket, ok := decodedS3BucketName(family, rawKey) + if !ok { + return false + } + if allowRawRouteMatch && kv.RouteKeyFilter(routeStart, routeEnd)(rawKey) { + return true + } + return decodedS3BucketRouteIntersects(bucket, routeStart, routeEnd) + } +} + +func s3BucketRouteBounds(routeStart, routeEnd []byte) bool { + return bytes.HasPrefix(routeStart, []byte(s3keys.RoutePrefix)) || + bytes.HasPrefix(routeEnd, []byte(s3keys.RoutePrefix)) +} + +func decodedS3BucketRouteIntersects(bucket string, routeStart, routeEnd []byte) bool { + bucketRouteStart := s3keys.RoutePrefixForBucketAnyGeneration(bucket) + return rangesIntersect(routeStart, routeEnd, bucketRouteStart, prefixScanEnd(bucketRouteStart)) +} + +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 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) error { if req == nil { return nil @@ -160,18 +567,60 @@ func (i *Internal) stampRawTimestamps(ctx context.Context, reqs []*pb.Request) e if r == nil { continue } - if r.Ts != 0 { + if r.Ts == 0 { + ts, err := i.nextTimestamp(ctx, "stampRawTimestamps") + if err != nil { + return err + } + r.Ts = ts + } + if err := i.rejectWriteTimestampFloorMutations(r.Mutations, r.Ts); err != nil { + return err + } + } + return nil +} + +func (i *Internal) rejectWriteTimestampFloorMutations(muts []*pb.Mutation, commitTS uint64) error { + if i == nil || i.routeEngine == nil || commitTS == 0 { + return nil + } + routes := i.routeEngine.Stats() + for _, mut := range muts { + if mut == nil || forwardedTxnControlMutation(mut) || (len(mut.Key) == 0 && mut.GetOp() != pb.Op_DEL_PREFIX) { continue } - ts, err := i.nextTimestamp(ctx, "stampRawTimestamps") - if err != nil { + if err := rejectMutationBelowRouteWriteFloor(routes, mut, commitTS); err != nil { return err } - r.Ts = ts } return nil } +func rejectMutationBelowRouteWriteFloor(routes []distribution.Route, mut *pb.Mutation, commitTS uint64) error { + for _, route := range routes { + if routeWriteTimestampFloorApplies(route, mut, commitTS) { + return errors.Wrapf(kv.ErrRouteWriteTimestampTooLow, "key %q commit_ts=%d floor=%d", mut.Key, commitTS, route.MinWriteTSExclusive) + } + } + return nil +} + +func forwardedTxnControlMutation(mut *pb.Mutation) bool { + return mut != nil && bytes.HasPrefix(mut.GetKey(), []byte(kv.TxnKeyPrefix)) +} + +func routeWriteTimestampFloorApplies(route distribution.Route, mut *pb.Mutation, commitTS uint64) bool { + if route.MinWriteTSExclusive == 0 || commitTS > route.MinWriteTSExclusive { + return false + } + if mut.GetOp() == pb.Op_DEL_PREFIX { + start, end := kv.RoutePrefixRange(mut.GetKey()) + return rangesIntersect(route.Start, route.End, start, end) + } + return kv.RouteKeyFilter(route.Start, route.End)(mut.GetKey()) +} + func (i *Internal) stampTxnTimestamps(ctx context.Context, reqs []*pb.Request) error { startTS := forwardedTxnStartTS(reqs) if startTS == 0 { @@ -192,7 +641,10 @@ func (i *Internal) stampTxnTimestamps(ctx context.Context, reqs []*pb.Request) e } } - return i.fillForwardedTxnCommitTS(ctx, reqs, startTS) + if err := i.fillForwardedTxnCommitTS(ctx, reqs, startTS); err != nil { + return err + } + return i.rejectWriteTimestampFloorTxnRequests(reqs) } func forwardedTxnStartTS(reqs []*pb.Request) uint64 { @@ -258,6 +710,34 @@ func (i *Internal) fillForwardedTxnCommitTS(ctx context.Context, reqs []*pb.Requ return nil } +func (i *Internal) rejectWriteTimestampFloorTxnRequests(reqs []*pb.Request) error { + for _, r := range reqs { + commitTS, ok := forwardedTxnRequestCommitTS(r) + if !ok { + continue + } + if err := i.rejectWriteTimestampFloorMutations(r.Mutations, commitTS); err != nil { + return err + } + } + return nil +} + +func forwardedTxnRequestCommitTS(r *pb.Request) (uint64, bool) { + if r == nil || (r.Phase != pb.Phase_COMMIT && r.Phase != pb.Phase_NONE) { + return 0, false + } + m, ok := forwardedTxnMetaMutation(r, []byte(kv.TxnMetaPrefix)) + if !ok { + return 0, false + } + meta, err := kv.DecodeTxnMeta(m.Value) + if err != nil || meta.CommitTS == 0 { + return 0, false + } + return meta.CommitTS, true +} + // forwardedTxnCommitTS allocates a commit timestamp for a forwarded // transaction, strictly greater than startTS. Pulled out of // fillForwardedTxnCommitTS so that function stays under cyclop. diff --git a/adapter/internal_migration_test.go b/adapter/internal_migration_test.go new file mode 100644 index 000000000..a3344b75b --- /dev/null +++ b/adapter/internal_migration_test.go @@ -0,0 +1,564 @@ +package adapter + +import ( + "context" + "encoding/binary" + "testing" + + "github.com/bootjp/elastickv/distribution" + "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 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)) +} + +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.RouteKey("bucket-b", 0, ""), + routeEnd: s3keys.RouteKey("bucket-c", 0, ""), + }, + { + name: "bucket generation", + family: distribution.MigrationFamilyS3BucketGeneration, + prefix: s3keys.BucketGenerationPrefix, + keyFor: s3keys.BucketGenerationKey, + value: []byte("generation"), + routeStart: s3keys.RouteKey("bucket-b", 0, ""), + routeEnd: s3keys.RouteKey("bucket-c", 0, ""), + }, + } { + 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 TestInternalExportRangeVersionsDecodedS3EmptyRouteEndIsUnbounded(t *testing.T) { + t.Parallel() + + ctx := context.Background() + st := store.NewMVCCStore() + internal := NewInternalWithEngine(nil, mockInternalLeader{}, nil, nil, WithInternalStore(st)) + + inRouteKey := s3keys.BucketMetaKey("bucket-z") + outRouteKey := s3keys.BucketMetaKey("bucket-a") + require.NoError(t, st.PutAt(ctx, inRouteKey, []byte("meta-z"), 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: s3keys.RouteKey("bucket-z", 0, ""), + 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: inRouteKey, CommitTs: 10, Value: []byte("meta-z"), KeyFamily: distribution.MigrationFamilyS3BucketMeta}, + }, stream.responses[0].GetVersions()) +} + +func TestInternalExportRangeVersionsIncludesS3BucketAuxiliaryForBucketRouteIntersection(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.Equal(t, []*pb.MVCCVersion{ + {Key: tc.key, CommitTs: 10, Value: tc.value, KeyFamily: tc.family}, + }, 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 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), + ) + + 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 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)) +} diff --git a/adapter/internal_test.go b/adapter/internal_test.go index 49802d20a..e4713a4fa 100644 --- a/adapter/internal_test.go +++ b/adapter/internal_test.go @@ -4,11 +4,18 @@ import ( "context" "testing" + "github.com/bootjp/elastickv/distribution" "github.com/bootjp/elastickv/kv" pb "github.com/bootjp/elastickv/proto" "github.com/stretchr/testify/require" ) +type fixedInternalTimestampAllocator uint64 + +func (a fixedInternalTimestampAllocator) Next(context.Context) (uint64, error) { + return uint64(a), nil +} + func TestStampTxnTimestamps_RejectsMaxStartTS(t *testing.T) { t.Parallel() @@ -170,3 +177,158 @@ func TestStampTxnTimestamps_UsesSingleTxnStartTS(t *testing.T) { require.NoError(t, err) require.Greater(t, meta.CommitTS, uint64(9)) } + +func TestStampRawTimestampsRejectsRouteWriteFloor(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, + MinWriteTSExclusive: 100, + }}, + })) + i := &Internal{ + tsAllocator: fixedInternalTimestampAllocator(100), + routeEngine: engine, + } + reqs := []*pb.Request{{ + Mutations: []*pb.Mutation{{Op: pb.Op_PUT, Key: []byte("k"), Value: []byte("v")}}, + }} + + err := i.stampRawTimestamps(context.Background(), reqs) + require.ErrorIs(t, err, kv.ErrRouteWriteTimestampTooLow) +} + +func TestStampRawTimestampsRejectsPreStampedRouteWriteFloor(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, + MinWriteTSExclusive: 100, + }}, + })) + i := &Internal{routeEngine: engine} + reqs := []*pb.Request{{ + Ts: 100, + Mutations: []*pb.Mutation{{Op: pb.Op_PUT, Key: []byte("k"), Value: []byte("v")}}, + }} + + err := i.stampRawTimestamps(context.Background(), reqs) + require.ErrorIs(t, err, kv.ErrRouteWriteTimestampTooLow) +} + +func TestStampRawTimestampsRejectsPreStampedDelPrefixRouteWriteFloor(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, + MinWriteTSExclusive: 0, + }, + { + RouteID: 2, + Start: []byte("m"), + End: nil, + GroupID: 2, + State: distribution.RouteStateActive, + MinWriteTSExclusive: 100, + }, + }, + })) + i := &Internal{routeEngine: engine} + reqs := []*pb.Request{{ + Ts: 100, + Mutations: []*pb.Mutation{{Op: pb.Op_DEL_PREFIX, Key: nil}}, + }} + + err := i.stampRawTimestamps(context.Background(), reqs) + require.ErrorIs(t, err, kv.ErrRouteWriteTimestampTooLow) +} + +func TestStampTxnTimestampsRejectsRouteWriteFloor(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, + MinWriteTSExclusive: 100, + }}, + })) + i := &Internal{routeEngine: engine} + reqs := []*pb.Request{{ + IsTxn: true, + Phase: pb.Phase_NONE, + Ts: 50, + Mutations: []*pb.Mutation{ + {Op: pb.Op_PUT, Key: []byte(kv.TxnMetaPrefix), Value: kv.EncodeTxnMeta(kv.TxnMeta{PrimaryKey: []byte("k"), CommitTS: 100})}, + {Op: pb.Op_PUT, Key: []byte("k"), Value: []byte("v")}, + }, + }} + + err := i.stampTxnTimestamps(context.Background(), reqs) + require.ErrorIs(t, err, kv.ErrRouteWriteTimestampTooLow) +} + +func TestStampTxnTimestampsIgnoresMetadataForRouteWriteFloor(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, + MinWriteTSExclusive: 100, + }, + { + RouteID: 2, + Start: []byte("m"), + End: nil, + GroupID: 2, + State: distribution.RouteStateActive, + }, + }, + })) + i := &Internal{routeEngine: engine} + reqs := []*pb.Request{{ + IsTxn: true, + Phase: pb.Phase_NONE, + Ts: 50, + Mutations: []*pb.Mutation{ + {Op: pb.Op_PUT, Key: []byte(kv.TxnMetaPrefix), Value: kv.EncodeTxnMeta(kv.TxnMeta{PrimaryKey: []byte("z"), CommitTS: 100})}, + {Op: pb.Op_PUT, Key: []byte("z"), Value: []byte("v")}, + }, + }} + + require.NoError(t, i.stampTxnTimestamps(context.Background(), reqs)) +} diff --git a/adapter/redis_lua_context.go b/adapter/redis_lua_context.go index 4a57051b5..447423ee4 100644 --- a/adapter/redis_lua_context.go +++ b/adapter/redis_lua_context.go @@ -478,20 +478,26 @@ func (c *luaScriptContext) keyType(key []byte) (redisValueType, error) { 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 + 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 { diff --git a/adapter/redis_lua_negative_type_cache_test.go b/adapter/redis_lua_negative_type_cache_test.go index 6253f7e30..f44aec35a 100644 --- a/adapter/redis_lua_negative_type_cache_test.go +++ b/adapter/redis_lua_negative_type_cache_test.go @@ -149,44 +149,31 @@ 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") } diff --git a/adapter/test_util.go b/adapter/test_util.go index 8fe5646fa..65fd07de8 100644 --- a/adapter/test_util.go +++ b/adapter/test_util.go @@ -641,7 +641,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 16b4ad4b8..f8449c4ae 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 @@ -376,17 +377,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 } @@ -412,17 +403,7 @@ func (e *Engine) GetIntersectingRoutes(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 b777cb0f3..b096db3db 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 TestEngineApplySnapshot_RejectsOldVersion(t *testing.T) { diff --git a/distribution/migrator.go b/distribution/migrator.go index 097bfbe2e..8a22ee61f 100644 --- a/distribution/migrator.go +++ b/distribution/migrator.go @@ -2,6 +2,7 @@ package distribution import ( "bytes" + "encoding/binary" "github.com/bootjp/elastickv/internal/s3keys" "github.com/bootjp/elastickv/store" @@ -69,6 +70,7 @@ const ( migrationDynamoGenPrefix = "!ddb|meta|gen|" migrationDynamoItemPrefix = "!ddb|item|" migrationDynamoGSIPrefix = "!ddb|gsi|" + migrationStagedDataPrefix = "!dist|migstage|" ) const ( @@ -82,6 +84,8 @@ const ( migrationSQSMsgGroupPrefix = "!sqs|msg|group|" migrationSQSMsgByAgePrefix = "!sqs|msg|byage|" migrationSQSPartitionedSuffix = "p|" + migrationStagedDataJobIDBytes = 8 + migrationStagedDataSeparator = byte('|') ) var ( @@ -270,6 +274,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 { diff --git a/distribution/migrator_export_plan_test.go b/distribution/migrator_export_plan_test.go index d87ed277d..a1e4ec893 100644 --- a/distribution/migrator_export_plan_test.go +++ b/distribution/migrator_export_plan_test.go @@ -350,6 +350,28 @@ 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.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() diff --git a/distribution/split_job_catalog.go b/distribution/split_job_catalog.go index 808ad8498..2470a1252 100644 --- a/distribution/split_job_catalog.go +++ b/distribution/split_job_catalog.go @@ -261,7 +261,10 @@ func EncodeSplitJob(job SplitJob) ([]byte, error) { if err != nil { return nil, errors.Wrapf(ErrCatalogInvalidSplitJobRecord, "marshal split job: %v", err) } - out := make([]byte, 1+len(body)) + if len(body) > math.MaxInt-1 { + return nil, errors.Wrap(ErrCatalogInvalidSplitJobRecord, "encoded split job exceeds maximum slice length") + } + out := make([]byte, len(body)+1) out[0] = catalogJobCodecVersion copy(out[1:], body) return out, nil @@ -543,9 +546,66 @@ func validateSplitJobPhases(job SplitJob) error { if !job.RetryPhase.valid() || !job.AbandonFromPhase.valid() { return errors.WithStack(ErrCatalogInvalidSplitJobPhase) } + if !job.retryPhaseValidForPhase() || !job.abandonFromPhaseValidForPhase() { + return errors.WithStack(ErrCatalogInvalidSplitJobPhase) + } return nil } +func (job SplitJob) retryPhaseValidForPhase() bool { + if job.Phase == SplitJobPhaseFailed { + return job.RetryPhase.retryable() + } + return job.RetryPhase == SplitJobPhaseNone +} + +func (job SplitJob) abandonFromPhaseValidForPhase() bool { + if job.Phase == SplitJobPhaseAbandoning { + return job.AbandonFromPhase.abandonable() + } + return job.AbandonFromPhase == SplitJobPhaseNone +} + +func (p SplitJobPhase) retryable() bool { + switch p { + case SplitJobPhaseBackfill, + SplitJobPhaseFence, + SplitJobPhaseDeltaCopy, + SplitJobPhaseCutover, + SplitJobPhaseCleanup: + return true + case SplitJobPhaseNone, + SplitJobPhasePlanned, + SplitJobPhaseDone, + SplitJobPhaseFailed, + SplitJobPhaseAbandoning, + SplitJobPhaseAbandoned: + return false + default: + return false + } +} + +func (p SplitJobPhase) abandonable() bool { + switch p { + case SplitJobPhaseBackfill, + SplitJobPhaseFence, + SplitJobPhaseDeltaCopy: + return true + case SplitJobPhaseNone, + SplitJobPhasePlanned, + SplitJobPhaseCutover, + SplitJobPhaseCleanup, + SplitJobPhaseDone, + SplitJobPhaseFailed, + SplitJobPhaseAbandoning, + SplitJobPhaseAbandoned: + return false + default: + return false + } +} + func validateSplitJobBarriers(job SplitJob) error { if !job.CutoverReadFenceState.valid() || !job.TargetStagedReadinessState.valid() { return errors.WithStack(ErrCatalogInvalidSplitJobBarrierState) @@ -554,10 +614,15 @@ func validateSplitJobBarriers(job SplitJob) error { } func validateSplitJobBracketProgress(progress []SplitJobBracketProgress) error { + seen := make(map[uint64]struct{}, len(progress)) for _, p := range progress { if !p.ExportPhase.valid() { return errors.WithStack(ErrCatalogInvalidSplitJobExportPhase) } + if _, ok := seen[p.BracketID]; ok { + return errors.WithStack(ErrCatalogInvalidSplitJobRecord) + } + seen[p.BracketID] = struct{}{} } return nil } @@ -727,7 +792,7 @@ func (s *CatalogStore) buildSplitJobPutMutations(ctx context.Context, readTS uin if jobID == math.MaxUint64 { return nil, errors.WithStack(ErrCatalogSplitJobIDOverflow) } - nextJobID, err := s.splitJobNextIDFloorAt(ctx, readTS, jobID) + nextJobID, err := s.splitJobNextIDAdvanceAt(ctx, readTS, jobID) if err != nil { return nil, err } @@ -747,7 +812,7 @@ func (s *CatalogStore) buildSplitJobPutMutations(ctx context.Context, readTS uin return mutations, nil } -func (s *CatalogStore) splitJobNextIDFloorAt(ctx context.Context, ts uint64, jobID uint64) (uint64, error) { +func (s *CatalogStore) splitJobNextIDAdvanceAt(ctx context.Context, ts uint64, jobID uint64) (uint64, error) { nextJobID, err := s.nextSplitJobIDAt(ctx, ts) if err != nil { return 0, err @@ -760,14 +825,13 @@ func (s *CatalogStore) splitJobNextIDFloorAt(ctx context.Context, ts uint64, job if err != nil { return 0, err } - if nextJobID < floor { - nextJobID = floor + if jobID+1 > floor { + floor = jobID + 1 } - floor = jobID + 1 - if nextJobID < floor { - nextJobID = floor + if nextJobID == 0 || nextJobID < floor { + return floor, nil } - return nextJobID, nil + return 0, nil } func (s *CatalogStore) applySplitJobMutations(ctx context.Context, readTS uint64, readKeys [][]byte, mutations []*store.KVPairMutation) error { diff --git a/distribution/split_job_catalog_test.go b/distribution/split_job_catalog_test.go index 0ac91a59b..d7ef808e4 100644 --- a/distribution/split_job_catalog_test.go +++ b/distribution/split_job_catalog_test.go @@ -102,6 +102,102 @@ func TestSplitJobCodecRejectsInvalidPhase(t *testing.T) { } } +func TestSplitJobCodecValidatesRestartPhases(t *testing.T) { + for _, phase := range []SplitJobPhase{ + SplitJobPhaseBackfill, + SplitJobPhaseFence, + SplitJobPhaseDeltaCopy, + SplitJobPhaseCutover, + SplitJobPhaseCleanup, + } { + job := sampleSplitJob(uint64(phase) + 10) + job.Phase = SplitJobPhaseFailed + job.RetryPhase = phase + if _, err := EncodeSplitJob(job); err != nil { + t.Fatalf("expected failed job with retry phase %v to encode: %v", phase, err) + } + } + + for _, phase := range []SplitJobPhase{ + SplitJobPhaseBackfill, + SplitJobPhaseFence, + SplitJobPhaseDeltaCopy, + } { + job := sampleSplitJob(uint64(phase) + 20) + job.Phase = SplitJobPhaseAbandoning + job.AbandonFromPhase = phase + if _, err := EncodeSplitJob(job); err != nil { + t.Fatalf("expected abandoning job from phase %v to encode: %v", phase, err) + } + } +} + +func TestSplitJobCodecRejectsInvalidRestartPhases(t *testing.T) { + for _, tc := range []struct { + name string + phase SplitJobPhase + retryPhase SplitJobPhase + abandonFromPhase SplitJobPhase + }{ + { + name: "retry phase outside failed", + phase: SplitJobPhaseDeltaCopy, + retryPhase: SplitJobPhaseBackfill, + }, + { + name: "failed without retry phase", + phase: SplitJobPhaseFailed, + retryPhase: SplitJobPhaseNone, + }, + { + name: "failed retry to planned", + phase: SplitJobPhaseFailed, + retryPhase: SplitJobPhasePlanned, + }, + { + name: "failed retry to terminal", + phase: SplitJobPhaseFailed, + retryPhase: SplitJobPhaseDone, + }, + { + name: "abandon phase outside abandoning", + phase: SplitJobPhaseDeltaCopy, + abandonFromPhase: SplitJobPhaseFence, + }, + { + name: "abandoning without source phase", + phase: SplitJobPhaseAbandoning, + abandonFromPhase: SplitJobPhaseNone, + }, + { + name: "abandoning from cutover", + phase: SplitJobPhaseAbandoning, + abandonFromPhase: SplitJobPhaseCutover, + }, + } { + t.Run(tc.name, func(t *testing.T) { + job := sampleSplitJob(2) + job.Phase = tc.phase + job.RetryPhase = tc.retryPhase + job.AbandonFromPhase = tc.abandonFromPhase + if _, err := EncodeSplitJob(job); !errors.Is(err, ErrCatalogInvalidSplitJobPhase) { + t.Fatalf("expected ErrCatalogInvalidSplitJobPhase, got %v", err) + } + }) + } +} + +func TestSplitJobCodecRejectsDuplicateBracketIDs(t *testing.T) { + job := sampleSplitJob(1) + dupe := job.BracketProgress[0] + dupe.Family++ + dupe.Cursor = []byte("other-cursor") + job.BracketProgress = append(job.BracketProgress, dupe) + if _, err := EncodeSplitJob(job); !errors.Is(err, ErrCatalogInvalidSplitJobRecord) { + t.Fatalf("expected ErrCatalogInvalidSplitJobRecord, got %v", err) + } +} + func TestSplitJobCodecRejectsUnknownProtoEnumValues(t *testing.T) { body, err := gproto.MarshalOptions{Deterministic: true}.Marshal(&pb.SplitJob{ JobId: 1, @@ -412,6 +508,36 @@ func TestCatalogStoreSaveSeedsNextSplitJobIDFromExistingJobsWhenMetaMissing(t *t } } +func TestCatalogStoreBuildSplitJobPutMutationsSkipsCurrentNextSplitJobID(t *testing.T) { + st := store.NewMVCCStore() + cs := NewCatalogStore(st) + ctx := context.Background() + + for _, jobID := range []uint64{1, 2} { + if err := cs.CreateSplitJob(ctx, sampleSplitJob(jobID)); err != nil { + t.Fatalf("create split job %d: %v", jobID, err) + } + } + updated := sampleSplitJob(1) + updated.Phase = SplitJobPhaseBackfill + updated.UpdatedAtMs++ + raw, err := EncodeSplitJob(updated) + if err != nil { + t.Fatalf("encode updated split job: %v", err) + } + + mutations, err := cs.buildSplitJobPutMutations(ctx, st.LastCommitTS(), CatalogSplitJobKey(updated.JobID), raw, updated.JobID) + if err != nil { + t.Fatalf("build split job mutations: %v", err) + } + if len(mutations) != 1 { + t.Fatalf("expected only job mutation, got %d", len(mutations)) + } + if string(mutations[0].Key) != string(CatalogSplitJobKey(updated.JobID)) { + t.Fatalf("expected job mutation, got key %q", mutations[0].Key) + } +} + func TestCatalogStoreNextSplitJobIDAt_FallsBackWhenMetaMissing(t *testing.T) { st := store.NewMVCCStore() cs := NewCatalogStore(st) @@ -461,8 +587,8 @@ func sampleSplitJob(jobID uint64) SplitJob { SplitKey: []byte("split-key"), TargetGroupID: 22, Phase: SplitJobPhaseDeltaCopy, - RetryPhase: SplitJobPhaseBackfill, - AbandonFromPhase: SplitJobPhaseFence, + RetryPhase: SplitJobPhaseNone, + AbandonFromPhase: SplitJobPhaseNone, SnapshotTS: 101, SnapshotMinAdmittedTS: 102, WriteTrackerArmed: true, diff --git a/internal/s3keys/keys_test.go b/internal/s3keys/keys_test.go index 2079fb76c..86063ad2d 100644 --- a/internal/s3keys/keys_test.go +++ b/internal/s3keys/keys_test.go @@ -29,6 +29,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 34450dab7..edc91470f 100644 --- a/kv/fsm.go +++ b/kv/fsm.go @@ -9,6 +9,7 @@ import ( "log/slog" "os" + "github.com/bootjp/elastickv/distribution" "github.com/bootjp/elastickv/internal/encryption/fsmwire" "github.com/bootjp/elastickv/internal/raftengine" "github.com/bootjp/elastickv/internal/s3keys" @@ -121,6 +122,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 @@ -273,6 +278,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 @@ -310,12 +317,11 @@ type fsmApplyResponse struct { } func (f *kvFSM) Apply(data []byte) any { - 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) @@ -356,13 +362,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] >= fsmwire.OpEncryptionMin && data[0] <= fsmwire.OpEncryptionMax: return f.applyEncryption(f.pendingApplyIdx, data[0], data[1:]), true default: @@ -394,6 +404,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) { @@ -492,17 +510,8 @@ func (f *kvFSM) handleRawRequest(ctx context.Context, r *pb.Request, commitTS ui return f.handleDelPrefix(ctx, prefix, commitTS) } - 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.assertNoConflictingTxnLock(ctx, mut.Key, nil, 0); err != nil { - return err - } + if err := f.validateRawMutationsForApply(ctx, r, commitTS); err != nil { + return err } muts, err := toStoreMutations(r.Mutations) @@ -518,6 +527,38 @@ 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) error { + bypassKeys := writeFenceBypassKeySet(r.GetWriteFenceBypassKeys()) + for _, mut := range r.GetMutations() { + if err := f.validateRawMutationForApply(ctx, mut, bypassKeys, commitTS); err != nil { + return err + } + } + return nil +} + +func (f *kvFSM) validateRawMutationForApply(ctx context.Context, mut *pb.Mutation, writeFenceBypassKeys map[string]struct{}, commitTS uint64) 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) + } + if _, bypass := writeFenceBypassKeys[string(mut.Key)]; !bypass { + if err := f.verifyRouteNotFencedForKey(mut.Key); err != nil { + return err + } + } + if err := f.verifyRouteWriteTimestampFloorForKey(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) { @@ -532,6 +573,12 @@ 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.verifyRouteNotFencedForPrefix(prefix); err != nil { + return err + } + if err := f.verifyRouteWriteTimestampFloorForPrefix(prefix, commitTS); err != nil { + return err + } if err := f.store.DeletePrefixAtRaftAt(ctx, prefix, txnCommonPrefix, commitTS, f.pendingApplyIdx); err != nil { return errors.WithStack(err) } @@ -539,6 +586,107 @@ func (f *kvFSM) handleDelPrefix(ctx context.Context, prefix []byte, commitTS uin return nil } +func (f *kvFSM) verifyRouteNotFencedForKey(key []byte) error { + if f.routes == nil { + return nil + } + snap, ok := f.routes.Current() + if !ok { + return nil + } + rkey := routeKey(key) + if snap.WriteFencedForKey(rkey) { + return errors.Wrapf(ErrRouteWriteFenced, "key %q routeKey %q", key, rkey) + } + if start, end, ok := s3BucketAuxiliaryRouteRange(key); ok && snap.WriteFencedIntersects(start, end) { + return errors.Wrapf(ErrRouteWriteFenced, "key %q route range [%q,%q)", key, start, end) + } + 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 (f *kvFSM) verifyRouteWriteTimestampFloorForKey(key []byte, commitTS uint64) error { + if f.routes == nil || commitTS == 0 { + return nil + } + snap, ok := f.routes.Current() + if !ok { + return nil + } + if start, end, ok := s3BucketAuxiliaryRouteRange(key); ok { + for _, route := range snap.IntersectingRoutes(start, end) { + if err := verifyRouteWriteTimestampFloorForRange(route, key, start, end, commitTS); err != nil { + return err + } + } + 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 (f *kvFSM) verifyRouteWriteTimestampFloorsForMutations(muts []*pb.Mutation, commitTS uint64) error { + for _, mut := range muts { + if mut == nil || len(mut.Key) == 0 || isTxnInternalKey(mut.Key) { + continue + } + if err := f.verifyRouteWriteTimestampFloorForKey(mut.Key, commitTS); err != nil { + return err + } + } + return nil +} + +func (f *kvFSM) verifyRouteWriteTimestampFloorForPrefix(prefix []byte, commitTS uint64) error { + if f.routes == nil || commitTS == 0 { + return nil + } + snap, ok := f.routes.Current() + if !ok { + 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 @@ -556,6 +704,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), @@ -953,21 +1106,24 @@ func (f *kvFSM) verifyOwnerFromSnapshot(mutations []*pb.Mutation, snap RouteSnap if isTxnInternalKey(mut.Key) { 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 { @@ -1029,7 +1185,7 @@ func (f *kvFSM) handlePrepareRequest(ctx context.Context, r *pb.Request) error { } startTS := r.Ts - uniq, err := uniqueMutations(muts) + uniq, err := f.uniqueMutationsAboveFloor(muts, startTS) if err != nil { return err } @@ -1099,7 +1255,7 @@ func (f *kvFSM) handleOnePhaseTxnRequest(ctx context.Context, r *pb.Request, com return nil } - uniq, err := uniqueMutations(muts) + uniq, err := f.uniqueMutationsAboveFloor(muts, commitTS) if err != nil { return err } @@ -1115,6 +1271,36 @@ 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) { + uniq, err := uniqueMutations(muts) + if err != nil { + return nil, err + } + if err := f.verifyRouteWriteTimestampFloorsForMutations(uniq, commitTS); err != nil { + return nil, err + } + return uniq, nil +} + +func (f *kvFSM) uniqueTxnMutationsAboveFloor(muts []*pb.Mutation, commitTS uint64) ([]*pb.Mutation, error) { + uniq, err := uniqueTxnMutations(muts) + if err != nil { + return nil, err + } + if err := f.verifyRouteWriteTimestampFloorsForMutations(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 @@ -1131,9 +1317,55 @@ func (f *kvFSM) dedupProbeOnePhase(ctx context.Context, meta TxnMeta) (bool, err if err != nil { return false, errors.WithStack(err) } + if landed { + return true, nil + } + route, ok := f.currentStagedVisibilityRouteForKey(meta.PrimaryKey) + if !ok { + return false, nil + } + stagedKey := distribution.MigrationStagedDataKey(route.MigrationJobID, meta.PrimaryKey) + landed, err = f.store.CommittedVersionAt(ctx, stagedKey, meta.PrevCommitTS) + if err != nil { + return false, errors.WithStack(err) + } return landed, nil } +func (f *kvFSM) currentStagedVisibilityRouteForKey(key []byte) (distribution.Route, bool) { + if f == nil || f.routes == nil || len(key) == 0 { + return distribution.Route{}, false + } + if _, _, ok := distribution.MigrationStagedDataKeyParts(key); ok { + return distribution.Route{}, false + } + snap, ok := f.routes.Current() + if !ok { + return distribution.Route{}, false + } + if route, ok := currentStagedVisibilityRouteForS3BucketAuxiliaryKey(snap, key, f.shardGroupID); ok { + return route, true + } + route, ok := snap.RouteOf(routeKey(key)) + if !ok || route.GroupID != f.shardGroupID || !routeHasStagedVisibility(route) { + return distribution.Route{}, false + } + return route, true +} + +func currentStagedVisibilityRouteForS3BucketAuxiliaryKey(snap RouteSnapshot, key []byte, shardGroupID uint64) (distribution.Route, bool) { + start, end, ok := s3BucketAuxiliaryRouteRange(key) + if !ok { + return distribution.Route{}, false + } + for _, route := range snap.IntersectingRoutes(start, end) { + if route.GroupID == shardGroupID && routeHasStagedVisibility(route) { + return route, true + } + } + return distribution.Route{}, false +} + func (f *kvFSM) handleCommitRequest(ctx context.Context, r *pb.Request) error { meta, muts, err := extractTxnMeta(r.Mutations) if err != nil { @@ -1154,7 +1386,7 @@ func (f *kvFSM) handleCommitRequest(ctx context.Context, r *pb.Request) error { if err != nil { return err } - uniq, err := uniqueMutations(muts) + uniq, err := f.uniqueTxnMutationsAboveFloor(muts, commitTS) if err != nil { return err } diff --git a/kv/fsm_migration_fence_test.go b/kv/fsm_migration_fence_test.go index ecb771b38..c29276bf4 100644 --- a/kv/fsm_migration_fence_test.go +++ b/kv/fsm_migration_fence_test.go @@ -7,6 +7,7 @@ import ( "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" ) @@ -21,6 +22,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() @@ -158,7 +169,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{ @@ -175,7 +186,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{ @@ -187,6 +198,55 @@ func TestFSMRejectsObservedWriteFencedS3BucketAuxiliaryPointWrite(t *testing.T) require.ErrorIs(t, err, ErrRouteWriteFenced) } +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 TestFSMRejectsCurrentWriteFencedDelPrefix(t *testing.T) { t.Parallel() @@ -295,3 +355,70 @@ 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{ + 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{ + 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, + 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, + 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) +} diff --git a/kv/fsm_migration_import.go b/kv/fsm_migration_import.go new file mode 100644 index 000000000..91dca3070 --- /dev/null +++ b/kv/fsm_migration_import.go @@ -0,0 +1,84 @@ +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" +) + +// 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 errors.WithStack(err) + } + result, err := f.store.ImportVersionsRaft(ctx, store.ImportVersionsOptions{ + JobID: req.GetJobId(), + BracketID: req.GetBracketId(), + BatchSeq: req.GetBatchSeq(), + Cursor: req.GetCursor(), + Versions: migrationStoreVersionsFromProto(req.GetJobId(), req.GetVersions()), + }) + if err != nil { + return errors.WithStack(err) + } + result.MaxImportedTS, err = f.migrationHLCFloorForApply(ctx, req, result) + if err != nil { + return errors.WithStack(err) + } + if f.hlc != nil && result.MaxImportedTS > 0 { + f.hlc.Observe(result.MaxImportedTS) + } + return result +} + +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..589b6e1b1 --- /dev/null +++ b/kv/fsm_migration_import_test.go @@ -0,0 +1,75 @@ +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" + "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) +} diff --git a/kv/fsm_migration_promote.go b/kv/fsm_migration_promote.go new file mode 100644 index 000000000..383426ae2 --- /dev/null +++ b/kv/fsm_migration_promote.go @@ -0,0 +1,98 @@ +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 +) + +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 { + maxVersions := int(req.GetMaxVersions()) + if maxVersions <= 0 { + maxVersions = defaultMigrationPromoteMaxVersions + } + maxBytes := req.GetMaxBytes() + if maxBytes == 0 { + maxBytes = defaultMigrationPromoteMaxBytes + } + maxScannedBytes := req.GetMaxScannedBytes() + if maxScannedBytes == 0 { + maxScannedBytes = defaultMigrationPromoteMaxScannedBytes + } + 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()), + } +} + +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..192b4af1a --- /dev/null +++ b/kv/fsm_migration_promote_test.go @@ -0,0 +1,96 @@ +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)) +} diff --git a/kv/fsm_onephase_dedup_test.go b/kv/fsm_onephase_dedup_test.go index aaf027c1d..b675e1079 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,70 @@ 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() + 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, + }}) + fsmIface := NewKvFSMWithHLC(st, NewHLC(), WithRouteHistory(WrapDistributionEngine(engine), 1)) + fsm, ok := fsmIface.(*kvFSM) + require.True(t, ok) + + key := []byte("list-item") + require.NoError(t, st.PutAt(ctx, distribution.MigrationStagedDataKey(9, key), []byte("v"), 20, 0)) + + require.NoError(t, applyFSMRequest(t, fsm, onePhaseReq(30, 40, 20, key, []byte("v")))) + + 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() + engine := distribution.NewEngine() + const bucket = "bucket-a" + routeStart := s3keys.RoutePrefixForBucketAnyGeneration(bucket) + applyComposed1Snapshot(t, engine, 1, []distribution.RouteDescriptor{{ + RouteID: 1, + Start: routeStart, + End: prefixScanEnd(routeStart), + GroupID: 1, + State: distribution.RouteStateActive, + StagedVisibilityActive: true, + MigrationJobID: 9, + }}) + fsmIface := NewKvFSMWithHLC(st, NewHLC(), WithRouteHistory(WrapDistributionEngine(engine), 1)) + fsm, ok := fsmIface.(*kvFSM) + require.True(t, ok) + + key := s3keys.BucketMetaKey(bucket) + require.NoError(t, st.PutAt(ctx, distribution.MigrationStagedDataKey(9, key), []byte("v"), 20, 0)) + + require.NoError(t, applyFSMRequest(t, fsm, onePhaseReq(30, 40, 20, key, []byte("v")))) + + 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/leader_routed_store.go b/kv/leader_routed_store.go index f92cc39f5..6fff18fdb 100644 --- a/kv/leader_routed_store.go +++ b/kv/leader_routed_store.go @@ -633,6 +633,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/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_router.go b/kv/shard_router.go index 4f6cd094c..4016fa452 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.stagedVisibilityRouteForS3BucketAuxiliaryKey(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,22 @@ func (s *ShardRouter) ResolveGroup(rawKey []byte) (uint64, bool) { return route.GroupID, true } +func (s *ShardRouter) stagedVisibilityRouteForS3BucketAuxiliaryKey(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 + } + for _, route := range s.engine.GetIntersectingRoutes(start, end) { + if routeHasStagedVisibility(route) { + return route, true + } + } + return distribution.Route{}, false +} + // 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 81b21e28d..a9086165b 100644 --- a/kv/shard_store.go +++ b/kv/shard_store.go @@ -26,6 +26,7 @@ type ShardStore struct { } var ErrCrossShardMutationBatchNotSupported = errors.New("cross-shard mutation batches are not supported") +var ErrExplicitGroupStagedVisibilityUnresolved = errors.New("explicit group read cannot resolve staged visibility route") // NewShardStore creates a sharded MVCC store wrapper. func NewShardStore(engine *distribution.Engine, groups map[uint64]*ShardGroup) *ShardStore { @@ -50,19 +51,19 @@ func (s *ShardStore) GetAtWithReadFence(ctx context.Context, key []byte, ts uint if groupID != 0 { return s.getGroupAtWithReadFence(ctx, groupID, key, ts, readRouteVersion) } - g, ok := s.groupForKey(key) + route, g, ok := s.routeAndGroupForKey(key) if !ok || g.Store == nil { return nil, store.ErrKeyNotFound } // Some tests use ShardStore without raft; in that case serve reads locally. if engineForGroup(g) == nil { - return s.localGetAt(ctx, g, key, ts) + return s.localGetAt(ctx, g, route, key, ts) } // Wait for a leader read fence before serving from local state. if isLinearizableRaftLeader(ctx, engineForGroup(g)) { - return s.leaderGetAt(ctx, g, key, ts) + return s.leaderGetAt(ctx, g, route, key, ts) } return s.proxyRawGet(ctx, g, key, ts, 0, readRouteVersion) } @@ -79,12 +80,16 @@ 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 + } if engineForGroup(g) == nil { - return s.localGetAt(ctx, g, key, ts) + return s.localGetAt(ctx, g, route, key, ts) } if isLinearizableRaftLeader(ctx, engineForGroup(g)) { - return s.leaderGetAt(ctx, g, key, ts) + return s.leaderGetAt(ctx, g, route, key, ts) } return s.proxyRawGet(ctx, g, key, ts, groupID, readRouteVersion) } @@ -106,16 +111,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) @@ -123,6 +131,115 @@ 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.stagedVisibilityRouteForS3BucketAuxiliaryKey(key); ok { + if route.GroupID == groupID { + return route, nil + } + return distribution.Route{}, errors.Wrapf(ErrExplicitGroupStagedVisibilityUnresolved, "group_id=%d key=%q", 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) + } + } + return fallback, nil +} + +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 + } + live, liveOK, err := latestMVCCVersionAt(ctx, g.Store, key, ts) + if err != nil { + return nil, err + } + stagedKey := distribution.MigrationStagedDataKey(route.MigrationJobID, key) + staged, stagedOK, err := latestMVCCVersionAt(ctx, g.Store, stagedKey, 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 +} + +func latestMVCCVersionAt(ctx context.Context, st store.MVCCStore, key []byte, ts uint64) (store.MVCCVersion, bool, error) { + opts := store.ExportVersionsOptions{ + StartKey: key, + EndKey: prefixScanEnd(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 { @@ -213,7 +330,7 @@ func (s *ShardStore) ScanAtWithReadFence(ctx context.Context, start []byte, end if reverse { if groupID != 0 { if routeScanBoundsPresent(routeStart, routeEnd) { - return s.scanRouteAtDirectionWithReadFence(ctx, distribution.Route{GroupID: groupID}, start, end, limit, ts, true, true, readRouteVersion, routeStart, routeEnd) + return s.scanExplicitGroupAtWithReadFence(ctx, groupID, start, end, limit, ts, true, readRouteVersion, routeStart, routeEnd) } return nil, errors.WithStack(store.ErrNotSupported) } @@ -228,7 +345,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, true, readRouteVersion, routeStart, routeEnd) + return s.scanExplicitGroupAtWithReadFence(ctx, groupID, start, end, limit, ts, false, readRouteVersion, routeStart, routeEnd) } routes, clampToRoutes := s.routesForFencedScan(start, end, routeStart, routeEnd) @@ -251,6 +368,9 @@ func (s *ShardStore) ScanAtPhysicalLimit(ctx context.Context, start []byte, end return []*store.KVPair{}, false, nil } routes, clampToRoutes := s.routesForScan(start, end) + if routesContainStagedVisibility(routes) { + return nil, true, nil + } if len(routes) != 1 || clampToRoutes { kvs, err := s.ScanAt(ctx, start, end, visibleLimit, ts) return kvs, false, err @@ -264,10 +384,53 @@ 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, true) + 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 && !dedupeByKey { + routes = dedupeRepeatedRawScanRoutes(routes) + } + return s.scanExplicitGroupRoutesAtWithReadFence(ctx, routes, start, end, limit, ts, reverse, readRouteVersion, routeStart, routeEnd, clampToRoutes, dedupeByKey) +} + +func (s *ShardStore) scanExplicitGroupRoutesAtWithReadFence(ctx context.Context, routes []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 { + 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.scanRouteAtDirectionWithReadFence(ctx, route, scanStart, scanEnd, limit, ts, reverse, true, readRouteVersion, routeStart, routeEnd) + 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 } func (s *ShardStore) ReverseScanAt(ctx context.Context, start []byte, end []byte, limit int, ts uint64) ([]*store.KVPair, error) { @@ -295,6 +458,9 @@ func (s *ShardStore) ReverseScanAtPhysicalLimit(ctx context.Context, start []byt return []*store.KVPair{}, false, nil } routes, clampToRoutes := s.routesForScan(start, end) + if routesContainStagedVisibility(routes) { + return nil, true, nil + } if len(routes) != 1 || clampToRoutes { kvs, err := s.ReverseScanAt(ctx, start, end, visibleLimit, ts) return kvs, false, err @@ -306,6 +472,9 @@ func (s *ShardStore) routesForScan(start []byte, end []byte) ([]distribution.Rou if routeStart, routeEnd, ok := s3keys.ManifestScanRouteBounds(start, end); ok { return s.engine.GetIntersectingRoutes(routeStart, routeEnd), false } + if routes, ok := s.routesForS3BucketAuxiliaryScan(start, end); ok { + return routes, false + } if isBroadLegacyListDeltaScan(start) { return s.engine.GetIntersectingRoutes(nil, nil), false } @@ -333,6 +502,159 @@ func (s *ShardStore) routesForScan(start []byte, end []byte) ([]distribution.Rou return routes, true } +func (s *ShardStore) routesForS3BucketAuxiliaryScan(start []byte, end []byte) ([]distribution.Route, bool) { + if s == nil || s.engine == nil || !s3BucketAuxiliaryScanBounds(start, end) { + return nil, false + } + routes := make([]distribution.Route, 0) + routes = append(routes, s.engine.GetIntersectingRoutes(start, end)...) + routeStart, routeEnd := s3BucketAuxiliaryScanRouteRange(start, end) + for _, route := range s.engine.GetIntersectingRoutes(routeStart, routeEnd) { + if routeHasStagedVisibility(route) { + routes = append(routes, route) + } + } + return routes, true +} + +func s3BucketAuxiliaryScanBounds(start []byte, end []byte) bool { + if !bytes.HasPrefix(start, []byte(s3keys.BucketMetaPrefix)) && + !bytes.HasPrefix(start, []byte(s3keys.BucketGenerationPrefix)) { + return false + } + 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 +} + +func dedupeRepeatedRawScanRoutes(routes []distribution.Route) []distribution.Route { + if len(routes) <= 1 { + return routes + } + stagedGroups := make(map[uint64]struct{}) + for _, route := range routes { + if routeHasStagedVisibility(route) { + stagedGroups[route.GroupID] = struct{}{} + } + } + out := make([]distribution.Route, 0, len(routes)) + seen := make(map[repeatedRawScanRouteKey]struct{}, len(routes)) + for _, route := range routes { + if _, hasStaged := stagedGroups[route.GroupID]; hasStaged && !routeHasStagedVisibility(route) { + continue + } + key := repeatedRawScanRouteKey{groupID: route.GroupID} + if routeHasStagedVisibility(route) { + key.staged = true + key.migrationJobID = route.MigrationJobID + } + if _, ok := seen[key]; ok { + continue + } + seen[key] = struct{}{} + out = append(out, route) + } + return out +} + +func routesContainStagedVisibility(routes []distribution.Route) bool { + for _, route := range routes { + if routeHasStagedVisibility(route) { + return 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 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 { + 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 + } + 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 +} + func (s *ShardStore) routesForLegacyListDeltaScan(start []byte, end []byte) ([]distribution.Route, bool) { logicalUserKey := store.ExtractLegacyListUserKeyFromDeltaScanPrefix(start) if logicalUserKey == nil { @@ -362,6 +684,10 @@ func isBroadLegacyListDeltaScan(start []byte) bool { return logicalUserKey == nil || !bytes.Equal(start, store.LegacyListMetaDeltaScanPrefix(logicalUserKey)) } +func shouldMarkRouteGroupOnScan(start []byte, explicitGroup bool, routeStart []byte, routeEnd []byte) bool { + return !explicitGroup && !routeScanBoundsPresent(routeStart, routeEnd) && isBroadLegacyListDeltaScan(start) +} + func scanRouteUserKey(start []byte) []byte { for _, extract := range scanRouteUserKeyExtractors { if userKey := extract(start); userKey != nil { @@ -407,19 +733,17 @@ func normalizedRouteScanEnd(routeEnd []byte) []byte { 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{}) routeFilterPresent := routeScanBoundsPresent(routeStart, routeEnd) + dedupeByKey := s3BucketAuxiliaryScanBounds(start, end) + if !clampToRoutes && !routeFilterPresent && !dedupeByKey { + routes = dedupeRepeatedRawScanRoutes(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.scanRouteAtDirectionWithReadFence(ctx, route, scanStart, scanEnd, limit, ts, false, false, readRouteVersion, routeStart, routeEnd) @@ -434,7 +758,7 @@ func (s *ShardStore) scanRoutesAtWithReadFence(ctx context.Context, routes []dis } continue } - out = mergeAndTrimScanResults(out, kvs, limit) + out = mergeAndTrimScanResultsWithOptions(out, kvs, limit, false, dedupeByKey) } return out, nil } @@ -452,8 +776,11 @@ func (s *ShardStore) reverseScanRoutesAtWithReadFence( routeEnd []byte, ) ([]*store.KVPair, error) { out := make([]*store.KVPair, 0) - seenGroups := make(map[uint64]struct{}) routeFilterPresent := routeScanBoundsPresent(routeStart, routeEnd) + dedupeByKey := s3BucketAuxiliaryScanBounds(start, end) + if !clampToRoutes && !routeFilterPresent && !dedupeByKey { + routes = dedupeRepeatedRawScanRoutes(routes) + } for i := len(routes) - 1; i >= 0; i-- { route := routes[i] if clampToRoutes { @@ -472,20 +799,11 @@ 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.scanRouteAtDirectionWithReadFence(ctx, route, start, end, limit, ts, true, false, readRouteVersion, routeStart, routeEnd) if err != nil { return nil, err } - out = mergeAndTrimReverseScanResults(out, kvs, limit) + out = mergeAndTrimScanResultsWithOptions(out, kvs, limit, true, dedupeByKey) } return out, nil } @@ -515,19 +833,6 @@ func (s *ShardStore) clampedReverseScanRouteAtWithReadFence( return kvs, false, nil } -func (s *ShardStore) scanRouteAtDirection( - ctx context.Context, - route distribution.Route, - start []byte, - end []byte, - limit int, - ts uint64, - reverse bool, - explicitGroup bool, -) ([]*store.KVPair, error) { - return s.scanRouteAtDirectionWithReadFence(ctx, route, start, end, limit, ts, reverse, explicitGroup, 0, nil, nil) -} - func (s *ShardStore) scanRouteAtDirectionWithReadFence( ctx context.Context, route distribution.Route, @@ -620,18 +925,19 @@ func (s *ShardStore) scanRouteAtDirectionWithReadFenceRouteFilterPage( if !ok || g == nil || g.Store == nil { return nil, nil, nil } + markRouteGroup := shouldMarkRouteGroupOnScan(start, explicitGroup, routeStart, routeEnd) if engineForGroup(g) == nil { - 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 markScanRouteGroup(filterTxnInternalKVs(kvs), route.GroupID), markScanRouteGroup(kvs, route.GroupID), nil + return markScanRouteGroup(filterTxnInternalKVs(kvs), route.GroupID, markRouteGroup), markScanRouteGroup(kvs, route.GroupID, markRouteGroup), nil } if isLinearizableRaftLeader(ctx, engineForGroup(g)) { - kvs, cursorKVs, err := s.scanRouteAtLeaderRouteFilter(ctx, g, start, end, limit, visibleLimit, ts, reverse, routeStart, routeEnd) - return markScanRouteGroup(kvs, route.GroupID), markScanRouteGroup(cursorKVs, route.GroupID), err + kvs, cursorKVs, err := s.scanRouteAtLeaderRouteFilter(ctx, g, route, start, end, limit, visibleLimit, ts, reverse, routeStart, routeEnd) + return markScanRouteGroup(kvs, route.GroupID, markRouteGroup), markScanRouteGroup(cursorKVs, route.GroupID, markRouteGroup), err } groupID := proxyScanGroupID(route, explicitGroup, routeStart, routeEnd) @@ -640,7 +946,7 @@ func (s *ShardStore) scanRouteAtDirectionWithReadFenceRouteFilterPage( return nil, nil, err } filtered := filterTxnInternalKVs(kvs) - return markScanRouteGroup(filtered, route.GroupID), markScanRouteGroup(kvs, route.GroupID), nil + return markScanRouteGroup(filtered, route.GroupID, markRouteGroup), markScanRouteGroup(kvs, route.GroupID, markRouteGroup), nil } func proxyScanGroupID(route distribution.Route, explicitGroup bool, routeStart []byte, routeEnd []byte) uint64 { @@ -667,18 +973,19 @@ func (s *ShardStore) scanRouteAtDirectionWithReadFenceOnce( if !ok || g == nil || g.Store == nil { return nil, nil } + markRouteGroup := shouldMarkRouteGroupOnScan(start, explicitGroup, routeStart, routeEnd) if engineForGroup(g) == nil { - 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, errors.WithStack(err) } - return markScanRouteGroup(filterTxnInternalKVs(kvs), route.GroupID), nil + return markScanRouteGroup(filterTxnInternalKVs(kvs), route.GroupID, markRouteGroup), nil } if isLinearizableRaftLeader(ctx, engineForGroup(g)) { - kvs, err := s.scanRouteAtLeader(ctx, g, start, end, limit, ts, reverse) - return markScanRouteGroup(kvs, route.GroupID), err + kvs, err := s.scanRouteAtLeader(ctx, g, route, start, end, limit, ts, reverse) + return markScanRouteGroup(kvs, route.GroupID, markRouteGroup), err } groupID := proxyScanGroupID(route, explicitGroup, routeStart, routeEnd) @@ -688,7 +995,7 @@ func (s *ShardStore) scanRouteAtDirectionWithReadFenceOnce( } // The leader's RawScanAt is expected to perform lock resolution and filtering // via ShardStore.ScanAt, so avoid N+1 proxy gets here. - return markScanRouteGroup(filterTxnInternalKVs(kvs), route.GroupID), nil + return markScanRouteGroup(filterTxnInternalKVs(kvs), route.GroupID, markRouteGroup), nil } const routeFilteredScanBatchMin = 128 @@ -773,6 +1080,9 @@ func minScanEnd(a []byte, b []byte) []byte { } func routeKeyInScanBounds(key []byte, routeStart []byte, routeEnd []byte) bool { + if s3BucketAuxiliaryRouteInRange(key, routeStart, routeEnd) { + return true + } key = routeKey(key) if len(routeStart) > 0 && bytes.Compare(key, routeStart) < 0 { return false @@ -802,18 +1112,25 @@ func (s *ShardStore) scanRouteAtDirectionPhysicalLimit( if !ok || g == nil || g.Store == nil { return nil, false, nil } + markRouteGroup := shouldMarkRouteGroupOnScan(start, false, nil, nil) if engineForGroup(g) == nil { + if routeHasStagedVisibility(route) { + return nil, true, nil + } kvs, limitReached, err := scanLocalPhysicalLimit(ctx, g.Store, start, end, visibleLimit, physicalLimit, ts, reverse) if err != nil { return nil, limitReached, errors.WithStack(err) } - return markScanRouteGroup(filterTxnInternalKVs(kvs), route.GroupID), limitReached, nil + return markScanRouteGroup(filterTxnInternalKVs(kvs), route.GroupID, markRouteGroup), limitReached, nil } if isLinearizableRaftLeader(ctx, engineForGroup(g)) { - kvs, limitReached, err := s.scanRouteAtLeaderPhysicalLimit(ctx, g, start, end, visibleLimit, physicalLimit, ts, reverse) - return markScanRouteGroup(kvs, route.GroupID), limitReached, err + if routeHasStagedVisibility(route) { + return nil, true, nil + } + 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 @@ -864,75 +1181,423 @@ func scanPhysicalLimitLocal( func (s *ShardStore) scanRouteLocal( ctx context.Context, g *ShardGroup, + route distribution.Route, start []byte, end []byte, limit int, ts uint64, reverse bool, ) ([]*store.KVPair, error) { + if routeHasStagedVisibility(route) { + return s.scanRouteWithStagedVisibility(ctx, g, route, start, end, limit, ts, reverse) + } 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) + kvs, err := g.Store.ScanAt(ctx, start, end, limit, ts) + return kvs, errors.WithStack(err) +} + +func (s *ShardStore) scanRouteAtLeaderPhysicalLimit( + ctx context.Context, + g *ShardGroup, + route distribution.Route, + 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) + } + lockStart, lockEnd := scanLockBoundsForKVsDirection(kvs, start, end, visibleLimit, reverse) + lockKVs, err := scanTxnLockRangeAt(ctx, g, lockStart, lockEnd, ts, visibleLimit) + if err != nil { + return nil, limitReached, err + } + resolved, err := s.resolveScanLocks(ctx, g, route, kvs, lockKVs, ts) + return resolved, limitReached, err +} + +func (s *ShardStore) scanRouteAtLeader( + ctx context.Context, + g *ShardGroup, + route distribution.Route, + start []byte, + end []byte, + limit int, + ts uint64, + 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 nil, errors.WithStack(err) + } + lockStart, lockEnd := scanLockBoundsForKVs(kvs, start, end, limit) + lockKVs, err := scanTxnLockRangeAt(ctx, g, lockStart, lockEnd, ts, limit) + if err != nil { + return nil, 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 + } + 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 { + liveKVs, err := scanVisibleCandidates(ctx, g.Store, start, end, window, ts, reverse) + if err != nil { + return nil, nil, false, err + } + stagedKVs, err := scanVisibleCandidates(ctx, g.Store, stagedStart, stagedEnd, 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 +} + +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) + } + kvs, err := st.ScanAt(ctx, start, end, limit, ts) + return kvs, errors.WithStack(err) +} + +func (s *ShardStore) latestStagedVisibilityCandidates( + ctx context.Context, + st store.MVCCStore, + route distribution.Route, + liveKVs []*store.KVPair, + stagedKVs []*store.KVPair, + ts uint64, +) (map[string]store.MVCCVersion, error) { + keys := stagedVisibilityCandidateKeys(liveKVs, stagedKVs) + out := make(map[string]store.MVCCVersion, len(keys)) + for _, key := range keys { + live, liveOK, err := latestMVCCVersionAt(ctx, st, key, ts) + if err != nil { + return nil, err + } + stagedKey := distribution.MigrationStagedDataKey(route.MigrationJobID, key) + staged, stagedOK, err := latestMVCCVersionAt(ctx, st, stagedKey, ts) + if err != nil { + return nil, err + } + 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) 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 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) } - lockStart, lockEnd := scanLockBoundsForKVsDirection(kvs, start, end, visibleLimit, reverse) - lockKVs, err := scanTxnLockRangeAt(ctx, g, lockStart, lockEnd, ts, visibleLimit) - if err != nil { - return nil, limitReached, err + for _, kvp := range stagedKVs { + if kvp == nil { + continue + } + _, rawKey, ok := distribution.MigrationStagedDataKeyParts(kvp.Key) + if !ok { + continue + } + seen[string(rawKey)] = bytes.Clone(rawKey) } - resolved, err := s.resolveScanLocks(ctx, g, kvs, lockKVs, ts) - return resolved, limitReached, err + out := make([][]byte, 0, len(seen)) + for _, key := range seen { + out = append(out, key) + } + return out } -func (s *ShardStore) scanRouteAtLeader( - ctx context.Context, - g *ShardGroup, - start []byte, - end []byte, - limit int, - ts uint64, - reverse bool, -) ([]*store.KVPair, error) { - var ( - kvs []*store.KVPair - err error - ) - if reverse { - kvs, err = g.Store.ReverseScanAt(ctx, start, end, limit, ts) - } else { - kvs, err = g.Store.ScanAt(ctx, start, end, limit, ts) +func isMigrationStagedDataKey(key []byte) bool { + _, _, ok := distribution.MigrationStagedDataKeyParts(key) + return ok +} + +func stagedVisibilityScanBounds(jobID uint64, start []byte, end []byte) ([]byte, []byte) { + prefix := distribution.MigrationStagedDataKeyPrefix(jobID) + scanStart := prefix + if start != nil { + scanStart = distribution.MigrationStagedDataKey(jobID, start) } - if err != nil { - return nil, errors.WithStack(err) + scanEnd := prefixScanEnd(prefix) + if end != nil { + scanEnd = distribution.MigrationStagedDataKey(jobID, end) } - lockStart, lockEnd := scanLockBoundsForKVs(kvs, start, end, limit) - lockKVs, err := scanTxnLockRangeAt(ctx, g, lockStart, lockEnd, ts, limit) - if err != nil { - return nil, 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), + }) } - return s.resolveScanLocks(ctx, g, kvs, lockKVs, ts) + 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, @@ -946,9 +1611,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 { @@ -960,7 +1628,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) return resolved, kvs, err } @@ -1034,34 +1702,54 @@ func scanUserKey(kvp *store.KVPair) ([]byte, bool) { 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...) - if len(out) <= limit { - return out + if dedupeByKey { + out = appendReplacingKVsByKey(out, kvs) + } else { + out = append(out, kvs...) } sort.Slice(out, func(i, j int) bool { - return bytes.Compare(out[i].Key, out[j].Key) < 0 + cmp := bytes.Compare(out[i].Key, out[j].Key) + if reverse { + return cmp > 0 + } + return cmp < 0 }) + if len(out) <= limit { + return out + } clear(out[limit:]) return out[:limit] } func mergeAndTrimReverseScanResults(out []*store.KVPair, kvs []*store.KVPair, limit int) []*store.KVPair { - if len(kvs) == 0 { - return out + 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.Slice(out, func(i, j int) bool { - return bytes.Compare(out[i].Key, out[j].Key) > 0 - }) - 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 countNonInternalKVs(kvs []*store.KVPair) int { @@ -1104,34 +1792,58 @@ func clampScanEnd(end []byte, routeEnd []byte) []byte { } func (s *ShardStore) PutAt(ctx context.Context, key []byte, value []byte, commitTS uint64, expireAt uint64) error { - g, ok := s.groupForKey(key) + 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, ok := s.groupForKey(key) + 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, ok := s.groupForKey(key) + 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, ok := s.groupForKey(key) + 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.ExpireAt(ctx, key, expireAt, commitTS)) } @@ -1140,13 +1852,13 @@ func (s *ShardStore) LatestCommitTS(ctx context.Context, key []byte) (uint64, bo } func (s *ShardStore) LatestCommitTSWithReadFence(ctx context.Context, key []byte, readRouteVersion uint64) (uint64, bool, error) { - g, ok := s.groupForKey(key) + route, g, ok := s.routeAndGroupForKey(key) 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) } @@ -1158,7 +1870,7 @@ func (s *ShardStore) LatestCommitTSWithReadFence(ctx context.Context, key []byte // 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) } @@ -1169,6 +1881,30 @@ func (s *ShardStore) LatestCommitTSWithReadFence(ctx context.Context, key []byte return s.proxyLatestCommitTS(ctx, g, key, readRouteVersion) } +func (s *ShardStore) localLatestCommitTS(ctx context.Context, g *ShardGroup, route distribution.Route, key []byte) (uint64, bool, error) { + liveTS, liveExists, err := g.Store.LatestCommitTS(ctx, key) + if err != nil { + return 0, false, errors.WithStack(err) + } + if !routeHasStagedVisibility(route) { + return liveTS, liveExists, nil + } + stagedTS, stagedExists, err := g.Store.LatestCommitTS(ctx, distribution.MigrationStagedDataKey(route.MigrationJobID, 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 { @@ -1294,7 +2030,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 } @@ -1309,7 +2045,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) { @@ -1607,7 +2343,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 { @@ -1617,7 +2353,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 @@ -1646,8 +2382,8 @@ func filterTxnInternalKVs(kvs []*store.KVPair) []*store.KVPair { return out } -func markScanRouteGroup(kvs []*store.KVPair, groupID uint64) []*store.KVPair { - if groupID == 0 { +func markScanRouteGroup(kvs []*store.KVPair, groupID uint64, mark bool) []*store.KVPair { + if !mark || groupID == 0 { return kvs } for _, kvp := range kvs { @@ -1838,6 +2574,11 @@ func (s *ShardStore) ApplyMutations(ctx context.Context, mutations []*store.KVPa 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)) } @@ -1848,6 +2589,11 @@ func (s *ShardStore) ApplyMutationsRaft(ctx context.Context, mutations []*store. 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)) } @@ -1859,9 +2605,105 @@ func (s *ShardStore) ApplyMutationsRaftAt(ctx context.Context, mutations []*stor 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 ensureRouteWriteTimestampFloor(route distribution.Route, key []byte, commitTS uint64) error { + if route.MinWriteTSExclusive == 0 || commitTS == 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 (s *ShardStore) ensureMutationWriteTimestampFloors(mutations []*store.KVPairMutation, commitTS uint64) error { + if commitTS == 0 { + return nil + } + 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 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 + } + for _, key := range readKeys { + readKeys = s.appendStagedVisibilityAlias(group, readKeys, key) + } + return readKeys +} + +func (s *ShardStore) readKeysWithStagedVisibilityMutationAliases(group *ShardGroup, readKeys [][]byte, mutations []*store.KVPairMutation) [][]byte { + for _, mut := range mutations { + if mut == nil { + continue + } + readKeys = s.appendStagedVisibilityAlias(group, readKeys, mut.Key) + } + return readKeys +} + +func (s *ShardStore) appendStagedVisibilityAlias(group *ShardGroup, readKeys [][]byte, key []byte) [][]byte { + if s == nil || s.engine == nil || group == nil { + return readKeys + } + alias, ok := s.stagedVisibilityReadKeyAlias(group, key) + if !ok { + return readKeys + } + out := append([][]byte(nil), readKeys...) + return append(out, alias) +} + +func (s *ShardStore) stagedVisibilityReadKeyAlias(group *ShardGroup, key []byte) ([]byte, bool) { + if 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 +} + // 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 @@ -1888,6 +2730,9 @@ func (s *ShardStore) resolveSingleShardGroup(mutations []*store.KVPairMutation) // 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.ensurePrefixWriteTimestampFloors(prefix, commitTS); err != nil { + return err + } for _, g := range s.groups { if g == nil || g.Store == nil { continue @@ -1896,11 +2741,69 @@ 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.ensurePrefixWriteTimestampFloors(prefix, commitTS); err != nil { + return err + } for _, g := range s.groups { if g == nil || g.Store == nil { continue @@ -1909,6 +2812,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 } @@ -1927,6 +2835,9 @@ 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.ensurePrefixWriteTimestampFloors(prefix, commitTS); err != nil { + return err + } for _, g := range s.groups { if g == nil || g.Store == nil { continue @@ -1945,6 +2856,11 @@ func (s *ShardStore) DeletePrefixAtRaftAt(ctx context.Context, prefix []byte, ex return errors.WithStack(err) } } + for _, del := range s.stagedVisibilityPrefixDeletes(prefix, excludePrefix) { + if err := del.group.Store.DeletePrefixAtRaftAt(ctx, del.prefix, del.excludePrefix, commitTS, appliedIndex); err != nil { + return errors.WithStack(err) + } + } return nil } @@ -2084,6 +3000,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 } @@ -2125,12 +3045,37 @@ func (s *ShardStore) closeGroup(g *ShardGroup) error { } func (s *ShardStore) groupForKey(key []byte) (*ShardGroup, bool) { + _, g, ok := s.routeAndGroupForKey(key) + return g, ok +} + +func (s *ShardStore) routeAndGroupForKey(key []byte) (distribution.Route, *ShardGroup, bool) { + if route, ok := s.stagedVisibilityRouteForS3BucketAuxiliaryKey(key); ok { + g, ok := s.groups[route.GroupID] + return route, g, ok + } route, ok := s.engine.GetRoute(routeKey(key)) if !ok { - return nil, false + return distribution.Route{}, nil, false } g, ok := s.groups[route.GroupID] - return g, ok + return route, g, ok +} + +func (s *ShardStore) stagedVisibilityRouteForS3BucketAuxiliaryKey(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 + } + for _, route := range s.engine.GetIntersectingRoutes(start, end) { + if routeHasStagedVisibility(route) { + return route, true + } + } + return distribution.Route{}, false } func (s *ShardStore) proxyRawGet(ctx context.Context, g *ShardGroup, key []byte, ts uint64, groupID uint64, readRouteVersion uint64) ([]byte, error) { diff --git a/kv/shard_store_test.go b/kv/shard_store_test.go index 33cccaad0..64c01751d 100644 --- a/kv/shard_store_test.go +++ b/kv/shard_store_test.go @@ -1,7 +1,9 @@ package kv import ( + "bytes" "context" + "fmt" "testing" "github.com/bootjp/elastickv/distribution" @@ -11,6 +13,750 @@ import ( "github.com/stretchr/testify/require" ) +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 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 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 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_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 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) + + 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) +} + +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{ + {RouteID: 1, Start: []byte("a"), End: []byte("m"), GroupID: 2, State: distribution.RouteStateActive}, + { + RouteID: 2, + Start: []byte("m"), + End: []byte("z"), + GroupID: 1, + 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 TestShardStorePhysicalLimitFailsClosedBeforeStagedVisibilityFallback(t *testing.T) { + t.Parallel() + + ctx := context.Background() + st, _ := newStagedVisibilityShardStore(t) + + kvs, limitReached, err := st.ScanAtPhysicalLimit(ctx, []byte("b"), []byte("c"), 10, 10, 50) + require.NoError(t, err) + require.True(t, limitReached) + require.Nil(t, kvs) + + kvs, limitReached, err = st.ReverseScanAtPhysicalLimit(ctx, []byte("b"), []byte("c"), 10, 10, 50) + require.NoError(t, err) + require.True(t, limitReached) + require.Nil(t, 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) +} + func TestShardStoreScanAt_IncludesListKeysAcrossShards(t *testing.T) { t.Parallel() @@ -190,6 +936,54 @@ func TestShardStoreScanGroupAt_UsesExplicitGroup(t *testing.T) { require.Equal(t, []byte("msg-2"), kvs[0].Value) } +func TestShardStoreScanGroupAt_DoesNotClampRouteMappedRawBounds(t *testing.T) { + t.Parallel() + + ctx := context.Background() + engine := distribution.NewEngine() + engine.UpdateRoute([]byte(""), nil, 42) + groups := map[uint64]*ShardGroup{ + 42: {Store: store.NewMVCCStore()}, + } + st := NewShardStore(engine, groups) + + 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)) + + 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 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: sqsGlobalRouteKey, End: split, GroupID: 42, State: distribution.RouteStateActive}, + {RouteID: 2, Start: split, End: routeEnd, GroupID: 42, State: distribution.RouteStateActive}, + }, + })) + groups := map[uint64]*ShardGroup{ + 42: {Store: store.NewMVCCStore()}, + } + st := NewShardStore(engine, groups) + + 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)) + + 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 TestShardStoreGetGroupAt_UsesExplicitGroup(t *testing.T) { t.Parallel() diff --git a/kv/sharded_coordinator.go b/kv/sharded_coordinator.go index 056aa63ab..a16d1e3e5 100644 --- a/kv/sharded_coordinator.go +++ b/kv/sharded_coordinator.go @@ -1084,6 +1084,9 @@ func (c *ShardedCoordinator) dispatchDelPrefixBroadcast(ctx context.Context, isT 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 { requests = append(requests, &pb.Request{ @@ -1197,6 +1200,56 @@ func (c *ShardedCoordinator) rejectWriteFencedDelPrefixes(elems []*Elem[OP]) err return nil } +func (c *ShardedCoordinator) rejectWriteTimestampFloorPointElems(elems []*Elem[OP], commitTS uint64) error { + if c == nil || c.engine == nil || commitTS == 0 { + return nil + } + for _, elem := range elems { + if elem == nil || len(elem.Key) == 0 { + continue + } + if err := c.rejectWriteTimestampFloorPointKey(elem.Key, commitTS); err != nil { + return err + } + } + return nil +} + +func (c *ShardedCoordinator) rejectWriteTimestampFloorPointKey(key []byte, commitTS uint64) error { + start, end, ok := s3BucketAuxiliaryRouteRange(key) + if ok { + for _, route := range c.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 + } + rkey := routeKey(key) + if route, ok := c.engine.GetRoute(rkey); ok && route.MinWriteTSExclusive != 0 && commitTS <= route.MinWriteTSExclusive { + return errors.Wrapf(ErrRouteWriteTimestampTooLow, "key %q routeKey %q commit_ts=%d floor=%d", key, rkey, commitTS, route.MinWriteTSExclusive) + } + 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 + } + start, end := routePrefixRange(elem.Key) + for _, route := range c.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", elem.Key, start, end, commitTS, route.MinWriteTSExclusive) + } + } + } + return nil +} + // broadcastToAllGroups sends the same set of requests to every shard group in // parallel and returns the maximum commit index. func (c *ShardedCoordinator) broadcastToAllGroups(ctx context.Context, requests []*pb.Request) (*CoordinateResponse, error) { @@ -1254,6 +1307,9 @@ func (c *ShardedCoordinator) dispatchTxn(ctx context.Context, startTS uint64, co if err != nil { return nil, err } + if err := c.rejectWriteTimestampFloorPointElems(elems, commitTS); err != nil { + return nil, err + } if len(gids) == 1 && c.allReadKeysInShard(readKeys, gids[0]) { // Fast path: all mutations and read keys are in a single shard. @@ -1294,6 +1350,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) if err != nil { return nil, err @@ -1368,6 +1428,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 @@ -1384,6 +1449,48 @@ func (c *ShardedCoordinator) dispatchSingleShardTxn(ctx context.Context, startTS return &CoordinateResponse{CommitIndex: resp.CommitIndex}, 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 @@ -1941,7 +2048,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 } @@ -1949,21 +2056,48 @@ 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.stagedVisibilityRouteForS3BucketAuxiliaryKey(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) stagedVisibilityRouteForS3BucketAuxiliaryKey(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 + } + for _, route := range c.engine.GetIntersectingRoutes(start, end) { + if routeHasStagedVisibility(route) { + return route, true + } + } + return distribution.Route{}, false +} + // 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 @@ -1995,7 +2129,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-"+ @@ -2003,10 +2137,61 @@ 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.stagedVisibilityRouteForS3BucketAuxiliaryKey(key); ok { + if route.GroupID != gid { + return nil, false + } + return distribution.MigrationStagedDataKey(route.MigrationJobID, key), true + } + route, ok := c.engine.GetRoute(routeKey(key)) + if !ok || 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. @@ -2056,7 +2241,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) } @@ -2067,6 +2252,43 @@ 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) { + liveTS, liveExists, err := g.Store.LatestCommitTS(ctx, key) + if err != nil { + return 0, false, errors.WithStack(err) + } + route, ok := c.stagedVisibilityRouteForReadKey(gid, key) + if !ok { + return liveTS, liveExists, nil + } + stagedTS, stagedExists, err := g.Store.LatestCommitTS(ctx, distribution.MigrationStagedDataKey(route.MigrationJobID, 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 + } + 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 { @@ -2100,6 +2322,9 @@ func (c *ShardedCoordinator) rawLogs(ctx context.Context, reqs *OperationGroup[O if err != nil { return nil, err } + if err := c.rejectWriteTimestampFloorMutations(grouped[gid], ts); err != nil { + return nil, err + } logs = append(logs, &pb.Request{ IsTxn: false, Phase: pb.Phase_NONE, @@ -2112,6 +2337,21 @@ func (c *ShardedCoordinator) rawLogs(ctx context.Context, reqs *OperationGroup[O return logs, 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 mut == nil || len(mut.Key) == 0 { + continue + } + if err := c.rejectWriteTimestampFloorPointKey(mut.Key, commitTS); err != nil { + return err + } + } + return nil +} + func (c *ShardedCoordinator) rawLogTimestamp(ctx context.Context) (uint64, error) { if c.tsAllocator != nil { return 0, nil @@ -2129,6 +2369,9 @@ func (c *ShardedCoordinator) stampRawRequestTimestamps(ctx context.Context, reqs return err } r.Ts = ts + if err := c.rejectWriteTimestampFloorMutations(r.Mutations, r.Ts); err != nil { + return err + } } return nil } @@ -2206,21 +2449,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 cd314eef6..adebba3bc 100644 --- a/kv/sharded_coordinator_del_prefix_test.go +++ b/kv/sharded_coordinator_del_prefix_test.go @@ -122,6 +122,56 @@ func TestShardedCoordinator_DelPrefixBroadcastsToAllGroups(t *testing.T) { "same DEL_PREFIX element must use the same timestamp across shards") } +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 TestShardedCoordinator_DelPrefixDoesNotAutoPinObservedRouteVersion(t *testing.T) { t.Parallel() @@ -354,6 +404,118 @@ func TestShardedCoordinatorRejectsS3BucketAuxiliaryPointWriteOnWriteFencedRoute( } } +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 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.stagedVisibilityRouteForS3BucketAuxiliaryKey(key) + require.True(t, ok) + require.Equal(t, uint64(2), route.GroupID) + + _, 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 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 TestShardedCoordinatorRejectsDelPrefixIntersectingWriteFencedRoute(t *testing.T) { t.Parallel() diff --git a/kv/sharded_coordinator_txn_test.go b/kv/sharded_coordinator_txn_test.go index d47871e91..d58820d5a 100644 --- a/kv/sharded_coordinator_txn_test.go +++ b/kv/sharded_coordinator_txn_test.go @@ -8,6 +8,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" @@ -48,6 +49,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 @@ -271,6 +424,26 @@ func TestShardedCoordinatorDispatchTxn_UsesProvidedCommitTS(t *testing.T) { require.Equal(t, commitTS, commitMeta2.CommitTS) } +func TestShardedCoordinatorDispatchTxn_RejectsMigrationTimestampFloor(t *testing.T) { + t.Parallel() + + g1Txn := &recordingTransactional{} + coord := NewShardedCoordinator(newMigrationFloorEngine(t, 100), map[uint64]*ShardGroup{ + 1: {Txn: g1Txn}, + }, 1, NewHLC(), nil) + + _, err := coord.Dispatch(context.Background(), &OperationGroup[OP]{ + IsTxn: true, + StartTS: 90, + CommitTS: 100, + 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 preparing a floor-violating txn") +} + func TestCommitSecondaryWithRetry_RetriesAndSucceeds(t *testing.T) { t.Parallel() @@ -400,6 +573,30 @@ 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]) +} + // --------------------------------------------------------------------------- // validateReadOnlyShards // --------------------------------------------------------------------------- diff --git a/kv/tso_test.go b/kv/tso_test.go index 8daa8651d..fa256fb83 100644 --- a/kv/tso_test.go +++ b/kv/tso_test.go @@ -284,6 +284,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 734436134..4a26e67c9 100644 --- a/main.go +++ b/main.go @@ -1582,6 +1582,7 @@ func prepareRuntimeServerRunner(waitRotateOnStartup startupRotationWaiter, in se shardStore: in.shardStore, coordinate: adapterCoordinate, distServer: in.distServer, + routeEngine: in.cfg.engine, adminServer: adminServer, adminGRPCOpts: adminGRPCOpts, redisAddress: *redisAddr, @@ -2190,6 +2191,7 @@ func startRaftServers( coordinate kv.Coordinator, readGate func() bool, distServer *adapter.DistributionServer, + routeEngine *distribution.Engine, relay *adapter.RedisPubSubRelay, proposalObserverForGroup func(uint64) kv.ProposalObserver, adminServer *adapter.AdminServer, @@ -2230,7 +2232,12 @@ func startRaftServers( rt.engine, coordinate.Clock(), relay, - internalTimestampOptions(coordinate)..., + append( + internalTimestampOptions(coordinate), + adapter.WithInternalStore(rt.store), + adapter.WithInternalMigrationProposer(proposerForGroup(rt, shardGroups)), + adapter.WithInternalRouteEngine(routeEngine), + )..., )) pb.RegisterDistributionServer(gs, distServer) if adminServer != nil { @@ -2566,6 +2573,7 @@ type runtimeServerRunner struct { shardStore *kv.ShardStore coordinate kv.Coordinator distServer *adapter.DistributionServer + routeEngine *distribution.Engine adminServer *adapter.AdminServer adminGRPCOpts adminGRPCInterceptors redisAddress string @@ -2682,6 +2690,7 @@ func (r *runtimeServerRunner) startRaftTransport() error { r.coordinate, readGate, r.distServer, + r.routeEngine, r.pubsubRelay, func(groupID uint64) kv.ProposalObserver { return r.metricsRegistry.RaftProposalObserver(groupID) diff --git a/main_bootstrap_e2e_test.go b/main_bootstrap_e2e_test.go index d962c8105..f4483c137 100644 --- a/main_bootstrap_e2e_test.go +++ b/main_bootstrap_e2e_test.go @@ -15,6 +15,7 @@ import ( "time" "github.com/bootjp/elastickv/adapter" + "github.com/bootjp/elastickv/distribution" internalraftadmin "github.com/bootjp/elastickv/internal/raftadmin" "github.com/bootjp/elastickv/internal/raftengine" "github.com/bootjp/elastickv/kv" @@ -577,6 +578,7 @@ func startBootstrapE2ENode( shardStore, coordinate, distServer, + cfg.engine, cfg.leaderRedis, listeners, ) @@ -660,6 +662,7 @@ func startBootstrapE2EMultiGroupNode( shardStore, coordinate, distServer, + cfg.engine, cfg.leaderRedis, listeners, ) @@ -688,6 +691,7 @@ func startRuntimeServersWithBoundListeners( shardStore *kv.ShardStore, coordinate kv.Coordinator, distServer *adapter.DistributionServer, + routeEngine *distribution.Engine, leaderRedis map[string]string, listeners bootstrapE2EListeners, ) error { @@ -701,7 +705,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, routeEngine, relay, listeners.grpc); err != nil { return waitErrgroupAfterStartupFailure(cancel, eg, err) } if err := startBoundDynamoDBServer(ctx, eg, listeners.dynamo, shardStore, coordinate); err != nil { @@ -718,6 +722,7 @@ func startRuntimeServersWithBoundMultiGroupListeners( shardStore *kv.ShardStore, coordinate kv.Coordinator, distServer *adapter.DistributionServer, + routeEngine *distribution.Engine, leaderRedis map[string]string, listeners bootstrapE2EMultiGroupListeners, ) error { @@ -730,7 +735,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, routeEngine, relay, listeners.grpc[rt.spec.id]); err != nil { return waitErrgroupAfterStartupFailure(cancel, eg, err) } } @@ -747,6 +752,7 @@ func startBoundGRPCServer( shardStore *kv.ShardStore, coordinate kv.Coordinator, distServer *adapter.DistributionServer, + routeEngine *distribution.Engine, relay *adapter.RedisPubSubRelay, listener net.Listener, ) error { @@ -762,7 +768,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.WithInternalRouteEngine(routeEngine), + )) 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 7490809ce..969ebcbc8 100644 --- a/proto/internal.pb.go +++ b/proto/internal.pb.go @@ -552,6 +552,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 } @@ -649,6 +658,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"` @@ -913,6 +943,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[12] + 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[12] + 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{12} +} + +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[13] + 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[13] + 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{13} +} + +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 = "" + @@ -942,7 +1116,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" + @@ -955,7 +1129,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" + @@ -977,7 +1156,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" + @@ -988,12 +1179,13 @@ const file_internal_proto_rawDesc = "" + "\aPREPARE\x10\x01\x12\n" + "\n" + "\x06COMMIT\x10\x02\x12\t\n" + - "\x05ABORT\x10\x032\xa3\x02\n" + + "\x05ABORT\x10\x032\xfd\x02\n" + "\bInternal\x12.\n" + "\aForward\x12\x0f.ForwardRequest\x1a\x10.ForwardResponse\"\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 @@ -1008,22 +1200,24 @@ func file_internal_proto_rawDescGZIP() []byte { } var file_internal_proto_enumTypes = make([]protoimpl.EnumInfo, 2) -var file_internal_proto_msgTypes = make([]protoimpl.MessageInfo, 12) +var file_internal_proto_msgTypes = make([]protoimpl.MessageInfo, 14) 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 - (*RelayPublishRequest)(nil), // 7: RelayPublishRequest - (*RelayPublishResponse)(nil), // 8: RelayPublishResponse - (*ExportRangeVersionsRequest)(nil), // 9: ExportRangeVersionsRequest - (*ExportRangeVersionsResponse)(nil), // 10: ExportRangeVersionsResponse - (*MVCCVersion)(nil), // 11: MVCCVersion - (*ImportRangeVersionsRequest)(nil), // 12: ImportRangeVersionsRequest - (*ImportRangeVersionsResponse)(nil), // 13: 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 + (*RelayPublishRequest)(nil), // 7: RelayPublishRequest + (*RelayPublishResponse)(nil), // 8: RelayPublishResponse + (*ExportRangeVersionsRequest)(nil), // 9: ExportRangeVersionsRequest + (*ExportRangeVersionsResponse)(nil), // 10: ExportRangeVersionsResponse + (*MVCCVersion)(nil), // 11: MVCCVersion + (*ImportRangeVersionsRequest)(nil), // 12: ImportRangeVersionsRequest + (*ImportRangeVersionsResponse)(nil), // 13: ImportRangeVersionsResponse + (*PromoteStagedVersionsRequest)(nil), // 14: PromoteStagedVersionsRequest + (*PromoteStagedVersionsResponse)(nil), // 15: PromoteStagedVersionsResponse } var file_internal_proto_depIdxs = []int32{ 0, // 0: Mutation.op:type_name -> Op @@ -1037,12 +1231,14 @@ var file_internal_proto_depIdxs = []int32{ 7, // 8: Internal.RelayPublish:input_type -> RelayPublishRequest 9, // 9: Internal.ExportRangeVersions:input_type -> ExportRangeVersionsRequest 12, // 10: Internal.ImportRangeVersions:input_type -> ImportRangeVersionsRequest - 6, // 11: Internal.Forward:output_type -> ForwardResponse - 8, // 12: Internal.RelayPublish:output_type -> RelayPublishResponse - 10, // 13: Internal.ExportRangeVersions:output_type -> ExportRangeVersionsResponse - 13, // 14: Internal.ImportRangeVersions:output_type -> ImportRangeVersionsResponse - 11, // [11:15] is the sub-list for method output_type - 7, // [7:11] is the sub-list for method input_type + 14, // 11: Internal.PromoteStagedVersions:input_type -> PromoteStagedVersionsRequest + 6, // 12: Internal.Forward:output_type -> ForwardResponse + 8, // 13: Internal.RelayPublish:output_type -> RelayPublishResponse + 10, // 14: Internal.ExportRangeVersions:output_type -> ExportRangeVersionsResponse + 13, // 15: Internal.ImportRangeVersions:output_type -> ImportRangeVersionsResponse + 15, // 16: Internal.PromoteStagedVersions:output_type -> PromoteStagedVersionsResponse + 12, // [12:17] is the sub-list for method output_type + 7, // [7:12] 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 @@ -1059,7 +1255,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: 12, + NumMessages: 14, NumExtensions: 0, NumServices: 1, }, diff --git a/proto/internal.proto b/proto/internal.proto index f58257cec..c9d121b1d 100644 --- a/proto/internal.proto +++ b/proto/internal.proto @@ -9,6 +9,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. @@ -96,6 +97,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 { @@ -124,3 +134,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 6a21b9eba..ba679ed80 100644 --- a/proto/internal_grpc.pb.go +++ b/proto/internal_grpc.pb.go @@ -19,10 +19,11 @@ import ( const _ = grpc.SupportPackageIsVersion9 const ( - Internal_Forward_FullMethodName = "/Internal/Forward" - Internal_RelayPublish_FullMethodName = "/Internal/RelayPublish" - Internal_ExportRangeVersions_FullMethodName = "/Internal/ExportRangeVersions" - Internal_ImportRangeVersions_FullMethodName = "/Internal/ImportRangeVersions" + Internal_Forward_FullMethodName = "/Internal/Forward" + 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. @@ -34,6 +35,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 { @@ -93,6 +95,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. @@ -102,6 +114,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() } @@ -124,6 +137,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() {} @@ -210,6 +226,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) @@ -229,6 +263,10 @@ var Internal_ServiceDesc = grpc.ServiceDesc{ MethodName: "ImportRangeVersions", Handler: _Internal_ImportRangeVersions_Handler, }, + { + MethodName: "PromoteStagedVersions", + Handler: _Internal_PromoteStagedVersions_Handler, + }, }, Streams: []grpc.StreamDesc{ { diff --git a/store/lsm_migration.go b/store/lsm_migration.go index 5a6d22f55..96ad1e082 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{Done: true}, nil } - s.dbMu.RLock() - defer s.dbMu.RUnlock() - iter, err := s.db.NewIter(pebbleExportIterOptions(opts)) if err != nil { return ExportVersionsResult{}, errors.WithStack(err) @@ -299,6 +303,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() @@ -314,7 +326,7 @@ func (s *pebbleStore) ImportVersions(ctx context.Context, opts ImportVersionsOpt } 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", @@ -347,10 +359,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{ @@ -364,7 +376,7 @@ func (s *pebbleStore) commitPebbleImportBatch(opts ImportVersionsOptions, batchM return err } defer unlock() - if err := batch.Commit(s.directApplyWriteOpts()); err != nil { + if err := batch.Commit(writeOpts); err != nil { return errors.WithStack(err) } if batchMax > 0 { @@ -392,7 +404,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 { @@ -402,7 +414,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 } diff --git a/store/lsm_store_applied_index_test.go b/store/lsm_store_applied_index_test.go index ff00ea950..4d4ecd4da 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" "os" "testing" @@ -96,6 +97,64 @@ 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) +} + // TestApplyMutationsRaftAt_ZeroIndexLeavesMetaKey covers the // appliedIndex==0 escape hatch — callers without a raft entry index // (test fakes, legacy ApplyMutationsRaft) MUST NOT bump the meta key. 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/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 6eb5f65db..be764e360 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 @@ -144,7 +204,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 { @@ -249,6 +310,61 @@ 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") @@ -517,6 +633,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() diff --git a/store/migration_versions_test.go b/store/migration_versions_test.go index 9035dd692..90b9ee472 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" ) @@ -404,6 +405,72 @@ func TestExportVersionsRejectsCursorOutsideRequestedRange(t *testing.T) { }) } +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 TestExportVersionsDoesNotTreatMigrationPrefixUserKeyAsMetadata(t *testing.T) { runMigrationStoreSuite(t, func(t *testing.T, st MVCCStore) { ctx := context.Background() @@ -817,6 +884,279 @@ func TestImportVersionsIdempotencyAndMetadata(t *testing.T) { }) } +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 TestPebbleImportMetadataPersistsAcrossReopen(t *testing.T) { ctx := context.Background() dir, err := os.MkdirTemp("", "migration-import-persist-*") diff --git a/store/mvcc_store.go b/store/mvcc_store.go index 8c8712eb9..88817d94d 100644 --- a/store/mvcc_store.go +++ b/store/mvcc_store.go @@ -60,13 +60,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 @@ -113,9 +114,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) @@ -936,6 +938,7 @@ func (s *mvccStore) restoreStreamingSnapshot(r io.Reader) error { s.minRetainedTS = minRetainedTS s.migrationAcks = make(map[migrationAckID]migrationImportAck) s.migrationHLCFloors = make(map[uint64]uint64) + s.migrationPromotions = make(map[uint64]PromotionState) return nil } diff --git a/store/mvcc_store_snapshot_test.go b/store/mvcc_store_snapshot_test.go index b16cbe68b..3032a1cc3 100644 --- a/store/mvcc_store_snapshot_test.go +++ b/store/mvcc_store_snapshot_test.go @@ -61,6 +61,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 +93,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, diff --git a/store/store.go b/store/store.go index bc80d6293..41b660bbf 100644 --- a/store/store.go +++ b/store/store.go @@ -121,6 +121,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 @@ -270,6 +317,9 @@ 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. + 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)