From 2a8585283c15e2ab50007b24eb4ff1d3de2b1b10 Mon Sep 17 00:00:00 2001 From: bootjp Date: Mon, 13 Jul 2026 22:51:40 +0900 Subject: [PATCH 01/13] migration: add range version RPC handlers --- adapter/internal.go | 138 +++++++++++++++++++++++++++++ adapter/internal_migration_test.go | 80 +++++++++++++++++ adapter/test_util.go | 8 +- distribution/engine.go | 55 ++++++------ distribution/engine_test.go | 24 ++++- main.go | 5 +- main_bootstrap_e2e_test.go | 8 +- 7 files changed, 286 insertions(+), 32 deletions(-) create mode 100644 adapter/internal_migration_test.go diff --git a/adapter/internal.go b/adapter/internal.go index aa98f16b9..ae9ac2e34 100644 --- a/adapter/internal.go +++ b/adapter/internal.go @@ -7,7 +7,10 @@ import ( "github.com/bootjp/elastickv/internal/raftengine" "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,6 +21,12 @@ func WithInternalTimestampAllocator(alloc kv.TimestampAllocator) InternalOption } } +func WithInternalStore(st store.MVCCStore) InternalOption { + return func(i *Internal) { + i.store = st + } +} + func NewInternalWithEngine(txm kv.Transactional, leader raftengine.LeaderView, clock *kv.HLC, relay *RedisPubSubRelay, opts ...InternalOption) *Internal { i := &Internal{ leader: leader, @@ -37,6 +46,7 @@ type Internal struct { clock *kv.HLC tsAllocator kv.TimestampAllocator relay *RedisPubSubRelay + store store.MVCCStore pb.UnimplementedInternalServer } @@ -47,6 +57,12 @@ var ErrNotLeader = errors.New("not leader") var ErrLeaderNotFound = errors.New("leader not found") var ErrTxnTimestampOverflow = errors.New("txn timestamp overflow") +const ( + defaultMigrationExportChunkBytes = 4 << 20 + defaultMigrationExportScanFactor = 4 + defaultMigrationExportMaxVersions = 1024 +) + func (i *Internal) Forward(ctx context.Context, req *pb.ForwardRequest) (*pb.ForwardResponse, error) { if i.leader == nil || i.leader.State() != raftengine.StateLeader { return nil, errors.WithStack(ErrNotLeader) @@ -85,6 +101,128 @@ func (i *Internal) RelayPublish(_ context.Context, req *pb.RelayPublishRequest) }, nil } +func (i *Internal) ExportRangeVersions(req *pb.ExportRangeVersionsRequest, stream pb.Internal_ExportRangeVersionsServer) 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 err := i.verifyInternalLeader(stream.Context()); err != nil { + return err + } + opts := exportRangeVersionsOptions(req) + for { + result, err := i.store.ExportVersions(stream.Context(), opts) + if err != nil { + return errors.WithStack(err) + } + 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 i.store == nil { + return nil, errors.WithStack(status.Error(codes.FailedPrecondition, "migration import store is not configured")) + } + if err := i.verifyInternalLeader(ctx); err != nil { + return nil, err + } + result, err := i.store.ImportVersions(ctx, store.ImportVersionsOptions{ + JobID: req.GetJobId(), + BracketID: req.GetBracketId(), + BatchSeq: req.GetBatchSeq(), + Cursor: req.GetCursor(), + Versions: storeMVCCVersionsFromProto(req.GetVersions()), + }) + if err != nil { + return nil, errors.WithStack(err) + } + return &pb.ImportRangeVersionsResponse{AckedCursor: result.AckedCursor}, nil +} + +func (i *Internal) verifyInternalLeader(ctx context.Context) error { + if i.leader == nil { + return nil + } + 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 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, + AcceptKey: kv.RouteKeyFilter(req.GetRouteStart(), req.GetRouteEnd()), + } + return opts +} + +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 storeMVCCVersionsFromProto(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: bytes.Clone(version.GetKey()), + CommitTS: version.GetCommitTs(), + Tombstone: version.GetTombstone(), + Value: bytes.Clone(version.GetValue()), + KeyFamily: version.GetKeyFamily(), + ExpireAt: version.GetExpireAt(), + }) + } + return out +} + func (i *Internal) stampTimestamps(ctx context.Context, req *pb.ForwardRequest) error { if req == nil { return nil diff --git a/adapter/internal_migration_test.go b/adapter/internal_migration_test.go new file mode 100644 index 000000000..5d2e30179 --- /dev/null +++ b/adapter/internal_migration_test.go @@ -0,0 +1,80 @@ +package adapter + +import ( + "context" + "testing" + + pb "github.com/bootjp/elastickv/proto" + "github.com/bootjp/elastickv/store" + "github.com/stretchr/testify/require" + "google.golang.org/grpc" +) + +type captureExportRangeVersionsStream struct { + grpc.ServerStream + ctx context.Context + responses []*pb.ExportRangeVersionsResponse +} + +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 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)) + internal := NewInternalWithEngine(nil, nil, nil, nil, WithInternalStore(st)) + stream := &captureExportRangeVersionsStream{ctx: ctx} + + err := internal.ExportRangeVersions(&pb.ExportRangeVersionsRequest{ + MaxCommitTs: 20, + RouteStart: []byte("a"), + RouteEnd: []byte("b"), + }, 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")}, + }, stream.responses[0].GetVersions()) +} + +func TestInternalImportRangeVersionsAppliesStoreBatch(t *testing.T) { + t.Parallel() + + ctx := context.Background() + st := store.NewMVCCStore() + internal := NewInternalWithEngine(nil, nil, nil, nil, WithInternalStore(st)) + + 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()) + + got, err := st.GetAt(ctx, []byte("k"), 30) + require.NoError(t, err) + require.Equal(t, []byte("v"), got) + floor, err := st.MigrationHLCFloor(ctx, 7) + require.NoError(t, err) + require.Equal(t, uint64(30), floor) +} diff --git a/adapter/test_util.go b/adapter/test_util.go index 8fe5646fa..6804657c5 100644 --- a/adapter/test_util.go +++ b/adapter/test_util.go @@ -641,7 +641,13 @@ 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), + )) internalraftadmin.RegisterOperationalServices(opsCtx, s, result.Engine, []string{"Example"}) grpcAdders = append(grpcAdders, port.grpcAddress) diff --git a/distribution/engine.go b/distribution/engine.go index dc55ed74e..8f6868021 100644 --- a/distribution/engine.go +++ b/distribution/engine.go @@ -25,6 +25,13 @@ type Route struct { GroupID uint64 // State tracks control-plane state for this route. State RouteState + // StagedVisibilityActive makes migrated versions visible through the + // staged/live merge path after cross-group CUTOVER. + StagedVisibilityActive bool + // MigrationJobID identifies the migration job that owns staged visibility. + MigrationJobID uint64 + // MinWriteTSExclusive is the post-migration write timestamp floor. + MinWriteTSExclusive uint64 // Load tracks the number of accesses served by this range. Load uint64 } @@ -365,14 +372,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, - Load: r.Load, - } + stats[i] = cloneRoute(r) } return stats } @@ -398,26 +398,22 @@ func (e *Engine) GetIntersectingRoutes(start, end []byte) []Route { break } // Route intersects with scan range - result = append(result, Route{ - RouteID: r.RouteID, - Start: CloneBytes(r.Start), - End: CloneBytes(r.End), - GroupID: r.GroupID, - State: r.State, - Load: r.Load, - }) + result = append(result, cloneRoute(*r)) } return result } func cloneRoute(r Route) Route { return Route{ - RouteID: r.RouteID, - Start: CloneBytes(r.Start), - End: CloneBytes(r.End), - GroupID: r.GroupID, - State: r.State, - Load: r.Load, + 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, } } @@ -454,12 +450,15 @@ func routesFromCatalog(routes []RouteDescriptor) ([]Route, error) { } seen[rd.RouteID] = struct{}{} out[i] = Route{ - RouteID: rd.RouteID, - Start: CloneBytes(rd.Start), - End: CloneBytes(rd.End), - GroupID: rd.GroupID, - State: rd.State, - Load: 0, + RouteID: rd.RouteID, + Start: CloneBytes(rd.Start), + End: CloneBytes(rd.End), + GroupID: rd.GroupID, + State: rd.State, + StagedVisibilityActive: rd.StagedVisibilityActive, + MigrationJobID: rd.MigrationJobID, + MinWriteTSExclusive: rd.MinWriteTSExclusive, + Load: 0, } } diff --git a/distribution/engine_test.go b/distribution/engine_test.go index 346e4fa2d..1794608c5 100644 --- a/distribution/engine_test.go +++ b/distribution/engine_test.go @@ -169,7 +169,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 { @@ -190,6 +199,19 @@ func TestEngineApplySnapshot_ReplacesRoutesAndVersion(t *testing.T) { if stats[1].RouteID != 11 || stats[1].State != RouteStateWriteFenced { t.Fatalf("unexpected second route metadata: %+v", stats[1]) } + assertStagedRouteMetadata(t, stats[1]) + route, ok := e.GetRoute([]byte("m")) + if !ok { + t.Fatalf("expected route for m") + } + assertStagedRouteMetadata(t, route) +} + +func assertStagedRouteMetadata(t *testing.T, route Route) { + t.Helper() + if !route.StagedVisibilityActive || route.MigrationJobID != 42 || route.MinWriteTSExclusive != 99 { + t.Fatalf("staged route metadata was not preserved: %+v", route) + } } func TestEngineApplySnapshot_RejectsOldVersion(t *testing.T) { diff --git a/main.go b/main.go index 530dbf4b7..bd111ea59 100644 --- a/main.go +++ b/main.go @@ -2093,7 +2093,10 @@ func startRaftServers( rt.engine, coordinate.Clock(), relay, - internalTimestampOptions(coordinate)..., + append( + internalTimestampOptions(coordinate), + adapter.WithInternalStore(rt.store), + )..., )) pb.RegisterDistributionServer(gs, distServer) if adminServer != nil { diff --git a/main_bootstrap_e2e_test.go b/main_bootstrap_e2e_test.go index d962c8105..c8434716a 100644 --- a/main_bootstrap_e2e_test.go +++ b/main_bootstrap_e2e_test.go @@ -762,7 +762,13 @@ 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), + )) pb.RegisterDistributionServer(gs, distServer) rt.registerGRPC(gs) internalraftadmin.RegisterOperationalServices(ctx, gs, rt.engine, []string{"RawKV"}) From fbd7f563b168cc03a183f27bf57ad4f767f9ccda Mon Sep 17 00:00:00 2001 From: bootjp Date: Mon, 13 Jul 2026 23:20:52 +0900 Subject: [PATCH 02/13] migration: raft-apply range imports --- adapter/internal.go | 124 +++++++++++++++++++++-------- adapter/internal_migration_test.go | 83 +++++++++++++++++-- adapter/test_util.go | 1 + kv/fsm.go | 13 ++- kv/fsm_migration_import.go | 83 +++++++++++++++++++ main.go | 1 + main_bootstrap_e2e_test.go | 1 + proto/internal.pb.go | 39 ++++++++- proto/internal.proto | 9 +++ 9 files changed, 309 insertions(+), 45 deletions(-) create mode 100644 kv/fsm_migration_import.go diff --git a/adapter/internal.go b/adapter/internal.go index ae9ac2e34..de61f5485 100644 --- a/adapter/internal.go +++ b/adapter/internal.go @@ -4,6 +4,7 @@ import ( "bytes" "context" + "github.com/bootjp/elastickv/distribution" "github.com/bootjp/elastickv/internal/raftengine" "github.com/bootjp/elastickv/kv" pb "github.com/bootjp/elastickv/proto" @@ -27,6 +28,12 @@ func WithInternalStore(st store.MVCCStore) InternalOption { } } +func WithInternalMigrationProposer(proposer raftengine.Proposer) InternalOption { + return func(i *Internal) { + i.migrationProposer = proposer + } +} + func NewInternalWithEngine(txm kv.Transactional, leader raftengine.LeaderView, clock *kv.HLC, relay *RedisPubSubRelay, opts ...InternalOption) *Internal { i := &Internal{ leader: leader, @@ -47,6 +54,7 @@ type Internal struct { tsAllocator kv.TimestampAllocator relay *RedisPubSubRelay store store.MVCCStore + migrationProposer raftengine.Proposer pb.UnimplementedInternalServer } @@ -102,21 +110,41 @@ func (i *Internal) RelayPublish(_ context.Context, req *pb.RelayPublishRequest) } func (i *Internal) ExportRangeVersions(req *pb.ExportRangeVersionsRequest, stream pb.Internal_ExportRangeVersionsServer) error { + if err := i.validateExportRangeVersionsRequest(req); err != nil { + return err + } + if err := i.verifyInternalLeader(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 err := i.verifyInternalLeader(stream.Context()); err != nil { - return err + 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")) + } + return nil +} + +func (i *Internal) streamExportRangeVersions(req *pb.ExportRangeVersionsRequest, stream pb.Internal_ExportRangeVersionsServer) error { opts := exportRangeVersionsOptions(req) for { result, err := i.store.ExportVersions(stream.Context(), opts) if err != nil { return errors.WithStack(err) } + if !result.Done && bytes.Equal(opts.Cursor, result.NextCursor) { + return errors.WithStack(status.Error(codes.Internal, "migration export cursor did not progress")) + } if err := stream.Send(&pb.ExportRangeVersionsResponse{ Versions: protoMVCCVersionsFromStore(result.Versions), NextCursor: result.NextCursor, @@ -135,29 +163,23 @@ func (i *Internal) ImportRangeVersions(ctx context.Context, req *pb.ImportRangeV if req == nil { return nil, errors.WithStack(status.Error(codes.InvalidArgument, "import range versions request is nil")) } - if i.store == nil { - return nil, errors.WithStack(status.Error(codes.FailedPrecondition, "migration import store is not configured")) + if i.migrationProposer == nil { + return nil, errors.WithStack(status.Error(codes.FailedPrecondition, "migration import proposer is not configured")) } if err := i.verifyInternalLeader(ctx); err != nil { return nil, err } - result, err := i.store.ImportVersions(ctx, store.ImportVersionsOptions{ - JobID: req.GetJobId(), - BracketID: req.GetBracketId(), - BatchSeq: req.GetBatchSeq(), - Cursor: req.GetCursor(), - Versions: storeMVCCVersionsFromProto(req.GetVersions()), - }) + 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 (i *Internal) verifyInternalLeader(ctx context.Context) error { - if i.leader == nil { - return nil - } if i.leader.State() != raftengine.StateLeader { return errors.WithStack(ErrNotLeader) } @@ -167,6 +189,33 @@ func (i *Internal) verifyInternalLeader(ctx context.Context) error { return nil } +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) + } + result, err := i.migrationProposer.Propose(ctx, cmd) + if err != nil { + return store.ImportVersionsResult{}, errors.WithStack(err) + } + if result == nil { + return store.ImportVersionsResult{}, errors.New("migration import proposal returned nil result") + } + switch resp := result.Response.(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 exportRangeVersionsOptions(req *pb.ExportRangeVersionsRequest) store.ExportVersionsOptions { chunkBytes := uint64(req.GetChunkBytes()) if chunkBytes == 0 { @@ -185,11 +234,38 @@ func exportRangeVersionsOptions(req *pb.ExportRangeVersionsRequest) store.Export MaxVersions: defaultMigrationExportMaxVersions, MaxBytes: chunkBytes, MaxScannedBytes: maxScannedBytes, - AcceptKey: kv.RouteKeyFilter(req.GetRouteStart(), req.GetRouteEnd()), + KeyFamily: req.GetKeyFamily(), + AcceptKey: migrationExportFilter(req), } return opts } +func migrationExportFilter(req *pb.ExportRangeVersionsRequest) func([]byte) bool { + routeFilter := kv.RouteKeyFilter(req.GetRouteStart(), req.GetRouteEnd()) + excludeKnownInternal := req.GetExcludeKnownInternal() || req.GetKeyFamily() == distribution.MigrationFamilyUser + bracket := distribution.MigrationBracket{ + Family: req.GetKeyFamily(), + Start: bytes.Clone(req.GetRangeStart()), + End: bytes.Clone(req.GetRangeEnd()), + ExcludeKnownInternal: excludeKnownInternal, + ExcludePrefixes: cloneByteSlices(req.GetExcludePrefixes()), + } + return func(rawKey []byte) bool { + return bracket.ContainsRawKey(rawKey) && routeFilter(rawKey) + } +} + +func 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 { @@ -205,24 +281,6 @@ func protoMVCCVersionsFromStore(in []store.MVCCVersion) []*pb.MVCCVersion { return out } -func storeMVCCVersionsFromProto(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: bytes.Clone(version.GetKey()), - CommitTS: version.GetCommitTs(), - Tombstone: version.GetTombstone(), - Value: bytes.Clone(version.GetValue()), - KeyFamily: version.GetKeyFamily(), - ExpireAt: version.GetExpireAt(), - }) - } - return out -} - func (i *Internal) stampTimestamps(ctx context.Context, req *pb.ForwardRequest) error { if req == nil { return nil diff --git a/adapter/internal_migration_test.go b/adapter/internal_migration_test.go index 5d2e30179..56d1b1350 100644 --- a/adapter/internal_migration_test.go +++ b/adapter/internal_migration_test.go @@ -4,12 +4,54 @@ import ( "context" "testing" + "github.com/bootjp/elastickv/distribution" + "github.com/bootjp/elastickv/internal/raftengine" + "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 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 @@ -35,29 +77,56 @@ func TestInternalExportRangeVersionsUsesStoreAndRouteFilter(t *testing.T) { 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)) - internal := NewInternalWithEngine(nil, nil, nil, nil, WithInternalStore(st)) + 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"), + 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")}, + {Key: []byte("a"), CommitTs: 10, Value: []byte("va"), KeyFamily: distribution.MigrationFamilyUser}, }, stream.responses[0].GetVersions()) } +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)) +} + func TestInternalImportRangeVersionsAppliesStoreBatch(t *testing.T) { t.Parallel() ctx := context.Background() st := store.NewMVCCStore() - internal := NewInternalWithEngine(nil, nil, nil, nil, WithInternalStore(st)) + clock := kv.NewHLC() + proposer := &applyingMigrationProposer{ + fsm: kv.NewKvFSMWithHLC(st, clock), + } + internal := NewInternalWithEngine(nil, mockInternalLeader{}, clock, nil, + WithInternalStore(st), + WithInternalMigrationProposer(proposer), + ) resp, err := internal.ImportRangeVersions(ctx, &pb.ImportRangeVersionsRequest{ JobId: 7, @@ -70,6 +139,7 @@ func TestInternalImportRangeVersionsAppliesStoreBatch(t *testing.T) { }) require.NoError(t, err) require.Equal(t, []byte("cursor-1"), resp.GetAckedCursor()) + require.Equal(t, uint64(1), proposer.calls) got, err := st.GetAt(ctx, []byte("k"), 30) require.NoError(t, err) @@ -77,4 +147,5 @@ func TestInternalImportRangeVersionsAppliesStoreBatch(t *testing.T) { floor, err := st.MigrationHLCFloor(ctx, 7) require.NoError(t, err) require.Equal(t, uint64(30), floor) + require.GreaterOrEqual(t, clock.Current(), uint64(30)) } diff --git a/adapter/test_util.go b/adapter/test_util.go index 6804657c5..65fd07de8 100644 --- a/adapter/test_util.go +++ b/adapter/test_util.go @@ -647,6 +647,7 @@ func setupNodes(t *testing.T, ctx context.Context, n int, ports []portsAdress) ( coordinator.Clock(), relay, WithInternalStore(st), + WithInternalMigrationProposer(result.Engine), )) internalraftadmin.RegisterOperationalServices(opsCtx, s, result.Engine, []string{"Example"}) diff --git a/kv/fsm.go b/kv/fsm.go index 5cd95d05d..00e4a5d3c 100644 --- a/kv/fsm.go +++ b/kv/fsm.go @@ -310,12 +310,11 @@ type fsmApplyResponse struct { } func (f *kvFSM) Apply(data []byte) any { - if resp, handled := f.applyReservedOpcode(data); handled { + ctx := context.TODO() + if resp, handled := f.applyReservedOpcode(ctx, data); handled { return resp } - ctx := context.TODO() - reqs, err := decodeRaftRequests(data) if err != nil { return errors.WithStack(err) @@ -356,13 +355,15 @@ 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] >= fsmwire.OpEncryptionMin && data[0] <= fsmwire.OpEncryptionMax: return f.applyEncryption(f.pendingApplyIdx, data[0], data[1:]), true default: @@ -394,6 +395,10 @@ 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 ) func decodeRaftRequests(data []byte) ([]*pb.Request, error) { diff --git a/kv/fsm_migration_import.go b/kv/fsm_migration_import.go new file mode 100644 index 000000000..d6a048130 --- /dev/null +++ b/kv/fsm_migration_import.go @@ -0,0 +1,83 @@ +package kv + +import ( + "bytes" + "context" + + 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.ImportVersions(ctx, store.ImportVersionsOptions{ + JobID: req.GetJobId(), + BracketID: req.GetBracketId(), + BatchSeq: req.GetBatchSeq(), + Cursor: req.GetCursor(), + Versions: migrationStoreVersionsFromProto(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(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: bytes.Clone(version.GetKey()), + CommitTS: version.GetCommitTs(), + Tombstone: version.GetTombstone(), + Value: bytes.Clone(version.GetValue()), + KeyFamily: version.GetKeyFamily(), + ExpireAt: version.GetExpireAt(), + }) + } + return out +} diff --git a/main.go b/main.go index bd111ea59..c236598d8 100644 --- a/main.go +++ b/main.go @@ -2096,6 +2096,7 @@ func startRaftServers( append( internalTimestampOptions(coordinate), adapter.WithInternalStore(rt.store), + adapter.WithInternalMigrationProposer(proposerForGroup(rt, shardGroups)), )..., )) pb.RegisterDistributionServer(gs, distServer) diff --git a/main_bootstrap_e2e_test.go b/main_bootstrap_e2e_test.go index c8434716a..81ca49dce 100644 --- a/main_bootstrap_e2e_test.go +++ b/main_bootstrap_e2e_test.go @@ -768,6 +768,7 @@ func startBoundGRPCServer( coordinate.Clock(), relay, adapter.WithInternalStore(rt.store), + adapter.WithInternalMigrationProposer(rt.engine), )) pb.RegisterDistributionServer(gs, distServer) rt.registerGRPC(gs) diff --git a/proto/internal.pb.go b/proto/internal.pb.go index 0e2bd4664..a7a9e30ce 100644 --- a/proto/internal.pb.go +++ b/proto/internal.pb.go @@ -540,6 +540,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 } @@ -637,6 +646,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"` @@ -929,7 +959,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" + @@ -942,7 +972,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" + diff --git a/proto/internal.proto b/proto/internal.proto index 1f7c77712..b60c83793 100644 --- a/proto/internal.proto +++ b/proto/internal.proto @@ -91,6 +91,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 { From 0fe341daea7fa1747c8e7e4a17e77608291747b0 Mon Sep 17 00:00:00 2001 From: bootjp Date: Mon, 13 Jul 2026 23:43:10 +0900 Subject: [PATCH 03/13] distribution: serve versioned ownership lookups --- adapter/distribution_server.go | 82 ++++++++++++++++-- adapter/distribution_server_test.go | 124 ++++++++++++++++++++++++++++ 2 files changed, 197 insertions(+), 9 deletions(-) diff --git a/adapter/distribution_server.go b/adapter/distribution_server.go index 3001b1329..9703c828b 100644 --- a/adapter/distribution_server.go +++ b/adapter/distribution_server.go @@ -68,15 +68,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. @@ -132,6 +133,56 @@ func (s *DistributionServer) ListRoutes(ctx context.Context, req *pb.ListRoutesR }, nil } +func (s *DistributionServer) GetRouteOwnership(ctx context.Context, req *pb.GetRouteOwnershipRequest) (*pb.GetRouteOwnershipResponse, error) { + 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) { + 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. @@ -525,6 +576,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 7604d0e91..f717dfb56 100644 --- a/adapter/distribution_server_test.go +++ b/adapter/distribution_server_test.go @@ -129,6 +129,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() From 7d3b01cd74c6221cc05da177e93a36134321ec4d Mon Sep 17 00:00:00 2001 From: bootjp Date: Mon, 13 Jul 2026 23:52:48 +0900 Subject: [PATCH 04/13] migration: stage imported versions --- adapter/internal_migration_test.go | 5 +- distribution/migrator.go | 36 +++++++++++ distribution/migrator_export_plan_test.go | 22 +++++++ kv/fsm_migration_import.go | 7 ++- kv/fsm_migration_import_test.go | 75 +++++++++++++++++++++++ 5 files changed, 141 insertions(+), 4 deletions(-) create mode 100644 kv/fsm_migration_import_test.go diff --git a/adapter/internal_migration_test.go b/adapter/internal_migration_test.go index 56d1b1350..dd8b48baf 100644 --- a/adapter/internal_migration_test.go +++ b/adapter/internal_migration_test.go @@ -141,9 +141,12 @@ func TestInternalImportRangeVersionsAppliesStoreBatch(t *testing.T) { require.Equal(t, []byte("cursor-1"), resp.GetAckedCursor()) require.Equal(t, uint64(1), proposer.calls) - got, err := st.GetAt(ctx, []byte("k"), 30) + 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) diff --git a/distribution/migrator.go b/distribution/migrator.go index c795f70f6..3b8df4876 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" @@ -68,6 +69,7 @@ const ( migrationDynamoGenPrefix = "!ddb|meta|gen|" migrationDynamoItemPrefix = "!ddb|item|" migrationDynamoGSIPrefix = "!ddb|gsi|" + migrationStagedDataPrefix = "!dist|migstage|" ) const ( @@ -81,6 +83,8 @@ const ( migrationSQSMsgGroupPrefix = "!sqs|msg|group|" migrationSQSMsgByAgePrefix = "!sqs|msg|byage|" migrationSQSPartitionedSuffix = "p|" + migrationStagedDataJobIDBytes = 8 + migrationStagedDataSeparator = byte('|') ) var ( @@ -237,6 +241,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 +} + // InitializeSplitJobPlan validates the source route and seeds the job's // bracket progress for the moving right child [SplitKey, source.End). func InitializeSplitJobPlan(job SplitJob, source RouteDescriptor, nowMs int64) (SplitJob, error) { diff --git a/distribution/migrator_export_plan_test.go b/distribution/migrator_export_plan_test.go index 6888aa3bf..d445876e4 100644 --- a/distribution/migrator_export_plan_test.go +++ b/distribution/migrator_export_plan_test.go @@ -177,6 +177,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/kv/fsm_migration_import.go b/kv/fsm_migration_import.go index d6a048130..9de47d897 100644 --- a/kv/fsm_migration_import.go +++ b/kv/fsm_migration_import.go @@ -4,6 +4,7 @@ import ( "bytes" "context" + "github.com/bootjp/elastickv/distribution" pb "github.com/bootjp/elastickv/proto" "github.com/bootjp/elastickv/store" "github.com/cockroachdb/errors" @@ -38,7 +39,7 @@ func (f *kvFSM) applyMigrationImport(ctx context.Context, data []byte) any { BracketID: req.GetBracketId(), BatchSeq: req.GetBatchSeq(), Cursor: req.GetCursor(), - Versions: migrationStoreVersionsFromProto(req.GetVersions()), + Versions: migrationStoreVersionsFromProto(req.GetJobId(), req.GetVersions()), }) if err != nil { return errors.WithStack(err) @@ -64,14 +65,14 @@ func (f *kvFSM) migrationHLCFloorForApply(ctx context.Context, req *pb.ImportRan return floor, nil } -func migrationStoreVersionsFromProto(in []*pb.MVCCVersion) []store.MVCCVersion { +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: bytes.Clone(version.GetKey()), + Key: distribution.MigrationStagedDataKey(jobID, version.GetKey()), CommitTS: version.GetCommitTs(), Tombstone: version.GetTombstone(), Value: bytes.Clone(version.GetValue()), diff --git a/kv/fsm_migration_import_test.go b/kv/fsm_migration_import_test.go new file mode 100644 index 000000000..589b6e1b1 --- /dev/null +++ b/kv/fsm_migration_import_test.go @@ -0,0 +1,75 @@ +package kv + +import ( + "context" + "testing" + + "github.com/bootjp/elastickv/distribution" + pb "github.com/bootjp/elastickv/proto" + "github.com/bootjp/elastickv/store" + "github.com/stretchr/testify/require" + "google.golang.org/protobuf/proto" +) + +func TestMigrationStoreVersionsFromProtoStagesKeys(t *testing.T) { + t.Parallel() + + rawKey := []byte("user|k") + value := []byte("value") + got := migrationStoreVersionsFromProto(7, []*pb.MVCCVersion{ + nil, + { + Key: rawKey, + CommitTs: 11, + Value: value, + KeyFamily: distribution.MigrationFamilyUser, + ExpireAt: 123, + }, + }) + + require.Len(t, got, 1) + require.Equal(t, distribution.MigrationStagedDataKey(7, []byte("user|k")), got[0].Key) + require.Equal(t, uint64(11), got[0].CommitTS) + require.Equal(t, []byte("value"), got[0].Value) + require.Equal(t, distribution.MigrationFamilyUser, got[0].KeyFamily) + require.Equal(t, uint64(123), got[0].ExpireAt) + + rawKey[0] = 'X' + value[0] = 'X' + require.Equal(t, distribution.MigrationStagedDataKey(7, []byte("user|k")), got[0].Key) + require.Equal(t, []byte("value"), got[0].Value) +} + +func TestApplyMigrationImportWritesOnlyStagedKeys(t *testing.T) { + t.Parallel() + + ctx := context.Background() + st := store.NewMVCCStore() + hlc := NewHLC() + fsm := &kvFSM{store: st, hlc: hlc} + req := &pb.ImportRangeVersionsRequest{ + JobId: 9, + BracketId: 1, + BatchSeq: 1, + Cursor: []byte("cursor"), + Versions: []*pb.MVCCVersion{ + {Key: []byte("user|k"), CommitTs: 10, Value: []byte("v")}, + }, + } + data, err := proto.Marshal(req) + require.NoError(t, err) + + applied := fsm.applyMigrationImport(ctx, data) + result, ok := applied.(store.ImportVersionsResult) + require.True(t, ok, "got %T: %v", applied, applied) + require.Equal(t, []byte("cursor"), result.AckedCursor) + require.Equal(t, uint64(10), result.MaxImportedTS) + require.GreaterOrEqual(t, hlc.Current(), uint64(10)) + + staged := distribution.MigrationStagedDataKey(9, []byte("user|k")) + got, err := st.GetAt(ctx, staged, 10) + require.NoError(t, err) + require.Equal(t, []byte("v"), got) + _, err = st.GetAt(ctx, []byte("user|k"), 10) + require.ErrorIs(t, err, store.ErrKeyNotFound) +} From e7f69efd9e822849426047307f30c1c349050201 Mon Sep 17 00:00:00 2001 From: bootjp Date: Tue, 14 Jul 2026 00:02:32 +0900 Subject: [PATCH 05/13] migration: merge staged visibility reads --- kv/shard_store.go | 304 +++++++++++++++++++++++++++++++++++++---- kv/shard_store_test.go | 92 +++++++++++++ 2 files changed, 373 insertions(+), 23 deletions(-) diff --git a/kv/shard_store.go b/kv/shard_store.go index 175c81224..ba5ef6ac9 100644 --- a/kv/shard_store.go +++ b/kv/shard_store.go @@ -36,19 +36,19 @@ func NewShardStore(engine *distribution.Engine, groups map[uint64]*ShardGroup) * } func (s *ShardStore) GetAt(ctx context.Context, key []byte, ts uint64) ([]byte, error) { - g, ok := s.groupForKey(key) + route, g, ok := s.routeAndGroupForKey(key) if !ok || g.Store == nil { return nil, store.ErrKeyNotFound } // Some tests use ShardStore without raft; in that case serve reads locally. if engineForGroup(g) == nil { - return s.localGetAt(ctx, g, key, ts) + return s.localGetAt(ctx, g, route, key, ts) } // Wait for a leader read fence before serving from local state. if isLinearizableRaftLeader(ctx, engineForGroup(g)) { - return s.leaderGetAt(ctx, g, key, ts) + return s.leaderGetAt(ctx, g, route, key, ts) } return s.proxyRawGet(ctx, g, key, ts, 0) } @@ -63,10 +63,10 @@ func (s *ShardStore) GetGroupAt(ctx context.Context, groupID uint64, key []byte, } if engineForGroup(g) == nil { - return s.localGetAt(ctx, g, key, ts) + return s.localGetAt(ctx, g, distribution.Route{}, key, ts) } if isLinearizableRaftLeader(ctx, engineForGroup(g)) { - return s.leaderGetAt(ctx, g, key, ts) + return s.leaderGetAt(ctx, g, distribution.Route{}, key, ts) } return s.proxyRawGet(ctx, g, key, ts, groupID) } @@ -88,16 +88,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) @@ -105,6 +108,70 @@ 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) getAtWithStagedVisibility(ctx context.Context, g *ShardGroup, route distribution.Route, key []byte, ts uint64) ([]byte, error) { + 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) { + result, err := st.ExportVersions(ctx, store.ExportVersionsOptions{ + StartKey: key, + EndKey: nextScanCursor(key), + MaxCommitTSInclusive: ts, + MaxVersions: 1, + MaxScannedBytes: 0, + MinCommitTSExclusive: 0, + MaxBytes: 0, + KeyFamily: 0, + AcceptKey: nil, + }) + 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 + } + } + return store.MVCCVersion{}, false, nil +} + +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 (s *ShardStore) ExistsAt(ctx context.Context, key []byte, ts uint64) (bool, error) { v, err := s.GetAt(ctx, key, ts) if err != nil { @@ -390,7 +457,7 @@ func (s *ShardStore) scanRouteAtDirection( } 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) } @@ -398,7 +465,7 @@ func (s *ShardStore) scanRouteAtDirection( } if isLinearizableRaftLeader(ctx, engineForGroup(g)) { - return s.scanRouteAtLeader(ctx, g, start, end, limit, ts, reverse) + return s.scanRouteAtLeader(ctx, g, route, start, end, limit, ts, reverse) } var groupID uint64 @@ -435,6 +502,9 @@ func (s *ShardStore) scanRouteAtDirectionPhysicalLimit( } if engineForGroup(g) == nil { + if routeHasStagedVisibility(route) { + return nil, true, nil + } kvs, limitReached, err := scanLocalPhysicalLimit(ctx, g.Store, start, end, visibleLimit, physicalLimit, ts, reverse) if err != nil { return nil, limitReached, errors.WithStack(err) @@ -443,6 +513,9 @@ func (s *ShardStore) scanRouteAtDirectionPhysicalLimit( } if isLinearizableRaftLeader(ctx, engineForGroup(g)) { + if routeHasStagedVisibility(route) { + return nil, true, nil + } return s.scanRouteAtLeaderPhysicalLimit(ctx, g, start, end, visibleLimit, physicalLimit, ts, reverse) } @@ -494,12 +567,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) @@ -527,13 +604,14 @@ func (s *ShardStore) scanRouteAtLeaderPhysicalLimit( if err != nil { return nil, limitReached, err } - resolved, err := s.resolveScanLocks(ctx, g, kvs, lockKVs, ts) + resolved, err := s.resolveScanLocks(ctx, g, distribution.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, @@ -544,9 +622,12 @@ func (s *ShardStore) scanRouteAtLeader( 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 { @@ -557,7 +638,155 @@ func (s *ShardStore) scanRouteAtLeader( if err != nil { return nil, err } - return s.resolveScanLocks(ctx, g, kvs, lockKVs, ts) + return s.resolveScanLocks(ctx, g, route, kvs, lockKVs, ts) +} + +const stagedVisibilityExportPageSize = 1024 + +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) { + live, err := collectLatestLogicalVersions(ctx, g.Store, start, end, start, end, ts, liveLogicalVersionKey) + if err != nil { + return nil, err + } + stagedStart, stagedEnd := stagedVisibilityScanBounds(route.MigrationJobID, start, end) + staged, err := collectLatestLogicalVersions(ctx, g.Store, stagedStart, stagedEnd, start, end, ts, stagedLogicalVersionKey) + if err != nil { + return nil, err + } + merged := mergeLogicalVersionMaps(live, staged) + out := visibleLogicalKVs(merged, ts, reverse) + if len(out) > limit { + clear(out[limit:]) + out = out[:limit] + } + return out, nil +} + +func stagedVisibilityScanBounds(jobID uint64, start []byte, end []byte) ([]byte, []byte) { + prefix := distribution.MigrationStagedDataKeyPrefix(jobID) + scanStart := prefix + if start != nil { + scanStart = distribution.MigrationStagedDataKey(jobID, start) + } + scanEnd := prefixScanEnd(prefix) + if end != nil { + scanEnd = distribution.MigrationStagedDataKey(jobID, end) + } + return scanStart, scanEnd +} + +type logicalVersionKeyFunc func([]byte) ([]byte, bool) + +func liveLogicalVersionKey(key []byte) ([]byte, bool) { + if distribution.IsMigrationStagedDataKey(key) { + return nil, false + } + return bytes.Clone(key), true +} + +func stagedLogicalVersionKey(key []byte) ([]byte, bool) { + _, rawKey, ok := distribution.MigrationStagedDataKeyParts(key) + return rawKey, ok +} + +func collectLatestLogicalVersions( + ctx context.Context, + st store.MVCCStore, + scanStart []byte, + scanEnd []byte, + logicalStart []byte, + logicalEnd []byte, + ts uint64, + logicalKey logicalVersionKeyFunc, +) (map[string]store.MVCCVersion, error) { + out := make(map[string]store.MVCCVersion) + var cursor []byte + for { + result, err := st.ExportVersions(ctx, store.ExportVersionsOptions{ + StartKey: scanStart, + EndKey: scanEnd, + MaxCommitTSInclusive: ts, + Cursor: cursor, + MaxVersions: stagedVisibilityExportPageSize, + }) + if err != nil { + return nil, errors.WithStack(err) + } + for _, version := range result.Versions { + key, ok := logicalKey(version.Key) + if !ok || !keyInRange(key, logicalStart, logicalEnd) { + continue + } + version.Key = key + recordLatestLogicalVersion(out, version) + } + if result.Done { + return out, nil + } + if bytes.Equal(cursor, result.NextCursor) { + return nil, errors.New("staged visibility export cursor did not progress") + } + cursor = bytes.Clone(result.NextCursor) + } +} + +func recordLatestLogicalVersion(out map[string]store.MVCCVersion, version store.MVCCVersion) { + key := string(version.Key) + if existing, ok := out[key]; ok && existing.CommitTS >= version.CommitTS { + return + } + out[key] = version +} + +func mergeLogicalVersionMaps(live map[string]store.MVCCVersion, staged map[string]store.MVCCVersion) map[string]store.MVCCVersion { + out := make(map[string]store.MVCCVersion, len(live)+len(staged)) + for key, version := range live { + out[key] = version + } + for key, stagedVersion := range staged { + liveVersion, liveOK := out[key] + if winner, ok := newerMigrationVersion(liveVersion, liveOK, stagedVersion, true); ok { + out[key] = winner + } + } + return out +} + +func visibleLogicalKVs(versions map[string]store.MVCCVersion, ts uint64, reverse bool) []*store.KVPair { + out := make([]*store.KVPair, 0, len(versions)) + for _, version := range versions { + if !migrationVersionVisible(version, ts) { + continue + } + out = append(out, &store.KVPair{ + Key: bytes.Clone(version.Key), + Value: bytes.Clone(version.Value), + }) + } + sort.Slice(out, func(i, j int) bool { + cmp := bytes.Compare(out[i].Key, out[j].Key) + if reverse { + return cmp > 0 + } + return cmp < 0 + }) + return out +} + +func keyInRange(key []byte, start []byte, end []byte) bool { + if start != nil && bytes.Compare(key, start) < 0 { + return false + } + return end == nil || bytes.Compare(key, end) < 0 } func scanLockBoundsForKVs(kvs []*store.KVPair, scanStart []byte, scanEnd []byte, limit int) ([]byte, []byte) { @@ -708,13 +937,13 @@ func (s *ShardStore) ExpireAt(ctx context.Context, key []byte, expireAt uint64, } func (s *ShardStore) LatestCommitTS(ctx context.Context, key []byte) (uint64, bool, error) { - g, ok := s.groupForKey(key) + route, g, ok := s.routeAndGroupForKey(key) if !ok || g.Store == nil { return 0, false, nil } if engineForGroup(g) == nil { - ts, exists, err := g.Store.LatestCommitTS(ctx, key) + ts, exists, err := s.localLatestCommitTS(ctx, g, route, key) if err != nil { return 0, false, errors.WithStack(err) } @@ -726,7 +955,7 @@ func (s *ShardStore) LatestCommitTS(ctx context.Context, key []byte) (uint64, bo // 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) } @@ -737,6 +966,30 @@ func (s *ShardStore) LatestCommitTS(ctx context.Context, key []byte) (uint64, bo return s.proxyLatestCommitTS(ctx, g, 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 + } + 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) (uint64, bool, error) { engine := engineForGroup(g) if engine == nil { @@ -862,7 +1115,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 } @@ -877,7 +1130,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) { @@ -1134,7 +1387,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 { @@ -1144,7 +1397,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 @@ -1640,12 +1893,17 @@ func (s *ShardStore) closeGroup(g *ShardGroup) error { } func (s *ShardStore) groupForKey(key []byte) (*ShardGroup, bool) { + _, g, ok := s.routeAndGroupForKey(key) + return g, ok +} + +func (s *ShardStore) routeAndGroupForKey(key []byte) (distribution.Route, *ShardGroup, bool) { route, ok := s.engine.GetRoute(routeKey(key)) if !ok { - return nil, false + return distribution.Route{}, nil, false } g, ok := s.groups[route.GroupID] - return g, ok + return route, g, ok } func (s *ShardStore) proxyRawGet(ctx context.Context, g *ShardGroup, key []byte, ts uint64, groupID uint64) ([]byte, error) { diff --git a/kv/shard_store_test.go b/kv/shard_store_test.go index 4f0da38b7..100b70e27 100644 --- a/kv/shard_store_test.go +++ b/kv/shard_store_test.go @@ -10,6 +10,98 @@ import ( "github.com/stretchr/testify/require" ) +func newStagedVisibilityShardStore(t *testing.T) (*ShardStore, *ShardGroup) { + t.Helper() + + engine := distribution.NewEngine() + require.NoError(t, engine.ApplySnapshot(distribution.CatalogSnapshot{ + Version: 1, + Routes: []distribution.RouteDescriptor{ + { + RouteID: 1, + Start: []byte("a"), + End: []byte("z"), + GroupID: 1, + State: distribution.RouteStateActive, + StagedVisibilityActive: true, + MigrationJobID: 9, + MinWriteTSExclusive: 100, + }, + }, + })) + group := &ShardGroup{Store: store.NewMVCCStore()} + return NewShardStore(engine, map[uint64]*ShardGroup{1: group}), group +} + +func 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 TestShardStoreScanAndLatestCommitTS_MergeStagedVisibility(t *testing.T) { + t.Parallel() + + ctx := context.Background() + st, group := newStagedVisibilityShardStore(t) + + require.NoError(t, group.Store.PutAt(ctx, []byte("b"), []byte("live-b"), 10, 0)) + require.NoError(t, group.Store.PutAt(ctx, []byte("c"), []byte("live-c"), 30, 0)) + require.NoError(t, group.Store.PutAt(ctx, distribution.MigrationStagedDataKey(9, []byte("b")), []byte("staged-b"), 20, 0)) + require.NoError(t, group.Store.PutAt(ctx, distribution.MigrationStagedDataKey(9, []byte("d")), []byte("staged-d"), 15, 0)) + require.NoError(t, group.Store.PutAt(ctx, distribution.MigrationStagedDataKey(9, []byte("e")), []byte("staged-e"), 40, 0)) + require.NoError(t, group.Store.DeleteAt(ctx, []byte("d"), 25)) + + kvs, err := st.ScanAt(ctx, []byte("a"), []byte("z"), 10, 50) + require.NoError(t, err) + require.Equal(t, []*store.KVPair{ + {Key: []byte("b"), Value: []byte("staged-b")}, + {Key: []byte("c"), Value: []byte("live-c")}, + {Key: []byte("e"), Value: []byte("staged-e")}, + }, kvs) + + kvs, err = st.ReverseScanAt(ctx, []byte("a"), []byte("z"), 10, 50) + require.NoError(t, err) + require.Equal(t, []*store.KVPair{ + {Key: []byte("e"), Value: []byte("staged-e")}, + {Key: []byte("c"), Value: []byte("live-c")}, + {Key: []byte("b"), Value: []byte("staged-b")}, + }, kvs) + + ts, exists, err := st.LatestCommitTS(ctx, []byte("b")) + require.NoError(t, err) + require.True(t, exists) + require.Equal(t, uint64(20), ts) + + ts, exists, err = st.LatestCommitTS(ctx, []byte("d")) + require.NoError(t, err) + require.True(t, exists) + require.Equal(t, uint64(25), ts) + + ts, exists, err = st.LatestCommitTS(ctx, []byte("e")) + require.NoError(t, err) + require.True(t, exists) + require.Equal(t, uint64(40), ts) +} + func TestShardStoreScanAt_IncludesListKeysAcrossShards(t *testing.T) { t.Parallel() From ca1a05070c7d18bff6741fa11e72e6b622116ee7 Mon Sep 17 00:00:00 2001 From: bootjp Date: Tue, 14 Jul 2026 02:40:05 +0900 Subject: [PATCH 06/13] migration: harden staged route writes --- adapter/internal.go | 34 +++ adapter/internal_migration_test.go | 74 ++++++ internal/s3keys/keys.go | 11 + internal/s3keys/keys_test.go | 19 ++ kv/fsm.go | 129 ++++++++-- kv/fsm_migration_fence_test.go | 118 +++++++++ kv/route_history.go | 8 + kv/shard_store.go | 300 +++++++++++++++------- kv/shard_store_test.go | 85 ++++++ kv/sharded_coordinator.go | 63 +++++ kv/sharded_coordinator_del_prefix_test.go | 50 ++++ kv/sharded_coordinator_txn_test.go | 20 ++ 12 files changed, 807 insertions(+), 104 deletions(-) diff --git a/adapter/internal.go b/adapter/internal.go index de61f5485..f72d116f2 100644 --- a/adapter/internal.go +++ b/adapter/internal.go @@ -6,6 +6,7 @@ import ( "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" @@ -242,6 +243,9 @@ func exportRangeVersionsOptions(req *pb.ExportRangeVersionsRequest) store.Export func migrationExportFilter(req *pb.ExportRangeVersionsRequest) func([]byte) bool { routeFilter := kv.RouteKeyFilter(req.GetRouteStart(), req.GetRouteEnd()) + if migrationFamilyRequiresDecodedS3(req.GetKeyFamily()) { + routeFilter = decodedS3BucketRouteFilter(req.GetKeyFamily(), req.GetRouteStart(), req.GetRouteEnd()) + } excludeKnownInternal := req.GetExcludeKnownInternal() || req.GetKeyFamily() == distribution.MigrationFamilyUser bracket := distribution.MigrationBracket{ Family: req.GetKeyFamily(), @@ -255,6 +259,36 @@ func migrationExportFilter(req *pb.ExportRangeVersionsRequest) func([]byte) bool } } +func migrationFamilyRequiresDecodedS3(family uint32) bool { + return family == distribution.MigrationFamilyS3BucketMeta || + family == distribution.MigrationFamilyS3BucketGeneration +} + +func decodedS3BucketRouteFilter(family uint32, routeStart, routeEnd []byte) func([]byte) bool { + return func(rawKey []byte) bool { + bucket, ok := decodedS3BucketName(family, rawKey) + if !ok { + return false + } + routeKey := s3keys.RouteKey(bucket, 0, "") + if routeStart != nil && bytes.Compare(routeKey, routeStart) < 0 { + return false + } + return routeEnd == nil || bytes.Compare(routeKey, routeEnd) < 0 + } +} + +func decodedS3BucketName(family uint32, rawKey []byte) (string, bool) { + switch family { + case distribution.MigrationFamilyS3BucketMeta: + return s3keys.ParseBucketMetaKey(rawKey) + case distribution.MigrationFamilyS3BucketGeneration: + return s3keys.ParseBucketGenerationKey(rawKey) + default: + return "", false + } +} + func cloneByteSlices(in [][]byte) [][]byte { if len(in) == 0 { return nil diff --git a/adapter/internal_migration_test.go b/adapter/internal_migration_test.go index dd8b48baf..b43a8f4b9 100644 --- a/adapter/internal_migration_test.go +++ b/adapter/internal_migration_test.go @@ -6,6 +6,7 @@ import ( "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" @@ -70,6 +71,17 @@ func (s *captureExportRangeVersionsStream) Send(resp *pb.ExportRangeVersionsResp return nil } +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() @@ -114,6 +126,68 @@ func TestInternalExportRangeVersionsRejectsUnboundedExport(t *testing.T) { 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 TestInternalImportRangeVersionsAppliesStoreBatch(t *testing.T) { t.Parallel() diff --git a/internal/s3keys/keys.go b/internal/s3keys/keys.go index 682105219..21326fb5e 100644 --- a/internal/s3keys/keys.go +++ b/internal/s3keys/keys.go @@ -80,6 +80,17 @@ func ParseBucketMetaKey(key []byte) (string, bool) { return string(segment), true } +func ParseBucketGenerationKey(key []byte) (string, bool) { + if !bytes.HasPrefix(key, bucketGenerationPrefixBytes) { + return "", false + } + segment, next, ok := decodeSegment(key, len(bucketGenerationPrefixBytes)) + if !ok || next != len(key) { + return "", false + } + return string(segment), true +} + func ObjectManifestKey(bucket string, generation uint64, object string) []byte { return buildObjectKey(objectManifestPrefixBytes, bucket, generation, object, "", 0, 0) } diff --git a/internal/s3keys/keys_test.go b/internal/s3keys/keys_test.go index e5e5fca5a..34449227a 100644 --- a/internal/s3keys/keys_test.go +++ b/internal/s3keys/keys_test.go @@ -18,6 +18,25 @@ func TestBucketMetaKey_RoundTripsZeroByteSegments(t *testing.T) { require.Equal(t, bucket, parsed) } +func TestBucketGenerationKey_RoundTripsZeroByteSegments(t *testing.T) { + t.Parallel() + + bucket := string([]byte{'b', 'u', 0x00, 'c', 'k', 'e', 't'}) + key := BucketGenerationKey(bucket) + + parsed, ok := ParseBucketGenerationKey(key) + require.True(t, ok) + 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 00e4a5d3c..1693d5eaa 100644 --- a/kv/fsm.go +++ b/kv/fsm.go @@ -9,6 +9,7 @@ import ( "log/slog" "os" + "github.com/bootjp/elastickv/distribution" "github.com/bootjp/elastickv/internal/encryption/fsmwire" "github.com/bootjp/elastickv/internal/raftengine" "github.com/bootjp/elastickv/internal/s3keys" @@ -121,6 +122,10 @@ type RouteSnapshot interface { // OwnerOf returns the Raft group ID that owned key at this // snapshot's version. (0, false) when no route covered key. OwnerOf(key []byte) (uint64, bool) + // RouteOf returns the complete route descriptor covering key. + RouteOf(key []byte) (distribution.Route, bool) + // IntersectingRoutes returns every route intersecting [start, end). + IntersectingRoutes(start, end []byte) []distribution.Route // WriteFencedForKey reports whether key is currently inside a // WriteFenced route in this snapshot. WriteFencedForKey(key []byte) bool @@ -273,6 +278,8 @@ var ErrUnknownRequestType = errors.New("unknown request type") // catches up to the promoted owner. var ErrRouteWriteFenced = errors.New("route is write-fenced; retry after route migration") +var ErrRouteWriteTimestampTooLow = errors.New("route write timestamp is below migration floor") + // ErrComposed1Violation is returned by verifyComposed1 when the // transaction's commit cannot proceed on this Raft group because the // txn's read-set or write-set keys are not owned by this group at @@ -494,20 +501,8 @@ func (f *kvFSM) handleRawRequest(ctx context.Context, r *pb.Request, commitTS ui return f.handleDelPrefix(ctx, prefix, commitTS) } - for _, mut := range r.Mutations { - if mut == nil || len(mut.Key) == 0 { - return errors.WithStack(ErrInvalidRequest) - } - // Raw requests should not mutate txn-internal keys. - if isTxnInternalKey(mut.Key) { - return errors.WithStack(ErrInvalidRequest) - } - if err := f.verifyRouteNotFencedForKey(mut.Key); err != nil { - return err - } - if err := f.assertNoConflictingTxnLock(ctx, mut.Key, nil, 0); err != nil { - return err - } + if err := f.validateRawMutationsForApply(ctx, r.Mutations, commitTS); err != nil { + return err } muts, err := toStoreMutations(r.Mutations) @@ -523,6 +518,35 @@ func (f *kvFSM) handleRawRequest(ctx context.Context, r *pb.Request, commitTS ui return nil } +func (f *kvFSM) validateRawMutationsForApply(ctx context.Context, muts []*pb.Mutation, commitTS uint64) error { + for _, mut := range muts { + if err := f.validateRawMutationForApply(ctx, mut, commitTS); err != nil { + return err + } + } + return nil +} + +func (f *kvFSM) validateRawMutationForApply(ctx context.Context, mut *pb.Mutation, 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 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 +564,9 @@ 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 { + return err + } if err := f.store.DeletePrefixAtRaftAt(ctx, prefix, txnCommonPrefix, commitTS, f.pendingApplyIdx); err != nil { return errors.WithStack(err) } @@ -547,6 +574,51 @@ func (f *kvFSM) handleDelPrefix(ctx context.Context, prefix []byte, commitTS uin return nil } +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 + } + rkey := routeKey(key) + route, ok := snap.RouteOf(rkey) + if !ok || route.MinWriteTSExclusive == 0 || commitTS > route.MinWriteTSExclusive { + return nil + } + return errors.Wrapf(ErrRouteWriteTimestampTooLow, "key %q routeKey %q commit_ts=%d floor=%d", key, rkey, commitTS, route.MinWriteTSExclusive) +} + +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 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 +} + +func (f *kvFSM) verifyRouteWriteTimestampFloorForMutations(muts []*pb.Mutation, commitTS uint64) error { + for _, mut := range muts { + if mut == nil || len(mut.Key) == 0 || isTxnInternalKey(mut.Key) { + continue + } + if err := f.verifyRouteWriteTimestampFloorForKey(mut.Key, commitTS); err != nil { + return err + } + } + return nil +} + func (f *kvFSM) verifyRouteNotFencedForMutations(muts []*pb.Mutation) error { for _, mut := range muts { if mut == nil || len(mut.Key) == 0 || isTxnInternalKey(mut.Key) { @@ -964,6 +1036,9 @@ func (f *kvFSM) handlePrepareRequest(ctx context.Context, r *pb.Request) error { if err := f.verifyRouteNotFencedForMutations(uniq); err != nil { return err } + if err := f.verifyRouteWriteTimestampFloorForMutations(uniq, meta.CommitTS); err != nil { + return err + } if err := f.validateConflicts(ctx, uniq, startTS); err != nil { return errors.WithStack(err) } @@ -1030,7 +1105,7 @@ func (f *kvFSM) handleOnePhaseTxnRequest(ctx context.Context, r *pb.Request, com return nil } - uniq, err := f.uniqueMutationsNotFenced(muts) + uniq, err := f.uniqueMutationsNotFencedAboveWriteFloor(muts, commitTS) if err != nil { return err } @@ -1057,6 +1132,28 @@ func (f *kvFSM) uniqueMutationsNotFenced(muts []*pb.Mutation) ([]*pb.Mutation, e return uniq, nil } +func (f *kvFSM) uniqueMutationsNotFencedAboveWriteFloor(muts []*pb.Mutation, commitTS uint64) ([]*pb.Mutation, error) { + uniq, err := f.uniqueMutationsNotFenced(muts) + if err != nil { + return nil, err + } + if err := f.verifyRouteWriteTimestampFloorForMutations(uniq, commitTS); err != nil { + return nil, err + } + return uniq, nil +} + +func (f *kvFSM) uniqueMutationsAboveWriteFloor(muts []*pb.Mutation, commitTS uint64) ([]*pb.Mutation, error) { + uniq, err := uniqueMutations(muts) + if err != nil { + return nil, err + } + if err := f.verifyRouteWriteTimestampFloorForMutations(uniq, commitTS); err != nil { + return nil, err + } + return uniq, nil +} + // dedupProbeOnePhase decides whether handleOnePhaseTxnRequest should no-op // because the entry is a retry whose prior attempt already landed. Extracted // to keep handleOnePhaseTxnRequest under the cyclop budget; the determinism @@ -1096,7 +1193,7 @@ func (f *kvFSM) handleCommitRequest(ctx context.Context, r *pb.Request) error { if err != nil { return err } - uniq, err := uniqueMutations(muts) + uniq, err := f.uniqueMutationsAboveWriteFloor(muts, commitTS) if err != nil { return err } diff --git a/kv/fsm_migration_fence_test.go b/kv/fsm_migration_fence_test.go index e47e71f95..045a9d04d 100644 --- a/kv/fsm_migration_fence_test.go +++ b/kv/fsm_migration_fence_test.go @@ -22,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 TestFSMRejectsRawPointWriteOnWriteFencedRoute(t *testing.T) { t.Parallel() @@ -112,3 +122,111 @@ func TestFSMRejectsPrepareOnWriteFencedRouteButAllowsAbort(t *testing.T) { err := fsm.handleTxnRequest(ctx, abort, 11) require.False(t, errors.Is(err, ErrRouteWriteFenced), "ABORT must keep the narrow cleanup lane open") } + +func TestFSMRejectsRawPointWriteAtMigrationTimestampFloor(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("low")}}, + }, 100) + require.ErrorIs(t, err, ErrRouteWriteTimestampTooLow) + + _, getErr := fsm.store.GetAt(ctx, []byte("z"), ^uint64(0)) + require.ErrorIs(t, getErr, store.ErrKeyNotFound) + + err = fsm.handleRawRequest(ctx, &pb.Request{ + Mutations: []*pb.Mutation{{Op: pb.Op_PUT, Key: []byte("z"), Value: []byte("ok")}}, + }, 101) + require.NoError(t, err) + got, getErr := fsm.store.GetAt(ctx, []byte("z"), ^uint64(0)) + require.NoError(t, getErr) + require.Equal(t, []byte("ok"), got) +} + +func TestFSMRejectsDelPrefixAtMigrationTimestampFloor(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 TestFSMRejectsOnePhaseTxnAtMigrationTimestampFloor(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) + + req.Mutations[0].Value = EncodeTxnMeta(TxnMeta{PrimaryKey: []byte("z"), CommitTS: 101}) + req.Mutations[1].Value = []byte("ok") + err = fsm.handleTxnRequest(ctx, req, 101) + require.NoError(t, err) + got, getErr := fsm.store.GetAt(ctx, []byte("z"), ^uint64(0)) + require.NoError(t, getErr) + require.Equal(t, []byte("ok"), got) +} + +func TestFSMRejectsCommitButNotPrepareAtMigrationTimestampFloor(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.NoError(t, fsm.handleTxnRequest(ctx, prepare, 90)) + + commit := &pb.Request{ + IsTxn: true, + Phase: pb.Phase_COMMIT, + 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, commit, 100) + require.ErrorIs(t, err, ErrRouteWriteTimestampTooLow) + + _, getErr := fsm.store.GetAt(ctx, []byte("z"), ^uint64(0)) + require.ErrorIs(t, getErr, store.ErrKeyNotFound) + + commit.Mutations[0].Value = EncodeTxnMeta(TxnMeta{PrimaryKey: []byte("z"), CommitTS: 101}) + commit.Mutations[1].Value = []byte("ignored") + err = fsm.handleTxnRequest(ctx, commit, 101) + require.NoError(t, err) + got, getErr := fsm.store.GetAt(ctx, []byte("z"), ^uint64(0)) + require.NoError(t, getErr) + require.Equal(t, []byte("v"), got) +} 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_store.go b/kv/shard_store.go index ba5ef6ac9..fcf47f15c 100644 --- a/kv/shard_store.go +++ b/kv/shard_store.go @@ -113,6 +113,9 @@ func routeHasStagedVisibility(route distribution.Route) bool { } 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 @@ -135,14 +138,16 @@ func (s *ShardStore) getAtWithStagedVisibility(ctx context.Context, g *ShardGrou func latestMVCCVersionAt(ctx context.Context, st store.MVCCStore, key []byte, ts uint64) (store.MVCCVersion, bool, error) { result, err := st.ExportVersions(ctx, store.ExportVersionsOptions{ StartKey: key, - EndKey: nextScanCursor(key), + EndKey: prefixScanEnd(key), MaxCommitTSInclusive: ts, MaxVersions: 1, MaxScannedBytes: 0, MinCommitTSExclusive: 0, MaxBytes: 0, KeyFamily: 0, - AcceptKey: nil, + AcceptKey: func(rawKey []byte) bool { + return bytes.Equal(rawKey, key) + }, }) if err != nil { return store.MVCCVersion{}, false, errors.WithStack(err) @@ -172,6 +177,18 @@ 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 { @@ -276,6 +293,9 @@ func (s *ShardStore) ScanAtPhysicalLimit(ctx context.Context, start []byte, end return []*store.KVPair{}, false, nil } routes, clampToRoutes := s.routesForScan(start, end) + if routesContainStagedVisibility(routes) { + return nil, true, nil + } if len(routes) != 1 || clampToRoutes { kvs, err := s.ScanAt(ctx, start, end, visibleLimit, ts) return kvs, false, err @@ -316,6 +336,9 @@ func (s *ShardStore) ReverseScanAtPhysicalLimit(ctx context.Context, start []byt return []*store.KVPair{}, false, nil } routes, clampToRoutes := s.routesForScan(start, end) + if routesContainStagedVisibility(routes) { + return nil, true, nil + } if len(routes) != 1 || clampToRoutes { kvs, err := s.ReverseScanAt(ctx, start, end, visibleLimit, ts) return kvs, false, err @@ -348,6 +371,15 @@ func (s *ShardStore) routesForScan(start []byte, end []byte) ([]distribution.Rou return routes, true } +func routesContainStagedVisibility(routes []distribution.Route) bool { + for _, route := range routes { + if routeHasStagedVisibility(route) { + return true + } + } + return false +} + func (s *ShardStore) scanRoutesAt(ctx context.Context, routes []distribution.Route, start []byte, end []byte, limit int, ts uint64, clampToRoutes bool) ([]*store.KVPair, error) { out := make([]*store.KVPair, 0) for _, route := range routes { @@ -641,7 +673,10 @@ func (s *ShardStore) scanRouteAtLeader( return s.resolveScanLocks(ctx, g, route, kvs, lockKVs, ts) } -const stagedVisibilityExportPageSize = 1024 +const ( + stagedVisibilityMaxCandidateWindow = 8192 + stagedVisibilityWindowGrowthFactor = 2 +) func (s *ShardStore) scanRouteWithStagedVisibility( ctx context.Context, @@ -653,112 +688,139 @@ func (s *ShardStore) scanRouteWithStagedVisibility( ts uint64, reverse bool, ) ([]*store.KVPair, error) { - live, err := collectLatestLogicalVersions(ctx, g.Store, start, end, start, end, ts, liveLogicalVersionKey) - if err != nil { + if err := ensureReadTSRetained(g.Store, ts); err != nil { return nil, err } stagedStart, stagedEnd := stagedVisibilityScanBounds(route.MigrationJobID, start, end) - staged, err := collectLatestLogicalVersions(ctx, g.Store, stagedStart, stagedEnd, start, end, ts, stagedLogicalVersionKey) - if err != nil { - return nil, err - } - merged := mergeLogicalVersionMaps(live, staged) - out := visibleLogicalKVs(merged, ts, reverse) - if len(out) > limit { - clear(out[limit:]) - out = out[:limit] + window := stagedVisibilityCandidateWindow(limit) + for { + liveKVs, err := scanVisibleCandidates(ctx, g.Store, start, end, window, ts, reverse) + if err != nil { + return nil, err + } + stagedKVs, err := scanVisibleCandidates(ctx, g.Store, stagedStart, stagedEnd, window, ts, reverse) + if err != nil { + return nil, err + } + versions, err := s.latestStagedVisibilityCandidates(ctx, g.Store, route, liveKVs, stagedKVs, ts) + if err != nil { + return nil, err + } + out := visibleLogicalKVs(versions, ts, reverse) + if len(out) >= limit { + clear(out[limit:]) + return out[:limit], nil + } + if len(liveKVs) < window && len(stagedKVs) < window { + return out, nil + } + nextWindow := nextStagedVisibilityCandidateWindow(window) + if nextWindow == window { + return out, nil + } + window = nextWindow } - return out, nil } -func stagedVisibilityScanBounds(jobID uint64, start []byte, end []byte) ([]byte, []byte) { - prefix := distribution.MigrationStagedDataKeyPrefix(jobID) - scanStart := prefix - if start != nil { - scanStart = distribution.MigrationStagedDataKey(jobID, start) +func stagedVisibilityCandidateWindow(limit int) int { + if limit <= 0 { + return 0 } - scanEnd := prefixScanEnd(prefix) - if end != nil { - scanEnd = distribution.MigrationStagedDataKey(jobID, end) + if limit > stagedVisibilityMaxCandidateWindow { + return stagedVisibilityMaxCandidateWindow } - return scanStart, scanEnd + return limit } -type logicalVersionKeyFunc func([]byte) ([]byte, bool) - -func liveLogicalVersionKey(key []byte) ([]byte, bool) { - if distribution.IsMigrationStagedDataKey(key) { - return nil, false +func nextStagedVisibilityCandidateWindow(window int) int { + if window >= stagedVisibilityMaxCandidateWindow { + return window + } + next := window * stagedVisibilityWindowGrowthFactor + if next < window || next > stagedVisibilityMaxCandidateWindow { + return stagedVisibilityMaxCandidateWindow } - return bytes.Clone(key), true + return next } -func stagedLogicalVersionKey(key []byte) ([]byte, bool) { - _, rawKey, ok := distribution.MigrationStagedDataKeyParts(key) - return rawKey, ok +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 collectLatestLogicalVersions( +func (s *ShardStore) latestStagedVisibilityCandidates( ctx context.Context, st store.MVCCStore, - scanStart []byte, - scanEnd []byte, - logicalStart []byte, - logicalEnd []byte, + route distribution.Route, + liveKVs []*store.KVPair, + stagedKVs []*store.KVPair, ts uint64, - logicalKey logicalVersionKeyFunc, ) (map[string]store.MVCCVersion, error) { - out := make(map[string]store.MVCCVersion) - var cursor []byte - for { - result, err := st.ExportVersions(ctx, store.ExportVersionsOptions{ - StartKey: scanStart, - EndKey: scanEnd, - MaxCommitTSInclusive: ts, - Cursor: cursor, - MaxVersions: stagedVisibilityExportPageSize, - }) + keys := stagedVisibilityCandidateKeys(liveKVs, stagedKVs) + out := make(map[string]store.MVCCVersion, len(keys)) + for _, key := range keys { + live, liveOK, err := latestMVCCVersionAt(ctx, st, key, ts) if err != nil { - return nil, errors.WithStack(err) + return nil, err } - for _, version := range result.Versions { - key, ok := logicalKey(version.Key) - if !ok || !keyInRange(key, logicalStart, logicalEnd) { - continue - } - version.Key = key - recordLatestLogicalVersion(out, version) + stagedKey := distribution.MigrationStagedDataKey(route.MigrationJobID, key) + staged, stagedOK, err := latestMVCCVersionAt(ctx, st, stagedKey, ts) + if err != nil { + return nil, err } - if result.Done { - return out, nil + if stagedOK { + staged.Key = bytes.Clone(key) } - if bytes.Equal(cursor, result.NextCursor) { - return nil, errors.New("staged visibility export cursor did not progress") + if winner, ok := newerMigrationVersion(live, liveOK, staged, stagedOK); ok { + out[string(key)] = winner } - cursor = bytes.Clone(result.NextCursor) } + return out, nil } -func recordLatestLogicalVersion(out map[string]store.MVCCVersion, version store.MVCCVersion) { - key := string(version.Key) - if existing, ok := out[key]; ok && existing.CommitTS >= version.CommitTS { - return +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 + } + 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) + } + out := make([][]byte, 0, len(seen)) + for _, key := range seen { + out = append(out, key) } - out[key] = version + return out } -func mergeLogicalVersionMaps(live map[string]store.MVCCVersion, staged map[string]store.MVCCVersion) map[string]store.MVCCVersion { - out := make(map[string]store.MVCCVersion, len(live)+len(staged)) - for key, version := range live { - out[key] = version +func stagedVisibilityScanBounds(jobID uint64, start []byte, end []byte) ([]byte, []byte) { + prefix := distribution.MigrationStagedDataKeyPrefix(jobID) + scanStart := prefix + if start != nil { + scanStart = distribution.MigrationStagedDataKey(jobID, start) } - for key, stagedVersion := range staged { - liveVersion, liveOK := out[key] - if winner, ok := newerMigrationVersion(liveVersion, liveOK, stagedVersion, true); ok { - out[key] = winner - } + scanEnd := prefixScanEnd(prefix) + if end != nil { + scanEnd = distribution.MigrationStagedDataKey(jobID, end) } - return out + return scanStart, scanEnd } func visibleLogicalKVs(versions map[string]store.MVCCVersion, ts uint64, reverse bool) []*store.KVPair { @@ -782,13 +844,6 @@ func visibleLogicalKVs(versions map[string]store.MVCCVersion, ts uint64, reverse return out } -func keyInRange(key []byte, start []byte, end []byte) bool { - if start != nil && bytes.Compare(key, start) < 0 { - return false - } - return end == nil || bytes.Compare(key, end) < 0 -} - func scanLockBoundsForKVs(kvs []*store.KVPair, scanStart []byte, scanEnd []byte, limit int) ([]byte, []byte) { if countNonInternalKVs(kvs) < limit { return scanStart, scanEnd @@ -905,34 +960,46 @@ 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 + } 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 + } 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 + } 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 + } return errors.WithStack(g.Store.ExpireAt(ctx, key, expireAt, commitTS)) } @@ -1606,6 +1673,9 @@ 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 + } return errors.WithStack(group.Store.ApplyMutations(ctx, mutations, readKeys, startTS, commitTS)) } @@ -1616,6 +1686,9 @@ 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 + } return errors.WithStack(group.Store.ApplyMutationsRaft(ctx, mutations, readKeys, startTS, commitTS)) } @@ -1627,9 +1700,38 @@ 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 + } 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 { + continue + } + route, _, ok := s.routeAndGroupForKey(mut.Key) + if !ok { + return store.ErrNotSupported + } + if err := ensureRouteWriteTimestampFloor(route, mut.Key, commitTS); err != nil { + return err + } + } + return nil +} + // resolveSingleShardGroup returns the shard group that owns every // mutation in the batch, or an error if the batch is cross-shard or // references an unknown group. A nil group with nil error means "empty @@ -1656,6 +1758,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 @@ -1667,8 +1772,24 @@ func (s *ShardStore) DeletePrefixAt(ctx context.Context, prefix []byte, excludeP return nil } +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 @@ -1695,6 +1816,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 diff --git a/kv/shard_store_test.go b/kv/shard_store_test.go index 100b70e27..e3a3ea8d5 100644 --- a/kv/shard_store_test.go +++ b/kv/shard_store_test.go @@ -33,6 +33,32 @@ func newStagedVisibilityShardStore(t *testing.T) (*ShardStore, *ShardGroup) { 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() @@ -57,6 +83,37 @@ func TestShardStoreGetAt_MergesStagedVisibility(t *testing.T) { 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 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() @@ -102,6 +159,34 @@ func TestShardStoreScanAndLatestCommitTS_MergeStagedVisibility(t *testing.T) { require.Equal(t, uint64(40), ts) } +func TestShardStorePhysicalLimitFailsClosedBeforeStagedVisibilityFallback(t *testing.T) { + t.Parallel() + + ctx := context.Background() + st, _ := newStagedVisibilityShardStore(t) + + kvs, limitReached, err := st.ScanAtPhysicalLimit(ctx, []byte("b"), []byte("c"), 10, 10, 50) + require.NoError(t, err) + require.True(t, limitReached) + require.Nil(t, kvs) + + kvs, limitReached, err = st.ReverseScanAtPhysicalLimit(ctx, []byte("b"), []byte("c"), 10, 10, 50) + require.NoError(t, err) + require.True(t, limitReached) + require.Nil(t, kvs) +} + +func TestShardStoreRejectsWritesAtMigrationTimestampFloor(t *testing.T) { + t.Parallel() + + ctx := context.Background() + st, _ := newStagedVisibilityShardStore(t) + + err := st.PutAt(ctx, []byte("k"), []byte("low"), 100, 0) + require.ErrorIs(t, err, ErrRouteWriteTimestampTooLow) + require.NoError(t, st.PutAt(ctx, []byte("k"), []byte("ok"), 101, 0)) +} + func TestShardStoreScanAt_IncludesListKeysAcrossShards(t *testing.T) { t.Parallel() diff --git a/kv/sharded_coordinator.go b/kv/sharded_coordinator.go index 7ed2ccb7c..a071a7e86 100644 --- a/kv/sharded_coordinator.go +++ b/kv/sharded_coordinator.go @@ -1036,6 +1036,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{ @@ -1084,6 +1087,42 @@ 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 + } + rkey := routeKey(elem.Key) + route, ok := c.engine.GetRoute(rkey) + if !ok || route.MinWriteTSExclusive == 0 || commitTS > route.MinWriteTSExclusive { + continue + } + return errors.Wrapf(ErrRouteWriteTimestampTooLow, "key %q routeKey %q commit_ts=%d floor=%d", elem.Key, rkey, commitTS, route.MinWriteTSExclusive) + } + return nil +} + +func (c *ShardedCoordinator) rejectWriteTimestampFloorDelPrefixes(elems []*Elem[OP], commitTS uint64) error { + if c == nil || c.engine == nil || commitTS == 0 { + return nil + } + for _, elem := range elems { + if elem == nil { + continue + } + start, end := routePrefixRange(elem.Key) + for _, route := range c.engine.GetIntersectingRoutes(start, end) { + if route.MinWriteTSExclusive != 0 && commitTS <= route.MinWriteTSExclusive { + return errors.Wrapf(ErrRouteWriteTimestampTooLow, "prefix %q route range [%q,%q) commit_ts=%d floor=%d", elem.Key, start, end, commitTS, route.MinWriteTSExclusive) + } + } + } + return nil +} + // broadcastToAllGroups sends the same set of requests to every shard group in // parallel and returns the maximum commit index. func (c *ShardedCoordinator) broadcastToAllGroups(ctx context.Context, requests []*pb.Request) (*CoordinateResponse, error) { @@ -1141,6 +1180,9 @@ func (c *ShardedCoordinator) dispatchTxn(ctx context.Context, startTS uint64, co if err != nil { return nil, err } + if err := c.rejectWriteTimestampFloorPointElems(elems, commitTS); err != nil { + return nil, err + } if len(gids) == 1 && c.allReadKeysInShard(readKeys, gids[0]) { // Fast path: all mutations and read keys are in a single shard. @@ -1984,6 +2026,9 @@ func (c *ShardedCoordinator) rawLogs(ctx context.Context, reqs *OperationGroup[O if err != nil { return nil, err } + if err := c.rejectWriteTimestampFloorMutations(grouped[gid], ts); err != nil { + return nil, err + } logs = append(logs, &pb.Request{ IsTxn: false, Phase: pb.Phase_NONE, @@ -1994,6 +2039,24 @@ func (c *ShardedCoordinator) rawLogs(ctx context.Context, reqs *OperationGroup[O return logs, nil } +func (c *ShardedCoordinator) rejectWriteTimestampFloorMutations(muts []*pb.Mutation, commitTS uint64) error { + if c == nil || c.engine == nil || commitTS == 0 { + return nil + } + for _, mut := range muts { + if mut == nil || len(mut.Key) == 0 { + continue + } + rkey := routeKey(mut.Key) + route, ok := c.engine.GetRoute(rkey) + if !ok || route.MinWriteTSExclusive == 0 || commitTS > route.MinWriteTSExclusive { + continue + } + return errors.Wrapf(ErrRouteWriteTimestampTooLow, "key %q routeKey %q commit_ts=%d floor=%d", mut.Key, rkey, commitTS, route.MinWriteTSExclusive) + } + return nil +} + func (c *ShardedCoordinator) rawLogTimestamp(ctx context.Context) (uint64, error) { if c.tsAllocator != nil { return 0, nil diff --git a/kv/sharded_coordinator_del_prefix_test.go b/kv/sharded_coordinator_del_prefix_test.go index 4b83dc131..a83a5ce95 100644 --- a/kv/sharded_coordinator_del_prefix_test.go +++ b/kv/sharded_coordinator_del_prefix_test.go @@ -121,6 +121,56 @@ func TestShardedCoordinator_DelPrefixBroadcastsToAllGroups(t *testing.T) { "same DEL_PREFIX element must use the same timestamp across shards") } +func newMigrationFloorEngine(t *testing.T, floor uint64) *distribution.Engine { + t.Helper() + + engine := distribution.NewEngine() + require.NoError(t, engine.ApplySnapshot(distribution.CatalogSnapshot{ + Version: 1, + Routes: []distribution.RouteDescriptor{ + { + RouteID: 1, + Start: []byte(""), + End: nil, + GroupID: 1, + State: distribution.RouteStateActive, + MinWriteTSExclusive: floor, + }, + }, + })) + return engine +} + +func TestShardedCoordinatorRejectsPointWriteAtMigrationTimestampFloor(t *testing.T) { + t.Parallel() + + g1Txn := &recordingTransactional{} + coord := NewShardedCoordinator(newMigrationFloorEngine(t, ^uint64(0)), map[uint64]*ShardGroup{ + 1: {Txn: g1Txn}, + }, 1, NewHLC(), nil) + + _, err := coord.Dispatch(context.Background(), &OperationGroup[OP]{ + Elems: []*Elem[OP]{{Op: Put, Key: []byte("z"), Value: []byte("v")}}, + }) + require.ErrorIs(t, err, ErrRouteWriteTimestampTooLow) + require.Empty(t, g1Txn.requests, "coordinator must reject before proposing a floor-violating point write") +} + +func TestShardedCoordinatorRejectsDelPrefixAtMigrationTimestampFloor(t *testing.T) { + t.Parallel() + + g1Txn := &recordingTransactional{} + coord := NewShardedCoordinator(newMigrationFloorEngine(t, ^uint64(0)), map[uint64]*ShardGroup{ + 1: {Txn: g1Txn}, + }, 1, NewHLC(), nil) + + _, err := coord.Dispatch(context.Background(), &OperationGroup[OP]{ + Elems: []*Elem[OP]{{Op: DelPrefix, Key: []byte("z")}}, + }) + require.ErrorIs(t, err, ErrRouteWriteTimestampTooLow) + require.Empty(t, g1Txn.requests, "coordinator must reject before broadcasting a floor-violating prefix delete") +} + func TestShardedCoordinatorRejectsPointWriteOnWriteFencedRoute(t *testing.T) { t.Parallel() diff --git a/kv/sharded_coordinator_txn_test.go b/kv/sharded_coordinator_txn_test.go index 964f2713f..c07733d16 100644 --- a/kv/sharded_coordinator_txn_test.go +++ b/kv/sharded_coordinator_txn_test.go @@ -246,6 +246,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() From d33bf3589a6a1f84ad05fa57073c8fa69dee4ed2 Mon Sep 17 00:00:00 2001 From: bootjp Date: Tue, 14 Jul 2026 05:26:57 +0900 Subject: [PATCH 07/13] Harden staged migration visibility --- adapter/internal.go | 14 +- adapter/internal_migration_test.go | 37 ++++ kv/fsm_migration_import.go | 2 +- kv/leader_routed_store.go | 8 + kv/shard_store.go | 226 +++++++++++++++++++++- kv/shard_store_test.go | 61 ++++++ kv/sharded_coordinator.go | 81 +++++++- kv/sharded_coordinator_txn_test.go | 69 +++++++ kv/tso_test.go | 15 ++ store/lsm_migration.go | 20 +- store/lsm_store_registration_gate_test.go | 38 ++++ store/migration_versions.go | 4 + store/store.go | 3 + 13 files changed, 558 insertions(+), 20 deletions(-) diff --git a/adapter/internal.go b/adapter/internal.go index f72d116f2..a14fb9134 100644 --- a/adapter/internal.go +++ b/adapter/internal.go @@ -133,9 +133,19 @@ 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")) + } return nil } +func exportRangeVersionsRequestFullyUnbounded(req *pb.ExportRangeVersionsRequest) bool { + return len(req.GetRangeStart()) == 0 && + len(req.GetRangeEnd()) == 0 && + len(req.GetRouteStart()) == 0 && + len(req.GetRouteEnd()) == 0 +} + func (i *Internal) streamExportRangeVersions(req *pb.ExportRangeVersionsRequest, stream pb.Internal_ExportRangeVersionsServer) error { opts := exportRangeVersionsOptions(req) for { @@ -271,10 +281,10 @@ func decodedS3BucketRouteFilter(family uint32, routeStart, routeEnd []byte) func return false } routeKey := s3keys.RouteKey(bucket, 0, "") - if routeStart != nil && bytes.Compare(routeKey, routeStart) < 0 { + if len(routeStart) > 0 && bytes.Compare(routeKey, routeStart) < 0 { return false } - return routeEnd == nil || bytes.Compare(routeKey, routeEnd) < 0 + return len(routeEnd) == 0 || bytes.Compare(routeKey, routeEnd) < 0 } } diff --git a/adapter/internal_migration_test.go b/adapter/internal_migration_test.go index b43a8f4b9..57ca63494 100644 --- a/adapter/internal_migration_test.go +++ b/adapter/internal_migration_test.go @@ -124,6 +124,14 @@ 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, + }, stream) + require.Error(t, err) + require.Equal(t, codes.InvalidArgument, status.Code(err)) } func TestInternalExportRangeVersionsUsesDecodedS3BucketRouteFilter(t *testing.T) { @@ -188,6 +196,35 @@ func TestInternalExportRangeVersionsUsesDecodedS3BucketRouteFilter(t *testing.T) } } +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 TestInternalImportRangeVersionsAppliesStoreBatch(t *testing.T) { t.Parallel() diff --git a/kv/fsm_migration_import.go b/kv/fsm_migration_import.go index 9de47d897..91dca3070 100644 --- a/kv/fsm_migration_import.go +++ b/kv/fsm_migration_import.go @@ -34,7 +34,7 @@ func (f *kvFSM) applyMigrationImport(ctx context.Context, data []byte) any { if err := proto.Unmarshal(data, req); err != nil { return errors.WithStack(err) } - result, err := f.store.ImportVersions(ctx, store.ImportVersionsOptions{ + result, err := f.store.ImportVersionsRaft(ctx, store.ImportVersionsOptions{ JobID: req.GetJobId(), BracketID: req.GetBracketId(), BatchSeq: req.GetBatchSeq(), diff --git a/kv/leader_routed_store.go b/kv/leader_routed_store.go index 3f6004242..45569789d 100644 --- a/kv/leader_routed_store.go +++ b/kv/leader_routed_store.go @@ -519,6 +519,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/shard_store.go b/kv/shard_store.go index fcf47f15c..c69b47e3b 100644 --- a/kv/shard_store.go +++ b/kv/shard_store.go @@ -26,6 +26,7 @@ type ShardStore struct { } var ErrCrossShardMutationBatchNotSupported = errors.New("cross-shard mutation batches are not supported") +var ErrExplicitGroupStagedVisibilityUnresolved = errors.New("explicit group read cannot resolve staged visibility route") // NewShardStore creates a sharded MVCC store wrapper. func NewShardStore(engine *distribution.Engine, groups map[uint64]*ShardGroup) *ShardStore { @@ -61,12 +62,16 @@ func (s *ShardStore) GetGroupAt(ctx context.Context, groupID uint64, key []byte, if !ok || g.Store == nil { return nil, store.ErrKeyNotFound } + route, err := s.routeForExplicitGroupKey(groupID, key) + if err != nil { + return nil, err + } if engineForGroup(g) == nil { - return s.localGetAt(ctx, g, distribution.Route{}, key, ts) + return s.localGetAt(ctx, g, route, key, ts) } if isLinearizableRaftLeader(ctx, engineForGroup(g)) { - return s.leaderGetAt(ctx, g, distribution.Route{}, key, ts) + return s.leaderGetAt(ctx, g, route, key, ts) } return s.proxyRawGet(ctx, g, key, ts, groupID) } @@ -112,6 +117,32 @@ 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.engine.GetRoute(routeKey(key)); ok && route.GroupID == groupID { + return route, nil + } + if s.groupHasStagedVisibility(groupID) { + return distribution.Route{}, errors.Wrapf(ErrExplicitGroupStagedVisibilityUnresolved, "group_id=%d key=%q", groupID, key) + } + return fallback, nil +} + +func (s *ShardStore) groupHasStagedVisibility(groupID uint64) bool { + if s == nil || s.engine == nil { + return false + } + for _, route := range s.engine.Stats() { + if route.GroupID == groupID && routeHasStagedVisibility(route) { + return true + } + } + return false +} + 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 @@ -312,7 +343,29 @@ func (s *ShardStore) ScanGroupAt(ctx context.Context, groupID uint64, start []by 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.routesForExplicitGroupScan(groupID, start, end) + if err != nil { + return nil, err + } + out := make([]*store.KVPair, 0) + for _, route := range routes { + scanStart := start + scanEnd := end + if clampToRoutes { + scanStart = clampScanStart(start, route.Start) + scanEnd = clampScanEnd(end, route.End) + } + kvs, err := s.scanRouteAtDirection(ctx, route, scanStart, scanEnd, limit, ts, false, true) + if err != nil { + return nil, err + } + out = append(out, kvs...) + if len(out) >= limit { + clear(out[limit:]) + return out[:limit], nil + } + } + return out, nil } func (s *ShardStore) ReverseScanAt(ctx context.Context, start []byte, end []byte, limit int, ts uint64) ([]*store.KVPair, error) { @@ -380,6 +433,27 @@ func routesContainStagedVisibility(routes []distribution.Route) bool { return false } +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 + } + routes := s.engine.GetIntersectingRoutes(start, end) + matched := make([]distribution.Route, 0, len(routes)) + for _, route := range routes { + if route.GroupID == groupID { + matched = append(matched, route) + } + } + if len(matched) > 0 { + return matched, true, nil + } + if s.groupHasStagedVisibility(groupID) { + return nil, false, errors.Wrapf(ErrExplicitGroupStagedVisibilityUnresolved, "group_id=%d range=[%q,%q)", groupID, start, end) + } + return fallback, false, nil +} + func (s *ShardStore) scanRoutesAt(ctx context.Context, routes []distribution.Route, start []byte, end []byte, limit int, ts uint64, clampToRoutes bool) ([]*store.KVPair, error) { out := make([]*store.KVPair, 0) for _, route := range routes { @@ -691,37 +765,127 @@ func (s *ShardStore) scanRouteWithStagedVisibility( 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, err + return nil, nil, false, err } stagedKVs, err := scanVisibleCandidates(ctx, g.Store, stagedStart, stagedEnd, window, ts, reverse) if err != nil { - return nil, err + return nil, nil, false, err } versions, err := s.latestStagedVisibilityCandidates(ctx, g.Store, route, liveKVs, stagedKVs, ts) if err != nil { - return nil, err + return nil, nil, false, err } out := visibleLogicalKVs(versions, ts, reverse) + boundary, hasBoundary := stagedVisibilityCandidateBoundary(liveKVs, stagedKVs, reverse) + exhausted := len(liveKVs) < window && len(stagedKVs) < window if len(out) >= limit { clear(out[limit:]) - return out[:limit], nil + return out[:limit], boundary, !exhausted && hasBoundary, nil } - if len(liveKVs) < window && len(stagedKVs) < window { - return out, nil + if exhausted { + return out, nil, false, nil } nextWindow := nextStagedVisibilityCandidateWindow(window) if nextWindow == window { - return out, nil + return out, boundary, hasBoundary, nil } window = nextWindow } } +func stagedVisibilityCandidateBoundary(liveKVs []*store.KVPair, stagedKVs []*store.KVPair, reverse bool) ([]byte, bool) { + boundary := stagedVisibilityBoundary{reverse: reverse} + for _, kvp := range liveKVs { + if kvp == nil { + continue + } + boundary.visit(kvp.Key) + } + for _, kvp := range stagedKVs { + rawKey, stagedOK := stagedVisibilityRawCandidateKey(kvp) + if !stagedOK { + continue + } + boundary.visit(rawKey) + } + return boundary.key, boundary.ok +} + +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 @@ -1676,6 +1840,7 @@ func (s *ShardStore) ApplyMutations(ctx context.Context, mutations []*store.KVPa if err := s.ensureMutationWriteTimestampFloors(mutations, commitTS); err != nil { return err } + readKeys = s.readKeysWithStagedVisibilityAliases(group, readKeys) return errors.WithStack(group.Store.ApplyMutations(ctx, mutations, readKeys, startTS, commitTS)) } @@ -1689,6 +1854,7 @@ func (s *ShardStore) ApplyMutationsRaft(ctx context.Context, mutations []*store. if err := s.ensureMutationWriteTimestampFloors(mutations, commitTS); err != nil { return err } + readKeys = s.readKeysWithStagedVisibilityAliases(group, readKeys) return errors.WithStack(group.Store.ApplyMutationsRaft(ctx, mutations, readKeys, startTS, commitTS)) } @@ -1703,6 +1869,7 @@ func (s *ShardStore) ApplyMutationsRaftAt(ctx context.Context, mutations []*stor if err := s.ensureMutationWriteTimestampFloors(mutations, commitTS); err != nil { return err } + readKeys = s.readKeysWithStagedVisibilityAliases(group, readKeys) return errors.WithStack(group.Store.ApplyMutationsRaftAt(ctx, mutations, readKeys, startTS, commitTS, appliedIndex)) } @@ -1732,6 +1899,41 @@ func (s *ShardStore) ensureMutationWriteTimestampFloors(mutations []*store.KVPai return nil } +func (s *ShardStore) readKeysWithStagedVisibilityAliases(group *ShardGroup, readKeys [][]byte) [][]byte { + if len(readKeys) == 0 || s == nil || s.engine == nil || group == nil { + return readKeys + } + var out [][]byte + for _, key := range readKeys { + alias, ok := s.stagedVisibilityReadKeyAlias(group, key) + if !ok { + continue + } + if out == nil { + out = append([][]byte(nil), readKeys...) + } + out = append(out, alias) + } + if out == nil { + return readKeys + } + return out +} + +func (s *ShardStore) stagedVisibilityReadKeyAlias(group *ShardGroup, key []byte) ([]byte, bool) { + if len(key) == 0 { + return nil, false + } + if _, _, ok := distribution.MigrationStagedDataKeyParts(key); ok { + return nil, false + } + route, g, ok := s.routeAndGroupForKey(key) + if !ok || g != group || !routeHasStagedVisibility(route) { + return nil, false + } + return distribution.MigrationStagedDataKey(route.MigrationJobID, key), true +} + // resolveSingleShardGroup returns the shard group that owns every // mutation in the batch, or an error if the batch is cross-shard or // references an unknown group. A nil group with nil error means "empty @@ -1976,6 +2178,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 } diff --git a/kv/shard_store_test.go b/kv/shard_store_test.go index e3a3ea8d5..9ff8677a4 100644 --- a/kv/shard_store_test.go +++ b/kv/shard_store_test.go @@ -2,6 +2,7 @@ package kv import ( "context" + "fmt" "testing" "github.com/bootjp/elastickv/distribution" @@ -159,6 +160,66 @@ func TestShardStoreScanAndLatestCommitTS_MergeStagedVisibility(t *testing.T) { require.Equal(t, uint64(40), ts) } +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) +} + +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 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 TestShardStorePhysicalLimitFailsClosedBeforeStagedVisibilityFallback(t *testing.T) { t.Parallel() diff --git a/kv/sharded_coordinator.go b/kv/sharded_coordinator.go index a071a7e86..9cb4b8eee 100644 --- a/kv/sharded_coordinator.go +++ b/kv/sharded_coordinator.go @@ -1297,6 +1297,7 @@ func (c *ShardedCoordinator) dispatchSingleShardTxn(ctx context.Context, startTS if err != nil { return nil, err } + readKeys = c.readKeysWithStagedVisibilityAliasesForGroup(gid, readKeys) // 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 @@ -1313,6 +1314,27 @@ func (c *ShardedCoordinator) dispatchSingleShardTxn(ctx context.Context, startTS return &CoordinateResponse{CommitIndex: resp.CommitIndex}, nil } +func (c *ShardedCoordinator) readKeysWithStagedVisibilityAliasesForGroup(gid uint64, readKeys [][]byte) [][]byte { + if len(readKeys) == 0 { + return readKeys + } + var out [][]byte + for _, key := range readKeys { + alias, ok := c.stagedVisibilityReadKeyAlias(gid, key) + if !ok { + continue + } + if out == nil { + out = append([][]byte(nil), readKeys...) + } + out = append(out, alias) + } + if out == nil { + return readKeys + } + return out +} + type preparedGroup struct { gid uint64 keys []*pb.Mutation @@ -1929,10 +1951,27 @@ 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) 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 + } + 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. @@ -1982,7 +2021,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) } @@ -1993,6 +2032,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 { @@ -2074,6 +2150,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 } diff --git a/kv/sharded_coordinator_txn_test.go b/kv/sharded_coordinator_txn_test.go index c07733d16..588404cfb 100644 --- a/kv/sharded_coordinator_txn_test.go +++ b/kv/sharded_coordinator_txn_test.go @@ -43,6 +43,75 @@ 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), + }, txn.requests[0].ReadKeys) +} + func cloneTxnRequest(req *pb.Request) *pb.Request { if req == nil { return nil diff --git a/kv/tso_test.go b/kv/tso_test.go index 8daa8651d..fa256fb83 100644 --- a/kv/tso_test.go +++ b/kv/tso_test.go @@ -284,6 +284,21 @@ func TestShardedCoordinatorRawFollowerDefersTSOAllocationToLeaderPath(t *testing require.EqualValues(t, 0, txn.requests[0].Ts) } +func TestShardedCoordinatorRejectsTSORawPointWriteAfterStamping(t *testing.T) { + t.Parallel() + + txn := &recordingTransactional{} + coord := NewShardedCoordinator(newMigrationFloorEngine(t, testTSOInitialBase), map[uint64]*ShardGroup{ + 1: {Txn: txn}, + }, 1, NewHLC(), nil).WithTSOAllocator(&fakeTSOAllocator{nextBase: testTSOInitialBase, leader: true}) + + _, err := coord.Dispatch(context.Background(), &OperationGroup[OP]{ + Elems: []*Elem[OP]{{Op: Put, Key: []byte("z"), Value: []byte("v")}}, + }) + require.ErrorIs(t, err, ErrRouteWriteTimestampTooLow) + require.Empty(t, txn.requests, "coordinator must reject after TSO stamping before proposing") +} + func TestShardedCoordinatorUsesTSOAllocatorForRawTxnAndDelPrefix(t *testing.T) { t.Parallel() diff --git a/store/lsm_migration.go b/store/lsm_migration.go index 3ae343afd..40c4efca8 100644 --- a/store/lsm_migration.go +++ b/store/lsm_migration.go @@ -192,6 +192,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() @@ -207,7 +215,7 @@ func (s *pebbleStore) ImportVersions(ctx context.Context, opts ImportVersionsOpt } batchMax := importBatchMaxTS(opts.Versions) - if err := s.commitPebbleImportBatch(opts, batchMax); err != nil { + if err := s.commitPebbleImportBatch(opts, batchMax, writeOpts, gateRegistration); err != nil { return ImportVersionsResult{}, errors.WithStack(err) } s.log.InfoContext(ctx, "import_versions", @@ -240,10 +248,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 := batch.Set(migrationAckKey(opts.JobID, opts.BracketID), encodeMigrationImportAck(migrationImportAck{ @@ -257,7 +265,7 @@ func (s *pebbleStore) commitPebbleImportBatch(opts ImportVersionsOptions, batchM return err } defer unlock() - if err := batch.Commit(s.directApplyWriteOpts()); err != nil { + if err := batch.Commit(writeOpts); err != nil { return errors.WithStack(err) } if batchMax > 0 { @@ -273,14 +281,14 @@ 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 := encodeKey(version.Key, version.CommitTS) var encoded []byte 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_registration_gate_test.go b/store/lsm_store_registration_gate_test.go index 4dbfeb6a5..8039a9afb 100644 --- a/store/lsm_store_registration_gate_test.go +++ b/store/lsm_store_registration_gate_test.go @@ -96,6 +96,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 +157,20 @@ 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") } // TestRegistrationGate_NotEncryptingIsUngated confirms the gate is only diff --git a/store/migration_versions.go b/store/migration_versions.go index 54b7dd19b..8211e060b 100644 --- a/store/migration_versions.go +++ b/store/migration_versions.go @@ -365,6 +365,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/store.go b/store/store.go index 78cb59354..5e4e75b53 100644 --- a/store/store.go +++ b/store/store.go @@ -265,6 +265,9 @@ type MVCCStore interface { // ImportVersions applies a migration import batch idempotently by // (jobID, bracketID, batchSeq), preserving tombstones and expireAt. ImportVersions(ctx context.Context, opts ImportVersionsOptions) (ImportVersionsResult, error) + // ImportVersionsRaft is the raft-apply variant of ImportVersions. It + // preserves the same idempotency contract while using the FSM write path. + ImportVersionsRaft(ctx context.Context, opts ImportVersionsOptions) (ImportVersionsResult, error) // MigrationHLCFloor returns the full-HLC target-local migration floor // persisted by ImportVersions for jobID. MigrationHLCFloor(ctx context.Context, jobID uint64) (uint64, error) From 4b8e345aef5629f57d151647a25c9da73249bca4 Mon Sep 17 00:00:00 2001 From: bootjp Date: Tue, 14 Jul 2026 07:02:01 +0900 Subject: [PATCH 08/13] Harden staged visibility OCC guards --- kv/shard_store.go | 47 +++++++++++++++++++++------- kv/shard_store_test.go | 71 ++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 106 insertions(+), 12 deletions(-) diff --git a/kv/shard_store.go b/kv/shard_store.go index c69b47e3b..ee7fa66a7 100644 --- a/kv/shard_store.go +++ b/kv/shard_store.go @@ -122,8 +122,13 @@ func (s *ShardStore) routeForExplicitGroupKey(groupID uint64, key []byte) (distr if s == nil || s.engine == nil { return fallback, nil } - if route, ok := s.engine.GetRoute(routeKey(key)); ok && route.GroupID == groupID { - return route, nil + if route, ok := s.engine.GetRoute(routeKey(key)); ok { + if route.GroupID == groupID { + return route, nil + } + if routeHasStagedVisibility(route) { + return distribution.Route{}, errors.Wrapf(ErrExplicitGroupStagedVisibilityUnresolved, "group_id=%d key=%q", groupID, key) + } } if s.groupHasStagedVisibility(groupID) { return distribution.Route{}, errors.Wrapf(ErrExplicitGroupStagedVisibilityUnresolved, "group_id=%d key=%q", groupID, key) @@ -443,6 +448,10 @@ func (s *ShardStore) routesForExplicitGroupScan(groupID uint64, start []byte, en 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 { @@ -1841,6 +1850,7 @@ func (s *ShardStore) ApplyMutations(ctx context.Context, mutations []*store.KVPa return err } readKeys = s.readKeysWithStagedVisibilityAliases(group, readKeys) + readKeys = s.readKeysWithStagedVisibilityMutationAliases(group, readKeys, mutations) return errors.WithStack(group.Store.ApplyMutations(ctx, mutations, readKeys, startTS, commitTS)) } @@ -1855,6 +1865,7 @@ func (s *ShardStore) ApplyMutationsRaft(ctx context.Context, mutations []*store. return err } readKeys = s.readKeysWithStagedVisibilityAliases(group, readKeys) + readKeys = s.readKeysWithStagedVisibilityMutationAliases(group, readKeys, mutations) return errors.WithStack(group.Store.ApplyMutationsRaft(ctx, mutations, readKeys, startTS, commitTS)) } @@ -1870,6 +1881,7 @@ func (s *ShardStore) ApplyMutationsRaftAt(ctx context.Context, mutations []*stor 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)) } @@ -1900,24 +1912,35 @@ func (s *ShardStore) ensureMutationWriteTimestampFloors(mutations []*store.KVPai } func (s *ShardStore) readKeysWithStagedVisibilityAliases(group *ShardGroup, readKeys [][]byte) [][]byte { - if len(readKeys) == 0 || s == nil || s.engine == nil || group == nil { + if len(readKeys) == 0 { return readKeys } - var out [][]byte for _, key := range readKeys { - alias, ok := s.stagedVisibilityReadKeyAlias(group, key) - if !ok { + readKeys = s.appendStagedVisibilityAlias(group, readKeys, key) + } + return readKeys +} + +func (s *ShardStore) readKeysWithStagedVisibilityMutationAliases(group *ShardGroup, readKeys [][]byte, mutations []*store.KVPairMutation) [][]byte { + for _, mut := range mutations { + if mut == nil { continue } - if out == nil { - out = append([][]byte(nil), readKeys...) - } - out = append(out, alias) + readKeys = s.appendStagedVisibilityAlias(group, readKeys, mut.Key) } - if out == nil { + return readKeys +} + +func (s *ShardStore) appendStagedVisibilityAlias(group *ShardGroup, readKeys [][]byte, key []byte) [][]byte { + if s == nil || s.engine == nil || group == nil { return readKeys } - return out + alias, ok := s.stagedVisibilityReadKeyAlias(group, key) + if !ok { + return readKeys + } + out := append([][]byte(nil), readKeys...) + return append(out, alias) } func (s *ShardStore) stagedVisibilityReadKeyAlias(group *ShardGroup, key []byte) ([]byte, bool) { diff --git a/kv/shard_store_test.go b/kv/shard_store_test.go index 9ff8677a4..2040402c3 100644 --- a/kv/shard_store_test.go +++ b/kv/shard_store_test.go @@ -182,6 +182,40 @@ func TestShardStoreExplicitGroupReads_MergeStagedVisibility(t *testing.T) { }, 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 TestShardStoreScanAt_ContinuesStagedVisibilityAfterCandidateWindow(t *testing.T) { t.Parallel() @@ -220,6 +254,43 @@ func TestShardStoreApplyMutations_ValidatesStagedReadKeys(t *testing.T) { require.ErrorIs(t, err, store.ErrWriteConflict) } +func TestShardStoreApplyMutations_ValidatesStagedWriteKeys(t *testing.T) { + t.Parallel() + + ctx := context.Background() + st, group := newStagedVisibilityShardStore(t) + writeKey := []byte("k") + require.NoError(t, group.Store.PutAt(ctx, distribution.MigrationStagedDataKey(9, writeKey), []byte("staged"), 20, 0)) + + apply := []struct { + name string + fn func(context.Context, []*store.KVPairMutation, [][]byte, uint64, uint64) error + }{ + { + name: "direct", + fn: st.ApplyMutations, + }, + { + name: "raft", + fn: st.ApplyMutationsRaft, + }, + { + name: "raft_at", + fn: func(ctx context.Context, muts []*store.KVPairMutation, readKeys [][]byte, startTS, commitTS uint64) error { + return st.ApplyMutationsRaftAt(ctx, muts, readKeys, startTS, commitTS, 1) + }, + }, + } + for _, tc := range apply { + t.Run(tc.name, func(t *testing.T) { + err := tc.fn(ctx, []*store.KVPairMutation{ + {Op: store.OpTypePut, Key: writeKey, Value: []byte("write")}, + }, nil, 10, 101) + require.ErrorIs(t, err, store.ErrWriteConflict) + }) + } +} + func TestShardStorePhysicalLimitFailsClosedBeforeStagedVisibilityFallback(t *testing.T) { t.Parallel() From ca503cff888dc1c5cdd979932cbbed0b537a5c7e Mon Sep 17 00:00:00 2001 From: bootjp Date: Tue, 14 Jul 2026 07:35:11 +0900 Subject: [PATCH 09/13] Fix staged migration read safety --- kv/fsm.go | 82 ++-------------------- kv/fsm_migration_fence_test.go | 45 +++---------- kv/shard_store.go | 94 ++++++++++++++++++-------- kv/shard_store_test.go | 105 +++++++++++++++++++++++++++++ kv/sharded_coordinator.go | 57 ++++++++++++++++ kv/sharded_coordinator_txn_test.go | 83 +++++++++++++++++++++++ 6 files changed, 329 insertions(+), 137 deletions(-) diff --git a/kv/fsm.go b/kv/fsm.go index 1693d5eaa..9e82f62b5 100644 --- a/kv/fsm.go +++ b/kv/fsm.go @@ -501,7 +501,7 @@ func (f *kvFSM) handleRawRequest(ctx context.Context, r *pb.Request, commitTS ui return f.handleDelPrefix(ctx, prefix, commitTS) } - if err := f.validateRawMutationsForApply(ctx, r.Mutations, commitTS); err != nil { + if err := f.validateRawMutationsForApply(ctx, r.Mutations); err != nil { return err } @@ -518,16 +518,16 @@ func (f *kvFSM) handleRawRequest(ctx context.Context, r *pb.Request, commitTS ui return nil } -func (f *kvFSM) validateRawMutationsForApply(ctx context.Context, muts []*pb.Mutation, commitTS uint64) error { +func (f *kvFSM) validateRawMutationsForApply(ctx context.Context, muts []*pb.Mutation) error { for _, mut := range muts { - if err := f.validateRawMutationForApply(ctx, mut, commitTS); err != nil { + if err := f.validateRawMutationForApply(ctx, mut); err != nil { return err } } return nil } -func (f *kvFSM) validateRawMutationForApply(ctx context.Context, mut *pb.Mutation, commitTS uint64) error { +func (f *kvFSM) validateRawMutationForApply(ctx context.Context, mut *pb.Mutation) error { if mut == nil || len(mut.Key) == 0 { return errors.WithStack(ErrInvalidRequest) } @@ -538,9 +538,6 @@ 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.assertNoConflictingTxnLock(ctx, mut.Key, nil, 0); err != nil { return err } @@ -564,9 +561,6 @@ 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 { - return err - } if err := f.store.DeletePrefixAtRaftAt(ctx, prefix, txnCommonPrefix, commitTS, f.pendingApplyIdx); err != nil { return errors.WithStack(err) } @@ -574,51 +568,6 @@ func (f *kvFSM) handleDelPrefix(ctx context.Context, prefix []byte, commitTS uin return nil } -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 - } - rkey := routeKey(key) - route, ok := snap.RouteOf(rkey) - if !ok || route.MinWriteTSExclusive == 0 || commitTS > route.MinWriteTSExclusive { - return nil - } - return errors.Wrapf(ErrRouteWriteTimestampTooLow, "key %q routeKey %q commit_ts=%d floor=%d", key, rkey, commitTS, route.MinWriteTSExclusive) -} - -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 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 -} - -func (f *kvFSM) verifyRouteWriteTimestampFloorForMutations(muts []*pb.Mutation, commitTS uint64) error { - for _, mut := range muts { - if mut == nil || len(mut.Key) == 0 || isTxnInternalKey(mut.Key) { - continue - } - if err := f.verifyRouteWriteTimestampFloorForKey(mut.Key, commitTS); err != nil { - return err - } - } - return nil -} - func (f *kvFSM) verifyRouteNotFencedForMutations(muts []*pb.Mutation) error { for _, mut := range muts { if mut == nil || len(mut.Key) == 0 || isTxnInternalKey(mut.Key) { @@ -1036,9 +985,6 @@ func (f *kvFSM) handlePrepareRequest(ctx context.Context, r *pb.Request) error { if err := f.verifyRouteNotFencedForMutations(uniq); err != nil { return err } - if err := f.verifyRouteWriteTimestampFloorForMutations(uniq, meta.CommitTS); err != nil { - return err - } if err := f.validateConflicts(ctx, uniq, startTS); err != nil { return errors.WithStack(err) } @@ -1105,7 +1051,7 @@ func (f *kvFSM) handleOnePhaseTxnRequest(ctx context.Context, r *pb.Request, com return nil } - uniq, err := f.uniqueMutationsNotFencedAboveWriteFloor(muts, commitTS) + uniq, err := f.uniqueMutationsNotFenced(muts) if err != nil { return err } @@ -1132,25 +1078,11 @@ func (f *kvFSM) uniqueMutationsNotFenced(muts []*pb.Mutation) ([]*pb.Mutation, e return uniq, nil } -func (f *kvFSM) uniqueMutationsNotFencedAboveWriteFloor(muts []*pb.Mutation, commitTS uint64) ([]*pb.Mutation, error) { - uniq, err := f.uniqueMutationsNotFenced(muts) - if err != nil { - return nil, err - } - if err := f.verifyRouteWriteTimestampFloorForMutations(uniq, commitTS); err != nil { - return nil, err - } - return uniq, nil -} - -func (f *kvFSM) uniqueMutationsAboveWriteFloor(muts []*pb.Mutation, commitTS uint64) ([]*pb.Mutation, error) { +func uniqueTxnMutations(muts []*pb.Mutation) ([]*pb.Mutation, error) { uniq, err := uniqueMutations(muts) if err != nil { return nil, err } - if err := f.verifyRouteWriteTimestampFloorForMutations(uniq, commitTS); err != nil { - return nil, err - } return uniq, nil } @@ -1193,7 +1125,7 @@ func (f *kvFSM) handleCommitRequest(ctx context.Context, r *pb.Request) error { if err != nil { return err } - uniq, err := f.uniqueMutationsAboveWriteFloor(muts, commitTS) + uniq, err := uniqueTxnMutations(muts) if err != nil { return err } diff --git a/kv/fsm_migration_fence_test.go b/kv/fsm_migration_fence_test.go index 045a9d04d..70be77390 100644 --- a/kv/fsm_migration_fence_test.go +++ b/kv/fsm_migration_fence_test.go @@ -123,29 +123,21 @@ func TestFSMRejectsPrepareOnWriteFencedRouteButAllowsAbort(t *testing.T) { require.False(t, errors.Is(err, ErrRouteWriteFenced), "ABORT must keep the narrow cleanup lane open") } -func TestFSMRejectsRawPointWriteAtMigrationTimestampFloor(t *testing.T) { +func TestFSMAllowsRawPointWriteAtMigrationTimestampFloorDuringReplay(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("low")}}, + 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) - - err = fsm.handleRawRequest(ctx, &pb.Request{ - Mutations: []*pb.Mutation{{Op: pb.Op_PUT, Key: []byte("z"), Value: []byte("ok")}}, - }, 101) require.NoError(t, err) got, getErr := fsm.store.GetAt(ctx, []byte("z"), ^uint64(0)) require.NoError(t, getErr) - require.Equal(t, []byte("ok"), got) + require.Equal(t, []byte("replayed"), got) } -func TestFSMRejectsDelPrefixAtMigrationTimestampFloor(t *testing.T) { +func TestFSMAllowsDelPrefixAtMigrationTimestampFloorDuringReplay(t *testing.T) { t.Parallel() ctx := context.Background() @@ -155,14 +147,13 @@ func TestFSMRejectsDelPrefixAtMigrationTimestampFloor(t *testing.T) { err := fsm.handleRawRequest(ctx, &pb.Request{ Mutations: []*pb.Mutation{{Op: pb.Op_DEL_PREFIX, Key: []byte("z")}}, }, 100) - require.ErrorIs(t, err, ErrRouteWriteTimestampTooLow) + require.NoError(t, err) - got, getErr := fsm.store.GetAt(ctx, []byte("z"), ^uint64(0)) - require.NoError(t, getErr) - require.Equal(t, []byte("v"), got) + _, getErr := fsm.store.GetAt(ctx, []byte("z"), ^uint64(0)) + require.ErrorIs(t, getErr, store.ErrKeyNotFound) } -func TestFSMRejectsOnePhaseTxnAtMigrationTimestampFloor(t *testing.T) { +func TestFSMAllowsOnePhaseTxnAtMigrationTimestampFloorDuringReplay(t *testing.T) { t.Parallel() ctx := context.Background() @@ -177,21 +168,13 @@ func TestFSMRejectsOnePhaseTxnAtMigrationTimestampFloor(t *testing.T) { }, } 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) - - req.Mutations[0].Value = EncodeTxnMeta(TxnMeta{PrimaryKey: []byte("z"), CommitTS: 101}) - req.Mutations[1].Value = []byte("ok") - err = fsm.handleTxnRequest(ctx, req, 101) require.NoError(t, err) got, getErr := fsm.store.GetAt(ctx, []byte("z"), ^uint64(0)) require.NoError(t, getErr) - require.Equal(t, []byte("ok"), got) + require.Equal(t, []byte("low"), got) } -func TestFSMRejectsCommitButNotPrepareAtMigrationTimestampFloor(t *testing.T) { +func TestFSMAllowsCommitAtMigrationTimestampFloorDuringReplay(t *testing.T) { t.Parallel() ctx := context.Background() @@ -217,14 +200,6 @@ func TestFSMRejectsCommitButNotPrepareAtMigrationTimestampFloor(t *testing.T) { }, } err := fsm.handleTxnRequest(ctx, commit, 100) - require.ErrorIs(t, err, ErrRouteWriteTimestampTooLow) - - _, getErr := fsm.store.GetAt(ctx, []byte("z"), ^uint64(0)) - require.ErrorIs(t, getErr, store.ErrKeyNotFound) - - commit.Mutations[0].Value = EncodeTxnMeta(TxnMeta{PrimaryKey: []byte("z"), CommitTS: 101}) - commit.Mutations[1].Value = []byte("ignored") - err = fsm.handleTxnRequest(ctx, commit, 101) require.NoError(t, err) got, getErr := fsm.store.GetAt(ctx, []byte("z"), ^uint64(0)) require.NoError(t, getErr) diff --git a/kv/shard_store.go b/kv/shard_store.go index ee7fa66a7..7e681d53e 100644 --- a/kv/shard_store.go +++ b/kv/shard_store.go @@ -130,24 +130,9 @@ func (s *ShardStore) routeForExplicitGroupKey(groupID uint64, key []byte) (distr return distribution.Route{}, errors.Wrapf(ErrExplicitGroupStagedVisibilityUnresolved, "group_id=%d key=%q", groupID, key) } } - if s.groupHasStagedVisibility(groupID) { - return distribution.Route{}, errors.Wrapf(ErrExplicitGroupStagedVisibilityUnresolved, "group_id=%d key=%q", groupID, key) - } return fallback, nil } -func (s *ShardStore) groupHasStagedVisibility(groupID uint64) bool { - if s == nil || s.engine == nil { - return false - } - for _, route := range s.engine.Stats() { - if route.GroupID == groupID && routeHasStagedVisibility(route) { - return true - } - } - return false -} - 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 @@ -443,7 +428,8 @@ func (s *ShardStore) routesForExplicitGroupScan(groupID uint64, start []byte, en if s == nil || s.engine == nil { return fallback, false, nil } - routes := s.engine.GetIntersectingRoutes(start, end) + 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 { @@ -455,14 +441,30 @@ func (s *ShardStore) routesForExplicitGroupScan(groupID uint64, start []byte, en } } if len(matched) > 0 { - return matched, true, nil - } - if s.groupHasStagedVisibility(groupID) { - return nil, false, errors.Wrapf(ErrExplicitGroupStagedVisibilityUnresolved, "group_id=%d range=[%q,%q)", groupID, start, end) + return matched, !routeMapped, nil } return fallback, false, nil } +func explicitGroupScanRouteBounds(start []byte, end []byte) ([]byte, []byte, bool) { + routeStart := routeKey(start) + if len(start) == 0 { + routeStart = []byte("") + } + routeEnd := end + routeMapped := !bytes.Equal(routeStart, start) + if end != nil { + normalizedEnd := routeKey(end) + if !bytes.Equal(normalizedEnd, end) { + routeMapped = true + } + } + if routeMapped && len(routeStart) != 0 { + routeEnd = prefixScanEnd(routeStart) + } + return routeStart, routeEnd, routeMapped +} + func (s *ShardStore) scanRoutesAt(ctx context.Context, routes []distribution.Route, start []byte, end []byte, limit int, ts uint64, clampToRoutes bool) ([]*store.KVPair, error) { out := make([]*store.KVPair, 0) for _, route := range routes { @@ -826,8 +828,10 @@ func (s *ShardStore) scanRouteWithStagedVisibilityPage( return nil, nil, false, err } out := visibleLogicalKVs(versions, ts, reverse) - boundary, hasBoundary := stagedVisibilityCandidateBoundary(liveKVs, stagedKVs, reverse) - exhausted := len(liveKVs) < window && len(stagedKVs) < window + liveExhausted := len(liveKVs) < window + stagedExhausted := len(stagedKVs) < window + boundary, hasBoundary := stagedVisibilityCandidateBoundary(liveKVs, stagedKVs, liveExhausted, stagedExhausted, reverse) + exhausted := liveExhausted && stagedExhausted if len(out) >= limit { clear(out[limit:]) return out[:limit], boundary, !exhausted && hasBoundary, nil @@ -843,22 +847,58 @@ func (s *ShardStore) scanRouteWithStagedVisibilityPage( } } -func stagedVisibilityCandidateBoundary(liveKVs []*store.KVPair, stagedKVs []*store.KVPair, reverse bool) ([]byte, bool) { - boundary := stagedVisibilityBoundary{reverse: reverse} +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 } - boundary.visit(kvp.Key) + liveBoundary.visit(kvp.Key) } + stagedBoundary := stagedVisibilityBoundary{reverse: reverse} for _, kvp := range stagedKVs { rawKey, stagedOK := stagedVisibilityRawCandidateKey(kvp) if !stagedOK { continue } - boundary.visit(rawKey) + 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 boundary.key, boundary.ok + return b } type stagedVisibilityBoundary struct { diff --git a/kv/shard_store_test.go b/kv/shard_store_test.go index 2040402c3..cfa7bc3e3 100644 --- a/kv/shard_store_test.go +++ b/kv/shard_store_test.go @@ -216,6 +216,74 @@ func TestShardStoreExplicitGroupReads_FailClosedWhenRouteMovedToStagedGroup(t *t 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() @@ -240,6 +308,23 @@ func TestShardStoreScanAt_ContinuesStagedVisibilityAfterCandidateWindow(t *testi require.Equal(t, []byte("k00000"), kvs[limit-1].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() @@ -402,6 +487,26 @@ 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 TestShardStoreGetGroupAt_UsesExplicitGroup(t *testing.T) { t.Parallel() diff --git a/kv/sharded_coordinator.go b/kv/sharded_coordinator.go index 9cb4b8eee..44805ce8d 100644 --- a/kv/sharded_coordinator.go +++ b/kv/sharded_coordinator.go @@ -1223,6 +1223,10 @@ func (c *ShardedCoordinator) dispatchMultiShardTxn(ctx context.Context, startTS, if err != nil { return nil, err } + groupedReadKeys = c.groupedReadKeysWithStagedVisibilityMutationAliases(groupedReadKeys, grouped) + if groupedReadKeyCount(groupedReadKeys) > maxReadKeys { + return nil, errors.WithStack(ErrInvalidRequest) + } prepared, err := c.prewriteTxn(ctx, startTS, commitTS, primaryKey, grouped, gids, groupedReadKeys, observedRouteVersion) if err != nil { return nil, err @@ -1298,6 +1302,10 @@ func (c *ShardedCoordinator) dispatchSingleShardTxn(ctx context.Context, startTS 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 @@ -1335,6 +1343,27 @@ func (c *ShardedCoordinator) readKeysWithStagedVisibilityAliasesForGroup(gid uin 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 @@ -1958,6 +1987,34 @@ func (c *ShardedCoordinator) groupReadKeysByShardID(readKeys [][]byte) (map[uint 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 diff --git a/kv/sharded_coordinator_txn_test.go b/kv/sharded_coordinator_txn_test.go index 588404cfb..e40318e00 100644 --- a/kv/sharded_coordinator_txn_test.go +++ b/kv/sharded_coordinator_txn_test.go @@ -109,9 +109,92 @@ func TestShardedCoordinatorDispatchTxn_AddsStagedReadKeyAlias(t *testing.T) { 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 From 2b501297445704217a48cc86b7315fe2a59d885b Mon Sep 17 00:00:00 2001 From: bootjp Date: Tue, 14 Jul 2026 09:07:17 +0900 Subject: [PATCH 10/13] Fix staged migration guard gaps --- adapter/internal.go | 35 +++++++++++++++++++++++++++++++++++ adapter/internal_test.go | 34 ++++++++++++++++++++++++++++++++++ kv/fsm.go | 30 ++++++++++++++++++++++++++++++ kv/fsm_onephase_dedup_test.go | 32 ++++++++++++++++++++++++++++++++ kv/shard_store.go | 27 +++++++++++++++++++++++++++ kv/shard_store_test.go | 22 ++++++++++++++++++++++ main.go | 5 +++++ main_bootstrap_e2e_test.go | 11 +++++++++-- 8 files changed, 194 insertions(+), 2 deletions(-) diff --git a/adapter/internal.go b/adapter/internal.go index a14fb9134..515faf9aa 100644 --- a/adapter/internal.go +++ b/adapter/internal.go @@ -35,6 +35,12 @@ func WithInternalMigrationProposer(proposer raftengine.Proposer) InternalOption } } +func WithInternalRouteEngine(engine *distribution.Engine) InternalOption { + return func(i *Internal) { + i.routeEngine = engine + } +} + func NewInternalWithEngine(txm kv.Transactional, leader raftengine.LeaderView, clock *kv.HLC, relay *RedisPubSubRelay, opts ...InternalOption) *Internal { i := &Internal{ leader: leader, @@ -56,6 +62,7 @@ type Internal struct { relay *RedisPubSubRelay store store.MVCCStore migrationProposer raftengine.Proposer + routeEngine *distribution.Engine pb.UnimplementedInternalServer } @@ -408,10 +415,38 @@ func (i *Internal) stampRawTimestamps(ctx context.Context, reqs []*pb.Request) e 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 || len(mut.Key) == 0 { + continue + } + for _, route := range routes { + if routeWriteTimestampFloorApplies(route, mut.Key, commitTS) { + return errors.Wrapf(kv.ErrRouteWriteTimestampTooLow, "key %q commit_ts=%d floor=%d", mut.Key, commitTS, route.MinWriteTSExclusive) + } + } } return nil } +func routeWriteTimestampFloorApplies(route distribution.Route, key []byte, commitTS uint64) bool { + if route.MinWriteTSExclusive == 0 || commitTS > route.MinWriteTSExclusive { + return false + } + return kv.RouteKeyFilter(route.Start, route.End)(key) +} + func (i *Internal) stampTxnTimestamps(ctx context.Context, reqs []*pb.Request) error { startTS := forwardedTxnStartTS(reqs) if startTS == 0 { diff --git a/adapter/internal_test.go b/adapter/internal_test.go index 49802d20a..a028aa65d 100644 --- a/adapter/internal_test.go +++ b/adapter/internal_test.go @@ -4,11 +4,18 @@ import ( "context" "testing" + "github.com/bootjp/elastickv/distribution" "github.com/bootjp/elastickv/kv" pb "github.com/bootjp/elastickv/proto" "github.com/stretchr/testify/require" ) +type fixedInternalTimestampAllocator uint64 + +func (a fixedInternalTimestampAllocator) Next(context.Context) (uint64, error) { + return uint64(a), nil +} + func TestStampTxnTimestamps_RejectsMaxStartTS(t *testing.T) { t.Parallel() @@ -170,3 +177,30 @@ func TestStampTxnTimestamps_UsesSingleTxnStartTS(t *testing.T) { require.NoError(t, err) require.Greater(t, meta.CommitTS, uint64(9)) } + +func TestStampRawTimestampsRejectsRouteWriteFloor(t *testing.T) { + t.Parallel() + + engine := distribution.NewEngine() + require.NoError(t, engine.ApplySnapshot(distribution.CatalogSnapshot{ + Version: 1, + Routes: []distribution.RouteDescriptor{{ + RouteID: 1, + Start: []byte(""), + End: nil, + GroupID: 1, + State: distribution.RouteStateActive, + MinWriteTSExclusive: 100, + }}, + })) + i := &Internal{ + tsAllocator: fixedInternalTimestampAllocator(100), + routeEngine: engine, + } + reqs := []*pb.Request{{ + Mutations: []*pb.Mutation{{Op: pb.Op_PUT, Key: []byte("k"), Value: []byte("v")}}, + }} + + err := i.stampRawTimestamps(context.Background(), reqs) + require.ErrorIs(t, err, kv.ErrRouteWriteTimestampTooLow) +} diff --git a/kv/fsm.go b/kv/fsm.go index 9e82f62b5..9e703a926 100644 --- a/kv/fsm.go +++ b/kv/fsm.go @@ -1102,9 +1102,39 @@ 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 + } + route, ok := snap.RouteOf(routeKey(key)) + if !ok || route.GroupID != f.shardGroupID || !routeHasStagedVisibility(route) { + return distribution.Route{}, false + } + return route, true +} + func (f *kvFSM) handleCommitRequest(ctx context.Context, r *pb.Request) error { meta, muts, err := extractTxnMeta(r.Mutations) if err != nil { diff --git a/kv/fsm_onephase_dedup_test.go b/kv/fsm_onephase_dedup_test.go index aaf027c1d..8b3f50f28 100644 --- a/kv/fsm_onephase_dedup_test.go +++ b/kv/fsm_onephase_dedup_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" @@ -67,6 +68,37 @@ 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) +} + // 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/shard_store.go b/kv/shard_store.go index 7e681d53e..217f630ad 100644 --- a/kv/shard_store.go +++ b/kv/shard_store.go @@ -832,6 +832,7 @@ func (s *ShardStore) scanRouteWithStagedVisibilityPage( 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 @@ -847,6 +848,32 @@ func (s *ShardStore) scanRouteWithStagedVisibilityPage( } } +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 { diff --git a/kv/shard_store_test.go b/kv/shard_store_test.go index cfa7bc3e3..cc58a4af6 100644 --- a/kv/shard_store_test.go +++ b/kv/shard_store_test.go @@ -308,6 +308,28 @@ func TestShardStoreScanAt_ContinuesStagedVisibilityAfterCandidateWindow(t *testi 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() diff --git a/main.go b/main.go index c236598d8..c77369312 100644 --- a/main.go +++ b/main.go @@ -1491,6 +1491,7 @@ func startServersAfterStartupRotation(waitRotateOnStartup startupRotationWaiter, shardStore: in.shardStore, coordinate: adapterCoordinate, distServer: in.distServer, + routeEngine: in.cfg.engine, adminServer: adminServer, adminGRPCOpts: adminGRPCOpts, redisAddress: *redisAddr, @@ -2053,6 +2054,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, @@ -2097,6 +2099,7 @@ func startRaftServers( internalTimestampOptions(coordinate), adapter.WithInternalStore(rt.store), adapter.WithInternalMigrationProposer(proposerForGroup(rt, shardGroups)), + adapter.WithInternalRouteEngine(routeEngine), )..., )) pb.RegisterDistributionServer(gs, distServer) @@ -2439,6 +2442,7 @@ type runtimeServerRunner struct { shardStore *kv.ShardStore coordinate kv.Coordinator distServer *adapter.DistributionServer + routeEngine *distribution.Engine adminServer *adapter.AdminServer adminGRPCOpts adminGRPCInterceptors redisAddress string @@ -2552,6 +2556,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) diff --git a/main_bootstrap_e2e_test.go b/main_bootstrap_e2e_test.go index 81ca49dce..f4483c137 100644 --- a/main_bootstrap_e2e_test.go +++ b/main_bootstrap_e2e_test.go @@ -15,6 +15,7 @@ import ( "time" "github.com/bootjp/elastickv/adapter" + "github.com/bootjp/elastickv/distribution" internalraftadmin "github.com/bootjp/elastickv/internal/raftadmin" "github.com/bootjp/elastickv/internal/raftengine" "github.com/bootjp/elastickv/kv" @@ -577,6 +578,7 @@ func startBootstrapE2ENode( shardStore, coordinate, distServer, + cfg.engine, cfg.leaderRedis, listeners, ) @@ -660,6 +662,7 @@ func startBootstrapE2EMultiGroupNode( shardStore, coordinate, distServer, + cfg.engine, cfg.leaderRedis, listeners, ) @@ -688,6 +691,7 @@ func startRuntimeServersWithBoundListeners( shardStore *kv.ShardStore, coordinate kv.Coordinator, distServer *adapter.DistributionServer, + routeEngine *distribution.Engine, leaderRedis map[string]string, listeners bootstrapE2EListeners, ) error { @@ -701,7 +705,7 @@ func startRuntimeServersWithBoundListeners( if err := startBoundRedisServer(ctx, eg, listeners.redis, shardStore, coordinate, leaderRedis, redisAddr, relay); err != nil { return waitErrgroupAfterStartupFailure(cancel, eg, err) } - if err := startBoundGRPCServer(ctx, eg, rt, shardStore, coordinate, distServer, relay, listeners.grpc); err != nil { + if err := startBoundGRPCServer(ctx, eg, rt, shardStore, coordinate, distServer, routeEngine, relay, listeners.grpc); err != nil { return waitErrgroupAfterStartupFailure(cancel, eg, err) } if err := startBoundDynamoDBServer(ctx, eg, listeners.dynamo, shardStore, coordinate); err != nil { @@ -718,6 +722,7 @@ func startRuntimeServersWithBoundMultiGroupListeners( shardStore *kv.ShardStore, coordinate kv.Coordinator, distServer *adapter.DistributionServer, + routeEngine *distribution.Engine, leaderRedis map[string]string, listeners bootstrapE2EMultiGroupListeners, ) error { @@ -730,7 +735,7 @@ func startRuntimeServersWithBoundMultiGroupListeners( return waitErrgroupAfterStartupFailure(cancel, eg, err) } for _, rt := range runtimes { - if err := startBoundGRPCServer(ctx, eg, rt, shardStore, coordinate, distServer, relay, listeners.grpc[rt.spec.id]); err != nil { + if err := startBoundGRPCServer(ctx, eg, rt, shardStore, coordinate, distServer, routeEngine, relay, listeners.grpc[rt.spec.id]); err != nil { return waitErrgroupAfterStartupFailure(cancel, eg, err) } } @@ -747,6 +752,7 @@ func startBoundGRPCServer( shardStore *kv.ShardStore, coordinate kv.Coordinator, distServer *adapter.DistributionServer, + routeEngine *distribution.Engine, relay *adapter.RedisPubSubRelay, listener net.Listener, ) error { @@ -769,6 +775,7 @@ func startBoundGRPCServer( relay, adapter.WithInternalStore(rt.store), adapter.WithInternalMigrationProposer(rt.engine), + adapter.WithInternalRouteEngine(routeEngine), )) pb.RegisterDistributionServer(gs, distServer) rt.registerGRPC(gs) From 295cdf6bde424cf9727205e5a6322b973e2ad3ef Mon Sep 17 00:00:00 2001 From: bootjp Date: Tue, 14 Jul 2026 19:01:08 +0900 Subject: [PATCH 11/13] migration: close cross-group floor gaps --- adapter/internal.go | 67 +++++++++--- adapter/internal_migration_test.go | 52 ++++++++++ adapter/internal_test.go | 55 ++++++++++ kv/shard_store.go | 64 +++++++++++- kv/shard_store_test.go | 121 ++++++++++++++++++++++ kv/sharded_coordinator.go | 31 ++++-- kv/sharded_coordinator_del_prefix_test.go | 36 +++++++ 7 files changed, 401 insertions(+), 25 deletions(-) diff --git a/adapter/internal.go b/adapter/internal.go index 515faf9aa..65bd16d95 100644 --- a/adapter/internal.go +++ b/adapter/internal.go @@ -287,12 +287,23 @@ func decodedS3BucketRouteFilter(family uint32, routeStart, routeEnd []byte) func if !ok { return false } - routeKey := s3keys.RouteKey(bucket, 0, "") - if len(routeStart) > 0 && bytes.Compare(routeKey, routeStart) < 0 { - return false - } - return len(routeEnd) == 0 || bytes.Compare(routeKey, routeEnd) < 0 + return decodedS3BucketRouteIntersects(bucket, routeStart, routeEnd) + } +} + +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) { @@ -407,14 +418,13 @@ func (i *Internal) stampRawTimestamps(ctx context.Context, reqs []*pb.Request) e if r == nil { continue } - if r.Ts != 0 { - continue - } - ts, err := i.nextTimestamp(ctx, "stampRawTimestamps") - if err != nil { - return err + if r.Ts == 0 { + ts, err := i.nextTimestamp(ctx, "stampRawTimestamps") + if err != nil { + return err + } + r.Ts = ts } - r.Ts = ts if err := i.rejectWriteTimestampFloorMutations(r.Mutations, r.Ts); err != nil { return err } @@ -467,7 +477,10 @@ func (i *Internal) stampTxnTimestamps(ctx context.Context, reqs []*pb.Request) e } } - return i.fillForwardedTxnCommitTS(ctx, reqs, startTS) + if err := i.fillForwardedTxnCommitTS(ctx, reqs, startTS); err != nil { + return err + } + return i.rejectWriteTimestampFloorTxnRequests(reqs) } func forwardedTxnStartTS(reqs []*pb.Request) uint64 { @@ -533,6 +546,34 @@ func (i *Internal) fillForwardedTxnCommitTS(ctx context.Context, reqs []*pb.Requ return nil } +func (i *Internal) rejectWriteTimestampFloorTxnRequests(reqs []*pb.Request) error { + for _, r := range reqs { + commitTS, ok := forwardedTxnRequestCommitTS(r) + if !ok { + continue + } + if err := i.rejectWriteTimestampFloorMutations(r.Mutations, commitTS); err != nil { + return err + } + } + return nil +} + +func forwardedTxnRequestCommitTS(r *pb.Request) (uint64, bool) { + if r == nil || (r.Phase != pb.Phase_COMMIT && r.Phase != pb.Phase_NONE) { + return 0, false + } + m, ok := forwardedTxnMetaMutation(r, []byte(kv.TxnMetaPrefix)) + if !ok { + return 0, false + } + meta, err := kv.DecodeTxnMeta(m.Value) + if err != nil || meta.CommitTS == 0 { + return 0, false + } + return meta.CommitTS, true +} + // forwardedTxnCommitTS allocates a commit timestamp for a forwarded // transaction, strictly greater than startTS. Pulled out of // fillForwardedTxnCommitTS so that function stays under cyclop. diff --git a/adapter/internal_migration_test.go b/adapter/internal_migration_test.go index 57ca63494..2a403c963 100644 --- a/adapter/internal_migration_test.go +++ b/adapter/internal_migration_test.go @@ -225,6 +225,58 @@ func TestInternalExportRangeVersionsDecodedS3EmptyRouteEndIsUnbounded(t *testing }, 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 TestInternalImportRangeVersionsAppliesStoreBatch(t *testing.T) { t.Parallel() diff --git a/adapter/internal_test.go b/adapter/internal_test.go index a028aa65d..fe4a0c5ea 100644 --- a/adapter/internal_test.go +++ b/adapter/internal_test.go @@ -204,3 +204,58 @@ func TestStampRawTimestampsRejectsRouteWriteFloor(t *testing.T) { 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 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) +} diff --git a/kv/shard_store.go b/kv/shard_store.go index e5fb6e1bb..107be5a49 100644 --- a/kv/shard_store.go +++ b/kv/shard_store.go @@ -140,6 +140,12 @@ 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.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 @@ -1457,9 +1463,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 { @@ -1622,6 +1631,9 @@ func (s *ShardStore) PutAt(ctx context.Context, key []byte, value []byte, commit 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)) } @@ -1633,6 +1645,9 @@ func (s *ShardStore) DeleteAt(ctx context.Context, key []byte, commitTS uint64) 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)) } @@ -1644,6 +1659,9 @@ func (s *ShardStore) PutWithTTLAt(ctx context.Context, key []byte, value []byte, 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)) } @@ -1655,6 +1673,9 @@ func (s *ShardStore) ExpireAt(ctx context.Context, key []byte, expireAt uint64, 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)) } @@ -2434,6 +2455,25 @@ func (s *ShardStore) ensureMutationWriteTimestampFloors(mutations []*store.KVPai 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 } @@ -2778,6 +2818,10 @@ func (s *ShardStore) groupForKey(key []byte) (*ShardGroup, bool) { } func (s *ShardStore) routeAndGroupForKey(key []byte) (distribution.Route, *ShardGroup, bool) { + if route, ok := s.stagedVisibilityRouteForS3BucketAuxiliaryKey(key); ok { + g, ok := s.groups[route.GroupID] + return route, g, ok + } route, ok := s.engine.GetRoute(routeKey(key)) if !ok { return distribution.Route{}, nil, false @@ -2786,6 +2830,22 @@ func (s *ShardStore) routeAndGroupForKey(key []byte) (distribution.Route, *Shard return route, g, ok } +func (s *ShardStore) stagedVisibilityRouteForS3BucketAuxiliaryKey(key []byte) (distribution.Route, bool) { + if s == nil || s.engine == nil { + return distribution.Route{}, false + } + start, end, ok := s3BucketAuxiliaryRouteRange(key) + if !ok { + return distribution.Route{}, false + } + for _, route := range s.engine.GetIntersectingRoutes(start, end) { + if routeHasStagedVisibility(route) { + return route, true + } + } + return distribution.Route{}, false +} + func (s *ShardStore) proxyRawGet(ctx context.Context, g *ShardGroup, key []byte, ts uint64, groupID uint64, readRouteVersion uint64) ([]byte, error) { engine := engineForGroup(g) if engine == nil { diff --git a/kv/shard_store_test.go b/kv/shard_store_test.go index 20cd9ba9e..c4d2615d0 100644 --- a/kv/shard_store_test.go +++ b/kv/shard_store_test.go @@ -101,6 +101,100 @@ func TestShardStoreGetAt_MergesStagedVisibilityPebbleExactKey(t *testing.T) { 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 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() @@ -161,6 +255,33 @@ func TestShardStoreScanAndLatestCommitTS_MergeStagedVisibility(t *testing.T) { require.Equal(t, uint64(40), ts) } +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() diff --git a/kv/sharded_coordinator.go b/kv/sharded_coordinator.go index 0fab4636c..edbac1140 100644 --- a/kv/sharded_coordinator.go +++ b/kv/sharded_coordinator.go @@ -1110,12 +1110,26 @@ func (c *ShardedCoordinator) rejectWriteTimestampFloorPointElems(elems []*Elem[O if elem == nil || len(elem.Key) == 0 { continue } - rkey := routeKey(elem.Key) - route, ok := c.engine.GetRoute(rkey) - if !ok || route.MinWriteTSExclusive == 0 || commitTS > route.MinWriteTSExclusive { - continue + if err := c.rejectWriteTimestampFloorPointKey(elem.Key, commitTS); err != nil { + return err + } + } + return nil +} + +func (c *ShardedCoordinator) rejectWriteTimestampFloorPointKey(key []byte, commitTS uint64) error { + 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) + } + start, end, ok := s3BucketAuxiliaryRouteRange(key) + if !ok { + return nil + } + 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 errors.Wrapf(ErrRouteWriteTimestampTooLow, "key %q routeKey %q commit_ts=%d floor=%d", elem.Key, rkey, commitTS, route.MinWriteTSExclusive) } return nil } @@ -2195,12 +2209,9 @@ func (c *ShardedCoordinator) rejectWriteTimestampFloorMutations(muts []*pb.Mutat if mut == nil || len(mut.Key) == 0 { continue } - rkey := routeKey(mut.Key) - route, ok := c.engine.GetRoute(rkey) - if !ok || route.MinWriteTSExclusive == 0 || commitTS > route.MinWriteTSExclusive { - continue + if err := c.rejectWriteTimestampFloorPointKey(mut.Key, commitTS); err != nil { + return err } - return errors.Wrapf(ErrRouteWriteTimestampTooLow, "key %q routeKey %q commit_ts=%d floor=%d", mut.Key, rkey, commitTS, route.MinWriteTSExclusive) } return nil } diff --git a/kv/sharded_coordinator_del_prefix_test.go b/kv/sharded_coordinator_del_prefix_test.go index 3ffc0a4b9..c6aa9192f 100644 --- a/kv/sharded_coordinator_del_prefix_test.go +++ b/kv/sharded_coordinator_del_prefix_test.go @@ -227,6 +227,42 @@ func TestShardedCoordinatorRejectsS3BucketAuxiliaryPointWriteOnWriteFencedRoute( } } +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() From 650c45072f0207f7705446244c9e0a4565125b26 Mon Sep 17 00:00:00 2001 From: bootjp Date: Tue, 14 Jul 2026 19:41:26 +0900 Subject: [PATCH 12/13] migration: route staged s3 auxiliaries --- adapter/internal.go | 34 ++++++++++++-- adapter/internal_migration_test.go | 49 ++++++++++++++++++++ adapter/internal_test.go | 35 ++++++++++++++ kv/fsm.go | 5 ++ kv/shard_router.go | 19 ++++++++ kv/sharded_coordinator.go | 56 +++++++++++++++++------ kv/sharded_coordinator_del_prefix_test.go | 39 ++++++++++++++++ kv/sharded_coordinator_txn_test.go | 25 ++++++++++ 8 files changed, 243 insertions(+), 19 deletions(-) diff --git a/adapter/internal.go b/adapter/internal.go index 65bd16d95..85b7caf29 100644 --- a/adapter/internal.go +++ b/adapter/internal.go @@ -181,6 +181,9 @@ func (i *Internal) ImportRangeVersions(ctx context.Context, req *pb.ImportRangeV 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")) } @@ -197,6 +200,16 @@ func (i *Internal) ImportRangeVersions(ctx context.Context, req *pb.ImportRangeV 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) verifyInternalLeader(ctx context.Context) error { if i.leader.State() != raftengine.StateLeader { return errors.WithStack(ErrNotLeader) @@ -282,15 +295,24 @@ func migrationFamilyRequiresDecodedS3(family uint32) bool { } 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)) @@ -438,11 +460,11 @@ func (i *Internal) rejectWriteTimestampFloorMutations(muts []*pb.Mutation, commi } routes := i.routeEngine.Stats() for _, mut := range muts { - if mut == nil || len(mut.Key) == 0 { + if mut == nil || (len(mut.Key) == 0 && mut.GetOp() != pb.Op_DEL_PREFIX) { continue } for _, route := range routes { - if routeWriteTimestampFloorApplies(route, mut.Key, commitTS) { + if routeWriteTimestampFloorApplies(route, mut, commitTS) { return errors.Wrapf(kv.ErrRouteWriteTimestampTooLow, "key %q commit_ts=%d floor=%d", mut.Key, commitTS, route.MinWriteTSExclusive) } } @@ -450,11 +472,15 @@ func (i *Internal) rejectWriteTimestampFloorMutations(muts []*pb.Mutation, commi return nil } -func routeWriteTimestampFloorApplies(route distribution.Route, key []byte, commitTS uint64) bool { +func routeWriteTimestampFloorApplies(route distribution.Route, mut *pb.Mutation, commitTS uint64) bool { if route.MinWriteTSExclusive == 0 || commitTS > route.MinWriteTSExclusive { return false } - return kv.RouteKeyFilter(route.Start, route.End)(key) + if mut.GetOp() == pb.Op_DEL_PREFIX { + start, end := kv.RoutePrefixRange(mut.GetKey()) + return rangesIntersect(route.Start, route.End, start, end) + } + return kv.RouteKeyFilter(route.Start, route.End)(mut.GetKey()) } func (i *Internal) stampTxnTimestamps(ctx context.Context, reqs []*pb.Request) error { diff --git a/adapter/internal_migration_test.go b/adapter/internal_migration_test.go index 2a403c963..05e74e76a 100644 --- a/adapter/internal_migration_test.go +++ b/adapter/internal_migration_test.go @@ -277,6 +277,33 @@ func TestInternalExportRangeVersionsIncludesS3BucketAuxiliaryForBucketRouteInter } } +func TestInternalExportRangeVersionsPreservesS3BucketRawRouteMatches(t *testing.T) { + t.Parallel() + + ctx := context.Background() + st := store.NewMVCCStore() + internal := NewInternalWithEngine(nil, mockInternalLeader{}, nil, nil, WithInternalStore(st)) + + key := s3keys.BucketMetaKey("bucket-raw") + require.NoError(t, st.PutAt(ctx, key, []byte("meta"), 10, 0)) + + stream := &captureExportRangeVersionsStream{ctx: ctx} + err := internal.ExportRangeVersions(&pb.ExportRangeVersionsRequest{ + MaxCommitTs: 20, + RouteStart: []byte("!s3|"), + RouteEnd: nil, + KeyFamily: distribution.MigrationFamilyS3BucketMeta, + RangeStart: []byte(s3keys.BucketMetaPrefix), + RangeEnd: testPrefixScanEnd([]byte(s3keys.BucketMetaPrefix)), + MaxScannedBytes: 1 << 20, + }, stream) + require.NoError(t, err) + require.Len(t, stream.responses, 1) + require.Equal(t, []*pb.MVCCVersion{ + {Key: key, CommitTs: 10, Value: []byte("meta"), KeyFamily: distribution.MigrationFamilyS3BucketMeta}, + }, stream.responses[0].GetVersions()) +} + func TestInternalImportRangeVersionsAppliesStoreBatch(t *testing.T) { t.Parallel() @@ -315,3 +342,25 @@ func TestInternalImportRangeVersionsAppliesStoreBatch(t *testing.T) { require.Equal(t, uint64(30), floor) require.GreaterOrEqual(t, clock.Current(), uint64(30)) } + +func TestInternalImportRangeVersionsRejectsMissingIdentifiers(t *testing.T) { + t.Parallel() + + internal := NewInternalWithEngine(nil, mockInternalLeader{}, nil, nil, + WithInternalMigrationProposer(&applyingMigrationProposer{}), + ) + + _, err := internal.ImportRangeVersions(context.Background(), &pb.ImportRangeVersionsRequest{ + BracketId: 1, + BatchSeq: 1, + }) + require.Error(t, err) + require.Equal(t, codes.InvalidArgument, status.Code(err)) + + _, err = internal.ImportRangeVersions(context.Background(), &pb.ImportRangeVersionsRequest{ + JobId: 1, + BatchSeq: 1, + }) + require.Error(t, err) + require.Equal(t, codes.InvalidArgument, status.Code(err)) +} diff --git a/adapter/internal_test.go b/adapter/internal_test.go index fe4a0c5ea..f7e922505 100644 --- a/adapter/internal_test.go +++ b/adapter/internal_test.go @@ -230,6 +230,41 @@ func TestStampRawTimestampsRejectsPreStampedRouteWriteFloor(t *testing.T) { require.ErrorIs(t, err, kv.ErrRouteWriteTimestampTooLow) } +func TestStampRawTimestampsRejectsPreStampedDelPrefixRouteWriteFloor(t *testing.T) { + t.Parallel() + + engine := distribution.NewEngine() + require.NoError(t, engine.ApplySnapshot(distribution.CatalogSnapshot{ + Version: 1, + Routes: []distribution.RouteDescriptor{ + { + RouteID: 1, + Start: []byte(""), + End: []byte("m"), + GroupID: 1, + State: distribution.RouteStateActive, + MinWriteTSExclusive: 0, + }, + { + RouteID: 2, + Start: []byte("m"), + End: nil, + GroupID: 2, + State: distribution.RouteStateActive, + MinWriteTSExclusive: 100, + }, + }, + })) + i := &Internal{routeEngine: engine} + reqs := []*pb.Request{{ + Ts: 100, + Mutations: []*pb.Mutation{{Op: pb.Op_DEL_PREFIX, Key: nil}}, + }} + + err := i.stampRawTimestamps(context.Background(), reqs) + require.ErrorIs(t, err, kv.ErrRouteWriteTimestampTooLow) +} + func TestStampTxnTimestampsRejectsRouteWriteFloor(t *testing.T) { t.Parallel() diff --git a/kv/fsm.go b/kv/fsm.go index 181ae515a..771f359c5 100644 --- a/kv/fsm.go +++ b/kv/fsm.go @@ -624,6 +624,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 routeKeyspaceWideRawPrefix(prefix []byte) bool { if !rawPrefixMayContainRouteMappedKeys(prefix) { return false 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/sharded_coordinator.go b/kv/sharded_coordinator.go index edbac1140..d943302ad 100644 --- a/kv/sharded_coordinator.go +++ b/kv/sharded_coordinator.go @@ -1947,7 +1947,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 } @@ -1955,21 +1955,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 @@ -2001,7 +2028,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-"+ @@ -2051,6 +2078,12 @@ 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 + } route, ok := c.engine.GetRoute(routeKey(key)) if !ok || route.GroupID != gid || !routeHasStagedVisibility(route) { return nil, false @@ -2308,17 +2341,10 @@ func (c *ShardedCoordinator) groupMutations(reqs []*Elem[OP], label keyviz.Label return nil, nil, ErrInvalidRequest } mut := elemToMutation(req) - gid, ok := c.router.ResolveGroup(mut.Key) + gid, routeID, ok := c.resolveGroupAndRouteForKey(mut.Key) if !ok { return nil, nil, errors.Wrapf(ErrInvalidRequest, "no route for key %q", 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 - } 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 c6aa9192f..26408e36b 100644 --- a/kv/sharded_coordinator_del_prefix_test.go +++ b/kv/sharded_coordinator_del_prefix_test.go @@ -227,6 +227,45 @@ 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 TestShardedCoordinatorRejectsS3BucketAuxiliaryPointWriteAtMigrationTimestampFloor(t *testing.T) { t.Parallel() diff --git a/kv/sharded_coordinator_txn_test.go b/kv/sharded_coordinator_txn_test.go index e40318e00..ee8cce672 100644 --- a/kv/sharded_coordinator_txn_test.go +++ b/kv/sharded_coordinator_txn_test.go @@ -8,6 +8,7 @@ import ( "github.com/bootjp/elastickv/distribution" "github.com/bootjp/elastickv/internal/raftengine" + "github.com/bootjp/elastickv/internal/s3keys" pb "github.com/bootjp/elastickv/proto" "github.com/bootjp/elastickv/store" "github.com/stretchr/testify/require" @@ -547,6 +548,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 // --------------------------------------------------------------------------- From 19a28e908e5ebcff0b8c04172f5bd0a3297a4ff5 Mon Sep 17 00:00:00 2001 From: bootjp Date: Tue, 14 Jul 2026 20:00:53 +0900 Subject: [PATCH 13/13] kv: keep floor checks out of raft replay --- kv/shard_store.go | 12 ------------ kv/shard_store_test.go | 16 ++++++++++++++++ 2 files changed, 16 insertions(+), 12 deletions(-) diff --git a/kv/shard_store.go b/kv/shard_store.go index 107be5a49..750e4b08b 100644 --- a/kv/shard_store.go +++ b/kv/shard_store.go @@ -2409,9 +2409,6 @@ 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)) @@ -2425,9 +2422,6 @@ 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)) @@ -2579,9 +2573,6 @@ func (s *ShardStore) ensurePrefixWriteTimestampFloors(prefix []byte, commitTS ui // 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 @@ -2608,9 +2599,6 @@ 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 diff --git a/kv/shard_store_test.go b/kv/shard_store_test.go index c4d2615d0..99a82eda0 100644 --- a/kv/shard_store_test.go +++ b/kv/shard_store_test.go @@ -548,6 +548,22 @@ func TestShardStoreRejectsWritesAtMigrationTimestampFloor(t *testing.T) { require.NoError(t, st.PutAt(ctx, []byte("k"), []byte("ok"), 101, 0)) } +func TestShardStoreRaftApplySkipsMigrationTimestampFloor(t *testing.T) { + t.Parallel() + + ctx := context.Background() + st, _ := newStagedVisibilityShardStore(t) + + require.NoError(t, st.ApplyMutationsRaft(ctx, []*store.KVPairMutation{ + {Op: store.OpTypePut, Key: []byte("k-raft"), Value: []byte("v")}, + }, nil, 90, 100)) + require.NoError(t, st.ApplyMutationsRaftAt(ctx, []*store.KVPairMutation{ + {Op: store.OpTypePut, Key: []byte("k-raft-at"), Value: []byte("v")}, + }, nil, 90, 100, 1)) + require.NoError(t, st.DeletePrefixAtRaft(ctx, []byte("k-raft"), nil, 100)) + require.NoError(t, st.DeletePrefixAtRaftAt(ctx, []byte("k-raft-at"), nil, 100, 2)) +} + func TestShardStoreScanAt_IncludesListKeysAcrossShards(t *testing.T) { t.Parallel()