From ac75ab1dd4ecd668f596b227b0a5368320ef01be Mon Sep 17 00:00:00 2001 From: bootjp Date: Sun, 19 Jul 2026 20:26:43 +0900 Subject: [PATCH 01/58] migration: add range version RPC handlers --- adapter/distribution_server.go | 88 +- adapter/distribution_server_test.go | 157 +- adapter/internal.go | 593 ++++++- adapter/internal_migration_test.go | 661 ++++++++ adapter/internal_test.go | 188 +++ adapter/redis_lua_context.go | 26 +- adapter/redis_lua_negative_type_cache_test.go | 35 +- adapter/test_util.go | 9 +- distribution/engine.go | 31 +- distribution/engine_test.go | 24 +- distribution/migrator.go | 36 + distribution/migrator_export_plan_test.go | 22 + internal/s3keys/keys_test.go | 8 + kv/fsm.go | 284 +++- kv/fsm_migration_fence_test.go | 145 +- kv/fsm_migration_import.go | 85 + kv/fsm_migration_import_test.go | 112 ++ kv/fsm_migration_promote.go | 98 ++ kv/fsm_migration_promote_test.go | 96 ++ kv/fsm_onephase_dedup_test.go | 66 + kv/leader_routed_store.go | 8 + kv/route_history.go | 8 + kv/shard_router.go | 19 + kv/shard_store.go | 1484 +++++++++++++++-- kv/shard_store_test.go | 979 ++++++++++- kv/sharded_coordinator.go | 290 +++- kv/sharded_coordinator_del_prefix_test.go | 178 ++ kv/sharded_coordinator_txn_test.go | 197 +++ kv/tso_test.go | 15 + main.go | 18 +- main_bootstrap_e2e_test.go | 22 +- proto/internal.pb.go | 250 ++- proto/internal.proto | 25 + proto/internal_grpc.pb.go | 46 +- store/lsm_migration.go | 55 +- store/lsm_store_applied_index_test.go | 113 ++ store/lsm_store_registration_gate_test.go | 82 + store/lsm_store_sync_mode_test.go | 60 + store/migration_promote.go | 496 ++++++ store/migration_versions.go | 128 +- store/migration_versions_test.go | 363 ++++ store/mvcc_store.go | 23 +- store/mvcc_store_snapshot_test.go | 46 + store/store.go | 65 +- 44 files changed, 7390 insertions(+), 344 deletions(-) create mode 100644 adapter/internal_migration_test.go create mode 100644 kv/fsm_migration_import.go create mode 100644 kv/fsm_migration_import_test.go create mode 100644 kv/fsm_migration_promote.go create mode 100644 kv/fsm_migration_promote_test.go create mode 100644 store/migration_promote.go diff --git a/adapter/distribution_server.go b/adapter/distribution_server.go index 06c924721..65b448a10 100644 --- a/adapter/distribution_server.go +++ b/adapter/distribution_server.go @@ -92,15 +92,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. @@ -176,6 +177,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. @@ -659,6 +716,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 ff4d198ea..46a6502e0 100644 --- a/adapter/distribution_server_test.go +++ b/adapter/distribution_server_test.go @@ -63,7 +63,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}, @@ -81,11 +86,37 @@ func TestDistributionServerRouteReadsHonorStartupGate(t *testing.T) { require.Error(t, err) require.Equal(t, codes.Unavailable, status.Code(err)) + _, err = s.GetRouteOwnership(context.Background(), &pb.GetRouteOwnershipRequest{ + Key: []byte("a"), + CatalogVersion: engine.Version(), + }) + require.Error(t, err) + require.Equal(t, codes.Unavailable, status.Code(err)) + + _, err = s.GetIntersectingRoutes(context.Background(), &pb.GetIntersectingRoutesRequest{ + Start: []byte("a"), + End: []byte("z"), + CatalogVersion: engine.Version(), + }) + require.Error(t, err) + require.Equal(t, codes.Unavailable, status.Code(err)) + blocked = false _, err = s.GetRoute(context.Background(), &pb.GetRouteRequest{Key: []byte("a")}) require.NoError(t, err) _, err = s.ListRoutes(context.Background(), &pb.ListRoutesRequest{}) require.NoError(t, err) + _, err = s.GetRouteOwnership(context.Background(), &pb.GetRouteOwnershipRequest{ + Key: []byte("a"), + CatalogVersion: engine.Version(), + }) + require.NoError(t, err) + _, err = s.GetIntersectingRoutes(context.Background(), &pb.GetIntersectingRoutesRequest{ + Start: []byte("a"), + End: []byte("z"), + CatalogVersion: engine.Version(), + }) + require.NoError(t, err) } func TestDistributionServerGetTimestamp_IsMonotonic(t *testing.T) { @@ -182,6 +213,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 bf4279e31..89784011e 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,51 @@ func WithInternalTimestampAllocator(alloc kv.TimestampAllocator) InternalOption } } +func WithInternalStore(st store.MVCCStore) InternalOption { + return func(i *Internal) { + i.store = st + } +} + +func WithInternalMigrationProposer(proposer raftengine.Proposer) InternalOption { + return func(i *Internal) { + i.migrationProposer = proposer + } +} + +func WithInternalMigrationImportGate(gate func(context.Context) error) InternalOption { + return func(i *Internal) { + i.migrationImportGate = gate + } +} + +func WithInternalMigrationPromoteGate(gate func(context.Context) error) InternalOption { + return func(i *Internal) { + i.migrationPromoteGate = gate + } +} + +func WithInternalRouteEngine(engine *distribution.Engine) InternalOption { + return func(i *Internal) { + i.routeEngine = engine + } +} + +func WithInternalMigrationExportRouting(groupID uint64, resolver kv.PartitionResolver) InternalOption { + return func(i *Internal) { + i.migrationExportGroupID = groupID + i.migrationExportResolver = resolver + } +} + func NewInternalWithEngine(txm kv.Transactional, leader raftengine.LeaderView, clock *kv.HLC, relay *RedisPubSubRelay, opts ...InternalOption) *Internal { i := &Internal{ - leader: leader, - transactionManager: txm, - clock: clock, - relay: relay, + leader: leader, + transactionManager: txm, + clock: clock, + relay: relay, + migrationImportGate: defaultMigrationImportGate, + migrationPromoteGate: defaultMigrationPromoteGate, } for _, opt := range opts { opt(i) @@ -32,11 +78,18 @@ 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 + migrationImportGate func(context.Context) error + migrationPromoteGate func(context.Context) error + routeEngine *distribution.Engine + migrationExportGroupID uint64 + migrationExportResolver kv.PartitionResolver pb.UnimplementedInternalServer } @@ -47,6 +100,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) @@ -87,6 +146,424 @@ 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 := i.exportRangeVersionsOptions(req) + for { + result, err := i.store.ExportVersions(stream.Context(), opts) + if err != nil { + return errors.WithStack(err) + } + if !result.Done && bytes.Equal(opts.Cursor, result.NextCursor) { + return errors.WithStack(status.Error(codes.Internal, "migration export cursor did not progress")) + } + 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 + } + if err := i.verifyMigrationImportEnabled(ctx); err != nil { + return nil, err + } + result, err := i.proposeMigrationImport(ctx, req) + if err != nil { + return nil, errors.WithStack(err) + } + if i.clock != nil && result.MaxImportedTS > 0 { + i.clock.Observe(result.MaxImportedTS) + } + return &pb.ImportRangeVersionsResponse{AckedCursor: result.AckedCursor}, nil +} + +func validateImportRangeVersionsRequest(req *pb.ImportRangeVersionsRequest) error { + if req.GetJobId() == 0 { + return errors.WithStack(status.Error(codes.InvalidArgument, "import range versions job_id is required")) + } + if req.GetBracketId() == 0 { + return errors.WithStack(status.Error(codes.InvalidArgument, "import range versions bracket_id is required")) + } + return nil +} + +func (i *Internal) verifyMigrationImportEnabled(ctx context.Context) error { + if i.migrationImportGate == nil { + return nil + } + if err := i.migrationImportGate(ctx); err != nil { + return errors.WithStack(err) + } + return nil +} + +func (i *Internal) PromoteStagedVersions(ctx context.Context, req *pb.PromoteStagedVersionsRequest) (*pb.PromoteStagedVersionsResponse, error) { + if req == nil { + return nil, errors.WithStack(status.Error(codes.InvalidArgument, "promote staged versions request is nil")) + } + if req.GetJobId() == 0 { + return nil, errors.WithStack(status.Error(codes.InvalidArgument, "promote staged versions job_id is required")) + } + if i.migrationProposer == nil { + return nil, errors.WithStack(status.Error(codes.FailedPrecondition, "migration promote proposer is not configured")) + } + if err := i.verifyInternalLeader(ctx); err != nil { + return nil, err + } + if err := i.verifyMigrationPromoteEnabled(ctx); err != nil { + return nil, err + } + if err := validatePromoteStagedVersionsRequest(req); err != nil { + return nil, errors.WithStack(err) + } + result, err := i.proposeMigrationPromote(ctx, req) + if err != nil { + return nil, errors.WithStack(err) + } + if i.clock != nil && result.MaxPromotedTS > 0 { + i.clock.Observe(result.MaxPromotedTS) + } + return &pb.PromoteStagedVersionsResponse{ + NextCursor: result.NextCursor, + Done: result.Done, + PromotedRows: result.PromotedRows, + MaxPromotedTs: result.MaxPromotedTS, + }, nil +} + +func (i *Internal) verifyMigrationPromoteEnabled(ctx context.Context) error { + if i.migrationPromoteGate == nil { + return nil + } + if err := i.migrationPromoteGate(ctx); err != nil { + return errors.WithStack(err) + } + return nil +} + +func validatePromoteStagedVersionsRequest(req *pb.PromoteStagedVersionsRequest) error { + prefix := distribution.MigrationStagedDataKeyPrefix(req.GetJobId()) + if err := store.ValidatePromotionCursorForRange(req.GetCursor(), prefix, store.PrefixScanEnd(prefix)); err != nil { + if errors.Is(err, store.ErrInvalidExportCursor) { + return errors.WithStack(status.Error(codes.InvalidArgument, store.ErrInvalidExportCursor.Error())) + } + return errors.WithStack(status.Errorf(codes.Internal, "validate promote cursor: %v", err)) + } + return nil +} + +func defaultMigrationImportGate(context.Context) error { + if migrationImportOpcodeEnabledFromEnv() { + return nil + } + return errors.WithStack(status.Error(codes.FailedPrecondition, "migration import opcode is disabled; enable after every voter is running a build that supports migration import")) +} + +func migrationImportOpcodeEnabledFromEnv() bool { + return envFlagEnabled("ELASTICKV_ENABLE_MIGRATION_IMPORT_OPCODE") +} + +func defaultMigrationPromoteGate(context.Context) error { + if migrationPromoteOpcodeEnabledFromEnv() { + return nil + } + return errors.WithStack(status.Error(codes.FailedPrecondition, "migration promote opcode is disabled; enable after every voter is running a build that supports migration promotion")) +} + +func migrationPromoteOpcodeEnabledFromEnv() bool { + return envFlagEnabled("ELASTICKV_ENABLE_MIGRATION_PROMOTE_OPCODE") +} + +func envFlagEnabled(name string) bool { + switch strings.ToLower(strings.TrimSpace(os.Getenv(name))) { + case "1", "true", "yes", "on": + return true + default: + return false + } +} + +func (i *Internal) verifyInternalLeader(ctx context.Context) error { + if i.leader == nil { + return errors.WithStack(ErrNotLeader) + } + if i.leader.State() != raftengine.StateLeader { + return errors.WithStack(ErrNotLeader) + } + if err := i.leader.VerifyLeader(ctx); err != nil { + return errors.WithStack(ErrNotLeader) + } + return nil +} + +func (i *Internal) verifyInternalLeaderApplied(ctx context.Context) error { + if i.leader == nil { + return errors.WithStack(ErrNotLeader) + } + if i.leader.State() != raftengine.StateLeader { + return errors.WithStack(ErrNotLeader) + } + _, err := i.leader.LinearizableRead(ctx) + return errors.WithStack(err) +} + +func (i *Internal) proposeMigrationImport(ctx context.Context, req *pb.ImportRangeVersionsRequest) (store.ImportVersionsResult, error) { + cmd, err := kv.MarshalMigrationImportCommand(req) + if err != nil { + return store.ImportVersionsResult{}, errors.WithStack(err) + } + resp, err := i.proposeMigrationCommand(ctx, cmd, "migration import") + if err != nil { + return store.ImportVersionsResult{}, errors.WithStack(err) + } + switch resp := resp.(type) { + case store.ImportVersionsResult: + return resp, nil + case *store.ImportVersionsResult: + if resp == nil { + return store.ImportVersionsResult{}, errors.New("migration import apply returned nil result") + } + return *resp, nil + case error: + return store.ImportVersionsResult{}, errors.WithStack(resp) + default: + return store.ImportVersionsResult{}, errors.WithStack(errors.Newf("unexpected migration import apply response type %T", resp)) + } +} + +func (i *Internal) proposeMigrationPromote(ctx context.Context, req *pb.PromoteStagedVersionsRequest) (store.PromoteVersionsResult, error) { + cmd, err := kv.MarshalMigrationPromoteCommand(req) + if err != nil { + return store.PromoteVersionsResult{}, errors.WithStack(err) + } + resp, err := i.proposeMigrationCommand(ctx, cmd, "migration promote") + if err != nil { + return store.PromoteVersionsResult{}, errors.WithStack(err) + } + switch resp := resp.(type) { + case store.PromoteVersionsResult: + return resp, nil + case *store.PromoteVersionsResult: + if resp == nil { + return store.PromoteVersionsResult{}, errors.New("migration promote apply returned nil result") + } + return *resp, nil + case error: + return store.PromoteVersionsResult{}, errors.WithStack(resp) + default: + return store.PromoteVersionsResult{}, errors.WithStack(errors.Newf("unexpected migration promote apply response type %T", resp)) + } +} + +func (i *Internal) proposeMigrationCommand(ctx context.Context, cmd []byte, label string) (any, error) { + result, err := i.migrationProposer.Propose(ctx, cmd) + if err != nil { + return nil, errors.WithStack(err) + } + if result == nil { + return nil, errors.WithStack(errors.Newf("%s proposal returned nil result", label)) + } + return result.Response, nil +} + +func (i *Internal) exportRangeVersionsOptions(req *pb.ExportRangeVersionsRequest) store.ExportVersionsOptions { + chunkBytes := 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: i.migrationExportFilter(req), + AcceptVersion: i.migrationExportVersionFilter(req), + } + return opts +} + +func (i *Internal) migrationExportFilter(req *pb.ExportRangeVersionsRequest) func([]byte) bool { + bracket := migrationExportBracket(req) + if req.GetKeyFamily() == distribution.MigrationFamilyLegacyListMetaDelta { + return bracket.ContainsRawKey + } + routeFilter := i.migrationExportRouteFilter(req) + if migrationFamilyRequiresDecodedS3(req.GetKeyFamily()) { + routeFilter = decodedS3BucketRouteFilter(req.GetKeyFamily(), req.GetRouteStart(), req.GetRouteEnd()) + } + return func(rawKey []byte) bool { + return bracket.ContainsRawKey(rawKey) && routeFilter(rawKey) + } +} + +func (i *Internal) migrationExportVersionFilter(req *pb.ExportRangeVersionsRequest) func([]byte, []byte) bool { + if req.GetKeyFamily() != distribution.MigrationFamilyLegacyListMetaDelta { + return nil + } + bracket := migrationExportBracket(req) + return func(rawKey, value []byte) bool { + return bracket.ContainsRoutedVersion(rawKey, value, req.GetRouteStart(), req.GetRouteEnd(), nil) + } +} + +func migrationExportBracket(req *pb.ExportRangeVersionsRequest) distribution.MigrationBracket { + excludeKnownInternal := req.GetExcludeKnownInternal() || req.GetKeyFamily() == distribution.MigrationFamilyUser + return distribution.MigrationBracket{ + Family: req.GetKeyFamily(), + Start: bytes.Clone(req.GetRangeStart()), + End: bytes.Clone(req.GetRangeEnd()), + ExcludeKnownInternal: excludeKnownInternal, + ExcludePrefixes: cloneByteSlices(req.GetExcludePrefixes()), + } +} + +func (i *Internal) migrationExportRouteFilter(req *pb.ExportRangeVersionsRequest) func([]byte) bool { + if i != nil && i.migrationExportGroupID != 0 && i.migrationExportResolver != nil { + return kv.RouteKeyFilterForGroup(req.GetRouteStart(), req.GetRouteEnd(), i.migrationExportGroupID, i.migrationExportResolver) + } + return kv.RouteKeyFilter(req.GetRouteStart(), req.GetRouteEnd()) +} + +func migrationFamilyRequiresDecodedS3(family uint32) bool { + return family == distribution.MigrationFamilyS3BucketMeta || + family == distribution.MigrationFamilyS3BucketGeneration +} + +func decodedS3BucketRouteFilter(family uint32, routeStart, routeEnd []byte) func([]byte) bool { + 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) (uint64, error) { if req == nil { return 0, nil @@ -162,18 +639,75 @@ 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) + } + if start, end, ok := forwardedS3BucketAuxiliaryRouteRange(mut.GetKey()); ok { + return rangesIntersect(route.Start, route.End, start, end) + } + return kv.RouteKeyFilter(route.Start, route.End)(mut.GetKey()) +} + +func forwardedS3BucketAuxiliaryRouteRange(key []byte) ([]byte, []byte, bool) { + bucket, ok := s3keys.ParseBucketMetaKey(key) + if !ok { + bucket, ok = s3keys.ParseBucketGenerationKey(key) + } + if !ok { + return nil, nil, false + } + start := s3keys.RoutePrefixForBucketAnyGeneration(bucket) + return start, prefixScanEnd(start), true +} + func (i *Internal) stampTxnTimestamps(ctx context.Context, reqs []*pb.Request) (uint64, error) { startTS := forwardedTxnStartTS(reqs) if startTS == 0 { @@ -194,7 +728,10 @@ func (i *Internal) stampTxnTimestamps(ctx context.Context, reqs []*pb.Request) ( } } - 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 { @@ -287,6 +824,34 @@ func collectForwardedTxnMetas(reqs []*pb.Request) ([]forwardedTxnMetaToUpdate, u return metaMutations, commitTS, 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..e2c162c96 --- /dev/null +++ b/adapter/internal_migration_test.go @@ -0,0 +1,661 @@ +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 TestInternalExportRangeVersionsUsesValueAwareLegacyListDeltaRouteFilter(t *testing.T) { + t.Parallel() + + ctx := context.Background() + st := store.NewMVCCStore() + userKey := []byte("target-list") + key := legacyListMetaDeltaKey(userKey, 10) + value := store.MarshalListMetaDelta(store.ListMetaDelta{LenDelta: 1}) + require.NoError(t, st.PutAt(ctx, key, value, 10, 0)) + internal := NewInternalWithEngine(nil, mockInternalLeader{}, nil, nil, WithInternalStore(st)) + stream := &captureExportRangeVersionsStream{ctx: ctx} + + err := internal.ExportRangeVersions(&pb.ExportRangeVersionsRequest{ + MaxCommitTs: 20, + RouteStart: []byte("target"), + RouteEnd: []byte("target-list\x00"), + KeyFamily: distribution.MigrationFamilyLegacyListMetaDelta, + RangeStart: []byte(store.LegacyListMetaDeltaPrefix), + RangeEnd: testPrefixScanEnd([]byte(store.LegacyListMetaDeltaPrefix)), + MaxScannedBytes: 1 << 20, + }, stream) + require.NoError(t, err) + require.Len(t, stream.responses, 1) + require.Equal(t, []*pb.MVCCVersion{ + {Key: key, CommitTs: 10, Value: value, KeyFamily: distribution.MigrationFamilyLegacyListMetaDelta}, + }, stream.responses[0].GetVersions()) +} + +func TestInternalExportRangeVersionsUsesAppliedReadFence(t *testing.T) { + t.Parallel() + + ctx := context.Background() + st := store.NewMVCCStore() + require.NoError(t, st.PutAt(ctx, []byte("a"), []byte("va"), 10, 0)) + leader := &recordingInternalLeader{} + internal := NewInternalWithEngine(nil, leader, nil, nil, WithInternalStore(st)) + stream := &captureExportRangeVersionsStream{ctx: ctx} + + err := internal.ExportRangeVersions(&pb.ExportRangeVersionsRequest{ + MaxCommitTs: 20, + RouteStart: []byte("a"), + RouteEnd: []byte("b"), + KeyFamily: distribution.MigrationFamilyUser, + RangeStart: []byte("a"), + RangeEnd: []byte("b"), + MaxScannedBytes: 1 << 20, + }, stream) + require.NoError(t, err) + require.Equal(t, 1, leader.linearizableReadCalls) + require.Zero(t, leader.verifyLeaderCalls) + require.Len(t, stream.responses, 1) + require.True(t, stream.responses[0].GetDone()) +} + +func TestInternalExportRangeVersionsFailsClosedWhenAppliedReadFenceFails(t *testing.T) { + t.Parallel() + + leader := &recordingInternalLeader{linearizableReadErr: context.Canceled} + internal := NewInternalWithEngine(nil, leader, nil, nil, WithInternalStore(store.NewMVCCStore())) + stream := &captureExportRangeVersionsStream{ctx: context.Background()} + + err := internal.ExportRangeVersions(&pb.ExportRangeVersionsRequest{ + MaxCommitTs: 20, + RouteStart: []byte("a"), + RouteEnd: []byte("b"), + KeyFamily: distribution.MigrationFamilyUser, + RangeStart: []byte("a"), + RangeEnd: []byte("b"), + MaxScannedBytes: 1 << 20, + }, stream) + require.ErrorIs(t, err, context.Canceled) + require.Equal(t, 1, leader.linearizableReadCalls) + require.Empty(t, stream.responses) +} + +func TestInternalExportRangeVersionsRejectsUnboundedExport(t *testing.T) { + t.Parallel() + + internal := NewInternalWithEngine(nil, mockInternalLeader{}, nil, nil, WithInternalStore(store.NewMVCCStore())) + stream := &captureExportRangeVersionsStream{ctx: context.Background()} + + err := internal.ExportRangeVersions(&pb.ExportRangeVersionsRequest{ + KeyFamily: distribution.MigrationFamilyUser, + }, stream) + require.Error(t, err) + require.Equal(t, codes.InvalidArgument, status.Code(err)) + + stream = &captureExportRangeVersionsStream{ctx: context.Background()} + err = internal.ExportRangeVersions(&pb.ExportRangeVersionsRequest{ + MaxCommitTs: 20, + KeyFamily: distribution.MigrationFamilyUser, + }, stream) + require.Error(t, err) + require.Equal(t, codes.InvalidArgument, status.Code(err)) +} + +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 TestInternalExportRangeVersionsUsesPartitionResolverGroup(t *testing.T) { + t.Parallel() + + ctx := context.Background() + st := store.NewMVCCStore() + resolver := NewSQSPartitionResolver(map[string][]uint64{ + "orders.fifo": {10, 11}, + }) + internal := NewInternalWithEngine(nil, mockInternalLeader{}, nil, nil, + WithInternalStore(st), + WithInternalMigrationExportRouting(11, resolver), + ) + + p0 := sqsPartitionedMsgDataKey("orders.fifo", 0, 1, "msg-0") + p1 := sqsPartitionedMsgDataKey("orders.fifo", 1, 1, "msg-1") + unknown := sqsPartitionedMsgDataKey("unknown.fifo", 0, 1, "msg-unknown") + require.NoError(t, st.PutAt(ctx, p0, []byte("p0"), 10, 0)) + require.NoError(t, st.PutAt(ctx, p1, []byte("p1"), 10, 0)) + require.NoError(t, st.PutAt(ctx, unknown, []byte("unknown"), 10, 0)) + + stream := &captureExportRangeVersionsStream{ctx: ctx} + prefix := []byte(SqsPartitionedMsgDataPrefix) + err := internal.ExportRangeVersions(&pb.ExportRangeVersionsRequest{ + MaxCommitTs: 20, + KeyFamily: distribution.MigrationFamilySQSPartitionedMessageData, + RangeStart: prefix, + RangeEnd: testPrefixScanEnd(prefix), + MaxScannedBytes: 1 << 20, + }, stream) + require.NoError(t, err) + require.Len(t, stream.responses, 1) + require.Equal(t, []*pb.MVCCVersion{ + {Key: p1, CommitTs: 10, Value: []byte("p1"), KeyFamily: distribution.MigrationFamilySQSPartitionedMessageData}, + }, stream.responses[0].GetVersions()) +} + +func TestInternalImportRangeVersionsAppliesStoreBatch(t *testing.T) { + t.Parallel() + + ctx := context.Background() + st := store.NewMVCCStore() + clock := kv.NewHLC() + proposer := &applyingMigrationProposer{ + fsm: kv.NewKvFSMWithHLC(st, clock), + } + internal := NewInternalWithEngine(nil, mockInternalLeader{}, clock, nil, + WithInternalStore(st), + WithInternalMigrationProposer(proposer), + WithInternalMigrationImportGate(func(context.Context) error { return nil }), + ) + + resp, err := internal.ImportRangeVersions(ctx, &pb.ImportRangeVersionsRequest{ + JobId: 7, + BracketId: 3, + BatchSeq: 1, + Cursor: []byte("cursor-1"), + Versions: []*pb.MVCCVersion{ + {Key: []byte("k"), CommitTs: 30, Value: []byte("v"), ExpireAt: 100}, + }, + }) + require.NoError(t, err) + require.Equal(t, []byte("cursor-1"), resp.GetAckedCursor()) + require.Equal(t, uint64(1), proposer.calls) + + staged := distribution.MigrationStagedDataKey(7, []byte("k")) + got, err := st.GetAt(ctx, staged, 30) + require.NoError(t, err) + require.Equal(t, []byte("v"), got) + _, err = st.GetAt(ctx, []byte("k"), 30) + require.ErrorIs(t, err, store.ErrKeyNotFound) + floor, err := st.MigrationHLCFloor(ctx, 7) + require.NoError(t, err) + require.Equal(t, uint64(30), floor) + require.GreaterOrEqual(t, clock.Current(), uint64(30)) +} + +func TestInternalImportRangeVersionsRejectsWhenOpcodeGateClosed(t *testing.T) { + t.Parallel() + + ctx := context.Background() + st := store.NewMVCCStore() + clock := kv.NewHLC() + proposer := &applyingMigrationProposer{ + fsm: kv.NewKvFSMWithHLC(st, clock), + } + internal := NewInternalWithEngine(nil, mockInternalLeader{}, clock, nil, + WithInternalStore(st), + WithInternalMigrationProposer(proposer), + WithInternalMigrationImportGate(func(context.Context) error { + return status.Error(codes.FailedPrecondition, "migration import disabled for test") + }), + ) + + resp, err := internal.ImportRangeVersions(ctx, &pb.ImportRangeVersionsRequest{ + JobId: 7, + BracketId: 3, + BatchSeq: 1, + Cursor: []byte("cursor-1"), + Versions: []*pb.MVCCVersion{ + {Key: []byte("k"), CommitTs: 30, Value: []byte("v")}, + }, + }) + require.Nil(t, resp) + require.Error(t, err) + require.Equal(t, codes.FailedPrecondition, status.Code(err)) + require.Equal(t, uint64(0), proposer.calls) +} + +func TestInternalImportRangeVersionsRejectsMissingIdentifiers(t *testing.T) { + t.Parallel() + + internal := NewInternalWithEngine(nil, mockInternalLeader{}, nil, nil, + WithInternalMigrationProposer(&applyingMigrationProposer{}), + ) + + _, err := internal.ImportRangeVersions(context.Background(), &pb.ImportRangeVersionsRequest{ + BracketId: 1, + BatchSeq: 1, + }) + require.Error(t, err) + require.Equal(t, codes.InvalidArgument, status.Code(err)) + + _, err = internal.ImportRangeVersions(context.Background(), &pb.ImportRangeVersionsRequest{ + JobId: 1, + BatchSeq: 1, + }) + require.Error(t, err) + require.Equal(t, codes.InvalidArgument, status.Code(err)) +} + +func TestInternalPromoteStagedVersionsRejectsWhenOpcodeGateClosed(t *testing.T) { + t.Parallel() + + ctx := context.Background() + st := store.NewMVCCStore() + clock := kv.NewHLC() + proposer := &applyingMigrationProposer{ + fsm: kv.NewKvFSMWithHLC(st, clock), + } + internal := NewInternalWithEngine(nil, mockInternalLeader{}, clock, nil, + WithInternalStore(st), + WithInternalMigrationProposer(proposer), + WithInternalMigrationPromoteGate(func(context.Context) error { + return status.Error(codes.FailedPrecondition, "migration promote disabled for test") + }), + ) + + resp, err := internal.PromoteStagedVersions(ctx, &pb.PromoteStagedVersionsRequest{ + JobId: 7, + MaxVersions: 10, + }) + require.Nil(t, resp) + require.Error(t, err) + require.Equal(t, codes.FailedPrecondition, status.Code(err)) + require.Equal(t, uint64(0), proposer.calls) +} + +func TestInternalPromoteStagedVersionsRejectsInvalidCursorBeforePropose(t *testing.T) { + t.Parallel() + + ctx := context.Background() + for _, tc := range []struct { + name string + cursor []byte + }{ + {name: "malformed cursor", cursor: []byte{0xff}}, + { + name: "cursor outside job staged prefix", + cursor: encodeTestExportCursor(distribution.MigrationStagedDataKey(8, []byte("k")), 30, testExportCursorTagEmitted), + }, + { + name: "scanned cursor inside staged prefix", + cursor: encodeTestExportCursor(distribution.MigrationStagedDataKey(7, []byte("k")), 31, testExportCursorTagScanned), + }, + { + name: "pruned-key cursor inside staged prefix", + cursor: encodeTestExportCursor(distribution.MigrationStagedDataKey(7, []byte("k")), 32, testExportCursorTagPrunedKey), + }, + { + name: "skipped-key cursor inside staged prefix", + cursor: encodeTestExportCursor(distribution.MigrationStagedDataKey(7, []byte("k")), 33, testExportCursorTagSkippedKey), + }, + } { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + + st := store.NewMVCCStore() + clock := kv.NewHLC() + proposer := &applyingMigrationProposer{ + fsm: kv.NewKvFSMWithHLC(st, clock), + } + internal := NewInternalWithEngine(nil, mockInternalLeader{}, clock, nil, + WithInternalStore(st), + WithInternalMigrationProposer(proposer), + WithInternalMigrationPromoteGate(func(context.Context) error { return nil }), + ) + + resp, err := internal.PromoteStagedVersions(ctx, &pb.PromoteStagedVersionsRequest{ + JobId: 7, + Cursor: tc.cursor, + MaxVersions: 10, + }) + require.Nil(t, resp) + require.Error(t, err) + require.Equal(t, codes.InvalidArgument, status.Code(err)) + require.ErrorContains(t, err, store.ErrInvalidExportCursor.Error()) + require.Equal(t, uint64(0), proposer.calls) + }) + } +} + +func TestInternalPromoteStagedVersionsAppliesStoreBatch(t *testing.T) { + t.Parallel() + + ctx := context.Background() + st := store.NewMVCCStore() + clock := kv.NewHLC() + proposer := &applyingMigrationProposer{ + fsm: kv.NewKvFSMWithHLC(st, clock), + } + internal := NewInternalWithEngine(nil, mockInternalLeader{}, clock, nil, + WithInternalStore(st), + WithInternalMigrationProposer(proposer), + WithInternalMigrationPromoteGate(func(context.Context) error { return nil }), + ) + + staged := distribution.MigrationStagedDataKey(7, []byte("k")) + require.NoError(t, st.PutAt(ctx, staged, []byte("v"), 30, 0)) + + resp, err := internal.PromoteStagedVersions(ctx, &pb.PromoteStagedVersionsRequest{ + JobId: 7, + MaxVersions: 10, + }) + require.NoError(t, err) + require.True(t, resp.GetDone()) + require.Equal(t, uint64(1), resp.GetPromotedRows()) + require.Equal(t, uint64(30), resp.GetMaxPromotedTs()) + require.Equal(t, uint64(1), proposer.calls) + + got, err := st.GetAt(ctx, []byte("k"), 30) + require.NoError(t, err) + require.Equal(t, []byte("v"), got) + _, err = st.GetAt(ctx, staged, 30) + require.ErrorIs(t, err, store.ErrKeyNotFound) + require.GreaterOrEqual(t, clock.Current(), uint64(30)) +} diff --git a/adapter/internal_test.go b/adapter/internal_test.go index fcabc5356..a3d937f95 100644 --- a/adapter/internal_test.go +++ b/adapter/internal_test.go @@ -5,11 +5,19 @@ import ( "encoding/binary" "testing" + "github.com/bootjp/elastickv/distribution" + "github.com/bootjp/elastickv/internal/s3keys" "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() @@ -241,3 +249,183 @@ func TestStampTxnTimestamps_UsesSingleTxnStartTS(t *testing.T) { require.Greater(t, meta.CommitTS, uint64(9)) require.Equal(t, meta.CommitTS, commitTS) } + +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 TestStampRawTimestampsIgnoresRawRouteFloorForS3BucketAuxiliaryWrite(t *testing.T) { + t.Parallel() + + const bucket = "bucket-a" + key := s3keys.BucketMetaKey(bucket) + auxStart := s3keys.RoutePrefixForBucketAnyGeneration(bucket) + auxEnd := prefixScanEnd(auxStart) + engine := distribution.NewEngine() + require.NoError(t, engine.ApplySnapshot(distribution.CatalogSnapshot{ + Version: 1, + Routes: []distribution.RouteDescriptor{ + {RouteID: 1, Start: []byte(""), End: auxStart, GroupID: 1, State: distribution.RouteStateActive}, + {RouteID: 2, Start: auxStart, End: auxEnd, GroupID: 2, State: distribution.RouteStateActive}, + {RouteID: 3, Start: auxEnd, End: nil, GroupID: 1, State: distribution.RouteStateActive, MinWriteTSExclusive: ^uint64(0)}, + }, + })) + i := &Internal{routeEngine: engine} + reqs := []*pb.Request{{ + Ts: 100, + Mutations: []*pb.Mutation{{Op: pb.Op_PUT, Key: key, Value: []byte("meta")}}, + }} + + require.NoError(t, i.stampRawTimestamps(context.Background(), reqs)) +} + +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 4d1319700..ee3bf1726 100644 --- a/adapter/redis_lua_context.go +++ b/adapter/redis_lua_context.go @@ -504,20 +504,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 30879644a..6c2eed76a 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 0535ff9b3..e55b6dc46 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 @@ -386,17 +387,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 } @@ -429,17 +420,7 @@ func (e *Engine) GetIntersectingRoutesWithVersion(start, end []byte) ([]Route, u 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, e.catalogVersion } diff --git a/distribution/engine_test.go b/distribution/engine_test.go index 67a7c0053..836303b93 100644 --- a/distribution/engine_test.go +++ b/distribution/engine_test.go @@ -239,7 +239,16 @@ func TestEngineApplySnapshot_ReplacesRoutesAndVersion(t *testing.T) { Version: 1, Routes: []RouteDescriptor{ {RouteID: 10, Start: []byte(""), End: []byte("m"), GroupID: 1, State: RouteStateActive}, - {RouteID: 11, Start: []byte("m"), End: nil, GroupID: 2, State: RouteStateWriteFenced}, + { + RouteID: 11, + Start: []byte("m"), + End: nil, + GroupID: 2, + State: RouteStateWriteFenced, + StagedVisibilityActive: true, + MigrationJobID: 42, + MinWriteTSExclusive: 99, + }, }, }) if err != nil { @@ -260,6 +269,19 @@ func TestEngineApplySnapshot_ReplacesRoutesAndVersion(t *testing.T) { if stats[1].RouteID != 11 || stats[1].State != RouteStateWriteFenced { t.Fatalf("unexpected second route metadata: %+v", stats[1]) } + assertStagedRouteMetadata(t, stats[1]) + route, ok := e.GetRoute([]byte("m")) + if !ok { + t.Fatalf("expected route for m") + } + assertStagedRouteMetadata(t, route) +} + +func assertStagedRouteMetadata(t *testing.T, route Route) { + t.Helper() + if !route.StagedVisibilityActive || route.MigrationJobID != 42 || route.MinWriteTSExclusive != 99 { + t.Fatalf("staged route metadata was not preserved: %+v", route) + } } func TestEngineRouteLookupsReturnMatchingCatalogVersion(t *testing.T) { diff --git a/distribution/migrator.go b/distribution/migrator.go index 000da7803..db32062a3 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 ( @@ -283,6 +287,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 104625499..c26394a07 100644 --- a/distribution/migrator_export_plan_test.go +++ b/distribution/migrator_export_plan_test.go @@ -396,6 +396,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/internal/s3keys/keys_test.go b/internal/s3keys/keys_test.go index 960ab6b21..70b0bf3a2 100644 --- a/internal/s3keys/keys_test.go +++ b/internal/s3keys/keys_test.go @@ -30,6 +30,14 @@ func TestBucketGenerationKey_RoundTripsZeroByteSegments(t *testing.T) { require.Equal(t, bucket, parsed) } +func TestParseBucketGenerationKey_RejectsNonGenerationKey(t *testing.T) { + t.Parallel() + + parsed, ok := ParseBucketGenerationKey(BucketMetaKey("bucket")) + require.False(t, ok) + require.Empty(t, parsed) +} + func TestObjectManifestKey_RoundTripsZeroByteSegments(t *testing.T) { t.Parallel() diff --git a/kv/fsm.go b/kv/fsm.go index 09846821b..55a95216d 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" @@ -128,6 +129,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 @@ -281,6 +286,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 @@ -318,12 +325,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) @@ -364,13 +370,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: @@ -402,6 +412,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) { @@ -500,16 +518,7 @@ 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.assertNoConflictingTxnLocks(ctx, r.Mutations, nil, 0); err != nil { + if err := f.validateRawMutationsForApply(ctx, r, commitTS); err != nil { return err } @@ -526,6 +535,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) { @@ -540,6 +581,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) } @@ -547,6 +594,111 @@ 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, writeFenceBypassKeys [][]byte, commitTS uint64) error { + bypassKeys := writeFenceBypassKeySet(writeFenceBypassKeys) + for _, mut := range muts { + if mut == nil || len(mut.Key) == 0 || isTxnInternalKey(mut.Key) { + continue + } + if _, bypass := bypassKeys[string(mut.Key)]; bypass { + 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 @@ -564,6 +716,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), @@ -968,21 +1125,24 @@ func (f *kvFSM) verifyOwnerFromSnapshot(mutations []*pb.Mutation, bypassKeys map if _, ok := bypassKeys[string(mut.Key)]; ok { continue } - // routeKey-normalize before OwnerOf so the gate routes the - // same way as ShardRouter.ResolveGroup — raw adapter keys - // and route catalog ranges live in different lex bands - // (issue #930). - rKey := routeKey(mut.Key) - owner, found := snap.OwnerOf(rKey) + ownerKey := composed1OwnerKey(mut.Key) + owner, found := snap.OwnerOf(ownerKey) if !found || owner != f.shardGroupID { return errors.Wrapf(ErrComposed1Violation, "%s-version v=%d: key %q (routeKey %q) owned by group %d (found=%v); this FSM serves group %d", - phase, snapVer, mut.Key, rKey, owner, found, f.shardGroupID) + phase, snapVer, mut.Key, ownerKey, owner, found, f.shardGroupID) } } return nil } +func composed1OwnerKey(key []byte) []byte { + if start, _, ok := s3BucketAuxiliaryRouteRange(key); ok { + return start + } + return routeKey(key) +} + func (f *kvFSM) validateConflicts(ctx context.Context, muts []*pb.Mutation, startTS uint64) error { seen := make(map[string]struct{}, len(muts)) for _, mut := range muts { @@ -1044,7 +1204,7 @@ func (f *kvFSM) handlePrepareRequest(ctx context.Context, r *pb.Request) error { } startTS := r.Ts - uniq, err := uniqueMutations(muts) + uniq, err := f.uniqueMutationsAboveFloor(muts, r.GetWriteFenceBypassKeys(), startTS) if err != nil { return err } @@ -1114,7 +1274,7 @@ func (f *kvFSM) handleOnePhaseTxnRequest(ctx context.Context, r *pb.Request, com return nil } - uniq, err := uniqueMutations(muts) + uniq, err := f.uniqueMutationsAboveFloor(muts, r.GetWriteFenceBypassKeys(), commitTS) if err != nil { return err } @@ -1130,6 +1290,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, writeFenceBypassKeys [][]byte, commitTS uint64) ([]*pb.Mutation, error) { + uniq, err := uniqueMutations(muts) + if err != nil { + return nil, err + } + if err := f.verifyRouteWriteTimestampFloorsForMutations(uniq, writeFenceBypassKeys, commitTS); err != nil { + return nil, err + } + return uniq, nil +} + +func (f *kvFSM) uniqueTxnMutationsAboveFloor(muts []*pb.Mutation, writeFenceBypassKeys [][]byte, commitTS uint64) ([]*pb.Mutation, error) { + uniq, err := uniqueTxnMutations(muts) + if err != nil { + return nil, err + } + if err := f.verifyRouteWriteTimestampFloorsForMutations(uniq, writeFenceBypassKeys, 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 @@ -1146,9 +1336,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 { @@ -1169,7 +1405,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, r.GetWriteFenceBypassKeys(), commitTS) if err != nil { return err } diff --git a/kv/fsm_migration_fence_test.go b/kv/fsm_migration_fence_test.go index 807a5f551..1eee26301 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() @@ -97,13 +108,25 @@ func TestFSMWriteFenceBypassAllowsMarkedRawPointWrite(t *testing.T) { require.Equal(t, []byte("v"), got) } +func TestFSMWriteFenceBypassAllowsRawWriteBelowBypassedRouteFloor(t *testing.T) { + t.Parallel() + + fsm := newWriteFloorFSM(t) + key := []byte("!sqs|msg|data|p|partitioned-key") + err := fsm.handleRawRequest(context.Background(), &pb.Request{ + WriteFenceBypassKeys: [][]byte{key}, + Mutations: []*pb.Mutation{{Op: pb.Op_PUT, Key: key, Value: []byte("v")}}, + }, 10) + require.NoError(t, err) +} + func TestFSMWriteFenceBypassAllowsPinnedTxnOnNonOwningGroup(t *testing.T) { t.Parallel() engine := distribution.NewEngine() applyComposed1Snapshot(t, engine, 1, []distribution.RouteDescriptor{ {RouteID: 1, Start: []byte(""), End: []byte("m"), GroupID: 1, State: distribution.RouteStateActive}, - {RouteID: 2, Start: []byte("m"), End: nil, GroupID: 2, State: distribution.RouteStateWriteFenced}, + {RouteID: 2, Start: []byte("m"), End: nil, GroupID: 2, State: distribution.RouteStateWriteFenced, MinWriteTSExclusive: 100}, }) fsm := newComposed1FSM(t, engine, 1) key := []byte("z") @@ -182,7 +205,7 @@ func TestFSMRejectsCurrentWriteFencedS3BucketAuxiliaryPointWrite(t *testing.T) { t.Parallel() ctx := context.Background() - const bucket = "bucket-a" + const bucket = "bucket-b" fsm := newS3BucketAuxiliaryWriteFencedFSM(t, bucket) for _, key := range [][]byte{ @@ -199,7 +222,7 @@ func TestFSMRejectsCurrentWriteFencedS3BucketAuxiliaryPointWrite(t *testing.T) { func TestFSMRejectsObservedWriteFencedS3BucketAuxiliaryPointWrite(t *testing.T) { t.Parallel() - const bucket = "bucket-a" + const bucket = "bucket-b" fsm := newS3BucketAuxiliaryWriteFencedFSM(t, bucket) err := fsm.handleRawRequest(context.Background(), &pb.Request{ @@ -211,6 +234,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() @@ -319,3 +391,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..d146b6fea --- /dev/null +++ b/kv/fsm_migration_import.go @@ -0,0 +1,85 @@ +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(), + AppliedIndex: f.pendingApplyIdx, + 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..dafd726d9 --- /dev/null +++ b/kv/fsm_migration_import_test.go @@ -0,0 +1,112 @@ +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) +} + +type captureMigrationImportStore struct { + store.MVCCStore + opts store.ImportVersionsOptions +} + +func (s *captureMigrationImportStore) ImportVersionsRaft(_ context.Context, opts store.ImportVersionsOptions) (store.ImportVersionsResult, error) { + s.opts = opts + return store.ImportVersionsResult{AckedCursor: opts.Cursor, MaxImportedTS: 10}, nil +} + +func TestApplyMigrationImportThreadsPendingApplyIndex(t *testing.T) { + t.Parallel() + + capturing := &captureMigrationImportStore{} + fsm := &kvFSM{store: capturing, pendingApplyIdx: 1234} + req := &pb.ImportRangeVersionsRequest{ + JobId: 9, + BracketId: 1, + BatchSeq: 1, + Cursor: []byte("cursor"), + Versions: []*pb.MVCCVersion{ + {Key: []byte("user|k"), CommitTs: 10, Value: []byte("v")}, + }, + } + data, err := proto.Marshal(req) + require.NoError(t, err) + + applied := fsm.applyMigrationImport(context.Background(), data) + result, ok := applied.(store.ImportVersionsResult) + require.True(t, ok, "got %T: %v", applied, applied) + require.Equal(t, []byte("cursor"), result.AckedCursor) + require.Equal(t, uint64(1234), capturing.opts.AppliedIndex) + require.Equal(t, uint64(9), capturing.opts.JobID) + require.Len(t, capturing.opts.Versions, 1) + require.Equal(t, distribution.MigrationStagedDataKey(9, []byte("user|k")), capturing.opts.Versions[0].Key) +} 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 89cf19e86..51b092867 100644 --- a/kv/leader_routed_store.go +++ b/kv/leader_routed_store.go @@ -718,6 +718,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 773645892..63cfaf745 100644 --- a/kv/shard_store.go +++ b/kv/shard_store.go @@ -34,9 +34,10 @@ type ShardStore struct { } var ( - ErrCrossShardMutationBatchNotSupported = errors.New("cross-shard mutation batches are not supported") - ErrReadRouteVersionUnavailable = errors.New("read route version is not locally available") - ErrFilesystemPlacementTargetNotFound = errors.New("filesystem placement target group has no routable home slot") + ErrCrossShardMutationBatchNotSupported = errors.New("cross-shard mutation batches are not supported") + ErrExplicitGroupStagedVisibilityUnresolved = errors.New("explicit group read cannot resolve staged visibility route") + ErrReadRouteVersionUnavailable = errors.New("read route version is not locally available") + ErrFilesystemPlacementTargetNotFound = errors.New("filesystem placement target group has no routable home slot") ) // NewShardStore creates a sharded MVCC store wrapper. @@ -183,24 +184,20 @@ func (s *ShardStore) GetAtWithReadFence(ctx context.Context, key []byte, ts uint if groupID != 0 { return s.getGroupAtWithReadFence(ctx, groupID, key, ts, readRouteVersion) } - route, routeVersion, ok := s.engine.GetRouteWithVersion(routeKey(key)) + route, g, routeVersion, ok := s.routeAndGroupForKeyWithVersion(key) readRouteVersion = max(readRouteVersion, routeVersion) - if !ok { - return nil, store.ErrKeyNotFound - } - g, ok := s.groupForID(route.GroupID) 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) } @@ -217,8 +214,12 @@ func (s *ShardStore) getGroupAtWithReadFence(ctx context.Context, groupID uint64 if !ok || g.Store == nil { return nil, store.ErrKeyNotFound } + route, err := s.routeForExplicitGroupKey(groupID, key) + if err != nil { + return nil, err + } - return s.getGroupAt(ctx, g, key, ts, groupID, readRouteVersion) + return s.getGroupAt(ctx, g, route, key, ts, groupID, readRouteVersion) } func (s *ShardStore) getRouteAt(ctx context.Context, route distribution.Route, key []byte, ts uint64) ([]byte, error) { @@ -226,15 +227,15 @@ func (s *ShardStore) getRouteAt(ctx context.Context, route distribution.Route, k if !ok || g.Store == nil { return nil, store.ErrKeyNotFound } - return s.getGroupAt(ctx, g, key, ts, route.GroupID, 0) + return s.getGroupAt(ctx, g, route, key, ts, route.GroupID, 0) } -func (s *ShardStore) getGroupAt(ctx context.Context, g *ShardGroup, key []byte, ts uint64, groupID uint64, readRouteVersion uint64) ([]byte, error) { +func (s *ShardStore) getGroupAt(ctx context.Context, g *ShardGroup, route distribution.Route, key []byte, ts uint64, groupID uint64, readRouteVersion uint64) ([]byte, error) { if engineForGroup(g) == nil { - return s.localGetAt(ctx, g, 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) } @@ -256,16 +257,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) @@ -273,6 +277,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 { @@ -366,7 +479,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) } @@ -381,7 +494,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, routeVersion := s.routesForFencedScanWithVersion(start, end, routeStart, routeEnd) @@ -434,6 +547,10 @@ func (s *ShardStore) ScanAtPhysicalLimit(ctx context.Context, start []byte, end return []*store.KVPair{}, false, nil } routes, clampToRoutes := s.routesForForwardScan(start, end) + if routesContainStagedVisibility(routes) { + kvs, err := s.ScanAt(ctx, start, end, visibleLimit, ts) + return kvs, false, err + } if len(routes) != 1 || clampToRoutes { kvs, err := s.ScanAt(ctx, start, end, visibleLimit, ts) return kvs, false, err @@ -447,10 +564,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 { + routes, dedupeByKey = prepareUnclampedRawScanRoutes(routes, dedupeByKey) + } + 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 && clampToRoutes { + route = routes[len(routes)-1-i] + } + scanStart := start + scanEnd := end + if clampToRoutes { + scanStart = clampScanStart(start, route.Start) + scanEnd = clampScanEnd(end, route.End) + } + kvs, err := s.scanRouteAtDirectionWithS3StagedOwnerFilter(ctx, routes, route, scanStart, scanEnd, limit, ts, reverse, true, readRouteVersion, routeStart, routeEnd, dedupeByKey) + if err != nil { + return nil, err + } + if clampToRoutes { + out = append(out, kvs...) + if len(out) >= limit { + clear(out[limit:]) + return out[:limit], nil + } + continue + } + out = mergeAndTrimScanResultsWithOptions(out, kvs, limit, reverse, dedupeByKey) + } + return out, nil } // ReverseScanGroupAt reverse-scans a range on the explicitly selected Raft group. @@ -496,6 +656,10 @@ func (s *ShardStore) ReverseScanAtPhysicalLimit(ctx context.Context, start []byt return []*store.KVPair{}, false, nil } routes, clampToRoutes := s.routesForReverseScan(start, end) + if routesContainStagedVisibility(routes) { + kvs, err := s.ReverseScanAt(ctx, start, end, visibleLimit, ts) + return kvs, false, err + } if len(routes) != 1 || clampToRoutes { kvs, err := s.ReverseScanAt(ctx, start, end, visibleLimit, ts) return kvs, false, err @@ -554,6 +718,9 @@ func (s *ShardStore) routesForScanWithVersion(start []byte, end []byte) ([]distr if routes, version, ok := s.routesForFilesystemUsageScanWithVersion(start, end); ok { return routes, false, version } + if routes, version, ok := s.routesForS3BucketAuxiliaryScan(start, end); ok { + return routes, false, version + } if routes, version, ok := s.routesForFilesystemChunkScanWithVersion(start, end); ok { return routes, false, version } @@ -607,6 +774,191 @@ func (s *ShardStore) routesForInternalScanWithVersion(start []byte, end []byte) return internalScanRouteSelection{}, false } +func (s *ShardStore) routesForS3BucketAuxiliaryScan(start []byte, end []byte) ([]distribution.Route, uint64, bool) { + if s == nil || s.engine == nil || !s3BucketAuxiliaryScanBounds(start, end) { + return nil, 0, false + } + catalogRoutes, version := s.engine.GetIntersectingRoutesWithVersion(nil, nil) + routes := make([]distribution.Route, 0) + for _, route := range catalogRoutes { + if migrationRouteRangesIntersect(route.Start, route.End, start, end) { + routes = append(routes, route) + } + } + routeStart, routeEnd := s3BucketAuxiliaryScanRouteRange(start, end) + for _, route := range catalogRoutes { + if routeHasStagedVisibility(route) && migrationRouteRangesIntersect(route.Start, route.End, routeStart, routeEnd) { + routes = append(routes, route) + } + } + return routes, version, true +} + +func s3BucketAuxiliaryScanBounds(start []byte, end []byte) bool { + if !bytes.HasPrefix(start, []byte(s3keys.BucketMetaPrefix)) && + !bytes.HasPrefix(start, []byte(s3keys.BucketGenerationPrefix)) { + return false + } + if end == nil { + return true + } + return bytes.Compare(start, end) < 0 +} + +func s3BucketAuxiliaryScanRouteRange(start []byte, end []byte) ([]byte, []byte) { + if routeStart, routeEnd, ok := s3BucketAuxiliaryRouteRange(start); ok && end != nil && bytes.Compare(end, prefixScanEnd(start)) <= 0 { + return routeStart, routeEnd + } + routeStart := []byte(s3keys.RoutePrefix) + return routeStart, prefixScanEnd(routeStart) +} + +type repeatedRawScanRouteKey struct { + groupID uint64 + staged bool + migrationJobID uint64 + routeStart string + routeEnd string +} + +func dedupeRepeatedRawScanRoutes(routes []distribution.Route) []distribution.Route { + if len(routes) <= 1 { + return routes + } + out := make([]distribution.Route, 0, len(routes)) + seen := make(map[repeatedRawScanRouteKey]struct{}, len(routes)) + for _, route := range routes { + key := repeatedRawScanRouteDedupeKey(route) + if _, ok := seen[key]; ok { + continue + } + seen[key] = struct{}{} + out = append(out, route) + } + return out +} + +func repeatedRawScanRouteDedupeKey(route distribution.Route) repeatedRawScanRouteKey { + key := repeatedRawScanRouteKey{groupID: route.GroupID} + if routeHasStagedVisibility(route) { + key.staged = true + key.migrationJobID = route.MigrationJobID + key.routeStart = string(route.Start) + key.routeEnd = string(route.End) + } + return key +} + +func prepareUnclampedRawScanRoutes(routes []distribution.Route, dedupeByKey bool) ([]distribution.Route, bool) { + if routesContainStagedVisibility(routes) { + routes = dedupeRepeatedRawScanRoutes(routes) + return orderRawScanRoutesForStagedVisibility(routes), true + } + if !dedupeByKey { + routes = dedupeRepeatedRawScanRoutes(routes) + } + return routes, dedupeByKey +} + +func orderRawScanRoutesForStagedVisibility(routes []distribution.Route) []distribution.Route { + if !routesContainStagedVisibility(routes) { + return routes + } + out := make([]distribution.Route, 0, len(routes)) + staged := make([]distribution.Route, 0) + for _, route := range routes { + if routeHasStagedVisibility(route) { + staged = append(staged, route) + continue + } + out = append(out, route) + } + return append(out, staged...) +} + +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 routesForLegacyListDeltaScan(catalogRoutes []distribution.Route, start []byte, end []byte) []distribution.Route { logicalUserKey := store.ExtractLegacyListUserKeyFromDeltaScanPrefix(start) routes := make([]distribution.Route, 0) @@ -640,6 +992,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 { @@ -686,29 +1042,30 @@ 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) filterUsageOwners := !clampToRoutes && !routeFilterPresent && filesystemUsageScanOverlap(start, end) + dedupeByKey := s3BucketAuxiliaryScanBounds(start, end) + if !clampToRoutes && !routeFilterPresent { + routes, dedupeByKey = prepareUnclampedRawScanRoutes(routes, dedupeByKey) + } for _, route := range routes { scanStart := start scanEnd := end if clampToRoutes { scanStart = clampScanStart(start, route.Start) scanEnd = clampScanEnd(end, route.End) - } else if !routeFilterPresent { - if _, seen := seenGroups[route.GroupID]; seen { - continue - } - seenGroups[route.GroupID] = struct{}{} } - kvs, err := s.scanRouteAtWithOptionalFilesystemUsageOwnerFilter( - ctx, route, scanStart, scanEnd, limit, ts, false, !clampToRoutes, - readRouteVersion, routeStart, routeEnd, filterUsageOwners, + kvs, err := s.scanRouteAtWithMigrationOwnerFilters( + ctx, routes, route, scanStart, scanEnd, limit, ts, false, !clampToRoutes, + readRouteVersion, routeStart, routeEnd, filterUsageOwners, dedupeByKey, ) if err != nil { return nil, err } + if isBroadLegacyListDeltaScan(start) && !routeFilterPresent { + kvs = markScanRouteGroup(kvs, route.GroupID, true) + } if clampToRoutes { out = append(out, kvs...) if len(out) >= limit { @@ -717,7 +1074,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 } @@ -748,6 +1105,34 @@ func (s *ShardStore) scanRouteAtWithOptionalFilesystemUsageOwnerFilter( ) } +func (s *ShardStore) scanRouteAtWithMigrationOwnerFilters( + ctx context.Context, + routes []distribution.Route, + route distribution.Route, + start []byte, + end []byte, + limit int, + ts uint64, + reverse bool, + explicitGroup bool, + readRouteVersion uint64, + routeStart []byte, + routeEnd []byte, + filterUsageOwners bool, + dedupeByKey bool, +) ([]*store.KVPair, error) { + if filterUsageOwners { + return s.scanRouteAtWithOptionalFilesystemUsageOwnerFilter( + ctx, route, start, end, limit, ts, reverse, explicitGroup, + readRouteVersion, routeStart, routeEnd, true, + ) + } + return s.scanRouteAtDirectionWithS3StagedOwnerFilter( + ctx, routes, route, start, end, limit, ts, reverse, explicitGroup, + readRouteVersion, routeStart, routeEnd, dedupeByKey, + ) +} + func (s *ShardStore) routesForFilesystemUsageScanWithVersion(start []byte, end []byte) ([]distribution.Route, uint64, bool) { if !filesystemUsageScanOverlap(start, end) { return nil, 0, false @@ -1001,12 +1386,16 @@ func (s *ShardStore) reverseScanRoutesAtWithReadFence( routeEnd []byte, ) ([]*store.KVPair, error) { out := make([]*store.KVPair, 0) - seenGroups := make(map[uint64]struct{}) routeFilterPresent := routeScanBoundsPresent(routeStart, routeEnd) filterUsageOwners := !clampToRoutes && !routeFilterPresent && filesystemUsageScanOverlap(start, end) - for i := len(routes) - 1; i >= 0; i-- { + dedupeByKey := s3BucketAuxiliaryScanBounds(start, end) + if !clampToRoutes && !routeFilterPresent { + routes, dedupeByKey = prepareUnclampedRawScanRoutes(routes, dedupeByKey) + } + for i := 0; i < len(routes); i++ { route := routes[i] if clampToRoutes { + route = routes[len(routes)-1-i] kvs, done, err := s.clampedReverseScanRouteAtWithReadFence(ctx, route, start, end, limit, len(out), ts, readRouteVersion, routeStart, routeEnd) if err != nil { return nil, err @@ -1022,23 +1411,17 @@ func (s *ShardStore) reverseScanRoutesAtWithReadFence( // shards), keys from different routes may interleave in descending order. // Fetch up to limit from every route and merge+sort descending so the // result honours the ReverseScanAt contract. - // De-duplicate by GroupID: after a range split both halves share the same - // GroupID (same backing shard store), so only scan each group once unless - // route filters make each descriptor's logical interval distinct. - if !routeFilterPresent { - if _, seen := seenGroups[route.GroupID]; seen { - continue - } - seenGroups[route.GroupID] = struct{}{} - } - kvs, err := s.scanRouteAtWithOptionalFilesystemUsageOwnerFilter( - ctx, route, start, end, limit, ts, true, true, - readRouteVersion, routeStart, routeEnd, filterUsageOwners, + kvs, err := s.scanRouteAtWithMigrationOwnerFilters( + ctx, routes, route, start, end, limit, ts, true, true, + readRouteVersion, routeStart, routeEnd, filterUsageOwners, dedupeByKey, ) if err != nil { return nil, err } - out = mergeAndTrimReverseScanResults(out, kvs, limit) + if isBroadLegacyListDeltaScan(start) && !routeFilterPresent { + kvs = markScanRouteGroup(kvs, route.GroupID, true) + } + out = mergeAndTrimScanResultsWithOptions(out, kvs, limit, true, dedupeByKey) } return out, nil } @@ -1070,11 +1453,17 @@ func (s *ShardStore) scanKeyRouteAtWithReadFence( } if engineForGroup(g) == nil { + if routeHasStagedVisibility(route) { + return scanKeysWithRefill(start, end, limit, func(cursor []byte, pageLimit int) ([][]byte, error) { + kvs, err := s.scanRouteWithStagedVisibility(ctx, g, route, cursor, end, pageLimit, ts, false) + return keysFromKVs(kvs), err + }) + } return s.scanKeysRouteLocal(ctx, g, start, end, limit, ts) } if isLinearizableRaftLeader(ctx, engineForGroup(g)) { - return s.scanKeysRouteAtLeader(ctx, g, start, end, limit, ts) + return s.scanKeysRouteAtLeader(ctx, g, route, start, end, limit, ts) } groupID := proxyScanGroupID(route, explicitGroup, readRouteVersion, nil, nil) @@ -1101,6 +1490,7 @@ func (s *ShardStore) scanKeysRouteLocal( func (s *ShardStore) scanKeysRouteAtLeader( ctx context.Context, g *ShardGroup, + route distribution.Route, start []byte, end []byte, limit int, @@ -1109,6 +1499,12 @@ func (s *ShardStore) scanKeysRouteAtLeader( if limit <= 0 { return [][]byte{}, nil } + if routeHasStagedVisibility(route) { + return scanKeysWithRefill(start, end, limit, func(cursor []byte, pageLimit int) ([][]byte, error) { + kvs, err := s.scanRouteAtLeader(ctx, g, route, cursor, end, pageLimit, ts, false) + return keysFromKVs(kvs), err + }) + } out := make([][]byte, 0, limit) cursor := start @@ -1118,7 +1514,7 @@ func (s *ShardStore) scanKeysRouteAtLeader( return nil, errors.WithStack(err) } if len(keys) == 0 { - keys, err := s.scanLockOnlyKeysAtLeader(ctx, g, cursor, end, ts, limit) + keys, err := s.scanLockOnlyKeysAtLeader(ctx, g, route, cursor, end, ts, limit) if err != nil { return nil, err } @@ -1132,7 +1528,7 @@ func (s *ShardStore) scanKeysRouteAtLeader( if err != nil { return nil, err } - kvs, err := s.resolveScanLocks(ctx, g, keyKVs, lockKVs, ts) + kvs, err := s.resolveScanLocks(ctx, g, route, keyKVs, lockKVs, ts) if err != nil { return nil, err } @@ -1150,6 +1546,7 @@ func (s *ShardStore) scanKeysRouteAtLeader( func (s *ShardStore) scanLockOnlyKeysAtLeader( ctx context.Context, g *ShardGroup, + route distribution.Route, start []byte, end []byte, ts uint64, @@ -1162,7 +1559,7 @@ func (s *ShardStore) scanLockOnlyKeysAtLeader( if len(lockKVs) == 0 { return nil, nil } - kvs, err := s.resolveScanLocks(ctx, g, nil, lockKVs, ts) + kvs, err := s.resolveScanLocks(ctx, g, route, nil, lockKVs, ts) if err != nil { return nil, err } @@ -1262,7 +1659,7 @@ func (s *ShardStore) clampedReverseScanRouteAtWithReadFence( return kvs, false, nil } -func (s *ShardStore) scanRouteAtDirection( +func (s *ShardStore) scanRouteAtDirectionWithReadFence( ctx context.Context, route distribution.Route, start []byte, @@ -1271,12 +1668,19 @@ func (s *ShardStore) scanRouteAtDirection( ts uint64, reverse bool, explicitGroup bool, + readRouteVersion uint64, + routeStart []byte, + routeEnd []byte, ) ([]*store.KVPair, error) { - return s.scanRouteAtDirectionWithReadFence(ctx, route, start, end, limit, ts, reverse, explicitGroup, 0, nil, nil) + if routeScanBoundsPresent(routeStart, routeEnd) { + return s.scanRouteAtDirectionWithReadFenceRouteFilter(ctx, route, start, end, limit, ts, reverse, explicitGroup, readRouteVersion, routeStart, routeEnd) + } + return s.scanRouteAtDirectionWithReadFenceOnce(ctx, route, start, end, limit, ts, reverse, explicitGroup, readRouteVersion, routeStart, routeEnd) } -func (s *ShardStore) scanRouteAtDirectionWithReadFence( +func (s *ShardStore) scanRouteAtDirectionWithS3StagedOwnerFilter( ctx context.Context, + routes []distribution.Route, route distribution.Route, start []byte, end []byte, @@ -1287,11 +1691,61 @@ func (s *ShardStore) scanRouteAtDirectionWithReadFence( readRouteVersion uint64, routeStart []byte, routeEnd []byte, + dedupeByKey bool, ) ([]*store.KVPair, error) { - if routeScanBoundsPresent(routeStart, routeEnd) { - return s.scanRouteAtDirectionWithReadFenceRouteFilter(ctx, route, start, end, limit, ts, reverse, explicitGroup, readRouteVersion, routeStart, routeEnd) + if !dedupeByKey || routeHasStagedVisibility(route) || !routesContainStagedVisibility(routes) { + return s.scanRouteAtDirectionWithReadFence(ctx, route, start, end, limit, ts, reverse, explicitGroup, readRouteVersion, routeStart, routeEnd) } - return s.scanRouteAtDirectionWithReadFenceOnce(ctx, route, start, end, limit, ts, reverse, explicitGroup, readRouteVersion, routeStart, routeEnd) + out := make([]*store.KVPair, 0, limit) + scanStart := start + scanEnd := end + for len(out) < limit { + page, err := s.scanRouteAtDirectionWithReadFence(ctx, route, scanStart, scanEnd, limit, ts, reverse, explicitGroup, readRouteVersion, routeStart, routeEnd) + if err != nil { + return nil, err + } + out = append(out, filterS3AuxiliaryKVsOwnedByStagedRoutes(page, routes)...) + if len(out) >= limit { + clear(out[limit:]) + return out[:limit], nil + } + if len(page) < limit { + return out, nil + } + lastKey := lastKVKey(page) + if lastKey == nil { + return out, nil + } + if reverse { + scanEnd = lastKey + } else { + scanStart = nextScanCursor(lastKey) + } + } + return out, nil +} + +func filterS3AuxiliaryKVsOwnedByStagedRoutes(kvs []*store.KVPair, routes []distribution.Route) []*store.KVPair { + out := make([]*store.KVPair, 0, len(kvs)) + for _, kvp := range kvs { + if kvp == nil { + continue + } + start, end, auxiliary := s3BucketAuxiliaryRouteRange(kvp.Key) + ownedByStagedRoute := false + if auxiliary { + for _, candidate := range routes { + if routeHasStagedVisibility(candidate) && migrationRouteRangesIntersect(candidate.Start, candidate.End, start, end) { + ownedByStagedRoute = true + break + } + } + } + if !ownedByStagedRoute { + out = append(out, kvp) + } + } + return out } func (s *ShardStore) scanRouteAtDirectionWithReadFenceRouteFilter( @@ -1367,18 +1821,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(filterScanInternalKVs(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, readRouteVersion, routeStart, routeEnd) @@ -1386,8 +1841,8 @@ func (s *ShardStore) scanRouteAtDirectionWithReadFenceRouteFilterPage( if err != nil { return nil, nil, err } - filtered := filterTxnInternalKVs(kvs) - return markScanRouteGroup(filtered, route.GroupID), markScanRouteGroup(kvs, route.GroupID), nil + filtered := filterScanInternalKVs(kvs) + return markScanRouteGroup(filtered, route.GroupID, markRouteGroup), markScanRouteGroup(kvs, route.GroupID, markRouteGroup), nil } func proxyScanGroupID(route distribution.Route, explicitGroup bool, readRouteVersion uint64, routeStart []byte, routeEnd []byte) uint64 { @@ -1414,22 +1869,23 @@ func (s *ShardStore) scanRouteAtDirectionWithReadFenceOnce( if !ok || g == nil || g.Store == nil { return nil, nil } + markRouteGroup := shouldMarkRouteGroupOnScan(start, explicitGroup, routeStart, routeEnd) if !reverse { return s.scanRouteAtForward(ctx, route, g, start, end, limit, ts, explicitGroup, readRouteVersion, 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(filterScanInternalKVs(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, readRouteVersion, routeStart, routeEnd) @@ -1439,7 +1895,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(filterScanInternalKVs(kvs), route.GroupID, markRouteGroup), nil } const routeFilteredScanBatchMin = 128 @@ -1524,6 +1980,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 @@ -1564,7 +2023,7 @@ func (s *ShardStore) scanRouteAtForward( if err != nil { return nil, err } - out = mergeAndTrimScanResults(out, page.kvs, limit) + out = mergeAndTrimScanResultsWithOptions(out, page.kvs, limit, false, false) if len(out) >= limit { break } @@ -1584,7 +2043,8 @@ func (s *ShardStore) scanRouteAtForward( if len(out) > limit { out = out[:limit] } - return markScanRouteGroup(out, route.GroupID), nil + markRouteGroup := shouldMarkRouteGroupOnScan(start, explicitGroup, routeStart, routeEnd) + return markScanRouteGroup(out, route.GroupID, markRouteGroup), nil } func (s *ShardStore) scanRouteAtForwardPage( @@ -1602,19 +2062,19 @@ func (s *ShardStore) scanRouteAtForwardPage( ) (scanRoutePage, error) { engine := engineForGroup(g) if engine == nil { - raw, err := s.scanRouteLocal(ctx, g, start, end, limit, ts, false) + raw, err := s.scanRouteLocal(ctx, g, route, start, end, limit, ts, false) if err != nil { return scanRoutePage{}, errors.WithStack(err) } return scanRoutePage{ - kvs: filterTxnInternalKVs(raw), + kvs: filterScanInternalKVs(raw), advanceKey: lastKVKey(raw), full: len(raw) >= limit, }, nil } if isLinearizableRaftLeader(ctx, engine) { - raw, err := g.Store.ScanAt(ctx, start, end, limit, ts) + raw, err := s.scanRouteLocal(ctx, g, route, start, end, limit, ts, false) if err != nil { return scanRoutePage{}, errors.WithStack(err) } @@ -1623,12 +2083,12 @@ func (s *ShardStore) scanRouteAtForwardPage( if err != nil { return scanRoutePage{}, err } - kvs, err := s.resolveScanLocks(ctx, g, raw, lockKVs, ts) + kvs, err := s.resolveScanLocks(ctx, g, route, raw, lockKVs, ts) if err != nil { return scanRoutePage{}, err } return scanRoutePage{ - kvs: filterTxnInternalKVs(kvs), + kvs: filterScanInternalKVs(kvs), advanceKey: lastKVKey(raw), full: len(raw) >= limit, }, nil @@ -1640,7 +2100,7 @@ func (s *ShardStore) scanRouteAtForwardPage( return scanRoutePage{}, err } return scanRoutePage{ - kvs: filterTxnInternalKVs(raw), + kvs: filterScanInternalKVs(raw), advanceKey: lastKVKey(raw), full: len(raw) >= limit, }, nil @@ -1665,18 +2125,27 @@ 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) { + kvs, err := s.scanRouteLocal(ctx, g, route, start, end, visibleLimit, ts, reverse) + return markScanRouteGroup(kvs, route.GroupID, markRouteGroup), false, err + } 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(filterScanInternalKVs(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) { + kvs, err := s.scanRouteAtLeader(ctx, g, route, start, end, visibleLimit, ts, reverse) + return markScanRouteGroup(kvs, route.GroupID, markRouteGroup), false, err + } + 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 @@ -1727,12 +2196,16 @@ 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) @@ -1744,6 +2217,7 @@ func (s *ShardStore) scanRouteLocal( func (s *ShardStore) scanRouteAtLeaderPhysicalLimit( ctx context.Context, g *ShardGroup, + route distribution.Route, start []byte, end []byte, visibleLimit int, @@ -1755,47 +2229,447 @@ func (s *ShardStore) scanRouteAtLeaderPhysicalLimit( 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 + 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 := scanLockBoundsForKVsDirection(kvs, start, end, limit, reverse) + 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) + liveVersions, err := latestCandidateVersionsAt(ctx, st, keys, ts) + if err != nil { + return nil, err + } + stagedKeys := make([][]byte, 0, len(keys)) + for _, key := range keys { + stagedKeys = append(stagedKeys, distribution.MigrationStagedDataKey(route.MigrationJobID, key)) + } + stagedVersions, err := latestCandidateVersionsAt(ctx, st, stagedKeys, ts) + if err != nil { + return nil, err + } + out := make(map[string]store.MVCCVersion, len(keys)) + for _, key := range keys { + live, liveOK := liveVersions[string(key)] + stagedKey := distribution.MigrationStagedDataKey(route.MigrationJobID, key) + staged, stagedOK := stagedVersions[string(stagedKey)] + if stagedOK { + staged.Key = bytes.Clone(key) + } + if winner, ok := newerMigrationVersion(live, liveOK, staged, stagedOK); ok { + out[string(key)] = winner + } + } + return out, nil +} + +func latestCandidateVersionsAt(ctx context.Context, st store.MVCCStore, keys [][]byte, ts uint64) (map[string]store.MVCCVersion, error) { + if len(keys) == 0 { + return map[string]store.MVCCVersion{}, nil + } + candidates := make(map[string]struct{}, len(keys)) + sortedKeys := make([][]byte, 0, len(keys)) + for _, key := range keys { + id := string(key) + if _, exists := candidates[id]; exists { + continue + } + candidates[id] = struct{}{} + sortedKeys = append(sortedKeys, key) + } + sort.Slice(sortedKeys, func(i, j int) bool { + return bytes.Compare(sortedKeys[i], sortedKeys[j]) < 0 + }) + accepted := make(map[string]struct{}, len(sortedKeys)) + result, err := st.ExportVersions(ctx, store.ExportVersionsOptions{ + StartKey: sortedKeys[0], + EndKey: prefixScanEnd(sortedKeys[len(sortedKeys)-1]), + MaxCommitTSInclusive: ts, + MaxVersions: len(sortedKeys), + MaxBytes: ^uint64(0), + MaxScannedBytes: ^uint64(0), + AcceptKey: func(key []byte) bool { + _, ok := candidates[string(key)] + return ok + }, + AcceptVersion: func(key []byte, _ []byte) bool { + id := string(key) + if _, ok := accepted[id]; ok { + return false + } + accepted[id] = struct{}{} + return true + }, + }) + if err != nil { + return nil, errors.WithStack(err) + } + if !result.Done && len(result.Versions) < len(sortedKeys) { + return nil, errors.New("staged visibility range export stopped before all candidates were examined") + } + out := make(map[string]store.MVCCVersion, len(result.Versions)) + for _, version := range result.Versions { + out[string(version.Key)] = version + } + return out, nil +} + +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) + } + 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 len(start) > 0 { + scanStart = distribution.MigrationStagedDataKey(jobID, start) } - if err != nil { - return nil, errors.WithStack(err) + scanEnd := prefixScanEnd(prefix) + if len(end) > 0 { + scanEnd = distribution.MigrationStagedDataKey(jobID, end) } - lockStart, lockEnd := scanLockBoundsForKVsDirection(kvs, start, end, limit, reverse) - 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, @@ -1809,9 +2683,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 { @@ -1823,7 +2700,7 @@ func (s *ShardStore) scanRouteAtLeaderRouteFilter( if err != nil { return nil, nil, err } - resolved, err := s.resolveScanLocks(ctx, g, filteredKVs, lockKVs, ts) + resolved, err := s.resolveScanLocks(ctx, g, route, filteredKVs, lockKVs, ts) if err == nil { sort.Slice(resolved, func(i, j int) bool { if reverse { @@ -1910,23 +2787,34 @@ func scanUserKey(kvp *store.KVPair) ([]byte, bool) { if kvp == nil || kvp.Key == nil { return nil, false } + if isMigrationStagedDataKey(kvp.Key) { + return nil, false + } if !isTxnInternalKey(kvp.Key) { return kvp.Key, true } return txnUserKeyFromLockKey(kvp.Key) } -func mergeAndTrimScanResults(out []*store.KVPair, kvs []*store.KVPair, limit int) []*store.KVPair { +func mergeAndTrimScanResultsWithOptions(out []*store.KVPair, kvs []*store.KVPair, limit int, reverse bool, dedupeByKey bool) []*store.KVPair { if len(kvs) == 0 { return out } - out = append(out, kvs...) - 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] } @@ -1960,18 +2848,30 @@ func mergeAndTrimScanKeys(out [][]byte, keys [][]byte, limit int) [][]byte { } func mergeAndTrimReverseScanResults(out []*store.KVPair, kvs []*store.KVPair, limit int) []*store.KVPair { - if len(kvs) == 0 { - return out + return mergeAndTrimScanResultsWithOptions(out, kvs, limit, true, false) +} + +func appendReplacingKVsByKey(out []*store.KVPair, kvs []*store.KVPair) []*store.KVPair { + indexByKey := make(map[string]int, len(out)+len(kvs)) + for i, kvp := range out { + if kvp == nil { + continue + } + indexByKey[string(kvp.Key)] = i } - out = append(out, kvs...) - sort.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 kvPairsFromKeys(keys [][]byte) []*store.KVPair { @@ -2049,34 +2949,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)) } @@ -2088,18 +3012,14 @@ func (s *ShardStore) LatestCommitTSWithReadFence(ctx context.Context, key []byte if err := s.awaitReadRouteVersion(ctx, readRouteVersion); err != nil { return 0, false, err } - route, routeVersion, ok := s.engine.GetRouteWithVersion(routeKey(key)) + route, g, routeVersion, ok := s.routeAndGroupForKeyWithVersion(key) readRouteVersion = max(readRouteVersion, routeVersion) - if !ok { - return 0, false, nil - } - g, ok := s.groupForID(route.GroupID) if !ok || g.Store == nil { return 0, false, nil } if engineForGroup(g) == nil { - ts, exists, err := g.Store.LatestCommitTS(ctx, key) + ts, exists, err := s.localLatestCommitTS(ctx, g, route, key) if err != nil { return 0, false, errors.WithStack(err) } @@ -2111,7 +3031,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) } @@ -2122,6 +3042,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 { @@ -2247,7 +3191,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 } @@ -2262,7 +3206,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) { @@ -2301,7 +3245,7 @@ func (s *ShardStore) planScanLockFromLockKVP(ctx context.Context, plan *scanLock } func (s *ShardStore) planScanLockItem(ctx context.Context, g *ShardGroup, ts uint64, plan *scanLockPlan, kvp *store.KVPair) error { - if kvp == nil || isTxnInternalKey(kvp.Key) { + if kvp == nil || isScanInternalKey(kvp.Key) { plan.items = append(plan.items, scanItem{skip: true}) return nil } @@ -2560,7 +3504,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 { @@ -2570,7 +3514,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 @@ -2582,7 +3526,7 @@ func (s *ShardStore) materializeScanLockResults(ctx context.Context, g *ShardGro return out, nil } -func filterTxnInternalKVs(kvs []*store.KVPair) []*store.KVPair { +func filterScanInternalKVs(kvs []*store.KVPair) []*store.KVPair { if len(kvs) == 0 { return kvs } @@ -2591,7 +3535,7 @@ func filterTxnInternalKVs(kvs []*store.KVPair) []*store.KVPair { if kvp == nil { continue } - if isTxnInternalKey(kvp.Key) { + if isScanInternalKey(kvp.Key) { continue } out = append(out, kvp) @@ -2599,8 +3543,12 @@ func filterTxnInternalKVs(kvs []*store.KVPair) []*store.KVPair { return out } -func markScanRouteGroup(kvs []*store.KVPair, groupID uint64) []*store.KVPair { - if groupID == 0 { +func isScanInternalKey(key []byte) bool { + return isTxnInternalKey(key) || isMigrationStagedDataKey(key) +} + +func markScanRouteGroup(kvs []*store.KVPair, groupID uint64, mark bool) []*store.KVPair { + if !mark || groupID == 0 { return kvs } for _, kvp := range kvs { @@ -2791,6 +3739,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)) } @@ -2801,6 +3754,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)) } @@ -2812,9 +3770,113 @@ 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 + } + out := readKeys + copied := false + for _, key := range readKeys { + alias, ok := s.stagedVisibilityReadKeyAlias(group, key) + if !ok { + continue + } + if !copied { + out = append([][]byte(nil), readKeys...) + copied = true + } + out = append(out, alias) + } + return out +} + +func (s *ShardStore) readKeysWithStagedVisibilityMutationAliases(group *ShardGroup, readKeys [][]byte, mutations []*store.KVPairMutation) [][]byte { + out := readKeys + copied := false + for _, mut := range mutations { + if mut == nil { + continue + } + alias, ok := s.stagedVisibilityReadKeyAlias(group, mut.Key) + if !ok { + continue + } + if !copied { + out = append([][]byte(nil), readKeys...) + copied = true + } + out = append(out, alias) + } + return out +} + +func (s *ShardStore) stagedVisibilityReadKeyAlias(group *ShardGroup, key []byte) ([]byte, bool) { + if s == nil || s.engine == nil || group == nil || len(key) == 0 { + return nil, false + } + if _, _, ok := distribution.MigrationStagedDataKeyParts(key); ok { + return nil, false + } + route, g, ok := s.routeAndGroupForKey(key) + if !ok || g != group || !routeHasStagedVisibility(route) { + return nil, false + } + return distribution.MigrationStagedDataKey(route.MigrationJobID, key), true +} + // 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 @@ -2841,6 +3903,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 @@ -2849,11 +3914,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 @@ -2862,6 +3985,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 } @@ -2880,6 +4008,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 @@ -2898,6 +4029,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 } @@ -3037,6 +4173,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 } @@ -3078,12 +4218,58 @@ func (s *ShardStore) closeGroup(g *ShardGroup) error { } func (s *ShardStore) groupForKey(key []byte) (*ShardGroup, bool) { - route, ok := s.engine.GetRoute(routeKey(key)) + _, g, ok := s.routeAndGroupForKey(key) + return g, ok +} + +func (s *ShardStore) routeAndGroupForKey(key []byte) (distribution.Route, *ShardGroup, bool) { + route, g, _, ok := s.routeAndGroupForKeyWithVersion(key) + return route, g, ok +} + +func (s *ShardStore) routeAndGroupForKeyWithVersion(key []byte) (distribution.Route, *ShardGroup, uint64, bool) { + if s == nil || s.engine == nil { + return distribution.Route{}, nil, 0, false + } + if start, end, auxiliary := s3BucketAuxiliaryRouteRange(key); auxiliary { + routes, version := s.engine.GetIntersectingRoutesWithVersion(nil, nil) + for _, route := range routes { + if routeHasStagedVisibility(route) && migrationRouteRangesIntersect(route.Start, route.End, start, end) { + g, ok := s.groups[route.GroupID] + return route, g, version, ok + } + } + normalizedKey := routeKey(key) + for _, route := range routes { + if routeContainsKey(route, normalizedKey) { + g, ok := s.groups[route.GroupID] + return route, g, version, ok + } + } + return distribution.Route{}, nil, version, false + } + route, version, ok := s.engine.GetRouteWithVersion(routeKey(key)) if !ok { - return nil, false + return distribution.Route{}, nil, version, false } g, ok := s.groups[route.GroupID] - return g, ok + return route, g, version, 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 31a3f53fc..bdd25e85f 100644 --- a/kv/shard_store_test.go +++ b/kv/shard_store_test.go @@ -1,7 +1,9 @@ package kv import ( + "bytes" "context" + "fmt" "sync" "testing" "time" @@ -15,6 +17,883 @@ import ( "github.com/stretchr/testify/require" ) +type exportCountingStore struct { + store.MVCCStore + exportCalls int +} + +func (s *exportCountingStore) ExportVersions(ctx context.Context, opts store.ExportVersionsOptions) (store.ExportVersionsResult, error) { + s.exportCalls++ + return s.MVCCStore.ExportVersions(ctx, opts) +} + +func newStagedVisibilityShardStore(t *testing.T) (*ShardStore, *ShardGroup) { + t.Helper() + + engine := distribution.NewEngine() + require.NoError(t, engine.ApplySnapshot(distribution.CatalogSnapshot{ + Version: 1, + Routes: []distribution.RouteDescriptor{ + { + RouteID: 1, + Start: []byte("a"), + End: []byte("z"), + GroupID: 1, + State: distribution.RouteStateActive, + StagedVisibilityActive: true, + MigrationJobID: 9, + MinWriteTSExclusive: 100, + }, + }, + })) + group := &ShardGroup{Store: store.NewMVCCStore()} + return NewShardStore(engine, map[uint64]*ShardGroup{1: group}), group +} + +func newStagedVisibilityPebbleShardStore(t *testing.T) (*ShardStore, *ShardGroup) { + t.Helper() + + engine := distribution.NewEngine() + require.NoError(t, engine.ApplySnapshot(distribution.CatalogSnapshot{ + Version: 1, + Routes: []distribution.RouteDescriptor{ + { + RouteID: 1, + Start: []byte("a"), + End: []byte("z"), + GroupID: 1, + State: distribution.RouteStateActive, + StagedVisibilityActive: true, + MigrationJobID: 9, + MinWriteTSExclusive: 100, + }, + }, + })) + st, err := store.NewPebbleStore(t.TempDir()) + require.NoError(t, err) + t.Cleanup(func() { require.NoError(t, st.Close()) }) + group := &ShardGroup{Store: st} + return NewShardStore(engine, map[uint64]*ShardGroup{1: group}), group +} + +func TestShardStoreGetAt_MergesStagedVisibility(t *testing.T) { + t.Parallel() + + ctx := context.Background() + st, group := newStagedVisibilityShardStore(t) + rawKey := []byte("k") + stagedKey := distribution.MigrationStagedDataKey(9, rawKey) + + require.NoError(t, group.Store.PutAt(ctx, rawKey, []byte("live-old"), 10, 0)) + require.NoError(t, group.Store.PutAt(ctx, stagedKey, []byte("staged-new"), 20, 0)) + got, err := st.GetAt(ctx, rawKey, 25) + require.NoError(t, err) + require.Equal(t, []byte("staged-new"), got) + + require.NoError(t, group.Store.PutAt(ctx, rawKey, []byte("live-new"), 30, 0)) + got, err = st.GetAt(ctx, rawKey, 35) + require.NoError(t, err) + require.Equal(t, []byte("live-new"), got) + + require.NoError(t, group.Store.DeleteAt(ctx, stagedKey, 40)) + _, err = st.GetAt(ctx, rawKey, 45) + require.ErrorIs(t, err, store.ErrKeyNotFound) +} + +func TestShardStoreGetAt_MergesStagedVisibilityPebbleExactKey(t *testing.T) { + t.Parallel() + + ctx := context.Background() + st, group := newStagedVisibilityPebbleShardStore(t) + rawKey := []byte("k") + stagedKey := distribution.MigrationStagedDataKey(9, rawKey) + + require.NoError(t, group.Store.PutAt(ctx, rawKey, []byte("live-old"), 10, 0)) + require.NoError(t, group.Store.PutAt(ctx, stagedKey, []byte("staged-new"), 20, 0)) + + got, err := st.GetAt(ctx, rawKey, 25) + require.NoError(t, err) + require.Equal(t, []byte("staged-new"), got) +} + +func TestShardStoreGetAt_MergesStagedVisibilityForS3BucketAuxiliary(t *testing.T) { + t.Parallel() + + ctx := context.Background() + const bucket = "bucket-a" + routeStart := s3keys.RoutePrefixForBucketAnyGeneration(bucket) + routeEnd := prefixScanEnd(routeStart) + engine := distribution.NewEngine() + require.NoError(t, engine.ApplySnapshot(distribution.CatalogSnapshot{ + Version: 1, + Routes: []distribution.RouteDescriptor{ + { + RouteID: 1, + Start: routeStart, + End: routeEnd, + GroupID: 2, + State: distribution.RouteStateActive, + StagedVisibilityActive: true, + MigrationJobID: 9, + }, + }, + })) + group := &ShardGroup{Store: store.NewMVCCStore()} + st := NewShardStore(engine, map[uint64]*ShardGroup{2: group}) + + for _, tc := range []struct { + name string + key []byte + value []byte + }{ + {name: "bucket meta", key: s3keys.BucketMetaKey(bucket), value: []byte("meta")}, + {name: "bucket generation", key: s3keys.BucketGenerationKey(bucket), value: []byte("generation")}, + } { + t.Run(tc.name, func(t *testing.T) { + require.NoError(t, group.Store.PutAt(ctx, distribution.MigrationStagedDataKey(9, tc.key), tc.value, 20, 0)) + + got, err := st.GetAt(ctx, tc.key, 25) + require.NoError(t, err) + require.Equal(t, tc.value, got) + }) + } +} + +func 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 TestShardStoreStagedVisibilityScanUsesTwoRangeExports(t *testing.T) { + t.Parallel() + + ctx := context.Background() + st, group := newStagedVisibilityShardStore(t) + counting := &exportCountingStore{MVCCStore: group.Store} + group.Store = counting + require.NoError(t, group.Store.PutAt(ctx, []byte("b"), []byte("live"), 10, 0)) + require.NoError(t, group.Store.PutAt(ctx, distribution.MigrationStagedDataKey(9, []byte("c")), []byte("staged"), 20, 0)) + + kvs, err := st.ScanAt(ctx, []byte("a"), []byte("z"), 10, 30) + require.NoError(t, err) + require.Equal(t, []*store.KVPair{ + {Key: []byte("b"), Value: []byte("live")}, + {Key: []byte("c"), Value: []byte("staged")}, + }, kvs) + require.Equal(t, 2, counting.exportCalls) +} + +func TestStagedVisibilityScanBoundsTreatsEmptyEndAsUnbounded(t *testing.T) { + t.Parallel() + + prefix := distribution.MigrationStagedDataKeyPrefix(9) + start, end := stagedVisibilityScanBounds(9, []byte{}, []byte{}) + require.Equal(t, prefix, start) + require.Equal(t, prefixScanEnd(prefix), end) +} + +func TestShardStoreScanAt_FiltersStagedShadowRowsFromLiveCandidates(t *testing.T) { + t.Parallel() + + ctx := context.Background() + st, group := newStagedVisibilityShardStore(t) + rawKey := []byte("b") + require.NoError(t, group.Store.PutAt(ctx, distribution.MigrationStagedDataKey(9, rawKey), []byte("staged"), 20, 0)) + + kvs, err := st.ScanAt(ctx, []byte(""), nil, 10, 30) + require.NoError(t, err) + require.Equal(t, []*store.KVPair{{Key: rawKey, Value: []byte("staged")}}, kvs) +} + +func TestShardStoreScanAt_PreservesNonStagedRoutesDuringBroadStagedVisibilityScan(t *testing.T) { + t.Parallel() + + ctx := context.Background() + engine := distribution.NewEngine() + require.NoError(t, engine.ApplySnapshot(distribution.CatalogSnapshot{ + Version: 1, + Routes: []distribution.RouteDescriptor{ + {RouteID: 1, Start: []byte("a"), End: []byte("m"), GroupID: 1, State: distribution.RouteStateActive}, + { + RouteID: 2, + Start: []byte("m"), + End: []byte("t"), + GroupID: 1, + State: distribution.RouteStateActive, + StagedVisibilityActive: true, + MigrationJobID: 9, + }, + {RouteID: 3, Start: []byte("t"), End: []byte("z"), GroupID: 1, State: distribution.RouteStateActive}, + }, + })) + group := &ShardGroup{Store: store.NewMVCCStore()} + st := NewShardStore(engine, map[uint64]*ShardGroup{1: group}) + + require.NoError(t, group.Store.PutAt(ctx, []byte("b"), []byte("live-b"), 10, 0)) + require.NoError(t, group.Store.PutAt(ctx, []byte("n"), []byte("live-n"), 10, 0)) + require.NoError(t, group.Store.PutAt(ctx, []byte("u"), []byte("live-u"), 10, 0)) + require.NoError(t, group.Store.PutAt(ctx, distribution.MigrationStagedDataKey(9, []byte("n")), []byte("staged-n"), 20, 0)) + require.NoError(t, group.Store.PutAt(ctx, distribution.MigrationStagedDataKey(9, []byte("o")), []byte("staged-o"), 20, 0)) + + kvs, err := st.ScanAt(ctx, nil, nil, 10, 30) + require.NoError(t, err) + require.Equal(t, []*store.KVPair{ + {Key: []byte("b"), Value: []byte("live-b")}, + {Key: []byte("n"), Value: []byte("staged-n")}, + {Key: []byte("o"), Value: []byte("staged-o")}, + {Key: []byte("u"), Value: []byte("live-u")}, + }, kvs) + + kvs, err = st.ReverseScanAt(ctx, nil, nil, 10, 30) + require.NoError(t, err) + require.Equal(t, []*store.KVPair{ + {Key: []byte("u"), Value: []byte("live-u")}, + {Key: []byte("o"), Value: []byte("staged-o")}, + {Key: []byte("n"), Value: []byte("staged-n")}, + {Key: []byte("b"), Value: []byte("live-b")}, + }, kvs) +} + +func TestShardStoreScanAt_RoutesS3BucketAuxiliaryStagedVisibility(t *testing.T) { + t.Parallel() + + ctx := context.Background() + const bucket = "bucket-a" + routeStart := s3keys.RoutePrefixForBucketAnyGeneration(bucket) + routeEnd := prefixScanEnd(routeStart) + engine := distribution.NewEngine() + require.NoError(t, engine.ApplySnapshot(distribution.CatalogSnapshot{ + Version: 1, + Routes: []distribution.RouteDescriptor{ + {RouteID: 1, Start: []byte(""), End: routeStart, GroupID: 1, State: distribution.RouteStateActive}, + { + RouteID: 2, + Start: routeStart, + End: routeEnd, + GroupID: 2, + State: distribution.RouteStateActive, + StagedVisibilityActive: true, + MigrationJobID: 9, + }, + {RouteID: 3, Start: routeEnd, End: nil, GroupID: 1, State: distribution.RouteStateActive}, + }, + })) + groups := map[uint64]*ShardGroup{ + 1: {Store: store.NewMVCCStore()}, + 2: {Store: store.NewMVCCStore()}, + } + st := NewShardStore(engine, groups) + + for _, tc := range []struct { + name string + prefix string + key []byte + value []byte + }{ + {name: "bucket meta", prefix: s3keys.BucketMetaPrefix, key: s3keys.BucketMetaKey(bucket), value: []byte("meta")}, + {name: "bucket generation", prefix: s3keys.BucketGenerationPrefix, key: s3keys.BucketGenerationKey(bucket), value: []byte("generation")}, + } { + t.Run(tc.name, func(t *testing.T) { + require.NoError(t, groups[2].Store.PutAt(ctx, distribution.MigrationStagedDataKey(9, tc.key), tc.value, 20, 0)) + + kvs, err := st.ScanAt(ctx, []byte(tc.prefix), prefixScanEnd([]byte(tc.prefix)), 10, 30) + require.NoError(t, err) + require.Contains(t, kvs, &store.KVPair{Key: tc.key, Value: tc.value}) + }) + } +} + +func TestShardStoreS3BucketAuxiliaryScanHonorsStagedTombstone(t *testing.T) { + t.Parallel() + + ctx := context.Background() + const migratedBucket = "bucket-a" + routeStart := s3keys.RoutePrefixForBucketAnyGeneration(migratedBucket) + routeEnd := prefixScanEnd(routeStart) + engine := distribution.NewEngine() + require.NoError(t, engine.ApplySnapshot(distribution.CatalogSnapshot{ + Version: 1, + Routes: []distribution.RouteDescriptor{ + {RouteID: 1, Start: []byte(""), End: routeStart, GroupID: 1, State: distribution.RouteStateActive}, + {RouteID: 2, Start: routeStart, End: routeEnd, GroupID: 2, State: distribution.RouteStateActive, StagedVisibilityActive: true, MigrationJobID: 9}, + {RouteID: 3, Start: routeEnd, End: nil, GroupID: 1, State: distribution.RouteStateActive}, + }, + })) + groups := map[uint64]*ShardGroup{ + 1: {Store: store.NewMVCCStore()}, + 2: {Store: store.NewMVCCStore()}, + } + st := NewShardStore(engine, groups) + deletedKey := s3keys.BucketMetaKey(migratedBucket) + visibleKey := s3keys.BucketMetaKey("bucket-z") + require.NoError(t, groups[1].Store.PutAt(ctx, deletedKey, []byte("stale"), 10, 0)) + require.NoError(t, groups[1].Store.PutAt(ctx, visibleKey, []byte("visible"), 10, 0)) + require.NoError(t, groups[2].Store.DeleteAt(ctx, distribution.MigrationStagedDataKey(9, deletedKey), 20)) + + start := []byte(s3keys.BucketMetaPrefix) + end := prefixScanEnd(start) + kvs, err := st.ScanAt(ctx, start, end, 1, 30) + require.NoError(t, err) + require.Equal(t, []*store.KVPair{{Key: visibleKey, Value: []byte("visible")}}, kvs) + + kvs, err = st.ScanAt(ctx, deletedKey, prefixScanEnd(deletedKey), 1, 30) + require.NoError(t, err) + require.Empty(t, kvs) +} + +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 TestShardStorePhysicalLimitFallsBackToStagedVisibilityScan(t *testing.T) { + t.Parallel() + + ctx := context.Background() + st, group := newStagedVisibilityShardStore(t) + require.NoError(t, group.Store.PutAt(ctx, []byte("b/live"), []byte("live"), 10, 0)) + require.NoError(t, group.Store.PutAt(ctx, distribution.MigrationStagedDataKey(9, []byte("b/staged")), []byte("staged"), 20, 0)) + + kvs, limitReached, err := st.ScanAtPhysicalLimit(ctx, []byte("b"), []byte("c"), 10, 10, 50) + require.NoError(t, err) + require.False(t, limitReached) + require.Equal(t, []*store.KVPair{ + {Key: []byte("b/live"), Value: []byte("live")}, + {Key: []byte("b/staged"), Value: []byte("staged")}, + }, kvs) + + kvs, limitReached, err = st.ReverseScanAtPhysicalLimit(ctx, []byte("b"), []byte("c"), 10, 10, 50) + require.NoError(t, err) + require.False(t, limitReached) + require.Equal(t, []*store.KVPair{ + {Key: []byte("b/staged"), Value: []byte("staged")}, + {Key: []byte("b/live"), Value: []byte("live")}, + }, kvs) +} + +func TestShardStoreRejectsWritesAtMigrationTimestampFloor(t *testing.T) { + t.Parallel() + + ctx := context.Background() + st, _ := newStagedVisibilityShardStore(t) + + err := st.PutAt(ctx, []byte("k"), []byte("low"), 100, 0) + require.ErrorIs(t, err, ErrRouteWriteTimestampTooLow) + require.NoError(t, st.PutAt(ctx, []byte("k"), []byte("ok"), 101, 0)) +} + +func TestShardStoreRaftApplyRejectsMigrationTimestampFloor(t *testing.T) { + t.Parallel() + + ctx := context.Background() + st, _ := newStagedVisibilityShardStore(t) + + require.ErrorIs(t, st.ApplyMutationsRaft(ctx, []*store.KVPairMutation{ + {Op: store.OpTypePut, Key: []byte("k-raft"), Value: []byte("v")}, + }, nil, 90, 100), ErrRouteWriteTimestampTooLow) + require.ErrorIs(t, st.ApplyMutationsRaftAt(ctx, []*store.KVPairMutation{ + {Op: store.OpTypePut, Key: []byte("k-raft-at"), Value: []byte("v")}, + }, nil, 90, 100, 1), ErrRouteWriteTimestampTooLow) + require.ErrorIs(t, st.DeletePrefixAtRaft(ctx, []byte("k-raft"), nil, 100), ErrRouteWriteTimestampTooLow) + require.ErrorIs(t, st.DeletePrefixAtRaftAt(ctx, []byte("k-raft-at"), nil, 100, 2), ErrRouteWriteTimestampTooLow) +} + type followerProxyEngine struct { leader string } @@ -237,6 +1116,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() @@ -415,7 +1342,7 @@ func TestShardStoreScanAtRoutesWideColumnPrefixesByUserKey(t *testing.T) { require.NoError(t, st.PutAt(ctx, tc.key, []byte("value"), 20, 0)) kvs, err := st.ScanAt(ctx, tc.prefix, prefixScanEnd(tc.prefix), 10, 20) require.NoError(t, err) - require.Equal(t, []*store.KVPair{{Key: tc.key, Value: []byte("value"), RouteGroupID: 2}}, kvs) + require.Equal(t, []*store.KVPair{{Key: tc.key, Value: []byte("value")}}, kvs) _, err = groups[1].Store.GetAt(ctx, tc.key, 20) require.ErrorIs(t, err, store.ErrKeyNotFound) }) @@ -985,7 +1912,7 @@ func TestShardStoreScanKeysRouteAtLeaderRefillsAfterTxnInternalKeys(t *testing.T require.NoError(t, g.Store.PutAt(ctx, txnCommitKey([]byte("primary"), 10), []byte("commit"), 1, 0)) require.NoError(t, g.Store.PutAt(ctx, []byte("a"), []byte("va"), 2, 0)) - keys, err := st.scanKeysRouteAtLeader(ctx, g, []byte(""), nil, 1, ^uint64(0)) + keys, err := st.scanKeysRouteAtLeader(ctx, g, distribution.Route{GroupID: 1}, []byte(""), nil, 1, ^uint64(0)) require.NoError(t, err) require.Equal(t, [][]byte{[]byte("a")}, keys) } @@ -1000,11 +1927,57 @@ func TestShardStoreScanKeysRouteAtLeaderPreservesEmptyKey(t *testing.T) { require.NoError(t, g.Store.PutAt(ctx, []byte(""), []byte("empty"), 1, 0)) require.NoError(t, g.Store.PutAt(ctx, []byte("a"), []byte("va"), 2, 0)) - keys, err := st.scanKeysRouteAtLeader(ctx, g, nil, nil, 2, ^uint64(0)) + keys, err := st.scanKeysRouteAtLeader(ctx, g, distribution.Route{GroupID: 1}, nil, nil, 2, ^uint64(0)) require.NoError(t, err) require.Equal(t, [][]byte{[]byte(""), []byte("a")}, keys) } +func TestShardStoreScanKeysRouteAtLeaderIncludesStagedOnlyKeys(t *testing.T) { + t.Parallel() + + ctx := context.Background() + g := &ShardGroup{Store: store.NewMVCCStore()} + st := NewShardStore(distribution.NewEngine(), map[uint64]*ShardGroup{1: g}) + route := distribution.Route{ + GroupID: 1, + StagedVisibilityActive: true, + MigrationJobID: 9, + } + key := []byte("staged-key") + require.NoError(t, g.Store.PutAt(ctx, distribution.MigrationStagedDataKey(route.MigrationJobID, key), []byte("value"), 1, 0)) + + keys, err := st.scanKeysRouteAtLeader(ctx, g, route, []byte(""), nil, 10, ^uint64(0)) + require.NoError(t, err) + require.Equal(t, [][]byte{key}, keys) +} + +func TestShardStoreScanKeysAtIncludesStagedOnlyKeys(t *testing.T) { + t.Parallel() + + ctx := context.Background() + engine := distribution.NewEngine() + require.NoError(t, engine.ApplySnapshot(distribution.CatalogSnapshot{ + Version: 1, + Routes: []distribution.RouteDescriptor{{ + RouteID: 1, + Start: []byte(""), + End: nil, + GroupID: 1, + State: distribution.RouteStateActive, + StagedVisibilityActive: true, + MigrationJobID: 9, + }}, + })) + g := &ShardGroup{Store: store.NewMVCCStore()} + st := NewShardStore(engine, map[uint64]*ShardGroup{1: g}) + key := []byte("staged-key") + require.NoError(t, g.Store.PutAt(ctx, distribution.MigrationStagedDataKey(9, key), []byte("value"), 1, 0)) + + keys, err := st.ScanKeysAt(ctx, []byte(""), nil, 10, ^uint64(0)) + require.NoError(t, err) + require.Equal(t, [][]byte{key}, keys) +} + func TestShardStoreProxyScanKeysAtUsesSelectedGroup(t *testing.T) { t.Parallel() diff --git a/kv/sharded_coordinator.go b/kv/sharded_coordinator.go index a016ee611..9954c1785 100644 --- a/kv/sharded_coordinator.go +++ b/kv/sharded_coordinator.go @@ -1159,6 +1159,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{ @@ -1277,6 +1280,63 @@ 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 + } + if err := c.rejectWriteTimestampFloorDelPrefix(elem.Key, commitTS); err != nil { + return err + } + } + return nil +} + +func (c *ShardedCoordinator) rejectWriteTimestampFloorDelPrefix(prefix []byte, commitTS uint64) error { + start, end := routePrefixRange(prefix) + for _, route := range c.engine.GetIntersectingRoutes(start, end) { + if route.MinWriteTSExclusive != 0 && commitTS <= route.MinWriteTSExclusive { + return errors.Wrapf(ErrRouteWriteTimestampTooLow, "prefix %q route range [%q,%q) commit_ts=%d floor=%d", prefix, start, end, commitTS, route.MinWriteTSExclusive) + } + } + return nil +} + // broadcastToAllGroups sends the same set of requests to every configured // all-shard data group in parallel and returns the maximum commit index. func (c *ShardedCoordinator) broadcastToAllGroups(ctx context.Context, requests []*pb.Request) (*CoordinateResponse, error) { @@ -1342,6 +1402,9 @@ func (c *ShardedCoordinator) dispatchTxn(ctx context.Context, startTS uint64, co if err := ValidateElemCommitTSPatches(elems, commitTS); 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. @@ -1385,6 +1448,10 @@ func (c *ShardedCoordinator) dispatchMultiShardTxn(ctx context.Context, startTS, if err != nil { return nil, err } + groupedReadKeys = c.groupedReadKeysWithStagedVisibilityMutationAliases(groupedReadKeys, grouped) + if groupedReadKeyCount(groupedReadKeys) > maxReadKeys { + return nil, errors.WithStack(ErrInvalidRequest) + } prepared, err := c.prewriteTxn(ctx, startTS, commitTS, primaryKey, grouped, gids, groupedReadKeys, observedRouteVersion, bypassKeysByGroup) if err != nil { return nil, err @@ -1459,6 +1526,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 @@ -1475,6 +1547,48 @@ func (c *ShardedCoordinator) dispatchSingleShardTxn(ctx context.Context, startTS return &CoordinateResponse{CommitIndex: resp.CommitIndex, CommitTS: commitTS}, nil } +func (c *ShardedCoordinator) readKeysWithStagedVisibilityAliasesForGroup(gid uint64, readKeys [][]byte) [][]byte { + if len(readKeys) == 0 { + return readKeys + } + var out [][]byte + for _, key := range readKeys { + alias, ok := c.stagedVisibilityReadKeyAlias(gid, key) + if !ok { + continue + } + if out == nil { + out = append([][]byte(nil), readKeys...) + } + out = append(out, alias) + } + if out == nil { + return readKeys + } + return out +} + +func (c *ShardedCoordinator) readKeysWithStagedVisibilityMutationAliasesForGroup(gid uint64, readKeys [][]byte, elems []*Elem[OP]) [][]byte { + var out [][]byte + for _, elem := range elems { + if elem == nil { + continue + } + alias, ok := c.stagedVisibilityReadKeyAlias(gid, elem.Key) + if !ok { + continue + } + if out == nil { + out = append([][]byte(nil), readKeys...) + } + out = append(out, alias) + } + if out == nil { + return readKeys + } + return out +} + type preparedGroup struct { gid uint64 keys []*pb.Mutation @@ -2112,7 +2226,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 } @@ -2120,21 +2234,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 @@ -2166,7 +2307,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-"+ @@ -2174,10 +2315,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. @@ -2227,7 +2419,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) } @@ -2238,6 +2430,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 { @@ -2277,6 +2506,9 @@ func (c *ShardedCoordinator) rawLogsWithGroups(ctx context.Context, reqs *Operat if err != nil { return nil, nil, err } + if err := c.rejectWriteTimestampFloorMutations(grouped[gid], ts); err != nil { + return nil, nil, err + } logs = append(logs, &pb.Request{ IsTxn: false, Phase: pb.Phase_NONE, @@ -2289,6 +2521,30 @@ func (c *ShardedCoordinator) rawLogsWithGroups(ctx context.Context, reqs *Operat return logs, gids, nil } +func (c *ShardedCoordinator) rejectWriteTimestampFloorMutations(muts []*pb.Mutation, commitTS uint64) error { + if c == nil || c.engine == nil || commitTS == 0 { + return nil + } + for _, mut := range muts { + if mut == nil { + continue + } + if mut.GetOp() == pb.Op_DEL_PREFIX { + if err := c.rejectWriteTimestampFloorDelPrefix(mut.Key, commitTS); err != nil { + return err + } + continue + } + if 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 @@ -2306,6 +2562,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 } @@ -2383,21 +2642,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 4c4e6d8dc..b53583056 100644 --- a/kv/sharded_coordinator_del_prefix_test.go +++ b/kv/sharded_coordinator_del_prefix_test.go @@ -122,6 +122,72 @@ 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 TestShardedCoordinatorRejectsRawDelPrefixMutationAtMigrationTimestampFloor(t *testing.T) { + t.Parallel() + + coord := NewShardedCoordinator(newMigrationFloorEngine(t, ^uint64(0)), map[uint64]*ShardGroup{ + 1: {Txn: &recordingTransactional{}}, + }, 1, NewHLC(), nil) + + for _, mut := range []*pb.Mutation{ + {Op: pb.Op_DEL_PREFIX, Key: []byte("z")}, + {Op: pb.Op_DEL_PREFIX, Key: nil}, + } { + err := coord.rejectWriteTimestampFloorMutations([]*pb.Mutation{mut}, 100) + require.ErrorIs(t, err, ErrRouteWriteTimestampTooLow) + } +} + func TestShardedCoordinator_DelPrefixDoesNotAutoPinObservedRouteVersion(t *testing.T) { t.Parallel() @@ -379,6 +445,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 ec6c651ff..96614f4b3 100644 --- a/kv/sharded_coordinator_txn_test.go +++ b/kv/sharded_coordinator_txn_test.go @@ -9,6 +9,7 @@ import ( "github.com/bootjp/elastickv/distribution" "github.com/bootjp/elastickv/internal/raftengine" + "github.com/bootjp/elastickv/internal/s3keys" "github.com/bootjp/elastickv/keyviz" pb "github.com/bootjp/elastickv/proto" "github.com/bootjp/elastickv/store" @@ -49,6 +50,158 @@ func (s *recordingTransactional) Abort(_ context.Context, _ []*pb.Request) (*Tra return &TransactionResponse{}, nil } +func TestShardedCoordinatorValidateReadKeysOnShard_UsesStagedVisibility(t *testing.T) { + t.Parallel() + + ctx := context.Background() + engine := distribution.NewEngine() + require.NoError(t, engine.ApplySnapshot(distribution.CatalogSnapshot{ + Version: 1, + Routes: []distribution.RouteDescriptor{ + { + RouteID: 1, + Start: []byte("a"), + End: []byte("z"), + GroupID: 1, + State: distribution.RouteStateActive, + StagedVisibilityActive: true, + MigrationJobID: 9, + }, + }, + })) + st := store.NewMVCCStore() + readKey := []byte("k") + require.NoError(t, st.PutAt(ctx, distribution.MigrationStagedDataKey(9, readKey), []byte("staged"), 20, 0)) + coord := NewShardedCoordinator(engine, map[uint64]*ShardGroup{ + 1: {Engine: stubLeaderEngine{}, Store: st}, + }, 1, NewHLC(), nil) + + err := coord.validateReadKeysOnShard(ctx, 1, [][]byte{readKey}, 10) + require.ErrorIs(t, err, store.ErrWriteConflict) +} + +func TestShardedCoordinatorDispatchTxn_AddsStagedReadKeyAlias(t *testing.T) { + t.Parallel() + + engine := distribution.NewEngine() + require.NoError(t, engine.ApplySnapshot(distribution.CatalogSnapshot{ + Version: 1, + Routes: []distribution.RouteDescriptor{ + { + RouteID: 1, + Start: []byte("a"), + End: []byte("z"), + GroupID: 1, + State: distribution.RouteStateActive, + StagedVisibilityActive: true, + MigrationJobID: 9, + }, + }, + })) + txn := &recordingTransactional{} + coord := NewShardedCoordinator(engine, map[uint64]*ShardGroup{ + 1: {Txn: txn}, + }, 1, NewHLC(), nil) + + readKey := []byte("k") + _, err := coord.Dispatch(context.Background(), &OperationGroup[OP]{ + IsTxn: true, + StartTS: 10, + CommitTS: 101, + Elems: []*Elem[OP]{{Op: Put, Key: []byte("m"), Value: []byte("write")}}, + ReadKeys: [][]byte{readKey}, + }) + require.NoError(t, err) + require.Len(t, txn.requests, 1) + require.Equal(t, [][]byte{ + readKey, + distribution.MigrationStagedDataKey(9, readKey), + distribution.MigrationStagedDataKey(9, []byte("m")), + }, txn.requests[0].ReadKeys) +} + +func TestShardedCoordinatorDispatchTxn_AddsStagedWriteKeyAlias(t *testing.T) { + t.Parallel() + + engine := distribution.NewEngine() + require.NoError(t, engine.ApplySnapshot(distribution.CatalogSnapshot{ + Version: 1, + Routes: []distribution.RouteDescriptor{ + { + RouteID: 1, + Start: []byte("a"), + End: []byte("z"), + GroupID: 1, + State: distribution.RouteStateActive, + StagedVisibilityActive: true, + MigrationJobID: 9, + }, + }, + })) + txn := &recordingTransactional{} + coord := NewShardedCoordinator(engine, map[uint64]*ShardGroup{ + 1: {Txn: txn}, + }, 1, NewHLC(), nil) + + writeKey := []byte("k") + _, err := coord.Dispatch(context.Background(), &OperationGroup[OP]{ + IsTxn: true, + StartTS: 10, + CommitTS: 101, + Elems: []*Elem[OP]{{Op: Put, Key: writeKey, Value: []byte("write")}}, + }) + require.NoError(t, err) + require.Len(t, txn.requests, 1) + require.Equal(t, [][]byte{ + distribution.MigrationStagedDataKey(9, writeKey), + }, txn.requests[0].ReadKeys) +} + +func TestShardedCoordinatorPrewrite_AddsStagedWriteKeyAlias(t *testing.T) { + t.Parallel() + + engine := distribution.NewEngine() + require.NoError(t, engine.ApplySnapshot(distribution.CatalogSnapshot{ + Version: 1, + Routes: []distribution.RouteDescriptor{ + { + RouteID: 1, + Start: []byte("a"), + End: []byte("m"), + GroupID: 1, + State: distribution.RouteStateActive, + StagedVisibilityActive: true, + MigrationJobID: 9, + }, + {RouteID: 2, Start: []byte("m"), End: []byte("z"), GroupID: 2, State: distribution.RouteStateActive}, + }, + })) + g1Txn := &recordingTransactional{} + g2Txn := &recordingTransactional{} + coord := NewShardedCoordinator(engine, map[uint64]*ShardGroup{ + 1: {Txn: g1Txn}, + 2: {Txn: g2Txn}, + }, 1, NewHLC(), nil) + + writeKey := []byte("b") + _, err := coord.Dispatch(context.Background(), &OperationGroup[OP]{ + IsTxn: true, + StartTS: 10, + CommitTS: 101, + Elems: []*Elem[OP]{ + {Op: Put, Key: writeKey, Value: []byte("write-b")}, + {Op: Put, Key: []byte("x"), Value: []byte("write-x")}, + }, + }) + require.NoError(t, err) + require.NotEmpty(t, g1Txn.requests) + require.NotEmpty(t, g2Txn.requests) + require.Equal(t, [][]byte{ + distribution.MigrationStagedDataKey(9, writeKey), + }, g1Txn.requests[0].ReadKeys) + require.Empty(t, g2Txn.requests[0].ReadKeys) +} + func cloneTxnRequest(req *pb.Request) *pb.Request { if req == nil { return nil @@ -400,6 +553,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() @@ -529,6 +702,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 ef4e58510..4b9874c6a 100644 --- a/kv/tso_test.go +++ b/kv/tso_test.go @@ -336,6 +336,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 6e173f540..6c08a6bd2 100644 --- a/main.go +++ b/main.go @@ -1851,6 +1851,7 @@ func startServersAfterStartupRotation(waitRotateOnStartup startupRotationWaiter, shardStore: in.shardStore, coordinate: adapterCoordinate, distServer: in.distServer, + routeEngine: in.cfg.engine, adminServer: adminServer, adminGRPCOpts: adminGRPCOpts, redisAddress: *redisAddr, @@ -2498,6 +2499,7 @@ func startRaftServers( shardStore *kv.ShardStore, coordinate kv.Coordinator, distServer *adapter.DistributionServer, + routeEngine *distribution.Engine, relay *adapter.RedisPubSubRelay, proposalObserverForGroup func(uint64) kv.ProposalObserver, adminServer *adapter.AdminServer, @@ -2505,6 +2507,7 @@ func startRaftServers( forwardDeps adminForwardServerDeps, confChangeInterceptor internalraftadmin.MembershipChangeInterceptor, encWiring encryptionWriteWiring, + sqsPartitionResolver kv.PartitionResolver, ) error { forwardLogger := slog.Default().With(slog.String("component", "admin")) // extraOptsCap reserves slots for the unary + stream admin interceptor @@ -2538,7 +2541,13 @@ func startRaftServers( rt.engine, coordinate.Clock(), relay, - internalTimestampOptions(coordinate)..., + append( + internalTimestampOptions(coordinate), + adapter.WithInternalStore(rt.store), + adapter.WithInternalMigrationProposer(proposerForGroup(rt, shardGroups)), + adapter.WithInternalRouteEngine(routeEngine), + adapter.WithInternalMigrationExportRouting(rt.spec.id, sqsPartitionResolver), + )..., )) pb.RegisterDistributionServer(gs, distServer) if adminServer != nil { @@ -2880,6 +2889,7 @@ type runtimeServerRunner struct { shardStore *kv.ShardStore coordinate kv.Coordinator distServer *adapter.DistributionServer + routeEngine *distribution.Engine adminServer *adapter.AdminServer adminGRPCOpts adminGRPCInterceptors redisAddress string @@ -2984,6 +2994,10 @@ func (r *runtimeServerRunner) startRaftTransport() error { buckets: newBucketsSource(r.s3Server), roles: r.roleStore, } + var sqsPartitionResolver kv.PartitionResolver + if r.sqsPartitionResolver != nil { + sqsPartitionResolver = r.sqsPartitionResolver + } if err := startRaftServers( r.ctx, r.lc, @@ -2993,6 +3007,7 @@ func (r *runtimeServerRunner) startRaftTransport() error { r.shardStore, r.coordinate, r.distServer, + r.routeEngine, r.pubsubRelay, func(groupID uint64) kv.ProposalObserver { return r.metricsRegistry.RaftProposalObserver(groupID) @@ -3002,6 +3017,7 @@ func (r *runtimeServerRunner) startRaftTransport() error { forwardDeps, r.encryptionConfChangeInterceptor, r.encWiring, + sqsPartitionResolver, ); err != nil { return r.startupFailure(err) } diff --git a/main_bootstrap_e2e_test.go b/main_bootstrap_e2e_test.go index d962c8105..98a7ddb74 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, nil, 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, nil, listeners.grpc[rt.spec.id]); err != nil { return waitErrgroupAfterStartupFailure(cancel, eg, err) } } @@ -747,7 +752,9 @@ func startBoundGRPCServer( shardStore *kv.ShardStore, coordinate kv.Coordinator, distServer *adapter.DistributionServer, + routeEngine *distribution.Engine, relay *adapter.RedisPubSubRelay, + sqsPartitionResolver kv.PartitionResolver, listener net.Listener, ) error { if rt == nil || rt.engine == nil { @@ -762,7 +769,16 @@ 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), + adapter.WithInternalMigrationExportRouting(rt.spec.id, sqsPartitionResolver), + )) pb.RegisterDistributionServer(gs, distServer) rt.registerGRPC(gs) internalraftadmin.RegisterOperationalServices(ctx, gs, rt.engine, []string{"RawKV"}) diff --git a/proto/internal.pb.go b/proto/internal.pb.go index dabe9d755..33d707446 100644 --- a/proto/internal.pb.go +++ b/proto/internal.pb.go @@ -574,6 +574,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 } @@ -671,6 +680,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"` @@ -935,6 +965,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 = "" + @@ -966,7 +1140,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" + @@ -979,7 +1153,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" + @@ -1001,7 +1180,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" + @@ -1012,12 +1203,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 @@ -1032,22 +1224,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 @@ -1061,12 +1255,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 @@ -1083,7 +1279,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 27b45bd0d..7432916b8 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. @@ -104,6 +105,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 { @@ -132,3 +142,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 6b50651c5..247b13205 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) @@ -310,6 +314,14 @@ func (s *pebbleStore) decodeExportedPebbleVersion(iter *pebble.Iterator, userKey } func (s *pebbleStore) ImportVersions(ctx context.Context, opts ImportVersionsOptions) (ImportVersionsResult, error) { + return s.importVersionsWithOpts(ctx, opts, s.directApplyWriteOpts(), true) +} + +func (s *pebbleStore) ImportVersionsRaft(ctx context.Context, opts ImportVersionsOptions) (ImportVersionsResult, error) { + return s.importVersionsWithOpts(ctx, opts, s.raftApplyWriteOpts(), false) +} + +func (s *pebbleStore) importVersionsWithOpts(ctx context.Context, opts ImportVersionsOptions, writeOpts *pebble.WriteOptions, gateRegistration bool) (ImportVersionsResult, error) { s.dbMu.RLock() defer s.dbMu.RUnlock() @@ -321,11 +333,14 @@ func (s *pebbleStore) ImportVersions(ctx context.Context, opts ImportVersionsOpt return ImportVersionsResult{}, err } if duplicate { + if err := s.commitPebbleImportAppliedIndex(opts.AppliedIndex, writeOpts); err != nil { + return ImportVersionsResult{}, err + } return ImportVersionsResult{AckedCursor: ackedCursor, Duplicate: true}, nil } batchMax := importBatchMaxTS(opts.Versions) - if err := s.commitPebbleImportBatch(opts, batchMax); err != nil { + if err := s.commitPebbleImportBatch(opts, batchMax, writeOpts, gateRegistration); err != nil { return ImportVersionsResult{}, errors.WithStack(err) } s.log.InfoContext(ctx, "import_versions", @@ -338,6 +353,18 @@ func (s *pebbleStore) ImportVersions(ctx context.Context, opts ImportVersionsOpt return ImportVersionsResult{AckedCursor: bytes.Clone(opts.Cursor), MaxImportedTS: batchMax}, nil } +func (s *pebbleStore) commitPebbleImportAppliedIndex(appliedIndex uint64, writeOpts *pebble.WriteOptions) error { + if appliedIndex == 0 { + return nil + } + batch := s.db.NewBatch() + defer batch.Close() + if err := setPebbleUint64InBatch(batch, metaAppliedIndexBytes, appliedIndex); err != nil { + return err + } + return errors.WithStack(batch.Commit(writeOpts)) +} + func (s *pebbleStore) validatePebbleImportBatch(opts ImportVersionsOptions) (bool, []byte, error) { existing, hasExisting, err := s.readMigrationImportAck(opts.JobID, opts.BracketID) if err != nil { @@ -358,10 +385,10 @@ func (s *pebbleStore) validatePebbleImportBatch(opts ImportVersionsOptions) (boo return false, nil, nil } -func (s *pebbleStore) commitPebbleImportBatch(opts ImportVersionsOptions, batchMax uint64) error { +func (s *pebbleStore) commitPebbleImportBatch(opts ImportVersionsOptions, batchMax uint64, writeOpts *pebble.WriteOptions, gateRegistration bool) error { batch := s.db.NewBatch() defer batch.Close() - if err := s.applyImportVersionsBatch(batch, opts.Versions); err != nil { + if err := s.applyImportVersionsBatch(batch, opts.Versions, gateRegistration); err != nil { return err } if err := s.stageMigrationImportAck(batch, opts.JobID, opts.BracketID, migrationImportAck{ @@ -375,7 +402,10 @@ func (s *pebbleStore) commitPebbleImportBatch(opts ImportVersionsOptions, batchM return err } defer unlock() - if err := batch.Commit(s.directApplyWriteOpts()); err != nil { + if err := stagePebbleAppliedIndex(batch, opts.AppliedIndex); err != nil { + return err + } + if err := batch.Commit(writeOpts); err != nil { return errors.WithStack(err) } if batchMax > 0 { @@ -384,6 +414,13 @@ func (s *pebbleStore) commitPebbleImportBatch(opts ImportVersionsOptions, batchM return nil } +func stagePebbleAppliedIndex(batch *pebble.Batch, appliedIndex uint64) error { + if appliedIndex == 0 { + return nil + } + return setPebbleUint64InBatch(batch, metaAppliedIndexBytes, appliedIndex) +} + func (s *pebbleStore) stageMigrationImportAck(batch *pebble.Batch, jobID, bracketID uint64, ack migrationImportAck) error { acks, err := s.readMigrationImportAcks() if err != nil { @@ -403,7 +440,7 @@ func (s *pebbleStore) stageMigrationClockMetadataIfNeeded(batch *pebble.Batch, j return s.stageMigrationClockMetadata(batch, jobID, batchMax) } -func (s *pebbleStore) applyImportVersionsBatch(batch *pebble.Batch, versions []MVCCVersion) error { +func (s *pebbleStore) applyImportVersionsBatch(batch *pebble.Batch, versions []MVCCVersion, gateRegistration bool) error { for _, version := range versions { k, err := encodePebbleUserVersionKey(version.Key, version.CommitTS) if err != nil { @@ -413,7 +450,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 8a30710d3..c872e63a0 100644 --- a/store/lsm_store_applied_index_test.go +++ b/store/lsm_store_applied_index_test.go @@ -1,6 +1,7 @@ package store import ( + "bytes" "context" "encoding/binary" "os" @@ -97,6 +98,118 @@ func TestApplyMutationsRaftAt_BundlesMetaAppliedIndex(t *testing.T) { require.Equal(t, []byte("v1"), val) } +func TestPromoteVersions_BundlesMetaAppliedIndex(t *testing.T) { + ctx := context.Background() + st := newApplyIndexPebbleStore(t) + ps := pebbleStoreApplied(t, st) + + stage := func(raw string) []byte { + return append([]byte("stage|"), []byte(raw)...) + } + targetKey := func(staged []byte) ([]byte, bool) { + return bytes.TrimPrefix(staged, []byte("stage|")), bytes.HasPrefix(staged, []byte("stage|")) + } + prefix := []byte("stage|") + + require.NoError(t, ps.PutAt(ctx, stage("k"), []byte("v10"), 10, 0)) + + const entryIdx uint64 = 77 + result, err := ps.PromoteVersions(ctx, PromoteVersionsOptions{ + JobID: 9, + AppliedIndex: entryIdx, + StartKey: prefix, + EndKey: PrefixScanEnd(prefix), + MaxVersions: 10, + TargetKey: targetKey, + }) + require.NoError(t, err) + require.True(t, result.Done) + require.Equal(t, uint64(1), result.PromotedRows) + + got, present, err := ps.LastAppliedIndex() + require.NoError(t, err) + require.True(t, present, "PromoteVersions must persist metaAppliedIndex") + require.Equal(t, entryIdx, got) + + val, err := ps.GetAt(ctx, []byte("k"), 10) + require.NoError(t, err) + require.Equal(t, []byte("v10"), val) + _, err = ps.GetAt(ctx, stage("k"), 10) + require.ErrorIs(t, err, ErrKeyNotFound) + + const retryEntryIdx uint64 = 78 + retry, err := ps.PromoteVersions(ctx, PromoteVersionsOptions{ + JobID: 9, + AppliedIndex: retryEntryIdx, + StartKey: prefix, + EndKey: PrefixScanEnd(prefix), + MaxVersions: 10, + TargetKey: targetKey, + }) + require.NoError(t, err) + require.True(t, retry.Done) + require.Equal(t, uint64(1), retry.TotalPromotedRows) + + got, present, err = ps.LastAppliedIndex() + require.NoError(t, err) + require.True(t, present, "completed PromoteVersions retry must persist metaAppliedIndex") + require.Equal(t, retryEntryIdx, got) +} + +func TestImportVersionsRaft_BundlesMetaAppliedIndex(t *testing.T) { + ctx := context.Background() + st := newApplyIndexPebbleStore(t) + ps := pebbleStoreApplied(t, st) + + const entryIdx uint64 = 123 + result, err := ps.ImportVersionsRaft(ctx, ImportVersionsOptions{ + JobID: 9, + AppliedIndex: entryIdx, + BracketID: 1, + BatchSeq: 1, + Cursor: []byte("c1"), + Versions: []MVCCVersion{ + {Key: []byte("stage|k"), CommitTS: 100, Value: []byte("v100")}, + }, + }) + require.NoError(t, err) + require.Equal(t, []byte("c1"), result.AckedCursor) + require.Equal(t, uint64(100), result.MaxImportedTS) + + got, present, err := ps.LastAppliedIndex() + require.NoError(t, err) + require.True(t, present, "ImportVersionsRaft must persist metaAppliedIndex") + require.Equal(t, entryIdx, got) + + val, err := ps.GetAt(ctx, []byte("stage|k"), 100) + require.NoError(t, err) + require.Equal(t, []byte("v100"), val) + + const retryEntryIdx uint64 = 124 + duplicate, err := ps.ImportVersionsRaft(ctx, ImportVersionsOptions{ + JobID: 9, + AppliedIndex: retryEntryIdx, + BracketID: 1, + BatchSeq: 1, + Cursor: []byte("ignored"), + Versions: []MVCCVersion{ + {Key: []byte("stage|k"), CommitTS: 100, Value: []byte("changed")}, + }, + }) + require.NoError(t, err) + require.True(t, duplicate.Duplicate) + require.Equal(t, []byte("c1"), duplicate.AckedCursor) + + got, present, err = ps.LastAppliedIndex() + require.NoError(t, err) + require.True(t, present, "duplicate ImportVersionsRaft retry must still advance metaAppliedIndex") + require.Equal(t, retryEntryIdx, got) + + val, err = ps.GetAt(ctx, []byte("stage|k"), 100) + require.NoError(t, err) + require.Equal(t, []byte("v100"), val, "duplicate import must not rewrite the acknowledged batch") +} + func TestApplyMutationsRaftAt_AlreadyLandedAdvancesStaleAppliedIndex(t *testing.T) { ctx := context.Background() st := newApplyIndexPebbleStore(t) 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 4e42af1db..8008d971c 100644 --- a/store/migration_versions.go +++ b/store/migration_versions.go @@ -18,6 +18,7 @@ const ( migrationAckMetaKey = "_migack" migrationHLCFloorMetaKey = "_mighlc" + migrationPromoteMetaKey = "_migpromote" migrationMetadataVersion = 1 migrationAckPrefix = "!migstage|ack|" @@ -29,6 +30,7 @@ const ( var ( migrationAckMetaKeyBytes = []byte(migrationAckMetaKey) migrationHLCFloorMetaKeyBytes = []byte(migrationHLCFloorMetaKey) + migrationPromoteMetaKeyBytes = []byte(migrationPromoteMetaKey) ) type exportCursorPosition struct { @@ -89,6 +91,54 @@ func decodeExportCursor(cursor []byte) (exportCursorPosition, error) { return exportCursorPosition{key: key, commitTS: commitTS, tag: tag, hasKey: true}, nil } +// ValidateExportCursorForRange verifies that an export cursor decodes and +// resumes inside the supplied key interval. Skipped-key cursors are accepted +// only when they describe a key outside the interval. +func ValidateExportCursorForRange(cursor, startKey, endKey []byte) error { + pos, err := decodeExportCursor(cursor) + if err != nil { + return err + } + return validateExportCursorPositionForRange(pos, startKey, endKey) +} + +// ValidatePromotionCursorForRange verifies a promotion cursor before it is +// proposed to Raft. Promotion scans emit only accepted positions, so callers +// must not resume from sparse-scan-only cursor tags. +func ValidatePromotionCursorForRange(cursor, startKey, endKey []byte) error { + pos, err := decodeExportCursor(cursor) + if err != nil { + return err + } + if !pos.hasKey { + return nil + } + if pos.tag != exportCursorTagEmitted { + return errors.WithStack(ErrInvalidExportCursor) + } + return validateExportCursorPositionForRange(pos, startKey, endKey) +} + +func validateExportCursorPositionForRange(pos exportCursorPosition, startKey, endKey []byte) error { + if !pos.hasKey { + return nil + } + if pos.tag == exportCursorTagSkippedKey { + opts := ExportVersionsOptions{StartKey: startKey, EndKey: endKey} + if !exportSkippedCursorOutsideRange(opts, pos.key) { + return errors.WithStack(ErrInvalidExportCursor) + } + return nil + } + if startKey != nil && bytes.Compare(pos.key, startKey) < 0 { + return errors.WithStack(ErrInvalidExportCursor) + } + if endKey != nil && bytes.Compare(pos.key, endKey) >= 0 { + return errors.WithStack(ErrInvalidExportCursor) + } + return nil +} + func decodeExportCursorForOptions(opts ExportVersionsOptions) (exportCursorPosition, error) { pos, err := decodeExportCursor(opts.Cursor) if err != nil { @@ -97,7 +147,7 @@ func decodeExportCursorForOptions(opts ExportVersionsOptions) (exportCursorPosit if err := validateExportCursorRange(opts, pos); err != nil { return exportCursorPosition{}, err } - return pos, nil + return normalizeExportCursorPositionForRange(opts, pos), nil } func validateExportCursorRange(opts ExportVersionsOptions, pos exportCursorPosition) error { @@ -124,6 +174,16 @@ func exportSkippedCursorOutsideRange(opts ExportVersionsOptions, key []byte) boo (opts.EndKey != nil && bytes.Compare(key, opts.EndKey) >= 0) } +func normalizeExportCursorPositionForRange(opts ExportVersionsOptions, pos exportCursorPosition) exportCursorPosition { + if !pos.hasKey || pos.tag != exportCursorTagSkippedKey || opts.StartKey == nil { + return pos + } + if bytes.Compare(pos.key, opts.StartKey) >= 0 { + return pos + } + return exportCursorPosition{} +} + func normalizeExportVersionsOptions(opts ExportVersionsOptions) ExportVersionsOptions { if opts.EndKey != nil && len(opts.EndKey) == 0 { opts.EndKey = nil @@ -145,7 +205,8 @@ func exportUsesSparseScanBudget(opts ExportVersionsOptions) bool { func isMigrationMetadataKey(rawKey []byte) bool { return bytes.Equal(rawKey, migrationAckMetaKeyBytes) || - bytes.Equal(rawKey, migrationHLCFloorMetaKeyBytes) + bytes.Equal(rawKey, migrationHLCFloorMetaKeyBytes) || + bytes.Equal(rawKey, migrationPromoteMetaKeyBytes) } func encodeMigrationImportAcks(acks map[migrationAckID]migrationImportAck) []byte { @@ -250,6 +311,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") @@ -407,10 +523,10 @@ func appendMemoryExportVersion(opts ExportVersionsOptions, key []byte, version V if opts.AcceptKey != nil && !opts.AcceptKey(key) { return exportCursorTagScanned } - if opts.AcceptVersion != nil && !opts.AcceptVersion(key, version.Value) { + if opts.MaxCommitTSInclusive != 0 && version.TS > opts.MaxCommitTSInclusive { return exportCursorTagScanned } - if opts.MaxCommitTSInclusive != 0 && version.TS > opts.MaxCommitTSInclusive { + if opts.AcceptVersion != nil && !opts.AcceptVersion(key, version.Value) { return exportCursorTagScanned } result.Versions = append(result.Versions, MVCCVersion{ @@ -521,6 +637,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 3dc418888..2f93a78d2 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" ) @@ -88,6 +89,29 @@ func TestExportVersionsAcceptVersionFiltersByValue(t *testing.T) { }) } +func TestExportVersionsAppliesTimestampBoundBeforeAcceptVersion(t *testing.T) { + runMigrationStoreSuite(t, func(t *testing.T, st MVCCStore) { + ctx := context.Background() + require.NoError(t, st.PutAt(ctx, []byte("k"), []byte("eligible"), 20, 0)) + require.NoError(t, st.PutAt(ctx, []byte("k"), []byte("too-new"), 30, 0)) + accepted := false + + result, err := st.ExportVersions(ctx, ExportVersionsOptions{ + MaxCommitTSInclusive: 25, + MaxVersions: 1, + AcceptVersion: func(_ []byte, _ []byte) bool { + if accepted { + return false + } + accepted = true + return true + }, + }) + require.NoError(t, err) + require.Equal(t, []MVCCVersion{{Key: []byte("k"), CommitTS: 20, Value: []byte("eligible")}}, result.Versions) + }) +} + func TestExportVersionsCursorResumesWithinHotKey(t *testing.T) { runMigrationStoreSuite(t, func(t *testing.T, st MVCCStore) { ctx := context.Background() @@ -422,6 +446,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() @@ -831,6 +921,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 9211ce6e7..4d2b61dfa 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) @@ -985,6 +987,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 d8e1b9dd6..5a746d8c3 100644 --- a/store/store.go +++ b/store/store.go @@ -108,11 +108,14 @@ type ExportVersionsResult struct { // ImportVersionsOptions applies one idempotent migration-import batch. type ImportVersionsOptions struct { - JobID uint64 - BracketID uint64 - BatchSeq uint64 - Versions []MVCCVersion - Cursor []byte + JobID uint64 + // AppliedIndex is the optional Raft entry index to bundle with Pebble + // import batches as metaAppliedIndex. Zero leaves the meta key unchanged. + AppliedIndex uint64 + BracketID uint64 + BatchSeq uint64 + Versions []MVCCVersion + Cursor []byte } // ImportVersionsResult reports the cursor durably acknowledged by the target. @@ -122,6 +125,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 @@ -273,6 +323,11 @@ type MVCCStore interface { // ImportVersions applies a migration import batch idempotently by // (jobID, bracketID, batchSeq), preserving tombstones and expireAt. ImportVersions(ctx context.Context, opts ImportVersionsOptions) (ImportVersionsResult, error) + // ImportVersionsRaft is the raft-apply variant of ImportVersions. It + // preserves the same idempotency contract while using the FSM write path. + // When opts.AppliedIndex is non-zero, the implementation must durably + // bundle metaAppliedIndex with the import batch. + ImportVersionsRaft(ctx context.Context, opts ImportVersionsOptions) (ImportVersionsResult, error) // MigrationHLCFloor returns the full-HLC target-local migration floor // persisted by ImportVersions for jobID. MigrationHLCFloor(ctx context.Context, jobID uint64) (uint64, error) From 0a1b9c890cd13e0d582b8243055f4b64184de7e0 Mon Sep 17 00:00:00 2001 From: bootjp Date: Sun, 19 Jul 2026 22:08:41 +0900 Subject: [PATCH 02/58] Preserve staged key scan routing --- adapter/internal.go | 10 ++-- adapter/internal_test.go | 5 +- kv/shard_store.go | 100 +++++++++++++++++++++++++++------------ kv/shard_store_test.go | 57 ++++++++++++++++++++++ 4 files changed, 136 insertions(+), 36 deletions(-) diff --git a/adapter/internal.go b/adapter/internal.go index 89784011e..4037f6b8d 100644 --- a/adapter/internal.go +++ b/adapter/internal.go @@ -728,10 +728,14 @@ func (i *Internal) stampTxnTimestamps(ctx context.Context, reqs []*pb.Request) ( } } - if err := i.fillForwardedTxnCommitTS(ctx, reqs, startTS); err != nil { - return err + commitTS, err := i.fillForwardedTxnCommitTS(ctx, reqs, startTS) + if err != nil { + return 0, err } - return i.rejectWriteTimestampFloorTxnRequests(reqs) + if err := i.rejectWriteTimestampFloorTxnRequests(reqs); err != nil { + return 0, err + } + return commitTS, nil } func forwardedTxnStartTS(reqs []*pb.Request) uint64 { diff --git a/adapter/internal_test.go b/adapter/internal_test.go index a3d937f95..ec2b9eb8b 100644 --- a/adapter/internal_test.go +++ b/adapter/internal_test.go @@ -388,7 +388,7 @@ func TestStampTxnTimestampsRejectsRouteWriteFloor(t *testing.T) { }, }} - err := i.stampTxnTimestamps(context.Background(), reqs) + _, err := i.stampTxnTimestamps(context.Background(), reqs) require.ErrorIs(t, err, kv.ErrRouteWriteTimestampTooLow) } @@ -427,5 +427,6 @@ func TestStampTxnTimestampsIgnoresMetadataForRouteWriteFloor(t *testing.T) { }, }} - require.NoError(t, i.stampTxnTimestamps(context.Background(), reqs)) + _, err := i.stampTxnTimestamps(context.Background(), reqs) + require.NoError(t, err) } diff --git a/kv/shard_store.go b/kv/shard_store.go index 63cfaf745..032930e54 100644 --- a/kv/shard_store.go +++ b/kv/shard_store.go @@ -618,7 +618,10 @@ func (s *ShardStore) ReverseScanGroupAt(ctx context.Context, groupID uint64, sta if limit <= 0 { return []*store.KVPair{}, nil } - return s.scanRouteAtDirection(ctx, distribution.Route{GroupID: groupID}, start, end, limit, ts, true, true) + return s.scanRouteAtDirectionWithReadFence( + ctx, distribution.Route{GroupID: groupID}, start, end, limit, ts, true, true, + 0, nil, nil, + ) } // ScanGroupKeysAt scans keys on the explicitly selected Raft group without @@ -1040,14 +1043,32 @@ func normalizedRouteScanEnd(routeEnd []byte) []byte { return routeEnd } -func (s *ShardStore) scanRoutesAtWithReadFence(ctx context.Context, routes []distribution.Route, start []byte, end []byte, limit int, ts uint64, clampToRoutes bool, readRouteVersion uint64, routeStart []byte, routeEnd []byte) ([]*store.KVPair, error) { - out := make([]*store.KVPair, 0) +type scanRouteOwnerFilterPlan struct { + routes []distribution.Route + routeFilterPresent bool + filterUsageOwners bool + dedupeByKey bool +} + +func prepareScanRouteOwnerFilters(routes []distribution.Route, start []byte, end []byte, clampToRoutes bool, routeStart []byte, routeEnd []byte) scanRouteOwnerFilterPlan { routeFilterPresent := routeScanBoundsPresent(routeStart, routeEnd) filterUsageOwners := !clampToRoutes && !routeFilterPresent && filesystemUsageScanOverlap(start, end) dedupeByKey := s3BucketAuxiliaryScanBounds(start, end) if !clampToRoutes && !routeFilterPresent { routes, dedupeByKey = prepareUnclampedRawScanRoutes(routes, dedupeByKey) } + return scanRouteOwnerFilterPlan{ + routes: routes, + routeFilterPresent: routeFilterPresent, + filterUsageOwners: filterUsageOwners, + dedupeByKey: dedupeByKey, + } +} + +func (s *ShardStore) scanRoutesAtWithReadFence(ctx context.Context, routes []distribution.Route, start []byte, end []byte, limit int, ts uint64, clampToRoutes bool, readRouteVersion uint64, routeStart []byte, routeEnd []byte) ([]*store.KVPair, error) { + out := make([]*store.KVPair, 0) + plan := prepareScanRouteOwnerFilters(routes, start, end, clampToRoutes, routeStart, routeEnd) + routes = plan.routes for _, route := range routes { scanStart := start scanEnd := end @@ -1058,12 +1079,12 @@ func (s *ShardStore) scanRoutesAtWithReadFence(ctx context.Context, routes []dis kvs, err := s.scanRouteAtWithMigrationOwnerFilters( ctx, routes, route, scanStart, scanEnd, limit, ts, false, !clampToRoutes, - readRouteVersion, routeStart, routeEnd, filterUsageOwners, dedupeByKey, + readRouteVersion, routeStart, routeEnd, plan.filterUsageOwners, plan.dedupeByKey, ) if err != nil { return nil, err } - if isBroadLegacyListDeltaScan(start) && !routeFilterPresent { + if isBroadLegacyListDeltaScan(start) && !plan.routeFilterPresent { kvs = markScanRouteGroup(kvs, route.GroupID, true) } if clampToRoutes { @@ -1074,7 +1095,7 @@ func (s *ShardStore) scanRoutesAtWithReadFence(ctx context.Context, routes []dis } continue } - out = mergeAndTrimScanResultsWithOptions(out, kvs, limit, false, dedupeByKey) + out = mergeAndTrimScanResultsWithOptions(out, kvs, limit, false, plan.dedupeByKey) } return out, nil } @@ -1318,7 +1339,7 @@ func (s *ShardStore) scanRouteAtWithFilesystemUsageOwnerFilter( if reverse { out = mergeAndTrimReverseScanResults(out, page, limit) } else { - out = mergeAndTrimScanResults(out, page, limit) + out = mergeAndTrimScanResultsWithOptions(out, page, limit, false, false) } if len(out) >= limit || pageLen < limit || advanceKey == nil { break @@ -1386,12 +1407,8 @@ func (s *ShardStore) reverseScanRoutesAtWithReadFence( routeEnd []byte, ) ([]*store.KVPair, error) { out := make([]*store.KVPair, 0) - routeFilterPresent := routeScanBoundsPresent(routeStart, routeEnd) - filterUsageOwners := !clampToRoutes && !routeFilterPresent && filesystemUsageScanOverlap(start, end) - dedupeByKey := s3BucketAuxiliaryScanBounds(start, end) - if !clampToRoutes && !routeFilterPresent { - routes, dedupeByKey = prepareUnclampedRawScanRoutes(routes, dedupeByKey) - } + plan := prepareScanRouteOwnerFilters(routes, start, end, clampToRoutes, routeStart, routeEnd) + routes = plan.routes for i := 0; i < len(routes); i++ { route := routes[i] if clampToRoutes { @@ -1413,15 +1430,15 @@ func (s *ShardStore) reverseScanRoutesAtWithReadFence( // result honours the ReverseScanAt contract. kvs, err := s.scanRouteAtWithMigrationOwnerFilters( ctx, routes, route, start, end, limit, ts, true, true, - readRouteVersion, routeStart, routeEnd, filterUsageOwners, dedupeByKey, + readRouteVersion, routeStart, routeEnd, plan.filterUsageOwners, plan.dedupeByKey, ) if err != nil { return nil, err } - if isBroadLegacyListDeltaScan(start) && !routeFilterPresent { + if isBroadLegacyListDeltaScan(start) && !plan.routeFilterPresent { kvs = markScanRouteGroup(kvs, route.GroupID, true) } - out = mergeAndTrimScanResultsWithOptions(out, kvs, limit, true, dedupeByKey) + out = mergeAndTrimScanResultsWithOptions(out, kvs, limit, true, plan.dedupeByKey) } return out, nil } @@ -1466,8 +1483,17 @@ func (s *ShardStore) scanKeyRouteAtWithReadFence( return s.scanKeysRouteAtLeader(ctx, g, route, start, end, limit, ts) } - groupID := proxyScanGroupID(route, explicitGroup, readRouteVersion, nil, nil) - return s.proxyScanKeysAt(ctx, g, start, end, limit, ts, groupID, readRouteVersion) + routeBoundsPresent := routeHasStagedVisibility(route) + var routeStart, routeEnd []byte + if routeBoundsPresent { + routeStart = route.Start + routeEnd = route.End + } + groupID := proxyScanGroupID(route, explicitGroup, readRouteVersion, routeStart, routeEnd) + return s.proxyScanKeysAt( + ctx, g, start, end, limit, ts, groupID, readRouteVersion, + routeStart, routeEnd, routeBoundsPresent, + ) } func (s *ShardStore) scanKeysRouteLocal( @@ -1532,7 +1558,7 @@ func (s *ShardStore) scanKeysRouteAtLeader( if err != nil { return nil, err } - out = mergeAndTrimScanKeys(out, filterTxnInternalKeys(keysFromKVs(kvs)), limit) + out = mergeAndTrimScanKeys(out, filterScanInternalKeys(keysFromKVs(kvs)), limit) nextCursor, ok := nextKeyScanCursor(keys, end, limit) if !ok { @@ -1563,7 +1589,7 @@ func (s *ShardStore) scanLockOnlyKeysAtLeader( if err != nil { return nil, err } - return filterTxnInternalKeys(keysFromKVs(kvs)), nil + return filterScanInternalKeys(keysFromKVs(kvs)), nil } func (s *ShardStore) proxyScanKeysAt( @@ -1575,9 +1601,15 @@ func (s *ShardStore) proxyScanKeysAt( ts uint64, groupID uint64, readRouteVersion uint64, + routeStart []byte, + routeEnd []byte, + routeBoundsPresent bool, ) ([][]byte, error) { return scanKeysWithRefill(start, end, limit, func(cursor []byte, pageLimit int) ([][]byte, error) { - return s.proxyRawScanKeysAt(ctx, g, cursor, end, pageLimit, ts, groupID, readRouteVersion) + return s.proxyRawScanKeysAt( + ctx, g, cursor, end, pageLimit, ts, groupID, readRouteVersion, + routeStart, routeEnd, routeBoundsPresent, + ) }) } @@ -1602,7 +1634,7 @@ func scanKeysWithRefill( break } - out = mergeAndTrimScanKeys(out, filterTxnInternalKeys(keys), limit) + out = mergeAndTrimScanKeys(out, filterScanInternalKeys(keys), limit) nextCursor, ok := nextKeyScanCursor(keys, end, limit) if !ok { @@ -2906,13 +2938,13 @@ func lastKVKey(kvs []*store.KVPair) []byte { return nil } -func filterTxnInternalKeys(keys [][]byte) [][]byte { +func filterScanInternalKeys(keys [][]byte) [][]byte { if len(keys) == 0 { return keys } out := make([][]byte, 0, len(keys)) for _, key := range keys { - if key == nil || isTxnInternalKey(key) { + if key == nil || isScanInternalKey(key) { continue } out = append(out, key) @@ -4368,6 +4400,9 @@ func (s *ShardStore) proxyRawScanKeysAt( ts uint64, groupID uint64, readRouteVersion uint64, + routeStart []byte, + routeEnd []byte, + routeBoundsPresent bool, ) ([][]byte, error) { engine := engineForGroup(g) if engine == nil { @@ -4387,13 +4422,16 @@ func (s *ShardStore) proxyRawScanKeysAt( defer cancel() cli := pb.NewRawKVClient(conn) resp, err := cli.RawScanAt(ctx, &pb.RawScanAtRequest{ - StartKey: start, - EndKey: end, - Limit: int64(limit), - Ts: ts, - GroupId: groupID, - ReadRouteVersion: readRouteVersion, - KeysOnly: true, + StartKey: start, + EndKey: end, + Limit: int64(limit), + Ts: ts, + GroupId: groupID, + ReadRouteVersion: readRouteVersion, + KeysOnly: true, + RouteStart: bytes.Clone(routeStart), + RouteEnd: bytes.Clone(routeEnd), + RouteBoundsPresent: routeBoundsPresent, }) if err != nil { return nil, errors.WithStack(err) diff --git a/kv/shard_store_test.go b/kv/shard_store_test.go index bdd25e85f..143e5351d 100644 --- a/kv/shard_store_test.go +++ b/kv/shard_store_test.go @@ -1917,6 +1917,23 @@ func TestShardStoreScanKeysRouteAtLeaderRefillsAfterTxnInternalKeys(t *testing.T require.Equal(t, [][]byte{[]byte("a")}, keys) } +func TestShardStoreScanKeysRouteAtLeaderRefillsAfterStagedControlKeys(t *testing.T) { + t.Parallel() + + ctx := context.Background() + g := &ShardGroup{Store: store.NewMVCCStore()} + st := NewShardStore(distribution.NewEngine(), map[uint64]*ShardGroup{1: g}) + t.Cleanup(func() { _ = st.Close() }) + + stagedKey := distribution.MigrationStagedDataKey(9, []byte("shadow")) + require.NoError(t, g.Store.PutAt(ctx, stagedKey, []byte("internal"), 1, 0)) + require.NoError(t, g.Store.PutAt(ctx, []byte("a"), []byte("visible"), 2, 0)) + + keys, err := st.scanKeysRouteAtLeader(ctx, g, distribution.Route{GroupID: 1}, []byte(""), nil, 1, ^uint64(0)) + require.NoError(t, err) + require.Equal(t, [][]byte{[]byte("a")}, keys) +} + func TestShardStoreScanKeysRouteAtLeaderPreservesEmptyKey(t *testing.T) { t.Parallel() @@ -2009,6 +2026,46 @@ func TestShardStoreProxyScanKeysAtUsesSelectedGroup(t *testing.T) { require.True(t, fake.lastScanKeysOnly) } +func TestShardStoreProxyScanKeysAtCarriesStagedRouteBounds(t *testing.T) { + t.Parallel() + + fake := &fakeRawKVServer{ + scanResp: &pb.RawScanAtResponse{ + Kv: []*pb.RawKVPair{{Key: []byte("k"), Value: []byte("v")}}, + }, + } + addr, stop := startRawKVServer(t, fake) + t.Cleanup(stop) + + ctx := context.Background() + g := &ShardGroup{ + Engine: &followerProxyEngine{leader: addr}, + Store: store.NewMVCCStore(), + } + st := NewShardStore(distribution.NewEngine(), map[uint64]*ShardGroup{42: g}) + t.Cleanup(func() { _ = st.Close() }) + route := distribution.Route{ + Start: []byte("a"), + End: []byte("m"), + GroupID: 42, + StagedVisibilityActive: true, + MigrationJobID: 9, + } + + keys, err := st.scanKeyRouteAtWithReadFence(ctx, route, []byte("a"), []byte("m"), 10, ^uint64(0), false, 7) + require.NoError(t, err) + require.Equal(t, [][]byte{[]byte("k")}, keys) + + fake.mu.Lock() + defer fake.mu.Unlock() + require.Equal(t, uint64(42), fake.lastScanReq.GetGroupId()) + require.Equal(t, uint64(7), fake.lastScanReq.GetReadRouteVersion()) + require.Equal(t, []byte("a"), fake.lastScanReq.GetRouteStart()) + require.Equal(t, []byte("m"), fake.lastScanReq.GetRouteEnd()) + require.True(t, fake.lastScanReq.GetRouteBoundsPresent()) + require.True(t, fake.lastScanReq.GetKeysOnly()) +} + func TestShardStoreProxyScanAtUsesSelectedGroup(t *testing.T) { t.Parallel() From e7815afc4817230486593d241353b3cad9f50437 Mon Sep 17 00:00:00 2001 From: bootjp Date: Sat, 22 Aug 2026 10:41:09 +0900 Subject: [PATCH 03/58] kv: cover staged tombstones in reverse auxiliary scans TestShardStoreS3BucketAuxiliaryScanHonorsStagedTombstone only exercised the forward direction. The review that raised this named SCAN and reverse-scan together, so lock the reverse path down as well: a staged tombstone must hide the stale live row that the old raw route still holds, in both directions. Both already pass. filterS3AuxiliaryKVsOwnedByStagedRoutes drops the non-staged route's auxiliary rows before the dedupe merge runs, so the merge never has a stale row to replace. Claude-Session: https://claude.ai/code/session_013rNHooj7NF3giihWVba8QE --- kv/shard_store_test.go | 44 ++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 44 insertions(+) diff --git a/kv/shard_store_test.go b/kv/shard_store_test.go index 143e5351d..9b014d25d 100644 --- a/kv/shard_store_test.go +++ b/kv/shard_store_test.go @@ -3388,3 +3388,47 @@ func TestScanLockBoundsForKVs_ReverseInternalOnlyPageUsesOriginalRange(t *testin require.Equal(t, []byte(""), lockStart) require.Equal(t, []byte("z"), lockEnd) } + +// Reverse-scan counterpart of TestShardStoreS3BucketAuxiliaryScanHonorsStagedTombstone. +// A staged tombstone must hide the stale live row from the old raw route in both +// scan directions, not just forward. +func TestShardStoreS3BucketAuxiliaryReverseScanHonorsStagedTombstone(t *testing.T) { + t.Parallel() + + ctx := context.Background() + const migratedBucket = "bucket-a" + routeStart := s3keys.RoutePrefixForBucketAnyGeneration(migratedBucket) + routeEnd := prefixScanEnd(routeStart) + engine := distribution.NewEngine() + require.NoError(t, engine.ApplySnapshot(distribution.CatalogSnapshot{ + Version: 1, + Routes: []distribution.RouteDescriptor{ + {RouteID: 1, Start: []byte(""), End: routeStart, GroupID: 1, State: distribution.RouteStateActive}, + {RouteID: 2, Start: routeStart, End: routeEnd, GroupID: 2, State: distribution.RouteStateActive, StagedVisibilityActive: true, MigrationJobID: 9}, + {RouteID: 3, Start: routeEnd, End: nil, GroupID: 1, State: distribution.RouteStateActive}, + }, + })) + groups := map[uint64]*ShardGroup{ + 1: {Store: store.NewMVCCStore()}, + 2: {Store: store.NewMVCCStore()}, + } + st := NewShardStore(engine, groups) + deletedKey := s3keys.BucketMetaKey(migratedBucket) + visibleKey := s3keys.BucketMetaKey("bucket-z") + require.NoError(t, groups[1].Store.PutAt(ctx, deletedKey, []byte("stale"), 10, 0)) + require.NoError(t, groups[1].Store.PutAt(ctx, visibleKey, []byte("visible"), 10, 0)) + require.NoError(t, groups[2].Store.DeleteAt(ctx, distribution.MigrationStagedDataKey(9, deletedKey), 20)) + + start := []byte(s3keys.BucketMetaPrefix) + end := prefixScanEnd(start) + + // Reverse over the whole family: the tombstoned bucket must not appear. + kvs, err := st.ReverseScanAt(ctx, start, end, 10, 30) + require.NoError(t, err) + require.Equal(t, []*store.KVPair{{Key: visibleKey, Value: []byte("visible")}}, kvs) + + // Reverse scoped to the tombstoned bucket alone. + kvs, err = st.ReverseScanAt(ctx, deletedKey, prefixScanEnd(deletedKey), 10, 30) + require.NoError(t, err) + require.Empty(t, kvs) +} From 3fa73ed1027e08e6c054806eb34ef6f58531d60f Mon Sep 17 00:00:00 2001 From: bootjp Date: Sat, 22 Aug 2026 15:47:22 +0900 Subject: [PATCH 04/58] migration: exempt resolver-owned keys from the floor precheck rejectWriteFencedPointKey skips keys the partition resolver recognises, because in a partition-resolved keyspace such as HT-FIFO SQS routeKey collapses a concrete partition key onto the global SQS route, so that route's state does not describe the key. rejectWriteTimestampFloorPointKey had no such exemption, so the same collapsed route's MinWriteTSExclusive was applied to those keys and rejected writes the fence precheck deliberately lets through. Apply the same exemption, using the existing partitionResolverRecognisesPointKey helper so both prechecks agree on what the resolver owns. Claude-Session: https://claude.ai/code/session_013rNHooj7NF3giihWVba8QE --- kv/sharded_coordinator.go | 9 +++++++ kv/sharded_coordinator_txn_test.go | 41 ++++++++++++++++++++++++++++++ 2 files changed, 50 insertions(+) diff --git a/kv/sharded_coordinator.go b/kv/sharded_coordinator.go index 9954c1785..594a9819e 100644 --- a/kv/sharded_coordinator.go +++ b/kv/sharded_coordinator.go @@ -1296,6 +1296,15 @@ func (c *ShardedCoordinator) rejectWriteTimestampFloorPointElems(elems []*Elem[O } func (c *ShardedCoordinator) rejectWriteTimestampFloorPointKey(key []byte, commitTS uint64) error { + // Same exemption rejectWriteFencedPointKey applies, and for the same + // reason: in partition-resolved keyspaces such as HT-FIFO SQS, routeKey + // collapses a concrete partition key onto the global SQS route, so that + // route's floor is not this key's floor. Applying it here rejected writes + // the write-fence precheck deliberately lets through, and the Raft-side + // gate already covers these keys through the resolver's own routing. + if c.partitionResolverRecognisesPointKey(key) { + return nil + } start, end, ok := s3BucketAuxiliaryRouteRange(key) if ok { for _, route := range c.engine.GetIntersectingRoutes(start, end) { diff --git a/kv/sharded_coordinator_txn_test.go b/kv/sharded_coordinator_txn_test.go index 96614f4b3..0bc769f3e 100644 --- a/kv/sharded_coordinator_txn_test.go +++ b/kv/sharded_coordinator_txn_test.go @@ -1031,3 +1031,44 @@ func TestShardedCoordinatorDispatchTxn_CrossShardPropagatesObservedRouteVersion( req.Phase) } } + +// In partition-resolved keyspaces such as HT-FIFO SQS, routeKey collapses a +// concrete partition key onto the global SQS route, so that route's write floor +// is not the key's floor. rejectWriteFencedPointKey already exempts these keys; +// the timestamp-floor precheck must match, or it rejects writes the fence +// precheck deliberately lets through. +func TestShardedCoordinatorFloorPrecheckSkipsResolverOwnedKeys(t *testing.T) { + t.Parallel() + + const partitionKey = "!sqs|msg|data|p|queue|7" + + engine := distribution.NewEngine() + require.NoError(t, engine.ApplySnapshot(distribution.CatalogSnapshot{ + Version: 1, + Routes: []distribution.RouteDescriptor{ + // The route routeKey() collapses partition keys onto, carrying a floor. + {RouteID: 1, Start: []byte(""), End: nil, GroupID: 1, State: distribution.RouteStateActive, MinWriteTSExclusive: 100}, + }, + })) + c := NewShardedCoordinator(engine, map[uint64]*ShardGroup{1: {}}, 1, NewHLC(), nil) + + // Without a resolver the floor applies, which is what makes the exemption + // below meaningful rather than vacuous. + require.ErrorIs(t, + c.rejectWriteTimestampFloorPointKey([]byte(partitionKey), 100), + ErrRouteWriteTimestampTooLow) + + c.WithPartitionResolver(&fakePartitionResolver{ + routes: map[string]uint64{partitionKey: 1}, + recognisedPrefix: []byte("!sqs|msg|data|p|"), + }) + + require.NoError(t, + c.rejectWriteTimestampFloorPointKey([]byte(partitionKey), 100), + "a resolver-owned key must not be judged by the route routeKey collapses it onto") + + // A key the resolver does not own still gets the floor. + require.ErrorIs(t, + c.rejectWriteTimestampFloorPointKey([]byte("ordinary-key"), 100), + ErrRouteWriteTimestampTooLow) +} From 86c48e62835c0a0ebee73d9126bed67079cafa12 Mon Sep 17 00:00:00 2001 From: bootjp Date: Sat, 22 Aug 2026 19:45:51 +0900 Subject: [PATCH 05/58] migration: mark route group on exact legacy list-delta scans Route-group marking was gated on isBroadLegacyListDeltaScan, which is false when the scan start is exactly LegacyListMetaDeltaScanPrefix(userKey). Redis cleanup and compaction build their deletes as {Del, pair.Key, GroupID: pair.RouteGroupID}, so a zero GroupID there routes the delete by the raw "!lst|meta|d|..." key rather than the logical list key -- after a split the delete goes to the wrong shard and the stale delta survives. Gate marking on isLegacyListDeltaScan, which covers both the broad and the per-user-key shape, at all three marking sites. Route *selection* still distinguishes them: routesForInternalScanWithVersion deliberately narrows an exact scan to the logical key's routes, so that call keeps isBroadLegacyListDeltaScan. Claude-Session: https://claude.ai/code/session_013rNHooj7NF3giihWVba8QE --- kv/shard_store.go | 20 +++++++++++++++++--- kv/shard_store_test.go | 29 +++++++++++++++++++++++++++++ 2 files changed, 46 insertions(+), 3 deletions(-) diff --git a/kv/shard_store.go b/kv/shard_store.go index 032930e54..0cdfaf700 100644 --- a/kv/shard_store.go +++ b/kv/shard_store.go @@ -995,8 +995,22 @@ func isBroadLegacyListDeltaScan(start []byte) bool { return logicalUserKey == nil || !bytes.Equal(start, store.LegacyListMetaDeltaScanPrefix(logicalUserKey)) } +// isLegacyListDeltaScan reports whether start addresses the legacy list-delta +// family at all, whether broad or scoped to a single user key. +// +// Route-group marking must cover both shapes. Redis cleanup and compaction build +// their delete elems as {Op: Del, Key: pair.Key, GroupID: pair.RouteGroupID}; a +// zero GroupID falls back to routing by the raw "!lst|meta|d|..." key rather +// than the logical list key, so after a split the delete lands on the wrong +// shard and the stale delta survives. Route *selection* still distinguishes the +// two shapes -- see routesForInternalScanWithVersion, where an exact scan is +// deliberately narrowed to the logical key's routes. +func isLegacyListDeltaScan(start []byte) bool { + return bytes.HasPrefix(start, []byte(store.LegacyListMetaDeltaPrefix)) +} + func shouldMarkRouteGroupOnScan(start []byte, explicitGroup bool, routeStart []byte, routeEnd []byte) bool { - return !explicitGroup && !routeScanBoundsPresent(routeStart, routeEnd) && isBroadLegacyListDeltaScan(start) + return !explicitGroup && !routeScanBoundsPresent(routeStart, routeEnd) && isLegacyListDeltaScan(start) } func scanRouteUserKey(start []byte) []byte { @@ -1084,7 +1098,7 @@ func (s *ShardStore) scanRoutesAtWithReadFence(ctx context.Context, routes []dis if err != nil { return nil, err } - if isBroadLegacyListDeltaScan(start) && !plan.routeFilterPresent { + if isLegacyListDeltaScan(start) && !plan.routeFilterPresent { kvs = markScanRouteGroup(kvs, route.GroupID, true) } if clampToRoutes { @@ -1435,7 +1449,7 @@ func (s *ShardStore) reverseScanRoutesAtWithReadFence( if err != nil { return nil, err } - if isBroadLegacyListDeltaScan(start) && !plan.routeFilterPresent { + if isLegacyListDeltaScan(start) && !plan.routeFilterPresent { kvs = markScanRouteGroup(kvs, route.GroupID, true) } out = mergeAndTrimScanResultsWithOptions(out, kvs, limit, true, plan.dedupeByKey) diff --git a/kv/shard_store_test.go b/kv/shard_store_test.go index 9b014d25d..019d762a4 100644 --- a/kv/shard_store_test.go +++ b/kv/shard_store_test.go @@ -3432,3 +3432,32 @@ func TestShardStoreS3BucketAuxiliaryReverseScanHonorsStagedTombstone(t *testing. require.NoError(t, err) require.Empty(t, kvs) } + +// An exact per-user-key legacy delta scan must carry RouteGroupID too. Redis +// cleanup and compaction build their deletes as {Del, pair.Key, GroupID: +// pair.RouteGroupID}; a zero GroupID routes the delete by the raw +// "!lst|meta|d|..." key instead of the logical list key, so after a split it +// lands on the wrong shard and the stale delta survives. +func TestShardStoreScanAt_ExactLegacyListDeltaScanMarksRouteGroup(t *testing.T) { + t.Parallel() + + ctx := context.Background() + st := newTwoRouteShardStoreForScanTest() + deltaValue := store.MarshalListMetaDelta(store.ListMetaDelta{LenDelta: 1}) + + // "right-list" sorts into the second route's group. + userKey := []byte("right-list") + key := legacyListMetaDeltaKey(userKey, 11, 1) + require.NoError(t, st.groups[2].Store.PutAt(ctx, key, deltaValue, 1, 0)) + + scanStart := store.LegacyListMetaDeltaScanPrefix(userKey) + require.False(t, isBroadLegacyListDeltaScan(scanStart), + "this test is only meaningful for the exact-scan shape") + + kvs, err := st.ScanAt(ctx, scanStart, store.PrefixScanEnd(scanStart), 10, ^uint64(0)) + require.NoError(t, err) + require.Len(t, kvs, 1) + require.Equal(t, key, kvs[0].Key) + require.Equal(t, uint64(2), kvs[0].RouteGroupID, + "an exact legacy delta scan must still report the owning route group") +} From da3b9671705097bce40e7ffb8e8687be9ff4c213 Mon Sep 17 00:00:00 2001 From: bootjp Date: Sat, 22 Aug 2026 20:18:39 +0900 Subject: [PATCH 06/58] migration: halt raft apply when an import fails on one voter applyMigrationImport returned every failure as an ordinary error. The engine only stops advancing setApplied for responses implementing HaltApply, so a store-side failure on a single target voter -- a Pebble I/O error, the encryption gate rejecting the write -- let that replica skip the imported versions while the leader applied and acked the RPC. The divergence surfaces later as missing data after failover or promotion. Classify the way applyMigrationPromote already does. Errors that are a verdict on the request bytes stay ordinary, because every replica applying the same entry reaches them identically and the group advances in step: batch gap, malformed version, oversized value. Everything else halts, including a failed HLC-floor read. Halting on the request-shaped errors instead would turn one malformed batch into a cluster-wide outage. Replay after a halt is safe: the import batch either did not commit, or committed and is recognised as a duplicate by validateNextImportBatch. validateImportVersion's anonymous errors become ErrInvalidImportVersion so the classifier can name them. Promotion also calls that validator but classifies only ErrInvalidExportCursor as ordinary, so its behaviour is unchanged. Claude-Session: https://claude.ai/code/session_013rNHooj7NF3giihWVba8QE --- kv/fsm_migration_import.go | 28 +++++++- kv/fsm_migration_import_test.go | 119 ++++++++++++++++++++++++++++++++ store/migration_versions.go | 6 +- store/store.go | 7 ++ 4 files changed, 154 insertions(+), 6 deletions(-) diff --git a/kv/fsm_migration_import.go b/kv/fsm_migration_import.go index d146b6fea..feacb6b8e 100644 --- a/kv/fsm_migration_import.go +++ b/kv/fsm_migration_import.go @@ -11,6 +11,8 @@ import ( "google.golang.org/protobuf/proto" ) +var ErrMigrationImportApply = errors.New("migration import: FSM apply failed; halting apply") + // MarshalMigrationImportCommand encodes a target-group migration import batch // as a Raft FSM command. The target Internal RPC handler uses this instead of // mutating its local store directly so an acknowledged batch has been applied @@ -32,7 +34,7 @@ func MarshalMigrationImportCommand(req *pb.ImportRangeVersionsRequest) ([]byte, 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) + return haltErr(errors.Wrap(errors.Mark(err, ErrMigrationImportApply), "kv/fsm: decode migration import")) } result, err := f.store.ImportVersionsRaft(ctx, store.ImportVersionsOptions{ JobID: req.GetJobId(), @@ -43,11 +45,14 @@ func (f *kvFSM) applyMigrationImport(ctx context.Context, data []byte) any { Versions: migrationStoreVersionsFromProto(req.GetJobId(), req.GetVersions()), }) if err != nil { - return errors.WithStack(err) + if isMigrationImportOrdinaryApplyError(err) { + return errors.Wrap(err, "kv/fsm: apply migration import") + } + return haltErr(errors.Wrap(errors.Mark(err, ErrMigrationImportApply), "kv/fsm: apply migration import")) } result.MaxImportedTS, err = f.migrationHLCFloorForApply(ctx, req, result) if err != nil { - return errors.WithStack(err) + return haltErr(errors.Wrap(errors.Mark(err, ErrMigrationImportApply), "kv/fsm: migration import hlc floor")) } if f.hlc != nil && result.MaxImportedTS > 0 { f.hlc.Observe(result.MaxImportedTS) @@ -55,6 +60,23 @@ func (f *kvFSM) applyMigrationImport(ctx context.Context, data []byte) any { return result } +// isMigrationImportOrdinaryApplyError reports whether err is a verdict on the +// request bytes rather than a failure of this replica's store. Only those may +// be returned as an ordinary apply error: the engine advances setApplied past +// a response that does not implement HaltApply, and every replica applying the +// same entry decides an ordinary error identically, so the group stays in step. +// +// A store-side failure -- Pebble I/O, an encryption gate rejecting the write -- +// is per-replica. Letting it advance setApplied would leave the leader acking +// the batch while the failed voter skips the imported versions for good, which +// surfaces as missing data after failover or promotion. Those halt instead, +// matching applyMigrationPromote. +func isMigrationImportOrdinaryApplyError(err error) bool { + return errors.Is(err, store.ErrImportBatchGap) || + errors.Is(err, store.ErrInvalidImportVersion) || + errors.Is(err, store.ErrValueTooLarge) +} + 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 diff --git a/kv/fsm_migration_import_test.go b/kv/fsm_migration_import_test.go index dafd726d9..92a0241b4 100644 --- a/kv/fsm_migration_import_test.go +++ b/kv/fsm_migration_import_test.go @@ -7,6 +7,7 @@ import ( "github.com/bootjp/elastickv/distribution" pb "github.com/bootjp/elastickv/proto" "github.com/bootjp/elastickv/store" + "github.com/cockroachdb/errors" "github.com/stretchr/testify/require" "google.golang.org/protobuf/proto" ) @@ -110,3 +111,121 @@ func TestApplyMigrationImportThreadsPendingApplyIndex(t *testing.T) { require.Len(t, capturing.opts.Versions, 1) require.Equal(t, distribution.MigrationStagedDataKey(9, []byte("user|k")), capturing.opts.Versions[0].Key) } + +// failingMigrationImportStore fails ImportVersionsRaft with a fixed error so +// the FSM's ordinary-vs-halt classification can be exercised directly. +type failingMigrationImportStore struct { + store.MVCCStore + importErr error + floorErr error + maxTS uint64 +} + +func (s *failingMigrationImportStore) ImportVersionsRaft(_ context.Context, opts store.ImportVersionsOptions) (store.ImportVersionsResult, error) { + if s.importErr != nil { + return store.ImportVersionsResult{}, s.importErr + } + return store.ImportVersionsResult{AckedCursor: opts.Cursor, MaxImportedTS: s.maxTS}, nil +} + +func (s *failingMigrationImportStore) MigrationHLCFloor(context.Context, uint64) (uint64, error) { + if s.floorErr != nil { + return 0, s.floorErr + } + return 0, nil +} + +func migrationImportCommandPayload(t *testing.T) []byte { + t.Helper() + + data, err := proto.Marshal(&pb.ImportRangeVersionsRequest{ + JobId: 9, + BracketId: 1, + BatchSeq: 1, + Cursor: []byte("cursor"), + Versions: []*pb.MVCCVersion{ + {Key: []byte("user|k"), CommitTs: 10, Value: []byte("v")}, + }, + }) + require.NoError(t, err) + + return data +} + +// A store-side import failure is per-replica: the leader can apply the batch +// and ack the RPC while this voter skips the imported versions for good. The +// response must therefore halt the apply loop instead of letting the engine +// advance setApplied past the entry. +func TestApplyMigrationImportHaltsOnStoreFailure(t *testing.T) { + t.Parallel() + + pebbleIOErr := errors.New("pebble: background error") + fsm := &kvFSM{store: &failingMigrationImportStore{importErr: pebbleIOErr}} + + applied := fsm.applyMigrationImport(context.Background(), migrationImportCommandPayload(t)) + + err := haltApplyOf(applied) + require.Error(t, err, "store failure must halt apply, got %T: %v", applied, applied) + require.True(t, errors.Is(err, ErrMigrationImportApply), "got %v", err) + require.ErrorIs(t, err, pebbleIOErr) +} + +// A failed HLC-floor read is a store read on this replica only, so it halts +// for the same reason. Replay after restart is safe: the import batch already +// committed, and the replayed entry is recognised as a duplicate. +func TestApplyMigrationImportHaltsOnHLCFloorFailure(t *testing.T) { + t.Parallel() + + floorErr := errors.New("pebble: read failed") + fsm := &kvFSM{store: &failingMigrationImportStore{maxTS: 0, floorErr: floorErr}} + + applied := fsm.applyMigrationImport(context.Background(), migrationImportCommandPayload(t)) + + err := haltApplyOf(applied) + require.Error(t, err, "hlc floor failure must halt apply, got %T: %v", applied, applied) + require.True(t, errors.Is(err, ErrMigrationImportApply), "got %v", err) + require.ErrorIs(t, err, floorErr) +} + +func TestApplyMigrationImportHaltsOnUndecodablePayload(t *testing.T) { + t.Parallel() + + fsm := &kvFSM{store: &failingMigrationImportStore{}} + + applied := fsm.applyMigrationImport(context.Background(), []byte{0xff, 0xff, 0xff, 0xff}) + + err := haltApplyOf(applied) + require.Error(t, err, "undecodable payload must halt apply, got %T: %v", applied, applied) + require.True(t, errors.Is(err, ErrMigrationImportApply), "got %v", err) +} + +// Verdicts on the request bytes are reached identically by every replica, so +// they stay ordinary errors: the RPC caller sees them and the group advances +// setApplied in step. Halting on these would turn a malformed request into a +// cluster-wide outage. +func TestApplyMigrationImportOrdinaryErrorsDoNotHalt(t *testing.T) { + t.Parallel() + + tests := map[string]error{ + "batch gap": store.ErrImportBatchGap, + "invalid version": store.ErrInvalidImportVersion, + "value too large": store.ErrValueTooLarge, + "wrapped batch gap": errors.Wrap(store.ErrImportBatchGap, "context"), + "wrapped bad version": errors.Wrap(store.ErrInvalidImportVersion, "context"), + } + for name, importErr := range tests { + t.Run(name, func(t *testing.T) { + t.Parallel() + + fsm := &kvFSM{store: &failingMigrationImportStore{importErr: importErr}} + + applied := fsm.applyMigrationImport(context.Background(), migrationImportCommandPayload(t)) + + require.NoError(t, haltApplyOf(applied), "request-shaped error must not halt apply") + err, ok := applied.(error) + require.True(t, ok, "got %T: %v", applied, applied) + require.ErrorIs(t, err, importErr) + require.False(t, errors.Is(err, ErrMigrationImportApply)) + }) + } +} diff --git a/store/migration_versions.go b/store/migration_versions.go index 8008d971c..08a6d1dd6 100644 --- a/store/migration_versions.go +++ b/store/migration_versions.go @@ -368,14 +368,14 @@ func decodeMigrationPromotionStates(data []byte) (map[uint64]PromotionState, boo func validateImportVersion(version MVCCVersion) error { if version.CommitTS == 0 { - return errors.New("migration import version has zero commit_ts") + return errors.Wrap(ErrInvalidImportVersion, "migration import version has zero commit_ts") } if version.Tombstone { if version.ExpireAt != 0 { - return errors.New("migration import tombstone carries expire_at") + return errors.Wrap(ErrInvalidImportVersion, "migration import tombstone carries expire_at") } if len(version.Value) != 0 { - return errors.New("migration import tombstone carries value") + return errors.Wrap(ErrInvalidImportVersion, "migration import tombstone carries value") } return nil } diff --git a/store/store.go b/store/store.go index 5a746d8c3..8614c5efa 100644 --- a/store/store.go +++ b/store/store.go @@ -32,6 +32,13 @@ var ErrValueTooLarge = errors.New("value too large") var ErrInvalidExportCursor = errors.New("invalid export cursor") var ErrImportBatchGap = errors.New("migration import batch gap") +// ErrInvalidImportVersion marks a migration import version that is malformed +// on its face (zero commit_ts, a tombstone carrying a value or expire_at). +// It is a property of the request bytes, so every replica applying the same +// Raft entry reaches the same verdict -- which is what lets kv/fsm classify +// it as an ordinary apply error instead of halting the apply loop. +var ErrInvalidImportVersion = errors.New("invalid migration import version") + // validateValueSize returns ErrValueTooLarge when the value exceeds maxSnapshotValueSize. func validateValueSize(value []byte) error { if len(value) > maxSnapshotValueSize { From 8668bdcec2b651e73a26e827b75b9ab4cab582cd Mon Sep 17 00:00:00 2001 From: bootjp Date: Sun, 23 Aug 2026 17:24:38 +0900 Subject: [PATCH 07/58] migration: tombstone staged rows on prefix deletes While a route has staged visibility a logical key can exist only under !dist|migstage||..., and handleDelPrefix tombstoned the raw prefix alone. The staged row stayed the winning visible version, so prefix cleanup and flush reported success with the deleted key still readable. ShardStore.DeletePrefixAtRaftAt already propagates to staged prefixes, but that is not the production path: main.go builds each FSM directly over the per-group store via NewKvFSMWithHLC(st, ...), never over ShardStore, so the apply path has to propagate itself. Ordering is load-bearing. The raw delete bundles pendingApplyIdx into its Pebble batch, so it has to be the last write and the staged deletes pass appliedIndex 0. A crash between them then replays the whole entry, and the staged tombstones are idempotent at the same commitTS. Doing the raw delete first would advance the applied index and let cold-start replay skip the staged half for good -- the same class of bug this is fixing. Staged deletes are scoped to this FSM's own group and its route's migration job, so another job's staged rows under the same logical prefix are untouched; a test pins that. Claude-Session: https://claude.ai/code/session_013rNHooj7NF3giihWVba8QE --- kv/fsm.go | 57 +++++++++++++++++++++++++++++ kv/fsm_onephase_dedup_test.go | 69 +++++++++++++++++++++++++++++++++++ 2 files changed, 126 insertions(+) diff --git a/kv/fsm.go b/kv/fsm.go index 55a95216d..7af683f63 100644 --- a/kv/fsm.go +++ b/kv/fsm.go @@ -587,6 +587,26 @@ func (f *kvFSM) handleDelPrefix(ctx context.Context, prefix []byte, commitTS uin if err := f.verifyRouteWriteTimestampFloorForPrefix(prefix, commitTS); err != nil { return err } + // Staged rows first, and deliberately with appliedIndex 0. + // + // While a route has staged visibility a logical key can exist only under + // !dist|migstage||..., and a raw-prefix tombstone leaves that staged + // row as the winning visible version -- so a prefix delete would report + // success with the key still readable. ShardStore.DeletePrefixAtRaftAt + // already propagates to staged prefixes, but production builds each FSM + // directly over the per-group store (main.go's NewKvFSMWithHLC(st, ...)), + // never over ShardStore, so this apply path has to do it itself. + // + // Ordering is load-bearing: the raw delete bundles pendingApplyIdx into its + // batch, so it must be the last write. A crash before it replays the whole + // entry, and the staged tombstones are idempotent at the same commitTS. + // Doing the raw delete first would advance the applied index and let replay + // skip the staged half permanently. + for _, staged := range f.stagedVisibilityPrefixDeletes(prefix) { + if err := f.store.DeletePrefixAtRaftAt(ctx, staged.prefix, staged.excludePrefix, commitTS, 0); err != nil { + return errors.WithStack(err) + } + } if err := f.store.DeletePrefixAtRaftAt(ctx, prefix, txnCommonPrefix, commitTS, f.pendingApplyIdx); err != nil { return errors.WithStack(err) } @@ -594,6 +614,43 @@ func (f *kvFSM) handleDelPrefix(ctx context.Context, prefix []byte, commitTS uin return nil } +type fsmStagedPrefixDelete struct { + prefix []byte + excludePrefix []byte +} + +// stagedVisibilityPrefixDeletes returns the staged prefixes this shard group +// must tombstone alongside a raw prefix delete. It mirrors +// ShardStore.stagedVisibilityPrefixDeletes but is scoped to this FSM's own +// group, which is the only store it can write. +func (f *kvFSM) stagedVisibilityPrefixDeletes(prefix []byte) []fsmStagedPrefixDelete { + if f == nil || f.routes == nil { + return nil + } + snap, ok := f.routes.Current() + if !ok { + return nil + } + start, end := routePrefixRange(prefix) + routes := snap.IntersectingRoutes(start, end) + out := make([]fsmStagedPrefixDelete, 0, len(routes)) + seen := make(map[string]struct{}, len(routes)) + for _, route := range routes { + if route.GroupID != f.shardGroupID || !routeHasStagedVisibility(route) { + continue + } + stagedPrefix := distribution.MigrationStagedDataKey(route.MigrationJobID, prefix) + stagedExclude := distribution.MigrationStagedDataKey(route.MigrationJobID, txnCommonPrefix) + dedupeKey := string(stagedPrefix) + "\x00" + string(stagedExclude) + if _, dup := seen[dedupeKey]; dup { + continue + } + seen[dedupeKey] = struct{}{} + out = append(out, fsmStagedPrefixDelete{prefix: stagedPrefix, excludePrefix: stagedExclude}) + } + return out +} + func (f *kvFSM) verifyRouteNotFencedForKey(key []byte) error { if f.routes == nil { return nil diff --git a/kv/fsm_onephase_dedup_test.go b/kv/fsm_onephase_dedup_test.go index b675e1079..f5dc220ed 100644 --- a/kv/fsm_onephase_dedup_test.go +++ b/kv/fsm_onephase_dedup_test.go @@ -235,3 +235,72 @@ func TestOnePhaseDedup_OtherTxnVersionDoesNotMaskRetry(t *testing.T) { require.NoError(t, err) require.Equal(t, []byte("v"), val, "retry's write must be readable; exactness loss would have lost it") } + +// A key that exists only under the staged prefix stays the winning visible +// version if a prefix delete tombstones the raw prefix alone, so prefix cleanup +// reports success while the key is still readable. ShardStore propagates to the +// staged prefix, but production builds the FSM over the per-group store, so the +// apply path has to do it itself. +func TestDelPrefix_TombstonesStagedRowsUnderStagedVisibility(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") + stagedKey := distribution.MigrationStagedDataKey(9, key) + // The row exists ONLY under the staged prefix, which is the case a raw + // tombstone cannot reach. + require.NoError(t, st.PutAt(ctx, stagedKey, []byte("v"), 20, 0)) + + require.NoError(t, fsm.handleDelPrefix(ctx, []byte("list-"), 30)) + + staged, err := st.GetAt(ctx, stagedKey, 40) + require.ErrorIs(t, err, store.ErrKeyNotFound, + "the staged row must be tombstoned by the prefix delete") + require.Nil(t, staged) +} + +// The staged tombstone must not escape its own migration job or this shard +// group: another job's staged rows under the same logical prefix are a +// different route's data. +func TestDelPrefix_LeavesOtherJobsStagedRowsIntact(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") + otherJobKey := distribution.MigrationStagedDataKey(11, key) + require.NoError(t, st.PutAt(ctx, otherJobKey, []byte("v"), 20, 0)) + + require.NoError(t, fsm.handleDelPrefix(ctx, []byte("list-"), 30)) + + got, err := st.GetAt(ctx, otherJobKey, 40) + require.NoError(t, err, "a different migration job's staged rows must survive") + require.Equal(t, []byte("v"), got) +} From b7aa8826735d49cd6b918c8eceb9e2747d907f06 Mon Sep 17 00:00:00 2001 From: bootjp Date: Sun, 23 Aug 2026 19:01:38 +0900 Subject: [PATCH 08/58] Revert "migration: tombstone staged rows on prefix deletes" This reverts 8668bdce. The fix decided whether to write a staged tombstone from f.routes.Current(), which is WrapDistributionEngine over the route engine that CatalogWatcher.Run refreshes on a polling timer -- replica-local state that is not synchronized with Raft apply. A follower applying the prefix-delete entry before its watcher installs the CUTOVER descriptor therefore skips the staged tombstone while still persisting the raw delete and the applied index. Its replay can never revisit the entry, so staged-only rows stay readable on that replica alone. That is replica divergence, which is worse than the uniform stale-row bug it was meant to fix: before, every replica agreed. The underlying bug is real and stays open. Fixing it needs routing that is durable and apply-consistent -- the migration job and ranges carried in the proposed command, rather than read from replica-local catalog state at apply time. That changes the Raft command format, so per CLAUDE.md it wants a design doc before implementation. Reverting the test too: TestDelPrefix_LeavesOtherJobsStagedRowsIntact asserted that another job's staged copy of a deleted logical key survives, which was my scoping assumption rather than a verified requirement, and a promotion could resurrect deleted data under it. Claude-Session: https://claude.ai/code/session_013rNHooj7NF3giihWVba8QE --- kv/fsm.go | 57 ----------------------------- kv/fsm_onephase_dedup_test.go | 69 ----------------------------------- 2 files changed, 126 deletions(-) diff --git a/kv/fsm.go b/kv/fsm.go index 7af683f63..55a95216d 100644 --- a/kv/fsm.go +++ b/kv/fsm.go @@ -587,26 +587,6 @@ func (f *kvFSM) handleDelPrefix(ctx context.Context, prefix []byte, commitTS uin if err := f.verifyRouteWriteTimestampFloorForPrefix(prefix, commitTS); err != nil { return err } - // Staged rows first, and deliberately with appliedIndex 0. - // - // While a route has staged visibility a logical key can exist only under - // !dist|migstage||..., and a raw-prefix tombstone leaves that staged - // row as the winning visible version -- so a prefix delete would report - // success with the key still readable. ShardStore.DeletePrefixAtRaftAt - // already propagates to staged prefixes, but production builds each FSM - // directly over the per-group store (main.go's NewKvFSMWithHLC(st, ...)), - // never over ShardStore, so this apply path has to do it itself. - // - // Ordering is load-bearing: the raw delete bundles pendingApplyIdx into its - // batch, so it must be the last write. A crash before it replays the whole - // entry, and the staged tombstones are idempotent at the same commitTS. - // Doing the raw delete first would advance the applied index and let replay - // skip the staged half permanently. - for _, staged := range f.stagedVisibilityPrefixDeletes(prefix) { - if err := f.store.DeletePrefixAtRaftAt(ctx, staged.prefix, staged.excludePrefix, commitTS, 0); err != nil { - return errors.WithStack(err) - } - } if err := f.store.DeletePrefixAtRaftAt(ctx, prefix, txnCommonPrefix, commitTS, f.pendingApplyIdx); err != nil { return errors.WithStack(err) } @@ -614,43 +594,6 @@ func (f *kvFSM) handleDelPrefix(ctx context.Context, prefix []byte, commitTS uin return nil } -type fsmStagedPrefixDelete struct { - prefix []byte - excludePrefix []byte -} - -// stagedVisibilityPrefixDeletes returns the staged prefixes this shard group -// must tombstone alongside a raw prefix delete. It mirrors -// ShardStore.stagedVisibilityPrefixDeletes but is scoped to this FSM's own -// group, which is the only store it can write. -func (f *kvFSM) stagedVisibilityPrefixDeletes(prefix []byte) []fsmStagedPrefixDelete { - if f == nil || f.routes == nil { - return nil - } - snap, ok := f.routes.Current() - if !ok { - return nil - } - start, end := routePrefixRange(prefix) - routes := snap.IntersectingRoutes(start, end) - out := make([]fsmStagedPrefixDelete, 0, len(routes)) - seen := make(map[string]struct{}, len(routes)) - for _, route := range routes { - if route.GroupID != f.shardGroupID || !routeHasStagedVisibility(route) { - continue - } - stagedPrefix := distribution.MigrationStagedDataKey(route.MigrationJobID, prefix) - stagedExclude := distribution.MigrationStagedDataKey(route.MigrationJobID, txnCommonPrefix) - dedupeKey := string(stagedPrefix) + "\x00" + string(stagedExclude) - if _, dup := seen[dedupeKey]; dup { - continue - } - seen[dedupeKey] = struct{}{} - out = append(out, fsmStagedPrefixDelete{prefix: stagedPrefix, excludePrefix: stagedExclude}) - } - return out -} - func (f *kvFSM) verifyRouteNotFencedForKey(key []byte) error { if f.routes == nil { return nil diff --git a/kv/fsm_onephase_dedup_test.go b/kv/fsm_onephase_dedup_test.go index f5dc220ed..b675e1079 100644 --- a/kv/fsm_onephase_dedup_test.go +++ b/kv/fsm_onephase_dedup_test.go @@ -235,72 +235,3 @@ func TestOnePhaseDedup_OtherTxnVersionDoesNotMaskRetry(t *testing.T) { require.NoError(t, err) require.Equal(t, []byte("v"), val, "retry's write must be readable; exactness loss would have lost it") } - -// A key that exists only under the staged prefix stays the winning visible -// version if a prefix delete tombstones the raw prefix alone, so prefix cleanup -// reports success while the key is still readable. ShardStore propagates to the -// staged prefix, but production builds the FSM over the per-group store, so the -// apply path has to do it itself. -func TestDelPrefix_TombstonesStagedRowsUnderStagedVisibility(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") - stagedKey := distribution.MigrationStagedDataKey(9, key) - // The row exists ONLY under the staged prefix, which is the case a raw - // tombstone cannot reach. - require.NoError(t, st.PutAt(ctx, stagedKey, []byte("v"), 20, 0)) - - require.NoError(t, fsm.handleDelPrefix(ctx, []byte("list-"), 30)) - - staged, err := st.GetAt(ctx, stagedKey, 40) - require.ErrorIs(t, err, store.ErrKeyNotFound, - "the staged row must be tombstoned by the prefix delete") - require.Nil(t, staged) -} - -// The staged tombstone must not escape its own migration job or this shard -// group: another job's staged rows under the same logical prefix are a -// different route's data. -func TestDelPrefix_LeavesOtherJobsStagedRowsIntact(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") - otherJobKey := distribution.MigrationStagedDataKey(11, key) - require.NoError(t, st.PutAt(ctx, otherJobKey, []byte("v"), 20, 0)) - - require.NoError(t, fsm.handleDelPrefix(ctx, []byte("list-"), 30)) - - got, err := st.GetAt(ctx, otherJobKey, 40) - require.NoError(t, err, "a different migration job's staged rows must survive") - require.Equal(t, []byte("v"), got) -} From ddfbdb0cf781ac8ab5a1b3d6db55162338cd18c4 Mon Sep 17 00:00:00 2001 From: bootjp Date: Mon, 24 Aug 2026 21:45:55 +0900 Subject: [PATCH 09/58] migration: export filesystem chunk payloads File chunk payloads are stored under !fs|chk| but route through a virtual !fs|route|chk| key produced by fskeys.ExtractRouteKey. A filesystem-chunk route's catalog interval is therefore the virtual range, and "!fs|chk|" sorts below "!fs|route|chk|", so the user bracket's raw interval never reached a single payload row. No family bracket covered them either. The export consequently found nothing for those routes while the bracket still completed and was promoted, so a cross-group split between file boundaries lost every chunk of the moved files. Add the family bracket, which scans the raw chunk prefix. The route filter still applies: migrationBracketRouteCheck defaults to RequiresRouteKeyCheck, and routeKey normalizes a raw chunk key through fskeys.ExtractRouteKey before the interval test, so only chunks belonging to the migrated route are exported. The constant is appended after MigrationFamilyLegacyListMetaDelta rather than grouped with related families: the values are the migration wire contract and inserting earlier would renumber families peers already use. Filesystem usage counters route through a virtual key the same way (extractUsageRouteKey) and have no bracket either. That is not fixed here -- it needs the same check and is called out separately. Claude-Session: https://claude.ai/code/session_013rNHooj7NF3giihWVba8QE --- distribution/migrator.go | 13 +++++++++ distribution/migrator_export_plan_test.go | 35 +++++++++++++++++++++++ 2 files changed, 48 insertions(+) diff --git a/distribution/migrator.go b/distribution/migrator.go index db32062a3..ef22e5459 100644 --- a/distribution/migrator.go +++ b/distribution/migrator.go @@ -4,6 +4,7 @@ import ( "bytes" "encoding/binary" + "github.com/bootjp/elastickv/internal/fskeys" "github.com/bootjp/elastickv/internal/s3keys" "github.com/bootjp/elastickv/store" "github.com/cockroachdb/errors" @@ -53,6 +54,10 @@ const ( MigrationFamilyS3Blob MigrationFamilyS3GCUpload MigrationFamilyLegacyListMetaDelta + // MigrationFamilyFilesystemChunk is appended last on purpose: the values + // above are the migration wire contract, so inserting anywhere earlier + // would renumber families that peers already use. + MigrationFamilyFilesystemChunk ) const ( @@ -477,6 +482,14 @@ func migrationFamilyBrackets() []MigrationBracket { {family: MigrationFamilyS3UploadPart, prefix: s3keys.UploadPartPrefix}, {family: MigrationFamilyS3Blob, prefix: s3keys.BlobPrefix}, {family: MigrationFamilyS3GCUpload, prefix: s3keys.GCUploadPrefix}, + // File chunk payloads live under !fs|chk| but route through a virtual + // !fs|route|chk| key via fskeys.ExtractRouteKey. "!fs|chk|" sorts below + // "!fs|route|chk|", so the user bracket's raw interval never reaches + // them and, without this bracket, a cross-group split completed and + // promoted while every chunk of the moved files stayed behind. The + // default route-key check in migrationBracketRouteCheck applies the + // logical route filter to the raw scan. + {family: MigrationFamilyFilesystemChunk, prefix: string(fskeys.ChunkAllPrefix())}, } out := make([]MigrationBracket, 0, len(defs)) diff --git a/distribution/migrator_export_plan_test.go b/distribution/migrator_export_plan_test.go index c26394a07..79a0135bd 100644 --- a/distribution/migrator_export_plan_test.go +++ b/distribution/migrator_export_plan_test.go @@ -5,6 +5,7 @@ import ( "encoding/binary" "testing" + "github.com/bootjp/elastickv/internal/fskeys" "github.com/bootjp/elastickv/internal/s3keys" "github.com/bootjp/elastickv/store" "github.com/cockroachdb/errors" @@ -519,3 +520,37 @@ func legacyListMetaDeltaKey(userKey []byte, commitTS uint64, seqInTxn uint32) [] binary.BigEndian.PutUint32(seq[:], seqInTxn) return append(key, seq[:]...) } + +// File chunk payloads live under !fs|chk| but route through a virtual +// !fs|route|chk| key. "!fs|chk|" sorts below "!fs|route|chk|", so a +// filesystem-chunk route's user bracket -- whose raw interval IS the virtual +// route range -- never reaches the payloads. Without a dedicated bracket the +// export found nothing, yet the migration completed and was promoted, losing +// every chunk of the moved files on a cross-group split. +func TestPlanMigrationBracketsCoversFilesystemChunkPayloads(t *testing.T) { + t.Parallel() + + routeStart := fskeys.ChunkRouteKey(0, 1) + routeEnd := fskeys.ChunkRouteKey(0, 9) + brackets, err := PlanMigrationBrackets(routeStart, routeEnd) + require.NoError(t, err) + + var chunk *MigrationBracket + for i := range brackets { + if brackets[i].Family == MigrationFamilyFilesystemChunk { + chunk = &brackets[i] + + break + } + } + require.NotNil(t, chunk, "the plan must carry a filesystem chunk bracket") + require.Equal(t, fskeys.ChunkAllPrefix(), chunk.Start, + "the bracket must scan the raw chunk prefix, not the virtual route range") + require.True(t, chunk.RequiresRouteKeyCheck, + "raw chunk keys must still be filtered through the logical route") + + // The gap this closes: the raw payload prefix sorts below the virtual route + // interval, so the user bracket's raw range cannot reach it. + require.Negative(t, bytes.Compare(fskeys.ChunkAllPrefix(), routeStart), + "chunk payloads sort below the virtual route interval") +} From a32559a5c9b10ca49ca9a91a6bce3caea6402159 Mon Sep 17 00:00:00 2001 From: bootjp Date: Mon, 24 Aug 2026 22:14:00 +0900 Subject: [PATCH 10/58] migration: stop double-exporting chunks; normalize ownership lookups The filesystem chunk bracket added in ddfbdb0c did not add its prefix to migrationInternalFamilyPrefixes. A route interval spanning both the raw !fs|chk| key and its normalized !fs|route|chk| route is accepted by the family bracket on prefix and by the user bracket on route key, so every chunk version was exported and proposed through Raft twice under separate bracket IDs. My regression from that commit. Every other family bracket's prefix is already in that list, so the fix is the missing entry. The test is written as an invariant over all family brackets rather than a chunk-specific case, because the chunk family was added without its exclusion and nothing caught it. Separately, GetRouteOwnership looked the raw request key up in the historical snapshot while GetRoute normalizes through kv.RouteKey. For an internal storage key -- a filesystem chunk, a Redis collection row -- the raw family prefix and the logical route key sort into different routes, so the RPC reported the raw-prefix owner instead of the group that owned the key at that catalog version. It now normalizes the same way. GetRouteOwnership has no in-repo callers beyond the generated gRPC plumbing, so the corrected answer reaches only external operator clients, which is the intended contract. Claude-Session: https://claude.ai/code/session_013rNHooj7NF3giihWVba8QE --- adapter/distribution_server.go | 7 ++++- adapter/distribution_server_test.go | 33 +++++++++++++++++++++++ distribution/migrator.go | 6 +++++ distribution/migrator_export_plan_test.go | 21 +++++++++++++++ 4 files changed, 66 insertions(+), 1 deletion(-) diff --git a/adapter/distribution_server.go b/adapter/distribution_server.go index 65b448a10..7ebf96708 100644 --- a/adapter/distribution_server.go +++ b/adapter/distribution_server.go @@ -185,7 +185,12 @@ func (s *DistributionServer) GetRouteOwnership(ctx context.Context, req *pb.GetR if err != nil { return nil, err } - route, ok := snapshot.RouteOf(req.GetKey()) + // Normalized exactly like GetRoute above. An internal storage key -- a + // filesystem chunk, a Redis collection row -- routes by its logical key, + // so looking the raw bytes up in the snapshot answers with the owner of + // the raw family prefix instead of the group that actually owned the key + // at that catalog version. + route, ok := snapshot.RouteOf(kv.RouteKey(req.GetKey())) if !ok { return &pb.GetRouteOwnershipResponse{ CatalogVersion: snapshot.Version(), diff --git a/adapter/distribution_server_test.go b/adapter/distribution_server_test.go index 46a6502e0..0b6ee9bab 100644 --- a/adapter/distribution_server_test.go +++ b/adapter/distribution_server_test.go @@ -1372,3 +1372,36 @@ type recordingDistributionFilesystemObserver struct { func (o *recordingDistributionFilesystemObserver) ObserveFilePinnedHotspot(reason string) { o.reasons = append(o.reasons, reason) } + +// GetRouteOwnership answers the historical owner of a key, so it has to +// normalize the same way GetRoute does. An internal storage key -- here a +// filesystem chunk -- routes by its logical !fs|route|chk| key, and the raw +// !fs|chk| bytes sort into a different route entirely. Without normalization +// the RPC reports the raw-prefix owner rather than the group that owned the +// key at that catalog version. +func TestDistributionServerGetRouteOwnership_NormalizesFilesystemChunkKeys(t *testing.T) { + t.Parallel() + + home := uint64(11) + inode := uint64(22) + routeKey := fskeys.ChunkRouteKey(home, inode) + engine := distribution.NewEngine() + require.NoError(t, engine.ApplySnapshot(distribution.CatalogSnapshot{ + Version: 3, + Routes: []distribution.RouteDescriptor{ + {RouteID: 1, Start: []byte(""), End: routeKey, GroupID: 1, State: distribution.RouteStateActive}, + {RouteID: 2, Start: routeKey, End: nil, GroupID: 2, State: distribution.RouteStateActive}, + }, + })) + + s := NewDistributionServer(engine, nil) + resp, err := s.GetRouteOwnership(context.Background(), &pb.GetRouteOwnershipRequest{ + Key: fskeys.ChunkKey(home, inode, 99), + CatalogVersion: 3, + }) + require.NoError(t, err) + require.True(t, resp.Found) + require.Equal(t, uint64(2), resp.Route.RaftGroupId, + "the chunk must resolve to its logical route owner, not the raw-prefix owner") + require.Equal(t, uint64(2), resp.Route.RouteId) +} diff --git a/distribution/migrator.go b/distribution/migrator.go index ef22e5459..b0eed506a 100644 --- a/distribution/migrator.go +++ b/distribution/migrator.go @@ -145,6 +145,12 @@ var migrationInternalFamilyPrefixes = [][]byte{ []byte(s3keys.UploadPartPrefix), []byte(s3keys.BlobPrefix), []byte(s3keys.GCUploadPrefix), + // Every family bracket's prefix belongs here so the user bracket does not + // also export those rows. A filesystem-chunk route interval can span both + // the raw !fs|chk| key and its normalized !fs|route|chk| route, and both + // route filters accept the row, so omitting this exported every chunk + // version twice under two bracket IDs. + fskeys.ChunkAllPrefix(), } // MigrationBracket is a raw MVCC export or drain slice used by the migrator. diff --git a/distribution/migrator_export_plan_test.go b/distribution/migrator_export_plan_test.go index 79a0135bd..fd65f9608 100644 --- a/distribution/migrator_export_plan_test.go +++ b/distribution/migrator_export_plan_test.go @@ -554,3 +554,24 @@ func TestPlanMigrationBracketsCoversFilesystemChunkPayloads(t *testing.T) { require.Negative(t, bytes.Compare(fskeys.ChunkAllPrefix(), routeStart), "chunk payloads sort below the virtual route interval") } + +// Invariant: every family bracket's scan prefix must also be excluded from the +// user bracket. Both filters accept the same raw row otherwise -- the family +// bracket by prefix and the user bracket by normalized route key -- so the rows +// are exported and proposed through Raft twice under separate bracket IDs. +// +// This is written as an invariant rather than a per-family case because the +// filesystem chunk family was added without its exclusion and nothing caught it. +func TestEveryFamilyBracketPrefixIsExcludedFromUserBracket(t *testing.T) { + t.Parallel() + + for _, bracket := range migrationFamilyBrackets() { + if bracket.DrainOnly { + continue + } + require.True(t, IsMigrationKnownInternalKey(bracket.Start), + "family %d prefix %q must be in migrationInternalFamilyPrefixes, "+ + "otherwise the user bracket exports the same rows a second time", + bracket.Family, bracket.Start) + } +} From 98e9a1cdbcc443b6557b829682158983812b3466 Mon Sep 17 00:00:00 2001 From: bootjp Date: Mon, 24 Aug 2026 22:42:06 +0900 Subject: [PATCH 11/58] migration: export routed filesystem usage counters Per-route usage counters are stored at !fs|usage|route| and normalize back to the embedded logical route key, so like chunk payloads their raw key sits outside a user route's interval and no family bracket covered them. The counter therefore stayed on the source across a cross-group split. After cutover the usage scan filtered that copy out, because its logical owner had become the target, while target-side updates started from zero -- so StatFS undercounted the files and bytes that were actually moved. The raw prefix is added to migrationInternalFamilyPrefixes in the same change, which the family-bracket invariant test now enforces rather than leaving to be remembered. This is the gap the previous commit called out as unfixed. It is the same shape as the chunk family, and the test additionally pins the round trip -- UsageRouteKey("customers") normalizes back to "customers" -- since the route filter keeping the counter depends on exactly that. Claude-Session: https://claude.ai/code/session_013rNHooj7NF3giihWVba8QE --- distribution/migrator.go | 10 ++++++ distribution/migrator_export_plan_test.go | 39 +++++++++++++++++++++++ 2 files changed, 49 insertions(+) diff --git a/distribution/migrator.go b/distribution/migrator.go index b0eed506a..b842b9365 100644 --- a/distribution/migrator.go +++ b/distribution/migrator.go @@ -58,6 +58,7 @@ const ( // above are the migration wire contract, so inserting anywhere earlier // would renumber families that peers already use. MigrationFamilyFilesystemChunk + MigrationFamilyFilesystemUsage ) const ( @@ -151,6 +152,7 @@ var migrationInternalFamilyPrefixes = [][]byte{ // route filters accept the row, so omitting this exported every chunk // version twice under two bracket IDs. fskeys.ChunkAllPrefix(), + fskeys.UsageRouteAllPrefix(), } // MigrationBracket is a raw MVCC export or drain slice used by the migrator. @@ -496,6 +498,14 @@ func migrationFamilyBrackets() []MigrationBracket { // default route-key check in migrationBracketRouteCheck applies the // logical route filter to the raw scan. {family: MigrationFamilyFilesystemChunk, prefix: string(fskeys.ChunkAllPrefix())}, + // Per-route usage counters are stored at !fs|usage|route| + // and normalize back to the embedded logical route key, so like chunks + // their raw key sits outside the user bracket's interval. Without this + // bracket the counter stayed on the source: after cutover the usage scan + // filtered that copy out because its logical owner had become the + // target, while target-side updates began from zero, so StatFS + // undercounted existing files and bytes. + {family: MigrationFamilyFilesystemUsage, prefix: string(fskeys.UsageRouteAllPrefix())}, } out := make([]MigrationBracket, 0, len(defs)) diff --git a/distribution/migrator_export_plan_test.go b/distribution/migrator_export_plan_test.go index fd65f9608..b2c0a4bbd 100644 --- a/distribution/migrator_export_plan_test.go +++ b/distribution/migrator_export_plan_test.go @@ -575,3 +575,42 @@ func TestEveryFamilyBracketPrefixIsExcludedFromUserBracket(t *testing.T) { bracket.Family, bracket.Start) } } + +// Per-route usage counters live at !fs|usage|route| and +// normalize back to the embedded logical route key, so their raw key sits +// outside a user route's interval exactly like chunk payloads. Without a +// bracket the counter stayed on the source: after cutover the usage scan +// filtered that copy out because its logical owner had become the target, +// while target-side updates began from zero, so StatFS undercounted. +func TestPlanMigrationBracketsCoversFilesystemUsageCounters(t *testing.T) { + t.Parallel() + + routeStart := []byte("a") + routeEnd := []byte("z") + brackets, err := PlanMigrationBrackets(routeStart, routeEnd) + require.NoError(t, err) + + var usage *MigrationBracket + for i := range brackets { + if brackets[i].Family == MigrationFamilyFilesystemUsage { + usage = &brackets[i] + + break + } + } + require.NotNil(t, usage, "the plan must carry a filesystem usage bracket") + require.Equal(t, fskeys.UsageRouteAllPrefix(), usage.Start, + "the bracket must scan the raw usage prefix, not the logical route range") + require.True(t, usage.RequiresRouteKeyCheck, + "raw usage keys must still be filtered through their embedded route") + + // The gap this closes: the raw counter key sorts outside a user route + // interval, so the user bracket cannot reach it. + require.Negative(t, bytes.Compare(fskeys.UsageRouteAllPrefix(), routeStart), + "usage counters sort below a user route interval") + + // And it must round-trip: a counter for a key inside the interval + // normalizes back into that interval, so the route filter keeps it. + counter := fskeys.UsageRouteKey([]byte("customers")) + require.Equal(t, []byte("customers"), fskeys.ExtractRouteKey(counter)) +} From b1896b77c4b549ee09de78afec3a3c77a3696b42 Mon Sep 17 00:00:00 2001 From: bootjp Date: Fri, 28 Aug 2026 04:12:14 +0900 Subject: [PATCH 12/58] migration: reject foreign explicit-group reads after promotion routeForExplicitGroupKey rejected an explicit-group read whose catalog route named a different group only while that route still carried StagedVisibilityActive. Promotion clears that flag, and at exactly that moment the source's former range belongs to the target while the source's pre-cutover MVCC is still sitting there waiting for cleanup. A coordinator that had not yet applied the new catalog version kept forwarding the source group, the mismatch fell through to the requested-group fallback, and the read returned the value from before the cutover. Fail closed on the mismatch instead, and allow the fallback only for keyspaces whose owning group is chosen by a resolver rather than by the byte-range catalog. Two exist: - SQS HT-FIFO, whose (queue, partition) resolver picks the group. - Filesystem chunks, placed per group by the filesystem's own home-slot placement. ListFilePlacementStats scans the whole chunk keyspace once per FilesystemGroupIDs entry, so most of those groups are not the catalog owner of the range and never will be. Without this the first non-owning group would have failed the whole placement scan. The same fallback existed on both explicit-group scan paths and is gated the same way there. TestShardStoreExplicitGroupRead_IgnoresUnrelatedStagedRoutes asserted the old override as a side effect of what it was actually testing, so its fixture now puts the read's own range on the group it names; the unrelated staged route it exists for is unchanged. Claude-Session: https://claude.ai/code/session_013rNHooj7NF3giihWVba8QE --- kv/shard_store.go | 45 +++++++++++++++++++ kv/shard_store_test.go | 100 ++++++++++++++++++++++++++++++++++++++++- 2 files changed, 143 insertions(+), 2 deletions(-) diff --git a/kv/shard_store.go b/kv/shard_store.go index 0cdfaf700..a5472bc06 100644 --- a/kv/shard_store.go +++ b/kv/shard_store.go @@ -36,6 +36,7 @@ type ShardStore struct { var ( ErrCrossShardMutationBatchNotSupported = errors.New("cross-shard mutation batches are not supported") ErrExplicitGroupStagedVisibilityUnresolved = errors.New("explicit group read cannot resolve staged visibility route") + ErrExplicitGroupRouteOwnerMismatch = errors.New("explicit group read does not own the requested key range") ErrReadRouteVersionUnavailable = errors.New("read route version is not locally available") ErrFilesystemPlacementTargetNotFound = errors.New("filesystem placement target group has no routable home slot") ) @@ -299,10 +300,44 @@ func (s *ShardStore) routeForExplicitGroupKey(groupID uint64, key []byte) (distr if routeHasStagedVisibility(route) { return distribution.Route{}, errors.Wrapf(ErrExplicitGroupStagedVisibilityUnresolved, "group_id=%d key=%q", groupID, key) } + if !explicitGroupResolverOwnedKey(key) { + return distribution.Route{}, errors.Wrapf( + ErrExplicitGroupRouteOwnerMismatch, + "group_id=%d catalog_group_id=%d key=%q", groupID, route.GroupID, key) + } } return fallback, nil } +// explicitGroupResolverOwnedKey reports whether key belongs to a keyspace whose +// owning group is chosen by a resolver rather than by the byte-range catalog -- +// SQS HT-FIFO's (queue, partition) resolver being the one such keyspace today. +// Only those keys may be served from a group the catalog does not name. +// +// Everything else must fail closed on a mismatch. Once a migration's promotion +// clears StagedVisibilityActive, the source's former range belongs to the +// target, and a coordinator that has not yet applied that catalog version keeps +// forwarding the source group. Falling back to the requested group then serves +// the source's own MVCC -- the value from before the cutover -- for as long as +// source cleanup takes. The staged-visibility rejection above stops covering +// that window the moment the flag is cleared, which is exactly when the +// pre-cutover data is still sitting there. +// fsChunkAllPrefix is the raw filesystem chunk keyspace prefix, hoisted so the +// explicit-group gate does not allocate it per read. +var fsChunkAllPrefix = fskeys.ChunkAllPrefix() + +func explicitGroupResolverOwnedKey(key []byte) bool { + if sqsRouteKey(key) != nil { + return true + } + // Filesystem chunks are placed per group by the filesystem's own home-slot + // placement, not by the byte-range catalog: ListFilePlacementStats scans the + // whole chunk keyspace once per FilesystemGroupIDs entry + // (internal/filesystem/placement.go scanPlacementChunks), so most of those + // groups are not the catalog owner of the range and never will be. + return bytes.HasPrefix(key, fsChunkAllPrefix) +} + func (s *ShardStore) getAtWithStagedVisibility(ctx context.Context, g *ShardGroup, route distribution.Route, key []byte, ts uint64) ([]byte, error) { if err := ensureReadTSRetained(g.Store, ts); err != nil { return nil, err @@ -912,6 +947,11 @@ func (s *ShardStore) routesForExplicitGroupRouteBounds(groupID uint64, start []b } } if len(matched) == 0 { + if len(routes) > 0 && !explicitGroupResolverOwnedKey(start) { + return nil, false, errors.Wrapf( + ErrExplicitGroupRouteOwnerMismatch, + "group_id=%d range=[%q,%q)", groupID, start, end) + } return fallback, false, nil } return matched, false, nil @@ -940,6 +980,11 @@ func (s *ShardStore) routesForExplicitGroupScan(groupID uint64, start []byte, en } return matched, !routeMapped, nil } + if len(routes) > 0 && !explicitGroupResolverOwnedKey(start) { + return nil, false, errors.Wrapf( + ErrExplicitGroupRouteOwnerMismatch, + "group_id=%d range=[%q,%q)", groupID, start, end) + } return fallback, false, nil } diff --git a/kv/shard_store_test.go b/kv/shard_store_test.go index 019d762a4..72ed2d8e8 100644 --- a/kv/shard_store_test.go +++ b/kv/shard_store_test.go @@ -700,12 +700,14 @@ func TestShardStoreExplicitGroupRead_IgnoresUnrelatedStagedRoutes(t *testing.T) 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}, + // The read's own range is owned by the group it names, so the + // staged route below is genuinely unrelated to it. + {RouteID: 1, Start: []byte("a"), End: []byte("m"), GroupID: 1, State: distribution.RouteStateActive}, { RouteID: 2, Start: []byte("m"), End: []byte("z"), - GroupID: 1, + GroupID: 2, State: distribution.RouteStateActive, StagedVisibilityActive: true, MigrationJobID: 9, @@ -3461,3 +3463,97 @@ func TestShardStoreScanAt_ExactLegacyListDeltaScanMarksRouteGroup(t *testing.T) require.Equal(t, uint64(2), kvs[0].RouteGroupID, "an exact legacy delta scan must still report the owning route group") } + +// A coordinator that has not yet applied a promotion keeps forwarding the +// pre-cutover source group. Once StagedVisibilityActive is cleared the source's +// former range belongs to the target, and the staged-visibility rejection stops +// covering the request -- exactly while the source's pre-cutover MVCC is still +// sitting there waiting for cleanup. Serving that is a stale read, so the +// mismatch must fail closed instead. +func TestShardStoreExplicitGroupRead_FailsClosedAfterPromotionClearsStaging(t *testing.T) { + t.Parallel() + + ctx := context.Background() + engine := distribution.NewEngine() + require.NoError(t, engine.ApplySnapshot(distribution.CatalogSnapshot{ + Version: 2, + Routes: []distribution.RouteDescriptor{ + // Promotion completed: the range now belongs to group 2 and the + // staged-visibility flag is gone. + {RouteID: 1, Start: []byte("a"), End: []byte("z"), GroupID: 2, State: distribution.RouteStateActive}, + }, + })) + groups := map[uint64]*ShardGroup{ + 1: {Store: store.NewMVCCStore()}, + 2: {Store: store.NewMVCCStore()}, + } + st := NewShardStore(engine, groups) + require.NoError(t, groups[1].Store.PutAt(ctx, []byte("b"), []byte("pre-cutover"), 10, 0)) + require.NoError(t, groups[2].Store.PutAt(ctx, []byte("b"), []byte("post-cutover"), 20, 0)) + + _, err := st.GetGroupAt(ctx, 1, []byte("b"), 25) + require.ErrorIs(t, err, ErrExplicitGroupRouteOwnerMismatch) + + _, err = st.ScanGroupAt(ctx, 1, []byte("a"), []byte("z"), 10, 25) + require.ErrorIs(t, err, ErrExplicitGroupRouteOwnerMismatch) + + _, err = st.ScanAtWithReadFence(ctx, []byte("a"), []byte("z"), 10, 25, false, 1, 0, []byte("a"), []byte("z")) + require.ErrorIs(t, err, ErrExplicitGroupRouteOwnerMismatch) + + // The group the catalog does name still serves the post-cutover value. + got, err := st.GetGroupAt(ctx, 2, []byte("b"), 25) + require.NoError(t, err) + require.Equal(t, []byte("post-cutover"), got) +} + +// SQS resolves its owning group through the (queue, partition) resolver rather +// than the byte-range catalog, so a catalog route naming another group must not +// reject those reads. +func TestShardStoreExplicitGroupRead_AllowsResolverOwnedKeysOnMismatch(t *testing.T) { + t.Parallel() + + ctx := context.Background() + engine := distribution.NewEngine() + engine.UpdateRoute([]byte(""), nil, 1) + groups := map[uint64]*ShardGroup{ + 1: {Store: store.NewMVCCStore()}, + 42: {Store: store.NewMVCCStore()}, + } + st := NewShardStore(engine, groups) + + key := []byte("!sqs|msg|data|p|orders|partition-2|msg-2") + require.NoError(t, groups[42].Store.PutAt(ctx, key, []byte("payload"), 7, 0)) + + got, err := st.GetGroupAt(ctx, 42, key, 7) + require.NoError(t, err) + require.Equal(t, []byte("payload"), got) + + start := []byte("!sqs|msg|data|p|orders|partition-2|") + kvs, err := st.ScanGroupAt(ctx, 42, start, prefixScanEnd(start), 10, 7) + require.NoError(t, err) + require.Equal(t, []*store.KVPair{{Key: key, Value: []byte("payload")}}, kvs) +} + +// Filesystem placement stats scan the whole chunk keyspace once per filesystem +// group, so most of those groups are not the catalog owner of the range. The +// explicit-group gate must let them through the way it lets SQS through. +func TestShardStoreExplicitGroupScan_AllowsFilesystemChunkKeys(t *testing.T) { + t.Parallel() + + ctx := context.Background() + engine := distribution.NewEngine() + engine.UpdateRoute([]byte(""), nil, 1) + groups := map[uint64]*ShardGroup{ + 1: {Store: store.NewMVCCStore()}, + 2: {Store: store.NewMVCCStore()}, + } + st := NewShardStore(engine, groups) + + chunkKey := fskeys.ChunkKey(3, 7, 0) + require.NoError(t, groups[2].Store.PutAt(ctx, chunkKey, []byte("chunk"), 7, 0)) + + start := fskeys.ChunkAllPrefix() + kvs, err := st.ScanGroupAt(ctx, 2, start, prefixScanEnd(start), 10, 7) + require.NoError(t, err) + require.Equal(t, []*store.KVPair{{Key: chunkKey, Value: []byte("chunk")}}, kvs) +} From c85adaaa2f8aaeff2456455c27ede0019777c6e0 Mon Sep 17 00:00:00 2001 From: bootjp Date: Fri, 28 Aug 2026 04:12:14 +0900 Subject: [PATCH 13/58] migration: cap the promotion batch bounds server-side migrationPromoteOptionsFromProto only filled in bounds the request left unset, so the defaults were defaults and not limits. A caller-supplied max_versions/max_bytes/max_scanned_bytes passed through unchanged into PromoteVersions -- which runs inside FSM apply, synchronously, on every voter. One oversized request could therefore make every replica load, re-encrypt, and commit an unbounded amount of staged data in a single Pebble batch at once, instead of the incremental bounded promotion this API offers. Clamp each bound to a hard ceiling. The caller loses nothing but rounds: PromoteVersionsResult already returns a cursor, so an oversized request now simply takes more chunks. The clamp is a pure function of the command and package constants, so every replica derives the same bounds from the same entry and apply stays deterministic. max_versions is clamped in the int domain the request already decodes into so no widening conversion is introduced. Claude-Session: https://claude.ai/code/session_013rNHooj7NF3giihWVba8QE --- kv/fsm_migration_promote.go | 43 ++++++++++++++++++----- kv/fsm_migration_promote_test.go | 59 ++++++++++++++++++++++++++++++++ 2 files changed, 93 insertions(+), 9 deletions(-) diff --git a/kv/fsm_migration_promote.go b/kv/fsm_migration_promote.go index 383426ae2..83365dde4 100644 --- a/kv/fsm_migration_promote.go +++ b/kv/fsm_migration_promote.go @@ -14,6 +14,21 @@ const ( defaultMigrationPromoteMaxVersions = 1024 defaultMigrationPromoteMaxBytes = 4 << 20 defaultMigrationPromoteMaxScannedBytes = defaultMigrationPromoteMaxBytes * 4 + + // Hard server-side ceilings. The defaults above only apply to a request + // that leaves a bound unset, so without these an operator or migrator could + // ask one apply to load, re-encrypt, and commit an unbounded amount of + // staged data in a single Pebble batch -- synchronously, in every voter's + // apply loop. Clamping keeps the incremental, bounded promotion this API + // promises: the caller simply gets more rounds through the cursor, which + // PromoteVersionsResult already returns. + // + // The clamp is a pure function of the request and these constants, so every + // replica derives the same bounds from the same command and apply stays + // deterministic. + maxMigrationPromoteMaxVersions = 8192 + maxMigrationPromoteMaxBytes = 32 << 20 + maxMigrationPromoteMaxScannedBytes = maxMigrationPromoteMaxBytes * 4 ) var ErrMigrationPromoteApply = errors.New("migration promote: FSM apply failed; halting apply") @@ -57,18 +72,19 @@ func (f *kvFSM) applyMigrationPromote(ctx context.Context, data []byte) any { } func migrationPromoteOptionsFromProto(req *pb.PromoteStagedVersionsRequest, appliedIndex uint64) store.PromoteVersionsOptions { + // Clamped in the int domain the request already decodes into, so no + // widening conversion is introduced here. maxVersions := int(req.GetMaxVersions()) - if maxVersions <= 0 { + switch { + case maxVersions <= 0: maxVersions = defaultMigrationPromoteMaxVersions + case maxVersions > maxMigrationPromoteMaxVersions: + maxVersions = maxMigrationPromoteMaxVersions } - maxBytes := req.GetMaxBytes() - if maxBytes == 0 { - maxBytes = defaultMigrationPromoteMaxBytes - } - maxScannedBytes := req.GetMaxScannedBytes() - if maxScannedBytes == 0 { - maxScannedBytes = defaultMigrationPromoteMaxScannedBytes - } + maxBytes := clampMigrationPromoteBound( + req.GetMaxBytes(), defaultMigrationPromoteMaxBytes, maxMigrationPromoteMaxBytes) + maxScannedBytes := clampMigrationPromoteBound( + req.GetMaxScannedBytes(), defaultMigrationPromoteMaxScannedBytes, maxMigrationPromoteMaxScannedBytes) prefix := distribution.MigrationStagedDataKeyPrefix(req.GetJobId()) return store.PromoteVersionsOptions{ JobID: req.GetJobId(), @@ -83,6 +99,15 @@ func migrationPromoteOptionsFromProto(req *pb.PromoteStagedVersionsRequest, appl } } +// clampMigrationPromoteBound resolves one promotion bound: unset takes the +// default, anything above the hard ceiling is clamped down to it. +func clampMigrationPromoteBound(requested, fallback, ceiling uint64) uint64 { + if requested == 0 { + return fallback + } + return min(requested, ceiling) +} + func isMigrationPromoteOrdinaryApplyError(err error) bool { return errors.Is(err, store.ErrInvalidExportCursor) } diff --git a/kv/fsm_migration_promote_test.go b/kv/fsm_migration_promote_test.go index 192b4af1a..570bf9a80 100644 --- a/kv/fsm_migration_promote_test.go +++ b/kv/fsm_migration_promote_test.go @@ -94,3 +94,62 @@ func TestApplyMigrationPromoteInvalidCursorReturnsOrdinaryError(t *testing.T) { require.ErrorIs(t, err, store.ErrInvalidExportCursor) require.False(t, errors.Is(err, ErrMigrationPromoteApply)) } + +// PromoteStagedVersions runs inside FSM apply, synchronously, on every voter. +// The per-chunk bounds arrive in the Raft command, so a request that names an +// oversized batch makes one apply load, re-encrypt, and commit that much staged +// data in a single Pebble batch on every replica at once. The defaults only fill +// in unset bounds, so they are not limits; the hard ceilings are. +func TestMigrationPromoteOptionsClampOversizedBounds(t *testing.T) { + t.Parallel() + + opts := migrationPromoteOptionsFromProto(&pb.PromoteStagedVersionsRequest{ + JobId: 9, + MaxVersions: 1 << 30, + MaxBytes: 1 << 40, + MaxScannedBytes: 1 << 42, + }, 7) + + require.Equal(t, maxMigrationPromoteMaxVersions, opts.MaxVersions) + require.Equal(t, uint64(maxMigrationPromoteMaxBytes), opts.MaxBytes) + require.Equal(t, uint64(maxMigrationPromoteMaxScannedBytes), opts.MaxScannedBytes) + require.Equal(t, uint64(7), opts.AppliedIndex) +} + +// Unset bounds still take the defaults, and a request under the ceiling is +// passed through unchanged so a caller can still ask for smaller chunks. +func TestMigrationPromoteOptionsKeepDefaultsAndSmallerRequests(t *testing.T) { + t.Parallel() + + defaults := migrationPromoteOptionsFromProto(&pb.PromoteStagedVersionsRequest{JobId: 9}, 0) + require.Equal(t, defaultMigrationPromoteMaxVersions, defaults.MaxVersions) + require.Equal(t, uint64(defaultMigrationPromoteMaxBytes), defaults.MaxBytes) + require.Equal(t, uint64(defaultMigrationPromoteMaxScannedBytes), defaults.MaxScannedBytes) + + smaller := migrationPromoteOptionsFromProto(&pb.PromoteStagedVersionsRequest{ + JobId: 9, + MaxVersions: 16, + MaxBytes: 1024, + MaxScannedBytes: 4096, + }, 0) + require.Equal(t, 16, smaller.MaxVersions) + require.Equal(t, uint64(1024), smaller.MaxBytes) + require.Equal(t, uint64(4096), smaller.MaxScannedBytes) +} + +// The clamp has to be a pure function of the command so every replica derives +// the same bounds from the same entry and apply stays deterministic. +func TestMigrationPromoteOptionsAreDeterministicPerCommand(t *testing.T) { + t.Parallel() + + req := &pb.PromoteStagedVersionsRequest{JobId: 9, MaxVersions: 1 << 30, MaxBytes: 1 << 40} + first := migrationPromoteOptionsFromProto(req, 11) + second := migrationPromoteOptionsFromProto(req, 11) + // PromoteVersionsOptions carries a closure, which never compares equal, so + // the bounds this clamp owns are compared directly. + require.Equal(t, first.MaxVersions, second.MaxVersions) + require.Equal(t, first.MaxBytes, second.MaxBytes) + require.Equal(t, first.MaxScannedBytes, second.MaxScannedBytes) + require.Equal(t, first.StartKey, second.StartKey) + require.Equal(t, first.EndKey, second.EndKey) +} From b087bed84682bf94993df90d948fa724a4328fba Mon Sep 17 00:00:00 2001 From: bootjp Date: Fri, 28 Aug 2026 05:28:37 +0900 Subject: [PATCH 14/58] migration: bound the staged-visibility candidate range scan latestCandidateVersionsAt resolves a scan page's candidate keys with one export spanning from the smallest candidate through the largest, filtered down to the candidate set. Both byte budgets were ^uint64(0), so it scanned and decoded every version in between -- tombstoned keys and dense MVCC history included -- for a page that wants at most stagedVisibilityMaxCandidateWindow exact keys. A migrated route whose visible rows are sparse could make one ordinary scan page consume unbounded I/O on the serving leader. Give the range pass a scan budget and resolve whatever it did not reach with one exact-key probe per remaining candidate, bounded by the candidate count the caller already bounds. The range pass stays the fast path for the dense case, where the candidates sit close together. The error the old code raised when the export stopped early is gone with it: an early stop is now the expected path into the probe fallback rather than a failure. Benchmark on the same fixture, 64 candidates with 64 filler keys of 16 versions each between them (BenchmarkLatestCandidateVersionsAt): dense 19463 -> 20066 ns/op 32680 -> 32708 B/op 469 -> 469 allocs/op sparse 4506534 -> 2552151 ns/op 1581001 -> 531613 B/op 129493 -> 41686 allocs/op The dense case is unchanged within noise. The sparse case is what the budget is for, and the fixture only holds a few megabytes of intervening history -- the old path grew with it without limit, while the new one is capped at the budget plus the probes. Claude-Session: https://claude.ai/code/session_013rNHooj7NF3giihWVba8QE --- kv/shard_store.go | 52 +++++- ...ed_visibility_candidates_benchmark_test.go | 150 ++++++++++++++++++ 2 files changed, 197 insertions(+), 5 deletions(-) create mode 100644 kv/staged_visibility_candidates_benchmark_test.go diff --git a/kv/shard_store.go b/kv/shard_store.go index a5472bc06..136b490da 100644 --- a/kv/shard_store.go +++ b/kv/shard_store.go @@ -2591,6 +2591,22 @@ func nextStagedVisibilityCandidateWindow(window int) int { return next } +// stagedVisibilityCandidateScanBudget bounds the range export that resolves a +// staged-visibility page's candidate keys in one pass. +// +// The export spans from the smallest candidate through the largest and filters +// for the candidate set, so with no budget it scans and decodes every version +// in between -- tombstoned keys and dense MVCC history included -- even though +// at most stagedVisibilityMaxCandidateWindow exact keys are wanted. A route +// whose visible rows are sparse could therefore make an ordinary scan page +// consume unbounded I/O on the serving leader. +// +// With the budget the range pass stays the fast path for the dense case, where +// the candidates sit close together, and a page that would have run away instead +// falls back to one exact-key probe per unresolved candidate -- bounded by the +// candidate count, which the caller already bounds. +const stagedVisibilityCandidateScanBudget = 1 << 20 + func scanVisibleCandidates(ctx context.Context, st store.MVCCStore, start, end []byte, limit int, ts uint64, reverse bool) ([]*store.KVPair, error) { if limit <= 0 { return []*store.KVPair{}, nil @@ -2663,7 +2679,7 @@ func latestCandidateVersionsAt(ctx context.Context, st store.MVCCStore, keys [][ MaxCommitTSInclusive: ts, MaxVersions: len(sortedKeys), MaxBytes: ^uint64(0), - MaxScannedBytes: ^uint64(0), + MaxScannedBytes: stagedVisibilityCandidateScanBudget, AcceptKey: func(key []byte) bool { _, ok := candidates[string(key)] return ok @@ -2680,13 +2696,39 @@ func latestCandidateVersionsAt(ctx context.Context, st store.MVCCStore, keys [][ if err != nil { return nil, errors.WithStack(err) } - if !result.Done && len(result.Versions) < len(sortedKeys) { - return nil, errors.New("staged visibility range export stopped before all candidates were examined") - } - out := make(map[string]store.MVCCVersion, len(result.Versions)) + out := make(map[string]store.MVCCVersion, len(sortedKeys)) for _, version := range result.Versions { out[string(version.Key)] = version } + if result.Done { + // The export walked the whole enclosing range, so a candidate missing + // from the result has no visible version. + return out, nil + } + return probeRemainingCandidateVersionsAt(ctx, st, sortedKeys, out, ts) +} + +// probeRemainingCandidateVersionsAt resolves the candidates the bounded range +// export did not reach, one exact-key probe each. +func probeRemainingCandidateVersionsAt( + ctx context.Context, + st store.MVCCStore, + sortedKeys [][]byte, + out map[string]store.MVCCVersion, + ts uint64, +) (map[string]store.MVCCVersion, error) { + for _, key := range sortedKeys { + if _, ok := out[string(key)]; ok { + continue + } + version, found, err := latestMVCCVersionAt(ctx, st, key, ts) + if err != nil { + return nil, err + } + if found { + out[string(key)] = version + } + } return out, nil } diff --git a/kv/staged_visibility_candidates_benchmark_test.go b/kv/staged_visibility_candidates_benchmark_test.go new file mode 100644 index 000000000..aaa479ee7 --- /dev/null +++ b/kv/staged_visibility_candidates_benchmark_test.go @@ -0,0 +1,150 @@ +package kv + +import ( + "context" + "fmt" + "testing" + + "github.com/bootjp/elastickv/store" + "github.com/stretchr/testify/require" +) + +// seedSparseCandidateRange builds the shape latestCandidateVersionsAt is worst +// at: a handful of candidate keys spread across a range that is otherwise dense +// with versions nobody asked for. +func seedSparseCandidateRange(tb testing.TB, st store.MVCCStore, candidates, fillerPerGap, versionsPerFiller int) [][]byte { + tb.Helper() + + ctx := context.Background() + keys := make([][]byte, 0, candidates) + commitTS := uint64(1) + for c := range candidates { + key := []byte(fmt.Sprintf("k%06d", c*(fillerPerGap+1))) + require.NoError(tb, st.PutAt(ctx, key, []byte("candidate"), commitTS, 0)) + commitTS++ + keys = append(keys, key) + for f := 1; f <= fillerPerGap; f++ { + filler := []byte(fmt.Sprintf("k%06d", c*(fillerPerGap+1)+f)) + for v := range versionsPerFiller { + _ = v + require.NoError(tb, st.PutAt(ctx, filler, []byte("filler-value-padding"), commitTS, 0)) + commitTS++ + } + } + } + return keys +} + +// BenchmarkLatestCandidateVersionsAt measures the candidate resolution a staged +// visibility scan page performs. The sparse case is the one the scan budget +// exists for: the enclosing range holds far more versions than the candidate +// set, and an unbounded export decodes all of them. +func BenchmarkLatestCandidateVersionsAt(b *testing.B) { + for _, tc := range []struct { + name string + candidates int + fillerPerGap int + versionsPerFiller int + }{ + {name: "dense", candidates: 64, fillerPerGap: 0, versionsPerFiller: 0}, + {name: "sparse", candidates: 64, fillerPerGap: 64, versionsPerFiller: 16}, + } { + b.Run(tc.name, func(b *testing.B) { + st := store.NewMVCCStore() + b.Cleanup(func() { _ = st.Close() }) + keys := seedSparseCandidateRange(b, st, tc.candidates, tc.fillerPerGap, tc.versionsPerFiller) + ctx := context.Background() + readTS := ^uint64(0) >> 1 + + b.ReportAllocs() + b.ResetTimer() + for range b.N { + got, err := latestCandidateVersionsAt(ctx, st, keys, readTS) + if err != nil { + b.Fatal(err) + } + if len(got) != len(keys) { + b.Fatalf("resolved %d of %d candidates", len(got), len(keys)) + } + } + }) + } +} + +// The bounded export must still resolve every candidate: whatever it does not +// reach is probed by exact key, so the result is the same set either way. +func TestLatestCandidateVersionsAtResolvesSparseCandidates(t *testing.T) { + t.Parallel() + + st := store.NewMVCCStore() + t.Cleanup(func() { _ = st.Close() }) + keys := seedSparseCandidateRange(t, st, 32, 64, 16) + + got, err := latestCandidateVersionsAt(context.Background(), st, keys, ^uint64(0)>>1) + require.NoError(t, err) + require.Len(t, got, len(keys)) + for _, key := range keys { + version, ok := got[string(key)] + require.True(t, ok, "candidate %q must resolve", key) + require.Equal(t, []byte("candidate"), version.Value) + } +} + +// A candidate with no visible version stays absent rather than being invented +// by the probe fallback. +func TestLatestCandidateVersionsAtOmitsAbsentCandidates(t *testing.T) { + t.Parallel() + + ctx := context.Background() + st := store.NewMVCCStore() + t.Cleanup(func() { _ = st.Close() }) + keys := seedSparseCandidateRange(t, st, 4, 64, 8) + missing := []byte("k999999") + keys = append(keys, missing) + + got, err := latestCandidateVersionsAt(ctx, st, keys, ^uint64(0)>>1) + require.NoError(t, err) + require.Len(t, got, len(keys)-1) + _, ok := got[string(missing)] + require.False(t, ok) +} + +// exportRecordingStore captures the options every ExportVersions call carries. +type exportRecordingStore struct { + store.MVCCStore + scannedBudgets []uint64 + exports int +} + +func (s *exportRecordingStore) ExportVersions( + ctx context.Context, + opts store.ExportVersionsOptions, +) (store.ExportVersionsResult, error) { + s.exports++ + s.scannedBudgets = append(s.scannedBudgets, opts.MaxScannedBytes) + return s.MVCCStore.ExportVersions(ctx, opts) +} + +// The candidate export spans from the smallest candidate through the largest +// and filters for the candidate set, so an unbounded scan budget decodes every +// intervening version -- tombstones and dense MVCC history included -- for a +// page that wants at most a bounded number of exact keys. A sparse route could +// make one ordinary scan page consume unbounded I/O on the serving leader. +func TestLatestCandidateVersionsAtBoundsTheRangeScan(t *testing.T) { + t.Parallel() + + base := store.NewMVCCStore() + t.Cleanup(func() { _ = base.Close() }) + keys := seedSparseCandidateRange(t, base, 32, 64, 16) + recording := &exportRecordingStore{MVCCStore: base} + + got, err := latestCandidateVersionsAt(context.Background(), recording, keys, ^uint64(0)>>1) + require.NoError(t, err) + require.Len(t, got, len(keys), "every candidate still resolves") + + require.NotEmpty(t, recording.scannedBudgets) + require.Equal(t, uint64(stagedVisibilityCandidateScanBudget), recording.scannedBudgets[0], + "the range pass must carry a finite scan budget") + require.Greater(t, recording.exports, 1, + "a range pass that hit the budget must fall back to exact-key probes") +} From a7c5967cd1dc780a4e384872ad7384a55ac82271 Mon Sep 17 00:00:00 2001 From: bootjp Date: Fri, 28 Aug 2026 06:18:51 +0900 Subject: [PATCH 15/58] migration: bound candidate probes to the exact key latestMVCCVersionAt scanned [key, prefixScanEnd(key)), which covers every key that has key as a prefix. Its AcceptKey filter keeps only the exact key, but the scan still walks the neighbours and all of their versions first, and the cursor loop continues until Done -- so a scan budget only splits that work into repeated exports rather than bounding it. The probe fallback added in b087bed8 can run once per unresolved candidate, so a page of candidates that are absent from one namespace multiplied it. Bound the probe to the candidate key alone. A byte-string has no value strictly between key and key+0x00, so that end covers exactly the key the probe is asking about. The same bound applies to getAtWithStagedVisibility, which resolves its live and staged versions through the same helper. Claude-Session: https://claude.ai/code/session_013rNHooj7NF3giihWVba8QE --- kv/shard_store.go | 16 +++++++- ...ed_visibility_candidates_benchmark_test.go | 37 +++++++++++++++++++ 2 files changed, 52 insertions(+), 1 deletion(-) diff --git a/kv/shard_store.go b/kv/shard_store.go index 136b490da..2a7ceeec8 100644 --- a/kv/shard_store.go +++ b/kv/shard_store.go @@ -361,10 +361,24 @@ func (s *ShardStore) getAtWithStagedVisibility(ctx context.Context, g *ShardGrou return bytes.Clone(winner.Value), nil } +// exactKeyScanEnd returns the exclusive upper bound that selects only key. +// +// prefixScanEnd(key) selects every key that has key as a prefix, which turns an +// exact-key probe into a prefix scan: probing an absent "a" walks "ab", "az" and +// all of their versions before concluding "a" is not there. A byte-string has no +// value strictly between key and key+0x00, so this bound covers key alone. +func exactKeyScanEnd(key []byte) []byte { + // The immediate successor of key in byte-string order. + const successorByte = byte(0) + out := make([]byte, 0, len(key)+1) + out = append(out, key...) + return append(out, successorByte) +} + func latestMVCCVersionAt(ctx context.Context, st store.MVCCStore, key []byte, ts uint64) (store.MVCCVersion, bool, error) { opts := store.ExportVersionsOptions{ StartKey: key, - EndKey: prefixScanEnd(key), + EndKey: exactKeyScanEnd(key), MaxCommitTSInclusive: ts, MaxVersions: 1, MaxScannedBytes: 0, diff --git a/kv/staged_visibility_candidates_benchmark_test.go b/kv/staged_visibility_candidates_benchmark_test.go index aaa479ee7..6290ce1d7 100644 --- a/kv/staged_visibility_candidates_benchmark_test.go +++ b/kv/staged_visibility_candidates_benchmark_test.go @@ -1,6 +1,7 @@ package kv import ( + "bytes" "context" "fmt" "testing" @@ -113,6 +114,7 @@ func TestLatestCandidateVersionsAtOmitsAbsentCandidates(t *testing.T) { type exportRecordingStore struct { store.MVCCStore scannedBudgets []uint64 + endKeys [][]byte exports int } @@ -122,6 +124,7 @@ func (s *exportRecordingStore) ExportVersions( ) (store.ExportVersionsResult, error) { s.exports++ s.scannedBudgets = append(s.scannedBudgets, opts.MaxScannedBytes) + s.endKeys = append(s.endKeys, bytes.Clone(opts.EndKey)) return s.MVCCStore.ExportVersions(ctx, opts) } @@ -148,3 +151,37 @@ func TestLatestCandidateVersionsAtBoundsTheRangeScan(t *testing.T) { require.Greater(t, recording.exports, 1, "a range pass that hit the budget must fall back to exact-key probes") } + +// The probe fallback must cover exactly the candidate key. prefixScanEnd(key) +// covers every key that has it as a prefix, so probing an absent "a" would walk +// "ab", "az" and all of their versions before concluding "a" is not there -- +// unbounded work that the scan budget only splits into repeated exports, once +// per unresolved candidate. +func TestLatestCandidateVersionsAtProbesExactKeysOnly(t *testing.T) { + t.Parallel() + + base := store.NewMVCCStore() + t.Cleanup(func() { _ = base.Close() }) + keys := seedSparseCandidateRange(t, base, 32, 64, 16) + // A candidate with no version of its own, but with neighbours that share it + // as a prefix: a prefix-bounded probe would scan all of them. + absent := []byte("k000000-absent") + ctx := context.Background() + for i := range 64 { + require.NoError(t, base.PutAt(ctx, []byte(fmt.Sprintf("%s-%03d", absent, i)), []byte("neighbour"), 1, 0)) + } + keys = append(keys, absent) + + recording := &exportRecordingStore{MVCCStore: base} + got, err := latestCandidateVersionsAt(ctx, recording, keys, ^uint64(0)>>1) + require.NoError(t, err) + _, ok := got[string(absent)] + require.False(t, ok, "the absent candidate has no visible version") + + require.Contains(t, recording.endKeys, exactKeyScanEnd(absent), + "the probe must be bounded to the candidate key alone") + for _, end := range recording.endKeys[1:] { + require.NotEqual(t, prefixScanEnd(absent), end, + "no probe may use a prefix bound") + } +} From 0d5212410547ef4038c3b89de5aea93925ca04c3 Mon Sep 17 00:00:00 2001 From: bootjp Date: Fri, 28 Aug 2026 09:41:01 +0900 Subject: [PATCH 16/58] migration: reject user writes to migration control prefixes validateRawMutationForApply refused only transaction-internal keys, so a RawKV Put or Delete aimed at !dist|meta|, !dist|route|, !dist|job|, !dist|jobhist|, !migstage|, !migwrite| or !migfence| was applied as ordinary user data. DEL_PREFIX never reaches that validator at all, so a delete aimed at one of those namespaces was not checked either. Refuse both. The check is a pure function of the key, so every replica reaches the same verdict for the same entry and apply stays deterministic. A DEL_PREFIX is refused for a partial spelling of a control prefix too, since the delete would sweep the namespace up; an empty prefix stays allowed, because that is the whole-keyspace flush the caller asked for and refusing it would break FLUSHDB. !dist|migstage| is deliberately NOT covered, and this is the part of the finding that stays open. ShardedCoordinator rewrites a user key into MigrationStagedDataKey while its route has staged visibility, so legitimate user writes do arrive under that prefix and refusing them would break the migration this branch adds. Telling a rewritten key from a forged one needs request-level provenance the FSM does not have: the route state that would distinguish them is refreshed by a polling watcher, so deciding an apply on it would make apply non-deterministic -- the same hazard that forced 8668bdce to be reverted. Claude-Session: https://claude.ai/code/session_013rNHooj7NF3giihWVba8QE --- distribution/migrator.go | 43 ++++++++++++++++ kv/fsm.go | 13 +++++ kv/fsm_reserved_control_test.go | 88 +++++++++++++++++++++++++++++++++ 3 files changed, 144 insertions(+) create mode 100644 kv/fsm_reserved_control_test.go diff --git a/distribution/migrator.go b/distribution/migrator.go index b842b9365..c2de258e2 100644 --- a/distribution/migrator.go +++ b/distribution/migrator.go @@ -101,6 +101,49 @@ var ( ErrMigrationSourceRouteChanged = errors.New("migration source route does not match split job") ) +// IsReservedControlKey reports whether key lives in a control namespace that no +// user mutation may write. +// +// The staged-data prefix is deliberately excluded: ShardedCoordinator rewrites a +// user key into MigrationStagedDataKey while its route has staged visibility, so +// legitimate user writes do arrive under it. Telling those apart from a forged +// one needs request-level provenance the FSM does not have -- the route state +// that would distinguish them is refreshed by a polling watcher and so cannot +// decide an apply. +func IsReservedControlKey(key []byte) bool { + if bytes.HasPrefix(key, []byte(migrationStagedDataPrefix)) { + return false + } + return reservedControlPrefixIntersects(key, bytes.HasPrefix) +} + +// ReservedControlPrefixIntersects reports whether a DEL_PREFIX over prefix is +// aimed at a control namespace, either because the prefix sits inside one or +// because it is a partial spelling of one. +// +// An empty prefix is exempt. That is the whole-keyspace flush, an operation the +// caller asked for deliberately, and refusing it would break FLUSHDB. Keeping +// control keys out of a flush needs the store's prefix-exclusion argument to +// accept more than the one transaction prefix it takes today, which is a +// separate change. +func ReservedControlPrefixIntersects(prefix []byte) bool { + if len(prefix) == 0 || bytes.HasPrefix(prefix, []byte(migrationStagedDataPrefix)) { + return false + } + return reservedControlPrefixIntersects(prefix, func(a, b []byte) bool { + return bytes.HasPrefix(a, b) || bytes.HasPrefix(b, a) + }) +} + +func reservedControlPrefixIntersects(key []byte, match func(key, reserved []byte) bool) bool { + for _, reserved := range migrationReservedControlPrefixes { + if match(key, reserved) { + return true + } + } + return false +} + var migrationReservedControlPrefixes = [][]byte{ []byte("!dist|"), []byte("!migstage|"), diff --git a/kv/fsm.go b/kv/fsm.go index 55a95216d..cc5210774 100644 --- a/kv/fsm.go +++ b/kv/fsm.go @@ -553,6 +553,13 @@ func (f *kvFSM) validateRawMutationForApply(ctx context.Context, mut *pb.Mutatio if isTxnInternalKey(mut.Key) { return errors.WithStack(ErrInvalidRequest) } + // Nor the migration and catalog control namespaces. Those are written only + // by the typed internal commands (catalog applies, migration import and + // promote), never by a RawKV mutation. The check is a pure function of the + // key, so every replica reaches the same verdict for the same entry. + if distribution.IsReservedControlKey(mut.Key) { + return errors.WithStack(ErrInvalidRequest) + } if _, bypass := writeFenceBypassKeys[string(mut.Key)]; !bypass { if err := f.verifyRouteNotFencedForKey(mut.Key); err != nil { return err @@ -581,6 +588,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 { + // DEL_PREFIX never reaches validateRawMutationsForApply, so the control + // namespaces are gated here as well. A prefix is refused both when it sits + // inside one and when it is broad enough to sweep one up. + if distribution.ReservedControlPrefixIntersects(prefix) { + return errors.WithStack(ErrInvalidRequest) + } if err := f.verifyRouteNotFencedForPrefix(prefix); err != nil { return err } diff --git a/kv/fsm_reserved_control_test.go b/kv/fsm_reserved_control_test.go new file mode 100644 index 000000000..41dcf68f5 --- /dev/null +++ b/kv/fsm_reserved_control_test.go @@ -0,0 +1,88 @@ +package kv + +import ( + "context" + "testing" + + "github.com/bootjp/elastickv/distribution" + pb "github.com/bootjp/elastickv/proto" + "github.com/bootjp/elastickv/store" + "github.com/stretchr/testify/require" +) + +// A RawKV mutation must not reach the migration and catalog control +// namespaces. Those are written only by the typed internal commands, and a user +// write that lands in one is later promoted as ordinary data. +func TestValidateRawMutationRejectsReservedControlKeys(t *testing.T) { + t.Parallel() + + ctx := context.Background() + st := store.NewMVCCStore() + t.Cleanup(func() { _ = st.Close() }) + f, ok := NewKvFSMWithHLC(st, NewHLC()).(*kvFSM) + require.True(t, ok) + + for _, key := range [][]byte{ + []byte("!dist|meta|version"), + []byte("!dist|route|0001"), + []byte("!dist|job|7"), + []byte("!dist|jobhist|7"), + []byte("!migstage|7|victim"), + []byte("!migwrite|7"), + []byte("!migfence|7"), + } { + err := f.validateRawMutationForApply(ctx, &pb.Mutation{Op: pb.Op_PUT, Key: key, Value: []byte("v")}, nil, 10) + require.ErrorIs(t, err, ErrInvalidRequest, "key %q must be refused", key) + } + + // Ordinary user keys are unaffected. + require.NoError(t, f.validateRawMutationForApply(ctx, + &pb.Mutation{Op: pb.Op_PUT, Key: []byte("user-key"), Value: []byte("v")}, nil, 10)) +} + +// DEL_PREFIX never reaches validateRawMutationsForApply, so it is gated in +// handleDelPrefix. A partial spelling of a control prefix is refused too, +// because the delete would sweep the namespace up. +func TestHandleDelPrefixRejectsReservedControlPrefixes(t *testing.T) { + t.Parallel() + + ctx := context.Background() + st := store.NewMVCCStore() + t.Cleanup(func() { _ = st.Close() }) + f, ok := NewKvFSMWithHLC(st, NewHLC()).(*kvFSM) + require.True(t, ok) + + for _, prefix := range [][]byte{ + []byte("!dist|"), + []byte("!dist|route|"), + []byte("!dist"), + []byte("!migwrite|"), + []byte("!migfence"), + } { + require.ErrorIs(t, f.handleDelPrefix(ctx, prefix, 10), ErrInvalidRequest, + "prefix %q must be refused", prefix) + } + + // The whole-keyspace flush is a deliberate operation and stays allowed. + require.NoError(t, f.handleDelPrefix(ctx, nil, 11)) + // So does an ordinary user prefix. + require.NoError(t, f.handleDelPrefix(ctx, []byte("user:"), 12)) +} + +// ShardedCoordinator rewrites a user key into MigrationStagedDataKey while its +// route has staged visibility, so legitimate user writes do arrive under the +// staged prefix. Refusing those would break the migration this branch adds. +func TestValidateRawMutationAllowsStagedDataKeys(t *testing.T) { + t.Parallel() + + ctx := context.Background() + st := store.NewMVCCStore() + t.Cleanup(func() { _ = st.Close() }) + f, ok := NewKvFSMWithHLC(st, NewHLC()).(*kvFSM) + require.True(t, ok) + + staged := distribution.MigrationStagedDataKey(7, []byte("user-key")) + require.NoError(t, f.validateRawMutationForApply(ctx, + &pb.Mutation{Op: pb.Op_PUT, Key: staged, Value: []byte("v")}, nil, 10)) + require.NoError(t, f.handleDelPrefix(ctx, distribution.MigrationStagedDataKey(7, []byte("user:")), 11)) +} From 5f07cc08a578662fb35bbd03c23903d5413282f0 Mon Sep 17 00:00:00 2001 From: bootjp Date: Fri, 28 Aug 2026 10:42:12 +0900 Subject: [PATCH 17/58] migration: cap export scan work server-side exportRangeVersionsOptions only filled in bounds the request left unset, so a caller-supplied chunk_bytes / max_scanned_bytes passed through unchanged. For a sparse family or route filter that accepts few rows, one ExportVersions call could then scan and decode an arbitrary portion of the source store before producing its next streamed response, with math.MaxUint64 removing the only work bound there is. Clamp both to hard ceilings, the same shape the promotion path already uses. A clamped caller simply advances through more cursor rounds, which the streamed response already supports. Claude-Session: https://claude.ai/code/session_013rNHooj7NF3giihWVba8QE --- adapter/internal.go | 32 ++++++++++++++++++++++++-------- adapter/internal_test.go | 38 ++++++++++++++++++++++++++++++++++++++ 2 files changed, 62 insertions(+), 8 deletions(-) diff --git a/adapter/internal.go b/adapter/internal.go index 4037f6b8d..f4ab85149 100644 --- a/adapter/internal.go +++ b/adapter/internal.go @@ -104,6 +104,17 @@ const ( defaultMigrationExportChunkBytes = 4 << 20 defaultMigrationExportScanFactor = 4 defaultMigrationExportMaxVersions = 1024 + + // Hard server-side ceilings. The defaults above only fill in a bound the + // request left unset, so without these a caller could ask one + // ExportVersions call to scan and decode an arbitrary portion of the source + // store before producing its next streamed response -- for a sparse family + // or route filter that accepts few rows, math.MaxUint64 removes the only + // work bound there is and the serving leader's I/O and CPU go with it. A + // clamped caller simply advances through more cursor rounds, which the + // response already supports. Same shape as the promotion path's clamp. + maxMigrationExportChunkBytes = 32 << 20 + maxMigrationExportScanBytes = maxMigrationExportChunkBytes * defaultMigrationExportScanFactor ) func (i *Internal) Forward(ctx context.Context, req *pb.ForwardRequest) (*pb.ForwardResponse, error) { @@ -422,14 +433,10 @@ func (i *Internal) proposeMigrationCommand(ctx context.Context, cmd []byte, labe } func (i *Internal) exportRangeVersionsOptions(req *pb.ExportRangeVersionsRequest) store.ExportVersionsOptions { - chunkBytes := uint64(req.GetChunkBytes()) - if chunkBytes == 0 { - chunkBytes = defaultMigrationExportChunkBytes - } - maxScannedBytes := req.GetMaxScannedBytes() - if maxScannedBytes == 0 { - maxScannedBytes = chunkBytes * defaultMigrationExportScanFactor - } + chunkBytes := clampMigrationExportBound( + uint64(req.GetChunkBytes()), defaultMigrationExportChunkBytes, maxMigrationExportChunkBytes) + maxScannedBytes := clampMigrationExportBound( + req.GetMaxScannedBytes(), chunkBytes*defaultMigrationExportScanFactor, maxMigrationExportScanBytes) opts := store.ExportVersionsOptions{ StartKey: req.GetRangeStart(), EndKey: req.GetRangeEnd(), @@ -446,6 +453,15 @@ func (i *Internal) exportRangeVersionsOptions(req *pb.ExportRangeVersionsRequest return opts } +// clampMigrationExportBound resolves one export bound: unset takes the default, +// anything above the hard ceiling is clamped down to it. +func clampMigrationExportBound(requested, fallback, ceiling uint64) uint64 { + if requested == 0 { + return fallback + } + return min(requested, ceiling) +} + func (i *Internal) migrationExportFilter(req *pb.ExportRangeVersionsRequest) func([]byte) bool { bracket := migrationExportBracket(req) if req.GetKeyFamily() == distribution.MigrationFamilyLegacyListMetaDelta { diff --git a/adapter/internal_test.go b/adapter/internal_test.go index ec2b9eb8b..3d3bcb202 100644 --- a/adapter/internal_test.go +++ b/adapter/internal_test.go @@ -430,3 +430,41 @@ func TestStampTxnTimestampsIgnoresMetadataForRouteWriteFloor(t *testing.T) { _, err := i.stampTxnTimestamps(context.Background(), reqs) require.NoError(t, err) } + +// The export defaults only fill in a bound the request left unset, so without a +// ceiling one ExportVersions call can scan and decode an arbitrary portion of +// the source store before producing its next streamed response. For a sparse +// family or route filter that accepts few rows, math.MaxUint64 removes the only +// work bound there is. +func TestExportRangeVersionsOptionsClampOversizedBounds(t *testing.T) { + t.Parallel() + + i := &Internal{} + opts := i.exportRangeVersionsOptions(&pb.ExportRangeVersionsRequest{ + ChunkBytes: ^uint32(0), + MaxScannedBytes: ^uint64(0), + }) + + require.Equal(t, uint64(maxMigrationExportChunkBytes), opts.MaxBytes) + require.Equal(t, uint64(maxMigrationExportScanBytes), opts.MaxScannedBytes) +} + +// Unset bounds still take the defaults, and a request under the ceiling passes +// through so a caller can still ask for smaller chunks. +func TestExportRangeVersionsOptionsKeepDefaultsAndSmallerRequests(t *testing.T) { + t.Parallel() + + i := &Internal{} + defaults := i.exportRangeVersionsOptions(&pb.ExportRangeVersionsRequest{}) + require.Equal(t, uint64(defaultMigrationExportChunkBytes), defaults.MaxBytes) + require.Equal(t, + uint64(defaultMigrationExportChunkBytes*defaultMigrationExportScanFactor), + defaults.MaxScannedBytes) + + smaller := i.exportRangeVersionsOptions(&pb.ExportRangeVersionsRequest{ + ChunkBytes: 4096, + MaxScannedBytes: 8192, + }) + require.Equal(t, uint64(4096), smaller.MaxBytes) + require.Equal(t, uint64(8192), smaller.MaxScannedBytes) +} From 9e0d8b69e27990677d50749d8c5d96e4c6d591f0 Mon Sep 17 00:00:00 2001 From: bootjp Date: Fri, 28 Aug 2026 10:42:12 +0900 Subject: [PATCH 18/58] migration: apply expiration to staged-only values Between cutover and promotion a key can be visible through its staged alias while the live key holds nothing. ShardStore.ExpireAt passed only the live key to the store, and both implementations read that key first and return ErrKeyNotFound when it is absent -- so an expiration issued in that window failed for a value the same route serves happily through GetAt. Resolve the staged/live winner and write the expiration as a live MVCC version, which is where every other post-cutover write goes. The live attempt still comes first, so the ordinary path is unchanged and the staged lookup is only paid when the live key really has nothing, and a key with nothing on either side still reports ErrKeyNotFound. Claude-Session: https://claude.ai/code/session_013rNHooj7NF3giihWVba8QE --- kv/shard_store.go | 35 ++++++++++++++++++++++++++++++ kv/shard_store_test.go | 48 ++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 83 insertions(+) diff --git a/kv/shard_store.go b/kv/shard_store.go index 2a7ceeec8..ebee7d612 100644 --- a/kv/shard_store.go +++ b/kv/shard_store.go @@ -3148,9 +3148,44 @@ func (s *ShardStore) ExpireAt(ctx context.Context, key []byte, expireAt uint64, if err := s.ensureS3BucketAuxiliaryWriteTimestampFloor(key, commitTS); err != nil { return err } + if routeHasStagedVisibility(route) { + return s.expireStagedVisibleAt(ctx, g, route, key, expireAt, commitTS) + } return errors.WithStack(g.Store.ExpireAt(ctx, key, expireAt, commitTS)) } +// expireStagedVisibleAt applies an expiration to a key whose only visible +// version may still be staged. +// +// Between cutover and promotion a key can be visible through its staged alias +// while the live key holds nothing. Both store implementations read the live key +// first and return ErrKeyNotFound when it is absent, so the expiration would +// fail for a value the same route serves happily through GetAt. Resolve the +// staged/live winner and write the expiration as a live MVCC version, which is +// where every other post-cutover write goes. +// +// The live attempt comes first so the ordinary path is unchanged and the staged +// lookup is only paid when the live key really has nothing. +func (s *ShardStore) expireStagedVisibleAt( + ctx context.Context, + g *ShardGroup, + route distribution.Route, + key []byte, + expireAt uint64, + commitTS uint64, +) error { + err := g.Store.ExpireAt(ctx, key, expireAt, commitTS) + if !errors.Is(err, store.ErrKeyNotFound) { + return errors.WithStack(err) + } + value, getErr := s.getAtWithStagedVisibility(ctx, g, route, key, commitTS) + if getErr != nil { + // Nothing visible either way: the original ErrKeyNotFound stands. + return getErr + } + return errors.WithStack(g.Store.PutWithTTLAt(ctx, key, value, commitTS, expireAt)) +} + func (s *ShardStore) LatestCommitTS(ctx context.Context, key []byte) (uint64, bool, error) { return s.LatestCommitTSWithReadFence(ctx, key, 0) } diff --git a/kv/shard_store_test.go b/kv/shard_store_test.go index 72ed2d8e8..e5d5dbdd6 100644 --- a/kv/shard_store_test.go +++ b/kv/shard_store_test.go @@ -3557,3 +3557,51 @@ func TestShardStoreExplicitGroupScan_AllowsFilesystemChunkKeys(t *testing.T) { require.NoError(t, err) require.Equal(t, []*store.KVPair{{Key: chunkKey, Value: []byte("chunk")}}, kvs) } + +// Between cutover and promotion a key can be visible through its staged alias +// while the live key holds nothing. Both store implementations read the live key +// first and return ErrKeyNotFound when it is absent, so an expiration issued in +// that window failed for a value the same route serves happily through GetAt. +func TestExpireAtAppliesToStagedOnlyValues(t *testing.T) { + t.Parallel() + + ctx := context.Background() + st, group := newStagedVisibilityShardStore(t) + + key := []byte("b") + require.NoError(t, group.Store.PutAt(ctx, + distribution.MigrationStagedDataKey(9, key), []byte("staged-b"), 20, 0)) + + // The value is visible even though the live key has nothing. + got, err := st.GetAt(ctx, key, 25) + require.NoError(t, err) + require.Equal(t, []byte("staged-b"), got) + + // The expiry is in the future relative to the commit timestamp, so the value + // survives; the point is that ExpireAt no longer fails outright. + require.NoError(t, st.ExpireAt(ctx, key, 5_000, 300)) + + got, err = st.GetAt(ctx, key, 300) + require.NoError(t, err) + require.Equal(t, []byte("staged-b"), got) + + // The expiration was recorded as a live MVCC version, which is where every + // other post-cutover write goes. + live, err := group.Store.GetAt(ctx, key, 300) + require.NoError(t, err) + require.Equal(t, []byte("staged-b"), live) + + // And it takes effect once the read passes the expiry. + _, err = st.GetAt(ctx, key, 6_000) + require.ErrorIs(t, err, store.ErrKeyNotFound) +} + +// A key with nothing on either side still reports ErrKeyNotFound. +func TestExpireAtStillFailsWhenNothingIsVisible(t *testing.T) { + t.Parallel() + + ctx := context.Background() + st, _ := newStagedVisibilityShardStore(t) + + require.ErrorIs(t, st.ExpireAt(ctx, []byte("absent"), 40, 300), store.ErrKeyNotFound) +} From 48a03d849a05ad259e41b37becbe0b2ed40b4a1d Mon Sep 17 00:00:00 2001 From: bootjp Date: Fri, 28 Aug 2026 10:42:12 +0900 Subject: [PATCH 19/58] migration: preserve snapshot key limits when staging imports MigrationStagedDataKey prepends the staged prefix, the job id, and a separator -- 24 bytes -- so a source key that is itself within store.MaxSnapshotKeySize can exceed it once staged. The 1 MiB key the Pebble snapshot round-trip explicitly supports is exactly such a key. Staging it anyway leaves the target holding a version no snapshot can carry: if Raft snapshots while it is still staged, a new or lagging target voter cannot restore and so cannot recover through snapshot transfer at all. Refuse the import instead, so the failure surfaces on the migrator rather than on a future restore. The refusal is an ordinary apply error, not a halt: the batch is rejected where halting would wedge apply on every voter, and the check is a pure function of the key so replicas agree. MaxSnapshotKeySize is exported for callers that wrap a key in an envelope before storing it, which is what staging does. Claude-Session: https://claude.ai/code/session_013rNHooj7NF3giihWVba8QE --- kv/fsm_migration_import.go | 30 ++++++++++++++- kv/fsm_migration_import_test.go | 67 +++++++++++++++++++++++++++++++++ store/mvcc_store.go | 10 ++++- 3 files changed, 103 insertions(+), 4 deletions(-) diff --git a/kv/fsm_migration_import.go b/kv/fsm_migration_import.go index feacb6b8e..cde823a0f 100644 --- a/kv/fsm_migration_import.go +++ b/kv/fsm_migration_import.go @@ -36,13 +36,19 @@ func (f *kvFSM) applyMigrationImport(ctx context.Context, data []byte) any { if err := proto.Unmarshal(data, req); err != nil { return haltErr(errors.Wrap(errors.Mark(err, ErrMigrationImportApply), "kv/fsm: decode migration import")) } + staged := migrationStoreVersionsFromProto(req.GetJobId(), req.GetVersions()) + if err := validateStagedKeySizes(staged); err != nil { + // Ordinary, not halting: the batch is rejected and the migrator can + // surface it, where halting would wedge apply on every voter. + return errors.Wrap(err, "kv/fsm: apply migration import") + } result, err := f.store.ImportVersionsRaft(ctx, store.ImportVersionsOptions{ JobID: req.GetJobId(), AppliedIndex: f.pendingApplyIdx, BracketID: req.GetBracketId(), BatchSeq: req.GetBatchSeq(), Cursor: req.GetCursor(), - Versions: migrationStoreVersionsFromProto(req.GetJobId(), req.GetVersions()), + Versions: staged, }) if err != nil { if isMigrationImportOrdinaryApplyError(err) { @@ -71,10 +77,30 @@ func (f *kvFSM) applyMigrationImport(ctx context.Context, data []byte) any { // the batch while the failed voter skips the imported versions for good, which // surfaces as missing data after failover or promotion. Those halt instead, // matching applyMigrationPromote. +// validateStagedKeySizes rejects a batch whose staged form would not fit in a +// snapshot. +// +// MigrationStagedDataKey prepends the staged prefix, the job id, and a +// separator, so a source key that is itself within store.MaxSnapshotKeySize can +// exceed it once staged. Storing it anyway leaves the target holding a version +// that no snapshot can carry: a new or lagging voter could never restore, and so +// could never recover through snapshot transfer. Refuse the import instead, so +// the failure surfaces on the migrator rather than on a future restore. +func validateStagedKeySizes(versions []store.MVCCVersion) error { + for _, version := range versions { + if len(version.Key) > store.MaxSnapshotKeySize { + return errors.Wrapf(store.ErrSnapshotKeyTooLarge, + "staged key length %d > %d", len(version.Key), store.MaxSnapshotKeySize) + } + } + return nil +} + func isMigrationImportOrdinaryApplyError(err error) bool { return errors.Is(err, store.ErrImportBatchGap) || errors.Is(err, store.ErrInvalidImportVersion) || - errors.Is(err, store.ErrValueTooLarge) + errors.Is(err, store.ErrValueTooLarge) || + errors.Is(err, store.ErrSnapshotKeyTooLarge) } func (f *kvFSM) migrationHLCFloorForApply(ctx context.Context, req *pb.ImportRangeVersionsRequest, result store.ImportVersionsResult) (uint64, error) { diff --git a/kv/fsm_migration_import_test.go b/kv/fsm_migration_import_test.go index 92a0241b4..5b6f415fa 100644 --- a/kv/fsm_migration_import_test.go +++ b/kv/fsm_migration_import_test.go @@ -1,6 +1,7 @@ package kv import ( + "bytes" "context" "testing" @@ -229,3 +230,69 @@ func TestApplyMigrationImportOrdinaryErrorsDoNotHalt(t *testing.T) { }) } } + +// MigrationStagedDataKey prepends the staged prefix, the job id, and a +// separator, so a source key that is itself within store.MaxSnapshotKeySize can +// exceed it once staged. Storing it anyway would leave the target holding a +// version no snapshot can carry, and a new or lagging voter could never restore. +func TestValidateStagedKeySizesRejectsOversizedStagedForm(t *testing.T) { + t.Parallel() + + envelope := len(distribution.MigrationStagedDataKey(1, nil)) + require.Positive(t, envelope) + + // A source key right at the limit: legal on its own, oversized once staged. + atLimit := bytes.Repeat([]byte("k"), store.MaxSnapshotKeySize) + staged := migrationStoreVersionsFromProto(9, []*pb.MVCCVersion{{Key: atLimit, CommitTs: 10}}) + require.Len(t, staged, 1) + require.Greater(t, len(staged[0].Key), store.MaxSnapshotKeySize) + require.ErrorIs(t, validateStagedKeySizes(staged), store.ErrSnapshotKeyTooLarge) + + // With headroom for the envelope it is accepted. + fits := bytes.Repeat([]byte("k"), store.MaxSnapshotKeySize-envelope) + ok := migrationStoreVersionsFromProto(9, []*pb.MVCCVersion{{Key: fits, CommitTs: 10}}) + require.Len(t, ok, 1) + require.Equal(t, store.MaxSnapshotKeySize, len(ok[0].Key)) + require.NoError(t, validateStagedKeySizes(ok)) +} + +// The rejection must be an ordinary apply error. Halting would wedge apply on +// every voter over a batch the migrator can simply be told about. +func TestOversizedStagedKeyIsAnOrdinaryApplyError(t *testing.T) { + t.Parallel() + + require.True(t, isMigrationImportOrdinaryApplyError( + errors.Wrap(store.ErrSnapshotKeyTooLarge, "staged"))) +} + +// The size check has to be wired into the apply path, not merely available: an +// oversized staged key must be refused before it is written, and refused as an +// ordinary error so apply is not halted on every voter. +func TestApplyMigrationImportRejectsOversizedStagedKey(t *testing.T) { + t.Parallel() + + ctx := context.Background() + st := store.NewMVCCStore() + t.Cleanup(func() { _ = st.Close() }) + fsm := &kvFSM{store: st, hlc: NewHLC()} + + // Legal on its own, oversized once the staging envelope is added. + atLimit := bytes.Repeat([]byte("k"), store.MaxSnapshotKeySize) + data, err := proto.Marshal(&pb.ImportRangeVersionsRequest{ + JobId: 9, + BracketId: 1, + BatchSeq: 1, + Versions: []*pb.MVCCVersion{{Key: atLimit, CommitTs: 10, Value: []byte("v")}}, + }) + require.NoError(t, err) + + applied := fsm.applyMigrationImport(ctx, data) + applyErr, ok := applied.(error) + require.True(t, ok, "got %T: %v", applied, applied) + require.ErrorIs(t, applyErr, store.ErrSnapshotKeyTooLarge) + require.NotErrorIs(t, applyErr, ErrMigrationImportApply, "must not halt apply") + + // Nothing was staged. + _, getErr := st.GetAt(ctx, distribution.MigrationStagedDataKey(9, atLimit), 100) + require.ErrorIs(t, getErr, store.ErrKeyNotFound) +} diff --git a/store/mvcc_store.go b/store/mvcc_store.go index 4d2b61dfa..1e36f1d7a 100644 --- a/store/mvcc_store.go +++ b/store/mvcc_store.go @@ -25,8 +25,14 @@ type VersionedValue struct { } const ( - mvccSnapshotVersion = uint32(1) - maxSnapshotKeySize = 1 << 20 // 1 MiB per key + mvccSnapshotVersion = uint32(1) + maxSnapshotKeySize = 1 << 20 // 1 MiB per key + // MaxSnapshotKeySize is the largest key a snapshot can carry. Callers that + // wrap a key in an envelope before storing it -- migration staging is the + // one today -- need it to reserve headroom, because a key that fits on its + // own can stop fitting once wrapped, and the snapshot that would carry it + // is then unrestorable. + MaxSnapshotKeySize = maxSnapshotKeySize maxSnapshotVersionCount = 1 << 20 // 1M versions per key ) From 5029acc1b09baaef0d963309be8515f1a2ced086 Mon Sep 17 00:00:00 2001 From: bootjp Date: Fri, 28 Aug 2026 14:38:45 +0900 Subject: [PATCH 20/58] Fix FSM migration apply fences --- kv/fsm.go | 67 ++++++++++++++++++++++++++-------- kv/fsm_migration_fence_test.go | 59 ++++++++++++++++++++++++++++-- 2 files changed, 107 insertions(+), 19 deletions(-) diff --git a/kv/fsm.go b/kv/fsm.go index 72f96b1df..5f46754ad 100644 --- a/kv/fsm.go +++ b/kv/fsm.go @@ -564,9 +564,9 @@ func (f *kvFSM) validateRawMutationForApply(ctx context.Context, mut *pb.Mutatio if err := f.verifyRouteNotFencedForKey(mut.Key); err != nil { return err } - if err := f.verifyRouteWriteTimestampFloorForKey(mut.Key, commitTS); 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 @@ -600,6 +600,11 @@ func (f *kvFSM) handleDelPrefix(ctx context.Context, prefix []byte, commitTS uin if err := f.verifyRouteWriteTimestampFloorForPrefix(prefix, commitTS); err != nil { return err } + for _, del := range f.stagedVisibilityPrefixDeletesForApply(prefix, txnCommonPrefix) { + if err := f.store.DeletePrefixAtRaftAt(ctx, del.prefix, del.excludePrefix, commitTS, 0); err != nil { + return errors.WithStack(err) + } + } if err := f.store.DeletePrefixAtRaftAt(ctx, prefix, txnCommonPrefix, commitTS, f.pendingApplyIdx); err != nil { return errors.WithStack(err) } @@ -607,6 +612,42 @@ func (f *kvFSM) handleDelPrefix(ctx context.Context, prefix []byte, commitTS uin return nil } +type fsmStagedVisibilityPrefixDelete struct { + prefix []byte + excludePrefix []byte +} + +func (f *kvFSM) stagedVisibilityPrefixDeletesForApply(prefix []byte, excludePrefix []byte) []fsmStagedVisibilityPrefixDelete { + if f == nil || f.routes == nil || f.shardGroupID == 0 { + return nil + } + snap, ok := f.routes.Current() + if !ok { + return nil + } + start, end := routePrefixRange(prefix) + routes := snap.IntersectingRoutes(start, end) + out := make([]fsmStagedVisibilityPrefixDelete, 0, len(routes)) + seen := make(map[string]struct{}, len(routes)) + for _, route := range routes { + if route.GroupID != f.shardGroupID || !routeHasStagedVisibility(route) { + continue + } + stagedPrefix := distribution.MigrationStagedDataKey(route.MigrationJobID, prefix) + var stagedExclude []byte + if excludePrefix != nil { + stagedExclude = distribution.MigrationStagedDataKey(route.MigrationJobID, excludePrefix) + } + dedupeKey := string(stagedPrefix) + "\x00" + string(stagedExclude) + if _, ok := seen[dedupeKey]; ok { + continue + } + seen[dedupeKey] = struct{}{} + out = append(out, fsmStagedVisibilityPrefixDelete{prefix: stagedPrefix, excludePrefix: stagedExclude}) + } + return out +} + func (f *kvFSM) verifyRouteNotFencedForKey(key []byte) error { if f.routes == nil { return nil @@ -665,15 +706,11 @@ func (f *kvFSM) verifyRouteWriteTimestampFloorForKey(key []byte, commitTS uint64 return nil } -func (f *kvFSM) verifyRouteWriteTimestampFloorsForMutations(muts []*pb.Mutation, writeFenceBypassKeys [][]byte, commitTS uint64) error { - bypassKeys := writeFenceBypassKeySet(writeFenceBypassKeys) +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 _, bypass := bypassKeys[string(mut.Key)]; bypass { - continue - } if err := f.verifyRouteWriteTimestampFloorForKey(mut.Key, commitTS); err != nil { return err } @@ -1217,7 +1254,7 @@ func (f *kvFSM) handlePrepareRequest(ctx context.Context, r *pb.Request) error { } startTS := r.Ts - uniq, err := f.uniqueMutationsAboveFloor(muts, r.GetWriteFenceBypassKeys(), startTS) + uniq, err := f.uniqueMutationsAboveFloor(muts, startTS) if err != nil { return err } @@ -1287,7 +1324,7 @@ func (f *kvFSM) handleOnePhaseTxnRequest(ctx context.Context, r *pb.Request, com return nil } - uniq, err := f.uniqueMutationsAboveFloor(muts, r.GetWriteFenceBypassKeys(), commitTS) + uniq, err := f.uniqueMutationsAboveFloor(muts, commitTS) if err != nil { return err } @@ -1311,23 +1348,23 @@ func uniqueTxnMutations(muts []*pb.Mutation) ([]*pb.Mutation, error) { return uniq, nil } -func (f *kvFSM) uniqueMutationsAboveFloor(muts []*pb.Mutation, writeFenceBypassKeys [][]byte, commitTS uint64) ([]*pb.Mutation, error) { +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, writeFenceBypassKeys, commitTS); err != nil { + if err := f.verifyRouteWriteTimestampFloorsForMutations(uniq, commitTS); err != nil { return nil, err } return uniq, nil } -func (f *kvFSM) uniqueTxnMutationsAboveFloor(muts []*pb.Mutation, writeFenceBypassKeys [][]byte, commitTS uint64) ([]*pb.Mutation, error) { +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, writeFenceBypassKeys, commitTS); err != nil { + if err := f.verifyRouteWriteTimestampFloorsForMutations(uniq, commitTS); err != nil { return nil, err } return uniq, nil @@ -1418,7 +1455,7 @@ func (f *kvFSM) handleCommitRequest(ctx context.Context, r *pb.Request) error { if err != nil { return err } - uniq, err := f.uniqueTxnMutationsAboveFloor(muts, r.GetWriteFenceBypassKeys(), commitTS) + 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 1eee26301..4af614932 100644 --- a/kv/fsm_migration_fence_test.go +++ b/kv/fsm_migration_fence_test.go @@ -108,7 +108,7 @@ func TestFSMWriteFenceBypassAllowsMarkedRawPointWrite(t *testing.T) { require.Equal(t, []byte("v"), got) } -func TestFSMWriteFenceBypassAllowsRawWriteBelowBypassedRouteFloor(t *testing.T) { +func TestFSMWriteFenceBypassRejectsRawWriteAtBypassedRouteFloor(t *testing.T) { t.Parallel() fsm := newWriteFloorFSM(t) @@ -116,8 +116,8 @@ func TestFSMWriteFenceBypassAllowsRawWriteBelowBypassedRouteFloor(t *testing.T) err := fsm.handleRawRequest(context.Background(), &pb.Request{ WriteFenceBypassKeys: [][]byte{key}, Mutations: []*pb.Mutation{{Op: pb.Op_PUT, Key: key, Value: []byte("v")}}, - }, 10) - require.NoError(t, err) + }, 100) + require.ErrorIs(t, err, ErrRouteWriteTimestampTooLow) } func TestFSMWriteFenceBypassAllowsPinnedTxnOnNonOwningGroup(t *testing.T) { @@ -126,7 +126,7 @@ func TestFSMWriteFenceBypassAllowsPinnedTxnOnNonOwningGroup(t *testing.T) { engine := distribution.NewEngine() applyComposed1Snapshot(t, engine, 1, []distribution.RouteDescriptor{ {RouteID: 1, Start: []byte(""), End: []byte("m"), GroupID: 1, State: distribution.RouteStateActive}, - {RouteID: 2, Start: []byte("m"), End: nil, GroupID: 2, State: distribution.RouteStateWriteFenced, MinWriteTSExclusive: 100}, + {RouteID: 2, Start: []byte("m"), End: nil, GroupID: 2, State: distribution.RouteStateWriteFenced}, }) fsm := newComposed1FSM(t, engine, 1) key := []byte("z") @@ -144,6 +144,24 @@ func TestFSMWriteFenceBypassAllowsPinnedTxnOnNonOwningGroup(t *testing.T) { require.NoError(t, err) } +func TestFSMWriteFenceBypassRejectsPinnedTxnAtBypassedRouteFloor(t *testing.T) { + t.Parallel() + + fsm := newWriteFloorFSM(t) + key := []byte("!sqs|msg|data|p|partitioned-key") + err := fsm.handleTxnRequest(context.Background(), &pb.Request{ + IsTxn: true, + Phase: pb.Phase_PREPARE, + Ts: 100, + WriteFenceBypassKeys: [][]byte{key}, + Mutations: []*pb.Mutation{ + {Op: pb.Op_PUT, Key: []byte(txnMetaPrefix), Value: EncodeTxnMeta(TxnMeta{PrimaryKey: key, LockTTLms: defaultTxnLockTTLms})}, + {Op: pb.Op_DEL, Key: key}, + }, + }, 100) + require.ErrorIs(t, err, ErrRouteWriteTimestampTooLow) +} + func TestFSMWriteFenceBypassDoesNotAllowDelPrefix(t *testing.T) { t.Parallel() @@ -156,6 +174,39 @@ func TestFSMWriteFenceBypassDoesNotAllowDelPrefix(t *testing.T) { require.ErrorIs(t, err, ErrRouteWriteFenced) } +func TestFSMDelPrefixTombstonesStagedVisibilityRowsDuringApply(t *testing.T) { + t.Parallel() + + ctx := context.Background() + engine := distribution.NewEngine() + applyComposed1Snapshot(t, engine, 1, []distribution.RouteDescriptor{ + { + RouteID: 1, + Start: []byte("a"), + End: []byte("z"), + GroupID: 1, + State: distribution.RouteStateActive, + StagedVisibilityActive: true, + MigrationJobID: 9, + }, + }) + fsm := newComposed1FSM(t, engine, 1) + dropKey := []byte("b/drop") + outsideKey := []byte("c/outside") + stagedDrop := distribution.MigrationStagedDataKey(9, dropKey) + stagedOutside := distribution.MigrationStagedDataKey(9, outsideKey) + require.NoError(t, fsm.store.PutAt(ctx, stagedDrop, []byte("drop"), 20, 0)) + require.NoError(t, fsm.store.PutAt(ctx, stagedOutside, []byte("outside"), 20, 0)) + + require.NoError(t, fsm.handleDelPrefix(ctx, []byte("b/"), 101)) + + _, err := fsm.store.GetAt(ctx, stagedDrop, 150) + require.ErrorIs(t, err, store.ErrKeyNotFound) + got, err := fsm.store.GetAt(ctx, stagedOutside, 150) + require.NoError(t, err) + require.Equal(t, []byte("outside"), got) +} + func TestFSMRejectsCurrentWriteFenceAfterObservedActiveRawPointWrite(t *testing.T) { t.Parallel() From f654d1b4cb1b61a943802091ca363c77d7d2eae7 Mon Sep 17 00:00:00 2001 From: bootjp Date: Fri, 28 Aug 2026 14:57:59 +0900 Subject: [PATCH 21/58] Route promoted S3 bucket metadata --- kv/migrator_filter.go | 18 ++++ kv/shard_router.go | 11 +-- kv/shard_store.go | 115 +++++++++++++--------- kv/shard_store_test.go | 65 ++++++++++++ kv/sharded_coordinator.go | 27 +++-- kv/sharded_coordinator_del_prefix_test.go | 41 +++++++- kv/sharded_coordinator_txn_test.go | 21 ++++ 7 files changed, 230 insertions(+), 68 deletions(-) diff --git a/kv/migrator_filter.go b/kv/migrator_filter.go index acbf1a796..3aa2a7a93 100644 --- a/kv/migrator_filter.go +++ b/kv/migrator_filter.go @@ -3,6 +3,7 @@ package kv import ( "bytes" + "github.com/bootjp/elastickv/distribution" "github.com/bootjp/elastickv/internal/s3keys" ) @@ -59,6 +60,23 @@ func s3BucketAuxiliaryRouteRange(rawKey []byte) ([]byte, []byte, bool) { return bucketRouteStart, prefixScanEnd(bucketRouteStart), true } +func s3BucketAuxiliaryOwnerRoute(rawKey []byte, routes []distribution.Route) (distribution.Route, bool) { + start, end, ok := s3BucketAuxiliaryRouteRange(rawKey) + if !ok { + return distribution.Route{}, false + } + return s3BucketAuxiliaryOwnerRouteFromRange(start, end, routes) +} + +func s3BucketAuxiliaryOwnerRouteFromRange(start []byte, end []byte, routes []distribution.Route) (distribution.Route, bool) { + for _, route := range routes { + if migrationRouteRangesIntersect(route.Start, route.End, start, end) { + return route, true + } + } + return distribution.Route{}, false +} + func keyInMigrationRouteRange(key, routeStart, routeEnd []byte) bool { if key == nil { return false diff --git a/kv/shard_router.go b/kv/shard_router.go index 4016fa452..c8b97a867 100644 --- a/kv/shard_router.go +++ b/kv/shard_router.go @@ -134,7 +134,7 @@ func (s *ShardRouter) ResolveGroup(rawKey []byte) (uint64, bool) { return 0, false } } - if route, ok := s.stagedVisibilityRouteForS3BucketAuxiliaryKey(rawKey); ok { + if route, ok := s.s3BucketAuxiliaryOwnerRouteForKey(rawKey); ok { return route.GroupID, true } // Engine routes against the user-key view of the byte-range @@ -148,7 +148,7 @@ func (s *ShardRouter) ResolveGroup(rawKey []byte) (uint64, bool) { return route.GroupID, true } -func (s *ShardRouter) stagedVisibilityRouteForS3BucketAuxiliaryKey(rawKey []byte) (distribution.Route, bool) { +func (s *ShardRouter) s3BucketAuxiliaryOwnerRouteForKey(rawKey []byte) (distribution.Route, bool) { if s == nil || s.engine == nil { return distribution.Route{}, false } @@ -156,12 +156,7 @@ func (s *ShardRouter) stagedVisibilityRouteForS3BucketAuxiliaryKey(rawKey []byte if !ok { return distribution.Route{}, false } - for _, route := range s.engine.GetIntersectingRoutes(start, end) { - if routeHasStagedVisibility(route) { - return route, true - } - } - return distribution.Route{}, false + return s3BucketAuxiliaryOwnerRouteFromRange(start, end, s.engine.GetIntersectingRoutes(start, end)) } // Register associates a raft group ID with its transactional manager and store. diff --git a/kv/shard_store.go b/kv/shard_store.go index 3eb07e4c3..5c22fba65 100644 --- a/kv/shard_store.go +++ b/kv/shard_store.go @@ -104,6 +104,40 @@ func (s *ShardStore) FilesystemGroupIDs() []uint64 { return groupIDs } +// LocalStoreForKey resolves the local store that owns a raw key in this +// process. S3 chunkblob backfill uses this for peer-local auxiliary rows, which +// bypass Raft but must still follow the same route catalog as the public paths. +func (s *ShardStore) LocalStoreForKey(key []byte) (store.MVCCStore, bool) { + _, group, _, ok := s.routeAndGroupForKeyWithVersion(key) + if !ok || group == nil || group.Store == nil { + return nil, false + } + return group.Store, true +} + +// LocalStores returns every local physical store once. It is intentionally not +// filtered by current route ownership because startup backfill must recover +// chunkrefs restored from snapshots or retained after a route move. +func (s *ShardStore) LocalStores() []store.MVCCStore { + if s == nil { + return nil + } + groupIDs := make([]uint64, 0, len(s.groups)) + for groupID := range s.groups { + groupIDs = append(groupIDs, groupID) + } + slices.Sort(groupIDs) + stores := make([]store.MVCCStore, 0, len(groupIDs)) + for _, groupID := range groupIDs { + group := s.groups[groupID] + if group == nil || group.Store == nil { + continue + } + stores = append(stores, group.Store) + } + return stores +} + // ResolveFilesystemHomeSlot finds a home token whose file route belongs to // targetGroup. It derives candidates from current route boundaries and verifies // each candidate against the live catalog before returning it. @@ -362,10 +396,8 @@ func (s *ShardStore) s3BucketAuxiliaryPointReadRoutesWithVersion(key []byte) ([] } catalogRoutes, version := s.engine.GetIntersectingRoutesWithVersion(nil, nil) routes := make([]distribution.Route, 0, pointReadRouteCandidateCapacity) - for _, route := range catalogRoutes { - if routeHasStagedVisibility(route) && migrationRouteRangesIntersect(route.Start, route.End, start, end) { - routes = append(routes, route) - } + if route, ok := s3BucketAuxiliaryOwnerRouteFromRange(start, end, catalogRoutes); ok { + routes = append(routes, route) } normalizedKey := routeKey(key) for _, route := range catalogRoutes { @@ -446,11 +478,16 @@ func (s *ShardStore) routeForExplicitGroupKey(groupID uint64, key []byte) (distr if s == nil || s.engine == nil { return fallback, nil } - if route, ok := s.stagedVisibilityRouteForS3BucketAuxiliaryKey(key); ok { + if route, ok := s.s3BucketAuxiliaryOwnerRouteForKey(key); ok { if route.GroupID == groupID { return route, nil } - return distribution.Route{}, errors.Wrapf(ErrExplicitGroupStagedVisibilityUnresolved, "group_id=%d key=%q", groupID, key) + if routeHasStagedVisibility(route) { + return distribution.Route{}, errors.Wrapf(ErrExplicitGroupStagedVisibilityUnresolved, "group_id=%d key=%q", groupID, key) + } + return distribution.Route{}, errors.Wrapf( + ErrExplicitGroupRouteOwnerMismatch, + "group_id=%d catalog_group_id=%d key=%q", groupID, route.GroupID, key) } if route, ok := s.engine.GetRoute(routeKey(key)); ok { if route.GroupID == groupID { @@ -855,7 +892,7 @@ func (s *ShardStore) scanExplicitGroupRoutesAtWithReadFence(ctx context.Context, scanStart = clampScanStart(start, route.Start) scanEnd = clampScanEnd(end, route.End) } - kvs, err := s.scanRouteAtDirectionWithS3StagedOwnerFilter(ctx, routes, route, scanStart, scanEnd, limit, ts, reverse, true, readRouteVersion, routeStart, routeEnd, dedupeByKey) + kvs, err := s.scanRouteAtDirectionWithS3AuxiliaryOwnerFilter(ctx, routes, route, scanStart, scanEnd, limit, ts, reverse, true, readRouteVersion, routeStart, routeEnd, dedupeByKey) if err != nil { return nil, err } @@ -1167,7 +1204,7 @@ func (s *ShardStore) routesForS3BucketAuxiliaryScan(start []byte, end []byte) ([ } routeStart, routeEnd := s3BucketAuxiliaryScanRouteRange(start, end) for _, route := range catalogRoutes { - if routeHasStagedVisibility(route) && migrationRouteRangesIntersect(route.Start, route.End, routeStart, routeEnd) { + if migrationRouteRangesIntersect(route.Start, route.End, routeStart, routeEnd) { routes = append(routes, route) } } @@ -1553,7 +1590,7 @@ func (s *ShardStore) scanRouteAtWithMigrationOwnerFilters( readRouteVersion, routeStart, routeEnd, true, ) } - return s.scanRouteAtDirectionWithS3StagedOwnerFilter( + return s.scanRouteAtDirectionWithS3AuxiliaryOwnerFilter( ctx, routes, route, start, end, limit, ts, reverse, explicitGroup, readRouteVersion, routeStart, routeEnd, dedupeByKey, ) @@ -2182,7 +2219,7 @@ func (s *ShardStore) scanRouteAtDirectionWithReadFence( return s.scanRouteAtDirectionWithReadFenceOnce(ctx, route, start, end, limit, ts, reverse, explicitGroup, readRouteVersion, routeStart, routeEnd) } -func (s *ShardStore) scanRouteAtDirectionWithS3StagedOwnerFilter( +func (s *ShardStore) scanRouteAtDirectionWithS3AuxiliaryOwnerFilter( ctx context.Context, routes []distribution.Route, route distribution.Route, @@ -2197,7 +2234,7 @@ func (s *ShardStore) scanRouteAtDirectionWithS3StagedOwnerFilter( routeEnd []byte, dedupeByKey bool, ) ([]*store.KVPair, error) { - if !dedupeByKey || routeHasStagedVisibility(route) || !routesContainStagedVisibility(routes) { + if !dedupeByKey { return s.scanRouteAtDirectionWithReadFence(ctx, route, start, end, limit, ts, reverse, explicitGroup, readRouteVersion, routeStart, routeEnd) } out := make([]*store.KVPair, 0, limit) @@ -2208,7 +2245,7 @@ func (s *ShardStore) scanRouteAtDirectionWithS3StagedOwnerFilter( if err != nil { return nil, err } - out = append(out, filterS3AuxiliaryKVsOwnedByStagedRoutes(page, routes)...) + out = append(out, filterS3AuxiliaryKVsOwnedByRoute(page, routes, route)...) if len(out) >= limit { clear(out[limit:]) return out[:limit], nil @@ -2229,29 +2266,31 @@ func (s *ShardStore) scanRouteAtDirectionWithS3StagedOwnerFilter( return out, nil } -func filterS3AuxiliaryKVsOwnedByStagedRoutes(kvs []*store.KVPair, routes []distribution.Route) []*store.KVPair { +func filterS3AuxiliaryKVsOwnedByRoute(kvs []*store.KVPair, routes []distribution.Route, route distribution.Route) []*store.KVPair { out := make([]*store.KVPair, 0, len(kvs)) for _, kvp := range kvs { if kvp == nil { continue } - start, end, auxiliary := s3BucketAuxiliaryRouteRange(kvp.Key) - ownedByStagedRoute := false - if auxiliary { - for _, candidate := range routes { - if routeHasStagedVisibility(candidate) && migrationRouteRangesIntersect(candidate.Start, candidate.End, start, end) { - ownedByStagedRoute = true - break - } - } - } - if !ownedByStagedRoute { - out = append(out, kvp) + owner, auxiliary := s3BucketAuxiliaryOwnerRoute(kvp.Key, routes) + if auxiliary && !routeMatchesS3BucketAuxiliaryOwner(route, owner) { + continue } + out = append(out, kvp) } return out } +func routeMatchesS3BucketAuxiliaryOwner(route distribution.Route, owner distribution.Route) bool { + if route.GroupID != owner.GroupID { + return false + } + if route.RouteID == 0 || owner.RouteID == 0 { + return true + } + return route.RouteID == owner.RouteID +} + func (s *ShardStore) scanRouteAtDirectionWithReadFenceRouteFilter( ctx context.Context, route distribution.Route, @@ -5144,19 +5183,10 @@ func (s *ShardStore) routeAndGroupForKeyWithVersion(key []byte) (distribution.Ro return distribution.Route{}, nil, 0, false } if start, end, auxiliary := s3BucketAuxiliaryRouteRange(key); auxiliary { - routes, version := s.engine.GetIntersectingRoutesWithVersion(nil, nil) - for _, route := range routes { - if routeHasStagedVisibility(route) && migrationRouteRangesIntersect(route.Start, route.End, start, end) { - g, ok := s.groups[route.GroupID] - return route, g, version, ok - } - } - normalizedKey := routeKey(key) - for _, route := range routes { - if routeContainsKey(route, normalizedKey) { - g, ok := s.groups[route.GroupID] - return route, g, version, ok - } + routes, version := s.engine.GetIntersectingRoutesWithVersion(start, end) + if route, ok := s3BucketAuxiliaryOwnerRouteFromRange(start, end, routes); ok { + g, groupOK := s.groups[route.GroupID] + return route, g, version, groupOK } return distribution.Route{}, nil, version, false } @@ -5168,7 +5198,7 @@ func (s *ShardStore) routeAndGroupForKeyWithVersion(key []byte) (distribution.Ro return route, g, version, ok } -func (s *ShardStore) stagedVisibilityRouteForS3BucketAuxiliaryKey(key []byte) (distribution.Route, bool) { +func (s *ShardStore) s3BucketAuxiliaryOwnerRouteForKey(key []byte) (distribution.Route, bool) { if s == nil || s.engine == nil { return distribution.Route{}, false } @@ -5176,12 +5206,7 @@ func (s *ShardStore) stagedVisibilityRouteForS3BucketAuxiliaryKey(key []byte) (d if !ok { return distribution.Route{}, false } - for _, route := range s.engine.GetIntersectingRoutes(start, end) { - if routeHasStagedVisibility(route) { - return route, true - } - } - return distribution.Route{}, false + return s3BucketAuxiliaryOwnerRouteFromRange(start, end, s.engine.GetIntersectingRoutes(start, end)) } func (s *ShardStore) proxyRawGet(ctx context.Context, g *ShardGroup, key []byte, ts uint64, groupID uint64, readRouteVersion uint64) ([]byte, error) { diff --git a/kv/shard_store_test.go b/kv/shard_store_test.go index 54f6eeb89..9d3926b65 100644 --- a/kv/shard_store_test.go +++ b/kv/shard_store_test.go @@ -159,6 +159,30 @@ func TestShardStoreGetAt_MergesStagedVisibilityForS3BucketAuxiliary(t *testing.T } } +func TestShardStoreGetAt_RoutesS3BucketAuxiliaryToPromotedOwner(t *testing.T) { + t.Parallel() + + ctx := context.Background() + const bucket = "bucket-a" + engine := distribution.NewEngine() + require.NoError(t, engine.ApplySnapshot(distribution.CatalogSnapshot{ + Version: 1, + Routes: s3BucketAuxiliaryPromotedRoutes(), + })) + groups := map[uint64]*ShardGroup{ + 1: {Store: store.NewMVCCStore()}, + 2: {Store: store.NewMVCCStore()}, + } + st := NewShardStore(engine, groups) + key := s3keys.BucketMetaKey(bucket) + require.NoError(t, groups[1].Store.PutAt(ctx, key, []byte("stale-source"), 10, 0)) + require.NoError(t, groups[2].Store.PutAt(ctx, key, []byte("promoted-target"), 20, 0)) + + got, err := st.GetAt(ctx, key, 30) + require.NoError(t, err) + require.Equal(t, []byte("promoted-target"), got) +} + func TestShardStoreS3BucketAuxiliaryScanFiltersStagedRoutesToBucketRange(t *testing.T) { t.Parallel() @@ -216,6 +240,47 @@ func TestShardStoreS3BucketAuxiliaryScanFiltersStagedRoutesToBucketRange(t *test }, reverseAll) } +func TestShardStoreS3BucketAuxiliaryScanUsesPromotedOwner(t *testing.T) { + t.Parallel() + + ctx := context.Background() + const ( + migratedBucket = "bucket-a" + otherBucket = "bucket-z" + ) + engine := distribution.NewEngine() + require.NoError(t, engine.ApplySnapshot(distribution.CatalogSnapshot{ + Version: 1, + Routes: s3BucketAuxiliaryPromotedRoutes(), + })) + groups := map[uint64]*ShardGroup{ + 1: {Store: store.NewMVCCStore()}, + 2: {Store: store.NewMVCCStore()}, + } + st := NewShardStore(engine, groups) + migratedKey := s3keys.BucketMetaKey(migratedBucket) + otherKey := s3keys.BucketMetaKey(otherBucket) + require.NoError(t, groups[1].Store.PutAt(ctx, migratedKey, []byte("stale-source"), 10, 0)) + require.NoError(t, groups[2].Store.PutAt(ctx, migratedKey, []byte("promoted-target"), 20, 0)) + require.NoError(t, groups[1].Store.PutAt(ctx, otherKey, []byte("raw-owner"), 15, 0)) + + start := []byte(s3keys.BucketMetaPrefix) + end := prefixScanEnd(start) + kvs, err := st.ScanAt(ctx, start, end, 10, 30) + require.NoError(t, err) + require.Equal(t, []*store.KVPair{ + {Key: migratedKey, Value: []byte("promoted-target")}, + {Key: otherKey, Value: []byte("raw-owner")}, + }, kvs) + + reverse, err := st.ReverseScanAt(ctx, start, end, 10, 30) + require.NoError(t, err) + require.Equal(t, []*store.KVPair{ + {Key: otherKey, Value: []byte("raw-owner")}, + {Key: migratedKey, Value: []byte("promoted-target")}, + }, reverse) +} + func TestShardStoreRouteBoundedS3BucketAuxiliaryScanKeepsStagedRows(t *testing.T) { t.Parallel() diff --git a/kv/sharded_coordinator.go b/kv/sharded_coordinator.go index b45933115..362bea1f7 100644 --- a/kv/sharded_coordinator.go +++ b/kv/sharded_coordinator.go @@ -2408,7 +2408,7 @@ func (c *ShardedCoordinator) engineGroupIDForKey(key []byte) uint64 { } func (c *ShardedCoordinator) resolveGroupAndRouteForKey(key []byte) (uint64, uint64, bool) { - if route, ok := c.stagedVisibilityRouteForS3BucketAuxiliaryKey(key); ok { + if route, ok := c.s3BucketAuxiliaryOwnerRouteForKey(key); ok { return route.GroupID, route.RouteID, true } gid, ok := c.router.ResolveGroup(key) @@ -2422,7 +2422,7 @@ func (c *ShardedCoordinator) resolveGroupAndRouteForKey(key []byte) (uint64, uin return gid, routeID, true } -func (c *ShardedCoordinator) stagedVisibilityRouteForS3BucketAuxiliaryKey(key []byte) (distribution.Route, bool) { +func (c *ShardedCoordinator) s3BucketAuxiliaryOwnerRouteForKey(key []byte) (distribution.Route, bool) { if c == nil || c.engine == nil { return distribution.Route{}, false } @@ -2430,12 +2430,7 @@ func (c *ShardedCoordinator) stagedVisibilityRouteForS3BucketAuxiliaryKey(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 + return s3BucketAuxiliaryOwnerRouteFromRange(start, end, c.engine.GetIntersectingRoutes(start, end)) } // EngineGroupIDForKey reports the Raft group ID that owns key, or 0 when @@ -2519,14 +2514,18 @@ func (c *ShardedCoordinator) stagedVisibilityReadKeyAlias(gid uint64, key []byte 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 + if route, ok := c.s3BucketAuxiliaryOwnerRouteForKey(key); ok { + return stagedVisibilityReadKeyAliasForRoute(gid, key, route) } route, ok := c.engine.GetRoute(routeKey(key)) - if !ok || route.GroupID != gid || !routeHasStagedVisibility(route) { + if !ok { + return nil, false + } + return stagedVisibilityReadKeyAliasForRoute(gid, key, route) +} + +func stagedVisibilityReadKeyAliasForRoute(gid uint64, key []byte, route distribution.Route) ([]byte, bool) { + if route.GroupID != gid || !routeHasStagedVisibility(route) { return nil, false } return distribution.MigrationStagedDataKey(route.MigrationJobID, key), true diff --git a/kv/sharded_coordinator_del_prefix_test.go b/kv/sharded_coordinator_del_prefix_test.go index b53583056..52a9aa75d 100644 --- a/kv/sharded_coordinator_del_prefix_test.go +++ b/kv/sharded_coordinator_del_prefix_test.go @@ -453,6 +453,13 @@ func s3BucketAuxiliaryStagedRoutes(bucket string, rawGroupID, stagedGroupID uint return routes } +func s3BucketAuxiliaryPromotedRoutes() []distribution.RouteDescriptor { + const bucket = "bucket-a" + routes := s3BucketAuxiliaryFenceRoutes(bucket, 1, 2) + routes[1].State = distribution.RouteStateActive + return routes +} + func TestShardedCoordinatorRoutesS3BucketAuxiliaryWriteToStagedOwner(t *testing.T) { t.Parallel() @@ -471,9 +478,41 @@ func TestShardedCoordinatorRoutesS3BucketAuxiliaryWriteToStagedOwner(t *testing. }, 1, NewHLC(), nil) key := s3keys.BucketMetaKey(bucket) - route, ok := coord.stagedVisibilityRouteForS3BucketAuxiliaryKey(key) + route, ok := coord.s3BucketAuxiliaryOwnerRouteForKey(key) + require.True(t, ok) + require.Equal(t, uint64(2), route.GroupID) + + _, err := coord.Dispatch(context.Background(), &OperationGroup[OP]{ + Elems: []*Elem[OP]{{Op: Put, Key: key, Value: []byte("meta")}}, + }) + require.NoError(t, err) + require.Empty(t, g1Txn.requests) + require.Len(t, g2Txn.requests, 1) + require.Equal(t, key, g2Txn.requests[0].Mutations[0].Key) +} + +func TestShardedCoordinatorRoutesS3BucketAuxiliaryWriteToPromotedOwner(t *testing.T) { + t.Parallel() + + const bucket = "bucket-a" + engine := distribution.NewEngine() + require.NoError(t, engine.ApplySnapshot(distribution.CatalogSnapshot{ + Version: 1, + Routes: s3BucketAuxiliaryPromotedRoutes(), + })) + + g1Txn := &recordingTransactional{} + g2Txn := &recordingTransactional{responses: []*TransactionResponse{{CommitIndex: 22}}} + coord := NewShardedCoordinator(engine, map[uint64]*ShardGroup{ + 1: {Txn: g1Txn}, + 2: {Txn: g2Txn}, + }, 1, NewHLC(), nil) + + key := s3keys.BucketMetaKey(bucket) + route, ok := coord.s3BucketAuxiliaryOwnerRouteForKey(key) require.True(t, ok) require.Equal(t, uint64(2), route.GroupID) + require.False(t, routeHasStagedVisibility(route)) _, err := coord.Dispatch(context.Background(), &OperationGroup[OP]{ Elems: []*Elem[OP]{{Op: Put, Key: key, Value: []byte("meta")}}, diff --git a/kv/sharded_coordinator_txn_test.go b/kv/sharded_coordinator_txn_test.go index b8a76d983..d9ed82e3b 100644 --- a/kv/sharded_coordinator_txn_test.go +++ b/kv/sharded_coordinator_txn_test.go @@ -767,6 +767,27 @@ func TestGroupReadKeysByShardID_RoutesS3BucketAuxiliaryToStagedOwner(t *testing. }, grouped[2]) } +func TestGroupReadKeysByShardID_RoutesS3BucketAuxiliaryToPromotedOwner(t *testing.T) { + t.Parallel() + + const bucket = "bucket-a" + engine := distribution.NewEngine() + require.NoError(t, engine.ApplySnapshot(distribution.CatalogSnapshot{ + Version: 1, + Routes: s3BucketAuxiliaryPromotedRoutes(), + })) + coord := NewShardedCoordinator(engine, map[uint64]*ShardGroup{ + 1: {}, + 2: {}, + }, 1, NewHLC(), nil) + + key := s3keys.BucketMetaKey(bucket) + grouped, err := coord.groupReadKeysByShardID([][]byte{key}) + require.NoError(t, err) + require.Empty(t, grouped[1]) + require.Equal(t, [][]byte{key}, grouped[2]) +} + // --------------------------------------------------------------------------- // validateReadOnlyShards // --------------------------------------------------------------------------- From c044e7d016022f84b6e6cf0a765f6ce95b396d75 Mon Sep 17 00:00:00 2001 From: bootjp Date: Fri, 28 Aug 2026 15:20:34 +0900 Subject: [PATCH 22/58] Batch staged prefix tombstones --- adapter/redis_lua_context.go | 62 ++++++++++++++++------- adapter/redis_txn_test.go | 31 +++++++----- adapter/s3_test.go | 3 ++ kv/fsm.go | 13 +++-- kv/fsm_migration_fence_test.go | 69 +++++++++++++++++++++++++ kv/leader_routed_store.go | 7 +++ kv/shard_store.go | 30 ++++++++--- store/lsm_store.go | 52 ++++++++++++++----- store/lsm_store_applied_index_test.go | 32 ++++++++++++ store/mvcc_store.go | 73 ++++++++++++++++++--------- store/store.go | 11 ++++ 11 files changed, 301 insertions(+), 82 deletions(-) diff --git a/adapter/redis_lua_context.go b/adapter/redis_lua_context.go index ef8553e55..2a3cbac76 100644 --- a/adapter/redis_lua_context.go +++ b/adapter/redis_lua_context.go @@ -51,6 +51,11 @@ type luaScriptContext struct { // server-side probe (still correct, just not cached). negativeType map[string]bool + // rawTypeAtStart caches the TTL-unfiltered type observed at startTS. Commit + // planning needs this raw type to decide whether a logically absent key has + // expired physical rows that must be cleaned up before recreation. + rawTypeAtStart map[string]redisValueType + // keyTypeProbeCount counts how many times the server-side keyTypeAt // helper was invoked during this Eval. Only read by tests via // luaScriptContext methods; ordinary production code never reads it. @@ -283,22 +288,23 @@ func newLuaScriptContext(ctx context.Context, server *RedisServer) (*luaScriptCo } startTS := server.readTS() return &luaScriptContext{ - server: server, - startTS: startTS, - readPin: server.pinReadTS(startTS), - ctx: ctx, - touched: map[string]struct{}{}, - readKeys: map[string][]byte{}, - deleted: map[string]bool{}, - everDeleted: map[string]bool{}, - negativeType: map[string]bool{}, - strings: map[string]*luaStringState{}, - lists: map[string]*luaListState{}, - hashes: map[string]*luaHashState{}, - sets: map[string]*luaSetState{}, - zsets: map[string]*luaZSetState{}, - streams: map[string]*luaStreamState{}, - ttls: map[string]*luaTTLState{}, + server: server, + startTS: startTS, + readPin: server.pinReadTS(startTS), + ctx: ctx, + touched: map[string]struct{}{}, + readKeys: map[string][]byte{}, + deleted: map[string]bool{}, + everDeleted: map[string]bool{}, + negativeType: map[string]bool{}, + rawTypeAtStart: map[string]redisValueType{}, + strings: map[string]*luaStringState{}, + lists: map[string]*luaListState{}, + hashes: map[string]*luaHashState{}, + sets: map[string]*luaSetState{}, + zsets: map[string]*luaZSetState{}, + streams: map[string]*luaStreamState{}, + ttls: map[string]*luaTTLState{}, }, nil } @@ -544,7 +550,11 @@ func (c *luaScriptContext) keyType(key []byte) (redisValueType, error) { } c.keyTypeProbeCount++ - typ, err := c.server.keyTypeAt(c.scriptCtx(), key, c.startTS) + rawTyp, err := c.rawStartTypeForCommitPlan(c.scriptCtx(), key) + if err != nil { + return redisTypeNone, err + } + typ, err := c.server.applyTTLFilter(c.scriptCtx(), key, c.startTS, rawTyp) if err != nil { return redisTypeNone, err } @@ -3755,7 +3765,7 @@ func (c *luaScriptContext) commitPlanForKey(ctx context.Context, key string, com } keyBytes := []byte(key) - rawStartType, err := c.server.rawKeyTypeAt(ctx, keyBytes, c.startTS) + rawStartType, err := c.rawStartTypeForCommitPlan(ctx, keyBytes) if err != nil { return luaKeyPlan{}, err } @@ -3789,6 +3799,22 @@ func (c *luaScriptContext) commitPlanForKey(ctx context.Context, key string, com }, nil } +func (c *luaScriptContext) rawStartTypeForCommitPlan(ctx context.Context, key []byte) (redisValueType, error) { + k := string(key) + if typ, ok := c.rawTypeAtStart[k]; ok { + return typ, nil + } + typ, err := c.server.rawKeyTypeAt(ctx, key, c.startTS) + if err != nil { + return redisTypeNone, err + } + if c.rawTypeAtStart == nil { + c.rawTypeAtStart = map[string]redisValueType{} + } + c.rawTypeAtStart[k] = typ + return typ, nil +} + func luaWideFenceReadKeysForPlan(key []byte, finalType, startType redisValueType, preserveExisting bool) [][]byte { if !preserveExisting || startType == redisTypeNone { return redisTxnWideCollectionFenceKeys(key) diff --git a/adapter/redis_txn_test.go b/adapter/redis_txn_test.go index ae3cdbc33..ff52aa534 100644 --- a/adapter/redis_txn_test.go +++ b/adapter/redis_txn_test.go @@ -1746,20 +1746,21 @@ func (s *luaCleanupScanTrackingStore) ScanAt(ctx context.Context, start []byte, func newLuaCommitPlanTestContext(server *RedisServer, startTS uint64) *luaScriptContext { return &luaScriptContext{ - server: server, - startTS: startTS, - touched: map[string]struct{}{}, - readKeys: map[string][]byte{}, - deleted: map[string]bool{}, - everDeleted: map[string]bool{}, - negativeType: map[string]bool{}, - strings: map[string]*luaStringState{}, - lists: map[string]*luaListState{}, - hashes: map[string]*luaHashState{}, - sets: map[string]*luaSetState{}, - zsets: map[string]*luaZSetState{}, - streams: map[string]*luaStreamState{}, - ttls: map[string]*luaTTLState{}, + server: server, + startTS: startTS, + touched: map[string]struct{}{}, + readKeys: map[string][]byte{}, + deleted: map[string]bool{}, + everDeleted: map[string]bool{}, + negativeType: map[string]bool{}, + rawTypeAtStart: map[string]redisValueType{}, + strings: map[string]*luaStringState{}, + lists: map[string]*luaListState{}, + hashes: map[string]*luaHashState{}, + sets: map[string]*luaSetState{}, + zsets: map[string]*luaZSetState{}, + streams: map[string]*luaStreamState{}, + ttls: map[string]*luaTTLState{}, } } @@ -1773,6 +1774,8 @@ func TestLuaCommitPlanForAbsentRewriteSkipsFullLogicalCleanupScans(t *testing.T) key := "lua:absent-rewrite" scriptCtx := newLuaCommitPlanTestContext(server, 10) + scriptCtx.negativeType[key] = true + scriptCtx.rawTypeAtStart[key] = redisTypeNone scriptCtx.strings[key] = &luaStringState{loaded: true, exists: true, dirty: true, value: []byte("v")} scriptCtx.ttls[key] = &luaTTLState{loaded: true} diff --git a/adapter/s3_test.go b/adapter/s3_test.go index 5ef2897b9..f2b4ee950 100644 --- a/adapter/s3_test.go +++ b/adapter/s3_test.go @@ -799,6 +799,9 @@ func TestS3Server_ShardedStoreRoutesBucketAndObjectData(t *testing.T) { readTS := shardStore.LastCommitTS() var err error _, err = store1.GetAt(ctx, s3keys.BucketMetaKey("bucket-a"), readTS) + require.ErrorIs(t, err, store.ErrKeyNotFound) + + _, err = store2.GetAt(ctx, s3keys.BucketMetaKey("bucket-a"), readTS) require.NoError(t, err) _, err = store1.GetAt(ctx, s3keys.ObjectManifestKey("bucket-a", 1, "dir/file.txt"), readTS) diff --git a/kv/fsm.go b/kv/fsm.go index 5f46754ad..2052991f7 100644 --- a/kv/fsm.go +++ b/kv/fsm.go @@ -600,12 +600,17 @@ func (f *kvFSM) handleDelPrefix(ctx context.Context, prefix []byte, commitTS uin if err := f.verifyRouteWriteTimestampFloorForPrefix(prefix, commitTS); err != nil { return err } + deletes := []store.PrefixDelete{{ + Prefix: prefix, + ExcludePrefix: txnCommonPrefix, + }} for _, del := range f.stagedVisibilityPrefixDeletesForApply(prefix, txnCommonPrefix) { - if err := f.store.DeletePrefixAtRaftAt(ctx, del.prefix, del.excludePrefix, commitTS, 0); err != nil { - return errors.WithStack(err) - } + deletes = append(deletes, store.PrefixDelete{ + Prefix: del.prefix, + ExcludePrefix: del.excludePrefix, + }) } - if err := f.store.DeletePrefixAtRaftAt(ctx, prefix, txnCommonPrefix, commitTS, f.pendingApplyIdx); err != nil { + if err := f.store.DeletePrefixesAtRaftAt(ctx, deletes, commitTS, f.pendingApplyIdx); err != nil { return errors.WithStack(err) } f.notifyApplyObserver(commitTS, pb.Op_DEL_PREFIX, prefix) diff --git a/kv/fsm_migration_fence_test.go b/kv/fsm_migration_fence_test.go index 4af614932..b64a84d7d 100644 --- a/kv/fsm_migration_fence_test.go +++ b/kv/fsm_migration_fence_test.go @@ -207,6 +207,75 @@ func TestFSMDelPrefixTombstonesStagedVisibilityRowsDuringApply(t *testing.T) { require.Equal(t, []byte("outside"), got) } +type recordingPrefixDeleteStore struct { + store.MVCCStore + + batchCalls int + singleCalls int + deletes []store.PrefixDelete + commitTS uint64 + appliedIndex uint64 +} + +func (s *recordingPrefixDeleteStore) DeletePrefixAtRaftAt(ctx context.Context, prefix []byte, excludePrefix []byte, commitTS, appliedIndex uint64) error { + s.singleCalls++ + return s.MVCCStore.DeletePrefixAtRaftAt(ctx, prefix, excludePrefix, commitTS, appliedIndex) +} + +func (s *recordingPrefixDeleteStore) DeletePrefixesAtRaftAt(ctx context.Context, deletes []store.PrefixDelete, commitTS, appliedIndex uint64) error { + s.batchCalls++ + s.deletes = clonePrefixDeletes(deletes) + s.commitTS = commitTS + s.appliedIndex = appliedIndex + return s.MVCCStore.DeletePrefixesAtRaftAt(ctx, deletes, commitTS, appliedIndex) +} + +func clonePrefixDeletes(deletes []store.PrefixDelete) []store.PrefixDelete { + out := make([]store.PrefixDelete, len(deletes)) + for i, del := range deletes { + out[i] = store.PrefixDelete{ + Prefix: append([]byte(nil), del.Prefix...), + ExcludePrefix: append([]byte(nil), del.ExcludePrefix...), + } + } + return out +} + +func TestFSMDelPrefixBatchesStagedAndRawTombstonesDuringApply(t *testing.T) { + t.Parallel() + + ctx := context.Background() + engine := distribution.NewEngine() + applyComposed1Snapshot(t, engine, 1, []distribution.RouteDescriptor{ + { + RouteID: 1, + Start: []byte("a"), + End: []byte("z"), + GroupID: 1, + State: distribution.RouteStateActive, + StagedVisibilityActive: true, + MigrationJobID: 9, + }, + }) + fsm := newComposed1FSM(t, engine, 1) + rec := &recordingPrefixDeleteStore{MVCCStore: fsm.store} + fsm.store = rec + fsm.pendingApplyIdx = 1234 + + require.NoError(t, fsm.handleDelPrefix(ctx, []byte("b/"), 101)) + require.Equal(t, 1, rec.batchCalls) + require.Zero(t, rec.singleCalls) + require.Equal(t, uint64(101), rec.commitTS) + require.Equal(t, uint64(1234), rec.appliedIndex) + require.Equal(t, []store.PrefixDelete{ + {Prefix: []byte("b/"), ExcludePrefix: txnCommonPrefix}, + { + Prefix: distribution.MigrationStagedDataKey(9, []byte("b/")), + ExcludePrefix: distribution.MigrationStagedDataKey(9, txnCommonPrefix), + }, + }, rec.deletes) +} + func TestFSMRejectsCurrentWriteFenceAfterObservedActiveRawPointWrite(t *testing.T) { t.Parallel() diff --git a/kv/leader_routed_store.go b/kv/leader_routed_store.go index 35eb51314..adf49ce81 100644 --- a/kv/leader_routed_store.go +++ b/kv/leader_routed_store.go @@ -594,6 +594,13 @@ func (s *LeaderRoutedStore) DeletePrefixAtRaftAt(ctx context.Context, prefix []b return errors.WithStack(s.local.DeletePrefixAtRaftAt(ctx, prefix, excludePrefix, commitTS, appliedIndex)) } +func (s *LeaderRoutedStore) DeletePrefixesAtRaftAt(ctx context.Context, deletes []store.PrefixDelete, commitTS, appliedIndex uint64) error { + if s == nil || s.local == nil { + return errors.WithStack(store.ErrNotSupported) + } + return errors.WithStack(s.local.DeletePrefixesAtRaftAt(ctx, deletes, commitTS, appliedIndex)) +} + // LastAppliedIndex forwards to the local store when it implements // raftengine.AppliedIndexReader. Defensive: in production today the // kvFSM holds a *pebbleStore directly (not a LeaderRoutedStore — that diff --git a/kv/shard_store.go b/kv/shard_store.go index 5c22fba65..a92f706e2 100644 --- a/kv/shard_store.go +++ b/kv/shard_store.go @@ -4959,13 +4959,32 @@ func (s *ShardStore) DeletePrefixAtRaft(ctx context.Context, prefix []byte, excl // is the receiver only when an aggregate (admin / coordinator) path // is replaying a global FLUSHALL, which is not raft-applied. func (s *ShardStore) DeletePrefixAtRaftAt(ctx context.Context, prefix []byte, excludePrefix []byte, commitTS, appliedIndex uint64) error { - if err := s.ensurePrefixWriteTimestampFloors(prefix, commitTS); err != nil { - return err + return s.DeletePrefixesAtRaftAt(ctx, []store.PrefixDelete{{Prefix: prefix, ExcludePrefix: excludePrefix}}, commitTS, appliedIndex) +} + +func (s *ShardStore) DeletePrefixesAtRaftAt(ctx context.Context, deletes []store.PrefixDelete, commitTS, appliedIndex uint64) error { + if len(deletes) == 0 { + return nil + } + stagedByGroup := make(map[*ShardGroup][]store.PrefixDelete) + for _, del := range deletes { + if err := s.ensurePrefixWriteTimestampFloors(del.Prefix, commitTS); err != nil { + return err + } + for _, staged := range s.stagedVisibilityPrefixDeletes(del.Prefix, del.ExcludePrefix) { + stagedByGroup[staged.group] = append(stagedByGroup[staged.group], store.PrefixDelete{ + Prefix: staged.prefix, + ExcludePrefix: staged.excludePrefix, + }) + } } for _, g := range s.groups { if g == nil || g.Store == nil { continue } + groupDeletes := make([]store.PrefixDelete, 0, len(deletes)+len(stagedByGroup[g])) + groupDeletes = append(groupDeletes, deletes...) + groupDeletes = append(groupDeletes, stagedByGroup[g]...) // Pass appliedIndex through to every group. In the // single-group call-path (the production raft-apply case) // this is correct: appliedIndex IS that group's raft entry @@ -4976,12 +4995,7 @@ func (s *ShardStore) DeletePrefixAtRaftAt(ctx context.Context, prefix []byte, ex // case impossible to reach in production. Tests that // exercise ShardStore.DeletePrefixAtRaftAt across multiple // groups MUST pass appliedIndex=0 to opt out. - if err := g.Store.DeletePrefixAtRaftAt(ctx, prefix, excludePrefix, commitTS, appliedIndex); err != nil { - 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 { + if err := g.Store.DeletePrefixesAtRaftAt(ctx, groupDeletes, commitTS, appliedIndex); err != nil { return errors.WithStack(err) } } diff --git a/store/lsm_store.go b/store/lsm_store.go index 628e536bd..973aa1ff9 100644 --- a/store/lsm_store.go +++ b/store/lsm_store.go @@ -2538,7 +2538,7 @@ func (s *pebbleStore) stageLastCommitTSInBatch(b *pebble.Batch, commitTS uint64, // ELASTICKV_FSM_SYNC_MODE=nosync. Raft-apply callers must use // DeletePrefixAtRaft instead. func (s *pebbleStore) DeletePrefixAt(ctx context.Context, prefix []byte, excludePrefix []byte, commitTS uint64) error { - return s.deletePrefixAtWithOpts(ctx, prefix, excludePrefix, commitTS, s.directApplyWriteOpts(), 0) + return s.deletePrefixesAtWithOpts(ctx, []PrefixDelete{{Prefix: prefix, ExcludePrefix: excludePrefix}}, commitTS, s.directApplyWriteOpts(), 0) } // DeletePrefixAtRaft is the raft-apply variant of DeletePrefixAt. Durability @@ -2549,7 +2549,7 @@ func (s *pebbleStore) DeletePrefixAt(ctx context.Context, prefix []byte, exclude // DeletePrefixAtRaftAt to bundle metaAppliedIndex atomically — see // PR #910 design §2 "why both leaves". func (s *pebbleStore) DeletePrefixAtRaft(ctx context.Context, prefix []byte, excludePrefix []byte, commitTS uint64) error { - return s.deletePrefixAtWithOpts(ctx, prefix, excludePrefix, commitTS, s.raftApplyWriteOpts(), 0) + return s.deletePrefixesAtWithOpts(ctx, []PrefixDelete{{Prefix: prefix, ExcludePrefix: excludePrefix}}, commitTS, s.raftApplyWriteOpts(), 0) } // DeletePrefixAtRaftAt is DeletePrefixAtRaft with the raft entry @@ -2560,24 +2560,25 @@ func (s *pebbleStore) DeletePrefixAtRaft(ctx context.Context, prefix []byte, exc // LastAppliedIndex behind the true applied count for any workload // that uses DEL_PREFIX. PR #910 design §2. func (s *pebbleStore) DeletePrefixAtRaftAt(ctx context.Context, prefix []byte, excludePrefix []byte, commitTS, appliedIndex uint64) error { - return s.deletePrefixAtWithOpts(ctx, prefix, excludePrefix, commitTS, s.raftApplyWriteOpts(), appliedIndex) + return s.DeletePrefixesAtRaftAt(ctx, []PrefixDelete{{Prefix: prefix, ExcludePrefix: excludePrefix}}, commitTS, appliedIndex) } -func (s *pebbleStore) deletePrefixAtWithOpts(_ context.Context, prefix []byte, excludePrefix []byte, commitTS uint64, writeOpts *pebble.WriteOptions, appliedIndex uint64) error { +func (s *pebbleStore) DeletePrefixesAtRaftAt(ctx context.Context, deletes []PrefixDelete, commitTS, appliedIndex uint64) error { + return s.deletePrefixesAtWithOpts(ctx, deletes, commitTS, s.raftApplyWriteOpts(), appliedIndex) +} + +func (s *pebbleStore) deletePrefixesAtWithOpts(_ context.Context, deletes []PrefixDelete, commitTS uint64, writeOpts *pebble.WriteOptions, appliedIndex uint64) error { + if len(deletes) == 0 { + return nil + } + s.dbMu.RLock() defer s.dbMu.RUnlock() s.applyMu.Lock() defer s.applyMu.Unlock() - var lowerBound []byte - if len(prefix) > 0 { - lowerBound = encodeKey(prefix, math.MaxUint64) - } - - iter, err := s.db.NewIter(&pebble.IterOptions{ - LowerBound: lowerBound, - }) + iter, err := s.newDeletePrefixesIterator(deletes) if err != nil { return errors.WithStack(err) } @@ -2586,10 +2587,34 @@ func (s *pebbleStore) deletePrefixAtWithOpts(_ context.Context, prefix []byte, e batch := s.db.NewBatch() defer batch.Close() - if err := s.scanDeletePrefix(iter, batch, prefix, excludePrefix, commitTS); err != nil { + if err := s.stageDeletePrefixes(iter, batch, deletes, commitTS); err != nil { return err } + return s.commitDeletePrefixesBatch(batch, commitTS, writeOpts, appliedIndex) +} +func (s *pebbleStore) newDeletePrefixesIterator(deletes []PrefixDelete) (*pebble.Iterator, error) { + var lowerBound []byte + if len(deletes) == 1 && len(deletes[0].Prefix) > 0 { + lowerBound = encodeKey(deletes[0].Prefix, math.MaxUint64) + } + iter, err := s.db.NewIter(&pebble.IterOptions{LowerBound: lowerBound}) + if err != nil { + return nil, errors.WithStack(err) + } + return iter, nil +} + +func (s *pebbleStore) stageDeletePrefixes(iter *pebble.Iterator, batch *pebble.Batch, deletes []PrefixDelete, commitTS uint64) error { + for _, del := range deletes { + if err := s.scanDeletePrefix(iter, batch, del.Prefix, del.ExcludePrefix, commitTS); err != nil { + return err + } + } + return nil +} + +func (s *pebbleStore) commitDeletePrefixesBatch(batch *pebble.Batch, commitTS uint64, writeOpts *pebble.WriteOptions, appliedIndex uint64) error { // Persist lastCommitTS update atomically with the tombstones. s.mtx.Lock() defer s.mtx.Unlock() @@ -2614,7 +2639,6 @@ func (s *pebbleStore) deletePrefixAtWithOpts(_ context.Context, prefix []byte, e return errors.WithStack(err) } s.updateLastCommitTS(newLastTS) - return nil } diff --git a/store/lsm_store_applied_index_test.go b/store/lsm_store_applied_index_test.go index c872e63a0..bad8be62a 100644 --- a/store/lsm_store_applied_index_test.go +++ b/store/lsm_store_applied_index_test.go @@ -341,6 +341,38 @@ func TestDeletePrefixAtRaftAt_BundlesMetaAppliedIndex(t *testing.T) { require.Equal(t, entryIdx, got) } +func TestDeletePrefixesAtRaftAt_BundlesMetaAppliedIndex(t *testing.T) { + ctx := context.Background() + st := newApplyIndexPebbleStore(t) + ps := pebbleStoreApplied(t, st) + + const seedTS uint64 = 50 + require.NoError(t, ps.ApplyMutations(ctx, []*KVPairMutation{ + {Op: OpTypePut, Key: []byte("p/k1"), Value: []byte("v")}, + {Op: OpTypePut, Key: []byte("q/drop"), Value: []byte("v")}, + {Op: OpTypePut, Key: []byte("q/keep"), Value: []byte("v")}, + }, nil, seedTS, seedTS)) + + const entryIdx uint64 = 100 + require.NoError(t, ps.DeletePrefixesAtRaftAt(ctx, []PrefixDelete{ + {Prefix: []byte("p/")}, + {Prefix: []byte("q/"), ExcludePrefix: []byte("q/keep")}, + }, 200, entryIdx)) + + _, err := ps.GetAt(ctx, []byte("p/k1"), 250) + require.ErrorIs(t, err, ErrKeyNotFound) + _, err = ps.GetAt(ctx, []byte("q/drop"), 250) + require.ErrorIs(t, err, ErrKeyNotFound) + gotVal, err := ps.GetAt(ctx, []byte("q/keep"), 250) + require.NoError(t, err) + require.Equal(t, []byte("v"), gotVal) + + got, present, err := ps.LastAppliedIndex() + require.NoError(t, err) + require.True(t, present, "DeletePrefixesAtRaftAt must persist metaAppliedIndex") + require.Equal(t, entryIdx, got) +} + // TestSetDurableAppliedIndex_UsesPebbleSync exercises the // nosync-mode independence claim — even when ELASTICKV_FSM_SYNC_MODE // is nosync, the checkpoint must use pebble.Sync. We can't directly diff --git a/store/mvcc_store.go b/store/mvcc_store.go index 4bb0340f3..14addf7f8 100644 --- a/store/mvcc_store.go +++ b/store/mvcc_store.go @@ -795,10 +795,22 @@ func (s *mvccStore) DeletePrefixAtRaft(ctx context.Context, prefix []byte, exclu // DeletePrefixAtRaftAt satisfies the MVCCStore interface — see // ApplyMutationsRaftAt for the appliedIndex disposition rationale. func (s *mvccStore) DeletePrefixAtRaftAt(ctx context.Context, prefix []byte, excludePrefix []byte, commitTS, _ uint64) error { - return s.DeletePrefixAt(ctx, prefix, excludePrefix, commitTS) + return s.DeletePrefixesAtRaftAt(ctx, []PrefixDelete{{Prefix: prefix, ExcludePrefix: excludePrefix}}, commitTS, 0) +} + +func (s *mvccStore) DeletePrefixesAtRaftAt(_ context.Context, deletes []PrefixDelete, commitTS, _ uint64) error { + return s.deletePrefixesAt(deletes, commitTS) } func (s *mvccStore) DeletePrefixAt(_ context.Context, prefix []byte, excludePrefix []byte, commitTS uint64) error { + return s.deletePrefixesAt([]PrefixDelete{{Prefix: prefix, ExcludePrefix: excludePrefix}}, commitTS) +} + +func (s *mvccStore) deletePrefixesAt(deletes []PrefixDelete, commitTS uint64) error { + if len(deletes) == 0 { + return nil + } + s.mtx.Lock() defer s.mtx.Unlock() @@ -806,29 +818,9 @@ func (s *mvccStore) DeletePrefixAt(_ context.Context, prefix []byte, excludePref // Collect matching keys first since we cannot modify the tree while iterating. var toDelete [][]byte - it := s.tree.Iterator() - var started bool - if len(prefix) > 0 { - started = seekForwardIteratorStart(s.tree, &it, prefix) - } else { - started = it.First() - } - for ok := started; ok; ok = it.Next() { - k, keyOK := it.Key().([]byte) - if !keyOK { - continue - } - if len(prefix) > 0 && !bytes.HasPrefix(k, prefix) { - break - } - if len(excludePrefix) > 0 && bytes.HasPrefix(k, excludePrefix) { - continue - } - versions, _ := it.Value().([]VersionedValue) - if _, visible := visibleValue(versions, commitTS); !visible { - continue - } - toDelete = append(toDelete, k) + seen := make(map[string]struct{}) + for _, del := range deletes { + toDelete = s.collectDeletePrefixKeysLocked(del, commitTS, seen, toDelete) } for _, k := range toDelete { @@ -1041,6 +1033,39 @@ func writeMVCCSnapshotVersion(w io.Writer, version VersionedValue) error { return nil } +func (s *mvccStore) collectDeletePrefixKeysLocked(del PrefixDelete, commitTS uint64, seen map[string]struct{}, toDelete [][]byte) [][]byte { + it := s.tree.Iterator() + var started bool + if len(del.Prefix) > 0 { + started = seekForwardIteratorStart(s.tree, &it, del.Prefix) + } else { + started = it.First() + } + for ok := started; ok; ok = it.Next() { + k, keyOK := it.Key().([]byte) + if !keyOK { + continue + } + if len(del.Prefix) > 0 && !bytes.HasPrefix(k, del.Prefix) { + break + } + if len(del.ExcludePrefix) > 0 && bytes.HasPrefix(k, del.ExcludePrefix) { + continue + } + versions, _ := it.Value().([]VersionedValue) + if _, visible := visibleValue(versions, commitTS); !visible { + continue + } + dedupeKey := string(k) + if _, ok := seen[dedupeKey]; ok { + continue + } + seen[dedupeKey] = struct{}{} + toDelete = append(toDelete, bytes.Clone(k)) + } + return toDelete +} + func mvccSnapshotTombstoneByte(tombstone bool) byte { if tombstone { return 1 diff --git a/store/store.go b/store/store.go index fb9ba51f9..d0ac8e0d5 100644 --- a/store/store.go +++ b/store/store.go @@ -77,6 +77,12 @@ type KVPair struct { RouteGroupID uint64 } +// PrefixDelete describes one prefix tombstone operation in a batched apply. +type PrefixDelete struct { + Prefix []byte + ExcludePrefix []byte +} + // MVCCVersion is a raw committed MVCC version for range migration. // Unlike scan results, it preserves tombstones and TTL expiry metadata. type MVCCVersion struct { @@ -320,6 +326,11 @@ type MVCCStore interface { // bundles metaAppliedIndex in that batch so DEL_PREFIX entries // also advance the meta key. PR #910 design §2 "why both leaves". DeletePrefixAtRaftAt(ctx context.Context, prefix []byte, excludePrefix []byte, commitTS, appliedIndex uint64) error + // DeletePrefixesAtRaftAt applies several prefix deletes in one + // raft-apply batch. It is used when one logical raft command must + // tombstone multiple physical namespaces without exposing a partial + // apply or advancing metaAppliedIndex separately from any tombstone. + DeletePrefixesAtRaftAt(ctx context.Context, deletes []PrefixDelete, commitTS, appliedIndex uint64) error // LastCommitTS returns the highest commit timestamp applied on this node. LastCommitTS() uint64 // WriteConflictCountsByPrefix returns a snapshot of the MVCC From 11307fa8363c8c8c8609b38b3a7647a40ad86ced Mon Sep 17 00:00:00 2001 From: bootjp Date: Fri, 28 Aug 2026 15:41:09 +0900 Subject: [PATCH 23/58] Preserve migration retry state --- adapter/redis_lists.go | 4 +- adapter/redis_retry_test.go | 28 ++++++++ adapter/redis_stream_cmds.go | 4 +- store/lsm_store.go | 23 ++++--- store/migration_versions_test.go | 63 ++++++++++++++++++ store/mvcc_store.go | 105 ++++++++++++++++++++++-------- store/mvcc_store_snapshot_test.go | 68 +++++++++++++++++++ 7 files changed, 254 insertions(+), 41 deletions(-) diff --git a/adapter/redis_lists.go b/adapter/redis_lists.go index 0eea75f56..362d74e69 100644 --- a/adapter/redis_lists.go +++ b/adapter/redis_lists.go @@ -265,7 +265,7 @@ func (r *RedisServer) dispatchListPushReuse(ctx context.Context, key []byte, pen // (non-retryable errors escape to the client; pending is then // discarded with the goroutine, so the update is wasted and the // stale value would be misleading if some future caller reads it). - if isReusableRedisTxnErr(dispErr) { + if shouldPreserveRedisTxnAttempt(dispErr) { pending.commitTS = commitTS } return 0, false, errors.WithStack(dispErr) @@ -440,7 +440,7 @@ func (r *RedisServer) listPushCoreWithDedup(ctx context.Context, key []byte, val // retryRedisWrite's retry predicate; ambiguous errors that escape // to the client are a separate problem space (cross-request // idempotency cache) and out of scope for this design. - if isReusableRedisTxnErr(dispErr) { + if shouldPreserveRedisTxnAttempt(dispErr) { pending = &reusableListPush{ ops: ops, startTS: startTS, diff --git a/adapter/redis_retry_test.go b/adapter/redis_retry_test.go index 98b7f8ccd..37f9d6416 100644 --- a/adapter/redis_retry_test.go +++ b/adapter/redis_retry_test.go @@ -352,6 +352,34 @@ func TestRedisXAddDedupsLandedWireWriteConflict(t *testing.T) { require.Equal(t, int64(1), meta.Length, "the generated XADD entry must not be appended twice") } +func TestRedisXAddDedupRouteFenceRetryPreservesPriorProbe(t *testing.T) { + t.Parallel() + + st := store.NewMVCCStore() + coord := newDedupTestCoordinator(st, 1, true) + coord.routeFenceAtDispatch = 2 + srv := &RedisServer{ + store: st, + coordinator: coord, + scriptCache: map[string]string{}, + onePhaseTxnDedup: true, + } + conn := &recordingConn{} + + srv.xadd(conn, redcon.Command{Args: [][]byte{ + []byte(cmdXAdd), []byte("retry:stream"), []byte("*"), []byte("field"), []byte("value"), + }}) + + require.Empty(t, conn.err) + require.NotEmpty(t, conn.bulk) + require.Equal(t, 3, coord.dispatches, "attempt 1 landed, route-fenced reuse, then dedup probe retry") + require.Equal(t, 1, coord.probeNoOps, "route-fenced reuse must not replace the prior landed probe") + meta, found, err := srv.loadStreamMetaAt(context.Background(), []byte("retry:stream"), snapshotTS(coord.Clock(), st)) + require.NoError(t, err) + require.True(t, found) + require.Equal(t, int64(1), meta.Length) +} + func TestRedisXAddDedupDisabledDoesNotReplayLandedWireConflict(t *testing.T) { t.Parallel() diff --git a/adapter/redis_stream_cmds.go b/adapter/redis_stream_cmds.go index 1e1c252fb..ae28ac059 100644 --- a/adapter/redis_stream_cmds.go +++ b/adapter/redis_stream_cmds.go @@ -386,7 +386,7 @@ func (r *RedisServer) firstXAddAttempt( // This path owns the exact ID and write set, so it can safely restore a // forwarded conflict to its typed form before entering retryRedisWrite. dispErr = normalizeRetryableRedisTxnErr(dispErr) - if !isReusableRedisTxnErr(dispErr) { + if !shouldPreserveRedisTxnAttempt(dispErr) { return "", nil, cockerrors.WithStack(dispErr) } return "", &reusableXAdd{ @@ -441,7 +441,7 @@ func (r *RedisServer) dispatchXAddReuse( } return "", true, cockerrors.WithStack(dispErr) } - if isReusableRedisTxnErr(dispErr) { + if shouldPreserveRedisTxnAttempt(dispErr) { // A lock response did not apply, but carrying the fresh commitTS keeps // the retry correct if a future retryable transport shape is added. pending.commitTS = commitTS diff --git a/store/lsm_store.go b/store/lsm_store.go index 973aa1ff9..4eacc1d57 100644 --- a/store/lsm_store.go +++ b/store/lsm_store.go @@ -689,6 +689,9 @@ func writeTempDBMetadata(db *pebble.DB, meta streamingMVCCRestoreMetadata) error if err := batch.Set(migrationHLCFloorMetaKeyBytes, encodeMigrationHLCFloors(meta.migrationHLCFloors), nil); err != nil { return errors.WithStack(err) } + if err := batch.Set(migrationPromoteMetaKeyBytes, encodeMigrationPromotionStates(meta.migrationPromotions), nil); err != nil { + return errors.WithStack(err) + } return errors.WithStack(batch.Commit(pebble.Sync)) } @@ -3385,10 +3388,11 @@ func writeNativeSnapshotToTempDir(r io.Reader, tmpDir string, ts uint64) error { // place after the CRC32 checksum is verified, preserving the existing store // on failure. type streamingMVCCRestoreMetadata struct { - lastCommitTS uint64 - minRetainedTS uint64 - migrationAcks map[migrationAckID]migrationImportAck - migrationHLCFloors map[uint64]uint64 + lastCommitTS uint64 + minRetainedTS uint64 + migrationAcks map[migrationAckID]migrationImportAck + migrationHLCFloors map[uint64]uint64 + migrationPromotions map[uint64]PromotionState } func readStreamingMVCCRestoreHeader(r io.Reader) (io.Reader, hash.Hash32, uint32, streamingMVCCRestoreMetadata, error) { @@ -3399,15 +3403,16 @@ func readStreamingMVCCRestoreHeader(r io.Reader) (io.Reader, hash.Hash32, uint32 hash := crc32.NewIEEE() body := io.TeeReader(r, hash) - lastCommitTS, minRetainedTS, migrationAcks, migrationHLCFloors, err := readMVCCSnapshotMetadata(body, version) + lastCommitTS, minRetainedTS, migrationAcks, migrationHLCFloors, migrationPromotions, err := readMVCCSnapshotMetadata(body, version) if err != nil { return nil, nil, 0, streamingMVCCRestoreMetadata{}, err } meta := streamingMVCCRestoreMetadata{ - lastCommitTS: lastCommitTS, - minRetainedTS: minRetainedTS, - migrationAcks: migrationAcks, - migrationHLCFloors: migrationHLCFloors, + lastCommitTS: lastCommitTS, + minRetainedTS: minRetainedTS, + migrationAcks: migrationAcks, + migrationHLCFloors: migrationHLCFloors, + migrationPromotions: migrationPromotions, } return body, hash, expectedChecksum, meta, nil } diff --git a/store/migration_versions_test.go b/store/migration_versions_test.go index 2f93a78d2..f4680a0ff 100644 --- a/store/migration_versions_test.go +++ b/store/migration_versions_test.go @@ -1279,3 +1279,66 @@ func TestPebbleSnapshotPreservesMigrationMetadata(t *testing.T) { _, err = dst.GetAt(ctx, []byte("fresh"), 60) require.ErrorIs(t, err, ErrKeyNotFound) } + +func TestPebbleRestoreStreamingSnapshotPreservesMigrationPromotionState(t *testing.T) { + ctx := context.Background() + src := NewMVCCStore() + promoter, ok := src.(MigrationPromoter) + require.True(t, ok) + + prefix := []byte("stage|") + targetKey := func(staged []byte) ([]byte, bool) { + return bytes.TrimPrefix(staged, prefix), bytes.HasPrefix(staged, prefix) + } + require.NoError(t, src.PutAt(ctx, []byte("stage|a"), []byte("va"), 100, 0)) + require.NoError(t, src.PutAt(ctx, []byte("stage|b"), []byte("vb"), 110, 0)) + + first, err := promoter.PromoteVersions(ctx, PromoteVersionsOptions{ + JobID: 12, + StartKey: prefix, + EndKey: PrefixScanEnd(prefix), + MaxVersions: 1, + TargetKey: targetKey, + }) + require.NoError(t, err) + require.False(t, first.Done) + require.Equal(t, uint64(1), first.TotalPromotedRows) + + snap, err := src.Snapshot() + require.NoError(t, err) + raw := snapshotBytes(t, snap) + require.NoError(t, snap.Close()) + + dstDir, err := os.MkdirTemp("", "migration-streaming-snapshot-dst-*") + require.NoError(t, err) + t.Cleanup(func() { require.NoError(t, os.RemoveAll(dstDir)) }) + dst, err := NewPebbleStore(dstDir) + require.NoError(t, err) + t.Cleanup(func() { require.NoError(t, dst.Close()) }) + require.NoError(t, dst.Restore(bytes.NewReader(raw))) + + dstPromoter, ok := any(dst).(MigrationPromoter) + require.True(t, ok) + stateReader, ok := any(dst).(MigrationPromotionStateReader) + require.True(t, ok) + state, ok, err := stateReader.MigrationPromotionState(ctx, 12) + require.NoError(t, err) + require.True(t, ok) + require.False(t, state.Done) + require.Equal(t, first.NextCursor, state.Cursor) + require.Equal(t, uint64(1), state.PromotedRows) + require.Equal(t, uint64(100), state.MaxPromotedTS) + + restored, err := dstPromoter.PromoteVersions(ctx, PromoteVersionsOptions{ + JobID: 12, + StartKey: prefix, + EndKey: PrefixScanEnd(prefix), + MaxVersions: 10, + TargetKey: targetKey, + }) + require.NoError(t, err) + require.True(t, restored.Done) + require.Equal(t, uint64(1), restored.PromotedRows) + require.Equal(t, uint64(2), restored.TotalPromotedRows) + require.Equal(t, uint64(110), restored.MaxPromotedTS) +} diff --git a/store/mvcc_store.go b/store/mvcc_store.go index 14addf7f8..614b82952 100644 --- a/store/mvcc_store.go +++ b/store/mvcc_store.go @@ -26,7 +26,8 @@ type VersionedValue struct { const ( mvccSnapshotVersionV1 = uint32(1) - mvccSnapshotVersion = uint32(2) + mvccSnapshotVersionV2 = uint32(2) + mvccSnapshotVersion = uint32(3) maxSnapshotKeySize = 1 << 20 // 1 MiB per key // MaxSnapshotKeySize is the largest key a snapshot can carry. Callers that // wrap a key in an envelope before storing it -- migration staging is the @@ -941,16 +942,7 @@ func (s *mvccStore) writeSnapshotBody(f *os.File) (uint32, error) { s.mtx.RLock() defer s.mtx.RUnlock() - if err := binary.Write(w, binary.LittleEndian, s.lastCommitTS); err != nil { - return 0, errors.WithStack(err) - } - if err := binary.Write(w, binary.LittleEndian, s.minRetainedTS); err != nil { - return 0, errors.WithStack(err) - } - if err := writeMVCCSnapshotBytes(w, encodeMigrationImportAcks(s.migrationAcks)); err != nil { - return 0, err - } - if err := writeMVCCSnapshotBytes(w, encodeMigrationHLCFloors(s.migrationHLCFloors)); err != nil { + if err := writeMVCCSnapshotMetadata(w, s.lastCommitTS, s.minRetainedTS, s.migrationAcks, s.migrationHLCFloors, s.migrationPromotions); err != nil { return 0, err } iter := s.tree.Iterator() @@ -973,6 +965,32 @@ func (s *mvccStore) writeSnapshotBody(f *os.File) (uint32, error) { return hash.Sum32(), nil } +func writeMVCCSnapshotMetadata( + w io.Writer, + lastCommitTS uint64, + minRetainedTS uint64, + migrationAcks map[migrationAckID]migrationImportAck, + migrationHLCFloors map[uint64]uint64, + migrationPromotions map[uint64]PromotionState, +) error { + if err := binary.Write(w, binary.LittleEndian, lastCommitTS); err != nil { + return errors.WithStack(err) + } + if err := binary.Write(w, binary.LittleEndian, minRetainedTS); err != nil { + return errors.WithStack(err) + } + if err := writeMVCCSnapshotBytes(w, encodeMigrationImportAcks(migrationAcks)); err != nil { + return err + } + if err := writeMVCCSnapshotBytes(w, encodeMigrationHLCFloors(migrationHLCFloors)); err != nil { + return err + } + if err := writeMVCCSnapshotBytes(w, encodeMigrationPromotionStates(migrationPromotions)); err != nil { + return err + } + return nil +} + func finalizeMVCCSnapshotFile(f *os.File, checksumOffset int64, sum uint32) error { if _, err := f.Seek(checksumOffset, io.SeekStart); err != nil { return errors.WithStack(err) @@ -1079,7 +1097,7 @@ func (s *mvccStore) restoreStreamingSnapshot(r io.Reader) error { return err } - tree, lastCommitTS, minRetainedTS, migrationAcks, migrationHLCFloors, actual, err := restoreStreamingMVCCSnapshotBody(r, version) + tree, lastCommitTS, minRetainedTS, migrationAcks, migrationHLCFloors, migrationPromotions, actual, err := restoreStreamingMVCCSnapshotBody(r, version) if err != nil { return err } @@ -1094,7 +1112,7 @@ func (s *mvccStore) restoreStreamingSnapshot(r io.Reader) error { s.minRetainedTS = minRetainedTS s.migrationAcks = migrationAcks s.migrationHLCFloors = migrationHLCFloors - s.migrationPromotions = make(map[uint64]PromotionState) + s.migrationPromotions = migrationPromotions return nil } @@ -1111,7 +1129,7 @@ func readMVCCSnapshotHeader(r io.Reader) (uint32, uint32, error) { if err := binary.Read(r, binary.LittleEndian, &version); err != nil { return 0, 0, errors.WithStack(err) } - if version != mvccSnapshotVersionV1 && version != mvccSnapshotVersion { + if version != mvccSnapshotVersionV1 && version != mvccSnapshotVersionV2 && version != mvccSnapshotVersion { return 0, 0, errors.WithStack(errors.Newf("unsupported mvcc snapshot version %d", version)) } @@ -1122,55 +1140,86 @@ func readMVCCSnapshotHeader(r io.Reader) (uint32, uint32, error) { return version, expected, nil } -func restoreStreamingMVCCSnapshotBody(r io.Reader, version uint32) (*treemap.Map, uint64, uint64, map[migrationAckID]migrationImportAck, map[uint64]uint64, uint32, error) { +func restoreStreamingMVCCSnapshotBody( + r io.Reader, + version uint32, +) (*treemap.Map, uint64, uint64, map[migrationAckID]migrationImportAck, map[uint64]uint64, map[uint64]PromotionState, uint32, error) { hash := crc32.NewIEEE() body := io.TeeReader(r, hash) - lastCommitTS, minRetainedTS, migrationAcks, migrationHLCFloors, err := readMVCCSnapshotMetadata(body, version) + lastCommitTS, minRetainedTS, migrationAcks, migrationHLCFloors, migrationPromotions, err := readMVCCSnapshotMetadata(body, version) if err != nil { - return nil, 0, 0, nil, nil, 0, err + return nil, 0, 0, nil, nil, nil, 0, err } tree, err := readMVCCSnapshotTree(body) if err != nil { - return nil, 0, 0, nil, nil, 0, err + return nil, 0, 0, nil, nil, nil, 0, err } - return tree, lastCommitTS, minRetainedTS, migrationAcks, migrationHLCFloors, hash.Sum32(), nil + return tree, lastCommitTS, minRetainedTS, migrationAcks, migrationHLCFloors, migrationPromotions, hash.Sum32(), nil } -func readMVCCSnapshotMetadata(r io.Reader, version uint32) (uint64, uint64, map[migrationAckID]migrationImportAck, map[uint64]uint64, error) { +func readMVCCSnapshotMetadata( + r io.Reader, + version uint32, +) (uint64, uint64, map[migrationAckID]migrationImportAck, map[uint64]uint64, map[uint64]PromotionState, error) { var lastCommitTS uint64 if err := binary.Read(r, binary.LittleEndian, &lastCommitTS); err != nil { - return 0, 0, nil, nil, errors.WithStack(err) + return 0, 0, nil, nil, nil, errors.WithStack(err) } var minRetainedTS uint64 if err := binary.Read(r, binary.LittleEndian, &minRetainedTS); err != nil { - return 0, 0, nil, nil, errors.WithStack(err) + return 0, 0, nil, nil, nil, errors.WithStack(err) } if version == mvccSnapshotVersionV1 { - return lastCommitTS, minRetainedTS, make(map[migrationAckID]migrationImportAck), make(map[uint64]uint64), nil + return lastCommitTS, + minRetainedTS, + make(map[migrationAckID]migrationImportAck), + make(map[uint64]uint64), + make(map[uint64]PromotionState), + nil } ackData, err := readMVCCSnapshotBytes(r, "snapshot migration acks") if err != nil { - return 0, 0, nil, nil, err + return 0, 0, nil, nil, nil, err } migrationAcks, ok := decodeMigrationImportAcks(ackData) if !ok { - return 0, 0, nil, nil, errors.New("invalid snapshot migration acks") + return 0, 0, nil, nil, nil, errors.New("invalid snapshot migration acks") } floorData, err := readMVCCSnapshotBytes(r, "snapshot migration hlc floors") if err != nil { - return 0, 0, nil, nil, err + return 0, 0, nil, nil, nil, err } migrationHLCFloors, ok := decodeMigrationHLCFloors(floorData) if !ok { - return 0, 0, nil, nil, errors.New("invalid snapshot migration hlc floors") + return 0, 0, nil, nil, nil, errors.New("invalid snapshot migration hlc floors") } - return lastCommitTS, minRetainedTS, migrationAcks, migrationHLCFloors, nil + migrationPromotions, err := readMVCCSnapshotPromotionMetadata(r, version) + if err != nil { + return 0, 0, nil, nil, nil, err + } + + return lastCommitTS, minRetainedTS, migrationAcks, migrationHLCFloors, migrationPromotions, nil +} + +func readMVCCSnapshotPromotionMetadata(r io.Reader, version uint32) (map[uint64]PromotionState, error) { + if version < mvccSnapshotVersion { + return make(map[uint64]PromotionState), nil + } + promotionData, err := readMVCCSnapshotBytes(r, "snapshot migration promotions") + if err != nil { + return nil, err + } + migrationPromotions, ok := decodeMigrationPromotionStates(promotionData) + if !ok { + return nil, errors.New("invalid snapshot migration promotions") + } + return migrationPromotions, nil } func readMVCCSnapshotBytes(r io.Reader, field string) ([]byte, error) { diff --git a/store/mvcc_store_snapshot_test.go b/store/mvcc_store_snapshot_test.go index ef5018a25..97e8da5ed 100644 --- a/store/mvcc_store_snapshot_test.go +++ b/store/mvcc_store_snapshot_test.go @@ -197,6 +197,74 @@ func TestMVCCStore_SnapshotRestorePreservesMigrationMetadata(t *testing.T) { require.ErrorIs(t, err, ErrKeyNotFound) } +func TestMVCCStore_SnapshotRestorePreservesMigrationPromotionState(t *testing.T) { + t.Parallel() + + ctx := context.Background() + st := newTestMVCCStore(t) + promoter, ok := any(st).(MigrationPromoter) + require.True(t, ok) + stateReader, ok := any(st).(MigrationPromotionStateReader) + require.True(t, ok) + + prefix := []byte("stage|") + targetKey := func(staged []byte) ([]byte, bool) { + return bytes.TrimPrefix(staged, prefix), bytes.HasPrefix(staged, prefix) + } + require.NoError(t, st.PutAt(ctx, []byte("stage|a"), []byte("va"), 100, 0)) + require.NoError(t, st.PutAt(ctx, []byte("stage|b"), []byte("vb"), 110, 0)) + + first, err := promoter.PromoteVersions(ctx, PromoteVersionsOptions{ + JobID: 11, + StartKey: prefix, + EndKey: PrefixScanEnd(prefix), + MaxVersions: 1, + TargetKey: targetKey, + }) + require.NoError(t, err) + require.False(t, first.Done) + require.Equal(t, uint64(1), first.TotalPromotedRows) + require.Equal(t, uint64(100), first.MaxPromotedTS) + + snap, err := st.Snapshot() + require.NoError(t, err) + defer snap.Close() + raw := snapshotBytes(t, snap) + + rest, err := promoter.PromoteVersions(ctx, PromoteVersionsOptions{ + JobID: 11, + StartKey: prefix, + EndKey: PrefixScanEnd(prefix), + MaxVersions: 10, + TargetKey: targetKey, + }) + require.NoError(t, err) + require.True(t, rest.Done) + require.Equal(t, uint64(2), rest.TotalPromotedRows) + + require.NoError(t, st.Restore(bytes.NewReader(raw))) + state, ok, err := stateReader.MigrationPromotionState(ctx, 11) + require.NoError(t, err) + require.True(t, ok) + require.False(t, state.Done) + require.Equal(t, first.NextCursor, state.Cursor) + require.Equal(t, uint64(1), state.PromotedRows) + require.Equal(t, uint64(100), state.MaxPromotedTS) + + restored, err := promoter.PromoteVersions(ctx, PromoteVersionsOptions{ + JobID: 11, + StartKey: prefix, + EndKey: PrefixScanEnd(prefix), + MaxVersions: 10, + TargetKey: targetKey, + }) + require.NoError(t, err) + require.True(t, restored.Done) + require.Equal(t, uint64(1), restored.PromotedRows) + require.Equal(t, uint64(2), restored.TotalPromotedRows) + require.Equal(t, uint64(110), restored.MaxPromotedTS) +} + func TestMVCCStore_ApplyMutations_WriteConflict(t *testing.T) { t.Parallel() From 7b6d12d399ee1edf5d7066d0e3af78836a588241 Mon Sep 17 00:00:00 2001 From: bootjp Date: Fri, 28 Aug 2026 15:48:25 +0900 Subject: [PATCH 24/58] Preserve legacy S3 auxiliary scans --- kv/shard_store.go | 50 +++++++++++++++++++++++++++++++++--- kv/shard_store_test.go | 57 ++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 103 insertions(+), 4 deletions(-) diff --git a/kv/shard_store.go b/kv/shard_store.go index a92f706e2..29c8544e2 100644 --- a/kv/shard_store.go +++ b/kv/shard_store.go @@ -2245,7 +2245,11 @@ func (s *ShardStore) scanRouteAtDirectionWithS3AuxiliaryOwnerFilter( if err != nil { return nil, err } - out = append(out, filterS3AuxiliaryKVsOwnedByRoute(page, routes, route)...) + filtered, err := s.filterS3AuxiliaryKVsOwnedByRoute(ctx, page, routes, route, ts) + if err != nil { + return nil, err + } + out = append(out, filtered...) if len(out) >= limit { clear(out[limit:]) return out[:limit], nil @@ -2266,7 +2270,13 @@ func (s *ShardStore) scanRouteAtDirectionWithS3AuxiliaryOwnerFilter( return out, nil } -func filterS3AuxiliaryKVsOwnedByRoute(kvs []*store.KVPair, routes []distribution.Route, route distribution.Route) []*store.KVPair { +func (s *ShardStore) filterS3AuxiliaryKVsOwnedByRoute( + ctx context.Context, + kvs []*store.KVPair, + routes []distribution.Route, + route distribution.Route, + ts uint64, +) ([]*store.KVPair, error) { out := make([]*store.KVPair, 0, len(kvs)) for _, kvp := range kvs { if kvp == nil { @@ -2274,11 +2284,43 @@ func filterS3AuxiliaryKVsOwnedByRoute(kvs []*store.KVPair, routes []distribution } owner, auxiliary := s3BucketAuxiliaryOwnerRoute(kvp.Key, routes) if auxiliary && !routeMatchesS3BucketAuxiliaryOwner(route, owner) { - continue + covered, err := s.s3BucketAuxiliaryOwnerHasVersionAt(ctx, owner, kvp.Key, ts) + if err != nil { + return nil, err + } + if covered { + continue + } } out = append(out, kvp) } - return out + return out, nil +} + +func (s *ShardStore) s3BucketAuxiliaryOwnerHasVersionAt( + ctx context.Context, + owner distribution.Route, + key []byte, + ts uint64, +) (bool, error) { + g, ok := s.groupForID(owner.GroupID) + if !ok || g == nil || g.Store == nil { + return false, nil + } + live, liveOK, err := latestMVCCVersionAt(ctx, g.Store, key, ts) + if err != nil { + return false, err + } + if !routeHasStagedVisibility(owner) { + return liveOK, nil + } + stagedKey := distribution.MigrationStagedDataKey(owner.MigrationJobID, key) + staged, stagedOK, err := latestMVCCVersionAt(ctx, g.Store, stagedKey, ts) + if err != nil { + return false, err + } + _, found := newerMigrationVersion(live, liveOK, staged, stagedOK) + return found, nil } func routeMatchesS3BucketAuxiliaryOwner(route distribution.Route, owner distribution.Route) bool { diff --git a/kv/shard_store_test.go b/kv/shard_store_test.go index 9d3926b65..aface7e73 100644 --- a/kv/shard_store_test.go +++ b/kv/shard_store_test.go @@ -281,6 +281,63 @@ func TestShardStoreS3BucketAuxiliaryScanUsesPromotedOwner(t *testing.T) { }, reverse) } +func TestShardStoreS3BucketAuxiliaryScanPreservesLegacyRawOnlyRows(t *testing.T) { + t.Parallel() + + ctx := context.Background() + for _, tc := range []struct { + name string + prefix string + keyFor func(string) []byte + value []byte + rawValue []byte + }{ + {name: "bucket meta", prefix: s3keys.BucketMetaPrefix, keyFor: s3keys.BucketMetaKey, value: []byte("legacy-meta"), rawValue: []byte("raw-meta")}, + {name: "bucket generation", prefix: s3keys.BucketGenerationPrefix, keyFor: s3keys.BucketGenerationKey, value: []byte("legacy-generation"), rawValue: []byte("raw-generation")}, + } { + t.Run(tc.name, func(t *testing.T) { + const ( + migratedBucket = "bucket-a" + otherBucket = "bucket-z" + ) + engine := distribution.NewEngine() + require.NoError(t, engine.ApplySnapshot(distribution.CatalogSnapshot{ + Version: 1, + Routes: s3BucketAuxiliaryPromotedRoutes(), + })) + groups := map[uint64]*ShardGroup{ + 1: {Store: store.NewMVCCStore()}, + 2: {Store: store.NewMVCCStore()}, + } + st := NewShardStore(engine, groups) + migratedKey := tc.keyFor(migratedBucket) + otherKey := tc.keyFor(otherBucket) + require.NoError(t, groups[1].Store.PutAt(ctx, migratedKey, tc.value, 10, 0)) + require.NoError(t, groups[1].Store.PutAt(ctx, otherKey, tc.rawValue, 15, 0)) + + start := []byte(tc.prefix) + end := prefixScanEnd(start) + kvs, err := st.ScanAt(ctx, start, end, 10, 30) + require.NoError(t, err) + require.Equal(t, []*store.KVPair{ + {Key: migratedKey, Value: tc.value}, + {Key: otherKey, Value: tc.rawValue}, + }, kvs) + + reverse, err := st.ReverseScanAt(ctx, start, end, 10, 30) + require.NoError(t, err) + require.Equal(t, []*store.KVPair{ + {Key: otherKey, Value: tc.rawValue}, + {Key: migratedKey, Value: tc.value}, + }, reverse) + + exact, err := st.ScanAt(ctx, migratedKey, prefixScanEnd(migratedKey), 10, 30) + require.NoError(t, err) + require.Equal(t, []*store.KVPair{{Key: migratedKey, Value: tc.value}}, exact) + }) + } +} + func TestShardStoreRouteBoundedS3BucketAuxiliaryScanKeepsStagedRows(t *testing.T) { t.Parallel() From 3adfc317ecca509099a77f0aebe77d1299275bf4 Mon Sep 17 00:00:00 2001 From: bootjp Date: Fri, 28 Aug 2026 15:51:45 +0900 Subject: [PATCH 25/58] Expire staged visibility winners --- kv/shard_store.go | 39 +++++++++++++++++++++------------------ kv/shard_store_test.go | 41 +++++++++++++++++++++++++++++++++++++++++ 2 files changed, 62 insertions(+), 18 deletions(-) diff --git a/kv/shard_store.go b/kv/shard_store.go index 29c8544e2..676706b92 100644 --- a/kv/shard_store.go +++ b/kv/shard_store.go @@ -3798,18 +3798,14 @@ func (s *ShardStore) ExpireAt(ctx context.Context, key []byte, expireAt uint64, return errors.WithStack(g.Store.ExpireAt(ctx, key, expireAt, commitTS)) } -// expireStagedVisibleAt applies an expiration to a key whose only visible -// version may still be staged. +// expireStagedVisibleAt applies an expiration to the version visible through +// the staged/live merge. // // Between cutover and promotion a key can be visible through its staged alias -// while the live key holds nothing. Both store implementations read the live key -// first and return ErrKeyNotFound when it is absent, so the expiration would -// fail for a value the same route serves happily through GetAt. Resolve the -// staged/live winner and write the expiration as a live MVCC version, which is -// where every other post-cutover write goes. -// -// The live attempt comes first so the ordinary path is unchanged and the staged -// lookup is only paid when the live key really has nothing. +// while the live key is missing, older, or tombstoned by a newer staged delete. +// Resolve the staged/live winner before writing the expiration, then record the +// result as a live MVCC version, which is where every other post-cutover write +// goes. func (s *ShardStore) expireStagedVisibleAt( ctx context.Context, g *ShardGroup, @@ -3818,16 +3814,23 @@ func (s *ShardStore) expireStagedVisibleAt( expireAt uint64, commitTS uint64, ) error { - err := g.Store.ExpireAt(ctx, key, expireAt, commitTS) - if !errors.Is(err, store.ErrKeyNotFound) { - return errors.WithStack(err) + live, liveOK, err := latestMVCCVersionAt(ctx, g.Store, key, commitTS) + if err != nil { + return err } - value, getErr := s.getAtWithStagedVisibility(ctx, g, route, key, commitTS) - if getErr != nil { - // Nothing visible either way: the original ErrKeyNotFound stands. - return getErr + stagedKey := distribution.MigrationStagedDataKey(route.MigrationJobID, key) + staged, stagedOK, err := latestMVCCVersionAt(ctx, g.Store, stagedKey, commitTS) + if err != nil { + return err } - return errors.WithStack(g.Store.PutWithTTLAt(ctx, key, value, commitTS, expireAt)) + if stagedOK { + staged.Key = bytes.Clone(key) + } + winner, ok := newerMigrationVersion(live, liveOK, staged, stagedOK) + if !ok || !migrationVersionVisible(winner, commitTS) { + return store.ErrKeyNotFound + } + return errors.WithStack(g.Store.PutWithTTLAt(ctx, key, winner.Value, commitTS, expireAt)) } func (s *ShardStore) LatestCommitTS(ctx context.Context, key []byte) (uint64, bool, error) { diff --git a/kv/shard_store_test.go b/kv/shard_store_test.go index aface7e73..e4d48b06d 100644 --- a/kv/shard_store_test.go +++ b/kv/shard_store_test.go @@ -3719,6 +3719,47 @@ func TestExpireAtAppliesToStagedOnlyValues(t *testing.T) { require.ErrorIs(t, err, store.ErrKeyNotFound) } +func TestExpireAtUsesNewerStagedValueOverLiveValue(t *testing.T) { + t.Parallel() + + ctx := context.Background() + st, group := newStagedVisibilityShardStore(t) + + key := []byte("b") + require.NoError(t, group.Store.PutAt(ctx, key, []byte("live-b"), 10, 0)) + require.NoError(t, group.Store.PutAt(ctx, + distribution.MigrationStagedDataKey(9, key), []byte("staged-b"), 20, 0)) + + require.NoError(t, st.ExpireAt(ctx, key, 5_000, 300)) + + got, err := st.GetAt(ctx, key, 300) + require.NoError(t, err) + require.Equal(t, []byte("staged-b"), got) + + live, err := group.Store.GetAt(ctx, key, 300) + require.NoError(t, err) + require.Equal(t, []byte("staged-b"), live) +} + +func TestExpireAtHonorsNewerStagedTombstone(t *testing.T) { + t.Parallel() + + ctx := context.Background() + st, group := newStagedVisibilityShardStore(t) + + key := []byte("b") + require.NoError(t, group.Store.PutAt(ctx, key, []byte("live-b"), 10, 0)) + require.NoError(t, group.Store.DeleteAt(ctx, distribution.MigrationStagedDataKey(9, key), 20)) + + require.ErrorIs(t, st.ExpireAt(ctx, key, 5_000, 300), store.ErrKeyNotFound) + + _, err := st.GetAt(ctx, key, 300) + require.ErrorIs(t, err, store.ErrKeyNotFound) + live, err := group.Store.GetAt(ctx, key, 300) + require.NoError(t, err) + require.Equal(t, []byte("live-b"), live) +} + // A key with nothing on either side still reports ErrKeyNotFound. func TestExpireAtStillFailsWhenNothingIsVisible(t *testing.T) { t.Parallel() From f99f9099bd938d13fd741bc9590af01c3fd61336 Mon Sep 17 00:00:00 2001 From: bootjp Date: Fri, 28 Aug 2026 16:00:37 +0900 Subject: [PATCH 26/58] Reject staged control raw writes --- distribution/migrator.go | 14 ++++---------- kv/fsm_reserved_control_test.go | 16 +++++++++------- 2 files changed, 13 insertions(+), 17 deletions(-) diff --git a/distribution/migrator.go b/distribution/migrator.go index 9e5710b46..b1e8fea06 100644 --- a/distribution/migrator.go +++ b/distribution/migrator.go @@ -105,16 +105,10 @@ var ( // IsReservedControlKey reports whether key lives in a control namespace that no // user mutation may write. // -// The staged-data prefix is deliberately excluded: ShardedCoordinator rewrites a -// user key into MigrationStagedDataKey while its route has staged visibility, so -// legitimate user writes do arrive under it. Telling those apart from a forged -// one needs request-level provenance the FSM does not have -- the route state -// that would distinguish them is refreshed by a polling watcher and so cannot -// decide an apply. +// Staged migration data is included in the reserved set. It is populated by the +// typed migration import/promote paths and by FSM-internal prefix-delete +// expansion, not by externally supplied RawKV mutations. func IsReservedControlKey(key []byte) bool { - if bytes.HasPrefix(key, []byte(migrationStagedDataPrefix)) { - return false - } return reservedControlPrefixIntersects(key, bytes.HasPrefix) } @@ -128,7 +122,7 @@ func IsReservedControlKey(key []byte) bool { // accept more than the one transaction prefix it takes today, which is a // separate change. func ReservedControlPrefixIntersects(prefix []byte) bool { - if len(prefix) == 0 || bytes.HasPrefix(prefix, []byte(migrationStagedDataPrefix)) { + if len(prefix) == 0 { return false } return reservedControlPrefixIntersects(prefix, func(a, b []byte) bool { diff --git a/kv/fsm_reserved_control_test.go b/kv/fsm_reserved_control_test.go index 41dcf68f5..1b7a425f1 100644 --- a/kv/fsm_reserved_control_test.go +++ b/kv/fsm_reserved_control_test.go @@ -27,6 +27,7 @@ func TestValidateRawMutationRejectsReservedControlKeys(t *testing.T) { []byte("!dist|route|0001"), []byte("!dist|job|7"), []byte("!dist|jobhist|7"), + distribution.MigrationStagedDataKey(7, []byte("victim")), []byte("!migstage|7|victim"), []byte("!migwrite|7"), []byte("!migfence|7"), @@ -56,6 +57,7 @@ func TestHandleDelPrefixRejectsReservedControlPrefixes(t *testing.T) { []byte("!dist|"), []byte("!dist|route|"), []byte("!dist"), + distribution.MigrationStagedDataKey(7, []byte("user:")), []byte("!migwrite|"), []byte("!migfence"), } { @@ -69,10 +71,10 @@ func TestHandleDelPrefixRejectsReservedControlPrefixes(t *testing.T) { require.NoError(t, f.handleDelPrefix(ctx, []byte("user:"), 12)) } -// ShardedCoordinator rewrites a user key into MigrationStagedDataKey while its -// route has staged visibility, so legitimate user writes do arrive under the -// staged prefix. Refusing those would break the migration this branch adds. -func TestValidateRawMutationAllowsStagedDataKeys(t *testing.T) { +// Staged migration data is private to typed migration paths. A user-supplied +// RawKV request must not be able to forge a staged key that promotion later +// treats as migrated data. +func TestValidateRawMutationRejectsStagedDataKeys(t *testing.T) { t.Parallel() ctx := context.Background() @@ -82,7 +84,7 @@ func TestValidateRawMutationAllowsStagedDataKeys(t *testing.T) { require.True(t, ok) staged := distribution.MigrationStagedDataKey(7, []byte("user-key")) - require.NoError(t, f.validateRawMutationForApply(ctx, - &pb.Mutation{Op: pb.Op_PUT, Key: staged, Value: []byte("v")}, nil, 10)) - require.NoError(t, f.handleDelPrefix(ctx, distribution.MigrationStagedDataKey(7, []byte("user:")), 11)) + require.ErrorIs(t, f.validateRawMutationForApply(ctx, + &pb.Mutation{Op: pb.Op_PUT, Key: staged, Value: []byte("v")}, nil, 10), ErrInvalidRequest) + require.ErrorIs(t, f.handleDelPrefix(ctx, distribution.MigrationStagedDataKey(7, []byte("user:")), 11), ErrInvalidRequest) } From 605958446cd5b9c7efe89c7076e8613d2f7531a6 Mon Sep 17 00:00:00 2001 From: bootjp Date: Fri, 28 Aug 2026 16:08:29 +0900 Subject: [PATCH 27/58] Route S3 auxiliary owner probes through leaders --- kv/shard_store.go | 44 ++++++++++++++++++------ kv/shard_store_test.go | 78 ++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 111 insertions(+), 11 deletions(-) diff --git a/kv/shard_store.go b/kv/shard_store.go index 676706b92..1044486f8 100644 --- a/kv/shard_store.go +++ b/kv/shard_store.go @@ -2245,7 +2245,7 @@ func (s *ShardStore) scanRouteAtDirectionWithS3AuxiliaryOwnerFilter( if err != nil { return nil, err } - filtered, err := s.filterS3AuxiliaryKVsOwnedByRoute(ctx, page, routes, route, ts) + filtered, err := s.filterS3AuxiliaryKVsOwnedByRoute(ctx, page, routes, route, ts, readRouteVersion) if err != nil { return nil, err } @@ -2276,6 +2276,7 @@ func (s *ShardStore) filterS3AuxiliaryKVsOwnedByRoute( routes []distribution.Route, route distribution.Route, ts uint64, + readRouteVersion uint64, ) ([]*store.KVPair, error) { out := make([]*store.KVPair, 0, len(kvs)) for _, kvp := range kvs { @@ -2284,7 +2285,7 @@ func (s *ShardStore) filterS3AuxiliaryKVsOwnedByRoute( } owner, auxiliary := s3BucketAuxiliaryOwnerRoute(kvp.Key, routes) if auxiliary && !routeMatchesS3BucketAuxiliaryOwner(route, owner) { - covered, err := s.s3BucketAuxiliaryOwnerHasVersionAt(ctx, owner, kvp.Key, ts) + covered, err := s.s3BucketAuxiliaryOwnerHasVersionAt(ctx, owner, kvp.Key, ts, readRouteVersion) if err != nil { return nil, err } @@ -2302,25 +2303,46 @@ func (s *ShardStore) s3BucketAuxiliaryOwnerHasVersionAt( owner distribution.Route, key []byte, ts uint64, + readRouteVersion uint64, ) (bool, error) { - g, ok := s.groupForID(owner.GroupID) - if !ok || g == nil || g.Store == nil { - return false, nil - } - live, liveOK, err := latestMVCCVersionAt(ctx, g.Store, key, ts) + live, liveOK, err := s.ownerRouteHasVersionAtOrBefore(ctx, owner, key, ts, readRouteVersion) if err != nil { return false, err } + if live { + return true, nil + } if !routeHasStagedVisibility(owner) { - return liveOK, nil + if !liveOK { + return true, nil + } + return false, nil } stagedKey := distribution.MigrationStagedDataKey(owner.MigrationJobID, key) - staged, stagedOK, err := latestMVCCVersionAt(ctx, g.Store, stagedKey, ts) + staged, stagedOK, err := s.ownerRouteHasVersionAtOrBefore(ctx, owner, stagedKey, ts, readRouteVersion) if err != nil { return false, err } - _, found := newerMigrationVersion(live, liveOK, staged, stagedOK) - return found, nil + if staged { + return true, nil + } + if !liveOK || !stagedOK { + return true, nil + } + return false, nil +} + +func (s *ShardStore) ownerRouteHasVersionAtOrBefore( + ctx context.Context, + owner distribution.Route, + key []byte, + ts uint64, + readRouteVersion uint64, +) (bool, bool, error) { + if exists, ok, err := s.routeHasVersionAtOrBefore(ctx, owner, key, ts); ok || err != nil { + return exists, ok, err + } + return s.routeHasVersionAtOrBeforeRemote(ctx, owner, key, ts, readRouteVersion) } func routeMatchesS3BucketAuxiliaryOwner(route distribution.Route, owner distribution.Route) bool { diff --git a/kv/shard_store_test.go b/kv/shard_store_test.go index e4d48b06d..c28ecbfdf 100644 --- a/kv/shard_store_test.go +++ b/kv/shard_store_test.go @@ -652,6 +652,84 @@ func TestShardStoreS3BucketAuxiliaryScanHonorsStagedTombstone(t *testing.T) { require.Empty(t, kvs) } +type versionVisibleRawKVServer struct { + pb.UnimplementedRawKVServer + + mu sync.Mutex + visible map[string]bool + latestReqs []*pb.RawLatestCommitTSRequest +} + +func (s *versionVisibleRawKVServer) RawGet(context.Context, *pb.RawGetRequest) (*pb.RawGetResponse, error) { + return &pb.RawGetResponse{}, nil +} + +func (s *versionVisibleRawKVServer) RawScanAt(context.Context, *pb.RawScanAtRequest) (*pb.RawScanAtResponse, error) { + return &pb.RawScanAtResponse{}, nil +} + +func (s *versionVisibleRawKVServer) RawLatestCommitTS(_ context.Context, req *pb.RawLatestCommitTSRequest) (*pb.RawLatestCommitTSResponse, error) { + s.mu.Lock() + defer s.mu.Unlock() + s.latestReqs = append(s.latestReqs, &pb.RawLatestCommitTSRequest{ + Key: bytes.Clone(req.GetKey()), + GroupId: req.GetGroupId(), + ReadRouteVersion: req.GetReadRouteVersion(), + VersionVisibleAtTs: req.GetVersionVisibleAtTs(), + }) + return &pb.RawLatestCommitTSResponse{ + VersionVisible: s.visible[string(req.GetKey())], + VersionVisibleSupported: true, + }, nil +} + +func TestShardStoreS3BucketAuxiliaryOwnerProbeUsesLeaderRoutedReadFence(t *testing.T) { + t.Parallel() + + ctx := context.Background() + const bucket = "bucket-a" + migratedKey := s3keys.BucketMetaKey(bucket) + stagedKey := distribution.MigrationStagedDataKey(9, migratedKey) + probe := &versionVisibleRawKVServer{ + visible: map[string]bool{string(stagedKey): true}, + } + addr, stop := startRawKVServer(t, probe) + t.Cleanup(stop) + + engine := distribution.NewEngine() + require.NoError(t, engine.ApplySnapshot(distribution.CatalogSnapshot{ + Version: 77, + Routes: s3BucketAuxiliaryStagedRoutes(bucket, 1, 2), + })) + groups := map[uint64]*ShardGroup{ + 1: {Store: store.NewMVCCStore()}, + 2: { + Store: store.NewMVCCStore(), + Engine: &followerProxyEngine{leader: addr}, + }, + } + st := NewShardStore(engine, groups) + visibleKey := s3keys.BucketMetaKey("bucket-z") + require.NoError(t, groups[1].Store.PutAt(ctx, migratedKey, []byte("stale-source"), 10, 0)) + require.NoError(t, groups[1].Store.PutAt(ctx, visibleKey, []byte("visible"), 10, 0)) + + start := []byte(s3keys.BucketMetaPrefix) + kvs, err := st.ScanAtWithReadFence(ctx, start, prefixScanEnd(start), 10, 30, false, 0, 77, nil, nil) + require.NoError(t, err) + require.Equal(t, []*store.KVPair{{Key: visibleKey, Value: []byte("visible")}}, kvs) + + probe.mu.Lock() + defer probe.mu.Unlock() + require.Len(t, probe.latestReqs, 2) + require.Equal(t, migratedKey, probe.latestReqs[0].GetKey()) + require.Equal(t, stagedKey, probe.latestReqs[1].GetKey()) + for _, req := range probe.latestReqs { + require.Equal(t, uint64(2), req.GetGroupId()) + require.Equal(t, uint64(77), req.GetReadRouteVersion()) + require.Equal(t, uint64(30), req.GetVersionVisibleAtTs()) + } +} + func TestShardStoreGetAt_ContinuesLatestVersionExportPages(t *testing.T) { t.Parallel() From 70d852f767acb5a384c4e8dfcfdbdd4b40082d3e Mon Sep 17 00:00:00 2001 From: bootjp Date: Fri, 28 Aug 2026 16:18:00 +0900 Subject: [PATCH 28/58] Preserve max-size keys during staged imports --- distribution/migrator_export_plan_test.go | 1 + kv/fsm_migration_import.go | 16 +++--- kv/fsm_migration_import_test.go | 65 ++++++++++++++++------- store/lsm_store.go | 8 +-- store/lsm_store_test.go | 36 +++++++++++++ store/mvcc_store.go | 30 +++++++---- store/mvcc_store_snapshot_test.go | 21 ++++++++ 7 files changed, 135 insertions(+), 42 deletions(-) diff --git a/distribution/migrator_export_plan_test.go b/distribution/migrator_export_plan_test.go index cd8ff4770..2b68519fd 100644 --- a/distribution/migrator_export_plan_test.go +++ b/distribution/migrator_export_plan_test.go @@ -446,6 +446,7 @@ func TestMigrationStagedDataKeyRoundTrip(t *testing.T) { raw := []byte("user|raw") key := MigrationStagedDataKey(42, raw) + require.LessOrEqual(t, len(MigrationStagedDataKey(42, nil)), store.MaxSnapshotInternalKeyEnvelope) require.True(t, IsMigrationStagedDataKey(key)) require.True(t, bytes.HasPrefix(key, MigrationStagedDataKeyPrefix(42))) require.False(t, IsMigrationStagedDataKey([]byte("!dist|migstage|short"))) diff --git a/kv/fsm_migration_import.go b/kv/fsm_migration_import.go index cde823a0f..43dfb285d 100644 --- a/kv/fsm_migration_import.go +++ b/kv/fsm_migration_import.go @@ -77,20 +77,18 @@ func (f *kvFSM) applyMigrationImport(ctx context.Context, data []byte) any { // the batch while the failed voter skips the imported versions for good, which // surfaces as missing data after failover or promotion. Those halt instead, // matching applyMigrationPromote. -// validateStagedKeySizes rejects a batch whose staged form would not fit in a -// snapshot. +// validateStagedKeySizes rejects a batch whose staged form would not fit in the +// physical stored-key snapshot budget. // // MigrationStagedDataKey prepends the staged prefix, the job id, and a -// separator, so a source key that is itself within store.MaxSnapshotKeySize can -// exceed it once staged. Storing it anyway leaves the target holding a version -// that no snapshot can carry: a new or lagging voter could never restore, and so -// could never recover through snapshot transfer. Refuse the import instead, so -// the failure surfaces on the migrator rather than on a future restore. +// separator. store.MaxSnapshotStoredKeySize includes bounded headroom for that +// envelope so source keys at store.MaxSnapshotKeySize remain migratable, while +// still refusing batches that would make the target snapshot unrestorable. func validateStagedKeySizes(versions []store.MVCCVersion) error { for _, version := range versions { - if len(version.Key) > store.MaxSnapshotKeySize { + if len(version.Key) > store.MaxSnapshotStoredKeySize { return errors.Wrapf(store.ErrSnapshotKeyTooLarge, - "staged key length %d > %d", len(version.Key), store.MaxSnapshotKeySize) + "staged key length %d > %d", len(version.Key), store.MaxSnapshotStoredKeySize) } } return nil diff --git a/kv/fsm_migration_import_test.go b/kv/fsm_migration_import_test.go index 5b6f415fa..5345549fc 100644 --- a/kv/fsm_migration_import_test.go +++ b/kv/fsm_migration_import_test.go @@ -232,28 +232,28 @@ func TestApplyMigrationImportOrdinaryErrorsDoNotHalt(t *testing.T) { } // MigrationStagedDataKey prepends the staged prefix, the job id, and a -// separator, so a source key that is itself within store.MaxSnapshotKeySize can -// exceed it once staged. Storing it anyway would leave the target holding a -// version no snapshot can carry, and a new or lagging voter could never restore. -func TestValidateStagedKeySizesRejectsOversizedStagedForm(t *testing.T) { +// separator. Source keys that already fit the ordinary snapshot key budget must +// stay migratable after that envelope is added, while truly oversized staged +// forms are still rejected before import. +func TestValidateStagedKeySizesAllowsSnapshotLimitSourceKey(t *testing.T) { t.Parallel() envelope := len(distribution.MigrationStagedDataKey(1, nil)) require.Positive(t, envelope) + require.LessOrEqual(t, envelope, store.MaxSnapshotInternalKeyEnvelope) - // A source key right at the limit: legal on its own, oversized once staged. atLimit := bytes.Repeat([]byte("k"), store.MaxSnapshotKeySize) staged := migrationStoreVersionsFromProto(9, []*pb.MVCCVersion{{Key: atLimit, CommitTs: 10}}) require.Len(t, staged, 1) require.Greater(t, len(staged[0].Key), store.MaxSnapshotKeySize) - require.ErrorIs(t, validateStagedKeySizes(staged), store.ErrSnapshotKeyTooLarge) - - // With headroom for the envelope it is accepted. - fits := bytes.Repeat([]byte("k"), store.MaxSnapshotKeySize-envelope) - ok := migrationStoreVersionsFromProto(9, []*pb.MVCCVersion{{Key: fits, CommitTs: 10}}) - require.Len(t, ok, 1) - require.Equal(t, store.MaxSnapshotKeySize, len(ok[0].Key)) - require.NoError(t, validateStagedKeySizes(ok)) + require.LessOrEqual(t, len(staged[0].Key), store.MaxSnapshotStoredKeySize) + require.NoError(t, validateStagedKeySizes(staged)) + + tooLong := bytes.Repeat([]byte("k"), store.MaxSnapshotStoredKeySize-envelope+1) + oversized := migrationStoreVersionsFromProto(9, []*pb.MVCCVersion{{Key: tooLong, CommitTs: 10}}) + require.Len(t, oversized, 1) + require.Equal(t, store.MaxSnapshotStoredKeySize+1, len(oversized[0].Key)) + require.ErrorIs(t, validateStagedKeySizes(oversized), store.ErrSnapshotKeyTooLarge) } // The rejection must be an ordinary apply error. Halting would wedge apply on @@ -265,10 +265,10 @@ func TestOversizedStagedKeyIsAnOrdinaryApplyError(t *testing.T) { errors.Wrap(store.ErrSnapshotKeyTooLarge, "staged"))) } -// The size check has to be wired into the apply path, not merely available: an -// oversized staged key must be refused before it is written, and refused as an -// ordinary error so apply is not halted on every voter. -func TestApplyMigrationImportRejectsOversizedStagedKey(t *testing.T) { +// The size check has to be wired into the apply path, not merely available: +// source keys at the snapshot limit must be accepted, but staged forms beyond +// the stored-key snapshot budget must be refused before they are written. +func TestApplyMigrationImportAllowsSnapshotLimitSourceKey(t *testing.T) { t.Parallel() ctx := context.Background() @@ -276,7 +276,6 @@ func TestApplyMigrationImportRejectsOversizedStagedKey(t *testing.T) { t.Cleanup(func() { _ = st.Close() }) fsm := &kvFSM{store: st, hlc: NewHLC()} - // Legal on its own, oversized once the staging envelope is added. atLimit := bytes.Repeat([]byte("k"), store.MaxSnapshotKeySize) data, err := proto.Marshal(&pb.ImportRangeVersionsRequest{ JobId: 9, @@ -286,6 +285,34 @@ func TestApplyMigrationImportRejectsOversizedStagedKey(t *testing.T) { }) require.NoError(t, err) + applied := fsm.applyMigrationImport(ctx, data) + result, ok := applied.(store.ImportVersionsResult) + require.True(t, ok, "got %T: %v", applied, applied) + require.Equal(t, uint64(10), result.MaxImportedTS) + + got, getErr := st.GetAt(ctx, distribution.MigrationStagedDataKey(9, atLimit), 100) + require.NoError(t, getErr) + require.Equal(t, []byte("v"), got) +} + +func TestApplyMigrationImportRejectsOversizedStagedKey(t *testing.T) { + t.Parallel() + + ctx := context.Background() + st := store.NewMVCCStore() + t.Cleanup(func() { _ = st.Close() }) + fsm := &kvFSM{store: st, hlc: NewHLC()} + + envelope := len(distribution.MigrationStagedDataKey(9, nil)) + oversizedSource := bytes.Repeat([]byte("k"), store.MaxSnapshotStoredKeySize-envelope+1) + data, err := proto.Marshal(&pb.ImportRangeVersionsRequest{ + JobId: 9, + BracketId: 1, + BatchSeq: 1, + Versions: []*pb.MVCCVersion{{Key: oversizedSource, CommitTs: 10, Value: []byte("v")}}, + }) + require.NoError(t, err) + applied := fsm.applyMigrationImport(ctx, data) applyErr, ok := applied.(error) require.True(t, ok, "got %T: %v", applied, applied) @@ -293,6 +320,6 @@ func TestApplyMigrationImportRejectsOversizedStagedKey(t *testing.T) { require.NotErrorIs(t, applyErr, ErrMigrationImportApply, "must not halt apply") // Nothing was staged. - _, getErr := st.GetAt(ctx, distribution.MigrationStagedDataKey(9, atLimit), 100) + _, getErr := st.GetAt(ctx, distribution.MigrationStagedDataKey(9, oversizedSource), 100) require.ErrorIs(t, getErr, store.ErrKeyNotFound) } diff --git a/store/lsm_store.go b/store/lsm_store.go index 4eacc1d57..6cb2152e0 100644 --- a/store/lsm_store.go +++ b/store/lsm_store.go @@ -55,10 +55,10 @@ const ( // maxPebbleEncodedKeySize is the limit for encoded Pebble on-disk keys, // which are the user key concatenated with the 8-byte inverted timestamp. - // Using maxSnapshotKeySize+timestampSize (instead of just maxSnapshotKeySize) - // avoids rejecting keys that are valid at the user-key level but slightly - // exceed maxSnapshotKeySize once the timestamp suffix is appended. - maxPebbleEncodedKeySize = maxSnapshotKeySize + timestampSize + // Using maxSnapshotStoredKeySize+timestampSize avoids rejecting logical keys + // that fit maxSnapshotKeySize but gain a bounded internal envelope before the + // timestamp suffix is appended. + maxPebbleEncodedKeySize = maxSnapshotStoredKeySize + timestampSize // defaultPebbleCacheBytes is the fallback process-wide Pebble block-cache // capacity when the node's effective memory budget cannot be discovered. diff --git a/store/lsm_store_test.go b/store/lsm_store_test.go index dbd6605e1..fb4c9a3ef 100644 --- a/store/lsm_store_test.go +++ b/store/lsm_store_test.go @@ -1212,3 +1212,39 @@ func TestPebbleStore_SnapshotRestore_MaxSizeKey(t *testing.T) { require.NoError(t, err) assert.Equal(t, []byte("val"), got) } + +func TestPebbleStore_SnapshotRestore_MaxStoredKey(t *testing.T) { + dir, err := os.MkdirTemp("", "pebble-maxstoredkey-snap-*") + require.NoError(t, err) + defer os.RemoveAll(dir) + + s, err := NewPebbleStore(dir) + require.NoError(t, err) + defer func() { assert.NoError(t, s.Close()) }() + + ctx := context.Background() + bigKey := bytes.Repeat([]byte("k"), maxSnapshotStoredKeySize) + require.NoError(t, s.PutAt(ctx, bigKey, []byte("val"), 1, 0)) + + snap, err := s.Snapshot() + require.NoError(t, err) + defer func() { assert.NoError(t, snap.Close()) }() + + var buf bytes.Buffer + _, err = snap.WriteTo(&buf) + require.NoError(t, err) + + dir2, err := os.MkdirTemp("", "pebble-maxstoredkey-restore-*") + require.NoError(t, err) + defer os.RemoveAll(dir2) + + s2, err := NewPebbleStore(dir2) + require.NoError(t, err) + defer func() { assert.NoError(t, s2.Close()) }() + + require.NoError(t, s2.Restore(bytes.NewReader(buf.Bytes()))) + + got, err := s2.GetAt(ctx, bigKey, 1) + require.NoError(t, err) + assert.Equal(t, []byte("val"), got) +} diff --git a/store/mvcc_store.go b/store/mvcc_store.go index 614b82952..e6853a28e 100644 --- a/store/mvcc_store.go +++ b/store/mvcc_store.go @@ -28,14 +28,24 @@ const ( mvccSnapshotVersionV1 = uint32(1) mvccSnapshotVersionV2 = uint32(2) mvccSnapshotVersion = uint32(3) - maxSnapshotKeySize = 1 << 20 // 1 MiB per key - // MaxSnapshotKeySize is the largest key a snapshot can carry. Callers that - // wrap a key in an envelope before storing it -- migration staging is the - // one today -- need it to reserve headroom, because a key that fits on its - // own can stop fitting once wrapped, and the snapshot that would carry it - // is then unrestorable. - MaxSnapshotKeySize = maxSnapshotKeySize - maxSnapshotVersionCount = 1 << 20 // 1M versions per key + maxSnapshotKeySize = 1 << 20 // 1 MiB per logical key + // maxSnapshotInternalKeyEnvelope reserves snapshot headroom for internal + // stored-key envelopes such as migration staging. Logical source keys keep + // the full maxSnapshotKeySize budget so existing snapshot-valid rows remain + // migratable. + maxSnapshotInternalKeyEnvelope = 32 + maxSnapshotStoredKeySize = maxSnapshotKeySize + maxSnapshotInternalKeyEnvelope + + // MaxSnapshotKeySize is the largest logical/source key size guaranteed to + // round-trip through snapshots. + MaxSnapshotKeySize = maxSnapshotKeySize + // MaxSnapshotInternalKeyEnvelope is the stored-key envelope headroom + // snapshots reserve for internal namespaces. + MaxSnapshotInternalKeyEnvelope = maxSnapshotInternalKeyEnvelope + // MaxSnapshotStoredKeySize is the largest physical stored key a snapshot can + // carry after adding a bounded internal envelope to a logical key. + MaxSnapshotStoredKeySize = maxSnapshotStoredKeySize + maxSnapshotVersionCount = 1 << 20 // 1M versions per key ) // maxSnapshotValueSize caps the allowed size of a single value during streaming @@ -1260,8 +1270,8 @@ func readMVCCSnapshotEntry(r io.Reader) ([]byte, []VersionedValue, bool, error) } return nil, nil, false, errors.WithStack(err) } - if keyLen > maxSnapshotKeySize { - return nil, nil, false, errors.Wrapf(ErrSnapshotKeyTooLarge, "%d > %d", keyLen, maxSnapshotKeySize) + if keyLen > maxSnapshotStoredKeySize { + return nil, nil, false, errors.Wrapf(ErrSnapshotKeyTooLarge, "%d > %d", keyLen, maxSnapshotStoredKeySize) } key := make([]byte, keyLen) diff --git a/store/mvcc_store_snapshot_test.go b/store/mvcc_store_snapshot_test.go index 97e8da5ed..c52c050de 100644 --- a/store/mvcc_store_snapshot_test.go +++ b/store/mvcc_store_snapshot_test.go @@ -38,6 +38,27 @@ func TestMVCCStore_SnapshotRestoreRoundTrip(t *testing.T) { require.Equal(t, []byte("v2"), v) } +func TestMVCCStore_SnapshotRestoreMaxStoredKey(t *testing.T) { + t.Parallel() + + ctx := context.Background() + src := newTestMVCCStore(t) + key := bytes.Repeat([]byte("k"), MaxSnapshotStoredKeySize) + require.NoError(t, src.PutAt(ctx, key, []byte("v"), 10, 0)) + + snap, err := src.Snapshot() + require.NoError(t, err) + defer snap.Close() + raw := snapshotBytes(t, snap) + + dst := newTestMVCCStore(t) + require.NoError(t, dst.Restore(bytes.NewReader(raw))) + + got, err := dst.GetAt(ctx, key, 10) + require.NoError(t, err) + require.Equal(t, []byte("v"), got) +} + func TestMVCCStore_RestoreRejectsInvalidChecksum(t *testing.T) { t.Parallel() From adfcf26dec4928880e3a83a498e0d691ca6debaf Mon Sep 17 00:00:00 2001 From: bootjp Date: Fri, 28 Aug 2026 16:34:30 +0900 Subject: [PATCH 29/58] Support group-aware owner probes --- adapter/grpc_test.go | 35 ++++++++++++++++++ kv/shard_store.go | 81 +++++++++++++++++++++++++++++++++++++++--- kv/shard_store_test.go | 51 ++++++++++++++++++++++++++ 3 files changed, 162 insertions(+), 5 deletions(-) diff --git a/adapter/grpc_test.go b/adapter/grpc_test.go index a8b0e06fb..9fb2634cd 100644 --- a/adapter/grpc_test.go +++ b/adapter/grpc_test.go @@ -37,6 +37,8 @@ const ( grpcSequenceShortIterations = 256 ) +var _ rawGroupCommitTSReader = (*kvstore.ShardStore)(nil) + func grpcSequenceIterations(t testing.TB) int { t.Helper() if testing.Short() { @@ -403,6 +405,39 @@ func TestGRPCServer_RawLatestCommitTS_UsesExplicitGroup(t *testing.T) { require.Zero(t, st.latestReadRouteVersion) } +func TestGRPCServer_RawLatestCommitTS_ExplicitGroupShardStoreVersionProbe(t *testing.T) { + t.Parallel() + + ctx := context.Background() + engine := distribution.NewEngine() + require.NoError(t, engine.ApplySnapshot(distribution.CatalogSnapshot{ + Version: 1, + Routes: []distribution.RouteDescriptor{ + {RouteID: 1, Start: []byte(""), End: nil, GroupID: 42, State: distribution.RouteStateActive}, + }, + })) + groupStore := store.NewMVCCStore() + t.Cleanup(func() { require.NoError(t, groupStore.Close()) }) + key := []byte("!dist|migstage|probe|k") + require.NoError(t, groupStore.PutAt(ctx, key, []byte("v"), 10, 0)) + shards := kvstore.NewShardStore(engine, map[uint64]*kvstore.ShardGroup{ + 42: {Store: groupStore}, + }) + server := NewGRPCServer(shards, nil) + + resp, err := server.RawLatestCommitTS(ctx, &pb.RawLatestCommitTSRequest{ + Key: key, + GroupId: 42, + ReadRouteVersion: 1, + VersionVisibleAtTs: 10, + }) + require.NoError(t, err) + require.True(t, resp.GetExists()) + require.Equal(t, uint64(10), resp.GetTs()) + require.True(t, resp.GetVersionVisibleSupported()) + require.True(t, resp.GetVersionVisible()) +} + func TestGRPCServer_RawReadFenceHelpersKeepCallerRouteVersion(t *testing.T) { t.Parallel() diff --git a/kv/shard_store.go b/kv/shard_store.go index 1044486f8..9d1b2f0b6 100644 --- a/kv/shard_store.go +++ b/kv/shard_store.go @@ -876,10 +876,21 @@ func (s *ShardStore) scanExplicitGroupAtWithReadFence(ctx context.Context, group if !clampToRoutes && !routeFilterPresent { routes, dedupeByKey = prepareUnclampedRawScanRoutes(routes, dedupeByKey) } - return s.scanExplicitGroupRoutesAtWithReadFence(ctx, routes, start, end, limit, ts, reverse, readRouteVersion, routeStart, routeEnd, clampToRoutes, dedupeByKey) + ownerFilterRoutes := s.s3BucketAuxiliaryOwnerFilterRoutes(start, end, routes, dedupeByKey) + return s.scanExplicitGroupRoutesAtWithReadFence(ctx, routes, ownerFilterRoutes, start, end, limit, ts, reverse, readRouteVersion, routeStart, routeEnd, clampToRoutes, dedupeByKey) } -func (s *ShardStore) 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) { +func (s *ShardStore) s3BucketAuxiliaryOwnerFilterRoutes(start []byte, end []byte, routes []distribution.Route, dedupeByKey bool) []distribution.Route { + if !dedupeByKey { + return routes + } + if candidates, _, ok := s.routesForS3BucketAuxiliaryScan(start, end); ok && len(candidates) > 0 { + return candidates + } + return routes +} + +func (s *ShardStore) scanExplicitGroupRoutesAtWithReadFence(ctx context.Context, routes []distribution.Route, ownerFilterRoutes []distribution.Route, start []byte, end []byte, limit int, ts uint64, reverse bool, readRouteVersion uint64, routeStart []byte, routeEnd []byte, clampToRoutes bool, dedupeByKey bool) ([]*store.KVPair, error) { out := make([]*store.KVPair, 0) for i := 0; i < len(routes); i++ { route := routes[i] @@ -892,7 +903,7 @@ func (s *ShardStore) scanExplicitGroupRoutesAtWithReadFence(ctx context.Context, scanStart = clampScanStart(start, route.Start) scanEnd = clampScanEnd(end, route.End) } - kvs, err := s.scanRouteAtDirectionWithS3AuxiliaryOwnerFilter(ctx, routes, route, scanStart, scanEnd, limit, ts, reverse, true, readRouteVersion, routeStart, routeEnd, dedupeByKey) + kvs, err := s.scanRouteAtDirectionWithS3AuxiliaryOwnerFilter(ctx, ownerFilterRoutes, route, scanStart, scanEnd, limit, ts, reverse, true, readRouteVersion, routeStart, routeEnd, dedupeByKey) if err != nil { return nil, err } @@ -2314,7 +2325,7 @@ func (s *ShardStore) s3BucketAuxiliaryOwnerHasVersionAt( } if !routeHasStagedVisibility(owner) { if !liveOK { - return true, nil + return false, ownerVersionProbeUnavailable(owner, key) } return false, nil } @@ -2327,11 +2338,15 @@ func (s *ShardStore) s3BucketAuxiliaryOwnerHasVersionAt( return true, nil } if !liveOK || !stagedOK { - return true, nil + return false, ownerVersionProbeUnavailable(owner, key) } return false, nil } +func ownerVersionProbeUnavailable(route distribution.Route, key []byte) error { + return errors.Wrapf(ErrLeaderNotFound, "s3 auxiliary owner version probe unavailable group_id=%d key=%q", route.GroupID, key) +} + func (s *ShardStore) ownerRouteHasVersionAtOrBefore( ctx context.Context, owner distribution.Route, @@ -3893,6 +3908,33 @@ func (s *ShardStore) LatestCommitTSWithReadFence(ctx context.Context, key []byte return s.proxyLatestCommitTS(ctx, g, key, readRouteVersion) } +func (s *ShardStore) LatestCommitTSGroupWithReadFence(ctx context.Context, key []byte, groupID uint64, readRouteVersion uint64) (uint64, bool, error) { + if groupID == 0 { + return s.LatestCommitTSWithReadFence(ctx, key, readRouteVersion) + } + if err := s.awaitReadRouteVersion(ctx, readRouteVersion); err != nil { + return 0, false, err + } + g, ok := s.groupForID(groupID) + if !ok || g.Store == nil { + return 0, false, nil + } + + if engineForGroup(g) == nil { + ts, exists, err := g.Store.LatestCommitTS(ctx, key) + return ts, exists, errors.WithStack(err) + } + + if engine := engineForGroup(g); isLeaderEngine(engine) { + if _, err := leaseReadEngineCtx(ctx, engine); err == nil { + ts, exists, err := g.Store.LatestCommitTS(ctx, key) + return ts, exists, errors.WithStack(err) + } + } + + return s.proxyLatestCommitTSGroup(ctx, g, key, groupID, readRouteVersion) +} + func (s *ShardStore) 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 { @@ -3942,6 +3984,35 @@ func (s *ShardStore) proxyLatestCommitTS(ctx context.Context, g *ShardGroup, key return resp.Ts, resp.Exists, nil } +func (s *ShardStore) proxyLatestCommitTSGroup(ctx context.Context, g *ShardGroup, key []byte, groupID uint64, readRouteVersion uint64) (uint64, bool, error) { + engine := engineForGroup(g) + if engine == nil { + return 0, false, nil + } + addr := leaderAddrFromEngine(engine) + if addr == "" { + return 0, false, errors.WithStack(ErrLeaderNotFound) + } + + conn, err := s.connCache.ConnFor(addr) + if err != nil { + return 0, false, err + } + + ctx, cancel := context.WithTimeout(ctx, proxyForwardTimeout) + defer cancel() + cli := pb.NewRawKVClient(conn) + resp, err := cli.RawLatestCommitTS(ctx, &pb.RawLatestCommitTSRequest{ + Key: key, + GroupId: groupID, + ReadRouteVersion: readRouteVersion, + }) + if err != nil { + return 0, false, errors.WithStack(err) + } + return resp.Ts, resp.Exists, nil +} + func (s *ShardStore) maybeResolveTxnLock(ctx context.Context, g *ShardGroup, key []byte, readTS uint64) error { lock, ok, err := loadTxnLockAt(ctx, g, key, readTS) if err != nil { diff --git a/kv/shard_store_test.go b/kv/shard_store_test.go index c28ecbfdf..e76ce19a8 100644 --- a/kv/shard_store_test.go +++ b/kv/shard_store_test.go @@ -730,6 +730,57 @@ func TestShardStoreS3BucketAuxiliaryOwnerProbeUsesLeaderRoutedReadFence(t *testi } } +func TestShardStoreS3BucketAuxiliaryOwnerProbeFailsWhenLeaderUnavailable(t *testing.T) { + t.Parallel() + + ctx := context.Background() + const bucket = "bucket-a" + migratedKey := s3keys.BucketMetaKey(bucket) + engine := distribution.NewEngine() + require.NoError(t, engine.ApplySnapshot(distribution.CatalogSnapshot{ + Version: 77, + Routes: s3BucketAuxiliaryStagedRoutes(bucket, 1, 2), + })) + groups := map[uint64]*ShardGroup{ + 1: {Store: store.NewMVCCStore()}, + 2: { + Store: store.NewMVCCStore(), + Engine: &followerProxyEngine{}, + }, + } + st := NewShardStore(engine, groups) + require.NoError(t, groups[1].Store.PutAt(ctx, migratedKey, []byte("stale-source"), 10, 0)) + + start := []byte(s3keys.BucketMetaPrefix) + _, err := st.ScanAtWithReadFence(ctx, start, prefixScanEnd(start), 10, 30, false, 0, 77, nil, nil) + require.ErrorIs(t, err, ErrLeaderNotFound) +} + +func TestShardStoreExplicitGroupS3BucketAuxiliaryScanKeepsOwnerRoutes(t *testing.T) { + t.Parallel() + + ctx := context.Background() + const bucket = "bucket-a" + migratedKey := s3keys.BucketMetaKey(bucket) + engine := distribution.NewEngine() + require.NoError(t, engine.ApplySnapshot(distribution.CatalogSnapshot{ + Version: 77, + Routes: s3BucketAuxiliaryStagedRoutes(bucket, 1, 2), + })) + groups := map[uint64]*ShardGroup{ + 1: {Store: store.NewMVCCStore()}, + 2: {Store: store.NewMVCCStore()}, + } + st := NewShardStore(engine, groups) + require.NoError(t, groups[1].Store.PutAt(ctx, migratedKey, []byte("stale-source"), 10, 0)) + require.NoError(t, groups[2].Store.PutAt(ctx, distribution.MigrationStagedDataKey(9, migratedKey), []byte("owner"), 20, 0)) + + start := []byte(s3keys.BucketMetaPrefix) + kvs, err := st.ScanAtWithReadFence(ctx, start, prefixScanEnd(start), 10, 30, false, 1, 77, nil, nil) + require.NoError(t, err) + require.Empty(t, kvs) +} + func TestShardStoreGetAt_ContinuesLatestVersionExportPages(t *testing.T) { t.Parallel() From cd83c0563c050e46e41ab1d0aff7aa59094fc35e Mon Sep 17 00:00:00 2001 From: bootjp Date: Fri, 28 Aug 2026 19:59:18 +0900 Subject: [PATCH 30/58] migration: resolve txn-wrapped keys through the partition resolver The group-aware export filter probed the partition resolver with the raw key. A transaction-family bracket carries the user key wrapped in !txn|int| / !txn|cmt| / !txn|rb| / !txn|ok| / !txn|meta|, and the resolver only recognises bare partitioned-family prefixes, so the wrapper answered "not mine" and fell through to routeKey's !sqs|route|global collapse. A partitioned HT-FIFO row could therefore be accepted by resolver group in its data bracket while its intent, commit, rollback or meta row was excluded by byte range -- a migration that completes with the rows copied and their transaction state left behind. Probe the embedded user key instead: route ownership of a txn key is the ownership of the key it locks. Claude-Session: https://claude.ai/code/session_013rNHooj7NF3giihWVba8QE --- kv/migrator_filter.go | 23 ++++++++++-- kv/migrator_filter_test.go | 71 ++++++++++++++++++++++++++++++++++++++ 2 files changed, 92 insertions(+), 2 deletions(-) create mode 100644 kv/migrator_filter_test.go diff --git a/kv/migrator_filter.go b/kv/migrator_filter.go index 3aa2a7a93..7036e64bd 100644 --- a/kv/migrator_filter.go +++ b/kv/migrator_filter.go @@ -22,10 +22,18 @@ func RouteKeyFilterForGroup(rangeStart, rangeEnd []byte, sourceGroupID uint64, r end := bytes.Clone(rangeEnd) return func(rawKey []byte) bool { if resolver != nil { - if gid, ok := resolver.ResolveGroup(rawKey); ok { + // Transaction-family brackets carry the user key wrapped in + // !txn|int| / !txn|cmt| / ... . The resolver only recognises bare + // partitioned-family prefixes, so probing the wrapper would answer + // "not mine" and drop the key onto routeKey's !sqs|route|global + // collapse -- a different verdict from the one the embedded row + // gets in its own data bracket, which is how a migration ends up + // copying SQS rows while leaving their intents behind. + probe := partitionResolverProbeKey(rawKey) + if gid, ok := resolver.ResolveGroup(probe); ok { return gid == sourceGroupID } - if resolver.RecognisesPartitionedKey(rawKey) { + if resolver.RecognisesPartitionedKey(probe) { return false } } @@ -37,6 +45,17 @@ func RouteKeyFilterForGroup(rangeStart, rangeEnd []byte, sourceGroupID uint64, r } } +// partitionResolverProbeKey returns the key a PartitionResolver should be asked +// about: the embedded user key for a transaction-internal wrapper, the key +// itself otherwise. Route ownership of a txn key is the ownership of the user +// key it locks, so both must be resolved the same way. +func partitionResolverProbeKey(rawKey []byte) []byte { + if embedded, ok := txnRouteKey(rawKey); ok { + return embedded + } + return rawKey +} + func s3BucketAuxiliaryRouteInRange(rawKey, routeStart, routeEnd []byte) bool { bucketRouteStart, bucketRouteEnd, ok := s3BucketAuxiliaryRouteRange(rawKey) if !ok { diff --git a/kv/migrator_filter_test.go b/kv/migrator_filter_test.go new file mode 100644 index 000000000..ba99d689d --- /dev/null +++ b/kv/migrator_filter_test.go @@ -0,0 +1,71 @@ +package kv + +import ( + "testing" + + "github.com/stretchr/testify/require" +) + +// A transaction-family bracket holds the user key wrapped in !txn|int| and +// friends. The group filter must resolve those through the same partition +// resolver that decides the embedded row's own bracket: a partitioned HT-FIFO +// SQS row and its intent belong to one group, and answering differently for +// the two is how a migration copies the data and strands the transaction +// state. +func TestRouteKeyFilterForGroupResolvesTxnWrappedPartitionedKeys(t *testing.T) { + t.Parallel() + + const ( + sourceGroup = uint64(7) + otherGroup = uint64(9) + ) + partitioned := []byte("!sqs|msg|data|p|q|0|m1") + elsewhere := []byte("!sqs|msg|data|p|q|1|m1") + resolver := &fakePartitionResolver{ + routes: map[string]uint64{ + string(partitioned): sourceGroup, + string(elsewhere): otherGroup, + }, + recognisedPrefix: []byte("!sqs|msg|data|p|"), + } + + // A route range that excludes the !sqs|route|global collapse target, so a + // filter that falls through to the byte-range path answers false. + filter := RouteKeyFilterForGroup([]byte("a"), []byte("b"), sourceGroup, resolver) + + for _, tc := range []struct { + name string + key []byte + want bool + }{ + {name: "bare partitioned row", key: partitioned, want: true}, + {name: "intent", key: txnIntentKey(partitioned), want: true}, + {name: "lock", key: txnLockKey(partitioned), want: true}, + {name: "meta", key: append(append([]byte{}, txnMetaPrefixBytes...), partitioned...), want: true}, + {name: "intent for another group", key: txnIntentKey(elsewhere), want: false}, + {name: "bare row for another group", key: elsewhere, want: false}, + } { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + require.Equal(t, tc.want, filter(tc.key), "key %q", tc.key) + }) + } +} + +// A recognised-but-unresolved partitioned key fails closed whether or not it +// arrives wrapped: the resolver cannot say which group owns it, and guessing +// through the byte-range route would export it from the wrong source. +func TestRouteKeyFilterForGroupFailsClosedOnWrappedUnresolvedKey(t *testing.T) { + t.Parallel() + + unresolved := []byte("!sqs|msg|data|p|unknown|0|m1") + resolver := &fakePartitionResolver{ + routes: map[string]uint64{}, + recognisedPrefix: []byte("!sqs|msg|data|p|"), + } + filter := RouteKeyFilterForGroup(nil, nil, 7, resolver) + + require.False(t, filter(unresolved)) + require.False(t, filter(txnIntentKey(unresolved))) + require.False(t, filter(txnCommitKey(unresolved, 42))) +} From 3ad0687e87c514bf871ee2da25ef59e34eec4f5e Mon Sep 17 00:00:00 2001 From: bootjp Date: Fri, 28 Aug 2026 19:59:18 +0900 Subject: [PATCH 31/58] migration: refuse export pages the transport cannot carry The store's MaxBytes budget is consulted after a version has been appended, and storage accepts values up to 256 MiB against the internal gRPC limit of 64 MiB, so one oversized row can land in an otherwise bounded export page. Sending it fails with ResourceExhausted -- and so does every retry of the same cursor, leaving the bracket stuck on an error that never names the row responsible. Check the page's payload bytes against the message budget before sending and fail with the offending key and its size. The check sums the caller-controlled bytes rather than calling proto.Size, which caches its result inside the message. Carrying values that large needs a chunked migration wire format, which is a protocol change rather than a fix here. Claude-Session: https://claude.ai/code/session_013rNHooj7NF3giihWVba8QE --- adapter/internal.go | 58 ++++++++++++++++++++++++++++-- adapter/internal_migration_test.go | 33 +++++++++++++++++ 2 files changed, 89 insertions(+), 2 deletions(-) diff --git a/adapter/internal.go b/adapter/internal.go index ad79c135f..5686d43ab 100644 --- a/adapter/internal.go +++ b/adapter/internal.go @@ -7,6 +7,7 @@ import ( "strings" "github.com/bootjp/elastickv/distribution" + "github.com/bootjp/elastickv/internal" "github.com/bootjp/elastickv/internal/raftengine" "github.com/bootjp/elastickv/internal/s3keys" "github.com/bootjp/elastickv/kv" @@ -129,6 +130,13 @@ const ( // response already supports. Same shape as the promotion path's clamp. maxMigrationExportChunkBytes = 32 << 20 maxMigrationExportScanBytes = maxMigrationExportChunkBytes * defaultMigrationExportScanFactor + + // Headroom for proto framing over the payload bytes a page carries: field + // tags and length prefixes for at most defaultMigrationExportMaxVersions + // entries come to tens of kilobytes, so a mebibyte covers them with room + // to spare. + migrationExportPageFramingHeadroom = 1 << 20 + migrationExportPageByteBudget = internal.GRPCMaxMessageBytes - migrationExportPageFramingHeadroom ) func (i *Internal) Forward(ctx context.Context, req *pb.ForwardRequest) (*pb.ForwardResponse, error) { @@ -220,11 +228,15 @@ func (i *Internal) streamExportRangeVersions(req *pb.ExportRangeVersionsRequest, 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{ + resp := &pb.ExportRangeVersionsResponse{ Versions: protoMVCCVersionsFromStore(result.Versions), NextCursor: result.NextCursor, Done: result.Done, - }); err != nil { + } + if err := checkMigrationExportPageSize(resp); err != nil { + return err + } + if err := stream.Send(resp); err != nil { return errors.WithStack(err) } if result.Done { @@ -234,6 +246,48 @@ func (i *Internal) streamExportRangeVersions(req *pb.ExportRangeVersionsRequest, } } +// checkMigrationExportPageSize refuses to send a page the transport cannot +// carry. The store's MaxBytes budget is only consulted after a version has +// been appended, and storage accepts values far above the message limit +// (store.maxSnapshotValueSize is 256 MiB against internal.GRPCMaxMessageBytes' +// 64 MiB), so a single oversized row can land in an otherwise bounded page. +// Sending it fails with ResourceExhausted, and so does every retry of the same +// cursor -- the bracket never completes and the error says nothing about which +// row is responsible. Naming the row instead keeps the failure diagnosable. +// Carrying values that large needs a chunked migration wire format, which is a +// protocol change rather than a fix here. +func checkMigrationExportPageSize(resp *pb.ExportRangeVersionsResponse) error { + if migrationExportPagePayloadBytes(resp) <= migrationExportPageByteBudget { + return nil + } + widest := widestMigrationExportVersion(resp.GetVersions()) + return errors.WithStack(status.Errorf(codes.FailedPrecondition, + "migration export page exceeds the %d byte message limit; key %q holds %d value bytes", + internal.GRPCMaxMessageBytes, widest.GetKey(), len(widest.GetValue()))) +} + +// migrationExportPagePayloadBytes sums the caller-controlled bytes in a page. +// It deliberately undercounts the proto framing rather than calling +// proto.Size, which caches its result inside the message; the budget carries +// enough headroom to cover the framing of a full page. +func migrationExportPagePayloadBytes(resp *pb.ExportRangeVersionsResponse) int { + total := len(resp.GetNextCursor()) + for _, version := range resp.GetVersions() { + total += len(version.GetKey()) + len(version.GetValue()) + } + return total +} + +func widestMigrationExportVersion(versions []*pb.MVCCVersion) *pb.MVCCVersion { + var widest *pb.MVCCVersion + for _, version := range versions { + if widest == nil || len(version.GetValue()) > len(widest.GetValue()) { + widest = version + } + } + return widest +} + func (i *Internal) ImportRangeVersions(ctx context.Context, req *pb.ImportRangeVersionsRequest) (*pb.ImportRangeVersionsResponse, error) { if req == nil { return nil, errors.WithStack(status.Error(codes.InvalidArgument, "import range versions request is nil")) diff --git a/adapter/internal_migration_test.go b/adapter/internal_migration_test.go index e2c162c96..91e395deb 100644 --- a/adapter/internal_migration_test.go +++ b/adapter/internal_migration_test.go @@ -1,11 +1,13 @@ package adapter import ( + "bytes" "context" "encoding/binary" "testing" "github.com/bootjp/elastickv/distribution" + internalutil "github.com/bootjp/elastickv/internal" "github.com/bootjp/elastickv/internal/raftengine" "github.com/bootjp/elastickv/internal/s3keys" "github.com/bootjp/elastickv/kv" @@ -659,3 +661,34 @@ func TestInternalPromoteStagedVersionsAppliesStoreBatch(t *testing.T) { require.ErrorIs(t, err, store.ErrKeyNotFound) require.GreaterOrEqual(t, clock.Current(), uint64(30)) } + +// Storage accepts values well past the internal gRPC message limit, and the +// store's byte budget is only consulted after a version has been appended, so +// one oversized row can land in an otherwise bounded export page. Sending it +// fails with ResourceExhausted, and so does every retry of the same cursor: +// the bracket stalls forever on an error that never names the row. The export +// must refuse the page and say which key is too large. +func TestInternalExportRangeVersionsRefusesOversizedPage(t *testing.T) { + t.Parallel() + + ctx := context.Background() + st := store.NewMVCCStore() + huge := bytes.Repeat([]byte{'x'}, internalutil.GRPCMaxMessageBytes+1) + require.NoError(t, st.PutAt(ctx, []byte("a"), huge, 10, 0)) + srv := NewInternalWithEngine(nil, mockInternalLeader{}, nil, nil, WithInternalStore(st)) + stream := &captureExportRangeVersionsStream{ctx: ctx} + + err := srv.ExportRangeVersions(&pb.ExportRangeVersionsRequest{ + MaxCommitTs: 20, + RouteStart: []byte("a"), + RouteEnd: []byte("b"), + KeyFamily: distribution.MigrationFamilyUser, + RangeStart: []byte("a"), + RangeEnd: []byte("b"), + }, stream) + + require.Error(t, err) + require.Equal(t, codes.FailedPrecondition, status.Code(err)) + require.Contains(t, err.Error(), `"a"`) + require.Empty(t, stream.responses, "an unsendable page must not be sent") +} From f06fdf836f77f81516b74ff4afa00c217159e2ad Mon Sep 17 00:00:00 2001 From: bootjp Date: Sat, 29 Aug 2026 00:46:45 +0900 Subject: [PATCH 32/58] migration: split export pages before they overshoot the budget MaxBytes was only consulted after a row had been appended, so a page could overshoot its budget by the whole size of its last row: a 3 MiB page followed by a 61 MiB row built a 64 MiB page the transport refuses, and the same cursor rebuilt it on every retry. Both rows fit on their own. Stop the page before a row that would push a non-empty page past MaxBytes, leaving the cursor on the previous row so the oversized one starts the next page alone. A row too large even alone still goes out, because a page holding nothing has to make progress -- that case is what the transport guard reports by name. The memory path carries the decision out as an explicit deferred flag rather than re-running the filters: AcceptVersion is not required to be a pure predicate, and an existing test pins that it is called once per version. Claude-Session: https://claude.ai/code/session_013rNHooj7NF3giihWVba8QE --- store/lsm_migration.go | 54 ++++++++++++++++++------- store/migration_versions.go | 69 +++++++++++++++++++++++++++----- store/migration_versions_test.go | 57 ++++++++++++++++++++++++++ 3 files changed, 155 insertions(+), 25 deletions(-) diff --git a/store/lsm_migration.go b/store/lsm_migration.go index 247b13205..d5f599cbd 100644 --- a/store/lsm_migration.go +++ b/store/lsm_migration.go @@ -252,26 +252,17 @@ func (s *pebbleStore) exportPebbleVersion( commitTS uint64, result *ExportVersionsResult, ) (bool, error) { - tag := exportCursorTagScanned rawValue := iter.Value() result.ScannedBytes += versionExportSize(userKey, len(rawValue)) + tag := exportCursorTagScanned if shouldExportPebbleVersion(opts, userKey, commitTS) { - version, err := s.decodeExportedPebbleVersion(iter, userKey, commitTS, opts.KeyFamily) - if err != nil { + emitted, cont, err := s.appendExportedPebbleVersion(iter, opts, userKey, commitTS, result) + if err != nil || !cont { return false, err } - if opts.AcceptVersion != nil && !opts.AcceptVersion(version.Key, version.Value) { - result.NextCursor = encodeExportCursor(userKey, commitTS, exportCursorTagScanned) - if finishExportIfLimited(opts, result) { - result.Done = false - return false, nil - } - return true, nil + if emitted { + tag = exportCursorTagEmitted } - result.Versions = append(result.Versions, version) - result.ExportedBytes += versionExportSize(userKey, len(version.Value)) - result.AcceptedRows++ - tag = exportCursorTagEmitted } result.NextCursor = encodeExportCursor(userKey, commitTS, tag) if finishExportIfLimited(opts, result) { @@ -281,6 +272,41 @@ func (s *pebbleStore) exportPebbleVersion( return true, nil } +// appendExportedPebbleVersion decodes the version at the iterator and puts it +// on the page. It returns emitted=false when a filter rejected the version, +// and cont=false when the page had to stop: either because the row would push +// a page that already holds rows past its byte budget -- the cursor stays on +// the previous row so this one starts the next page -- or because a rejected +// version finished the chunk. +func (s *pebbleStore) appendExportedPebbleVersion( + iter *pebble.Iterator, + opts ExportVersionsOptions, + userKey []byte, + commitTS uint64, + result *ExportVersionsResult, +) (emitted bool, cont bool, err error) { + version, err := s.decodeExportedPebbleVersion(iter, userKey, commitTS, opts.KeyFamily) + if err != nil { + return false, false, err + } + if opts.AcceptVersion != nil && !opts.AcceptVersion(version.Key, version.Value) { + result.NextCursor = encodeExportCursor(userKey, commitTS, exportCursorTagScanned) + if finishExportIfLimited(opts, result) { + result.Done = false + return false, false, nil + } + return false, true, nil + } + if exportPageWouldOverflow(opts, result, userKey, len(version.Value)) { + result.Done = false + return false, false, nil + } + result.Versions = append(result.Versions, version) + result.ExportedBytes += versionExportSize(userKey, len(version.Value)) + result.AcceptedRows++ + return true, true, nil +} + func shouldExportPebbleVersion(opts ExportVersionsOptions, userKey []byte, commitTS uint64) bool { if shouldSkipMigrationExportKey(userKey) { return false diff --git a/store/migration_versions.go b/store/migration_versions.go index 08a6d1dd6..eb00802a8 100644 --- a/store/migration_versions.go +++ b/store/migration_versions.go @@ -510,24 +510,53 @@ func exportMemoryIteratorKey( return exportMemoryVersionsForKey(ctx, opts, cursorCommitTS, key, versions, result) } +// exportPageWouldOverflow reports whether adding one more row of valueLen +// bytes would push a page that already holds rows past its byte budget. +// MaxBytes is otherwise only consulted after a row has been appended, so a +// page can overshoot it by the whole size of its last row. That overshoot is +// what pushes an otherwise ordinary page past the migration transport limit -- +// a 3 MiB page followed by a 61 MiB row -- and the same cursor rebuilds the +// same page on every retry. Stopping first leaves the oversized row to start +// the next page, where it is alone and within the transport limit; a row too +// large even alone still goes through, because a page that holds nothing yet +// has to make progress. +func exportPageWouldOverflow(opts ExportVersionsOptions, result *ExportVersionsResult, key []byte, valueLen int) bool { + if opts.MaxBytes == 0 || len(result.Versions) == 0 { + return false + } + return result.ExportedBytes+versionExportSize(key, valueLen) > opts.MaxBytes +} + func finishExportIfLimited(opts ExportVersionsOptions, result *ExportVersionsResult) bool { return len(result.Versions) >= opts.MaxVersions || (opts.MaxBytes > 0 && result.ExportedBytes >= opts.MaxBytes) || (opts.MaxScannedBytes > 0 && result.ScannedBytes >= opts.MaxScannedBytes) } -func appendMemoryExportVersion(opts ExportVersionsOptions, key []byte, version VersionedValue, result *ExportVersionsResult) byte { +// appendMemoryExportVersion returns the cursor tag for this version and, as +// deferred, whether the page had to stop before it. AcceptVersion is called at +// most once per version here: it is not required to be a pure predicate, so +// the page-size guard sits after the filters rather than re-running them. +func appendMemoryExportVersion( + opts ExportVersionsOptions, + key []byte, + version VersionedValue, + result *ExportVersionsResult, +) (byte, bool) { if shouldSkipMigrationExportKey(key) { - return exportCursorTagScanned + return exportCursorTagScanned, false } if opts.AcceptKey != nil && !opts.AcceptKey(key) { - return exportCursorTagScanned + return exportCursorTagScanned, false } if opts.MaxCommitTSInclusive != 0 && version.TS > opts.MaxCommitTSInclusive { - return exportCursorTagScanned + return exportCursorTagScanned, false } if opts.AcceptVersion != nil && !opts.AcceptVersion(key, version.Value) { - return exportCursorTagScanned + return exportCursorTagScanned, false + } + if exportPageWouldOverflow(opts, result, key, len(version.Value)) { + return exportCursorTagScanned, true } result.Versions = append(result.Versions, MVCCVersion{ Key: bytes.Clone(key), @@ -539,7 +568,7 @@ func appendMemoryExportVersion(opts ExportVersionsOptions, key []byte, version V }) result.ExportedBytes += versionExportSize(key, len(version.Value)) result.AcceptedRows++ - return exportCursorTagEmitted + return exportCursorTagEmitted, false } func shouldSkipMigrationExportKey(key []byte) bool { @@ -560,12 +589,23 @@ func shouldSkipMemoryVersion(cursorCommitTS uint64, version VersionedValue) bool return cursorCommitTS != 0 && version.TS >= cursorCommitTS } -func exportMemoryVersion(opts ExportVersionsOptions, cursorCommitTS uint64, key []byte, version VersionedValue, result *ExportVersionsResult) bool { +// exportMemoryVersion returns (continue, deferred): deferred means the page is +// full for this version and it must start the next one. +func exportMemoryVersion( + opts ExportVersionsOptions, + cursorCommitTS uint64, + key []byte, + version VersionedValue, + result *ExportVersionsResult, +) (bool, bool) { if shouldSkipMemoryVersion(cursorCommitTS, version) { - return true + return true, false } - tag := appendMemoryExportVersion(opts, key, version, result) - return finishMemoryExportPosition(opts, key, version, tag, result) + tag, deferred := appendMemoryExportVersion(opts, key, version, result) + if deferred { + return false, true + } + return finishMemoryExportPosition(opts, key, version, tag, result), false } func exportMemoryVersionsForKey( @@ -589,7 +629,14 @@ func exportMemoryVersionsForKey( } return true, nil } - if !exportMemoryVersion(opts, cursorCommitTS, key, versions[i], result) { + cont, deferred := exportMemoryVersion(opts, cursorCommitTS, key, versions[i], result) + if deferred { + // The page is full for this row. The cursor still points at the row + // before it, so the next page starts here and carries it alone. + result.Done = false + return false, nil + } + if !cont { return !finishExportIfLimited(opts, result), nil } if !result.Done && finishExportIfLimited(opts, result) { diff --git a/store/migration_versions_test.go b/store/migration_versions_test.go index f4680a0ff..05bb1aaf8 100644 --- a/store/migration_versions_test.go +++ b/store/migration_versions_test.go @@ -1342,3 +1342,60 @@ func TestPebbleRestoreStreamingSnapshotPreservesMigrationPromotionState(t *testi require.Equal(t, uint64(2), restored.TotalPromotedRows) require.Equal(t, uint64(110), restored.MaxPromotedTS) } + +// MaxBytes is checked after a row is appended, so a page can overshoot its +// budget by the whole size of its last row. A 3 MiB page followed by a 61 MiB +// row therefore builds a 64 MiB page that the migration transport refuses -- +// and the same cursor rebuilds it on every retry. Each row fits on its own, so +// the page must stop before the row that would overshoot. +func TestExportVersionsSplitsBeforeOverflowingTheByteBudget(t *testing.T) { + runMigrationStoreSuite(t, func(t *testing.T, st MVCCStore) { + ctx := context.Background() + + small := bytes.Repeat([]byte("s"), 1<<10) + large := bytes.Repeat([]byte("l"), 8<<10) + require.NoError(t, st.PutAt(ctx, []byte("a"), small, 10, 0)) + require.NoError(t, st.PutAt(ctx, []byte("b"), large, 11, 0)) + + opts := ExportVersionsOptions{ + MaxVersions: 16, + MaxBytes: 4 << 10, + MaxCommitTSInclusive: 100, + EndKey: []byte("z"), + } + first, err := st.ExportVersions(ctx, opts) + require.NoError(t, err) + require.False(t, first.Done) + require.Len(t, first.Versions, 1, "the oversized row must not join this page") + require.Equal(t, []byte("a"), first.Versions[0].Key) + require.LessOrEqual(t, first.ExportedBytes, opts.MaxBytes, "the page stays inside its budget") + + opts.Cursor = first.NextCursor + second, err := st.ExportVersions(ctx, opts) + require.NoError(t, err) + require.Len(t, second.Versions, 1, "the oversized row goes out alone") + require.Equal(t, []byte("b"), second.Versions[0].Key) + require.Equal(t, large, second.Versions[0].Value) + }) +} + +// A single row larger than the whole budget still goes out: a page holding +// nothing has to make progress, or the cursor never advances. +func TestExportVersionsEmitsSingleOversizedRow(t *testing.T) { + runMigrationStoreSuite(t, func(t *testing.T, st MVCCStore) { + ctx := context.Background() + + huge := bytes.Repeat([]byte("h"), 8<<10) + require.NoError(t, st.PutAt(ctx, []byte("a"), huge, 10, 0)) + + got, err := st.ExportVersions(ctx, ExportVersionsOptions{ + MaxVersions: 16, + MaxBytes: 1 << 10, + MaxCommitTSInclusive: 100, + EndKey: []byte("z"), + }) + require.NoError(t, err) + require.Len(t, got.Versions, 1) + require.Equal(t, huge, got.Versions[0].Value) + }) +} From d0dea91d8a3e7115768a9630c5e0af8ca2f1b723 Mon Sep 17 00:00:00 2001 From: bootjp Date: Sat, 29 Aug 2026 10:08:47 +0900 Subject: [PATCH 33/58] kv: reject reserved control keys on the transactional paths validateRawMutationForApply refuses !dist| / !migstage| / !migwrite| / !migfence| writes, but the transactional apply paths build store mutations through their own helpers and never ran that check. A TransactionalKV request could therefore write catalog or staged migration state directly, and when the same group is also the migration target a forged !dist|migstage|| row is later promoted as user data. Reject them in uniqueMutationsAboveFloor and uniqueTxnMutationsAboveFloor, which between them cover PREPARE, one-phase, and COMMIT. Caller audit: those two helpers have exactly three production call sites (fsm.go:1262, :1332, :1491), all on the transactional apply path; the typed internal commands that legitimately write these namespaces -- catalog apply, migration import and promote -- do not pass through either. Claude-Session: https://claude.ai/code/session_013rNHooj7NF3giihWVba8QE --- kv/fsm.go | 28 +++++++++++++++++++++ kv/fsm_reserved_control_test.go | 44 +++++++++++++++++++++++++++++++++ 2 files changed, 72 insertions(+) diff --git a/kv/fsm.go b/kv/fsm.go index 2052991f7..feb7fa8f7 100644 --- a/kv/fsm.go +++ b/kv/fsm.go @@ -1358,17 +1358,45 @@ func (f *kvFSM) uniqueMutationsAboveFloor(muts []*pb.Mutation, commitTS uint64) if err != nil { return nil, err } + if err := rejectReservedControlMutations(uniq); err != nil { + return nil, err + } if err := f.verifyRouteWriteTimestampFloorsForMutations(uniq, commitTS); err != nil { return nil, err } return uniq, nil } +// rejectReservedControlMutations keeps the catalog and migration control +// namespaces out of client writes. validateRawMutationForApply already refuses +// them on the raw path; the transactional paths reached the store without the +// check, so a TransactionalKV request could write !dist| or !migstage| state +// directly -- and a forged !dist|migstage|| row is promoted as user data +// when the same group is a migration target. +// +// The verdict is a pure function of the key, so every replica reaches it for +// the same entry. Typed internal commands (catalog apply, migration import and +// promote) do not pass through here and keep writing these namespaces. +func rejectReservedControlMutations(muts []*pb.Mutation) error { + for _, mut := range muts { + if mut == nil { + continue + } + if distribution.IsReservedControlKey(mut.Key) { + return errors.WithStack(ErrInvalidRequest) + } + } + return 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 := rejectReservedControlMutations(uniq); err != nil { + return nil, err + } if err := f.verifyRouteWriteTimestampFloorsForMutations(uniq, commitTS); err != nil { return nil, err } diff --git a/kv/fsm_reserved_control_test.go b/kv/fsm_reserved_control_test.go index 1b7a425f1..9d008293d 100644 --- a/kv/fsm_reserved_control_test.go +++ b/kv/fsm_reserved_control_test.go @@ -88,3 +88,47 @@ func TestValidateRawMutationRejectsStagedDataKeys(t *testing.T) { &pb.Mutation{Op: pb.Op_PUT, Key: staged, Value: []byte("v")}, nil, 10), ErrInvalidRequest) require.ErrorIs(t, f.handleDelPrefix(ctx, distribution.MigrationStagedDataKey(7, []byte("user:")), 11), ErrInvalidRequest) } + +// The transactional paths reach the store through their own helpers, which did +// not carry the reserved-key check the raw path has. A TransactionalKV request +// could therefore write catalog or staged-migration state directly, and a +// forged !migstage|| row is promoted as user data when the same group is +// the migration target. +func TestTxnMutationHelpersRejectReservedControlKeys(t *testing.T) { + t.Parallel() + + st := store.NewMVCCStore() + t.Cleanup(func() { _ = st.Close() }) + f, ok := NewKvFSMWithHLC(st, NewHLC()).(*kvFSM) + require.True(t, ok) + + for _, key := range [][]byte{ + []byte("!dist|meta|version"), + []byte("!dist|route|0001"), + distribution.MigrationStagedDataKey(7, []byte("victim")), + []byte("!migwrite|7"), + []byte("!migfence|7"), + } { + muts := []*pb.Mutation{{Op: pb.Op_PUT, Key: key, Value: []byte("v")}} + _, err := f.uniqueMutationsAboveFloor(muts, 10) + require.ErrorIs(t, err, ErrInvalidRequest, + "prepare/one-phase must refuse %q", key) + _, err = f.uniqueTxnMutationsAboveFloor(muts, 10) + require.ErrorIs(t, err, ErrInvalidRequest, + "commit must refuse %q", key) + } + + // Ordinary user keys, and the transaction-internal keys the txn paths + // legitimately write, are unaffected. + for _, key := range [][]byte{ + []byte("user-key"), + txnLockKey([]byte("user-key")), + txnIntentKey([]byte("user-key")), + } { + muts := []*pb.Mutation{{Op: pb.Op_PUT, Key: key, Value: []byte("v")}} + _, err := f.uniqueMutationsAboveFloor(muts, 10) + require.NoError(t, err, "key %q must be accepted", key) + _, err = f.uniqueTxnMutationsAboveFloor(muts, 10) + require.NoError(t, err, "key %q must be accepted", key) + } +} From 76e6b2146b2278e936ebd7c0e4962bff3d1c7b8e Mon Sep 17 00:00:00 2001 From: bootjp Date: Sat, 29 Aug 2026 10:08:48 +0900 Subject: [PATCH 34/58] migration: order the staged and live probes against promotion getAtWithStagedVisibility makes two separate store reads, so a promotion batch can land between them. Reading live first misses both sides of a staged-only key -- live has not been written yet at the first probe, and the staged alias is gone by the second -- and can return a stale live value when the newer staged version is promoted in between. Read staged first, then live. Promotion only moves a row staged -> live, so the live set only gains keys and the staged set only loses them; reading the shrinking side first and the growing side second means a key present in either at any instant during the pair is observed by at least one probe. Claude-Session: https://claude.ai/code/session_013rNHooj7NF3giihWVba8QE --- kv/shard_store.go | 17 +++++++++--- kv/shard_store_test.go | 59 ++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 73 insertions(+), 3 deletions(-) diff --git a/kv/shard_store.go b/kv/shard_store.go index 9d1b2f0b6..ec8c4fa42 100644 --- a/kv/shard_store.go +++ b/kv/shard_store.go @@ -538,12 +538,23 @@ func (s *ShardStore) getAtWithStagedVisibility(ctx context.Context, g *ShardGrou if err := ensureReadTSRetained(g.Store, ts); err != nil { return nil, err } - live, liveOK, err := latestMVCCVersionAt(ctx, g.Store, key, ts) + // Staged first, then live. The two probes are separate store reads, so a + // promotion batch can land between them, and the order decides whether that + // is observable. Promotion only ever moves a row staged -> live: the live + // set only gains keys, the staged set only loses them. Reading the shrinking + // side first and the growing side second means a key present in either at + // any instant during the pair is seen by at least one probe. + // + // The reverse order has two holes. A staged-only key read live-then-staged + // misses both sides and returns not-found for a key that existed + // throughout; and a key whose staged version is newer than its live one can + // return the stale live value when promotion removes the alias in between. + stagedKey := distribution.MigrationStagedDataKey(route.MigrationJobID, key) + staged, stagedOK, err := latestMVCCVersionAt(ctx, g.Store, stagedKey, ts) if err != nil { return nil, err } - stagedKey := distribution.MigrationStagedDataKey(route.MigrationJobID, key) - staged, stagedOK, err := latestMVCCVersionAt(ctx, g.Store, stagedKey, ts) + live, liveOK, err := latestMVCCVersionAt(ctx, g.Store, key, ts) if err != nil { return nil, err } diff --git a/kv/shard_store_test.go b/kv/shard_store_test.go index e76ce19a8..e705a4bd5 100644 --- a/kv/shard_store_test.go +++ b/kv/shard_store_test.go @@ -3898,3 +3898,62 @@ func TestExpireAtStillFailsWhenNothingIsVisible(t *testing.T) { require.ErrorIs(t, st.ExpireAt(ctx, []byte("absent"), 40, 300), store.ErrKeyNotFound) } + +// promotingExportStore runs a hook after the first ExportVersions call, which +// is how a promotion batch is landed exactly between the two probes +// getAtWithStagedVisibility makes. +type promotingExportStore struct { + store.MVCCStore + afterFirst func() + calls int +} + +func (s *promotingExportStore) ExportVersions( + ctx context.Context, + opts store.ExportVersionsOptions, +) (store.ExportVersionsResult, error) { + res, err := s.MVCCStore.ExportVersions(ctx, opts) + s.calls++ + if s.calls == 1 && s.afterFirst != nil { + s.afterFirst() + } + return res, err +} + +// A promotion batch landing between the staged and live probes must not make a +// key disappear. Promotion writes the live version and drops the staged alias, +// so a live-then-staged order misses both sides of a staged-only key: live has +// not been written yet at the first probe, and the alias is gone by the second. +func TestShardStoreGetAt_StagedVisibilitySurvivesPromotionBetweenProbes(t *testing.T) { + t.Parallel() + + ctx := context.Background() + 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, + }}, + })) + + inner := store.NewMVCCStore() + t.Cleanup(func() { _ = inner.Close() }) + rawKey := []byte("k") + stagedKey := distribution.MigrationStagedDataKey(9, rawKey) + require.NoError(t, inner.PutAt(ctx, stagedKey, []byte("staged-only"), 20, 0)) + + promoting := &promotingExportStore{MVCCStore: inner} + promoting.afterFirst = func() { + // The promotion batch: the row becomes live at its original commit ts + // and the staged alias goes away. + require.NoError(t, inner.PutAt(ctx, rawKey, []byte("staged-only"), 20, 0)) + require.NoError(t, inner.DeleteAt(ctx, stagedKey, 21)) + } + st := NewShardStore(engine, map[uint64]*ShardGroup{1: {Store: promoting}}) + + got, err := st.GetAt(ctx, rawKey, 25) + require.NoError(t, err, "a promotion between the probes must not hide the key") + require.Equal(t, []byte("staged-only"), got) +} From 58104e060cc935c955618acaf035d74984a137e2 Mon Sep 17 00:00:00 2001 From: bootjp Date: Sat, 29 Aug 2026 10:25:20 +0900 Subject: [PATCH 35/58] kv: narrow the transactional reserved-key test to migration-only keys Refusing every reserved control namespace on the transactional paths broke the control plane: SplitRange commits catalog route and meta records through the coordinator, so those writes are transactions and the blanket check turned them into "invalid request". TestMilestone1SplitRange_EndToEndRefreshAndDataPath caught it. Reject only the namespaces no coordinator dispatch ever writes -- !dist|migstage|, !migwrite|, !migfence| -- which still closes the case the finding is about: a client write landing under a staged-data key is promoted as user data when the same group is the migration target. The raw path keeps the wider test, because nothing legitimate writes any control namespace as a raw mutation. Claude-Session: https://claude.ai/code/session_013rNHooj7NF3giihWVba8QE --- distribution/migrator.go | 27 +++++++++++++++++++++++++++ kv/fsm.go | 20 +++++++++++--------- kv/fsm_reserved_control_test.go | 11 +++++++---- 3 files changed, 45 insertions(+), 13 deletions(-) diff --git a/distribution/migrator.go b/distribution/migrator.go index 3bc5c7dd9..65d4aa484 100644 --- a/distribution/migrator.go +++ b/distribution/migrator.go @@ -108,6 +108,33 @@ var ( // Staged migration data is included in the reserved set. It is populated by the // typed migration import/promote paths and by FSM-internal prefix-delete // expansion, not by externally supplied RawKV mutations. +// migrationOnlyControlPrefixes are the control namespaces that no coordinator +// dispatch ever writes: they are produced solely by the typed migration +// commands (import, promote, fence). The rest of !dist| is excluded, because +// the control plane commits route and job records through the transactional +// coordinator -- SplitRange does exactly that. +var migrationOnlyControlPrefixes = [][]byte{ + []byte(migrationStagedDataPrefix), + []byte("!migwrite|"), + []byte("!migfence|"), +} + +// IsMigrationOnlyControlKey reports whether key belongs to a namespace that +// only the typed migration commands may write. It is the transactional apply +// path's reserved-key test: the raw path can refuse every control namespace +// because nothing legitimate writes one as a raw mutation, but the catalog +// reaches the store through a transaction, so the transactional test has to be +// the narrower one. A staged-data row is the case that matters -- a client +// write landing there is promoted as user data. +func IsMigrationOnlyControlKey(key []byte) bool { + for _, prefix := range migrationOnlyControlPrefixes { + if bytes.HasPrefix(key, prefix) { + return true + } + } + return false +} + func IsReservedControlKey(key []byte) bool { return reservedControlPrefixIntersects(key, bytes.HasPrefix) } diff --git a/kv/fsm.go b/kv/fsm.go index feb7fa8f7..82c481f60 100644 --- a/kv/fsm.go +++ b/kv/fsm.go @@ -1367,22 +1367,24 @@ func (f *kvFSM) uniqueMutationsAboveFloor(muts []*pb.Mutation, commitTS uint64) return uniq, nil } -// rejectReservedControlMutations keeps the catalog and migration control -// namespaces out of client writes. validateRawMutationForApply already refuses -// them on the raw path; the transactional paths reached the store without the -// check, so a TransactionalKV request could write !dist| or !migstage| state -// directly -- and a forged !dist|migstage|| row is promoted as user data -// when the same group is a migration target. +// rejectReservedControlMutations keeps the migration-only control namespaces +// out of transactional writes. validateRawMutationForApply refuses every +// control namespace on the raw path, but the transactional paths cannot: the +// control plane commits catalog route and job records through the coordinator, +// so SplitRange itself is a transaction writing !dist|route| and !dist|meta|. +// The narrower test still closes the case that matters -- a client write +// landing under !dist|migstage|| is promoted as user data when the same +// group is a migration target. // // The verdict is a pure function of the key, so every replica reaches it for -// the same entry. Typed internal commands (catalog apply, migration import and -// promote) do not pass through here and keep writing these namespaces. +// the same entry. The typed migration commands (import, promote, fence) do not +// pass through here and keep writing these namespaces. func rejectReservedControlMutations(muts []*pb.Mutation) error { for _, mut := range muts { if mut == nil { continue } - if distribution.IsReservedControlKey(mut.Key) { + if distribution.IsMigrationOnlyControlKey(mut.Key) { return errors.WithStack(ErrInvalidRequest) } } diff --git a/kv/fsm_reserved_control_test.go b/kv/fsm_reserved_control_test.go index 9d008293d..bedabdebf 100644 --- a/kv/fsm_reserved_control_test.go +++ b/kv/fsm_reserved_control_test.go @@ -103,8 +103,6 @@ func TestTxnMutationHelpersRejectReservedControlKeys(t *testing.T) { require.True(t, ok) for _, key := range [][]byte{ - []byte("!dist|meta|version"), - []byte("!dist|route|0001"), distribution.MigrationStagedDataKey(7, []byte("victim")), []byte("!migwrite|7"), []byte("!migfence|7"), @@ -118,12 +116,17 @@ func TestTxnMutationHelpersRejectReservedControlKeys(t *testing.T) { "commit must refuse %q", key) } - // Ordinary user keys, and the transaction-internal keys the txn paths - // legitimately write, are unaffected. + // Ordinary user keys, the transaction-internal keys the txn paths + // legitimately write, and the catalog records the control plane commits + // through the coordinator are unaffected. SplitRange is a transaction that + // writes !dist|route| and !dist|meta|, so refusing those here would break + // the control plane itself. for _, key := range [][]byte{ []byte("user-key"), txnLockKey([]byte("user-key")), txnIntentKey([]byte("user-key")), + []byte("!dist|route|0001"), + []byte("!dist|meta|version"), } { muts := []*pb.Mutation{{Op: pb.Op_PUT, Key: key, Value: []byte("v")}} _, err := f.uniqueMutationsAboveFloor(muts, 10) From ff2918873296c868ff346bde33b6ec88b1b587c7 Mon Sep 17 00:00:00 2001 From: bootjp Date: Sat, 29 Aug 2026 11:16:07 +0900 Subject: [PATCH 36/58] migration: read staged before live everywhere, not just the point read Reordering only getAtWithStagedVisibility left four other places reading the two namespaces live-first, each losing the key a different way when a promotion batch lands between the pair: - the candidate scan page, where a live-then-staged pair misses a staged-only key entirely and the scan returns without it; - latestStagedVisibilityCandidates, which can resolve an older live value; - expireStagedVisibleAt, where the miss copies the older live payload into a fresh version at commitTS -- a permanent rollback rather than a TTL; - localLatestCommitTS and the S3 auxiliary owner existence probe, which report a missing or stale watermark. All five now read staged first. The argument is the same in each: promotion only moves a row staged -> live, so the live set only gains keys and the staged set only loses them, and reading the shrinking side first keeps a row that is mid-promotion visible to one of the two reads. The test drives the real PromoteVersions between the two scans rather than modelling promotion as a Delete -- a delete writes a tombstone above the promoted row and legitimately hides the key, which is a different scenario and made the first version of this test fail for the wrong reason. Claude-Session: https://claude.ai/code/session_013rNHooj7NF3giihWVba8QE --- kv/shard_store.go | 74 ++++++++++++++++++++++------------ kv/shard_store_test.go | 90 +++++++++++++++++++++++++++++++++++++++++- 2 files changed, 138 insertions(+), 26 deletions(-) diff --git a/kv/shard_store.go b/kv/shard_store.go index ec8c4fa42..be6028788 100644 --- a/kv/shard_store.go +++ b/kv/shard_store.go @@ -2327,19 +2327,17 @@ func (s *ShardStore) s3BucketAuxiliaryOwnerHasVersionAt( ts uint64, readRouteVersion uint64, ) (bool, error) { - live, liveOK, err := s.ownerRouteHasVersionAtOrBefore(ctx, owner, key, ts, readRouteVersion) - if err != nil { - return false, err - } - if live { - return true, nil - } if !routeHasStagedVisibility(owner) { - if !liveOK { - return false, ownerVersionProbeUnavailable(owner, key) + live, liveOK, err := s.ownerRouteHasVersionAtOrBefore(ctx, owner, key, ts, readRouteVersion) + if err != nil { + return false, err } - return false, nil + return existenceProbeResult(live, liveOK, owner, key) } + // Staged before live: a live-first existence probe can miss a row that + // promotion moves between the two calls -- absent from live before the + // move, absent from staged after it -- and report a key that existed + // throughout as missing. stagedKey := distribution.MigrationStagedDataKey(owner.MigrationJobID, key) staged, stagedOK, err := s.ownerRouteHasVersionAtOrBefore(ctx, owner, stagedKey, ts, readRouteVersion) if err != nil { @@ -2348,7 +2346,21 @@ func (s *ShardStore) s3BucketAuxiliaryOwnerHasVersionAt( if staged { return true, nil } - if !liveOK || !stagedOK { + live, liveOK, err := s.ownerRouteHasVersionAtOrBefore(ctx, owner, key, ts, readRouteVersion) + if err != nil { + return false, err + } + return existenceProbeResult(live, liveOK && stagedOK, owner, key) +} + +// existenceProbeResult turns a probe outcome into the (found, error) pair the +// callers expect: an unanswered probe is an error rather than a "no", because +// treating an unreachable owner as absent is what silently drops a row. +func existenceProbeResult(found bool, answered bool, owner distribution.Route, key []byte) (bool, error) { + if found { + return true, nil + } + if !answered { return false, ownerVersionProbeUnavailable(owner, key) } return false, nil @@ -3109,11 +3121,15 @@ func (s *ShardStore) scanRouteWithStagedVisibilityPage( stagedStart, stagedEnd := stagedVisibilityScanBounds(route.MigrationJobID, start, end) window := stagedVisibilityCandidateWindow(limit) for { - liveKVs, err := scanVisibleCandidates(ctx, g.Store, start, end, window, ts, reverse) + // Staged before live, for the reason on getAtWithStagedVisibility: + // promotion only moves rows staged -> live, so reading the shrinking + // side first keeps a key that is mid-promotion visible to one of the + // two scans. + stagedKVs, err := scanVisibleCandidates(ctx, g.Store, stagedStart, stagedEnd, window, ts, reverse) if err != nil { return nil, nil, false, err } - stagedKVs, err := scanVisibleCandidates(ctx, g.Store, stagedStart, stagedEnd, window, ts, reverse) + liveKVs, err := scanVisibleCandidates(ctx, g.Store, start, end, window, ts, reverse) if err != nil { return nil, nil, false, err } @@ -3318,18 +3334,19 @@ func (s *ShardStore) latestStagedVisibilityCandidates( ts uint64, ) (map[string]store.MVCCVersion, error) { keys := stagedVisibilityCandidateKeys(liveKVs, stagedKVs) - liveVersions, err := latestCandidateVersionsAt(ctx, st, keys, ts) - if err != nil { - return nil, err - } stagedKeys := make([][]byte, 0, len(keys)) for _, key := range keys { stagedKeys = append(stagedKeys, distribution.MigrationStagedDataKey(route.MigrationJobID, key)) } + // Staged before live: same promotion ordering as the point read. stagedVersions, err := latestCandidateVersionsAt(ctx, st, stagedKeys, ts) if err != nil { return nil, err } + liveVersions, err := latestCandidateVersionsAt(ctx, st, keys, ts) + if err != nil { + return nil, err + } out := make(map[string]store.MVCCVersion, len(keys)) for _, key := range keys { live, liveOK := liveVersions[string(key)] @@ -3862,12 +3879,16 @@ func (s *ShardStore) expireStagedVisibleAt( expireAt uint64, commitTS uint64, ) error { - live, liveOK, err := latestMVCCVersionAt(ctx, g.Store, key, commitTS) + // Staged before live. Reading live first lets a promotion land in between + // and leave the staged probe empty, so the older live payload would be + // copied into a fresh version at commitTS -- a permanent rollback of the + // value rather than a TTL on the current one. + stagedKey := distribution.MigrationStagedDataKey(route.MigrationJobID, key) + staged, stagedOK, err := latestMVCCVersionAt(ctx, g.Store, stagedKey, commitTS) if err != nil { return err } - stagedKey := distribution.MigrationStagedDataKey(route.MigrationJobID, key) - staged, stagedOK, err := latestMVCCVersionAt(ctx, g.Store, stagedKey, commitTS) + live, liveOK, err := latestMVCCVersionAt(ctx, g.Store, key, commitTS) if err != nil { return err } @@ -3947,17 +3968,20 @@ func (s *ShardStore) LatestCommitTSGroupWithReadFence(ctx context.Context, key [ } 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 + liveTS, liveExists, err := g.Store.LatestCommitTS(ctx, key) + return liveTS, liveExists, errors.WithStack(err) } + // Staged before live: same promotion ordering as every other pair of + // probes across the two namespaces. stagedTS, stagedExists, err := g.Store.LatestCommitTS(ctx, distribution.MigrationStagedDataKey(route.MigrationJobID, key)) if err != nil { return 0, false, errors.WithStack(err) } + liveTS, liveExists, err := g.Store.LatestCommitTS(ctx, key) + if err != nil { + return 0, false, errors.WithStack(err) + } switch { case !liveExists: return stagedTS, stagedExists, nil diff --git a/kv/shard_store_test.go b/kv/shard_store_test.go index e705a4bd5..78a2321cf 100644 --- a/kv/shard_store_test.go +++ b/kv/shard_store_test.go @@ -3913,11 +3913,32 @@ func (s *promotingExportStore) ExportVersions( opts store.ExportVersionsOptions, ) (store.ExportVersionsResult, error) { res, err := s.MVCCStore.ExportVersions(ctx, opts) + s.fireAfterFirst() + return res, err +} + +// ScanAt fires the promotion when the *live* range is scanned. That is the +// only interleaving that distinguishes the two orderings: live-first means the +// live scan misses the key and the staged scan that follows misses it too, +// while staged-first has already captured it before promotion runs. +func (s *promotingExportStore) ScanAt( + ctx context.Context, + start, end []byte, + limit int, + ts uint64, +) ([]*store.KVPair, error) { + kvs, err := s.MVCCStore.ScanAt(ctx, start, end, limit, ts) + if !isMigrationStagedDataKey(start) && s.afterFirst != nil { + s.fireAfterFirst() + } + return kvs, err +} + +func (s *promotingExportStore) fireAfterFirst() { s.calls++ if s.calls == 1 && s.afterFirst != nil { s.afterFirst() } - return res, err } // A promotion batch landing between the staged and live probes must not make a @@ -3957,3 +3978,70 @@ func TestShardStoreGetAt_StagedVisibilitySurvivesPromotionBetweenProbes(t *testi require.NoError(t, err, "a promotion between the probes must not hide the key") require.Equal(t, []byte("staged-only"), got) } + +// Every place that reads the live and staged namespaces as two separate store +// calls has to read staged first, for the reason getAtWithStagedVisibility +// documents. Fixing only the point read left the scan, the TTL winner, and the +// watermark on the old order, each with its own way of losing the key. +func TestStagedVisibilityProbesReadStagedFirst(t *testing.T) { + t.Parallel() + + ctx := context.Background() + engine := distribution.NewEngine() + require.NoError(t, engine.ApplySnapshot(distribution.CatalogSnapshot{ + Version: 1, + Routes: []distribution.RouteDescriptor{{ + RouteID: 1, Start: []byte("a"), End: []byte("z"), GroupID: 1, + State: distribution.RouteStateActive, StagedVisibilityActive: true, + MigrationJobID: 9, MinWriteTSExclusive: 100, + }}, + })) + + rawKey := []byte("k") + stagedKey := distribution.MigrationStagedDataKey(9, rawKey) + + newStore := func(t *testing.T) (*ShardStore, *promotingExportStore, store.MVCCStore) { + t.Helper() + inner := store.NewMVCCStore() + t.Cleanup(func() { _ = inner.Close() }) + require.NoError(t, inner.PutAt(ctx, stagedKey, []byte("staged-only"), 20, 0)) + promoting := &promotingExportStore{MVCCStore: inner} + promoting.afterFirst = func() { + // The real promotion batch: PromoteVersions moves the staged row to + // its live key and removes the staged version physically. Modelling + // it with a Delete would write a tombstone that legitimately hides + // the key, which is a different scenario. + promoter, ok := inner.(store.MigrationPromoter) + require.True(t, ok) + _, err := promoter.PromoteVersions(ctx, store.PromoteVersionsOptions{ + JobID: 9, + StartKey: distribution.MigrationStagedDataKeyPrefix(9), + EndKey: prefixScanEnd(distribution.MigrationStagedDataKeyPrefix(9)), + MaxVersions: 16, + TargetKey: func(staged []byte) ([]byte, bool) { + _, raw, ok := distribution.MigrationStagedDataKeyParts(staged) + return raw, ok + }, + }) + require.NoError(t, err) + } + return NewShardStore(engine, map[uint64]*ShardGroup{1: {Store: promoting}}), promoting, inner + } + + t.Run("scan", func(t *testing.T) { + t.Parallel() + st, _, _ := newStore(t) + kvs, err := st.ScanAt(ctx, []byte("a"), []byte("z"), 10, 25) + require.NoError(t, err) + require.Len(t, kvs, 1, "a promotion between the scans must not drop the key") + require.Equal(t, rawKey, kvs[0].Key) + }) + + t.Run("point read", func(t *testing.T) { + t.Parallel() + st, _, _ := newStore(t) + got, err := st.GetAt(ctx, rawKey, 25) + require.NoError(t, err) + require.Equal(t, []byte("staged-only"), got) + }) +} From d9892936b59a12abd940c0e07f0ac3e7fff0e03e Mon Sep 17 00:00:00 2001 From: bootjp Date: Sat, 29 Aug 2026 11:18:22 +0900 Subject: [PATCH 37/58] kv: expect the staged owner probe first The S3 auxiliary owner probe now asks staged before live and a staged hit answers the question outright, so the second request the old expectation counted is no longer issued. The test pins the new order rather than the count. Claude-Session: https://claude.ai/code/session_013rNHooj7NF3giihWVba8QE --- kv/shard_store_test.go | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/kv/shard_store_test.go b/kv/shard_store_test.go index 78a2321cf..d9361b4b8 100644 --- a/kv/shard_store_test.go +++ b/kv/shard_store_test.go @@ -720,9 +720,12 @@ func TestShardStoreS3BucketAuxiliaryOwnerProbeUsesLeaderRoutedReadFence(t *testi probe.mu.Lock() defer probe.mu.Unlock() - require.Len(t, probe.latestReqs, 2) - require.Equal(t, migratedKey, probe.latestReqs[0].GetKey()) - require.Equal(t, stagedKey, probe.latestReqs[1].GetKey()) + // Staged is probed first, and a staged hit answers the question, so the + // live probe is never issued. Probing live first would have to fall + // through to staged anyway, and that pair is what a concurrent promotion + // slips between. + require.Len(t, probe.latestReqs, 1) + require.Equal(t, stagedKey, probe.latestReqs[0].GetKey()) for _, req := range probe.latestReqs { require.Equal(t, uint64(2), req.GetGroupId()) require.Equal(t, uint64(77), req.GetReadRouteVersion()) From 2fb90adbf23863022d45250898ed2743189a918c Mon Sep 17 00:00:00 2001 From: bootjp Date: Sat, 29 Aug 2026 15:06:16 +0900 Subject: [PATCH 38/58] migration: read staged before live in OCC read validation too latestCommitTSForReadKeyOnShard was the last live-first pair, and it is the one with the worst outcome: a promotion landing between the two probes hides the version from both, so validateReadKeysOnShard sees no commit above startTS and lets the transaction commit on a read that should have conflicted. Swept the repository for the pattern this time rather than one file -- kv/shard_store.go and kv/sharded_coordinator.go were the only two holding paired probes, and all of them now read the shrinking staged namespace first. Claude-Session: https://claude.ai/code/session_013rNHooj7NF3giihWVba8QE --- kv/shard_store_test.go | 10 +++++ kv/sharded_coordinator.go | 15 +++++--- kv/sharded_coordinator_test.go | 69 ++++++++++++++++++++++++++++++++++ 3 files changed, 89 insertions(+), 5 deletions(-) create mode 100644 kv/sharded_coordinator_test.go diff --git a/kv/shard_store_test.go b/kv/shard_store_test.go index d9361b4b8..35b3f90e3 100644 --- a/kv/shard_store_test.go +++ b/kv/shard_store_test.go @@ -3937,6 +3937,16 @@ func (s *promotingExportStore) ScanAt( return kvs, err } +// LatestCommitTS fires the promotion on the live-key probe, which is the +// interleaving that separates the two orderings for the watermark reads. +func (s *promotingExportStore) LatestCommitTS(ctx context.Context, key []byte) (uint64, bool, error) { + ts, exists, err := s.MVCCStore.LatestCommitTS(ctx, key) + if !isMigrationStagedDataKey(key) && s.afterFirst != nil { + s.fireAfterFirst() + } + return ts, exists, err +} + func (s *promotingExportStore) fireAfterFirst() { s.calls++ if s.calls == 1 && s.afterFirst != nil { diff --git a/kv/sharded_coordinator.go b/kv/sharded_coordinator.go index d753acd82..679aee24c 100644 --- a/kv/sharded_coordinator.go +++ b/kv/sharded_coordinator.go @@ -2592,18 +2592,23 @@ func (c *ShardedCoordinator) validateReadKeysOnShard(ctx context.Context, gid ui } 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 + liveTS, liveExists, err := g.Store.LatestCommitTS(ctx, key) + return liveTS, liveExists, errors.WithStack(err) } + // Staged before live, for the reason on getAtWithStagedVisibility. + // Missing the version here is worse than a stale read: OCC would see no + // commit above startTS and let the transaction commit on a read it should + // have conflicted with. stagedTS, stagedExists, err := g.Store.LatestCommitTS(ctx, distribution.MigrationStagedDataKey(route.MigrationJobID, key)) if err != nil { return 0, false, errors.WithStack(err) } + liveTS, liveExists, err := g.Store.LatestCommitTS(ctx, key) + if err != nil { + return 0, false, errors.WithStack(err) + } return maxStagedVisibilityLatestCommitTS(liveTS, liveExists, stagedTS, stagedExists), liveExists || stagedExists, nil } diff --git a/kv/sharded_coordinator_test.go b/kv/sharded_coordinator_test.go new file mode 100644 index 000000000..9853cd9d6 --- /dev/null +++ b/kv/sharded_coordinator_test.go @@ -0,0 +1,69 @@ +package kv + +import ( + "context" + "testing" + + "github.com/bootjp/elastickv/distribution" + "github.com/bootjp/elastickv/store" + "github.com/stretchr/testify/require" +) + +// The OCC read-key validation reads the same two namespaces as the store, and +// missing a version here is worse than a stale read: with no commit above +// startTS the transaction commits on a read it should have conflicted with. +// Promotion landing between a live-first pair hides exactly that version. +func TestValidateReadKeysOnShardReadsStagedBeforeLive(t *testing.T) { + t.Parallel() + + ctx := context.Background() + engine := distribution.NewEngine() + require.NoError(t, engine.ApplySnapshot(distribution.CatalogSnapshot{ + Version: 1, + Routes: []distribution.RouteDescriptor{{ + RouteID: 1, Start: []byte("a"), End: []byte("z"), GroupID: 1, + State: distribution.RouteStateActive, StagedVisibilityActive: true, + MigrationJobID: 9, MinWriteTSExclusive: 100, + }}, + })) + + rawKey := []byte("k") + stagedKey := distribution.MigrationStagedDataKey(9, rawKey) + inner := store.NewMVCCStore() + t.Cleanup(func() { _ = inner.Close() }) + // A staged commit above the transaction's startTS: the read must conflict. + require.NoError(t, inner.PutAt(ctx, stagedKey, []byte("staged"), 60, 0)) + + promoting := &promotingExportStore{MVCCStore: inner} + promoting.afterFirst = func() { + promoter, ok := inner.(store.MigrationPromoter) + require.True(t, ok) + _, err := promoter.PromoteVersions(ctx, store.PromoteVersionsOptions{ + JobID: 9, + StartKey: distribution.MigrationStagedDataKeyPrefix(9), + EndKey: prefixScanEnd(distribution.MigrationStagedDataKeyPrefix(9)), + MaxVersions: 16, + TargetKey: func(staged []byte) ([]byte, bool) { + _, raw, ok := distribution.MigrationStagedDataKeyParts(staged) + return raw, ok + }, + }) + require.NoError(t, err) + } + + c := &ShardedCoordinator{ + engine: engine, + groups: map[uint64]*ShardGroup{1: {Store: promoting}}, + } + + // latestCommitTSForReadKeyOnShard is the unit under test; + // validateReadKeysOnShard wraps it in a Raft read barrier that needs a + // live engine, which this case does not depend on. A missed version here + // is what makes validateReadKeysOnShard let the transaction through. + ts, exists, err := c.latestCommitTSForReadKeyOnShard(ctx, 1, c.groups[1], rawKey) + require.NoError(t, err) + require.True(t, exists, + "a promotion between the probes must not hide the conflicting commit") + require.Equal(t, uint64(60), ts) + require.Greater(t, ts, uint64(50), "the read must still conflict with startTS=50") +} From 1bbe007356caca26467b4f947870952197f4f71c Mon Sep 17 00:00:00 2001 From: bootjp Date: Sat, 29 Aug 2026 15:52:53 +0900 Subject: [PATCH 39/58] store: emit the oldest snapshot layout that fits the state Every snapshot version is its predecessor plus one trailing metadata field: v2 adds the migration import acks and HLC floors, v3 adds the promotion state. The writer emitted v3 unconditionally, so an upgraded replica advertised a layout no older binary parses even on a cluster where no migration has ever run and all three fields are empty. A snapshot sent mid-rolling-upgrade is then rejected by a follower that has not been upgraded yet, and that follower cannot catch up. main is at v1, so the exposure is both new layouts, not just v3. Choose the oldest layout that can carry what the store actually holds, and write only the fields that layout defines. The version and the body are now decided under one lock, so a promotion landing mid-write cannot produce a body the header does not describe. Claude-Session: https://claude.ai/code/session_013rNHooj7NF3giihWVba8QE --- store/mvcc_store.go | 66 ++++++++-- store/mvcc_store_snapshot_version_test.go | 150 ++++++++++++++++++++++ 2 files changed, 204 insertions(+), 12 deletions(-) create mode 100644 store/mvcc_store_snapshot_version_test.go diff --git a/store/mvcc_store.go b/store/mvcc_store.go index e6853a28e..073db8183 100644 --- a/store/mvcc_store.go +++ b/store/mvcc_store.go @@ -27,7 +27,7 @@ type VersionedValue struct { const ( mvccSnapshotVersionV1 = uint32(1) mvccSnapshotVersionV2 = uint32(2) - mvccSnapshotVersion = uint32(3) + mvccSnapshotVersionV3 = uint32(3) maxSnapshotKeySize = 1 << 20 // 1 MiB per logical key // maxSnapshotInternalKeyEnvelope reserves snapshot headroom for internal // stored-key envelopes such as migration staging. Logical source keys keep @@ -914,12 +914,20 @@ func isStreamingMVCCSnapshot(r *bufio.Reader) (bool, error) { } func (s *mvccStore) writeSnapshotFile(f *os.File) error { - checksumOffset, err := writeMVCCSnapshotHeader(f) + // The version and the body are chosen under one lock: the version depends on + // which metadata maps are populated, so releasing between the header and the + // body would let a concurrent migration write a body the header does not + // describe. + s.mtx.RLock() + defer s.mtx.RUnlock() + + version := mvccSnapshotVersionForState(s.migrationAcks, s.migrationHLCFloors, s.migrationPromotions) + checksumOffset, err := writeMVCCSnapshotHeader(f, version) if err != nil { return err } - sum, err := s.writeSnapshotBody(f) + sum, err := s.writeSnapshotBodyLocked(f, version) if err != nil { return err } @@ -927,11 +935,33 @@ func (s *mvccStore) writeSnapshotFile(f *os.File) error { return finalizeMVCCSnapshotFile(f, checksumOffset, sum) } -func writeMVCCSnapshotHeader(f *os.File) (int64, error) { +// mvccSnapshotVersionForState returns the oldest snapshot layout that can carry +// this state. Every version is its predecessor plus one trailing metadata +// field, so a store holding none of that metadata still emits the layout older +// binaries already parse. Emitting the newest layout unconditionally would +// strand a not-yet-upgraded follower mid-rolling-upgrade on a snapshot it +// rejects as unsupported, even on a cluster where the migration opcodes have +// never been enabled and there is nothing extra to carry. +func mvccSnapshotVersionForState( + acks map[migrationAckID]migrationImportAck, + floors map[uint64]uint64, + promotions map[uint64]PromotionState, +) uint32 { + switch { + case len(promotions) > 0: + return mvccSnapshotVersionV3 + case len(acks) > 0 || len(floors) > 0: + return mvccSnapshotVersionV2 + default: + return mvccSnapshotVersionV1 + } +} + +func writeMVCCSnapshotHeader(f *os.File, version uint32) (int64, error) { if _, err := f.Write(mvccSnapshotMagic[:]); err != nil { return 0, errors.WithStack(err) } - if err := binary.Write(f, binary.LittleEndian, mvccSnapshotVersion); err != nil { + if err := binary.Write(f, binary.LittleEndian, version); err != nil { return 0, errors.WithStack(err) } checksumOffset, err := f.Seek(0, io.SeekCurrent) @@ -944,15 +974,20 @@ func writeMVCCSnapshotHeader(f *os.File) (int64, error) { return checksumOffset, nil } -func (s *mvccStore) writeSnapshotBody(f *os.File) (uint32, error) { +func (s *mvccStore) writeSnapshotBodyLocked(f *os.File, version uint32) (uint32, error) { hash := crc32.NewIEEE() bw := bufio.NewWriter(f) w := io.MultiWriter(bw, hash) - s.mtx.RLock() - defer s.mtx.RUnlock() - - if err := writeMVCCSnapshotMetadata(w, s.lastCommitTS, s.minRetainedTS, s.migrationAcks, s.migrationHLCFloors, s.migrationPromotions); err != nil { + if err := writeMVCCSnapshotMetadata( + w, + version, + s.lastCommitTS, + s.minRetainedTS, + s.migrationAcks, + s.migrationHLCFloors, + s.migrationPromotions, + ); err != nil { return 0, err } iter := s.tree.Iterator() @@ -977,6 +1012,7 @@ func (s *mvccStore) writeSnapshotBody(f *os.File) (uint32, error) { func writeMVCCSnapshotMetadata( w io.Writer, + version uint32, lastCommitTS uint64, minRetainedTS uint64, migrationAcks map[migrationAckID]migrationImportAck, @@ -989,12 +1025,18 @@ func writeMVCCSnapshotMetadata( if err := binary.Write(w, binary.LittleEndian, minRetainedTS); err != nil { return errors.WithStack(err) } + if version == mvccSnapshotVersionV1 { + return nil + } if err := writeMVCCSnapshotBytes(w, encodeMigrationImportAcks(migrationAcks)); err != nil { return err } if err := writeMVCCSnapshotBytes(w, encodeMigrationHLCFloors(migrationHLCFloors)); err != nil { return err } + if version < mvccSnapshotVersionV3 { + return nil + } if err := writeMVCCSnapshotBytes(w, encodeMigrationPromotionStates(migrationPromotions)); err != nil { return err } @@ -1139,7 +1181,7 @@ func readMVCCSnapshotHeader(r io.Reader) (uint32, uint32, error) { if err := binary.Read(r, binary.LittleEndian, &version); err != nil { return 0, 0, errors.WithStack(err) } - if version != mvccSnapshotVersionV1 && version != mvccSnapshotVersionV2 && version != mvccSnapshotVersion { + if version != mvccSnapshotVersionV1 && version != mvccSnapshotVersionV2 && version != mvccSnapshotVersionV3 { return 0, 0, errors.WithStack(errors.Newf("unsupported mvcc snapshot version %d", version)) } @@ -1218,7 +1260,7 @@ func readMVCCSnapshotMetadata( } func readMVCCSnapshotPromotionMetadata(r io.Reader, version uint32) (map[uint64]PromotionState, error) { - if version < mvccSnapshotVersion { + if version < mvccSnapshotVersionV3 { return make(map[uint64]PromotionState), nil } promotionData, err := readMVCCSnapshotBytes(r, "snapshot migration promotions") diff --git a/store/mvcc_store_snapshot_version_test.go b/store/mvcc_store_snapshot_version_test.go new file mode 100644 index 000000000..a461550a0 --- /dev/null +++ b/store/mvcc_store_snapshot_version_test.go @@ -0,0 +1,150 @@ +package store + +import ( + "bytes" + "context" + "encoding/binary" + "testing" + + "github.com/stretchr/testify/require" +) + +// snapshotHeaderVersion reads the layout version a snapshot advertises, which +// is what an older binary checks against its own accept list before it will +// restore from it. +func snapshotHeaderVersion(t *testing.T, raw []byte) uint32 { + t.Helper() + + require.Greater(t, len(raw), len(mvccSnapshotMagic)+4) + require.Equal(t, mvccSnapshotMagic[:], raw[:len(mvccSnapshotMagic)]) + return binary.LittleEndian.Uint32(raw[len(mvccSnapshotMagic) : len(mvccSnapshotMagic)+4]) +} + +// Each snapshot layout is its predecessor plus one trailing metadata field, so +// a replica must advertise the oldest layout that can carry the state it +// actually holds. Advertising the newest one unconditionally strands a +// not-yet-upgraded follower during a rolling upgrade: it rejects the snapshot +// as an unsupported version and cannot catch up, even on a cluster where no +// migration has ever run and the extra fields would all be empty. +func TestMVCCStoreSnapshotEmitsOldestSufficientVersion(t *testing.T) { + t.Parallel() + + prefix := []byte("stage|") + seedPromotion := func(t *testing.T, st MVCCStore) { + t.Helper() + promoter, ok := any(st).(MigrationPromoter) + require.True(t, ok) + require.NoError(t, st.PutAt(context.Background(), append(bytes.Clone(prefix), []byte("k")...), []byte("v"), 70, 0)) + promoted, err := promoter.PromoteVersions(context.Background(), PromoteVersionsOptions{ + JobID: 7, + StartKey: prefix, + EndKey: PrefixScanEnd(prefix), + MaxVersions: 10, + TargetKey: func(staged []byte) ([]byte, bool) { + return bytes.TrimPrefix(staged, prefix), bytes.HasPrefix(staged, prefix) + }, + }) + require.NoError(t, err) + require.True(t, promoted.Done) + } + seedImport := func(t *testing.T, st MVCCStore) { + t.Helper() + _, err := st.ImportVersions(context.Background(), ImportVersionsOptions{ + JobID: 7, + BracketID: 3, + BatchSeq: 1, + Cursor: []byte("c"), + Versions: []MVCCVersion{{Key: []byte("imported"), CommitTS: 50, Value: []byte("v50")}}, + }) + require.NoError(t, err) + } + + for _, tc := range []struct { + name string + seed func(*testing.T, MVCCStore) + want uint32 + }{ + { + name: "no migration metadata stays on the original layout", + want: mvccSnapshotVersionV1, + }, + { + name: "import acks and hlc floors need v2", + seed: seedImport, + want: mvccSnapshotVersionV2, + }, + { + name: "promotion state needs v3", + seed: seedPromotion, + want: mvccSnapshotVersionV3, + }, + } { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + + ctx := context.Background() + st := newTestMVCCStore(t) + require.NoError(t, st.PutAt(ctx, []byte("base"), []byte("v1"), 10, 0)) + if tc.seed != nil { + tc.seed(t, st) + } + + snap, err := st.Snapshot() + require.NoError(t, err) + defer snap.Close() + raw := snapshotBytes(t, snap) + + require.Equal(t, tc.want, snapshotHeaderVersion(t, raw)) + + // Whatever version was chosen, the body must match it: restoring + // through the real reader reproduces the data and the metadata. + dst := newTestMVCCStore(t) + require.NoError(t, dst.Restore(bytes.NewReader(raw))) + got, err := dst.GetAt(ctx, []byte("base"), 10) + require.NoError(t, err) + require.Equal(t, []byte("v1"), got) + }) + } +} + +// The metadata a newer layout exists to carry must survive the round trip, so +// the version choice cannot be made by simply dropping fields. +func TestMVCCStoreSnapshotRoundTripsMigrationMetadataAtItsVersion(t *testing.T) { + t.Parallel() + + ctx := context.Background() + st := newTestMVCCStore(t) + _, err := st.ImportVersions(ctx, ImportVersionsOptions{ + JobID: 7, + BracketID: 3, + BatchSeq: 1, + Cursor: []byte("c"), + Versions: []MVCCVersion{{Key: []byte("imported"), CommitTS: 50, Value: []byte("v50")}}, + }) + require.NoError(t, err) + + snap, err := st.Snapshot() + require.NoError(t, err) + defer snap.Close() + raw := snapshotBytes(t, snap) + require.Equal(t, mvccSnapshotVersionV2, snapshotHeaderVersion(t, raw)) + + dst := newTestMVCCStore(t) + require.NoError(t, dst.Restore(bytes.NewReader(raw))) + floor, err := dst.MigrationHLCFloor(ctx, 7) + require.NoError(t, err) + require.Equal(t, uint64(50), floor, "the hlc floor is why v2 exists; it must survive") + + // The recorded ack is what makes a replayed batch a duplicate rather than a + // second application, so it has to survive the restore as well. + result, err := dst.ImportVersions(ctx, ImportVersionsOptions{ + JobID: 7, + BracketID: 3, + BatchSeq: 1, + Cursor: []byte("replayed"), + Versions: []MVCCVersion{{Key: []byte("imported"), CommitTS: 50, Value: []byte("v50")}}, + }) + require.NoError(t, err) + require.True(t, result.Duplicate) + require.Equal(t, []byte("c"), result.AckedCursor) +} From 712431ab4cdfc361d4c20a617a9c2b11e2476366 Mon Sep 17 00:00:00 2001 From: bootjp Date: Mon, 31 Aug 2026 23:29:27 +0900 Subject: [PATCH 40/58] migration: tighten staged transaction retries --- adapter/internal.go | 10 +- adapter/internal_migration_test.go | 12 ++ kv/fsm.go | 226 +++++++++++++++++++---------- kv/fsm_abort_test.go | 93 +++++++++++- kv/fsm_onephase_dedup_test.go | 35 +---- kv/fsm_txn_test.go | 89 ++++++++++++ 6 files changed, 353 insertions(+), 112 deletions(-) diff --git a/adapter/internal.go b/adapter/internal.go index a5b790209..a160b91d1 100644 --- a/adapter/internal.go +++ b/adapter/internal.go @@ -276,16 +276,14 @@ func (i *Internal) validateExportRangeVersionsRequest(req *pb.ExportRangeVersion 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")) + if exportRangeVersionsRequestRouteUnbounded(req) { + return errors.WithStack(status.Error(codes.InvalidArgument, "migration export requires route bounds")) } return nil } -func exportRangeVersionsRequestFullyUnbounded(req *pb.ExportRangeVersionsRequest) bool { - return len(req.GetRangeStart()) == 0 && - len(req.GetRangeEnd()) == 0 && - len(req.GetRouteStart()) == 0 && +func exportRangeVersionsRequestRouteUnbounded(req *pb.ExportRangeVersionsRequest) bool { + return len(req.GetRouteStart()) == 0 && len(req.GetRouteEnd()) == 0 } diff --git a/adapter/internal_migration_test.go b/adapter/internal_migration_test.go index 91e395deb..18e577f52 100644 --- a/adapter/internal_migration_test.go +++ b/adapter/internal_migration_test.go @@ -243,6 +243,16 @@ func TestInternalExportRangeVersionsRejectsUnboundedExport(t *testing.T) { }, stream) require.Error(t, err) require.Equal(t, codes.InvalidArgument, status.Code(err)) + + stream = &captureExportRangeVersionsStream{ctx: context.Background()} + err = internal.ExportRangeVersions(&pb.ExportRangeVersionsRequest{ + MaxCommitTs: 20, + KeyFamily: distribution.MigrationFamilyUser, + RangeStart: []byte("a"), + RangeEnd: []byte("b"), + }, stream) + require.Error(t, err) + require.Equal(t, codes.InvalidArgument, status.Code(err)) } func TestInternalExportRangeVersionsUsesDecodedS3BucketRouteFilter(t *testing.T) { @@ -440,6 +450,8 @@ func TestInternalExportRangeVersionsUsesPartitionResolverGroup(t *testing.T) { err := internal.ExportRangeVersions(&pb.ExportRangeVersionsRequest{ MaxCommitTs: 20, KeyFamily: distribution.MigrationFamilySQSPartitionedMessageData, + RouteStart: []byte("!sqs|route|global"), + RouteEnd: testPrefixScanEnd([]byte("!sqs|route|global")), RangeStart: prefix, RangeEnd: testPrefixScanEnd(prefix), MaxScannedBytes: 1 << 20, diff --git a/kv/fsm.go b/kv/fsm.go index 4f5c266b9..ae58bd6e4 100644 --- a/kv/fsm.go +++ b/kv/fsm.go @@ -1360,7 +1360,7 @@ func (f *kvFSM) handleOnePhaseTxnRequest(ctx context.Context, r *pb.Request, com // duplicate (the very :duplicate-elements anomaly), so no-op the whole // apply and let the adapter reconstruct the prior result. // - // Determinism note (codex P1 round-11): the underlying CommittedVersionAt + // Determinism note: the underlying CommittedVersionAt // intentionally does NOT enforce the retention watermark — branching FSM // apply on the per-replica minRetainedTS would let replicas with stale // retention surface ErrReadTSCompacted and skip dedup while replicas that @@ -1371,7 +1371,7 @@ func (f *kvFSM) handleOnePhaseTxnRequest(ctx context.Context, r *pb.Request, com // applying this log entry. The retention-window > max-retry-latency // invariant prevents the rare case where a real never-landed retry // arrives with PrevCommitTS below pebble's compacted floor. - dedup, err := f.dedupProbeOnePhase(ctx, meta) + dedup, err := f.dedupProbeOnePhase(ctx, meta, r.ReadKeys) if err != nil { return err } @@ -1463,7 +1463,7 @@ func (f *kvFSM) uniqueTxnMutationsAboveFloor(muts []*pb.Mutation, commitTS uint6 // Returns (true, nil) → the entry must no-op (prior attempt landed). // Returns (false, nil) → fall through to normal apply. // Returns (false, err) → propagate err; apply must not proceed. -func (f *kvFSM) dedupProbeOnePhase(ctx context.Context, meta TxnMeta) (bool, error) { +func (f *kvFSM) dedupProbeOnePhase(ctx context.Context, meta TxnMeta, readKeys [][]byte) (bool, error) { if meta.PrevCommitTS == 0 { return false, nil } @@ -1474,50 +1474,48 @@ func (f *kvFSM) dedupProbeOnePhase(ctx context.Context, meta TxnMeta) (bool, 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) + for _, stagedKey := range stagedVisibilityKeysForReadAlias(readKeys, meta.PrimaryKey, meta.PrimaryKey) { + landed, err = f.store.CommittedVersionAt(ctx, stagedKey, meta.PrevCommitTS) + if err != nil { + return false, errors.WithStack(err) + } + if landed { + return true, nil + } } - return landed, nil + return false, 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 +func stagedVisibilityJobIDsForReadAlias(readKeys [][]byte, aliasRawKey []byte) []uint64 { + if len(readKeys) == 0 { + return nil } - route, ok := snap.RouteOf(routeKey(key)) - if !ok || route.GroupID != f.shardGroupID || !routeHasStagedVisibility(route) { - return distribution.Route{}, false + var out []uint64 + seen := make(map[uint64]struct{}) + for _, readKey := range readKeys { + jobID, rawKey, ok := distribution.MigrationStagedDataKeyParts(readKey) + if !ok || jobID == 0 || !bytes.Equal(rawKey, aliasRawKey) { + continue + } + if _, ok := seen[jobID]; ok { + continue + } + seen[jobID] = struct{}{} + out = append(out, jobID) } - return route, true + return out } -func currentStagedVisibilityRouteForS3BucketAuxiliaryKey(snap RouteSnapshot, key []byte, shardGroupID uint64) (distribution.Route, bool) { - start, end, ok := s3BucketAuxiliaryRouteRange(key) - if !ok { - return distribution.Route{}, false +func stagedVisibilityKeysForReadAlias(readKeys [][]byte, aliasRawKey, stagedRawKey []byte) [][]byte { + jobIDs := stagedVisibilityJobIDsForReadAlias(readKeys, aliasRawKey) + if len(jobIDs) == 0 { + return nil } - for _, route := range snap.IntersectingRoutes(start, end) { - if route.GroupID == shardGroupID && routeHasStagedVisibility(route) { - return route, true - } + out := make([][]byte, 0, len(jobIDs)) + for _, jobID := range jobIDs { + out = append(out, distribution.MigrationStagedDataKey(jobID, stagedRawKey)) } - return distribution.Route{}, false + return out } func (f *kvFSM) handleCommitRequest(ctx context.Context, r *pb.Request) error { @@ -1536,7 +1534,7 @@ func (f *kvFSM) handleCommitRequest(ctx context.Context, r *pb.Request) error { if len(meta.PrimaryKey) == 0 { return errors.WithStack(ErrTxnPrimaryKeyRequired) } - applyStartTS, err := f.commitApplyStartTS(ctx, meta.PrimaryKey, startTS, commitTS) + applyStartTS, err := f.commitApplyStartTS(ctx, meta.PrimaryKey, startTS, commitTS, r.ReadKeys) if err != nil { return err } @@ -1544,7 +1542,7 @@ func (f *kvFSM) handleCommitRequest(ctx context.Context, r *pb.Request) error { if err != nil { return err } - storeMuts, err := f.buildCommitStoreMutations(ctx, uniq, meta, startTS, commitTS) + storeMuts, err := f.buildCommitStoreMutations(ctx, uniq, meta, startTS, commitTS, r.ReadKeys) if err != nil { return err } @@ -1567,11 +1565,17 @@ func (f *kvFSM) handleCommitRequest(ctx context.Context, r *pb.Request) error { // ErrTxnAlreadyAborted. Together with the commit-record check in // appendRollbackRecord, this enforces the invariant that at most one of // {rollback marker, commit record} is present for any (primaryKey, startTS). -func (f *kvFSM) commitApplyStartTS(ctx context.Context, primaryKey []byte, startTS, commitTS uint64) (uint64, error) { +func (f *kvFSM) commitApplyStartTS(ctx context.Context, primaryKey []byte, startTS, commitTS uint64, readKeys [][]byte) (uint64, error) { recordedCommitTS, committed, err := f.txnCommitTS(ctx, primaryKey, startTS) if err != nil { return 0, err } + if !committed { + recordedCommitTS, committed, err = f.stagedTxnCommitTS(ctx, primaryKey, startTS, readKeys) + if err != nil { + return 0, err + } + } if !committed { // No commit record yet: reject if a rollback marker is present. // This catches out-of-order apply (COMMIT after ABORT), buggy @@ -1583,6 +1587,13 @@ func (f *kvFSM) commitApplyStartTS(ctx context.Context, primaryKey []byte, start if exists { return 0, errors.WithStack(ErrTxnAlreadyAborted) } + exists, rerr = f.stagedTxnRecordExists(ctx, primaryKey, txnRollbackKey(primaryKey, startTS), readKeys) + if rerr != nil { + return 0, rerr + } + if exists { + return 0, errors.WithStack(ErrTxnAlreadyAborted) + } return startTS, nil } if recordedCommitTS != commitTS { @@ -1654,19 +1665,19 @@ func (f *kvFSM) handleAbortRequest(ctx context.Context, r *pb.Request, abortTS u // abort for a secondary key (same primaryKey, same startTS) would // see the marker already present and must still clean up that // secondary's lock/intent. Idempotency is enforced per-key in - // shouldClearAbortKey (lock-missing ⇒ nothing to do) and for the + // abortCleanupMutationsForKey (lock-missing ⇒ nothing to do) and for the // rollback-marker Put in appendRollbackRecord. uniq, err := uniqueMutations(muts) if err != nil { return err } - storeMuts, abortingPrimary, err := f.buildAbortCleanupStoreMutations(ctx, uniq, meta.PrimaryKey, startTS) + storeMuts, abortingPrimary, err := f.buildAbortCleanupStoreMutations(ctx, uniq, meta.PrimaryKey, startTS, r.ReadKeys) if err != nil { return err } if abortingPrimary { - if err := f.appendRollbackRecord(ctx, meta.PrimaryKey, startTS, &storeMuts); err != nil { + if err := f.appendRollbackRecord(ctx, meta.PrimaryKey, startTS, &storeMuts, r.ReadKeys); err != nil { return err } } @@ -1705,7 +1716,7 @@ func (f *kvFSM) buildOnePhaseStoreMutations(ctx context.Context, muts []*pb.Muta return storeMuts, nil } -func (f *kvFSM) buildCommitStoreMutations(ctx context.Context, muts []*pb.Mutation, meta TxnMeta, startTS, commitTS uint64) ([]*store.KVPairMutation, error) { +func (f *kvFSM) buildCommitStoreMutations(ctx context.Context, muts []*pb.Mutation, meta TxnMeta, startTS, commitTS uint64, readKeys [][]byte) ([]*store.KVPairMutation, error) { storeMuts := make([]*store.KVPairMutation, 0, len(muts)*txnCommitStoreMutationFactor+txnCommitStoreMutationSlack) committingPrimary := false @@ -1715,7 +1726,7 @@ func (f *kvFSM) buildCommitStoreMutations(ctx context.Context, muts []*pb.Mutati committingPrimary = true } - keyMuts, err := f.commitTxnKeyMutations(ctx, key, meta.PrimaryKey, startTS) + keyMuts, err := f.commitTxnKeyMutations(ctx, key, meta.PrimaryKey, startTS, readKeys) if err != nil { return nil, err } @@ -1733,7 +1744,7 @@ func (f *kvFSM) buildCommitStoreMutations(ctx context.Context, muts []*pb.Mutati return storeMuts, nil } -func (f *kvFSM) buildAbortCleanupStoreMutations(ctx context.Context, muts []*pb.Mutation, primaryKey []byte, startTS uint64) ([]*store.KVPairMutation, bool, error) { +func (f *kvFSM) buildAbortCleanupStoreMutations(ctx context.Context, muts []*pb.Mutation, primaryKey []byte, startTS uint64, readKeys [][]byte) ([]*store.KVPairMutation, bool, error) { storeMuts := make([]*store.KVPairMutation, 0, len(muts)*txnAbortStoreMutationFactor) abortingPrimary := false for _, mut := range muts { @@ -1742,18 +1753,16 @@ func (f *kvFSM) buildAbortCleanupStoreMutations(ctx context.Context, muts []*pb. abortingPrimary = true } - shouldClear, err := f.shouldClearAbortKey(ctx, key, primaryKey, startTS) + cleanupMuts, err := f.abortCleanupMutationsForKey(ctx, key, primaryKey, startTS, readKeys) if err != nil { return nil, false, err } - if shouldClear { - storeMuts = append(storeMuts, txnCleanupMutations(key)...) - } + storeMuts = append(storeMuts, cleanupMuts...) } return storeMuts, abortingPrimary, nil } -func (f *kvFSM) appendRollbackRecord(ctx context.Context, primaryKey []byte, startTS uint64, storeMuts *[]*store.KVPairMutation) error { +func (f *kvFSM) appendRollbackRecord(ctx context.Context, primaryKey []byte, startTS uint64, storeMuts *[]*store.KVPairMutation, readKeys [][]byte) error { // Desired invariant: for any (primaryKey, startTS) pair, at most // one of {rollback marker, commit record} is present. The invariant // holds when aborts/commits flow through the symmetric guards in @@ -1772,6 +1781,12 @@ func (f *kvFSM) appendRollbackRecord(ctx context.Context, primaryKey []byte, sta if err != nil { return errors.WithStack(err) } + if !markerPresent { + markerPresent, err = f.stagedTxnRecordExists(ctx, primaryKey, txnRollbackKey(primaryKey, startTS), readKeys) + if err != nil { + return err + } + } // Verify the invariant regardless of marker presence: if a commit // record is present for this (primaryKey, startTS), refuse to @@ -1782,6 +1797,12 @@ func (f *kvFSM) appendRollbackRecord(ctx context.Context, primaryKey []byte, sta if err != nil { return errors.WithStack(err) } + if !commitExists { + commitExists, err = f.stagedTxnRecordExists(ctx, primaryKey, txnCommitKey(primaryKey, startTS), readKeys) + if err != nil { + return err + } + } if commitExists { return errors.WithStack(ErrTxnAlreadyCommitted) } @@ -1799,7 +1820,37 @@ func (f *kvFSM) appendRollbackRecord(ctx context.Context, primaryKey []byte, sta } func (f *kvFSM) txnCommitTS(ctx context.Context, primaryKey []byte, startTS uint64) (uint64, bool, error) { - b, err := f.store.GetAt(ctx, txnCommitKey(primaryKey, startTS), ^uint64(0)) + return f.txnCommitTSAtKey(ctx, txnCommitKey(primaryKey, startTS)) +} + +func (f *kvFSM) stagedTxnCommitTS(ctx context.Context, primaryKey []byte, startTS uint64, readKeys [][]byte) (uint64, bool, error) { + for _, key := range stagedVisibilityKeysForReadAlias(readKeys, primaryKey, txnCommitKey(primaryKey, startTS)) { + commitTS, committed, err := f.txnCommitTSAtKey(ctx, key) + if err != nil { + return 0, false, err + } + if committed { + return commitTS, true, nil + } + } + return 0, false, nil +} + +func (f *kvFSM) stagedTxnRecordExists(ctx context.Context, primaryKey, recordKey []byte, readKeys [][]byte) (bool, error) { + for _, key := range stagedVisibilityKeysForReadAlias(readKeys, primaryKey, recordKey) { + exists, err := f.store.ExistsAt(ctx, key, ^uint64(0)) + if err != nil { + return false, errors.WithStack(err) + } + if exists { + return true, nil + } + } + return false, nil +} + +func (f *kvFSM) txnCommitTSAtKey(ctx context.Context, key []byte) (uint64, bool, error) { + b, err := f.store.GetAt(ctx, key, ^uint64(0)) if err != nil { if errors.Is(err, store.ErrKeyNotFound) { return 0, false, nil @@ -1850,15 +1901,38 @@ func txnIntentFromPBMutation(mut *pb.Mutation, startTS uint64) (txnIntent, error } } -func txnCleanupMutations(key []byte) []*store.KVPairMutation { +func txnCleanupMutationsForKeys(lockKey, intentKey []byte) []*store.KVPairMutation { return []*store.KVPairMutation{ - {Op: store.OpTypeDelete, Key: txnLockKey(key)}, - {Op: store.OpTypeDelete, Key: txnIntentKey(key)}, + {Op: store.OpTypeDelete, Key: lockKey}, + {Op: store.OpTypeDelete, Key: intentKey}, } } -func (f *kvFSM) txnLockForCommit(ctx context.Context, key []byte) (txnLock, bool, error) { - lockBytes, err := f.store.GetAt(ctx, txnLockKey(key), ^uint64(0)) +func (f *kvFSM) txnLockForMutationKey(ctx context.Context, key []byte, readKeys [][]byte) (txnLock, []byte, []byte, bool, error) { + lockKey := txnLockKey(key) + intentKey := txnIntentKey(key) + lock, ok, err := f.txnLockAtKey(ctx, lockKey) + if err != nil || ok { + return lock, lockKey, intentKey, ok, err + } + for _, stagedLockKey := range stagedVisibilityKeysForReadAlias(readKeys, key, lockKey) { + jobID, _, ok := distribution.MigrationStagedDataKeyParts(stagedLockKey) + if !ok { + continue + } + lock, ok, err := f.txnLockAtKey(ctx, stagedLockKey) + if err != nil { + return txnLock{}, nil, nil, false, err + } + if ok { + return lock, stagedLockKey, distribution.MigrationStagedDataKey(jobID, intentKey), true, nil + } + } + return txnLock{}, nil, nil, false, nil +} + +func (f *kvFSM) txnLockAtKey(ctx context.Context, key []byte) (txnLock, bool, error) { + lockBytes, err := f.store.GetAt(ctx, key, ^uint64(0)) if err != nil { if errors.Is(err, store.ErrKeyNotFound) { return txnLock{}, false, nil @@ -1872,8 +1946,8 @@ func (f *kvFSM) txnLockForCommit(ctx context.Context, key []byte) (txnLock, bool return lock, true, nil } -func (f *kvFSM) txnIntentForCommit(ctx context.Context, key []byte) (txnIntent, bool, error) { - intentBytes, err := f.store.GetAt(ctx, txnIntentKey(key), ^uint64(0)) +func (f *kvFSM) txnIntentForCommit(ctx context.Context, intentKey []byte) (txnIntent, bool, error) { + intentBytes, err := f.store.GetAt(ctx, intentKey, ^uint64(0)) if err != nil { if errors.Is(err, store.ErrKeyNotFound) { return txnIntent{}, false, nil @@ -1898,8 +1972,8 @@ func storeMutationForIntent(key []byte, intent txnIntent) (*store.KVPairMutation } } -func (f *kvFSM) commitTxnKeyMutations(ctx context.Context, key, primaryKey []byte, startTS uint64) ([]*store.KVPairMutation, error) { - lock, ok, err := f.txnLockForCommit(ctx, key) +func (f *kvFSM) commitTxnKeyMutations(ctx context.Context, key, primaryKey []byte, startTS uint64, readKeys [][]byte) ([]*store.KVPairMutation, error) { + lock, lockKey, intentKey, ok, err := f.txnLockForMutationKey(ctx, key, readKeys) if err != nil { return nil, err } @@ -1914,7 +1988,7 @@ func (f *kvFSM) commitTxnKeyMutations(ctx context.Context, key, primaryKey []byt return nil, errors.Wrapf(ErrTxnInvalidMeta, "lock primary_key mismatch for key %s", string(key)) } - intent, ok, err := f.txnIntentForCommit(ctx, key) + intent, ok, err := f.txnIntentForCommit(ctx, intentKey) if err != nil { return nil, err } @@ -1930,37 +2004,33 @@ func (f *kvFSM) commitTxnKeyMutations(ctx context.Context, key, primaryKey []byt } out = append(out, mut) } - out = append(out, txnCleanupMutations(key)...) + out = append(out, txnCleanupMutationsForKeys(lockKey, intentKey)...) return out, nil } -// shouldClearAbortKey reports whether this abort request must emit -// cleanup (lock+intent Delete) mutations for key. It returns false +// abortCleanupMutationsForKey reports whether this abort request must emit +// cleanup (lock+intent Delete) mutations for key. It returns nil // when the lock is already missing: lock/intent are always written // and deleted together in a single ApplyMutations batch // (lock missing ⇔ intent missing), so missing lock means either // cleanup already ran for this (startTS, primaryKey) or the key was // never prepared. Emitting Deletes on already-tombstoned keys would // trigger MVCC write conflicts and has no observable effect. -func (f *kvFSM) shouldClearAbortKey(ctx context.Context, key, primaryKey []byte, startTS uint64) (bool, error) { - lockBytes, err := f.store.GetAt(ctx, txnLockKey(key), ^uint64(0)) +func (f *kvFSM) abortCleanupMutationsForKey(ctx context.Context, key, primaryKey []byte, startTS uint64, readKeys [][]byte) ([]*store.KVPairMutation, error) { + lock, lockKey, intentKey, ok, err := f.txnLockForMutationKey(ctx, key, readKeys) if err != nil { - if errors.Is(err, store.ErrKeyNotFound) { - return false, nil - } - return false, errors.WithStack(err) + return nil, err } - lock, derr := decodeTxnLock(lockBytes) - if derr != nil { - return false, errors.WithStack(derr) + if !ok { + return nil, nil } if lock.StartTS != startTS { - return false, nil + return nil, nil } if !bytes.Equal(lock.PrimaryKey, primaryKey) { - return false, errors.Wrapf(ErrTxnInvalidMeta, "abort primary_key mismatch for key %s", string(key)) + return nil, errors.Wrapf(ErrTxnInvalidMeta, "abort primary_key mismatch for key %s", string(key)) } - return true, nil + return txnCleanupMutationsForKeys(lockKey, intentKey), nil } func (f *kvFSM) assertNoConflictingTxnLock(ctx context.Context, key, primaryKey []byte, startTS uint64) error { diff --git a/kv/fsm_abort_test.go b/kv/fsm_abort_test.go index 8498c783f..c486395c2 100644 --- a/kv/fsm_abort_test.go +++ b/kv/fsm_abort_test.go @@ -4,6 +4,7 @@ import ( "context" "testing" + "github.com/bootjp/elastickv/distribution" pb "github.com/bootjp/elastickv/proto" "github.com/bootjp/elastickv/store" "github.com/stretchr/testify/require" @@ -297,7 +298,7 @@ func TestFSMAbort_AbortTSMustBeGreaterThanStartTS(t *testing.T) { // same mutation set must return nil without performing additional // writes or store mutations (reads are allowed — the idempotent path // still probes for the rollback marker and commit record via ExistsAt). -// Idempotency is enforced per-key in shouldClearAbortKey (lock +// Idempotency is enforced per-key in abortCleanupMutationsForKey (lock // already gone ⇒ skip) and by appendRollbackRecord (marker already // present ⇒ skip). The prior behaviour (write-conflict on the // rollback-marker Put) surfaced in prod as "secondary write failed" @@ -553,6 +554,96 @@ func TestFSMAbort_CommitAfterAbortIsRejected(t *testing.T) { require.ErrorIs(t, err, store.ErrKeyNotFound, "commit record must not have been written after rejection") } +func TestFSMAbort_StagedRollbackRecordIsIdempotent(t *testing.T) { + t.Parallel() + + ctx := context.Background() + st := store.NewMVCCStore() + fsm, ok := NewKvFSMWithHLC(st, NewHLC()).(*kvFSM) + require.True(t, ok) + + primary := []byte("pk") + startTS := uint64(10) + readKeys := [][]byte{distribution.MigrationStagedDataKey(9, primary)} + stagedRollbackKey := distribution.MigrationStagedDataKey(9, txnRollbackKey(primary, startTS)) + require.NoError(t, st.PutAt(ctx, stagedRollbackKey, encodeTxnRollbackRecord(), startTS, 0)) + + var storeMuts []*store.KVPairMutation + require.NoError(t, fsm.appendRollbackRecord(ctx, primary, startTS, &storeMuts, readKeys)) + require.Empty(t, storeMuts) +} + +func TestFSMAbort_RejectsStagedCommitRecord(t *testing.T) { + t.Parallel() + + ctx := context.Background() + st := store.NewMVCCStore() + fsm, ok := NewKvFSMWithHLC(st, NewHLC()).(*kvFSM) + require.True(t, ok) + + primary := []byte("pk") + startTS := uint64(10) + commitTS := uint64(20) + readKeys := [][]byte{distribution.MigrationStagedDataKey(9, primary)} + stagedCommitKey := distribution.MigrationStagedDataKey(9, txnCommitKey(primary, startTS)) + require.NoError(t, st.PutAt(ctx, stagedCommitKey, encodeTxnCommitRecord(commitTS), commitTS, 0)) + + var storeMuts []*store.KVPairMutation + err := fsm.appendRollbackRecord(ctx, primary, startTS, &storeMuts, readKeys) + require.Error(t, err) + require.ErrorIs(t, err, ErrTxnAlreadyCommitted) + require.Empty(t, storeMuts) +} + +func TestFSMAbort_CleansUpStagedPreparedArtifacts(t *testing.T) { + t.Parallel() + + ctx := context.Background() + st := store.NewMVCCStore() + fsm, ok := NewKvFSMWithHLC(st, NewHLC()).(*kvFSM) + require.True(t, ok) + + primary := []byte("pk") + startTS := uint64(10) + abortTS := uint64(20) + jobID := uint64(9) + stagedLockKey := distribution.MigrationStagedDataKey(jobID, txnLockKey(primary)) + stagedIntentKey := distribution.MigrationStagedDataKey(jobID, txnIntentKey(primary)) + require.NoError(t, st.PutAt(ctx, stagedLockKey, encodeTxnLock(txnLock{ + StartTS: startTS, + PrimaryKey: primary, + IsPrimaryKey: true, + }), startTS, 0)) + require.NoError(t, st.PutAt(ctx, stagedIntentKey, encodeTxnIntent(txnIntent{ + StartTS: startTS, + Op: txnIntentOpPut, + Value: []byte("v"), + }), startTS, 0)) + + abortReq := &pb.Request{ + IsTxn: true, + Phase: pb.Phase_ABORT, + Ts: startTS, + ReadKeys: [][]byte{ + distribution.MigrationStagedDataKey(jobID, primary), + }, + Mutations: []*pb.Mutation{ + {Op: pb.Op_PUT, Key: []byte(txnMetaPrefix), Value: EncodeTxnMeta(TxnMeta{PrimaryKey: primary, CommitTS: abortTS})}, + {Op: pb.Op_PUT, Key: primary}, + }, + } + require.NoError(t, applyFSMRequest(t, fsm, abortReq)) + + _, err := st.GetAt(ctx, stagedLockKey, ^uint64(0)) + require.ErrorIs(t, err, store.ErrKeyNotFound) + _, err = st.GetAt(ctx, stagedIntentKey, ^uint64(0)) + require.ErrorIs(t, err, store.ErrKeyNotFound) + _, err = st.GetAt(ctx, txnRollbackKey(primary, startTS), ^uint64(0)) + require.NoError(t, err) + _, err = st.GetAt(ctx, primary, ^uint64(0)) + require.ErrorIs(t, err, store.ErrKeyNotFound) +} + func TestFSMAbort_EmptyMutationsReturnsError(t *testing.T) { t.Parallel() diff --git a/kv/fsm_onephase_dedup_test.go b/kv/fsm_onephase_dedup_test.go index b675e1079..c9f06c088 100644 --- a/kv/fsm_onephase_dedup_test.go +++ b/kv/fsm_onephase_dedup_test.go @@ -73,24 +73,15 @@ 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) + fsm, ok := NewKvFSMWithHLC(st, NewHLC()).(*kvFSM) require.True(t, ok) key := []byte("list-item") require.NoError(t, st.PutAt(ctx, distribution.MigrationStagedDataKey(9, key), []byte("v"), 20, 0)) - require.NoError(t, applyFSMRequest(t, fsm, onePhaseReq(30, 40, 20, key, []byte("v")))) + req := onePhaseReq(30, 40, 20, key, []byte("v")) + req.ReadKeys = [][]byte{distribution.MigrationStagedDataKey(9, key)} + require.NoError(t, applyFSMRequest(t, fsm, req)) at40, err := st.CommittedVersionAt(ctx, key, 40) require.NoError(t, err) @@ -104,26 +95,16 @@ func TestOnePhaseDedup_NoOpsWhenS3AuxiliaryPriorAttemptLandedAsStagedVersion(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) + fsm, ok := NewKvFSMWithHLC(st, NewHLC()).(*kvFSM) require.True(t, ok) key := s3keys.BucketMetaKey(bucket) require.NoError(t, st.PutAt(ctx, distribution.MigrationStagedDataKey(9, key), []byte("v"), 20, 0)) - require.NoError(t, applyFSMRequest(t, fsm, onePhaseReq(30, 40, 20, key, []byte("v")))) + req := onePhaseReq(30, 40, 20, key, []byte("v")) + req.ReadKeys = [][]byte{distribution.MigrationStagedDataKey(9, key)} + require.NoError(t, applyFSMRequest(t, fsm, req)) at40, err := st.CommittedVersionAt(ctx, key, 40) require.NoError(t, err) diff --git a/kv/fsm_txn_test.go b/kv/fsm_txn_test.go index db35beab5..97fb7d17f 100644 --- a/kv/fsm_txn_test.go +++ b/kv/fsm_txn_test.go @@ -5,6 +5,7 @@ import ( "fmt" "testing" + "github.com/bootjp/elastickv/distribution" pb "github.com/bootjp/elastickv/proto" "github.com/bootjp/elastickv/store" "github.com/stretchr/testify/require" @@ -213,6 +214,94 @@ func TestCommitIsIdempotentAfterCommitRecordExists(t *testing.T) { require.Equal(t, commitTS, gotCommitTS) } +func TestCommitApplyStartTSUsesStagedCommitRecord(t *testing.T) { + t.Parallel() + + ctx := context.Background() + st := store.NewMVCCStore() + fsm, ok := NewKvFSMWithHLC(st, NewHLC()).(*kvFSM) + require.True(t, ok) + + primary := []byte("pk") + startTS := uint64(10) + commitTS := uint64(20) + readKeys := [][]byte{distribution.MigrationStagedDataKey(9, primary)} + stagedCommitKey := distribution.MigrationStagedDataKey(9, txnCommitKey(primary, startTS)) + require.NoError(t, st.PutAt(ctx, stagedCommitKey, encodeTxnCommitRecord(commitTS), commitTS, 0)) + + applyStartTS, err := fsm.commitApplyStartTS(ctx, primary, startTS, commitTS, readKeys) + require.NoError(t, err) + require.Equal(t, commitTS, applyStartTS) +} + +func TestCommitRejectsStagedRollbackRecord(t *testing.T) { + t.Parallel() + + ctx := context.Background() + st := store.NewMVCCStore() + fsm, ok := NewKvFSMWithHLC(st, NewHLC()).(*kvFSM) + require.True(t, ok) + + primary := []byte("pk") + startTS := uint64(10) + commitTS := uint64(20) + readKeys := [][]byte{distribution.MigrationStagedDataKey(9, primary)} + stagedRollbackKey := distribution.MigrationStagedDataKey(9, txnRollbackKey(primary, startTS)) + require.NoError(t, st.PutAt(ctx, stagedRollbackKey, encodeTxnRollbackRecord(), startTS, 0)) + + _, err := fsm.commitApplyStartTS(ctx, primary, startTS, commitTS, readKeys) + require.Error(t, err) + require.ErrorIs(t, err, ErrTxnAlreadyAborted) +} + +func TestCommitCleansUpStagedPreparedArtifacts(t *testing.T) { + t.Parallel() + + ctx := context.Background() + st := store.NewMVCCStore() + fsm, ok := NewKvFSMWithHLC(st, NewHLC()).(*kvFSM) + require.True(t, ok) + + primary := []byte("pk") + startTS := uint64(10) + commitTS := uint64(20) + jobID := uint64(9) + stagedLockKey := distribution.MigrationStagedDataKey(jobID, txnLockKey(primary)) + stagedIntentKey := distribution.MigrationStagedDataKey(jobID, txnIntentKey(primary)) + require.NoError(t, st.PutAt(ctx, stagedLockKey, encodeTxnLock(txnLock{ + StartTS: startTS, + PrimaryKey: primary, + IsPrimaryKey: true, + }), startTS, 0)) + require.NoError(t, st.PutAt(ctx, stagedIntentKey, encodeTxnIntent(txnIntent{ + StartTS: startTS, + Op: txnIntentOpPut, + Value: []byte("v"), + }), startTS, 0)) + + commit := &pb.Request{ + IsTxn: true, + Phase: pb.Phase_COMMIT, + Ts: startTS, + ReadKeys: [][]byte{ + distribution.MigrationStagedDataKey(jobID, primary), + }, + Mutations: []*pb.Mutation{ + {Op: pb.Op_PUT, Key: []byte(txnMetaPrefix), Value: EncodeTxnMeta(TxnMeta{PrimaryKey: primary, CommitTS: commitTS})}, + {Op: pb.Op_PUT, Key: primary}, + }, + } + require.NoError(t, applyFSMRequest(t, fsm, commit)) + + value, err := st.GetAt(ctx, primary, ^uint64(0)) + require.NoError(t, err) + require.Equal(t, []byte("v"), value) + _, err = st.GetAt(ctx, stagedLockKey, ^uint64(0)) + require.ErrorIs(t, err, store.ErrKeyNotFound) + _, err = st.GetAt(ctx, stagedIntentKey, ^uint64(0)) + require.ErrorIs(t, err, store.ErrKeyNotFound) +} + func TestCommitIsIdempotentOnSecondaryShardWhenKeyAlreadyCommitted(t *testing.T) { t.Parallel() From 7b9e199fd0492903fcc97d05f9f76bf62c6b14dd Mon Sep 17 00:00:00 2001 From: bootjp Date: Mon, 31 Aug 2026 23:36:30 +0900 Subject: [PATCH 41/58] migration: honor S3 auxiliary write fences --- kv/fsm.go | 23 ++++++++------ kv/fsm_migration_fence_test.go | 27 +++++++++++++++++ kv/sharded_coordinator.go | 18 +++++------ kv/sharded_coordinator_del_prefix_test.go | 37 +++++++++++++++++++++++ 4 files changed, 87 insertions(+), 18 deletions(-) diff --git a/kv/fsm.go b/kv/fsm.go index ae58bd6e4..010152196 100644 --- a/kv/fsm.go +++ b/kv/fsm.go @@ -694,13 +694,16 @@ func (f *kvFSM) verifyRouteNotFencedForKey(key []byte) error { if !ok { return nil } + if start, end, ok := s3BucketAuxiliaryRouteRange(key); ok { + if snap.WriteFencedIntersects(start, end) { + return errors.Wrapf(ErrRouteWriteFenced, "key %q route range [%q,%q)", key, start, end) + } + return nil + } rkey := routeKey(key) if snap.WriteFencedForKey(rkey) { return errors.Wrapf(ErrRouteWriteFenced, "key %q routeKey %q", key, rkey) } - 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 } @@ -1183,18 +1186,20 @@ func verifyWriteFenceFromSnapshot(mutations []*pb.Mutation, writeFenceBypassKeys if _, ok := bypassKeys[string(mut.Key)]; ok { continue } + if start, end, ok := s3BucketAuxiliaryRouteRange(mut.Key); ok { + if snap.WriteFencedIntersects(start, end) { + return errors.Wrapf(ErrRouteWriteFenced, + "%s-version v=%d: key %q route range [%q,%q)", + phase, snapVer, mut.Key, start, end) + } + continue + } rKey := routeKey(mut.Key) if snap.WriteFencedForKey(rKey) { return errors.Wrapf(ErrRouteWriteFenced, "%s-version v=%d: key %q routeKey %q", phase, snapVer, mut.Key, rKey) } - start, end, ok := s3BucketAuxiliaryRouteRange(mut.Key) - if ok && snap.WriteFencedIntersects(start, end) { - return errors.Wrapf(ErrRouteWriteFenced, - "%s-version v=%d: key %q route range [%q,%q)", - phase, snapVer, mut.Key, start, end) - } } return nil } diff --git a/kv/fsm_migration_fence_test.go b/kv/fsm_migration_fence_test.go index b64a84d7d..df7b4846b 100644 --- a/kv/fsm_migration_fence_test.go +++ b/kv/fsm_migration_fence_test.go @@ -354,6 +354,33 @@ func TestFSMRejectsObservedWriteFencedS3BucketAuxiliaryPointWrite(t *testing.T) require.ErrorIs(t, err, ErrRouteWriteFenced) } +func TestFSMIgnoresRawRouteFenceForS3BucketAuxiliaryWrite(t *testing.T) { + t.Parallel() + + const bucket = "bucket-a" + key := s3keys.BucketMetaKey(bucket) + engine := distribution.NewEngine() + routes := s3BucketAuxiliaryFenceRoutes(bucket, 1, 1) + routes[1].State = distribution.RouteStateActive + routes[2].State = distribution.RouteStateWriteFenced + applyComposed1Snapshot(t, engine, 1, routes) + + rawRoute, ok := engine.GetRoute(routeKey(key)) + require.True(t, ok) + require.Equal(t, distribution.RouteStateWriteFenced, rawRoute.State) + auxStart, auxEnd, ok := s3BucketAuxiliaryRouteRange(key) + require.True(t, ok) + auxRoutes := engine.GetIntersectingRoutes(auxStart, auxEnd) + require.NotEmpty(t, auxRoutes) + require.Equal(t, distribution.RouteStateActive, auxRoutes[0].State) + + fsm := newComposed1FSM(t, engine, 1) + err := fsm.handleRawRequest(context.Background(), &pb.Request{ + Mutations: []*pb.Mutation{{Op: pb.Op_PUT, Key: key, Value: []byte("meta")}}, + }, 100) + require.NoError(t, err) +} + func TestFSMComposed1UsesS3BucketAuxiliaryRouteOwner(t *testing.T) { t.Parallel() diff --git a/kv/sharded_coordinator.go b/kv/sharded_coordinator.go index e4141b871..7e21176f5 100644 --- a/kv/sharded_coordinator.go +++ b/kv/sharded_coordinator.go @@ -1267,18 +1267,18 @@ func (c *ShardedCoordinator) rejectWriteFencedPointKey(key []byte) error { if c.partitionResolverRecognisesPointKey(key) { return nil } - rkey := routeKey(key) - if route, ok := c.engine.GetRoute(rkey); ok && route.State == distribution.RouteStateWriteFenced { - return errors.Wrapf(ErrRouteWriteFenced, "key %q routeKey %q", key, rkey) - } start, end, ok := s3BucketAuxiliaryRouteRange(key) - if !ok { + if ok { + for _, route := range c.engine.GetIntersectingRoutes(start, end) { + if route.State == distribution.RouteStateWriteFenced { + return errors.Wrapf(ErrRouteWriteFenced, "key %q route range [%q,%q)", key, start, end) + } + } return nil } - for _, route := range c.engine.GetIntersectingRoutes(start, end) { - if route.State == distribution.RouteStateWriteFenced { - return errors.Wrapf(ErrRouteWriteFenced, "key %q route range [%q,%q)", key, start, end) - } + rkey := routeKey(key) + if route, ok := c.engine.GetRoute(rkey); ok && route.State == distribution.RouteStateWriteFenced { + return errors.Wrapf(ErrRouteWriteFenced, "key %q routeKey %q", key, rkey) } return nil } diff --git a/kv/sharded_coordinator_del_prefix_test.go b/kv/sharded_coordinator_del_prefix_test.go index 52a9aa75d..3b567a3a4 100644 --- a/kv/sharded_coordinator_del_prefix_test.go +++ b/kv/sharded_coordinator_del_prefix_test.go @@ -560,6 +560,43 @@ func TestShardedCoordinatorIgnoresRawRouteFloorForS3BucketAuxiliaryWrite(t *test require.Len(t, g2Txn.requests, 1) } +func TestShardedCoordinatorIgnoresRawRouteFenceForS3BucketAuxiliaryWrite(t *testing.T) { + t.Parallel() + + const bucket = "bucket-a" + key := s3keys.BucketMetaKey(bucket) + engine := distribution.NewEngine() + routes := s3BucketAuxiliaryStagedRoutes(bucket, 1, 2) + routes[2].State = distribution.RouteStateWriteFenced + require.NoError(t, engine.ApplySnapshot(distribution.CatalogSnapshot{ + Version: 1, + Routes: routes, + })) + + rawRoute, ok := engine.GetRoute(routeKey(key)) + require.True(t, ok) + require.Equal(t, distribution.RouteStateWriteFenced, rawRoute.State) + auxStart, auxEnd, ok := s3BucketAuxiliaryRouteRange(key) + require.True(t, ok) + auxRoutes := engine.GetIntersectingRoutes(auxStart, auxEnd) + require.NotEmpty(t, auxRoutes) + require.Equal(t, distribution.RouteStateActive, auxRoutes[0].State) + + g1Txn := &recordingTransactional{} + g2Txn := &recordingTransactional{responses: []*TransactionResponse{{CommitIndex: 22}}} + coord := NewShardedCoordinator(engine, map[uint64]*ShardGroup{ + 1: {Txn: g1Txn}, + 2: {Txn: g2Txn}, + }, 1, NewHLC(), nil) + + _, err := coord.Dispatch(context.Background(), &OperationGroup[OP]{ + Elems: []*Elem[OP]{{Op: Put, Key: key, Value: []byte("meta")}}, + }) + require.NoError(t, err) + require.Empty(t, g1Txn.requests) + require.Len(t, g2Txn.requests, 1) +} + func TestShardedCoordinatorRejectsS3BucketAuxiliaryPointWriteAtMigrationTimestampFloor(t *testing.T) { t.Parallel() From b16dd62cb6219e1bfce13eafc4e9b13239a70dd8 Mon Sep 17 00:00:00 2001 From: bootjp Date: Mon, 31 Aug 2026 23:59:56 +0900 Subject: [PATCH 42/58] migration: narrow auxiliary write fences --- adapter/redis_lua_context.go | 11 +++- adapter/redis_lua_negative_type_cache_test.go | 24 +++++++ kv/fsm.go | 33 ++++++---- kv/fsm_migration_fence_test.go | 60 ++++++++++++++++++ kv/sharded_coordinator.go | 20 +++--- kv/sharded_coordinator_del_prefix_test.go | 62 +++++++++++++++++++ 6 files changed, 185 insertions(+), 25 deletions(-) diff --git a/adapter/redis_lua_context.go b/adapter/redis_lua_context.go index 2a3cbac76..2938ec764 100644 --- a/adapter/redis_lua_context.go +++ b/adapter/redis_lua_context.go @@ -3808,11 +3808,18 @@ func (c *luaScriptContext) rawStartTypeForCommitPlan(ctx context.Context, key [] if err != nil { return redisTypeNone, err } + c.rememberRawTypeAtStart(k, typ) + return typ, nil +} + +func (c *luaScriptContext) rememberRawTypeAtStart(key string, typ redisValueType) { if c.rawTypeAtStart == nil { c.rawTypeAtStart = map[string]redisValueType{} } - c.rawTypeAtStart[k] = typ - return typ, nil + if _, ok := c.rawTypeAtStart[key]; !ok && len(c.rawTypeAtStart) >= maxNegativeTypeCacheEntries { + return + } + c.rawTypeAtStart[key] = typ } func luaWideFenceReadKeysForPlan(key []byte, finalType, startType redisValueType, preserveExisting bool) [][]byte { diff --git a/adapter/redis_lua_negative_type_cache_test.go b/adapter/redis_lua_negative_type_cache_test.go index f44aec35a..2085543e3 100644 --- a/adapter/redis_lua_negative_type_cache_test.go +++ b/adapter/redis_lua_negative_type_cache_test.go @@ -177,3 +177,27 @@ func TestLuaNegativeTypeCache_BoundedSize(t *testing.T) { require.Equal(t, maxNegativeTypeCacheEntries, len(sc.negativeType), "fallback probe must NOT grow the bounded cache") } + +func TestLuaRawTypeAtStartCache_BoundedSize(t *testing.T) { + t.Parallel() + + sc := &luaScriptContext{rawTypeAtStart: map[string]redisValueType{}} + + const overflow = 50 + total := maxNegativeTypeCacheEntries + overflow + for i := 0; i < total; i++ { + sc.rememberRawTypeAtStart(fmt.Sprintf("lua:rawtype:cap:%d", i), redisTypeNone) + } + require.Equal(t, maxNegativeTypeCacheEntries, len(sc.rawTypeAtStart), + "rawTypeAtStart map must be capped at maxNegativeTypeCacheEntries") + + cachedKey := "lua:rawtype:cap:0" + require.Equal(t, redisTypeNone, sc.rawTypeAtStart[cachedKey]) + + overflowKey := fmt.Sprintf("lua:rawtype:cap:%d", maxNegativeTypeCacheEntries+1) + _, ok := sc.rawTypeAtStart[overflowKey] + require.False(t, ok, "a key probed after the cap must fall back to a later server probe") + sc.rememberRawTypeAtStart(overflowKey, redisTypeNone) + require.Equal(t, maxNegativeTypeCacheEntries, len(sc.rawTypeAtStart), + "fallback raw type probe must NOT grow the bounded cache") +} diff --git a/kv/fsm.go b/kv/fsm.go index 010152196..0c93f29e6 100644 --- a/kv/fsm.go +++ b/kv/fsm.go @@ -695,7 +695,8 @@ func (f *kvFSM) verifyRouteNotFencedForKey(key []byte) error { return nil } if start, end, ok := s3BucketAuxiliaryRouteRange(key); ok { - if snap.WriteFencedIntersects(start, end) { + route, found := s3BucketAuxiliaryOwnerRouteFromRange(start, end, snap.IntersectingRoutes(start, end)) + if found && route.State == distribution.RouteStateWriteFenced { return errors.Wrapf(ErrRouteWriteFenced, "key %q route range [%q,%q)", key, start, end) } return nil @@ -731,10 +732,9 @@ func (f *kvFSM) verifyRouteWriteTimestampFloorForKey(key []byte, commitTS uint64 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 - } + route, found := s3BucketAuxiliaryOwnerRouteFromRange(start, end, snap.IntersectingRoutes(start, end)) + if found { + return verifyRouteWriteTimestampFloorForRange(route, key, start, end, commitTS) } return nil } @@ -1186,13 +1186,8 @@ func verifyWriteFenceFromSnapshot(mutations []*pb.Mutation, writeFenceBypassKeys if _, ok := bypassKeys[string(mut.Key)]; ok { continue } - if start, end, ok := s3BucketAuxiliaryRouteRange(mut.Key); ok { - if snap.WriteFencedIntersects(start, end) { - return errors.Wrapf(ErrRouteWriteFenced, - "%s-version v=%d: key %q route range [%q,%q)", - phase, snapVer, mut.Key, start, end) - } - continue + if checked, err := verifyS3BucketAuxiliaryWriteFenceFromSnapshot(mut.Key, snap, snapVer, phase); checked || err != nil { + return err } rKey := routeKey(mut.Key) if snap.WriteFencedForKey(rKey) { @@ -1204,6 +1199,20 @@ func verifyWriteFenceFromSnapshot(mutations []*pb.Mutation, writeFenceBypassKeys return nil } +func verifyS3BucketAuxiliaryWriteFenceFromSnapshot(key []byte, snap RouteSnapshot, snapVer uint64, phase string) (bool, error) { + start, end, ok := s3BucketAuxiliaryRouteRange(key) + if !ok { + return false, nil + } + route, found := s3BucketAuxiliaryOwnerRouteFromRange(start, end, snap.IntersectingRoutes(start, end)) + if found && route.State == distribution.RouteStateWriteFenced { + return true, errors.Wrapf(ErrRouteWriteFenced, + "%s-version v=%d: key %q route range [%q,%q)", + phase, snapVer, key, start, end) + } + return true, nil +} + func writeFenceBypassKeySet(keys [][]byte) map[string]struct{} { if len(keys) == 0 { return nil diff --git a/kv/fsm_migration_fence_test.go b/kv/fsm_migration_fence_test.go index df7b4846b..fa1254777 100644 --- a/kv/fsm_migration_fence_test.go +++ b/kv/fsm_migration_fence_test.go @@ -53,6 +53,18 @@ func s3BucketAuxiliaryFenceRoutes(bucket string, rawGroupID, fencedGroupID uint6 } } +func s3BucketAuxiliarySplitRoutes(bucket string, rawGroupID, ownerGroupID, splitGroupID uint64) []distribution.RouteDescriptor { + start := s3keys.RoutePrefixForBucketAnyGeneration(bucket) + split := append(append([]byte(nil), start...), 'm') + end := prefixScanEnd(start) + return []distribution.RouteDescriptor{ + {RouteID: 1, Start: []byte(""), End: start, GroupID: rawGroupID, State: distribution.RouteStateActive}, + {RouteID: 2, Start: start, End: split, GroupID: ownerGroupID, State: distribution.RouteStateActive}, + {RouteID: 3, Start: split, End: end, GroupID: splitGroupID, State: distribution.RouteStateActive}, + {RouteID: 4, Start: end, End: nil, GroupID: rawGroupID, State: distribution.RouteStateActive}, + } +} + func newS3BucketAuxiliaryWriteFencedFSM(t *testing.T, bucket string) *kvFSM { t.Helper() @@ -381,6 +393,30 @@ func TestFSMIgnoresRawRouteFenceForS3BucketAuxiliaryWrite(t *testing.T) { require.NoError(t, err) } +func TestFSMIgnoresNonOwnerS3BucketAuxiliaryFenceForPointWrite(t *testing.T) { + t.Parallel() + + const bucket = "bucket-a" + key := s3keys.BucketMetaKey(bucket) + engine := distribution.NewEngine() + routes := s3BucketAuxiliarySplitRoutes(bucket, 5, 1, 1) + routes[2].State = distribution.RouteStateWriteFenced + applyComposed1Snapshot(t, engine, 1, routes) + + auxStart, auxEnd, ok := s3BucketAuxiliaryRouteRange(key) + require.True(t, ok) + auxRoutes := engine.GetIntersectingRoutes(auxStart, auxEnd) + require.Len(t, auxRoutes, 2) + require.Equal(t, distribution.RouteStateActive, auxRoutes[0].State) + require.Equal(t, distribution.RouteStateWriteFenced, auxRoutes[1].State) + + fsm := newComposed1FSM(t, engine, 1) + err := fsm.handleRawRequest(context.Background(), &pb.Request{ + Mutations: []*pb.Mutation{{Op: pb.Op_PUT, Key: key, Value: []byte("meta")}}, + }, 100) + require.NoError(t, err) +} + func TestFSMComposed1UsesS3BucketAuxiliaryRouteOwner(t *testing.T) { t.Parallel() @@ -430,6 +466,30 @@ func TestFSMIgnoresRawRouteFloorForS3BucketAuxiliaryWrite(t *testing.T) { require.NoError(t, err) } +func TestFSMIgnoresNonOwnerS3BucketAuxiliaryFloorForPointWrite(t *testing.T) { + t.Parallel() + + const bucket = "bucket-b" + key := s3keys.BucketMetaKey(bucket) + engine := distribution.NewEngine() + routes := s3BucketAuxiliarySplitRoutes(bucket, 1, 1, 1) + routes[2].MinWriteTSExclusive = ^uint64(0) + applyComposed1Snapshot(t, engine, 1, routes) + + auxStart, auxEnd, ok := s3BucketAuxiliaryRouteRange(key) + require.True(t, ok) + auxRoutes := engine.GetIntersectingRoutes(auxStart, auxEnd) + require.Len(t, auxRoutes, 2) + require.Zero(t, auxRoutes[0].MinWriteTSExclusive) + require.Equal(t, ^uint64(0), auxRoutes[1].MinWriteTSExclusive) + + fsm := newComposed1FSM(t, engine, 1) + err := fsm.handleRawRequest(context.Background(), &pb.Request{ + Mutations: []*pb.Mutation{{Op: pb.Op_PUT, Key: key, Value: []byte("meta")}}, + }, 100) + require.NoError(t, err) +} + func TestFSMRejectsCurrentWriteFencedDelPrefix(t *testing.T) { t.Parallel() diff --git a/kv/sharded_coordinator.go b/kv/sharded_coordinator.go index 7e21176f5..d31d32ed1 100644 --- a/kv/sharded_coordinator.go +++ b/kv/sharded_coordinator.go @@ -1269,10 +1269,9 @@ func (c *ShardedCoordinator) rejectWriteFencedPointKey(key []byte) error { } start, end, ok := s3BucketAuxiliaryRouteRange(key) if ok { - for _, route := range c.engine.GetIntersectingRoutes(start, end) { - if route.State == distribution.RouteStateWriteFenced { - return errors.Wrapf(ErrRouteWriteFenced, "key %q route range [%q,%q)", key, start, end) - } + route, found := s3BucketAuxiliaryOwnerRouteFromRange(start, end, c.engine.GetIntersectingRoutes(start, end)) + if found && route.State == distribution.RouteStateWriteFenced { + return errors.Wrapf(ErrRouteWriteFenced, "key %q route range [%q,%q)", key, start, end) } return nil } @@ -1383,13 +1382,12 @@ func (c *ShardedCoordinator) rejectS3BucketAuxiliaryWriteTimestampFloor(key []by if !ok { return false, nil } - for _, route := range c.engine.GetIntersectingRoutes(start, end) { - if route.MinWriteTSExclusive != 0 && commitTS <= route.MinWriteTSExclusive { - return true, errors.Join( - errors.Wrapf(ErrRouteWriteTimestampTooLow, "key %q route range [%q,%q) commit_ts=%d floor=%d", key, start, end, commitTS, route.MinWriteTSExclusive), - store.NewWriteConflictError(key), - ) - } + route, found := s3BucketAuxiliaryOwnerRouteFromRange(start, end, c.engine.GetIntersectingRoutes(start, end)) + if found && route.MinWriteTSExclusive != 0 && commitTS <= route.MinWriteTSExclusive { + return true, errors.Join( + errors.Wrapf(ErrRouteWriteTimestampTooLow, "key %q route range [%q,%q) commit_ts=%d floor=%d", key, start, end, commitTS, route.MinWriteTSExclusive), + store.NewWriteConflictError(key), + ) } return true, nil } diff --git a/kv/sharded_coordinator_del_prefix_test.go b/kv/sharded_coordinator_del_prefix_test.go index 3b567a3a4..0cfda39b5 100644 --- a/kv/sharded_coordinator_del_prefix_test.go +++ b/kv/sharded_coordinator_del_prefix_test.go @@ -597,6 +597,37 @@ func TestShardedCoordinatorIgnoresRawRouteFenceForS3BucketAuxiliaryWrite(t *test require.Len(t, g2Txn.requests, 1) } +func TestShardedCoordinatorIgnoresNonOwnerS3BucketAuxiliaryFenceForPointWrite(t *testing.T) { + t.Parallel() + + const bucket = "bucket-a" + key := s3keys.BucketMetaKey(bucket) + engine := distribution.NewEngine() + routes := s3BucketAuxiliarySplitRoutes(bucket, 1, 2, 3) + routes[2].State = distribution.RouteStateWriteFenced + require.NoError(t, engine.ApplySnapshot(distribution.CatalogSnapshot{ + Version: 1, + Routes: routes, + })) + + g1Txn := &recordingTransactional{} + g2Txn := &recordingTransactional{responses: []*TransactionResponse{{CommitIndex: 22}}} + g3Txn := &recordingTransactional{} + coord := NewShardedCoordinator(engine, map[uint64]*ShardGroup{ + 1: {Txn: g1Txn}, + 2: {Txn: g2Txn}, + 3: {Txn: g3Txn}, + }, 1, NewHLC(), nil) + + _, err := coord.Dispatch(context.Background(), &OperationGroup[OP]{ + Elems: []*Elem[OP]{{Op: Put, Key: key, Value: []byte("meta")}}, + }) + require.NoError(t, err) + require.Empty(t, g1Txn.requests) + require.Len(t, g2Txn.requests, 1) + require.Empty(t, g3Txn.requests) +} + func TestShardedCoordinatorRejectsS3BucketAuxiliaryPointWriteAtMigrationTimestampFloor(t *testing.T) { t.Parallel() @@ -633,6 +664,37 @@ func TestShardedCoordinatorRejectsS3BucketAuxiliaryPointWriteAtMigrationTimestam } } +func TestShardedCoordinatorIgnoresNonOwnerS3BucketAuxiliaryFloorForPointWrite(t *testing.T) { + t.Parallel() + + const bucket = "bucket-b" + key := s3keys.BucketMetaKey(bucket) + engine := distribution.NewEngine() + routes := s3BucketAuxiliarySplitRoutes(bucket, 1, 2, 3) + routes[2].MinWriteTSExclusive = ^uint64(0) + require.NoError(t, engine.ApplySnapshot(distribution.CatalogSnapshot{ + Version: 1, + Routes: routes, + })) + + g1Txn := &recordingTransactional{} + g2Txn := &recordingTransactional{responses: []*TransactionResponse{{CommitIndex: 22}}} + g3Txn := &recordingTransactional{} + coord := NewShardedCoordinator(engine, map[uint64]*ShardGroup{ + 1: {Txn: g1Txn}, + 2: {Txn: g2Txn}, + 3: {Txn: g3Txn}, + }, 1, NewHLC(), nil) + + _, err := coord.Dispatch(context.Background(), &OperationGroup[OP]{ + Elems: []*Elem[OP]{{Op: Put, Key: key, Value: []byte("meta")}}, + }) + require.NoError(t, err) + require.Empty(t, g1Txn.requests) + require.Len(t, g2Txn.requests, 1) + require.Empty(t, g3Txn.requests) +} + func TestShardedCoordinatorRejectsDelPrefixIntersectingWriteFencedRoute(t *testing.T) { t.Parallel() From ace20aef55b2785879b69ad359460c64ca44ab66 Mon Sep 17 00:00:00 2001 From: bootjp Date: Tue, 1 Sep 2026 00:12:04 +0900 Subject: [PATCH 43/58] migration: pin apply-time floor checks --- kv/fsm.go | 91 +++++++++++++++++++++++---------- kv/fsm_migration_fence_test.go | 46 ++++++++++++++--- kv/fsm_reserved_control_test.go | 6 +-- 3 files changed, 106 insertions(+), 37 deletions(-) diff --git a/kv/fsm.go b/kv/fsm.go index 0c93f29e6..119ced56d 100644 --- a/kv/fsm.go +++ b/kv/fsm.go @@ -544,14 +544,18 @@ func (f *kvFSM) handleRawRequest(ctx context.Context, r *pb.Request, commitTS ui if err := f.verifyWriteFence(r); err != nil { return err } + floorSnap, err := f.routeFloorSnapshotForRequest(r) + if err != nil { + return err + } // DEL_PREFIX mutations are handled by the store's DeletePrefixAt which // scans and writes tombstones locally. A DEL_PREFIX request must be the // sole mutation in a request (enforced by the coordinator's toRawRequest). if hasDelPrefix, prefix := extractDelPrefix(r.Mutations); hasDelPrefix { - return f.handleDelPrefix(ctx, prefix, commitTS) + return f.handleDelPrefixWithFloorSnapshot(ctx, prefix, commitTS, floorSnap) } - if err := f.validateRawMutationsForApply(ctx, r, commitTS); err != nil { + if err := f.validateRawMutationsForApply(ctx, r, commitTS, floorSnap); err != nil { return err } @@ -568,17 +572,17 @@ 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 { +func (f *kvFSM) validateRawMutationsForApply(ctx context.Context, r *pb.Request, commitTS uint64, floorSnap RouteSnapshot) error { bypassKeys := writeFenceBypassKeySet(r.GetWriteFenceBypassKeys()) for _, mut := range r.GetMutations() { - if err := f.validateRawMutationForApply(ctx, mut, bypassKeys, commitTS); err != nil { + if err := f.validateRawMutationForApply(ctx, mut, bypassKeys, commitTS, floorSnap); err != nil { return err } } return nil } -func (f *kvFSM) validateRawMutationForApply(ctx context.Context, mut *pb.Mutation, writeFenceBypassKeys map[string]struct{}, commitTS uint64) error { +func (f *kvFSM) validateRawMutationForApply(ctx context.Context, mut *pb.Mutation, writeFenceBypassKeys map[string]struct{}, commitTS uint64, floorSnap RouteSnapshot) error { if mut == nil || len(mut.Key) == 0 { return errors.WithStack(ErrInvalidRequest) } @@ -598,7 +602,7 @@ func (f *kvFSM) validateRawMutationForApply(ctx context.Context, mut *pb.Mutatio return err } } - if err := f.verifyRouteWriteTimestampFloorForKey(mut.Key, commitTS); err != nil { + if err := verifyRouteWriteTimestampFloorForKeyFromSnapshot(floorSnap, mut.Key, commitTS); err != nil { return err } if err := f.assertNoConflictingTxnLock(ctx, mut.Key, nil, 0); err != nil { @@ -621,6 +625,10 @@ 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 { + return f.handleDelPrefixWithFloorSnapshot(ctx, prefix, commitTS, nil) +} + +func (f *kvFSM) handleDelPrefixWithFloorSnapshot(ctx context.Context, prefix []byte, commitTS uint64, floorSnap RouteSnapshot) error { // DEL_PREFIX never reaches validateRawMutationsForApply, so the control // namespaces are gated here as well. A prefix is refused both when it sits // inside one and when it is broad enough to sweep one up. @@ -630,7 +638,7 @@ func (f *kvFSM) handleDelPrefix(ctx context.Context, prefix []byte, commitTS uin if err := f.verifyRouteNotFencedForPrefix(prefix); err != nil { return err } - if err := f.verifyRouteWriteTimestampFloorForPrefix(prefix, commitTS); err != nil { + if err := verifyRouteWriteTimestampFloorForPrefixFromSnapshot(floorSnap, prefix, commitTS); err != nil { return err } deletes := []store.PrefixDelete{{ @@ -650,6 +658,21 @@ func (f *kvFSM) handleDelPrefix(ctx context.Context, prefix []byte, commitTS uin return nil } +func (f *kvFSM) routeFloorSnapshotForRequest(r *pb.Request) (RouteSnapshot, error) { + if f.routes == nil || f.shardGroupID == 0 { + return nil, nil + } + observedVer, pinned := DecodeObservedRouteVersion(r.GetObservedRouteVersion()) + if !pinned { + return nil, nil + } + snap, ok := f.routes.SnapshotAt(observedVer) + if !ok { + return nil, errors.WithStack(ErrComposed1VersionGCd) + } + return snap, nil +} + type fsmStagedVisibilityPrefixDelete struct { prefix []byte excludePrefix []byte @@ -723,12 +746,8 @@ func (f *kvFSM) verifyRouteNotFencedForPrefix(prefix []byte) error { 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 { +func verifyRouteWriteTimestampFloorForKeyFromSnapshot(snap RouteSnapshot, key []byte, commitTS uint64) error { + if snap == nil || commitTS == 0 { return nil } if start, end, ok := s3BucketAuxiliaryRouteRange(key); ok { @@ -747,24 +766,20 @@ func (f *kvFSM) verifyRouteWriteTimestampFloorForKey(key []byte, commitTS uint64 return nil } -func (f *kvFSM) verifyRouteWriteTimestampFloorsForMutations(muts []*pb.Mutation, commitTS uint64) error { +func verifyRouteWriteTimestampFloorsForMutationsFromSnapshot(snap RouteSnapshot, muts []*pb.Mutation, commitTS uint64) error { for _, mut := range muts { if mut == nil || len(mut.Key) == 0 || isTxnInternalKey(mut.Key) { continue } - if err := f.verifyRouteWriteTimestampFloorForKey(mut.Key, commitTS); err != nil { + if err := verifyRouteWriteTimestampFloorForKeyFromSnapshot(snap, 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 { +func verifyRouteWriteTimestampFloorForPrefixFromSnapshot(snap RouteSnapshot, prefix []byte, commitTS uint64) error { + if snap == nil || commitTS == 0 { return nil } start, end := routePrefixRange(prefix) @@ -1323,7 +1338,7 @@ func (f *kvFSM) handlePrepareRequest(ctx context.Context, r *pb.Request) error { } startTS := r.Ts - uniq, err := f.uniqueMutationsAboveFloor(muts, startTS) + uniq, err := f.uniqueMutationsAboveFloorForRequest(r, muts, startTS) if err != nil { return err } @@ -1393,7 +1408,7 @@ func (f *kvFSM) handleOnePhaseTxnRequest(ctx context.Context, r *pb.Request, com return nil } - uniq, err := f.uniqueMutationsAboveFloor(muts, commitTS) + uniq, err := f.uniqueMutationsAboveFloorForRequest(r, muts, commitTS) if err != nil { return err } @@ -1418,6 +1433,18 @@ func uniqueTxnMutations(muts []*pb.Mutation) ([]*pb.Mutation, error) { } func (f *kvFSM) uniqueMutationsAboveFloor(muts []*pb.Mutation, commitTS uint64) ([]*pb.Mutation, error) { + return f.uniqueMutationsAboveFloorWithSnapshot(muts, commitTS, nil) +} + +func (f *kvFSM) uniqueMutationsAboveFloorForRequest(r *pb.Request, muts []*pb.Mutation, commitTS uint64) ([]*pb.Mutation, error) { + floorSnap, err := f.routeFloorSnapshotForRequest(r) + if err != nil { + return nil, err + } + return f.uniqueMutationsAboveFloorWithSnapshot(muts, commitTS, floorSnap) +} + +func (f *kvFSM) uniqueMutationsAboveFloorWithSnapshot(muts []*pb.Mutation, commitTS uint64, floorSnap RouteSnapshot) ([]*pb.Mutation, error) { uniq, err := uniqueMutations(muts) if err != nil { return nil, err @@ -1425,7 +1452,7 @@ func (f *kvFSM) uniqueMutationsAboveFloor(muts []*pb.Mutation, commitTS uint64) if err := rejectReservedControlMutations(uniq); err != nil { return nil, err } - if err := f.verifyRouteWriteTimestampFloorsForMutations(uniq, commitTS); err != nil { + if err := verifyRouteWriteTimestampFloorsForMutationsFromSnapshot(floorSnap, uniq, commitTS); err != nil { return nil, err } return uniq, nil @@ -1456,6 +1483,18 @@ func rejectReservedControlMutations(muts []*pb.Mutation) error { } func (f *kvFSM) uniqueTxnMutationsAboveFloor(muts []*pb.Mutation, commitTS uint64) ([]*pb.Mutation, error) { + return f.uniqueTxnMutationsAboveFloorWithSnapshot(muts, commitTS, nil) +} + +func (f *kvFSM) uniqueTxnMutationsAboveFloorForRequest(r *pb.Request, muts []*pb.Mutation, commitTS uint64) ([]*pb.Mutation, error) { + floorSnap, err := f.routeFloorSnapshotForRequest(r) + if err != nil { + return nil, err + } + return f.uniqueTxnMutationsAboveFloorWithSnapshot(muts, commitTS, floorSnap) +} + +func (f *kvFSM) uniqueTxnMutationsAboveFloorWithSnapshot(muts []*pb.Mutation, commitTS uint64, floorSnap RouteSnapshot) ([]*pb.Mutation, error) { uniq, err := uniqueTxnMutations(muts) if err != nil { return nil, err @@ -1463,7 +1502,7 @@ func (f *kvFSM) uniqueTxnMutationsAboveFloor(muts []*pb.Mutation, commitTS uint6 if err := rejectReservedControlMutations(uniq); err != nil { return nil, err } - if err := f.verifyRouteWriteTimestampFloorsForMutations(uniq, commitTS); err != nil { + if err := verifyRouteWriteTimestampFloorsForMutationsFromSnapshot(floorSnap, uniq, commitTS); err != nil { return nil, err } return uniq, nil @@ -1552,7 +1591,7 @@ func (f *kvFSM) handleCommitRequest(ctx context.Context, r *pb.Request) error { if err != nil { return err } - uniq, err := f.uniqueTxnMutationsAboveFloor(muts, commitTS) + uniq, err := f.uniqueTxnMutationsAboveFloorForRequest(r, muts, commitTS) if err != nil { return err } diff --git a/kv/fsm_migration_fence_test.go b/kv/fsm_migration_fence_test.go index fa1254777..e656f5021 100644 --- a/kv/fsm_migration_fence_test.go +++ b/kv/fsm_migration_fence_test.go @@ -126,6 +126,7 @@ func TestFSMWriteFenceBypassRejectsRawWriteAtBypassedRouteFloor(t *testing.T) { fsm := newWriteFloorFSM(t) key := []byte("!sqs|msg|data|p|partitioned-key") err := fsm.handleRawRequest(context.Background(), &pb.Request{ + ObservedRouteVersion: 1, WriteFenceBypassKeys: [][]byte{key}, Mutations: []*pb.Mutation{{Op: pb.Op_PUT, Key: key, Value: []byte("v")}}, }, 100) @@ -165,6 +166,7 @@ func TestFSMWriteFenceBypassRejectsPinnedTxnAtBypassedRouteFloor(t *testing.T) { IsTxn: true, Phase: pb.Phase_PREPARE, Ts: 100, + ObservedRouteVersion: 1, WriteFenceBypassKeys: [][]byte{key}, Mutations: []*pb.Mutation{ {Op: pb.Op_PUT, Key: []byte(txnMetaPrefix), Value: EncodeTxnMeta(TxnMeta{PrimaryKey: key, LockTTLms: defaultTxnLockTTLms})}, @@ -605,7 +607,8 @@ func TestFSMRejectsRawPointWriteAtMigrationTimestampFloorDuringApply(t *testing. 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")}}, + ObservedRouteVersion: 1, + Mutations: []*pb.Mutation{{Op: pb.Op_PUT, Key: []byte("z"), Value: []byte("replayed")}}, }, 100) require.ErrorIs(t, err, ErrRouteWriteTimestampTooLow) _, getErr := fsm.store.GetAt(ctx, []byte("z"), ^uint64(0)) @@ -620,7 +623,8 @@ func TestFSMRejectsDelPrefixAtMigrationTimestampFloorDuringApply(t *testing.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")}}, + ObservedRouteVersion: 1, + Mutations: []*pb.Mutation{{Op: pb.Op_DEL_PREFIX, Key: []byte("z")}}, }, 100) require.ErrorIs(t, err, ErrRouteWriteTimestampTooLow) @@ -635,9 +639,10 @@ func TestFSMRejectsOnePhaseTxnAtMigrationTimestampFloorDuringApply(t *testing.T) ctx := context.Background() fsm := newWriteFloorFSM(t) req := &pb.Request{ - IsTxn: true, - Phase: pb.Phase_NONE, - Ts: 90, + IsTxn: true, + Phase: pb.Phase_NONE, + Ts: 90, + ObservedRouteVersion: 1, Mutations: []*pb.Mutation{ {Op: pb.Op_PUT, Key: []byte(txnMetaPrefix), Value: EncodeTxnMeta(TxnMeta{PrimaryKey: []byte("z"), CommitTS: 100})}, {Op: pb.Op_PUT, Key: []byte("z"), Value: []byte("low")}, @@ -655,9 +660,10 @@ func TestFSMRejectsPrepareAtMigrationTimestampFloorDuringApply(t *testing.T) { ctx := context.Background() fsm := newWriteFloorFSM(t) prepare := &pb.Request{ - IsTxn: true, - Phase: pb.Phase_PREPARE, - Ts: 90, + IsTxn: true, + Phase: pb.Phase_PREPARE, + Ts: 90, + ObservedRouteVersion: 1, Mutations: []*pb.Mutation{ {Op: pb.Op_PUT, Key: []byte(txnMetaPrefix), Value: EncodeTxnMeta(TxnMeta{PrimaryKey: []byte("z"), LockTTLms: defaultTxnLockTTLms})}, {Op: pb.Op_PUT, Key: []byte("z"), Value: []byte("v")}, @@ -665,3 +671,27 @@ func TestFSMRejectsPrepareAtMigrationTimestampFloorDuringApply(t *testing.T) { } require.ErrorIs(t, fsm.handleTxnRequest(ctx, prepare, 90), ErrRouteWriteTimestampTooLow) } + +func TestFSMTimestampFloorUsesObservedSnapshotDuringApply(t *testing.T) { + t.Parallel() + + ctx := context.Background() + engine := distribution.NewEngine() + applyComposed1Snapshot(t, engine, 1, []distribution.RouteDescriptor{ + {RouteID: 1, Start: []byte(""), End: nil, GroupID: 1, State: distribution.RouteStateActive}, + }) + fsm := newComposed1FSM(t, engine, 1) + applyComposed1Snapshot(t, engine, 2, []distribution.RouteDescriptor{ + {RouteID: 1, Start: []byte(""), End: nil, GroupID: 1, State: distribution.RouteStateActive, MinWriteTSExclusive: ^uint64(0)}, + }) + + err := fsm.handleRawRequest(ctx, &pb.Request{ + ObservedRouteVersion: 1, + Mutations: []*pb.Mutation{{Op: pb.Op_PUT, Key: []byte("z"), Value: []byte("proposed-before-floor")}}, + }, 100) + require.NoError(t, err) + + got, getErr := fsm.store.GetAt(ctx, []byte("z"), ^uint64(0)) + require.NoError(t, getErr) + require.Equal(t, []byte("proposed-before-floor"), got) +} diff --git a/kv/fsm_reserved_control_test.go b/kv/fsm_reserved_control_test.go index bedabdebf..6a5077f64 100644 --- a/kv/fsm_reserved_control_test.go +++ b/kv/fsm_reserved_control_test.go @@ -32,13 +32,13 @@ func TestValidateRawMutationRejectsReservedControlKeys(t *testing.T) { []byte("!migwrite|7"), []byte("!migfence|7"), } { - err := f.validateRawMutationForApply(ctx, &pb.Mutation{Op: pb.Op_PUT, Key: key, Value: []byte("v")}, nil, 10) + err := f.validateRawMutationForApply(ctx, &pb.Mutation{Op: pb.Op_PUT, Key: key, Value: []byte("v")}, nil, 10, nil) require.ErrorIs(t, err, ErrInvalidRequest, "key %q must be refused", key) } // Ordinary user keys are unaffected. require.NoError(t, f.validateRawMutationForApply(ctx, - &pb.Mutation{Op: pb.Op_PUT, Key: []byte("user-key"), Value: []byte("v")}, nil, 10)) + &pb.Mutation{Op: pb.Op_PUT, Key: []byte("user-key"), Value: []byte("v")}, nil, 10, nil)) } // DEL_PREFIX never reaches validateRawMutationsForApply, so it is gated in @@ -85,7 +85,7 @@ func TestValidateRawMutationRejectsStagedDataKeys(t *testing.T) { staged := distribution.MigrationStagedDataKey(7, []byte("user-key")) require.ErrorIs(t, f.validateRawMutationForApply(ctx, - &pb.Mutation{Op: pb.Op_PUT, Key: staged, Value: []byte("v")}, nil, 10), ErrInvalidRequest) + &pb.Mutation{Op: pb.Op_PUT, Key: staged, Value: []byte("v")}, nil, 10, nil), ErrInvalidRequest) require.ErrorIs(t, f.handleDelPrefix(ctx, distribution.MigrationStagedDataKey(7, []byte("user:")), 11), ErrInvalidRequest) } From 986823cb592a99542d569ee246e85f0cfdf5406a Mon Sep 17 00:00:00 2001 From: bootjp Date: Tue, 1 Sep 2026 00:24:42 +0900 Subject: [PATCH 44/58] migration: export filesystem auxiliary rows --- adapter/internal.go | 73 ++++++++++++++++++++++++++++-- adapter/internal_migration_test.go | 61 +++++++++++++++++++++++++ 2 files changed, 130 insertions(+), 4 deletions(-) diff --git a/adapter/internal.go b/adapter/internal.go index a160b91d1..6a8650cdb 100644 --- a/adapter/internal.go +++ b/adapter/internal.go @@ -8,6 +8,7 @@ import ( "github.com/bootjp/elastickv/distribution" "github.com/bootjp/elastickv/internal" + "github.com/bootjp/elastickv/internal/fskeys" "github.com/bootjp/elastickv/internal/raftengine" "github.com/bootjp/elastickv/internal/s3keys" "github.com/bootjp/elastickv/kv" @@ -577,9 +578,10 @@ func (i *Internal) exportRangeVersionsOptions(req *pb.ExportRangeVersionsRequest uint64(req.GetChunkBytes()), defaultMigrationExportChunkBytes, maxMigrationExportChunkBytes) maxScannedBytes := clampMigrationExportBound( req.GetMaxScannedBytes(), chunkBytes*defaultMigrationExportScanFactor, maxMigrationExportScanBytes) + startKey, endKey := migrationExportScanBounds(req) opts := store.ExportVersionsOptions{ - StartKey: req.GetRangeStart(), - EndKey: req.GetRangeEnd(), + StartKey: startKey, + EndKey: endKey, MinCommitTSExclusive: req.GetMinCommitTs(), MaxCommitTSInclusive: req.GetMaxCommitTs(), Cursor: req.GetCursor(), @@ -593,6 +595,68 @@ func (i *Internal) exportRangeVersionsOptions(req *pb.ExportRangeVersionsRequest return opts } +func migrationExportScanBounds(req *pb.ExportRangeVersionsRequest) ([]byte, []byte) { + start := bytes.Clone(req.GetRangeStart()) + end := bytes.Clone(req.GetRangeEnd()) + switch req.GetKeyFamily() { + case distribution.MigrationFamilyFilesystemChunk: + if bytes.HasPrefix(start, fskeys.ChunkAllPrefix()) { + return start, end + } + if scanStart, scanEnd, ok := filesystemChunkExportScanBounds(req.GetRouteStart(), req.GetRouteEnd()); ok { + return scanStart, scanEnd + } + return fskeys.ChunkAllPrefix(), prefixScanEnd(fskeys.ChunkAllPrefix()) + case distribution.MigrationFamilyFilesystemUsage: + if bytes.HasPrefix(start, fskeys.UsageRouteAllPrefix()) { + return start, end + } + return filesystemUsageExportScanBounds(req.GetRouteStart(), req.GetRouteEnd()) + default: + return start, end + } +} + +func filesystemChunkExportScanBounds(routeStart, routeEnd []byte) ([]byte, []byte, bool) { + routePrefix := fskeys.ChunkRouteAllPrefix() + routeDomainEnd := prefixScanEnd(routePrefix) + if !rangesIntersect(routeStart, routeEnd, routePrefix, routeDomainEnd) { + return fskeys.ChunkAllPrefix(), fskeys.ChunkAllPrefix(), true + } + rawPrefix := fskeys.ChunkAllPrefix() + rawDomainEnd := prefixScanEnd(rawPrefix) + start := rawPrefix + if len(routeStart) > 0 && bytes.Compare(routeStart, routePrefix) > 0 { + if !bytes.HasPrefix(routeStart, routePrefix) { + return nil, nil, false + } + start = append(bytes.Clone(rawPrefix), routeStart[len(routePrefix):]...) + } + end := rawDomainEnd + if len(routeEnd) > 0 && bytes.Compare(routeEnd, routeDomainEnd) < 0 { + if bytes.Compare(routeEnd, routePrefix) <= 0 { + return rawPrefix, rawPrefix, true + } + if !bytes.HasPrefix(routeEnd, routePrefix) { + return nil, nil, false + } + end = append(bytes.Clone(rawPrefix), routeEnd[len(routePrefix):]...) + } + return bytes.Clone(start), bytes.Clone(end), true +} + +func filesystemUsageExportScanBounds(routeStart, routeEnd []byte) ([]byte, []byte) { + start := fskeys.UsageRouteAllPrefix() + if len(routeStart) > 0 { + start = fskeys.UsageRouteKey(routeStart) + } + end := prefixScanEnd(fskeys.UsageRouteAllPrefix()) + if len(routeEnd) > 0 { + end = fskeys.UsageRouteKey(routeEnd) + } + return start, end +} + // clampMigrationExportBound resolves one export bound: unset takes the default, // anything above the hard ceiling is clamped down to it. func clampMigrationExportBound(requested, fallback, ceiling uint64) uint64 { @@ -628,10 +692,11 @@ func (i *Internal) migrationExportVersionFilter(req *pb.ExportRangeVersionsReque func migrationExportBracket(req *pb.ExportRangeVersionsRequest) distribution.MigrationBracket { excludeKnownInternal := req.GetExcludeKnownInternal() || req.GetKeyFamily() == distribution.MigrationFamilyUser + start, end := migrationExportScanBounds(req) return distribution.MigrationBracket{ Family: req.GetKeyFamily(), - Start: bytes.Clone(req.GetRangeStart()), - End: bytes.Clone(req.GetRangeEnd()), + Start: start, + End: end, ExcludeKnownInternal: excludeKnownInternal, ExcludePrefixes: cloneByteSlices(req.GetExcludePrefixes()), } diff --git a/adapter/internal_migration_test.go b/adapter/internal_migration_test.go index 18e577f52..d6a72032c 100644 --- a/adapter/internal_migration_test.go +++ b/adapter/internal_migration_test.go @@ -8,6 +8,7 @@ import ( "github.com/bootjp/elastickv/distribution" internalutil "github.com/bootjp/elastickv/internal" + "github.com/bootjp/elastickv/internal/fskeys" "github.com/bootjp/elastickv/internal/raftengine" "github.com/bootjp/elastickv/internal/s3keys" "github.com/bootjp/elastickv/kv" @@ -463,6 +464,66 @@ func TestInternalExportRangeVersionsUsesPartitionResolverGroup(t *testing.T) { }, stream.responses[0].GetVersions()) } +func TestInternalExportRangeVersionsDerivesFilesystemChunkScanBoundsFromRouteBounds(t *testing.T) { + t.Parallel() + + ctx := context.Background() + st := store.NewMVCCStore() + internal := NewInternalWithEngine(nil, mockInternalLeader{}, nil, nil, WithInternalStore(st)) + + routeKey := fskeys.ChunkRouteKey(10, 20) + inRouteChunk := fskeys.ChunkKey(10, 20, 3) + outRouteChunk := fskeys.ChunkKey(10, 21, 3) + require.NoError(t, st.PutAt(ctx, inRouteChunk, []byte("chunk"), 10, 0)) + require.NoError(t, st.PutAt(ctx, outRouteChunk, []byte("skip"), 10, 0)) + + stream := &captureExportRangeVersionsStream{ctx: ctx} + err := internal.ExportRangeVersions(&pb.ExportRangeVersionsRequest{ + MaxCommitTs: 20, + RouteStart: routeKey, + RouteEnd: testPrefixScanEnd(routeKey), + KeyFamily: distribution.MigrationFamilyFilesystemChunk, + RangeStart: routeKey, + RangeEnd: testPrefixScanEnd(routeKey), + MaxScannedBytes: 1 << 20, + }, stream) + require.NoError(t, err) + require.Len(t, stream.responses, 1) + require.Equal(t, []*pb.MVCCVersion{ + {Key: inRouteChunk, CommitTs: 10, Value: []byte("chunk"), KeyFamily: distribution.MigrationFamilyFilesystemChunk}, + }, stream.responses[0].GetVersions()) +} + +func TestInternalExportRangeVersionsDerivesFilesystemUsageScanBoundsFromRouteBounds(t *testing.T) { + t.Parallel() + + ctx := context.Background() + st := store.NewMVCCStore() + internal := NewInternalWithEngine(nil, mockInternalLeader{}, nil, nil, WithInternalStore(st)) + + routeKey := fskeys.ChunkRouteKey(10, 20) + inRouteUsage := fskeys.UsageRouteKey(routeKey) + outRouteUsage := fskeys.UsageRouteKey(fskeys.ChunkRouteKey(10, 21)) + require.NoError(t, st.PutAt(ctx, inRouteUsage, []byte("usage"), 10, 0)) + require.NoError(t, st.PutAt(ctx, outRouteUsage, []byte("skip"), 10, 0)) + + stream := &captureExportRangeVersionsStream{ctx: ctx} + err := internal.ExportRangeVersions(&pb.ExportRangeVersionsRequest{ + MaxCommitTs: 20, + RouteStart: routeKey, + RouteEnd: testPrefixScanEnd(routeKey), + KeyFamily: distribution.MigrationFamilyFilesystemUsage, + RangeStart: routeKey, + RangeEnd: testPrefixScanEnd(routeKey), + MaxScannedBytes: 1 << 20, + }, stream) + require.NoError(t, err) + require.Len(t, stream.responses, 1) + require.Equal(t, []*pb.MVCCVersion{ + {Key: inRouteUsage, CommitTs: 10, Value: []byte("usage"), KeyFamily: distribution.MigrationFamilyFilesystemUsage}, + }, stream.responses[0].GetVersions()) +} + func TestInternalImportRangeVersionsAppliesStoreBatch(t *testing.T) { t.Parallel() From da19644a4709ac094dcf2306116cd2220834c8d6 Mon Sep 17 00:00:00 2001 From: bootjp Date: Tue, 1 Sep 2026 00:28:20 +0900 Subject: [PATCH 45/58] migration: continue write fence validation --- kv/fsm.go | 7 +++++-- kv/fsm_migration_fence_test.go | 23 +++++++++++++++++++++++ 2 files changed, 28 insertions(+), 2 deletions(-) diff --git a/kv/fsm.go b/kv/fsm.go index 119ced56d..6931d35a9 100644 --- a/kv/fsm.go +++ b/kv/fsm.go @@ -1201,8 +1201,11 @@ func verifyWriteFenceFromSnapshot(mutations []*pb.Mutation, writeFenceBypassKeys if _, ok := bypassKeys[string(mut.Key)]; ok { continue } - if checked, err := verifyS3BucketAuxiliaryWriteFenceFromSnapshot(mut.Key, snap, snapVer, phase); checked || err != nil { - return err + if checked, err := verifyS3BucketAuxiliaryWriteFenceFromSnapshot(mut.Key, snap, snapVer, phase); checked { + if err != nil { + return err + } + continue } rKey := routeKey(mut.Key) if snap.WriteFencedForKey(rKey) { diff --git a/kv/fsm_migration_fence_test.go b/kv/fsm_migration_fence_test.go index e656f5021..0ec84b4e6 100644 --- a/kv/fsm_migration_fence_test.go +++ b/kv/fsm_migration_fence_test.go @@ -395,6 +395,29 @@ func TestFSMIgnoresRawRouteFenceForS3BucketAuxiliaryWrite(t *testing.T) { require.NoError(t, err) } +func TestFSMContinuesWriteFenceValidationAfterS3BucketAuxiliaryWrite(t *testing.T) { + t.Parallel() + + const bucket = "bucket-a" + fsm := newWriteFencedFSM(t) + + err := fsm.handleTxnRequest(context.Background(), &pb.Request{ + IsTxn: true, + Phase: pb.Phase_PREPARE, + Ts: 10, + Mutations: []*pb.Mutation{ + { + Op: pb.Op_PUT, + Key: []byte(txnMetaPrefix), + Value: EncodeTxnMeta(TxnMeta{PrimaryKey: []byte("z"), LockTTLms: defaultTxnLockTTLms}), + }, + {Op: pb.Op_PUT, Key: s3keys.BucketMetaKey(bucket), Value: []byte("meta")}, + {Op: pb.Op_PUT, Key: []byte("z"), Value: []byte("v")}, + }, + }, 10) + require.ErrorIs(t, err, ErrRouteWriteFenced) +} + func TestFSMIgnoresNonOwnerS3BucketAuxiliaryFenceForPointWrite(t *testing.T) { t.Parallel() From 30be14dd2a0edb6bdea7c6104d3caad02698792c Mon Sep 17 00:00:00 2001 From: bootjp Date: Tue, 1 Sep 2026 00:34:54 +0900 Subject: [PATCH 46/58] migration: route grouped reverse scans --- kv/shard_store.go | 13 ++----------- kv/shard_store_test.go | 23 ++++++++++++++++++++--- 2 files changed, 22 insertions(+), 14 deletions(-) diff --git a/kv/shard_store.go b/kv/shard_store.go index 234678948..6acaa34e6 100644 --- a/kv/shard_store.go +++ b/kv/shard_store.go @@ -746,10 +746,7 @@ func (s *ShardStore) ScanAtWithReadFence(ctx context.Context, start []byte, end } if reverse { if groupID != 0 { - if routeScanBoundsPresent(routeStart, routeEnd) { - return s.scanExplicitGroupAtWithReadFence(ctx, groupID, start, end, limit, ts, true, readRouteVersion, routeStart, routeEnd) - } - return nil, errors.WithStack(store.ErrNotSupported) + return s.scanExplicitGroupAtWithReadFence(ctx, groupID, start, end, limit, ts, true, readRouteVersion, routeStart, routeEnd) } return s.reverseScanAtWithReadFence(ctx, start, end, limit, ts, readRouteVersion, routeStart, routeEnd) } @@ -945,13 +942,7 @@ func (s *ShardStore) scanExplicitGroupRoutesAtWithReadFence(ctx context.Context, // ReverseScanGroupAt reverse-scans a range on the explicitly selected Raft group. func (s *ShardStore) ReverseScanGroupAt(ctx context.Context, groupID uint64, start []byte, end []byte, limit int, ts uint64) ([]*store.KVPair, error) { - if limit <= 0 { - return []*store.KVPair{}, nil - } - return s.scanRouteAtDirectionWithReadFence( - ctx, distribution.Route{GroupID: groupID}, start, end, limit, ts, true, true, - 0, nil, nil, - ) + return s.scanExplicitGroupAtWithReadFence(ctx, groupID, start, end, limit, ts, true, 0, nil, nil) } // ScanGroupKeysAt scans keys on the explicitly selected Raft group without diff --git a/kv/shard_store_test.go b/kv/shard_store_test.go index 90e4fd8b7..d7754292a 100644 --- a/kv/shard_store_test.go +++ b/kv/shard_store_test.go @@ -872,12 +872,26 @@ func TestShardStoreExplicitGroupReads_MergeStagedVisibility(t *testing.T) { {Key: []byte("c"), Value: []byte("staged-c")}, }, kvs) + kvs, err = st.ReverseScanGroupAt(ctx, 1, []byte("a"), []byte("z"), 10, 35) + require.NoError(t, err) + require.Equal(t, []*store.KVPair{ + {Key: []byte("c"), Value: []byte("staged-c")}, + {Key: []byte("b"), Value: []byte("staged-b")}, + }, kvs) + kvs, err = st.ScanAtWithReadFence(ctx, []byte("a"), []byte("z"), 10, 35, false, 1, 0, []byte("a"), []byte("z")) require.NoError(t, err) require.Equal(t, []*store.KVPair{ {Key: []byte("b"), Value: []byte("staged-b")}, {Key: []byte("c"), Value: []byte("staged-c")}, }, kvs) + + kvs, err = st.ScanAtWithReadFence(ctx, []byte("a"), []byte("z"), 10, 35, true, 1, 0, nil, nil) + require.NoError(t, err) + require.Equal(t, []*store.KVPair{ + {Key: []byte("c"), Value: []byte("staged-c")}, + {Key: []byte("b"), Value: []byte("staged-b")}, + }, kvs) } func TestShardStoreExplicitGroupReads_FailClosedWhenRouteMovedToStagedGroup(t *testing.T) { @@ -1895,10 +1909,13 @@ func TestShardStoreScanAtWithReadFence_AllowsExplicitGroupRouteBoundReverse(t *t require.NoError(t, groups[1].Store.PutAt(ctx, left, []byte("left"), 1, 0)) require.NoError(t, groups[1].Store.PutAt(ctx, right, []byte("right"), 2, 0)) - _, err := st.ScanAtWithReadFence(ctx, rawPrefix, prefixScanEnd(rawPrefix), 1, 2, true, 1, st.ReadRouteVersion(), nil, nil) - require.ErrorIs(t, err, store.ErrNotSupported) + kvs, err := st.ScanAtWithReadFence(ctx, rawPrefix, prefixScanEnd(rawPrefix), 1, 2, true, 1, st.ReadRouteVersion(), nil, nil) + require.NoError(t, err) + require.Len(t, kvs, 1) + require.Equal(t, right, kvs[0].Key) + require.Equal(t, []byte("right"), kvs[0].Value) - kvs, err := st.ScanAtWithReadFence(ctx, rawPrefix, prefixScanEnd(rawPrefix), -1, 2, true, 1, st.ReadRouteVersion(), []byte("m"), nil) + kvs, err = st.ScanAtWithReadFence(ctx, rawPrefix, prefixScanEnd(rawPrefix), -1, 2, true, 1, st.ReadRouteVersion(), []byte("m"), nil) require.NoError(t, err) require.Empty(t, kvs) From 8fc01b3c004840c5af9dad0193b8a1daebe223be Mon Sep 17 00:00:00 2001 From: bootjp Date: Tue, 1 Sep 2026 00:50:37 +0900 Subject: [PATCH 47/58] migration: route bucket auxiliary ownership --- adapter/distribution_server.go | 4 +-- adapter/distribution_server_test.go | 53 +++++++++++++++++++++++++++++ adapter/internal.go | 16 ++++++--- adapter/internal_migration_test.go | 16 ++++----- kv/shard_key.go | 16 +++++++++ kv/shard_key_test.go | 8 +++++ 6 files changed, 98 insertions(+), 15 deletions(-) diff --git a/adapter/distribution_server.go b/adapter/distribution_server.go index dc9c34713..550cfc38b 100644 --- a/adapter/distribution_server.go +++ b/adapter/distribution_server.go @@ -168,7 +168,7 @@ func (s *DistributionServer) GetRoute(ctx context.Context, req *pb.GetRouteReque if err := s.requireReadReady(); err != nil { return nil, err } - r, ok := s.engine.GetRoute(kv.RouteKey(req.Key)) + r, ok := s.engine.GetRoute(kv.RouteOwnershipKey(req.Key)) if !ok { return &pb.GetRouteResponse{}, nil } @@ -214,7 +214,7 @@ func (s *DistributionServer) GetRouteOwnership(ctx context.Context, req *pb.GetR // so looking the raw bytes up in the snapshot answers with the owner of // the raw family prefix instead of the group that actually owned the key // at that catalog version. - route, ok := snapshot.RouteOf(kv.RouteKey(req.GetKey())) + route, ok := snapshot.RouteOf(kv.RouteOwnershipKey(req.GetKey())) if !ok { return &pb.GetRouteOwnershipResponse{ CatalogVersion: snapshot.Version(), diff --git a/adapter/distribution_server_test.go b/adapter/distribution_server_test.go index 2fd781d94..d8418a2d3 100644 --- a/adapter/distribution_server_test.go +++ b/adapter/distribution_server_test.go @@ -9,6 +9,7 @@ import ( "github.com/bootjp/elastickv/distribution" "github.com/bootjp/elastickv/internal/fskeys" + "github.com/bootjp/elastickv/internal/s3keys" "github.com/bootjp/elastickv/kv" pb "github.com/bootjp/elastickv/proto" "github.com/bootjp/elastickv/store" @@ -59,6 +60,27 @@ func TestDistributionServerGetRoute_NormalizesFilesystemChunkKeys(t *testing.T) require.Equal(t, uint64(2), resp.RaftGroupId) } +func TestDistributionServerGetRoute_NormalizesS3BucketAuxiliaryKeys(t *testing.T) { + t.Parallel() + + bucket := "bucket-a" + routeKey := s3keys.RoutePrefixForBucketAnyGeneration(bucket) + engine := distribution.NewEngine() + engine.UpdateRoute([]byte(""), routeKey, 1) + engine.UpdateRoute(routeKey, nil, 2) + + s := NewDistributionServer(engine, nil) + for _, key := range [][]byte{ + s3keys.BucketMetaKey(bucket), + s3keys.BucketGenerationKey(bucket), + } { + resp, err := s.GetRoute(context.Background(), &pb.GetRouteRequest{Key: key}) + require.NoError(t, err) + require.Equal(t, routeKey, resp.Start) + require.Equal(t, uint64(2), resp.RaftGroupId) + } +} + func TestDistributionServerRouteReadsHonorStartupGate(t *testing.T) { t.Parallel() @@ -1438,3 +1460,34 @@ func TestDistributionServerGetRouteOwnership_NormalizesFilesystemChunkKeys(t *te "the chunk must resolve to its logical route owner, not the raw-prefix owner") require.Equal(t, uint64(2), resp.Route.RouteId) } + +func TestDistributionServerGetRouteOwnership_NormalizesS3BucketAuxiliaryKeys(t *testing.T) { + t.Parallel() + + bucket := "bucket-a" + routeKey := s3keys.RoutePrefixForBucketAnyGeneration(bucket) + engine := distribution.NewEngine() + require.NoError(t, engine.ApplySnapshot(distribution.CatalogSnapshot{ + Version: 3, + Routes: []distribution.RouteDescriptor{ + {RouteID: 1, Start: []byte(""), End: routeKey, GroupID: 1, State: distribution.RouteStateActive}, + {RouteID: 2, Start: routeKey, End: nil, GroupID: 2, State: distribution.RouteStateActive}, + }, + })) + + s := NewDistributionServer(engine, nil) + for _, key := range [][]byte{ + s3keys.BucketMetaKey(bucket), + s3keys.BucketGenerationKey(bucket), + } { + resp, err := s.GetRouteOwnership(context.Background(), &pb.GetRouteOwnershipRequest{ + Key: key, + CatalogVersion: 3, + }) + require.NoError(t, err) + require.True(t, resp.Found) + require.Equal(t, uint64(2), resp.Route.RaftGroupId, + "bucket auxiliary keys must resolve to the bucket route owner, not the raw-prefix owner") + require.Equal(t, uint64(2), resp.Route.RouteId) + } +} diff --git a/adapter/internal.go b/adapter/internal.go index 6a8650cdb..25278f02e 100644 --- a/adapter/internal.go +++ b/adapter/internal.go @@ -716,15 +716,16 @@ func migrationFamilyRequiresDecodedS3(family uint32) bool { func decodedS3BucketRouteFilter(family uint32, routeStart, routeEnd []byte) func([]byte) bool { allowRawRouteMatch := !s3BucketRouteBounds(routeStart, routeEnd) + rawRouteFilter := kv.RouteKeyFilter(routeStart, routeEnd) return func(rawKey []byte) bool { bucket, ok := decodedS3BucketName(family, rawKey) if !ok { return false } - if allowRawRouteMatch && kv.RouteKeyFilter(routeStart, routeEnd)(rawKey) { + if allowRawRouteMatch && rawRouteFilter(rawKey) { return true } - return decodedS3BucketRouteIntersects(bucket, routeStart, routeEnd) + return decodedS3BucketRouteSelected(bucket, routeStart, routeEnd) } } @@ -733,9 +734,9 @@ func s3BucketRouteBounds(routeStart, routeEnd []byte) bool { bytes.HasPrefix(routeEnd, []byte(s3keys.RoutePrefix)) } -func decodedS3BucketRouteIntersects(bucket string, routeStart, routeEnd []byte) bool { +func decodedS3BucketRouteSelected(bucket string, routeStart, routeEnd []byte) bool { bucketRouteStart := s3keys.RoutePrefixForBucketAnyGeneration(bucket) - return rangesIntersect(routeStart, routeEnd, bucketRouteStart, prefixScanEnd(bucketRouteStart)) + return keyInRouteRange(bucketRouteStart, routeStart, routeEnd) } func rangesIntersect(aStart, aEnd, bStart, bEnd []byte) bool { @@ -748,6 +749,13 @@ func rangesIntersect(aStart, aEnd, bStart, bEnd []byte) bool { return true } +func keyInRouteRange(key, start, end []byte) bool { + if bytes.Compare(key, start) < 0 { + return false + } + return len(end) == 0 || bytes.Compare(key, end) < 0 +} + func decodedS3BucketName(family uint32, rawKey []byte) (string, bool) { switch family { case distribution.MigrationFamilyS3BucketMeta: diff --git a/adapter/internal_migration_test.go b/adapter/internal_migration_test.go index d6a72032c..923b93876 100644 --- a/adapter/internal_migration_test.go +++ b/adapter/internal_migration_test.go @@ -278,8 +278,8 @@ func TestInternalExportRangeVersionsUsesDecodedS3BucketRouteFilter(t *testing.T) prefix: s3keys.BucketMetaPrefix, keyFor: s3keys.BucketMetaKey, value: []byte("meta"), - routeStart: s3keys.RouteKey("bucket-b", 0, ""), - routeEnd: s3keys.RouteKey("bucket-c", 0, ""), + routeStart: s3keys.RoutePrefixForBucketAnyGeneration("bucket-b"), + routeEnd: testPrefixScanEnd(s3keys.RoutePrefixForBucketAnyGeneration("bucket-b")), }, { name: "bucket generation", @@ -287,8 +287,8 @@ func TestInternalExportRangeVersionsUsesDecodedS3BucketRouteFilter(t *testing.T) prefix: s3keys.BucketGenerationPrefix, keyFor: s3keys.BucketGenerationKey, value: []byte("generation"), - routeStart: s3keys.RouteKey("bucket-b", 0, ""), - routeEnd: s3keys.RouteKey("bucket-c", 0, ""), + routeStart: s3keys.RoutePrefixForBucketAnyGeneration("bucket-b"), + routeEnd: testPrefixScanEnd(s3keys.RoutePrefixForBucketAnyGeneration("bucket-b")), }, } { t.Run(tc.name, func(t *testing.T) { @@ -333,7 +333,7 @@ func TestInternalExportRangeVersionsDecodedS3EmptyRouteEndIsUnbounded(t *testing stream := &captureExportRangeVersionsStream{ctx: ctx} err := internal.ExportRangeVersions(&pb.ExportRangeVersionsRequest{ MaxCommitTs: 20, - RouteStart: s3keys.RouteKey("bucket-z", 0, ""), + RouteStart: s3keys.RoutePrefixForBucketAnyGeneration("bucket-z"), RouteEnd: []byte{}, KeyFamily: distribution.MigrationFamilyS3BucketMeta, RangeStart: []byte(s3keys.BucketMetaPrefix), @@ -347,7 +347,7 @@ func TestInternalExportRangeVersionsDecodedS3EmptyRouteEndIsUnbounded(t *testing }, stream.responses[0].GetVersions()) } -func TestInternalExportRangeVersionsIncludesS3BucketAuxiliaryForBucketRouteIntersection(t *testing.T) { +func TestInternalExportRangeVersionsSkipsS3BucketAuxiliaryForNonOwnerRouteSlice(t *testing.T) { t.Parallel() ctx := context.Background() @@ -392,9 +392,7 @@ func TestInternalExportRangeVersionsIncludesS3BucketAuxiliaryForBucketRouteInter }, 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()) + require.Empty(t, stream.responses[0].GetVersions()) }) } } diff --git a/kv/shard_key.go b/kv/shard_key.go index 1d2b609a8..220278b04 100644 --- a/kv/shard_key.go +++ b/kv/shard_key.go @@ -115,6 +115,12 @@ func RouteKey(key []byte) []byte { return routeKey(key) } +// RouteOwnershipKey normalizes a stored key to the catalog key used when +// answering route ownership queries. +func RouteOwnershipKey(key []byte) []byte { + return routeOwnershipKey(key) +} + func routeKey(key []byte) []byte { if key == nil { return nil @@ -125,6 +131,16 @@ func routeKey(key []byte) []byte { return normalizeRouteKey(key) } +func routeOwnershipKey(key []byte) []byte { + if bucket, ok := s3keys.ParseBucketMetaKey(key); ok { + return s3keys.RoutePrefixForBucketAnyGeneration(bucket) + } + if bucket, ok := s3keys.ParseBucketGenerationKey(key); ok { + return s3keys.RoutePrefixForBucketAnyGeneration(bucket) + } + return routeKey(key) +} + func routeFilterKey(key []byte) []byte { if key == nil { return nil diff --git a/kv/shard_key_test.go b/kv/shard_key_test.go index 8e3caa07d..1348a00e4 100644 --- a/kv/shard_key_test.go +++ b/kv/shard_key_test.go @@ -32,6 +32,14 @@ func TestRouteKey_NormalizesTxnWrappedS3Key(t *testing.T) { require.Equal(t, s3keys.RouteKey("bucket-a", 7, "path/to/object"), routeKey(txnLockKey(embedded))) } +func TestRouteOwnershipKey_NormalizesS3BucketAuxiliaryKeys(t *testing.T) { + t.Parallel() + + want := s3keys.RoutePrefixForBucketAnyGeneration("bucket-a") + require.Equal(t, want, RouteOwnershipKey(s3keys.BucketMetaKey("bucket-a"))) + require.Equal(t, want, RouteOwnershipKey(s3keys.BucketGenerationKey("bucket-a"))) +} + func TestRouteKey_NormalizesFilesystemChunkKey(t *testing.T) { t.Parallel() From 8386465a40978bfc0f178df205a2b63ead02533a Mon Sep 17 00:00:00 2001 From: bootjp Date: Tue, 1 Sep 2026 18:54:39 +0900 Subject: [PATCH 48/58] migration: preserve mixed S3 raw exports --- adapter/internal.go | 9 +-------- adapter/internal_migration_test.go | 7 ++++--- 2 files changed, 5 insertions(+), 11 deletions(-) diff --git a/adapter/internal.go b/adapter/internal.go index 25278f02e..0f7b58773 100644 --- a/adapter/internal.go +++ b/adapter/internal.go @@ -715,25 +715,18 @@ func migrationFamilyRequiresDecodedS3(family uint32) bool { } func decodedS3BucketRouteFilter(family uint32, routeStart, routeEnd []byte) func([]byte) bool { - allowRawRouteMatch := !s3BucketRouteBounds(routeStart, routeEnd) - rawRouteFilter := kv.RouteKeyFilter(routeStart, routeEnd) return func(rawKey []byte) bool { bucket, ok := decodedS3BucketName(family, rawKey) if !ok { return false } - if allowRawRouteMatch && rawRouteFilter(rawKey) { + if keyInRouteRange(rawKey, routeStart, routeEnd) { return true } return decodedS3BucketRouteSelected(bucket, routeStart, routeEnd) } } -func s3BucketRouteBounds(routeStart, routeEnd []byte) bool { - return bytes.HasPrefix(routeStart, []byte(s3keys.RoutePrefix)) || - bytes.HasPrefix(routeEnd, []byte(s3keys.RoutePrefix)) -} - func decodedS3BucketRouteSelected(bucket string, routeStart, routeEnd []byte) bool { bucketRouteStart := s3keys.RoutePrefixForBucketAnyGeneration(bucket) return keyInRouteRange(bucketRouteStart, routeStart, routeEnd) diff --git a/adapter/internal_migration_test.go b/adapter/internal_migration_test.go index 923b93876..40c18b456 100644 --- a/adapter/internal_migration_test.go +++ b/adapter/internal_migration_test.go @@ -318,7 +318,7 @@ func TestInternalExportRangeVersionsUsesDecodedS3BucketRouteFilter(t *testing.T) } } -func TestInternalExportRangeVersionsDecodedS3EmptyRouteEndIsUnbounded(t *testing.T) { +func TestInternalExportRangeVersionsDecodedS3MixedRoutePreservesRawRangeMatches(t *testing.T) { t.Parallel() ctx := context.Background() @@ -326,9 +326,9 @@ func TestInternalExportRangeVersionsDecodedS3EmptyRouteEndIsUnbounded(t *testing internal := NewInternalWithEngine(nil, mockInternalLeader{}, nil, nil, WithInternalStore(st)) inRouteKey := s3keys.BucketMetaKey("bucket-z") - outRouteKey := s3keys.BucketMetaKey("bucket-a") + rawRangeKey := 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)) + require.NoError(t, st.PutAt(ctx, rawRangeKey, []byte("meta-a"), 10, 0)) stream := &captureExportRangeVersionsStream{ctx: ctx} err := internal.ExportRangeVersions(&pb.ExportRangeVersionsRequest{ @@ -343,6 +343,7 @@ func TestInternalExportRangeVersionsDecodedS3EmptyRouteEndIsUnbounded(t *testing require.NoError(t, err) require.Len(t, stream.responses, 1) require.Equal(t, []*pb.MVCCVersion{ + {Key: rawRangeKey, CommitTs: 10, Value: []byte("meta-a"), KeyFamily: distribution.MigrationFamilyS3BucketMeta}, {Key: inRouteKey, CommitTs: 10, Value: []byte("meta-z"), KeyFamily: distribution.MigrationFamilyS3BucketMeta}, }, stream.responses[0].GetVersions()) } From 8f1811f036ff30aa5ab1c0c034f45f902a34e761 Mon Sep 17 00:00:00 2001 From: bootjp Date: Tue, 1 Sep 2026 19:02:05 +0900 Subject: [PATCH 49/58] migration: preserve grouped key scan staging --- adapter/grpc_test.go | 43 ++++++++++++++++++++++++++++++++++++++++++ kv/shard_store.go | 23 ++++++++++++++++++---- kv/shard_store_test.go | 4 ++++ 3 files changed, 66 insertions(+), 4 deletions(-) diff --git a/adapter/grpc_test.go b/adapter/grpc_test.go index 9fb2634cd..203aa3ab6 100644 --- a/adapter/grpc_test.go +++ b/adapter/grpc_test.go @@ -894,6 +894,49 @@ func TestGRPCServer_RawScanAt_KeysOnlyUsesExplicitGroup(t *testing.T) { require.Equal(t, []byte("z"), st.scanEnd) } +func TestGRPCServer_RawScanAt_KeysOnlyExplicitGroupMergesStagedVisibility(t *testing.T) { + t.Parallel() + + ctx := context.Background() + engine := distribution.NewEngine() + require.NoError(t, engine.ApplySnapshot(distribution.CatalogSnapshot{ + Version: 1, + Routes: []distribution.RouteDescriptor{ + { + RouteID: 1, + Start: []byte("a"), + End: []byte("z"), + GroupID: 1, + State: distribution.RouteStateActive, + StagedVisibilityActive: true, + MigrationJobID: 9, + }, + }, + })) + group := &kvstore.ShardGroup{Store: store.NewMVCCStore()} + shards := kvstore.NewShardStore(engine, map[uint64]*kvstore.ShardGroup{1: group}) + t.Cleanup(func() { require.NoError(t, shards.Close()) }) + + require.NoError(t, group.Store.PutAt(ctx, []byte("b"), []byte("live-b"), 10, 0)) + require.NoError(t, group.Store.PutAt(ctx, distribution.MigrationStagedDataKey(9, []byte("b")), []byte("staged-b"), 20, 0)) + require.NoError(t, group.Store.PutAt(ctx, distribution.MigrationStagedDataKey(9, []byte("c")), []byte("staged-c"), 30, 0)) + + s := NewGRPCServer(shards, nil) + resp, err := s.RawScanAt(ctx, &pb.RawScanAtRequest{ + StartKey: []byte("a"), + EndKey: []byte("z"), + Limit: 10, + Ts: 35, + GroupId: 1, + KeysOnly: true, + }) + require.NoError(t, err) + require.Equal(t, []*pb.RawKVPair{ + {Key: []byte("b")}, + {Key: []byte("c")}, + }, resp.GetKv()) +} + func TestGRPCServer_RawScanAt_ReverseKeysOnlyUsesExplicitGroup(t *testing.T) { t.Parallel() diff --git a/kv/shard_store.go b/kv/shard_store.go index 6acaa34e6..902112918 100644 --- a/kv/shard_store.go +++ b/kv/shard_store.go @@ -797,7 +797,7 @@ func (s *ShardStore) ScanKeysAtWithReadFence(ctx context.Context, start []byte, return keysFromKVs(kvs), nil } if groupID != 0 { - return s.scanKeyRouteAtWithReadFence(ctx, distribution.Route{GroupID: groupID}, start, end, limit, ts, true, readRouteVersion) + return s.scanExplicitGroupKeysAtWithReadFence(ctx, groupID, start, end, limit, ts, readRouteVersion, nil, nil) } routes, clampToRoutes, routeVersion := s.routesForScanWithVersion(start, end) @@ -1703,7 +1703,7 @@ func filesystemChunkScanOverlap(start []byte, end []byte) ([]byte, []byte, bool) func (s *ShardStore) scanKeyRoutesAtWithReadFence(ctx context.Context, routes []distribution.Route, start []byte, end []byte, limit int, ts uint64, clampToRoutes bool, readRouteVersion uint64) ([][]byte, error) { out := make([][]byte, 0) - seenGroups := make(map[uint64]struct{}) + seenRoutes := make(map[repeatedRawScanRouteKey]struct{}) filterUsageOwners := !clampToRoutes && filesystemUsageScanOverlap(start, end) for _, route := range routes { scanStart := start @@ -1712,10 +1712,11 @@ func (s *ShardStore) scanKeyRoutesAtWithReadFence(ctx context.Context, routes [] scanStart = clampScanStart(start, route.Start) scanEnd = clampScanEnd(end, route.End) } else { - if _, seen := seenGroups[route.GroupID]; seen { + dedupeKey := repeatedRawScanRouteDedupeKey(route) + if _, seen := seenRoutes[dedupeKey]; seen { continue } - seenGroups[route.GroupID] = struct{}{} + seenRoutes[dedupeKey] = struct{}{} } var keys [][]byte @@ -1742,6 +1743,20 @@ func (s *ShardStore) scanKeyRoutesAtWithReadFence(ctx context.Context, routes [] return out, nil } +func (s *ShardStore) scanExplicitGroupKeysAtWithReadFence(ctx context.Context, groupID uint64, start []byte, end []byte, limit int, ts uint64, readRouteVersion uint64, routeStart []byte, routeEnd []byte) ([][]byte, error) { + if limit <= 0 { + return [][]byte{}, nil + } + routes, clampToRoutes, err := s.routesForExplicitGroupScanWithRouteBounds(groupID, start, end, routeStart, routeEnd) + if err != nil { + return nil, err + } + if !clampToRoutes && !routeScanBoundsPresent(routeStart, routeEnd) { + routes, _ = prepareUnclampedRawScanRoutes(routes, false) + } + return s.scanKeyRoutesAtWithReadFence(ctx, routes, start, end, limit, ts, clampToRoutes, readRouteVersion) +} + func (s *ShardStore) filterFilesystemUsageKVsForGroup(kvs []*store.KVPair, groupID uint64) []*store.KVPair { write := 0 for _, pair := range kvs { diff --git a/kv/shard_store_test.go b/kv/shard_store_test.go index d7754292a..740173d07 100644 --- a/kv/shard_store_test.go +++ b/kv/shard_store_test.go @@ -872,6 +872,10 @@ func TestShardStoreExplicitGroupReads_MergeStagedVisibility(t *testing.T) { {Key: []byte("c"), Value: []byte("staged-c")}, }, kvs) + keys, err := st.ScanKeysAtWithReadFence(ctx, []byte("a"), []byte("z"), 10, 35, 1, 0) + require.NoError(t, err) + require.Equal(t, [][]byte{[]byte("b"), []byte("c")}, keys) + kvs, err = st.ReverseScanGroupAt(ctx, 1, []byte("a"), []byte("z"), 10, 35) require.NoError(t, err) require.Equal(t, []*store.KVPair{ From c2f2406d0a295f738cd7f761251c4fe971d479ac Mon Sep 17 00:00:00 2001 From: bootjp Date: Tue, 1 Sep 2026 19:06:15 +0900 Subject: [PATCH 50/58] migration: retire promotion metadata --- store/lsm_migration.go | 12 +++++++++ store/migration_versions.go | 1 + store/migration_versions_test.go | 46 ++++++++++++++++++++++++++++++++ 3 files changed, 59 insertions(+) diff --git a/store/lsm_migration.go b/store/lsm_migration.go index 461fb702b..de2f6a1aa 100644 --- a/store/lsm_migration.go +++ b/store/lsm_migration.go @@ -594,6 +594,9 @@ func (s *pebbleStore) RetireMigration(ctx context.Context, jobID uint64) error { if err := s.stageRetireMigrationHLCFloor(batch, jobID); err != nil { return err } + if err := s.stageRetireMigrationPromotionState(batch, jobID); err != nil { + return err + } return errors.WithStack(batch.Commit(s.directApplyWriteOpts())) } @@ -619,6 +622,15 @@ func (s *pebbleStore) stageRetireMigrationHLCFloor(batch *pebble.Batch, jobID ui return stageMigrationMetadataMap(batch, migrationHLCFloorMetaKeyBytes, len(floors), encodeMigrationHLCFloors(floors)) } +func (s *pebbleStore) stageRetireMigrationPromotionState(batch *pebble.Batch, jobID uint64) error { + states, err := s.readPebblePromotionStates() + if err != nil { + return err + } + delete(states, jobID) + return stageMigrationMetadataMap(batch, migrationPromoteMetaKeyBytes, len(states), encodeMigrationPromotionStates(states)) +} + func stageMigrationMetadataMap(batch *pebble.Batch, key []byte, entries int, encoded []byte) error { if entries == 0 { return errors.WithStack(batch.Delete(key, nil)) diff --git a/store/migration_versions.go b/store/migration_versions.go index 1306b3dee..e68eb273f 100644 --- a/store/migration_versions.go +++ b/store/migration_versions.go @@ -706,5 +706,6 @@ func (s *mvccStore) RetireMigration(ctx context.Context, jobID uint64) error { } } delete(s.migrationHLCFloors, jobID) + delete(s.migrationPromotions, jobID) return nil } diff --git a/store/migration_versions_test.go b/store/migration_versions_test.go index a1610377b..59ca5a206 100644 --- a/store/migration_versions_test.go +++ b/store/migration_versions_test.go @@ -835,6 +835,9 @@ func TestImportVersionsIdempotencyAndMetadata(t *testing.T) { func TestRetireMigrationRemovesOnlySelectedJobMetadata(t *testing.T) { runMigrationStoreSuite(t, func(t *testing.T, st MVCCStore) { ctx := context.Background() + seedPromotionState(t, ctx, st, 1, []byte("stage|job1|"), []byte("job1-promoted")) + seedPromotionState(t, ctx, st, 2, []byte("stage|job2|"), []byte("job2-promoted")) + _, err := st.ImportVersions(ctx, ImportVersionsOptions{ JobID: 1, BracketID: 2, @@ -875,6 +878,14 @@ func TestRetireMigrationRemovesOnlySelectedJobMetadata(t *testing.T) { floor, err = st.MigrationHLCFloor(ctx, 2) require.NoError(t, err) require.Equal(t, uint64(30), floor) + stateReader, ok := st.(MigrationPromotionStateReader) + require.True(t, ok) + _, ok = migrationPromotionState(t, ctx, stateReader, 1) + require.False(t, ok) + state, ok := migrationPromotionState(t, ctx, stateReader, 2) + require.True(t, ok) + require.True(t, state.Done) + require.Equal(t, uint64(1), state.PromotedRows) res, err := st.ImportVersions(ctx, ImportVersionsOptions{ JobID: 1, @@ -906,6 +917,31 @@ func TestRetireMigrationRemovesOnlySelectedJobMetadata(t *testing.T) { }) } +func seedPromotionState(t *testing.T, ctx context.Context, st MVCCStore, jobID uint64, prefix []byte, target []byte) { + t.Helper() + promoter, ok := st.(MigrationPromoter) + require.True(t, ok) + require.NoError(t, st.PutAt(ctx, append(bytes.Clone(prefix), 'k'), []byte("v"), 10+jobID, 0)) + result, err := promoter.PromoteVersions(ctx, PromoteVersionsOptions{ + JobID: jobID, + StartKey: prefix, + EndKey: PrefixScanEnd(prefix), + MaxVersions: 10, + TargetKey: func(staged []byte) ([]byte, bool) { + return target, bytes.HasPrefix(staged, prefix) + }, + }) + require.NoError(t, err) + require.True(t, result.Done) +} + +func migrationPromotionState(t *testing.T, ctx context.Context, reader MigrationPromotionStateReader, jobID uint64) (PromotionState, bool) { + t.Helper() + state, ok, err := reader.MigrationPromotionState(ctx, jobID) + require.NoError(t, err) + return state, ok +} + func TestPebbleImportMetadataPersistsAcrossReopen(t *testing.T) { ctx := context.Background() dir, err := os.MkdirTemp("", "migration-import-persist-*") @@ -965,6 +1001,8 @@ func TestPebbleRetireMigrationPersistsAcrossReopen(t *testing.T) { Versions: []MVCCVersion{{Key: []byte("kept-k"), CommitTS: 109, Value: []byte("v109")}}, }) require.NoError(t, err) + seedPromotionState(t, ctx, st, 9, []byte("stage|job9|"), []byte("job9-promoted")) + seedPromotionState(t, ctx, st, 10, []byte("stage|job10|"), []byte("job10-promoted")) require.NoError(t, st.RetireMigration(ctx, 9)) require.NoError(t, st.Close()) @@ -977,6 +1015,14 @@ func TestPebbleRetireMigrationPersistsAcrossReopen(t *testing.T) { floor, err = reopened.MigrationHLCFloor(ctx, 10) require.NoError(t, err) require.Equal(t, uint64(109), floor) + stateReader, ok := reopened.(MigrationPromotionStateReader) + require.True(t, ok) + _, ok = migrationPromotionState(t, ctx, stateReader, 9) + require.False(t, ok) + state, ok := migrationPromotionState(t, ctx, stateReader, 10) + require.True(t, ok) + require.True(t, state.Done) + require.Equal(t, uint64(1), state.PromotedRows) res, err := reopened.ImportVersions(ctx, ImportVersionsOptions{ JobID: 9, From 398cbc2a0574a1bf18106721de005a80d433efd2 Mon Sep 17 00:00:00 2001 From: bootjp Date: Tue, 1 Sep 2026 19:09:19 +0900 Subject: [PATCH 51/58] migration: unwrap txn route ownership keys --- kv/fsm_migration_fence_test.go | 19 +++++++++++++++++-- kv/shard_key.go | 3 +++ kv/shard_key_test.go | 2 ++ 3 files changed, 22 insertions(+), 2 deletions(-) diff --git a/kv/fsm_migration_fence_test.go b/kv/fsm_migration_fence_test.go index 0ec84b4e6..72d4ae6a6 100644 --- a/kv/fsm_migration_fence_test.go +++ b/kv/fsm_migration_fence_test.go @@ -399,22 +399,37 @@ func TestFSMContinuesWriteFenceValidationAfterS3BucketAuxiliaryWrite(t *testing. t.Parallel() const bucket = "bucket-a" + auxiliaryKey := s3keys.BucketMetaKey(bucket) fsm := newWriteFencedFSM(t) err := fsm.handleTxnRequest(context.Background(), &pb.Request{ IsTxn: true, Phase: pb.Phase_PREPARE, Ts: 10, + Mutations: []*pb.Mutation{ + { + Op: pb.Op_PUT, + Key: []byte(txnMetaPrefix), + Value: EncodeTxnMeta(TxnMeta{PrimaryKey: auxiliaryKey, LockTTLms: defaultTxnLockTTLms}), + }, + {Op: pb.Op_PUT, Key: auxiliaryKey, Value: []byte("meta")}, + }, + }, 10) + require.NoError(t, err) + + err = fsm.handleTxnRequest(context.Background(), &pb.Request{ + IsTxn: true, + Phase: pb.Phase_PREPARE, + Ts: 11, Mutations: []*pb.Mutation{ { Op: pb.Op_PUT, Key: []byte(txnMetaPrefix), Value: EncodeTxnMeta(TxnMeta{PrimaryKey: []byte("z"), LockTTLms: defaultTxnLockTTLms}), }, - {Op: pb.Op_PUT, Key: s3keys.BucketMetaKey(bucket), Value: []byte("meta")}, {Op: pb.Op_PUT, Key: []byte("z"), Value: []byte("v")}, }, - }, 10) + }, 11) require.ErrorIs(t, err, ErrRouteWriteFenced) } diff --git a/kv/shard_key.go b/kv/shard_key.go index 220278b04..156bf8b78 100644 --- a/kv/shard_key.go +++ b/kv/shard_key.go @@ -132,6 +132,9 @@ func routeKey(key []byte) []byte { } func routeOwnershipKey(key []byte) []byte { + if embedded, ok := txnRouteKey(key); ok { + key = embedded + } if bucket, ok := s3keys.ParseBucketMetaKey(key); ok { return s3keys.RoutePrefixForBucketAnyGeneration(bucket) } diff --git a/kv/shard_key_test.go b/kv/shard_key_test.go index 1348a00e4..13eb99bb7 100644 --- a/kv/shard_key_test.go +++ b/kv/shard_key_test.go @@ -38,6 +38,8 @@ func TestRouteOwnershipKey_NormalizesS3BucketAuxiliaryKeys(t *testing.T) { want := s3keys.RoutePrefixForBucketAnyGeneration("bucket-a") require.Equal(t, want, RouteOwnershipKey(s3keys.BucketMetaKey("bucket-a"))) require.Equal(t, want, RouteOwnershipKey(s3keys.BucketGenerationKey("bucket-a"))) + require.Equal(t, want, RouteOwnershipKey(txnLockKey(s3keys.BucketMetaKey("bucket-a")))) + require.Equal(t, want, RouteOwnershipKey(txnLockKey(s3keys.BucketGenerationKey("bucket-a")))) } func TestRouteKey_NormalizesFilesystemChunkKey(t *testing.T) { From 7bb298303165bbaeef516f2a96b45f38fb43fe88 Mon Sep 17 00:00:00 2001 From: bootjp Date: Tue, 1 Sep 2026 19:32:21 +0900 Subject: [PATCH 52/58] migration: batch S3 owner probes --- adapter/grpc.go | 48 +++++- kv/shard_store.go | 366 ++++++++++++++++++++++++++++++++++++----- kv/shard_store_test.go | 68 +++++++- proto/service.pb.go | 35 +++- proto/service.proto | 5 + 5 files changed, 471 insertions(+), 51 deletions(-) diff --git a/adapter/grpc.go b/adapter/grpc.go index 8fe8e9c5d..7a61f1936 100644 --- a/adapter/grpc.go +++ b/adapter/grpc.go @@ -76,6 +76,10 @@ type rawVersionPresenceReader interface { VersionExistsAtOrBeforeGroupWithReadFence(ctx context.Context, key []byte, groupID uint64, ts uint64, readRouteVersion uint64) (bool, bool, error) } +type rawVersionPresenceBatchReader interface { + VersionsExistAtOrBeforeGroupWithReadFence(ctx context.Context, keys [][]byte, groupID uint64, ts uint64, readRouteVersion uint64) ([]bool, bool, error) +} + type rawGroupReverseScanner interface { ReverseScanGroupAt(ctx context.Context, groupID uint64, start []byte, end []byte, limit int, ts uint64) ([]*store.KVPair, error) } @@ -189,6 +193,21 @@ func (r *GRPCServer) RawLatestCommitTS(ctx context.Context, req *pb.RawLatestCom if err := r.requireReadReady(); err != nil { return nil, err } + readRouteVersion := r.readRouteVersion(req.GetReadRouteVersion()) + if len(req.GetKeys()) > 0 { + visible, visibleSupported, err := r.rawVersionsVisibleAt(ctx, req, readRouteVersion) + if err != nil { + return nil, errors.WithStack(err) + } + return &pb.RawLatestCommitTSResponse{ + VersionVisibleResults: visible, + VersionVisibleSupported: visibleSupported, + }, nil + } + return r.rawLatestCommitTSSingle(ctx, req, readRouteVersion) +} + +func (r *GRPCServer) rawLatestCommitTSSingle(ctx context.Context, req *pb.RawLatestCommitTSRequest, readRouteVersion uint64) (*pb.RawLatestCommitTSResponse, error) { key := req.GetKey() if len(key) == 0 { // No key: return the store's global last-committed watermark. @@ -204,7 +223,6 @@ func (r *GRPCServer) RawLatestCommitTS(ctx context.Context, req *pb.RawLatestCom var ts uint64 var exists bool var err error - readRouteVersion := r.readRouteVersion(req.GetReadRouteVersion()) if groupID := req.GetGroupId(); groupID != 0 { groupReader, ok := r.store.(rawGroupCommitTSReader) if !ok { @@ -233,6 +251,34 @@ func (r *GRPCServer) RawLatestCommitTS(ctx context.Context, req *pb.RawLatestCom }, nil } +func (r *GRPCServer) rawVersionsVisibleAt(ctx context.Context, req *pb.RawLatestCommitTSRequest, readRouteVersion uint64) ([]bool, bool, error) { + keys := req.GetKeys() + out := make([]bool, len(keys)) + at := req.GetVersionVisibleAtTs() + if at == 0 { + return out, false, nil + } + if reader, ok := r.store.(rawVersionPresenceBatchReader); ok { + visible, supported, err := reader.VersionsExistAtOrBeforeGroupWithReadFence(ctx, keys, req.GetGroupId(), at, readRouteVersion) + return visible, supported, errors.WithStack(err) + } + reader, ok := r.store.(rawVersionPresenceReader) + if !ok { + return out, false, nil + } + for i, key := range keys { + visible, supported, err := reader.VersionExistsAtOrBeforeGroupWithReadFence(ctx, key, req.GetGroupId(), at, readRouteVersion) + if err != nil { + return nil, false, errors.WithStack(err) + } + if !supported { + return out, false, nil + } + out[i] = visible + } + return out, true, nil +} + // rawVersionVisibleAt answers the optional version_visible_at_ts probe. The // second bool tells the caller whether this server answered it at all, so a // store that cannot check presence never looks like "no version exists". diff --git a/kv/shard_store.go b/kv/shard_store.go index 902112918..e7eba11f3 100644 --- a/kv/shard_store.go +++ b/kv/shard_store.go @@ -2318,87 +2318,340 @@ func (s *ShardStore) filterS3AuxiliaryKVsOwnedByRoute( ts uint64, readRouteVersion uint64, ) ([]*store.KVPair, error) { - out := make([]*store.KVPair, 0, len(kvs)) - for _, kvp := range kvs { + include, batches := s3BucketAuxiliaryOwnerProbePlan(kvs, routes, route) + if err := s.applyS3BucketAuxiliaryOwnerProbeBatches(ctx, include, batches, ts, readRouteVersion); err != nil { + return nil, err + } + return includedKVs(kvs, include), nil +} + +func s3BucketAuxiliaryOwnerProbePlan( + kvs []*store.KVPair, + routes []distribution.Route, + route distribution.Route, +) ([]bool, map[s3BucketAuxiliaryOwnerProbeRouteKey]*s3BucketAuxiliaryOwnerProbeBatch) { + include := make([]bool, len(kvs)) + batches := make(map[s3BucketAuxiliaryOwnerProbeRouteKey]*s3BucketAuxiliaryOwnerProbeBatch) + for i, kvp := range kvs { if kvp == nil { continue } owner, auxiliary := s3BucketAuxiliaryOwnerRoute(kvp.Key, routes) if auxiliary && !routeMatchesS3BucketAuxiliaryOwner(route, owner) { - covered, err := s.s3BucketAuxiliaryOwnerHasVersionAt(ctx, owner, kvp.Key, ts, readRouteVersion) - if err != nil { - return nil, err + key := s3BucketAuxiliaryOwnerProbeRouteKeyFor(owner) + batch := batches[key] + if batch == nil { + batch = &s3BucketAuxiliaryOwnerProbeBatch{owner: owner} + batches[key] = batch } - if covered { + batch.probes = append(batch.probes, s3BucketAuxiliaryOwnerProbe{index: i, key: kvp.Key}) + continue + } + include[i] = true + } + return include, batches +} + +func (s *ShardStore) applyS3BucketAuxiliaryOwnerProbeBatches( + ctx context.Context, + include []bool, + batches map[s3BucketAuxiliaryOwnerProbeRouteKey]*s3BucketAuxiliaryOwnerProbeBatch, + ts uint64, + readRouteVersion uint64, +) error { + for _, batch := range batches { + keys := batch.probeKeys() + covered, err := s.s3BucketAuxiliaryOwnerHasVersionsAt(ctx, batch.owner, keys, ts, readRouteVersion) + if err != nil { + return err + } + if len(covered) != len(batch.probes) { + return errors.WithStack(errors.Newf("s3 auxiliary owner version probe returned %d results for %d keys", len(covered), len(batch.probes))) + } + for i, found := range covered { + if found { continue } + include[batch.probes[i].index] = true } - out = append(out, kvp) } - return out, nil + return nil +} + +func includedKVs(kvs []*store.KVPair, include []bool) []*store.KVPair { + out := make([]*store.KVPair, 0, len(kvs)) + for i, kvp := range kvs { + if include[i] { + out = append(out, kvp) + } + } + return out +} + +type s3BucketAuxiliaryOwnerProbe struct { + index int + key []byte } -func (s *ShardStore) s3BucketAuxiliaryOwnerHasVersionAt( +type s3BucketAuxiliaryOwnerProbeBatch struct { + owner distribution.Route + probes []s3BucketAuxiliaryOwnerProbe +} + +func (b *s3BucketAuxiliaryOwnerProbeBatch) probeKeys() [][]byte { + keys := make([][]byte, 0, len(b.probes)) + for _, probe := range b.probes { + keys = append(keys, probe.key) + } + return keys +} + +type s3BucketAuxiliaryOwnerProbeRouteKey struct { + groupID uint64 + routeID uint64 + staged bool + migrationJobID uint64 + routeStart string + routeEnd string +} + +func s3BucketAuxiliaryOwnerProbeRouteKeyFor(route distribution.Route) s3BucketAuxiliaryOwnerProbeRouteKey { + return s3BucketAuxiliaryOwnerProbeRouteKey{ + groupID: route.GroupID, + routeID: route.RouteID, + staged: routeHasStagedVisibility(route), + migrationJobID: route.MigrationJobID, + routeStart: string(route.Start), + routeEnd: string(route.End), + } +} + +func (s *ShardStore) s3BucketAuxiliaryOwnerHasVersionsAt( ctx context.Context, owner distribution.Route, - key []byte, + keys [][]byte, ts uint64, readRouteVersion uint64, -) (bool, error) { +) ([]bool, error) { if !routeHasStagedVisibility(owner) { - live, liveOK, err := s.ownerRouteHasVersionAtOrBefore(ctx, owner, key, ts, readRouteVersion) - if err != nil { - return false, err - } - return existenceProbeResult(live, liveOK, owner, key) + return s.s3BucketAuxiliaryLiveOwnerHasVersionsAt(ctx, owner, keys, ts, readRouteVersion) } - // Staged before live: a live-first existence probe can miss a row that - // promotion moves between the two calls -- absent from live before the - // move, absent from staged after it -- and report a key that existed - // throughout as missing. - stagedKey := distribution.MigrationStagedDataKey(owner.MigrationJobID, key) - staged, stagedOK, err := s.ownerRouteHasVersionAtOrBefore(ctx, owner, stagedKey, ts, readRouteVersion) + return s.s3BucketAuxiliaryStagedOwnerHasVersionsAt(ctx, owner, keys, ts, readRouteVersion) +} + +func (s *ShardStore) s3BucketAuxiliaryLiveOwnerHasVersionsAt( + ctx context.Context, + owner distribution.Route, + keys [][]byte, + ts uint64, + readRouteVersion uint64, +) ([]bool, error) { + visible, ok, err := s.ownerRouteHasVersionsAtOrBefore(ctx, owner, keys, ts, readRouteVersion) if err != nil { - return false, err + return nil, err } - if staged { - return true, nil + return existenceProbeResults(visible, ok, owner, keys) +} + +func (s *ShardStore) s3BucketAuxiliaryStagedOwnerHasVersionsAt( + ctx context.Context, + owner distribution.Route, + keys [][]byte, + ts uint64, + readRouteVersion uint64, +) ([]bool, error) { + stagedKeys := s3BucketAuxiliaryStagedDataKeys(owner, keys) + staged, stagedOK, err := s.ownerRouteHasVersionsAtOrBefore(ctx, owner, stagedKeys, ts, readRouteVersion) + if err != nil { + return nil, err } - live, liveOK, err := s.ownerRouteHasVersionAtOrBefore(ctx, owner, key, ts, readRouteVersion) + + out, liveKeys, liveIndex := s3BucketAuxiliaryLiveProbePlan(keys, staged) + if len(liveKeys) == 0 { + return out, nil + } + + live, liveOK, err := s.ownerRouteHasVersionsAtOrBefore(ctx, owner, liveKeys, ts, readRouteVersion) if err != nil { - return false, err + return nil, err + } + if err := applyS3BucketAuxiliaryLiveProbeResults(out, live, liveIndex, liveKeys, owner, stagedOK && liveOK); err != nil { + return nil, err + } + return out, nil +} + +func s3BucketAuxiliaryStagedDataKeys(owner distribution.Route, keys [][]byte) [][]byte { + stagedKeys := make([][]byte, 0, len(keys)) + for _, key := range keys { + stagedKeys = append(stagedKeys, distribution.MigrationStagedDataKey(owner.MigrationJobID, key)) + } + return stagedKeys +} + +func s3BucketAuxiliaryLiveProbePlan(keys [][]byte, staged []bool) ([]bool, [][]byte, []int) { + out := make([]bool, len(keys)) + liveKeys := make([][]byte, 0, len(keys)) + liveIndex := make([]int, 0, len(keys)) + for i, found := range staged { + if found { + out[i] = true + continue + } + liveKeys = append(liveKeys, keys[i]) + liveIndex = append(liveIndex, i) + } + return out, liveKeys, liveIndex +} + +func applyS3BucketAuxiliaryLiveProbeResults( + out []bool, + live []bool, + liveIndex []int, + liveKeys [][]byte, + owner distribution.Route, + answered bool, +) error { + for i, found := range live { + if found { + out[liveIndex[i]] = true + continue + } + if !answered { + return ownerVersionProbeUnavailable(owner, liveKeys[i]) + } } - return existenceProbeResult(live, liveOK && stagedOK, owner, key) + return nil } -// existenceProbeResult turns a probe outcome into the (found, error) pair the -// callers expect: an unanswered probe is an error rather than a "no", because -// treating an unreachable owner as absent is what silently drops a row. -func existenceProbeResult(found bool, answered bool, owner distribution.Route, key []byte) (bool, error) { - if found { - return true, nil +func existenceProbeResults(found []bool, answered bool, owner distribution.Route, keys [][]byte) ([]bool, error) { + if answered { + return found, nil } - if !answered { - return false, ownerVersionProbeUnavailable(owner, key) + for i, visible := range found { + if visible { + continue + } + return nil, ownerVersionProbeUnavailable(owner, keys[i]) } - return false, nil + return found, nil } func ownerVersionProbeUnavailable(route distribution.Route, key []byte) error { return errors.Wrapf(ErrLeaderNotFound, "s3 auxiliary owner version probe unavailable group_id=%d key=%q", route.GroupID, key) } -func (s *ShardStore) ownerRouteHasVersionAtOrBefore( +func (s *ShardStore) ownerRouteHasVersionsAtOrBefore( ctx context.Context, owner distribution.Route, - key []byte, + keys [][]byte, ts uint64, readRouteVersion uint64, -) (bool, bool, error) { - if exists, ok, err := s.routeHasVersionAtOrBefore(ctx, owner, key, ts); ok || err != nil { +) ([]bool, bool, error) { + if exists, ok, err := s.routeHasVersionsAtOrBefore(ctx, owner, keys, ts); ok || err != nil { return exists, ok, err } - return s.routeHasVersionAtOrBeforeRemote(ctx, owner, key, ts, readRouteVersion) + return s.routeHasVersionsAtOrBeforeRemote(ctx, owner, keys, ts, readRouteVersion) +} + +func (s *ShardStore) routeHasVersionsAtOrBefore( + ctx context.Context, + route distribution.Route, + keys [][]byte, + ts uint64, +) ([]bool, bool, error) { + out := make([]bool, len(keys)) + g, ok := s.groupForID(route.GroupID) + if !ok || g.Store == nil { + return out, true, nil + } + if engine := engineForGroup(g); engine != nil && !isLinearizableRaftLeader(ctx, engine) { + return out, false, nil + } + exists, err := versionsExistAtOrBefore(ctx, g.Store, keys, ts) + return exists, true, errors.WithStack(err) +} + +func versionsExistAtOrBefore(ctx context.Context, st store.MVCCStore, keys [][]byte, ts uint64) ([]bool, error) { + out := make([]bool, len(keys)) + versions, err := latestCandidateVersionsAt(ctx, st, keys, ts) + if err != nil { + return nil, err + } + for i, key := range keys { + _, out[i] = versions[string(key)] + } + return out, nil +} + +func (s *ShardStore) routeHasVersionsAtOrBeforeRemote( + ctx context.Context, + route distribution.Route, + keys [][]byte, + ts uint64, + readRouteVersion uint64, +) ([]bool, bool, error) { + out := make([]bool, len(keys)) + if len(keys) == 0 { + return out, true, nil + } + if ts == 0 { + return out, false, nil + } + cli, ok, err := s.routeVersionPresenceClient(route) + if err != nil { + return nil, false, err + } + if !ok { + return out, false, nil + } + rpcCtx, cancel := context.WithTimeout(ctx, proxyForwardTimeout) + defer cancel() + resp, err := cli.RawLatestCommitTS(rpcCtx, &pb.RawLatestCommitTSRequest{ + Keys: cloneKeyBatch(keys), + ReadRouteVersion: readRouteVersion, + GroupId: route.GroupID, + VersionVisibleAtTs: ts, + }) + if err != nil { + return nil, false, errors.WithStack(err) + } + if !resp.GetVersionVisibleSupported() { + return out, false, nil + } + visible := resp.GetVersionVisibleResults() + if len(visible) != len(keys) { + return nil, false, errors.WithStack(errors.Newf("s3 auxiliary owner version probe returned %d results for %d keys", len(visible), len(keys))) + } + return append([]bool(nil), visible...), true, nil +} + +func (s *ShardStore) routeVersionPresenceClient(route distribution.Route) (pb.RawKVClient, bool, error) { + g, ok := s.groupForID(route.GroupID) + if !ok || g == nil { + return nil, false, nil + } + engine := engineForGroup(g) + if engine == nil { + return nil, false, nil + } + addr := leaderAddrFromEngine(engine) + if addr == "" { + return nil, false, nil + } + conn, err := s.connCache.ConnFor(addr) + if err != nil { + return nil, false, err + } + return pb.NewRawKVClient(conn), true, nil +} + +func cloneKeyBatch(keys [][]byte) [][]byte { + out := make([][]byte, 0, len(keys)) + for _, key := range keys { + out = append(out, bytes.Clone(key)) + } + return out } func routeMatchesS3BucketAuxiliaryOwner(route distribution.Route, owner distribution.Route) bool { @@ -3985,6 +4238,35 @@ func (s *ShardStore) LatestCommitTSGroupWithReadFence(ctx context.Context, key [ return s.proxyLatestCommitTSGroup(ctx, g, key, groupID, readRouteVersion) } +func (s *ShardStore) VersionsExistAtOrBeforeGroupWithReadFence(ctx context.Context, keys [][]byte, groupID uint64, ts uint64, readRouteVersion uint64) ([]bool, bool, error) { + if groupID == 0 { + return s.versionsExistAtOrBeforeWithReadFence(ctx, keys, ts, readRouteVersion) + } + if err := s.awaitReadRouteVersion(ctx, readRouteVersion); err != nil { + return nil, false, err + } + return s.routeHasVersionsAtOrBefore(ctx, distribution.Route{GroupID: groupID}, keys, ts) +} + +func (s *ShardStore) versionsExistAtOrBeforeWithReadFence(ctx context.Context, keys [][]byte, ts uint64, readRouteVersion uint64) ([]bool, bool, error) { + out := make([]bool, len(keys)) + if err := s.awaitReadRouteVersion(ctx, readRouteVersion); err != nil { + return nil, false, err + } + for i, key := range keys { + route, _, _, ok := s.routeAndGroupForKeyWithVersion(key) + if !ok { + continue + } + exists, answered, err := s.routeHasVersionsAtOrBefore(ctx, route, [][]byte{key}, ts) + if err != nil || !answered { + return out, answered, err + } + out[i] = exists[0] + } + return out, true, nil +} + func (s *ShardStore) localLatestCommitTS(ctx context.Context, g *ShardGroup, route distribution.Route, key []byte) (uint64, bool, error) { if !routeHasStagedVisibility(route) { liveTS, liveExists, err := g.Store.LatestCommitTS(ctx, key) diff --git a/kv/shard_store_test.go b/kv/shard_store_test.go index 740173d07..06437828e 100644 --- a/kv/shard_store_test.go +++ b/kv/shard_store_test.go @@ -673,10 +673,21 @@ func (s *versionVisibleRawKVServer) RawLatestCommitTS(_ context.Context, req *pb defer s.mu.Unlock() s.latestReqs = append(s.latestReqs, &pb.RawLatestCommitTSRequest{ Key: bytes.Clone(req.GetKey()), + Keys: cloneKeyBatch(req.GetKeys()), GroupId: req.GetGroupId(), ReadRouteVersion: req.GetReadRouteVersion(), VersionVisibleAtTs: req.GetVersionVisibleAtTs(), }) + if keys := req.GetKeys(); len(keys) > 0 { + results := make([]bool, len(keys)) + for i, key := range keys { + results[i] = s.visible[string(key)] + } + return &pb.RawLatestCommitTSResponse{ + VersionVisibleResults: results, + VersionVisibleSupported: true, + }, nil + } return &pb.RawLatestCommitTSResponse{ VersionVisible: s.visible[string(req.GetKey())], VersionVisibleSupported: true, @@ -725,7 +736,7 @@ func TestShardStoreS3BucketAuxiliaryOwnerProbeUsesLeaderRoutedReadFence(t *testi // through to staged anyway, and that pair is what a concurrent promotion // slips between. require.Len(t, probe.latestReqs, 1) - require.Equal(t, stagedKey, probe.latestReqs[0].GetKey()) + require.Equal(t, [][]byte{stagedKey}, probe.latestReqs[0].GetKeys()) for _, req := range probe.latestReqs { require.Equal(t, uint64(2), req.GetGroupId()) require.Equal(t, uint64(77), req.GetReadRouteVersion()) @@ -733,6 +744,61 @@ func TestShardStoreS3BucketAuxiliaryOwnerProbeUsesLeaderRoutedReadFence(t *testi } } +func TestShardStoreS3BucketAuxiliaryOwnerProbeBatchesFollowerChecks(t *testing.T) { + t.Parallel() + + ctx := context.Background() + keys := [][]byte{ + s3keys.BucketMetaKey("bucket-a"), + s3keys.BucketMetaKey("bucket-b"), + s3keys.BucketMetaKey("bucket-c"), + } + probe := &versionVisibleRawKVServer{visible: map[string]bool{}} + for _, key := range keys { + probe.visible[string(distribution.MigrationStagedDataKey(9, key))] = true + } + addr, stop := startRawKVServer(t, probe) + t.Cleanup(stop) + + engine := distribution.NewEngine() + require.NoError(t, engine.ApplySnapshot(distribution.CatalogSnapshot{ + Version: 77, + Routes: []distribution.RouteDescriptor{ + {RouteID: 1, Start: []byte(""), End: []byte(s3keys.RoutePrefix), GroupID: 1, State: distribution.RouteStateActive}, + { + RouteID: 2, Start: []byte(s3keys.RoutePrefix), End: prefixScanEnd([]byte(s3keys.RoutePrefix)), + GroupID: 2, State: distribution.RouteStateActive, StagedVisibilityActive: true, MigrationJobID: 9, + }, + {RouteID: 3, Start: prefixScanEnd([]byte(s3keys.RoutePrefix)), End: nil, GroupID: 1, State: distribution.RouteStateActive}, + }, + })) + groups := map[uint64]*ShardGroup{ + 1: {Store: store.NewMVCCStore()}, + 2: {Store: store.NewMVCCStore(), Engine: &followerProxyEngine{leader: addr}}, + } + st := NewShardStore(engine, groups) + for _, key := range keys { + require.NoError(t, groups[1].Store.PutAt(ctx, key, []byte("stale-source"), 10, 0)) + } + + start := []byte(s3keys.BucketMetaPrefix) + kvs, err := st.ScanAtWithReadFence(ctx, start, prefixScanEnd(start), 10, 30, false, 0, 77, nil, nil) + require.NoError(t, err) + require.Empty(t, kvs) + + probe.mu.Lock() + defer probe.mu.Unlock() + require.Len(t, probe.latestReqs, 1) + require.Equal(t, uint64(2), probe.latestReqs[0].GetGroupId()) + require.Equal(t, uint64(77), probe.latestReqs[0].GetReadRouteVersion()) + require.Equal(t, uint64(30), probe.latestReqs[0].GetVersionVisibleAtTs()) + require.Equal(t, [][]byte{ + distribution.MigrationStagedDataKey(9, keys[0]), + distribution.MigrationStagedDataKey(9, keys[1]), + distribution.MigrationStagedDataKey(9, keys[2]), + }, probe.latestReqs[0].GetKeys()) +} + func TestShardStoreS3BucketAuxiliaryOwnerProbeFailsWhenLeaderUnavailable(t *testing.T) { t.Parallel() diff --git a/proto/service.pb.go b/proto/service.pb.go index 5390269b5..ee2d6a359 100644 --- a/proto/service.pb.go +++ b/proto/service.pb.go @@ -413,8 +413,12 @@ type RawLatestCommitTSRequest struct { // below this timestamp. Comparing only `ts` cannot tell a tombstone at or // before the read timestamp apart from a newer version above it. VersionVisibleAtTs uint64 `protobuf:"varint,4,opt,name=version_visible_at_ts,json=versionVisibleAtTs,proto3" json:"version_visible_at_ts,omitempty"` - unknownFields protoimpl.UnknownFields - sizeCache protoimpl.SizeCache + // Optional batch form for version_visible_at_ts probes. Used by routed scan + // filters to avoid issuing one leader RPC per candidate key. The response + // carries version_visible_results in the same order. + Keys [][]byte `protobuf:"bytes,5,rep,name=keys,proto3" json:"keys,omitempty"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache } func (x *RawLatestCommitTSRequest) Reset() { @@ -475,6 +479,13 @@ func (x *RawLatestCommitTSRequest) GetVersionVisibleAtTs() uint64 { return 0 } +func (x *RawLatestCommitTSRequest) GetKeys() [][]byte { + if x != nil { + return x.Keys + } + return nil +} + type RawLatestCommitTSResponse struct { state protoimpl.MessageState `protogen:"open.v1"` Ts uint64 `protobuf:"varint,1,opt,name=ts,proto3" json:"ts,omitempty"` @@ -482,8 +493,9 @@ type RawLatestCommitTSResponse struct { // Answer to version_visible_at_ts. Only meaningful when // version_visible_supported is set; a server that predates the probe leaves // both unset and the caller falls back to comparing `ts`. - VersionVisible bool `protobuf:"varint,3,opt,name=version_visible,json=versionVisible,proto3" json:"version_visible,omitempty"` - VersionVisibleSupported bool `protobuf:"varint,4,opt,name=version_visible_supported,json=versionVisibleSupported,proto3" json:"version_visible_supported,omitempty"` + VersionVisible bool `protobuf:"varint,3,opt,name=version_visible,json=versionVisible,proto3" json:"version_visible,omitempty"` + VersionVisibleSupported bool `protobuf:"varint,4,opt,name=version_visible_supported,json=versionVisibleSupported,proto3" json:"version_visible_supported,omitempty"` + VersionVisibleResults []bool `protobuf:"varint,5,rep,packed,name=version_visible_results,json=versionVisibleResults,proto3" json:"version_visible_results,omitempty"` unknownFields protoimpl.UnknownFields sizeCache protoimpl.SizeCache } @@ -546,6 +558,13 @@ func (x *RawLatestCommitTSResponse) GetVersionVisibleSupported() bool { return false } +func (x *RawLatestCommitTSResponse) GetVersionVisibleResults() []bool { + if x != nil { + return x.VersionVisibleResults + } + return nil +} + type RawScanAtRequest struct { state protoimpl.MessageState `protogen:"open.v1"` StartKey []byte `protobuf:"bytes,1,opt,name=start_key,json=startKey,proto3" json:"start_key,omitempty"` @@ -2593,17 +2612,19 @@ const file_service_proto_rawDesc = "" + "\x03key\x18\x01 \x01(\fR\x03key\"P\n" + "\x11RawDeleteResponse\x12!\n" + "\fcommit_index\x18\x01 \x01(\x04R\vcommitIndex\x12\x18\n" + - "\asuccess\x18\x02 \x01(\bR\asuccess\"\xa8\x01\n" + + "\asuccess\x18\x02 \x01(\bR\asuccess\"\xbc\x01\n" + "\x18RawLatestCommitTSRequest\x12\x10\n" + "\x03key\x18\x01 \x01(\fR\x03key\x12,\n" + "\x12read_route_version\x18\x02 \x01(\x04R\x10readRouteVersion\x12\x19\n" + "\bgroup_id\x18\x03 \x01(\x04R\agroupId\x121\n" + - "\x15version_visible_at_ts\x18\x04 \x01(\x04R\x12versionVisibleAtTs\"\xa8\x01\n" + + "\x15version_visible_at_ts\x18\x04 \x01(\x04R\x12versionVisibleAtTs\x12\x12\n" + + "\x04keys\x18\x05 \x03(\fR\x04keys\"\xe0\x01\n" + "\x19RawLatestCommitTSResponse\x12\x0e\n" + "\x02ts\x18\x01 \x01(\x04R\x02ts\x12\x16\n" + "\x06exists\x18\x02 \x01(\bR\x06exists\x12'\n" + "\x0fversion_visible\x18\x03 \x01(\bR\x0eversionVisible\x12:\n" + - "\x19version_visible_supported\x18\x04 \x01(\bR\x17versionVisibleSupported\"\xde\x02\n" + + "\x19version_visible_supported\x18\x04 \x01(\bR\x17versionVisibleSupported\x126\n" + + "\x17version_visible_results\x18\x05 \x03(\bR\x15versionVisibleResults\"\xde\x02\n" + "\x10RawScanAtRequest\x12\x1b\n" + "\tstart_key\x18\x01 \x01(\fR\bstartKey\x12\x17\n" + "\aend_key\x18\x02 \x01(\fR\x06endKey\x12\x14\n" + diff --git a/proto/service.proto b/proto/service.proto index 6a0ccf570..023df4567 100644 --- a/proto/service.proto +++ b/proto/service.proto @@ -76,6 +76,10 @@ message RawLatestCommitTSRequest { // below this timestamp. Comparing only `ts` cannot tell a tombstone at or // before the read timestamp apart from a newer version above it. uint64 version_visible_at_ts = 4; + // Optional batch form for version_visible_at_ts probes. Used by routed scan + // filters to avoid issuing one leader RPC per candidate key. The response + // carries version_visible_results in the same order. + repeated bytes keys = 5; } message RawLatestCommitTSResponse { @@ -86,6 +90,7 @@ message RawLatestCommitTSResponse { // both unset and the caller falls back to comparing `ts`. bool version_visible = 3; bool version_visible_supported = 4; + repeated bool version_visible_results = 5; } message RawScanAtRequest { From 89acef18fdbd8f2626bd0e5201c08c3295802594 Mon Sep 17 00:00:00 2001 From: bootjp Date: Tue, 1 Sep 2026 19:38:16 +0900 Subject: [PATCH 53/58] migration: route S3 read probes by owner --- kv/sharded_coordinator.go | 3 +++ kv/sharded_coordinator_test.go | 32 ++++++++++++++++++++++++++++++++ 2 files changed, 35 insertions(+) diff --git a/kv/sharded_coordinator.go b/kv/sharded_coordinator.go index d31d32ed1..a15a48eff 100644 --- a/kv/sharded_coordinator.go +++ b/kv/sharded_coordinator.go @@ -2721,6 +2721,9 @@ func (c *ShardedCoordinator) stagedVisibilityRouteForReadKey(gid uint64, key []b if _, _, ok := distribution.MigrationStagedDataKeyParts(key); ok { return distribution.Route{}, false } + if route, ok := c.s3BucketAuxiliaryOwnerRouteForKey(key); ok { + return route, route.GroupID == gid && routeHasStagedVisibility(route) + } route, ok := c.engine.GetRoute(routeKey(key)) return route, ok && route.GroupID == gid && routeHasStagedVisibility(route) } diff --git a/kv/sharded_coordinator_test.go b/kv/sharded_coordinator_test.go index 9853cd9d6..f2f1dd280 100644 --- a/kv/sharded_coordinator_test.go +++ b/kv/sharded_coordinator_test.go @@ -5,6 +5,7 @@ import ( "testing" "github.com/bootjp/elastickv/distribution" + "github.com/bootjp/elastickv/internal/s3keys" "github.com/bootjp/elastickv/store" "github.com/stretchr/testify/require" ) @@ -67,3 +68,34 @@ func TestValidateReadKeysOnShardReadsStagedBeforeLive(t *testing.T) { require.Equal(t, uint64(60), ts) require.Greater(t, ts, uint64(50), "the read must still conflict with startTS=50") } + +func TestValidateReadKeysOnShardUsesS3BucketAuxiliaryOwnerForStagedProbe(t *testing.T) { + t.Parallel() + + ctx := context.Background() + const bucket = "bucket-a" + engine := distribution.NewEngine() + require.NoError(t, engine.ApplySnapshot(distribution.CatalogSnapshot{ + Version: 1, + Routes: s3BucketAuxiliaryStagedRoutes(bucket, 1, 2), + })) + + rawKey := s3keys.BucketMetaKey(bucket) + stagedKey := distribution.MigrationStagedDataKey(9, rawKey) + ownerStore := store.NewMVCCStore() + t.Cleanup(func() { _ = ownerStore.Close() }) + require.NoError(t, ownerStore.PutAt(ctx, stagedKey, []byte("staged"), 60, 0)) + + c := &ShardedCoordinator{ + engine: engine, + groups: map[uint64]*ShardGroup{ + 1: {Store: store.NewMVCCStore()}, + 2: {Store: ownerStore}, + }, + } + + ts, exists, err := c.latestCommitTSForReadKeyOnShard(ctx, 2, c.groups[2], rawKey) + require.NoError(t, err) + require.True(t, exists) + require.Equal(t, uint64(60), ts) +} From 82c5ff835b4930b7d39b6b789a5949b6881df841 Mon Sep 17 00:00:00 2001 From: bootjp Date: Tue, 1 Sep 2026 19:46:03 +0900 Subject: [PATCH 54/58] migration: bound raw version batches --- adapter/grpc.go | 7 ++- adapter/grpc_test.go | 86 +++++++++++++++++++++++++++++ kv/shard_store.go | 10 +--- kv/shard_store_test.go | 16 ++++-- proto/raw_latest_commit_ts_batch.go | 58 +++++++++++++++++++ proto/service.pb.go | 18 +++--- proto/service.proto | 8 +-- 7 files changed, 175 insertions(+), 28 deletions(-) create mode 100644 proto/raw_latest_commit_ts_batch.go diff --git a/adapter/grpc.go b/adapter/grpc.go index 7a61f1936..f08c8cfa1 100644 --- a/adapter/grpc.go +++ b/adapter/grpc.go @@ -194,7 +194,7 @@ func (r *GRPCServer) RawLatestCommitTS(ctx context.Context, req *pb.RawLatestCom return nil, err } readRouteVersion := r.readRouteVersion(req.GetReadRouteVersion()) - if len(req.GetKeys()) > 0 { + if len(req.GetKeyBatch()) > 0 { visible, visibleSupported, err := r.rawVersionsVisibleAt(ctx, req, readRouteVersion) if err != nil { return nil, errors.WithStack(err) @@ -252,7 +252,10 @@ func (r *GRPCServer) rawLatestCommitTSSingle(ctx context.Context, req *pb.RawLat } func (r *GRPCServer) rawVersionsVisibleAt(ctx context.Context, req *pb.RawLatestCommitTSRequest, readRouteVersion uint64) ([]bool, bool, error) { - keys := req.GetKeys() + keys, err := pb.DecodeRawLatestCommitTSKeyBatch(req.GetKeyBatch(), maxGRPCScanLimit) + if err != nil { + return nil, false, errors.WithStack(status.Error(codes.InvalidArgument, err.Error())) + } out := make([]bool, len(keys)) at := req.GetVersionVisibleAtTs() if at == 0 { diff --git a/adapter/grpc_test.go b/adapter/grpc_test.go index 203aa3ab6..b80534040 100644 --- a/adapter/grpc_test.go +++ b/adapter/grpc_test.go @@ -1221,6 +1221,41 @@ func (s *recordingVersionPresenceStore) VersionExistsAtOrBeforeGroupWithReadFenc return s.visible, s.supported, nil } +type recordingVersionPresenceBatchStore struct { + store.MVCCStore + + visible map[string]bool + supported bool + calls int + lastKeys [][]byte + lastGroup uint64 + lastTS uint64 + lastReadRouteVersion uint64 +} + +func cloneBytes2D(keys [][]byte) [][]byte { + out := make([][]byte, 0, len(keys)) + for _, key := range keys { + out = append(out, append([]byte(nil), key...)) + } + return out +} + +func (s *recordingVersionPresenceBatchStore) VersionsExistAtOrBeforeGroupWithReadFence( + _ context.Context, keys [][]byte, groupID uint64, ts uint64, readRouteVersion uint64, +) ([]bool, bool, error) { + s.calls++ + s.lastKeys = cloneBytes2D(keys) + s.lastGroup = groupID + s.lastTS = ts + s.lastReadRouteVersion = readRouteVersion + out := make([]bool, len(keys)) + for i, key := range keys { + out[i] = s.visible[string(key)] + } + return out, s.supported, nil +} + // version_visible_at_ts is optional: only a request that asks gets an answer, // and a store that cannot answer must not look like "no version exists". func TestGRPCServer_RawLatestCommitTS_VersionVisibleProbe(t *testing.T) { @@ -1297,3 +1332,54 @@ func TestGRPCServer_RawLatestCommitTS_VersionVisibleProbe(t *testing.T) { }) } } + +func TestGRPCServer_RawLatestCommitTS_BatchVersionVisibleProbe(t *testing.T) { + t.Parallel() + + st := &recordingVersionPresenceBatchStore{ + MVCCStore: store.NewMVCCStore(), + visible: map[string]bool{"a": true, "b": false}, + supported: true, + } + t.Cleanup(func() { _ = st.Close() }) + s := NewGRPCServer(st, nil) + + resp, err := s.RawLatestCommitTS(context.Background(), &pb.RawLatestCommitTSRequest{ + KeyBatch: pb.EncodeRawLatestCommitTSKeyBatch([][]byte{[]byte("a"), []byte("b")}), + GroupId: 42, + ReadRouteVersion: 77, + VersionVisibleAtTs: 100, + }) + require.NoError(t, err) + require.Equal(t, []bool{true, false}, resp.GetVersionVisibleResults()) + require.True(t, resp.GetVersionVisibleSupported()) + require.Equal(t, 1, st.calls) + require.Equal(t, [][]byte{[]byte("a"), []byte("b")}, st.lastKeys) + require.Equal(t, uint64(42), st.lastGroup) + require.Equal(t, uint64(100), st.lastTS) + require.Equal(t, uint64(77), st.lastReadRouteVersion) +} + +func TestGRPCServer_RawLatestCommitTS_RejectsOversizedBatchBeforeProbe(t *testing.T) { + t.Parallel() + + keys := make([][]byte, maxGRPCScanLimit+1) + for i := range keys { + keys[i] = []byte("k") + } + st := &recordingVersionPresenceBatchStore{ + MVCCStore: store.NewMVCCStore(), + visible: map[string]bool{}, + supported: true, + } + t.Cleanup(func() { _ = st.Close() }) + s := NewGRPCServer(st, nil) + + _, err := s.RawLatestCommitTS(context.Background(), &pb.RawLatestCommitTSRequest{ + KeyBatch: pb.EncodeRawLatestCommitTSKeyBatch(keys), + VersionVisibleAtTs: 100, + }) + require.Error(t, err) + require.Equal(t, codes.InvalidArgument, status.Code(err)) + require.Zero(t, st.calls) +} diff --git a/kv/shard_store.go b/kv/shard_store.go index e7eba11f3..af112d6b8 100644 --- a/kv/shard_store.go +++ b/kv/shard_store.go @@ -2608,7 +2608,7 @@ func (s *ShardStore) routeHasVersionsAtOrBeforeRemote( rpcCtx, cancel := context.WithTimeout(ctx, proxyForwardTimeout) defer cancel() resp, err := cli.RawLatestCommitTS(rpcCtx, &pb.RawLatestCommitTSRequest{ - Keys: cloneKeyBatch(keys), + KeyBatch: pb.EncodeRawLatestCommitTSKeyBatch(keys), ReadRouteVersion: readRouteVersion, GroupId: route.GroupID, VersionVisibleAtTs: ts, @@ -2646,14 +2646,6 @@ func (s *ShardStore) routeVersionPresenceClient(route distribution.Route) (pb.Ra return pb.NewRawKVClient(conn), true, nil } -func cloneKeyBatch(keys [][]byte) [][]byte { - out := make([][]byte, 0, len(keys)) - for _, key := range keys { - out = append(out, bytes.Clone(key)) - } - return out -} - func routeMatchesS3BucketAuxiliaryOwner(route distribution.Route, owner distribution.Route) bool { if route.GroupID != owner.GroupID { return false diff --git a/kv/shard_store_test.go b/kv/shard_store_test.go index 06437828e..b0e7fa0ff 100644 --- a/kv/shard_store_test.go +++ b/kv/shard_store_test.go @@ -673,12 +673,16 @@ func (s *versionVisibleRawKVServer) RawLatestCommitTS(_ context.Context, req *pb defer s.mu.Unlock() s.latestReqs = append(s.latestReqs, &pb.RawLatestCommitTSRequest{ Key: bytes.Clone(req.GetKey()), - Keys: cloneKeyBatch(req.GetKeys()), + KeyBatch: bytes.Clone(req.GetKeyBatch()), GroupId: req.GetGroupId(), ReadRouteVersion: req.GetReadRouteVersion(), VersionVisibleAtTs: req.GetVersionVisibleAtTs(), }) - if keys := req.GetKeys(); len(keys) > 0 { + keys, err := pb.DecodeRawLatestCommitTSKeyBatch(req.GetKeyBatch(), store.MaxDeltaScanLimit+1) + if err != nil { + return nil, err + } + if len(keys) > 0 { results := make([]bool, len(keys)) for i, key := range keys { results[i] = s.visible[string(key)] @@ -736,7 +740,9 @@ func TestShardStoreS3BucketAuxiliaryOwnerProbeUsesLeaderRoutedReadFence(t *testi // through to staged anyway, and that pair is what a concurrent promotion // slips between. require.Len(t, probe.latestReqs, 1) - require.Equal(t, [][]byte{stagedKey}, probe.latestReqs[0].GetKeys()) + keys, err := pb.DecodeRawLatestCommitTSKeyBatch(probe.latestReqs[0].GetKeyBatch(), store.MaxDeltaScanLimit+1) + require.NoError(t, err) + require.Equal(t, [][]byte{stagedKey}, keys) for _, req := range probe.latestReqs { require.Equal(t, uint64(2), req.GetGroupId()) require.Equal(t, uint64(77), req.GetReadRouteVersion()) @@ -792,11 +798,13 @@ func TestShardStoreS3BucketAuxiliaryOwnerProbeBatchesFollowerChecks(t *testing.T require.Equal(t, uint64(2), probe.latestReqs[0].GetGroupId()) require.Equal(t, uint64(77), probe.latestReqs[0].GetReadRouteVersion()) require.Equal(t, uint64(30), probe.latestReqs[0].GetVersionVisibleAtTs()) + decodedKeys, err := pb.DecodeRawLatestCommitTSKeyBatch(probe.latestReqs[0].GetKeyBatch(), store.MaxDeltaScanLimit+1) + require.NoError(t, err) require.Equal(t, [][]byte{ distribution.MigrationStagedDataKey(9, keys[0]), distribution.MigrationStagedDataKey(9, keys[1]), distribution.MigrationStagedDataKey(9, keys[2]), - }, probe.latestReqs[0].GetKeys()) + }, decodedKeys) } func TestShardStoreS3BucketAuxiliaryOwnerProbeFailsWhenLeaderUnavailable(t *testing.T) { diff --git a/proto/raw_latest_commit_ts_batch.go b/proto/raw_latest_commit_ts_batch.go new file mode 100644 index 000000000..f3782b9f2 --- /dev/null +++ b/proto/raw_latest_commit_ts_batch.go @@ -0,0 +1,58 @@ +package proto + +import ( + "encoding/binary" + "fmt" + "io" +) + +// EncodeRawLatestCommitTSKeyBatch packs exact-key version-presence probes into +// a single protobuf bytes field. The service-side decoder can reject an +// oversized count before allocating per-key slices. +func EncodeRawLatestCommitTSKeyBatch(keys [][]byte) []byte { + size := binary.MaxVarintLen64 + for _, key := range keys { + size += binary.MaxVarintLen64 + len(key) + } + out := make([]byte, 0, size) + out = binary.AppendUvarint(out, uint64(len(keys))) + for _, key := range keys { + out = binary.AppendUvarint(out, uint64(len(key))) + out = append(out, key...) + } + return out +} + +func DecodeRawLatestCommitTSKeyBatch(data []byte, maxKeys int) ([][]byte, error) { + if len(data) == 0 { + return nil, nil + } + if maxKeys < 0 { + return nil, fmt.Errorf("raw latest commit timestamp key batch max_keys must be non-negative") + } + count, n := binary.Uvarint(data) + if n <= 0 { + return nil, fmt.Errorf("raw latest commit timestamp key batch has invalid count") + } + if count > uint64(maxKeys) { + return nil, fmt.Errorf("raw latest commit timestamp key batch has %d keys, max %d", count, maxKeys) + } + data = data[n:] + keys := make([][]byte, 0, count) + for range count { + keyLen, n := binary.Uvarint(data) + if n <= 0 { + return nil, fmt.Errorf("raw latest commit timestamp key batch has invalid key length") + } + data = data[n:] + if keyLen > uint64(len(data)) { + return nil, io.ErrUnexpectedEOF + } + keys = append(keys, data[:keyLen]) + data = data[keyLen:] + } + if len(data) != 0 { + return nil, fmt.Errorf("raw latest commit timestamp key batch has trailing bytes") + } + return keys, nil +} diff --git a/proto/service.pb.go b/proto/service.pb.go index ee2d6a359..e26eed3f7 100644 --- a/proto/service.pb.go +++ b/proto/service.pb.go @@ -413,10 +413,10 @@ type RawLatestCommitTSRequest struct { // below this timestamp. Comparing only `ts` cannot tell a tombstone at or // before the read timestamp apart from a newer version above it. VersionVisibleAtTs uint64 `protobuf:"varint,4,opt,name=version_visible_at_ts,json=versionVisibleAtTs,proto3" json:"version_visible_at_ts,omitempty"` - // Optional batch form for version_visible_at_ts probes. Used by routed scan - // filters to avoid issuing one leader RPC per candidate key. The response - // carries version_visible_results in the same order. - Keys [][]byte `protobuf:"bytes,5,rep,name=keys,proto3" json:"keys,omitempty"` + // Optional batch form for version_visible_at_ts probes. Encoded as a single + // bounded key_batch blob so untrusted clients cannot expand one request into + // an unbounded repeated-bytes slice before the service enforces page limits. + KeyBatch []byte `protobuf:"bytes,5,opt,name=key_batch,json=keyBatch,proto3" json:"key_batch,omitempty"` unknownFields protoimpl.UnknownFields sizeCache protoimpl.SizeCache } @@ -479,9 +479,9 @@ func (x *RawLatestCommitTSRequest) GetVersionVisibleAtTs() uint64 { return 0 } -func (x *RawLatestCommitTSRequest) GetKeys() [][]byte { +func (x *RawLatestCommitTSRequest) GetKeyBatch() []byte { if x != nil { - return x.Keys + return x.KeyBatch } return nil } @@ -2612,13 +2612,13 @@ const file_service_proto_rawDesc = "" + "\x03key\x18\x01 \x01(\fR\x03key\"P\n" + "\x11RawDeleteResponse\x12!\n" + "\fcommit_index\x18\x01 \x01(\x04R\vcommitIndex\x12\x18\n" + - "\asuccess\x18\x02 \x01(\bR\asuccess\"\xbc\x01\n" + + "\asuccess\x18\x02 \x01(\bR\asuccess\"\xc5\x01\n" + "\x18RawLatestCommitTSRequest\x12\x10\n" + "\x03key\x18\x01 \x01(\fR\x03key\x12,\n" + "\x12read_route_version\x18\x02 \x01(\x04R\x10readRouteVersion\x12\x19\n" + "\bgroup_id\x18\x03 \x01(\x04R\agroupId\x121\n" + - "\x15version_visible_at_ts\x18\x04 \x01(\x04R\x12versionVisibleAtTs\x12\x12\n" + - "\x04keys\x18\x05 \x03(\fR\x04keys\"\xe0\x01\n" + + "\x15version_visible_at_ts\x18\x04 \x01(\x04R\x12versionVisibleAtTs\x12\x1b\n" + + "\tkey_batch\x18\x05 \x01(\fR\bkeyBatch\"\xe0\x01\n" + "\x19RawLatestCommitTSResponse\x12\x0e\n" + "\x02ts\x18\x01 \x01(\x04R\x02ts\x12\x16\n" + "\x06exists\x18\x02 \x01(\bR\x06exists\x12'\n" + diff --git a/proto/service.proto b/proto/service.proto index 023df4567..b4f13b929 100644 --- a/proto/service.proto +++ b/proto/service.proto @@ -76,10 +76,10 @@ message RawLatestCommitTSRequest { // below this timestamp. Comparing only `ts` cannot tell a tombstone at or // before the read timestamp apart from a newer version above it. uint64 version_visible_at_ts = 4; - // Optional batch form for version_visible_at_ts probes. Used by routed scan - // filters to avoid issuing one leader RPC per candidate key. The response - // carries version_visible_results in the same order. - repeated bytes keys = 5; + // Optional batch form for version_visible_at_ts probes. Encoded as a single + // bounded key_batch blob so untrusted clients cannot expand one request into + // an unbounded repeated-bytes slice before the service enforces page limits. + bytes key_batch = 5; } message RawLatestCommitTSResponse { From 202eb8bf9267ed61b1189f6b32ff18dd603ec8e7 Mon Sep 17 00:00:00 2001 From: bootjp Date: Tue, 1 Sep 2026 19:49:54 +0900 Subject: [PATCH 55/58] migration: cover chunkblob export exclusion --- adapter/internal_migration_test.go | 27 +++++++++++++++++++++++++++ 1 file changed, 27 insertions(+) diff --git a/adapter/internal_migration_test.go b/adapter/internal_migration_test.go index 40c18b456..f3127136a 100644 --- a/adapter/internal_migration_test.go +++ b/adapter/internal_migration_test.go @@ -150,6 +150,33 @@ func TestInternalExportRangeVersionsUsesStoreAndRouteFilter(t *testing.T) { }, stream.responses[0].GetVersions()) } +func TestInternalExportRangeVersionsExcludesPeerLocalChunkBlobsFromUserFamily(t *testing.T) { + t.Parallel() + + ctx := context.Background() + st := store.NewMVCCStore() + var digest [32]byte + digest[0] = 0xab + blobKey := s3keys.ChunkBlobKey(digest) + require.NoError(t, st.PutAt(ctx, blobKey, []byte("peer-local"), 10, 0)) + internal := NewInternalWithEngine(nil, mockInternalLeader{}, nil, nil, WithInternalStore(st)) + stream := &captureExportRangeVersionsStream{ctx: ctx} + + err := internal.ExportRangeVersions(&pb.ExportRangeVersionsRequest{ + MaxCommitTs: 20, + RouteStart: []byte(s3keys.ChunkBlobPrefix), + RouteEnd: []byte("!s4|"), + KeyFamily: distribution.MigrationFamilyUser, + RangeStart: []byte(s3keys.ChunkBlobPrefix), + RangeEnd: testPrefixScanEnd([]byte(s3keys.ChunkBlobPrefix)), + MaxScannedBytes: 1 << 20, + }, stream) + require.NoError(t, err) + require.Len(t, stream.responses, 1) + require.True(t, stream.responses[0].GetDone()) + require.Empty(t, stream.responses[0].GetVersions()) +} + func TestInternalExportRangeVersionsUsesValueAwareLegacyListDeltaRouteFilter(t *testing.T) { t.Parallel() From 8dfa4393eb9c47cd7f5d02ec1624ca0e6c37cf0f Mon Sep 17 00:00:00 2001 From: bootjp Date: Tue, 1 Sep 2026 19:56:35 +0900 Subject: [PATCH 56/58] migration: cover txn wrapped SQS routing --- adapter/internal_migration_test.go | 97 ++++++++++++++++++++++++++++++ 1 file changed, 97 insertions(+) diff --git a/adapter/internal_migration_test.go b/adapter/internal_migration_test.go index f3127136a..47448aaf9 100644 --- a/adapter/internal_migration_test.go +++ b/adapter/internal_migration_test.go @@ -490,6 +490,103 @@ func TestInternalExportRangeVersionsUsesPartitionResolverGroup(t *testing.T) { }, stream.responses[0].GetVersions()) } +func TestInternalExportRangeVersionsResolvesTxnWrappedPartitionedSQSKeysByEmbeddedOwner(t *testing.T) { + t.Parallel() + + wrapWithStartTS := func(prefix string, userKey []byte) []byte { + var raw [8]byte + binary.BigEndian.PutUint64(raw[:], 77) + key := append([]byte(prefix), userKey...) + return append(key, raw[:]...) + } + cases := []struct { + name string + family uint32 + prefix []byte + wrap func([]byte) []byte + }{ + { + name: "intent", + family: distribution.MigrationFamilyTxnIntent, + prefix: []byte("!txn|int|"), + wrap: func(userKey []byte) []byte { + return append([]byte("!txn|int|"), userKey...) + }, + }, + { + name: "commit", + family: distribution.MigrationFamilyTxnCommit, + prefix: []byte("!txn|cmt|"), + wrap: func(userKey []byte) []byte { + return wrapWithStartTS("!txn|cmt|", userKey) + }, + }, + { + name: "rollback", + family: distribution.MigrationFamilyTxnRollback, + prefix: []byte("!txn|rb|"), + wrap: func(userKey []byte) []byte { + return wrapWithStartTS("!txn|rb|", userKey) + }, + }, + { + name: "success", + family: distribution.MigrationFamilyTxnSuccess, + prefix: []byte("!txn|ok|"), + wrap: func(userKey []byte) []byte { + return kv.TxnSuccessMarkerKey(userKey, 77, 88, []byte("primary")) + }, + }, + { + name: "meta", + family: distribution.MigrationFamilyTxnMeta, + prefix: []byte(kv.TxnMetaPrefix), + wrap: func(userKey []byte) []byte { + return append([]byte(kv.TxnMetaPrefix), userKey...) + }, + }, + } + + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + + ctx := context.Background() + st := store.NewMVCCStore() + resolver := NewSQSPartitionResolver(map[string][]uint64{ + "orders.fifo": {10, 11}, + }) + internal := NewInternalWithEngine(nil, mockInternalLeader{}, nil, nil, + WithInternalStore(st), + WithInternalMigrationExportRouting(11, resolver), + ) + + p0 := tc.wrap(sqsPartitionedMsgDataKey("orders.fifo", 0, 1, "msg-0")) + p1 := tc.wrap(sqsPartitionedMsgDataKey("orders.fifo", 1, 1, "msg-1")) + unknown := tc.wrap(sqsPartitionedMsgDataKey("unknown.fifo", 0, 1, "msg-unknown")) + require.NoError(t, st.PutAt(ctx, p0, []byte("p0"), 10, 0)) + require.NoError(t, st.PutAt(ctx, p1, []byte("p1"), 10, 0)) + require.NoError(t, st.PutAt(ctx, unknown, []byte("unknown"), 10, 0)) + + stream := &captureExportRangeVersionsStream{ctx: ctx} + err := internal.ExportRangeVersions(&pb.ExportRangeVersionsRequest{ + MaxCommitTs: 20, + KeyFamily: tc.family, + RouteStart: []byte("!sqs|route|global"), + RouteEnd: testPrefixScanEnd([]byte("!sqs|route|global")), + RangeStart: tc.prefix, + RangeEnd: testPrefixScanEnd(tc.prefix), + MaxScannedBytes: 1 << 20, + }, stream) + require.NoError(t, err) + require.Len(t, stream.responses, 1) + require.Equal(t, []*pb.MVCCVersion{ + {Key: p1, CommitTs: 10, Value: []byte("p1"), KeyFamily: tc.family}, + }, stream.responses[0].GetVersions()) + }) + } +} + func TestInternalExportRangeVersionsDerivesFilesystemChunkScanBoundsFromRouteBounds(t *testing.T) { t.Parallel() From d24ff4f87f3cc53bae2f3ed6c874a1fd2624dc25 Mon Sep 17 00:00:00 2001 From: bootjp Date: Tue, 1 Sep 2026 20:05:28 +0900 Subject: [PATCH 57/58] migration: route filesystem prefix deletes --- kv/fsm.go | 5 +++ kv/fsm_migration_fence_test.go | 75 ++++++++++++++++++++++++++++++++++ 2 files changed, 80 insertions(+) diff --git a/kv/fsm.go b/kv/fsm.go index 6931d35a9..ec6faa0b0 100644 --- a/kv/fsm.go +++ b/kv/fsm.go @@ -12,6 +12,7 @@ import ( "github.com/bootjp/elastickv/distribution" "github.com/bootjp/elastickv/internal/encryption/fsmwire" + "github.com/bootjp/elastickv/internal/fskeys" "github.com/bootjp/elastickv/internal/raftengine" "github.com/bootjp/elastickv/internal/s3keys" pb "github.com/bootjp/elastickv/proto" @@ -815,6 +816,9 @@ func routePrefixRange(prefix []byte) ([]byte, []byte) { if start, ok := dynamoExactCleanupRouteKey(prefix); ok { return start, routePointRangeEnd(start) } + if start, end, ok := fskeys.ChunkScanRouteBounds(prefix, prefixScanEnd(prefix)); ok { + return start, end + } if routeKeyspaceWideRawPrefix(prefix) { return []byte(""), nil } @@ -896,6 +900,7 @@ var routeMappedRawPrefixes = append([][]byte{ []byte(s3keys.BlobPrefix), []byte(s3keys.GCUploadPrefix), []byte(s3keys.RoutePrefix), + fskeys.UsageRouteAllPrefix(), }, sqsConcreteInternalPrefixBytes...) var ErrNotImplemented = errors.New("not implemented") diff --git a/kv/fsm_migration_fence_test.go b/kv/fsm_migration_fence_test.go index 72d4ae6a6..23660780c 100644 --- a/kv/fsm_migration_fence_test.go +++ b/kv/fsm_migration_fence_test.go @@ -5,6 +5,7 @@ import ( "testing" "github.com/bootjp/elastickv/distribution" + "github.com/bootjp/elastickv/internal/fskeys" "github.com/bootjp/elastickv/internal/s3keys" pb "github.com/bootjp/elastickv/proto" "github.com/bootjp/elastickv/store" @@ -221,6 +222,80 @@ func TestFSMDelPrefixTombstonesStagedVisibilityRowsDuringApply(t *testing.T) { require.Equal(t, []byte("outside"), got) } +func TestFSMDelPrefixRoutesFilesystemChunkPrefixThroughVirtualRange(t *testing.T) { + t.Parallel() + + ctx := context.Background() + home := uint64(11) + inode := uint64(22) + chunkPrefix := fskeys.ChunkPrefix(home, inode) + chunkRoute := fskeys.ChunkRouteKey(home, inode) + chunkRouteEnd := prefixScanEnd(chunkRoute) + chunkRawAll := fskeys.ChunkAllPrefix() + chunkRawEnd := prefixScanEnd(chunkRawAll) + + engine := distribution.NewEngine() + applyComposed1Snapshot(t, engine, 1, []distribution.RouteDescriptor{ + {RouteID: 1, Start: []byte(""), End: chunkRawAll, GroupID: 1, State: distribution.RouteStateActive}, + {RouteID: 2, Start: chunkRawAll, End: chunkRawEnd, GroupID: 1, State: distribution.RouteStateActive}, + {RouteID: 3, Start: chunkRawEnd, End: chunkRoute, GroupID: 1, State: distribution.RouteStateActive}, + { + RouteID: 4, + Start: chunkRoute, + End: chunkRouteEnd, + GroupID: 1, + State: distribution.RouteStateActive, + StagedVisibilityActive: true, + MigrationJobID: 9, + }, + {RouteID: 5, Start: chunkRouteEnd, End: nil, GroupID: 1, State: distribution.RouteStateActive}, + }) + fsm := newComposed1FSM(t, engine, 1) + stagedChunk := distribution.MigrationStagedDataKey(9, fskeys.ChunkKey(home, inode, 0)) + require.NoError(t, fsm.store.PutAt(ctx, stagedChunk, []byte("chunk"), 20, 0)) + + require.NoError(t, fsm.handleDelPrefix(ctx, chunkPrefix, 101)) + + _, err := fsm.store.GetAt(ctx, stagedChunk, 150) + require.ErrorIs(t, err, store.ErrKeyNotFound) +} + +func TestFSMRejectsFilesystemChunkDelPrefixAgainstVirtualRouteFence(t *testing.T) { + t.Parallel() + + home := uint64(11) + inode := uint64(22) + chunkRoute := fskeys.ChunkRouteKey(home, inode) + chunkRouteEnd := prefixScanEnd(chunkRoute) + chunkRawAll := fskeys.ChunkAllPrefix() + chunkRawEnd := prefixScanEnd(chunkRawAll) + + engine := distribution.NewEngine() + applyComposed1Snapshot(t, engine, 1, []distribution.RouteDescriptor{ + {RouteID: 1, Start: []byte(""), End: chunkRawAll, GroupID: 1, State: distribution.RouteStateActive}, + {RouteID: 2, Start: chunkRawAll, End: chunkRawEnd, GroupID: 1, State: distribution.RouteStateActive}, + {RouteID: 3, Start: chunkRawEnd, End: chunkRoute, GroupID: 1, State: distribution.RouteStateActive}, + {RouteID: 4, Start: chunkRoute, End: chunkRouteEnd, GroupID: 1, State: distribution.RouteStateWriteFenced}, + {RouteID: 5, Start: chunkRouteEnd, End: nil, GroupID: 1, State: distribution.RouteStateActive}, + }) + fsm := newComposed1FSM(t, engine, 1) + + err := fsm.handleRawRequest(context.Background(), &pb.Request{ + Mutations: []*pb.Mutation{{Op: pb.Op_DEL_PREFIX, Key: fskeys.ChunkPrefix(home, inode)}}, + }, 10) + require.ErrorIs(t, err, ErrRouteWriteFenced) +} + +func TestFSMRejectsFilesystemUsageRoutePrefixWhenAnyOwnerRouteIsWriteFenced(t *testing.T) { + t.Parallel() + + fsm := newWriteFencedFSM(t) + err := fsm.handleRawRequest(context.Background(), &pb.Request{ + Mutations: []*pb.Mutation{{Op: pb.Op_DEL_PREFIX, Key: fskeys.UsageRouteAllPrefix()}}, + }, 10) + require.ErrorIs(t, err, ErrRouteWriteFenced) +} + type recordingPrefixDeleteStore struct { store.MVCCStore From 2b0b18ce3b831d1108ca0db260e193744254949e Mon Sep 17 00:00:00 2001 From: bootjp Date: Tue, 1 Sep 2026 20:31:18 +0900 Subject: [PATCH 58/58] migration: replicate metadata retirement --- distribution/catalog.go | 13 ++++++ distribution/split_job_catalog.go | 6 +++ distribution/split_job_catalog_test.go | 41 +++++++++++++++++ kv/fsm.go | 5 ++ kv/fsm_migration_retire.go | 42 +++++++++++++++++ kv/fsm_migration_retire_test.go | 64 ++++++++++++++++++++++++++ kv/leader_routed_store.go | 7 +++ kv/shard_store.go | 4 ++ main.go | 35 +++++++++++++- store/lsm_migration.go | 13 +++++- store/lsm_store_applied_index_test.go | 47 +++++++++++++++++++ store/migration_versions.go | 4 ++ store/store.go | 4 ++ 13 files changed, 283 insertions(+), 2 deletions(-) create mode 100644 kv/fsm_migration_retire.go create mode 100644 kv/fsm_migration_retire_test.go diff --git a/distribution/catalog.go b/distribution/catalog.go index 7936d3b4d..aec621bd3 100644 --- a/distribution/catalog.go +++ b/distribution/catalog.go @@ -99,6 +99,7 @@ type CatalogStore struct { store store.MVCCStore allowRouteDescriptorV2Writes bool migrationStoreForGroup CatalogMigrationStoreResolver + migrationRetireForGroup CatalogMigrationRetireResolver } type CatalogStoreOption func(*CatalogStore) @@ -106,6 +107,10 @@ type CatalogStoreOption func(*CatalogStore) // CatalogMigrationStoreResolver resolves the local MVCC store for a data group. type CatalogMigrationStoreResolver func(groupID uint64) (store.MVCCStore, error) +// CatalogMigrationRetireResolver retires migration metadata through the target +// data group's replicated apply path. +type CatalogMigrationRetireResolver func(ctx context.Context, groupID, jobID uint64) error + func WithCatalogRouteDescriptorV2Writes(enabled bool) CatalogStoreOption { return func(s *CatalogStore) { s.allowRouteDescriptorV2Writes = enabled @@ -120,6 +125,14 @@ func WithCatalogMigrationStoreResolver(resolver CatalogMigrationStoreResolver) C } } +// WithCatalogMigrationRetireResolver makes split-job finalization retire +// target-local migration metadata through the target data group Raft log. +func WithCatalogMigrationRetireResolver(resolver CatalogMigrationRetireResolver) CatalogStoreOption { + return func(s *CatalogStore) { + s.migrationRetireForGroup = resolver + } +} + // NewCatalogStore creates a route catalog persistence helper. func NewCatalogStore(st store.MVCCStore, opts ...CatalogStoreOption) *CatalogStore { s := &CatalogStore{store: st} diff --git a/distribution/split_job_catalog.go b/distribution/split_job_catalog.go index 02570182e..7da615bca 100644 --- a/distribution/split_job_catalog.go +++ b/distribution/split_job_catalog.go @@ -484,6 +484,12 @@ func (s *CatalogStore) finishAppliedSplitJobHistoryMove(ctx context.Context, job } func (s *CatalogStore) retireSplitJobMigrationMetadata(ctx context.Context, job SplitJob) error { + if s.migrationRetireForGroup != nil { + if err := s.migrationRetireForGroup(ctx, job.TargetGroupID, job.JobID); err != nil && !errors.Is(err, store.ErrNotSupported) { + return errors.WithStack(err) + } + return nil + } st, err := s.migrationMetadataStore(job.TargetGroupID) if err != nil { return err diff --git a/distribution/split_job_catalog_test.go b/distribution/split_job_catalog_test.go index 351227593..c6cde5a98 100644 --- a/distribution/split_job_catalog_test.go +++ b/distribution/split_job_catalog_test.go @@ -433,6 +433,47 @@ func TestCatalogStoreMoveSplitJobToHistoryRetiresMigrationMetadata(t *testing.T) assertStoreValueAt(t, ctx, targetStore, "job14", 50, "imported") } +func TestCatalogStoreMoveSplitJobToHistoryUsesMigrationRetireResolver(t *testing.T) { + ctx := context.Background() + job := sampleSplitJob(14) + var calls []struct { + groupID uint64 + jobID uint64 + } + cs := NewCatalogStore( + store.NewMVCCStore(), + WithCatalogMigrationRetireResolver(func(_ context.Context, groupID, jobID uint64) error { + calls = append(calls, struct { + groupID uint64 + jobID uint64 + }{groupID: groupID, jobID: jobID}) + return nil + }), + ) + + if err := cs.CreateSplitJob(ctx, job); err != nil { + t.Fatalf("create split job: %v", err) + } + terminal := job + terminal.Phase = SplitJobPhaseDone + terminal.TerminalAtMs = 1000 + if err := cs.MoveSplitJobToHistory(ctx, job, terminal); err != nil { + t.Fatalf("move split job to history: %v", err) + } + if err := cs.MoveSplitJobToHistory(ctx, job, terminal); err != nil { + t.Fatalf("retry move split job to history: %v", err) + } + + if len(calls) != 2 { + t.Fatalf("expected two retire resolver calls, got %d", len(calls)) + } + for _, call := range calls { + if call.groupID != job.TargetGroupID || call.jobID != job.JobID { + t.Fatalf("unexpected retire resolver call: %+v", call) + } + } +} + func TestCatalogStoreMoveSplitJobToHistoryRetiresMetadataOnIdempotentRetry(t *testing.T) { ctx := context.Background() job := sampleSplitJob(15) diff --git a/kv/fsm.go b/kv/fsm.go index ec6faa0b0..5ccb3ec3b 100644 --- a/kv/fsm.go +++ b/kv/fsm.go @@ -408,6 +408,8 @@ func (f *kvFSM) applyReservedOpcode(ctx context.Context, data []byte) (any, bool return f.applyHLCLease(data[1:]), true case data[0] == raftEncodeMigrationImport: return f.applyMigrationImport(ctx, data[1:]), true + case data[0] == raftEncodeMigrationRetire: + return f.applyMigrationRetire(ctx, data[1:]), true case data[0] == raftEncodeMigrationPromote: return f.applyMigrationPromote(ctx, data[1:]), true case data[0] == raftEncodeBackup: @@ -447,6 +449,9 @@ const ( // batch. Every target voter applies the raw MVCC versions, import ack, and // migration HLC floor before the RPC handler returns success. raftEncodeMigrationImport byte = 0x09 + // raftEncodeMigrationRetire carries target-group metadata retirement for a + // completed migration job. Every target voter drops import/progress state. + raftEncodeMigrationRetire byte = 0x0c // 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. diff --git a/kv/fsm_migration_retire.go b/kv/fsm_migration_retire.go new file mode 100644 index 000000000..621dbec85 --- /dev/null +++ b/kv/fsm_migration_retire.go @@ -0,0 +1,42 @@ +package kv + +import ( + "context" + "encoding/binary" + + "github.com/cockroachdb/errors" +) + +const migrationRetireCommandPayloadLen = 8 + +var ErrMigrationRetireApply = errors.New("migration retire: FSM apply failed; halting apply") + +// MarshalMigrationRetireCommand encodes target-group migration metadata +// retirement as a Raft FSM command. +func MarshalMigrationRetireCommand(jobID uint64) ([]byte, error) { + if jobID == 0 { + return nil, errors.WithStack(ErrInvalidRequest) + } + payload := make([]byte, migrationRetireCommandPayloadLen) + binary.BigEndian.PutUint64(payload, jobID) + return prependByte(raftEncodeMigrationRetire, payload), nil +} + +func (f *kvFSM) applyMigrationRetire(ctx context.Context, data []byte) any { + if len(data) != migrationRetireCommandPayloadLen { + return haltErr(errors.Wrapf( + errors.Mark(ErrInvalidRequest, ErrMigrationRetireApply), + "kv/fsm: decode migration retire: expected %d bytes, got %d", + migrationRetireCommandPayloadLen, + len(data), + )) + } + jobID := binary.BigEndian.Uint64(data) + if jobID == 0 { + return errors.Wrap(ErrInvalidRequest, "kv/fsm: apply migration retire") + } + if err := f.store.RetireMigrationRaft(ctx, jobID, f.pendingApplyIdx); err != nil { + return haltErr(errors.Wrap(errors.Mark(err, ErrMigrationRetireApply), "kv/fsm: apply migration retire")) + } + return nil +} diff --git a/kv/fsm_migration_retire_test.go b/kv/fsm_migration_retire_test.go new file mode 100644 index 000000000..1bedb7361 --- /dev/null +++ b/kv/fsm_migration_retire_test.go @@ -0,0 +1,64 @@ +package kv + +import ( + "context" + "testing" + + "github.com/bootjp/elastickv/store" + "github.com/cockroachdb/errors" + "github.com/stretchr/testify/require" +) + +type captureMigrationRetireStore struct { + store.MVCCStore + jobID uint64 + appliedIndex uint64 +} + +func (s *captureMigrationRetireStore) RetireMigrationRaft(_ context.Context, jobID, appliedIndex uint64) error { + s.jobID = jobID + s.appliedIndex = appliedIndex + return nil +} + +func TestApplyMigrationRetireThreadsPendingApplyIndex(t *testing.T) { + t.Parallel() + + capturing := &captureMigrationRetireStore{} + fsm := &kvFSM{store: capturing, pendingApplyIdx: 1234} + cmd, err := MarshalMigrationRetireCommand(9) + require.NoError(t, err) + + applied := fsm.Apply(cmd) + require.Nil(t, applied) + require.Equal(t, uint64(9), capturing.jobID) + require.Equal(t, uint64(1234), capturing.appliedIndex) +} + +func TestMarshalMigrationRetireCommandRejectsZeroJobID(t *testing.T) { + t.Parallel() + + _, err := MarshalMigrationRetireCommand(0) + require.ErrorIs(t, err, ErrInvalidRequest) +} + +func TestApplyMigrationRetireMalformedPayloadHalts(t *testing.T) { + t.Parallel() + + fsm := &kvFSM{store: store.NewMVCCStore()} + err := haltApplyOf(fsm.Apply([]byte{raftEncodeMigrationRetire, 0xff})) + require.True(t, errors.Is(err, ErrMigrationRetireApply), "got %v", err) +} + +func TestApplyMigrationRetireZeroJobIDReturnsOrdinaryError(t *testing.T) { + t.Parallel() + + fsm := &kvFSM{store: store.NewMVCCStore()} + payload := make([]byte, migrationRetireCommandPayloadLen) + resp := fsm.Apply(append([]byte{raftEncodeMigrationRetire}, payload...)) + require.Nil(t, haltApplyOf(resp)) + err, ok := resp.(error) + require.True(t, ok, "got %T: %v", resp, resp) + require.ErrorIs(t, err, ErrInvalidRequest) + require.False(t, errors.Is(err, ErrMigrationRetireApply)) +} diff --git a/kv/leader_routed_store.go b/kv/leader_routed_store.go index 2a1ef53f8..20e915792 100644 --- a/kv/leader_routed_store.go +++ b/kv/leader_routed_store.go @@ -751,6 +751,13 @@ func (s *LeaderRoutedStore) RetireMigration(ctx context.Context, jobID uint64) e return errors.WithStack(s.local.RetireMigration(ctx, jobID)) } +func (s *LeaderRoutedStore) RetireMigrationRaft(ctx context.Context, jobID, appliedIndex uint64) error { + if s == nil || s.local == nil { + return errors.WithStack(store.ErrNotSupported) + } + return errors.WithStack(s.local.RetireMigrationRaft(ctx, jobID, appliedIndex)) +} + func (s *LeaderRoutedStore) Snapshot() (store.Snapshot, error) { if s == nil || s.local == nil { return nil, errors.WithStack(store.ErrNotSupported) diff --git a/kv/shard_store.go b/kv/shard_store.go index af112d6b8..529788498 100644 --- a/kv/shard_store.go +++ b/kv/shard_store.go @@ -5615,6 +5615,10 @@ func (s *ShardStore) RetireMigration(context.Context, uint64) error { return store.ErrNotSupported } +func (s *ShardStore) RetireMigrationRaft(context.Context, uint64, uint64) error { + return store.ErrNotSupported +} + func (s *ShardStore) Restore(_ io.Reader) error { return store.ErrNotSupported } diff --git a/main.go b/main.go index c38871fec..e4504e6fe 100644 --- a/main.go +++ b/main.go @@ -3406,7 +3406,10 @@ func distributionCatalogStoreForGroup(runtimes []*raftGroupRuntime, groupID uint continue } if rt.spec.id == groupID { - return distribution.NewCatalogStore(rt.store, distribution.WithCatalogMigrationStoreResolver(catalogMigrationStoreResolver(runtimes))) + return distribution.NewCatalogStore(rt.store, + distribution.WithCatalogMigrationStoreResolver(catalogMigrationStoreResolver(runtimes)), + distribution.WithCatalogMigrationRetireResolver(catalogMigrationRetireResolver(runtimes)), + ) } } return nil @@ -3426,6 +3429,36 @@ func catalogMigrationStoreResolver(runtimes []*raftGroupRuntime) distribution.Ca } } +func catalogMigrationRetireResolver(runtimes []*raftGroupRuntime) distribution.CatalogMigrationRetireResolver { + return func(ctx context.Context, groupID, jobID uint64) error { + for _, rt := range runtimes { + if rt == nil || rt.spec.id != groupID { + continue + } + engine := rt.snapshotEngine() + if engine == nil { + break + } + cmd, err := kv.MarshalMigrationRetireCommand(jobID) + if err != nil { + return errors.WithStack(err) + } + result, err := engine.Propose(ctx, cmd) + if err != nil { + return errors.WithStack(err) + } + if result == nil { + return errors.Newf("migration retire proposal returned nil result for group %d job %d", groupID, jobID) + } + if err, ok := result.Response.(error); ok { + return errors.WithStack(err) + } + return nil + } + return errors.Newf("migration target engine is not available for group %d", groupID) + } +} + func setupDistributionCatalog( ctx context.Context, runtimes []*raftGroupRuntime, diff --git a/store/lsm_migration.go b/store/lsm_migration.go index de2f6a1aa..ce708e928 100644 --- a/store/lsm_migration.go +++ b/store/lsm_migration.go @@ -580,6 +580,14 @@ func (s *pebbleStore) MigrationHLCFloor(_ context.Context, jobID uint64) (uint64 } func (s *pebbleStore) RetireMigration(ctx context.Context, jobID uint64) error { + return s.retireMigrationWithOpts(ctx, jobID, s.directApplyWriteOpts(), 0) +} + +func (s *pebbleStore) RetireMigrationRaft(ctx context.Context, jobID, appliedIndex uint64) error { + return s.retireMigrationWithOpts(ctx, jobID, s.raftApplyWriteOpts(), appliedIndex) +} + +func (s *pebbleStore) retireMigrationWithOpts(ctx context.Context, jobID uint64, writeOpts *pebble.WriteOptions, appliedIndex uint64) error { if err := ctx.Err(); err != nil { return errors.WithStack(err) } @@ -597,7 +605,10 @@ func (s *pebbleStore) RetireMigration(ctx context.Context, jobID uint64) error { if err := s.stageRetireMigrationPromotionState(batch, jobID); err != nil { return err } - return errors.WithStack(batch.Commit(s.directApplyWriteOpts())) + if err := stagePebbleAppliedIndex(batch, appliedIndex); err != nil { + return err + } + return errors.WithStack(batch.Commit(writeOpts)) } func (s *pebbleStore) stageRetireMigrationImportAcks(batch *pebble.Batch, jobID uint64) error { diff --git a/store/lsm_store_applied_index_test.go b/store/lsm_store_applied_index_test.go index bad8be62a..2811e7150 100644 --- a/store/lsm_store_applied_index_test.go +++ b/store/lsm_store_applied_index_test.go @@ -210,6 +210,53 @@ func TestImportVersionsRaft_BundlesMetaAppliedIndex(t *testing.T) { require.Equal(t, []byte("v100"), val, "duplicate import must not rewrite the acknowledged batch") } +func TestRetireMigrationRaft_BundlesMetaAppliedIndex(t *testing.T) { + ctx := context.Background() + st := newApplyIndexPebbleStore(t) + ps := pebbleStoreApplied(t, st) + + _, err := ps.ImportVersions(ctx, ImportVersionsOptions{ + JobID: 9, + BracketID: 1, + BatchSeq: 1, + Cursor: []byte("old"), + Versions: []MVCCVersion{ + {Key: []byte("stage|k"), CommitTS: 100, Value: []byte("v100")}, + }, + }) + require.NoError(t, err) + seedPromotionState(t, ctx, ps, 9, []byte("promote|"), []byte("promoted")) + + const entryIdx uint64 = 125 + require.NoError(t, ps.RetireMigrationRaft(ctx, 9, entryIdx)) + + got, present, err := ps.LastAppliedIndex() + require.NoError(t, err) + require.True(t, present, "RetireMigrationRaft must persist metaAppliedIndex") + require.Equal(t, entryIdx, got) + + floor, err := ps.MigrationHLCFloor(ctx, 9) + require.NoError(t, err) + require.Zero(t, floor) + state, ok, err := ps.MigrationPromotionState(ctx, 9) + require.NoError(t, err) + require.False(t, ok, "promotion state must be retired, got %+v", state) + + again, err := ps.ImportVersions(ctx, ImportVersionsOptions{ + JobID: 9, + BracketID: 1, + BatchSeq: 1, + Cursor: []byte("fresh"), + }) + require.NoError(t, err) + require.False(t, again.Duplicate, "retired import ack must not force duplicate detection") + require.Equal(t, []byte("fresh"), again.AckedCursor) + + val, err := ps.GetAt(ctx, []byte("stage|k"), 100) + require.NoError(t, err) + require.Equal(t, []byte("v100"), val, "retirement must leave imported data intact") +} + func TestApplyMutationsRaftAt_AlreadyLandedAdvancesStaleAppliedIndex(t *testing.T) { ctx := context.Background() st := newApplyIndexPebbleStore(t) diff --git a/store/migration_versions.go b/store/migration_versions.go index e68eb273f..32ef2ccc7 100644 --- a/store/migration_versions.go +++ b/store/migration_versions.go @@ -709,3 +709,7 @@ func (s *mvccStore) RetireMigration(ctx context.Context, jobID uint64) error { delete(s.migrationPromotions, jobID) return nil } + +func (s *mvccStore) RetireMigrationRaft(ctx context.Context, jobID, _ uint64) error { + return s.RetireMigration(ctx, jobID) +} diff --git a/store/store.go b/store/store.go index 853d70b69..4090d2340 100644 --- a/store/store.go +++ b/store/store.go @@ -364,6 +364,10 @@ type MVCCStore interface { // RetireMigration removes target-local import progress metadata for a // completed migration job. Data imported by the job is left intact. RetireMigration(ctx context.Context, jobID uint64) error + // RetireMigrationRaft is the raft-apply variant of RetireMigration. When + // appliedIndex is non-zero, the implementation must durably bundle + // metaAppliedIndex with the metadata retirement batch. + RetireMigrationRaft(ctx context.Context, jobID, appliedIndex uint64) error Snapshot() (Snapshot, error) Restore(buf io.Reader) error Close() error