From b4f247751bd017bd0381256db704a7f92a97fcfa Mon Sep 17 00:00:00 2001 From: bootjp Date: Tue, 14 Jul 2026 00:44:05 +0900 Subject: [PATCH 01/10] migration: promote staged versions --- adapter/internal.go | 70 ++++- adapter/internal_migration_test.go | 35 +++ kv/fsm.go | 6 + kv/fsm_migration_promote.go | 88 +++++++ kv/fsm_migration_promote_test.go | 69 +++++ proto/internal.pb.go | 211 +++++++++++++-- proto/internal.proto | 16 ++ proto/internal_grpc.pb.go | 46 +++- store/lsm_migration.go | 10 +- store/migration_promote.go | 397 +++++++++++++++++++++++++++++ store/migration_versions.go | 11 +- store/migration_versions_test.go | 98 +++++++ store/mvcc_store.go | 22 +- store/store.go | 43 ++++ 14 files changed, 1074 insertions(+), 48 deletions(-) create mode 100644 kv/fsm_migration_promote.go create mode 100644 kv/fsm_migration_promote_test.go create mode 100644 store/migration_promote.go diff --git a/adapter/internal.go b/adapter/internal.go index de61f5485..efacdab71 100644 --- a/adapter/internal.go +++ b/adapter/internal.go @@ -179,6 +179,34 @@ func (i *Internal) ImportRangeVersions(ctx context.Context, req *pb.ImportRangeV return &pb.ImportRangeVersionsResponse{AckedCursor: result.AckedCursor}, nil } +func (i *Internal) PromoteStagedVersions(ctx context.Context, req *pb.PromoteStagedVersionsRequest) (*pb.PromoteStagedVersionsResponse, error) { + if req == nil { + return nil, errors.WithStack(status.Error(codes.InvalidArgument, "promote staged versions request is nil")) + } + if req.GetJobId() == 0 { + return nil, errors.WithStack(status.Error(codes.InvalidArgument, "promote staged versions job_id is required")) + } + if i.migrationProposer == nil { + return nil, errors.WithStack(status.Error(codes.FailedPrecondition, "migration promote proposer is not configured")) + } + if err := i.verifyInternalLeader(ctx); err != nil { + return nil, err + } + result, err := i.proposeMigrationPromote(ctx, req) + if err != nil { + return nil, errors.WithStack(err) + } + if i.clock != nil && result.MaxPromotedTS > 0 { + i.clock.Observe(result.MaxPromotedTS) + } + return &pb.PromoteStagedVersionsResponse{ + NextCursor: result.NextCursor, + Done: result.Done, + PromotedRows: result.PromotedRows, + MaxPromotedTs: result.MaxPromotedTS, + }, nil +} + func (i *Internal) verifyInternalLeader(ctx context.Context) error { if i.leader.State() != raftengine.StateLeader { return errors.WithStack(ErrNotLeader) @@ -194,14 +222,11 @@ func (i *Internal) proposeMigrationImport(ctx context.Context, req *pb.ImportRan if err != nil { return store.ImportVersionsResult{}, errors.WithStack(err) } - result, err := i.migrationProposer.Propose(ctx, cmd) + resp, err := i.proposeMigrationCommand(ctx, cmd, "migration import") 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) { + switch resp := resp.(type) { case store.ImportVersionsResult: return resp, nil case *store.ImportVersionsResult: @@ -216,6 +241,41 @@ func (i *Internal) proposeMigrationImport(ctx context.Context, req *pb.ImportRan } } +func (i *Internal) proposeMigrationPromote(ctx context.Context, req *pb.PromoteStagedVersionsRequest) (store.PromoteVersionsResult, error) { + cmd, err := kv.MarshalMigrationPromoteCommand(req) + if err != nil { + return store.PromoteVersionsResult{}, errors.WithStack(err) + } + resp, err := i.proposeMigrationCommand(ctx, cmd, "migration promote") + if err != nil { + return store.PromoteVersionsResult{}, errors.WithStack(err) + } + switch resp := resp.(type) { + case store.PromoteVersionsResult: + return resp, nil + case *store.PromoteVersionsResult: + if resp == nil { + return store.PromoteVersionsResult{}, errors.New("migration promote apply returned nil result") + } + return *resp, nil + case error: + return store.PromoteVersionsResult{}, errors.WithStack(resp) + default: + return store.PromoteVersionsResult{}, errors.WithStack(errors.Newf("unexpected migration promote apply response type %T", resp)) + } +} + +func (i *Internal) proposeMigrationCommand(ctx context.Context, cmd []byte, label string) (any, error) { + result, err := i.migrationProposer.Propose(ctx, cmd) + if err != nil { + return nil, errors.WithStack(err) + } + if result == nil { + return nil, errors.WithStack(errors.Newf("%s proposal returned nil result", label)) + } + return result.Response, nil +} + func exportRangeVersionsOptions(req *pb.ExportRangeVersionsRequest) store.ExportVersionsOptions { chunkBytes := uint64(req.GetChunkBytes()) if chunkBytes == 0 { diff --git a/adapter/internal_migration_test.go b/adapter/internal_migration_test.go index dd8b48baf..54227d418 100644 --- a/adapter/internal_migration_test.go +++ b/adapter/internal_migration_test.go @@ -152,3 +152,38 @@ func TestInternalImportRangeVersionsAppliesStoreBatch(t *testing.T) { require.Equal(t, uint64(30), floor) require.GreaterOrEqual(t, clock.Current(), uint64(30)) } + +func TestInternalPromoteStagedVersionsAppliesStoreBatch(t *testing.T) { + t.Parallel() + + ctx := context.Background() + st := store.NewMVCCStore() + clock := kv.NewHLC() + proposer := &applyingMigrationProposer{ + fsm: kv.NewKvFSMWithHLC(st, clock), + } + internal := NewInternalWithEngine(nil, mockInternalLeader{}, clock, nil, + WithInternalStore(st), + WithInternalMigrationProposer(proposer), + ) + + staged := distribution.MigrationStagedDataKey(7, []byte("k")) + require.NoError(t, st.PutAt(ctx, staged, []byte("v"), 30, 0)) + + resp, err := internal.PromoteStagedVersions(ctx, &pb.PromoteStagedVersionsRequest{ + JobId: 7, + MaxVersions: 10, + }) + require.NoError(t, err) + require.True(t, resp.GetDone()) + require.Equal(t, uint64(1), resp.GetPromotedRows()) + require.Equal(t, uint64(30), resp.GetMaxPromotedTs()) + require.Equal(t, uint64(1), proposer.calls) + + got, err := st.GetAt(ctx, []byte("k"), 30) + require.NoError(t, err) + require.Equal(t, []byte("v"), got) + _, err = st.GetAt(ctx, staged, 30) + require.ErrorIs(t, err, store.ErrKeyNotFound) + require.GreaterOrEqual(t, clock.Current(), uint64(30)) +} diff --git a/kv/fsm.go b/kv/fsm.go index 00e4a5d3c..269b13954 100644 --- a/kv/fsm.go +++ b/kv/fsm.go @@ -364,6 +364,8 @@ func (f *kvFSM) applyReservedOpcode(ctx context.Context, data []byte) (any, bool return f.applyHLCLease(data[1:]), true case data[0] == raftEncodeMigrationImport: return f.applyMigrationImport(ctx, data[1:]), true + case data[0] == raftEncodeMigrationPromote: + return f.applyMigrationPromote(ctx, data[1:]), true case data[0] >= fsmwire.OpEncryptionMin && data[0] <= fsmwire.OpEncryptionMax: return f.applyEncryption(f.pendingApplyIdx, data[0], data[1:]), true default: @@ -399,6 +401,10 @@ const ( // batch. Every target voter applies the raw MVCC versions, import ack, and // migration HLC floor before the RPC handler returns success. raftEncodeMigrationImport byte = 0x09 + // raftEncodeMigrationPromote carries a target-group range-migration staged + // data promotion chunk. Every target voter atomically copies staged MVCC + // versions into the live keyspace and removes the promoted staged rows. + raftEncodeMigrationPromote byte = 0x0b ) func decodeRaftRequests(data []byte) ([]*pb.Request, error) { diff --git a/kv/fsm_migration_promote.go b/kv/fsm_migration_promote.go new file mode 100644 index 000000000..0d9d1ed86 --- /dev/null +++ b/kv/fsm_migration_promote.go @@ -0,0 +1,88 @@ +package kv + +import ( + "context" + + "github.com/bootjp/elastickv/distribution" + pb "github.com/bootjp/elastickv/proto" + "github.com/bootjp/elastickv/store" + "github.com/cockroachdb/errors" + "google.golang.org/protobuf/proto" +) + +const ( + defaultMigrationPromoteMaxVersions = 1024 + defaultMigrationPromoteMaxBytes = 4 << 20 + defaultMigrationPromoteMaxScannedBytes = defaultMigrationPromoteMaxBytes * 4 +) + +// MarshalMigrationPromoteCommand encodes a target-group staged-data promotion +// chunk as a Raft FSM command. +func MarshalMigrationPromoteCommand(req *pb.PromoteStagedVersionsRequest) ([]byte, error) { + if req == nil { + return nil, errors.WithStack(ErrInvalidRequest) + } + b, err := proto.Marshal(req) + if err != nil { + return nil, errors.WithStack(err) + } + if len(b) >= maxMarshaledCommandSize { + return nil, errors.New("marshaled migration promote request too large") + } + return prependByte(raftEncodeMigrationPromote, b), nil +} + +func (f *kvFSM) applyMigrationPromote(ctx context.Context, data []byte) any { + req := &pb.PromoteStagedVersionsRequest{} + if err := proto.Unmarshal(data, req); err != nil { + return errors.WithStack(err) + } + promoter, ok := f.store.(store.MigrationPromoter) + if !ok { + return errors.WithStack(store.ErrNotSupported) + } + result, err := promoter.PromoteVersions(ctx, migrationPromoteOptionsFromProto(req)) + if err != nil { + return errors.WithStack(err) + } + if f.hlc != nil && result.MaxPromotedTS > 0 { + f.hlc.Observe(result.MaxPromotedTS) + } + return result +} + +func migrationPromoteOptionsFromProto(req *pb.PromoteStagedVersionsRequest) store.PromoteVersionsOptions { + maxVersions := int(req.GetMaxVersions()) + if maxVersions <= 0 { + maxVersions = defaultMigrationPromoteMaxVersions + } + maxBytes := req.GetMaxBytes() + if maxBytes == 0 { + maxBytes = defaultMigrationPromoteMaxBytes + } + maxScannedBytes := req.GetMaxScannedBytes() + if maxScannedBytes == 0 { + maxScannedBytes = defaultMigrationPromoteMaxScannedBytes + } + prefix := distribution.MigrationStagedDataKeyPrefix(req.GetJobId()) + return store.PromoteVersionsOptions{ + JobID: req.GetJobId(), + StartKey: prefix, + EndKey: store.PrefixScanEnd(prefix), + Cursor: req.GetCursor(), + MaxVersions: maxVersions, + MaxBytes: maxBytes, + MaxScannedBytes: maxScannedBytes, + TargetKey: migrationPromoteTargetKey(req.GetJobId()), + } +} + +func migrationPromoteTargetKey(jobID uint64) func([]byte) ([]byte, bool) { + return func(stagedKey []byte) ([]byte, bool) { + gotJobID, rawKey, ok := distribution.MigrationStagedDataKeyParts(stagedKey) + if !ok || gotJobID != jobID { + return nil, false + } + return rawKey, true + } +} diff --git a/kv/fsm_migration_promote_test.go b/kv/fsm_migration_promote_test.go new file mode 100644 index 000000000..2470c485d --- /dev/null +++ b/kv/fsm_migration_promote_test.go @@ -0,0 +1,69 @@ +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" +) + +func TestMigrationPromoteTargetKeyRestoresRawKey(t *testing.T) { + t.Parallel() + + targetKey := migrationPromoteTargetKey(9) + raw, ok := targetKey(distribution.MigrationStagedDataKey(9, []byte("user|k"))) + require.True(t, ok) + require.Equal(t, []byte("user|k"), raw) + + _, ok = targetKey(distribution.MigrationStagedDataKey(10, []byte("user|k"))) + require.False(t, ok) +} + +func TestApplyMigrationPromoteMovesStagedVersions(t *testing.T) { + t.Parallel() + + ctx := context.Background() + st := store.NewMVCCStore() + hlc := NewHLC() + fsm := &kvFSM{store: st, hlc: hlc} + staged := distribution.MigrationStagedDataKey(9, []byte("user|k")) + require.NoError(t, st.PutAt(ctx, staged, []byte("v10"), 10, 0)) + require.NoError(t, st.DeleteAt(ctx, staged, 20)) + + cmd, err := MarshalMigrationPromoteCommand(&pb.PromoteStagedVersionsRequest{ + JobId: 9, + MaxVersions: 10, + }) + require.NoError(t, err) + applied := fsm.Apply(cmd) + result, ok := applied.(store.PromoteVersionsResult) + require.True(t, ok, "got %T: %v", applied, applied) + require.True(t, result.Done) + require.Equal(t, uint64(2), result.PromotedRows) + require.Equal(t, uint64(2), result.TotalPromotedRows) + require.Equal(t, uint64(20), result.MaxPromotedTS) + require.GreaterOrEqual(t, hlc.Current(), uint64(20)) + stateReader, ok := st.(store.MigrationPromotionStateReader) + require.True(t, ok) + state, ok, err := stateReader.MigrationPromotionState(ctx, 9) + require.NoError(t, err) + require.True(t, ok) + require.True(t, state.Done) + require.Equal(t, uint64(2), state.PromotedRows) + + got, err := st.GetAt(ctx, []byte("user|k"), 10) + require.NoError(t, err) + require.Equal(t, []byte("v10"), got) + _, err = st.GetAt(ctx, []byte("user|k"), 20) + require.ErrorIs(t, err, store.ErrKeyNotFound) + left, err := st.ExportVersions(ctx, store.ExportVersionsOptions{ + StartKey: distribution.MigrationStagedDataKeyPrefix(9), + EndKey: store.PrefixScanEnd(distribution.MigrationStagedDataKeyPrefix(9)), + MaxVersions: 10, + }) + require.NoError(t, err) + require.Empty(t, left.Versions) +} diff --git a/proto/internal.pb.go b/proto/internal.pb.go index a7a9e30ce..91bcff36b 100644 --- a/proto/internal.pb.go +++ b/proto/internal.pb.go @@ -931,6 +931,150 @@ func (x *ImportRangeVersionsResponse) GetAckedCursor() []byte { return nil } +type PromoteStagedVersionsRequest struct { + state protoimpl.MessageState `protogen:"open.v1"` + JobId uint64 `protobuf:"varint,1,opt,name=job_id,json=jobId,proto3" json:"job_id,omitempty"` + Cursor []byte `protobuf:"bytes,2,opt,name=cursor,proto3" json:"cursor,omitempty"` + MaxVersions uint32 `protobuf:"varint,3,opt,name=max_versions,json=maxVersions,proto3" json:"max_versions,omitempty"` + MaxBytes uint64 `protobuf:"varint,4,opt,name=max_bytes,json=maxBytes,proto3" json:"max_bytes,omitempty"` + MaxScannedBytes uint64 `protobuf:"varint,5,opt,name=max_scanned_bytes,json=maxScannedBytes,proto3" json:"max_scanned_bytes,omitempty"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *PromoteStagedVersionsRequest) Reset() { + *x = PromoteStagedVersionsRequest{} + mi := &file_internal_proto_msgTypes[12] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *PromoteStagedVersionsRequest) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*PromoteStagedVersionsRequest) ProtoMessage() {} + +func (x *PromoteStagedVersionsRequest) ProtoReflect() protoreflect.Message { + mi := &file_internal_proto_msgTypes[12] + if x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use PromoteStagedVersionsRequest.ProtoReflect.Descriptor instead. +func (*PromoteStagedVersionsRequest) Descriptor() ([]byte, []int) { + return file_internal_proto_rawDescGZIP(), []int{12} +} + +func (x *PromoteStagedVersionsRequest) GetJobId() uint64 { + if x != nil { + return x.JobId + } + return 0 +} + +func (x *PromoteStagedVersionsRequest) GetCursor() []byte { + if x != nil { + return x.Cursor + } + return nil +} + +func (x *PromoteStagedVersionsRequest) GetMaxVersions() uint32 { + if x != nil { + return x.MaxVersions + } + return 0 +} + +func (x *PromoteStagedVersionsRequest) GetMaxBytes() uint64 { + if x != nil { + return x.MaxBytes + } + return 0 +} + +func (x *PromoteStagedVersionsRequest) GetMaxScannedBytes() uint64 { + if x != nil { + return x.MaxScannedBytes + } + return 0 +} + +type PromoteStagedVersionsResponse struct { + state protoimpl.MessageState `protogen:"open.v1"` + NextCursor []byte `protobuf:"bytes,1,opt,name=next_cursor,json=nextCursor,proto3" json:"next_cursor,omitempty"` + Done bool `protobuf:"varint,2,opt,name=done,proto3" json:"done,omitempty"` + PromotedRows uint64 `protobuf:"varint,3,opt,name=promoted_rows,json=promotedRows,proto3" json:"promoted_rows,omitempty"` + MaxPromotedTs uint64 `protobuf:"varint,4,opt,name=max_promoted_ts,json=maxPromotedTs,proto3" json:"max_promoted_ts,omitempty"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *PromoteStagedVersionsResponse) Reset() { + *x = PromoteStagedVersionsResponse{} + mi := &file_internal_proto_msgTypes[13] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *PromoteStagedVersionsResponse) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*PromoteStagedVersionsResponse) ProtoMessage() {} + +func (x *PromoteStagedVersionsResponse) ProtoReflect() protoreflect.Message { + mi := &file_internal_proto_msgTypes[13] + if x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use PromoteStagedVersionsResponse.ProtoReflect.Descriptor instead. +func (*PromoteStagedVersionsResponse) Descriptor() ([]byte, []int) { + return file_internal_proto_rawDescGZIP(), []int{13} +} + +func (x *PromoteStagedVersionsResponse) GetNextCursor() []byte { + if x != nil { + return x.NextCursor + } + return nil +} + +func (x *PromoteStagedVersionsResponse) GetDone() bool { + if x != nil { + return x.Done + } + return false +} + +func (x *PromoteStagedVersionsResponse) GetPromotedRows() uint64 { + if x != nil { + return x.PromotedRows + } + return 0 +} + +func (x *PromoteStagedVersionsResponse) GetMaxPromotedTs() uint64 { + if x != nil { + return x.MaxPromotedTs + } + return 0 +} + var File_internal_proto protoreflect.FileDescriptor const file_internal_proto_rawDesc = "" + @@ -999,7 +1143,19 @@ const file_internal_proto_rawDesc = "" + "bracket_id\x18\x04 \x01(\x04R\tbracketId\x12\x1b\n" + "\tbatch_seq\x18\x05 \x01(\x04R\bbatchSeq\"@\n" + "\x1bImportRangeVersionsResponse\x12!\n" + - "\facked_cursor\x18\x01 \x01(\fR\vackedCursor*&\n" + + "\facked_cursor\x18\x01 \x01(\fR\vackedCursor\"\xb9\x01\n" + + "\x1cPromoteStagedVersionsRequest\x12\x15\n" + + "\x06job_id\x18\x01 \x01(\x04R\x05jobId\x12\x16\n" + + "\x06cursor\x18\x02 \x01(\fR\x06cursor\x12!\n" + + "\fmax_versions\x18\x03 \x01(\rR\vmaxVersions\x12\x1b\n" + + "\tmax_bytes\x18\x04 \x01(\x04R\bmaxBytes\x12*\n" + + "\x11max_scanned_bytes\x18\x05 \x01(\x04R\x0fmaxScannedBytes\"\xa1\x01\n" + + "\x1dPromoteStagedVersionsResponse\x12\x1f\n" + + "\vnext_cursor\x18\x01 \x01(\fR\n" + + "nextCursor\x12\x12\n" + + "\x04done\x18\x02 \x01(\bR\x04done\x12#\n" + + "\rpromoted_rows\x18\x03 \x01(\x04R\fpromotedRows\x12&\n" + + "\x0fmax_promoted_ts\x18\x04 \x01(\x04R\rmaxPromotedTs*&\n" + "\x02Op\x12\a\n" + "\x03PUT\x10\x00\x12\a\n" + "\x03DEL\x10\x01\x12\x0e\n" + @@ -1010,12 +1166,13 @@ const file_internal_proto_rawDesc = "" + "\aPREPARE\x10\x01\x12\n" + "\n" + "\x06COMMIT\x10\x02\x12\t\n" + - "\x05ABORT\x10\x032\xa3\x02\n" + + "\x05ABORT\x10\x032\xfd\x02\n" + "\bInternal\x12.\n" + "\aForward\x12\x0f.ForwardRequest\x1a\x10.ForwardResponse\"\x00\x12=\n" + "\fRelayPublish\x12\x14.RelayPublishRequest\x1a\x15.RelayPublishResponse\"\x00\x12T\n" + "\x13ExportRangeVersions\x12\x1b.ExportRangeVersionsRequest\x1a\x1c.ExportRangeVersionsResponse\"\x000\x01\x12R\n" + - "\x13ImportRangeVersions\x12\x1b.ImportRangeVersionsRequest\x1a\x1c.ImportRangeVersionsResponse\"\x00B#Z!github.com/bootjp/elastickv/protob\x06proto3" + "\x13ImportRangeVersions\x12\x1b.ImportRangeVersionsRequest\x1a\x1c.ImportRangeVersionsResponse\"\x00\x12X\n" + + "\x15PromoteStagedVersions\x12\x1d.PromoteStagedVersionsRequest\x1a\x1e.PromoteStagedVersionsResponse\"\x00B#Z!github.com/bootjp/elastickv/protob\x06proto3" var ( file_internal_proto_rawDescOnce sync.Once @@ -1030,22 +1187,24 @@ func file_internal_proto_rawDescGZIP() []byte { } var file_internal_proto_enumTypes = make([]protoimpl.EnumInfo, 2) -var file_internal_proto_msgTypes = make([]protoimpl.MessageInfo, 12) +var file_internal_proto_msgTypes = make([]protoimpl.MessageInfo, 14) var file_internal_proto_goTypes = []any{ - (Op)(0), // 0: Op - (Phase)(0), // 1: Phase - (*Mutation)(nil), // 2: Mutation - (*Request)(nil), // 3: Request - (*RaftCommand)(nil), // 4: RaftCommand - (*ForwardRequest)(nil), // 5: ForwardRequest - (*ForwardResponse)(nil), // 6: ForwardResponse - (*RelayPublishRequest)(nil), // 7: RelayPublishRequest - (*RelayPublishResponse)(nil), // 8: RelayPublishResponse - (*ExportRangeVersionsRequest)(nil), // 9: ExportRangeVersionsRequest - (*ExportRangeVersionsResponse)(nil), // 10: ExportRangeVersionsResponse - (*MVCCVersion)(nil), // 11: MVCCVersion - (*ImportRangeVersionsRequest)(nil), // 12: ImportRangeVersionsRequest - (*ImportRangeVersionsResponse)(nil), // 13: ImportRangeVersionsResponse + (Op)(0), // 0: Op + (Phase)(0), // 1: Phase + (*Mutation)(nil), // 2: Mutation + (*Request)(nil), // 3: Request + (*RaftCommand)(nil), // 4: RaftCommand + (*ForwardRequest)(nil), // 5: ForwardRequest + (*ForwardResponse)(nil), // 6: ForwardResponse + (*RelayPublishRequest)(nil), // 7: RelayPublishRequest + (*RelayPublishResponse)(nil), // 8: RelayPublishResponse + (*ExportRangeVersionsRequest)(nil), // 9: ExportRangeVersionsRequest + (*ExportRangeVersionsResponse)(nil), // 10: ExportRangeVersionsResponse + (*MVCCVersion)(nil), // 11: MVCCVersion + (*ImportRangeVersionsRequest)(nil), // 12: ImportRangeVersionsRequest + (*ImportRangeVersionsResponse)(nil), // 13: ImportRangeVersionsResponse + (*PromoteStagedVersionsRequest)(nil), // 14: PromoteStagedVersionsRequest + (*PromoteStagedVersionsResponse)(nil), // 15: PromoteStagedVersionsResponse } var file_internal_proto_depIdxs = []int32{ 0, // 0: Mutation.op:type_name -> Op @@ -1059,12 +1218,14 @@ var file_internal_proto_depIdxs = []int32{ 7, // 8: Internal.RelayPublish:input_type -> RelayPublishRequest 9, // 9: Internal.ExportRangeVersions:input_type -> ExportRangeVersionsRequest 12, // 10: Internal.ImportRangeVersions:input_type -> ImportRangeVersionsRequest - 6, // 11: Internal.Forward:output_type -> ForwardResponse - 8, // 12: Internal.RelayPublish:output_type -> RelayPublishResponse - 10, // 13: Internal.ExportRangeVersions:output_type -> ExportRangeVersionsResponse - 13, // 14: Internal.ImportRangeVersions:output_type -> ImportRangeVersionsResponse - 11, // [11:15] is the sub-list for method output_type - 7, // [7:11] is the sub-list for method input_type + 14, // 11: Internal.PromoteStagedVersions:input_type -> PromoteStagedVersionsRequest + 6, // 12: Internal.Forward:output_type -> ForwardResponse + 8, // 13: Internal.RelayPublish:output_type -> RelayPublishResponse + 10, // 14: Internal.ExportRangeVersions:output_type -> ExportRangeVersionsResponse + 13, // 15: Internal.ImportRangeVersions:output_type -> ImportRangeVersionsResponse + 15, // 16: Internal.PromoteStagedVersions:output_type -> PromoteStagedVersionsResponse + 12, // [12:17] is the sub-list for method output_type + 7, // [7:12] is the sub-list for method input_type 7, // [7:7] is the sub-list for extension type_name 7, // [7:7] is the sub-list for extension extendee 0, // [0:7] is the sub-list for field type_name @@ -1081,7 +1242,7 @@ func file_internal_proto_init() { GoPackagePath: reflect.TypeOf(x{}).PkgPath(), RawDescriptor: unsafe.Slice(unsafe.StringData(file_internal_proto_rawDesc), len(file_internal_proto_rawDesc)), NumEnums: 2, - NumMessages: 12, + NumMessages: 14, NumExtensions: 0, NumServices: 1, }, diff --git a/proto/internal.proto b/proto/internal.proto index b60c83793..b45bf8fb5 100644 --- a/proto/internal.proto +++ b/proto/internal.proto @@ -9,6 +9,7 @@ service Internal { rpc RelayPublish(RelayPublishRequest) returns (RelayPublishResponse) {} rpc ExportRangeVersions(ExportRangeVersionsRequest) returns (stream ExportRangeVersionsResponse) {} rpc ImportRangeVersions(ImportRangeVersionsRequest) returns (ImportRangeVersionsResponse) {} + rpc PromoteStagedVersions(PromoteStagedVersionsRequest) returns (PromoteStagedVersionsResponse) {} } // internal.proto is node to node communication message in raft replication. @@ -128,3 +129,18 @@ message ImportRangeVersionsRequest { message ImportRangeVersionsResponse { bytes acked_cursor = 1; } + +message PromoteStagedVersionsRequest { + uint64 job_id = 1; + bytes cursor = 2; + uint32 max_versions = 3; + uint64 max_bytes = 4; + uint64 max_scanned_bytes = 5; +} + +message PromoteStagedVersionsResponse { + bytes next_cursor = 1; + bool done = 2; + uint64 promoted_rows = 3; + uint64 max_promoted_ts = 4; +} diff --git a/proto/internal_grpc.pb.go b/proto/internal_grpc.pb.go index 6a21b9eba..ba679ed80 100644 --- a/proto/internal_grpc.pb.go +++ b/proto/internal_grpc.pb.go @@ -19,10 +19,11 @@ import ( const _ = grpc.SupportPackageIsVersion9 const ( - Internal_Forward_FullMethodName = "/Internal/Forward" - Internal_RelayPublish_FullMethodName = "/Internal/RelayPublish" - Internal_ExportRangeVersions_FullMethodName = "/Internal/ExportRangeVersions" - Internal_ImportRangeVersions_FullMethodName = "/Internal/ImportRangeVersions" + Internal_Forward_FullMethodName = "/Internal/Forward" + Internal_RelayPublish_FullMethodName = "/Internal/RelayPublish" + Internal_ExportRangeVersions_FullMethodName = "/Internal/ExportRangeVersions" + Internal_ImportRangeVersions_FullMethodName = "/Internal/ImportRangeVersions" + Internal_PromoteStagedVersions_FullMethodName = "/Internal/PromoteStagedVersions" ) // InternalClient is the client API for Internal service. @@ -34,6 +35,7 @@ type InternalClient interface { RelayPublish(ctx context.Context, in *RelayPublishRequest, opts ...grpc.CallOption) (*RelayPublishResponse, error) ExportRangeVersions(ctx context.Context, in *ExportRangeVersionsRequest, opts ...grpc.CallOption) (grpc.ServerStreamingClient[ExportRangeVersionsResponse], error) ImportRangeVersions(ctx context.Context, in *ImportRangeVersionsRequest, opts ...grpc.CallOption) (*ImportRangeVersionsResponse, error) + PromoteStagedVersions(ctx context.Context, in *PromoteStagedVersionsRequest, opts ...grpc.CallOption) (*PromoteStagedVersionsResponse, error) } type internalClient struct { @@ -93,6 +95,16 @@ func (c *internalClient) ImportRangeVersions(ctx context.Context, in *ImportRang return out, nil } +func (c *internalClient) PromoteStagedVersions(ctx context.Context, in *PromoteStagedVersionsRequest, opts ...grpc.CallOption) (*PromoteStagedVersionsResponse, error) { + cOpts := append([]grpc.CallOption{grpc.StaticMethod()}, opts...) + out := new(PromoteStagedVersionsResponse) + err := c.cc.Invoke(ctx, Internal_PromoteStagedVersions_FullMethodName, in, out, cOpts...) + if err != nil { + return nil, err + } + return out, nil +} + // InternalServer is the server API for Internal service. // All implementations must embed UnimplementedInternalServer // for forward compatibility. @@ -102,6 +114,7 @@ type InternalServer interface { RelayPublish(context.Context, *RelayPublishRequest) (*RelayPublishResponse, error) ExportRangeVersions(*ExportRangeVersionsRequest, grpc.ServerStreamingServer[ExportRangeVersionsResponse]) error ImportRangeVersions(context.Context, *ImportRangeVersionsRequest) (*ImportRangeVersionsResponse, error) + PromoteStagedVersions(context.Context, *PromoteStagedVersionsRequest) (*PromoteStagedVersionsResponse, error) mustEmbedUnimplementedInternalServer() } @@ -124,6 +137,9 @@ func (UnimplementedInternalServer) ExportRangeVersions(*ExportRangeVersionsReque func (UnimplementedInternalServer) ImportRangeVersions(context.Context, *ImportRangeVersionsRequest) (*ImportRangeVersionsResponse, error) { return nil, status.Error(codes.Unimplemented, "method ImportRangeVersions not implemented") } +func (UnimplementedInternalServer) PromoteStagedVersions(context.Context, *PromoteStagedVersionsRequest) (*PromoteStagedVersionsResponse, error) { + return nil, status.Error(codes.Unimplemented, "method PromoteStagedVersions not implemented") +} func (UnimplementedInternalServer) mustEmbedUnimplementedInternalServer() {} func (UnimplementedInternalServer) testEmbeddedByValue() {} @@ -210,6 +226,24 @@ func _Internal_ImportRangeVersions_Handler(srv interface{}, ctx context.Context, return interceptor(ctx, in, info, handler) } +func _Internal_PromoteStagedVersions_Handler(srv interface{}, ctx context.Context, dec func(interface{}) error, interceptor grpc.UnaryServerInterceptor) (interface{}, error) { + in := new(PromoteStagedVersionsRequest) + if err := dec(in); err != nil { + return nil, err + } + if interceptor == nil { + return srv.(InternalServer).PromoteStagedVersions(ctx, in) + } + info := &grpc.UnaryServerInfo{ + Server: srv, + FullMethod: Internal_PromoteStagedVersions_FullMethodName, + } + handler := func(ctx context.Context, req interface{}) (interface{}, error) { + return srv.(InternalServer).PromoteStagedVersions(ctx, req.(*PromoteStagedVersionsRequest)) + } + return interceptor(ctx, in, info, handler) +} + // Internal_ServiceDesc is the grpc.ServiceDesc for Internal service. // It's only intended for direct use with grpc.RegisterService, // and not to be introspected or modified (even as a copy) @@ -229,6 +263,10 @@ var Internal_ServiceDesc = grpc.ServiceDesc{ MethodName: "ImportRangeVersions", Handler: _Internal_ImportRangeVersions_Handler, }, + { + MethodName: "PromoteStagedVersions", + Handler: _Internal_PromoteStagedVersions_Handler, + }, }, Streams: []grpc.StreamDesc{ { diff --git a/store/lsm_migration.go b/store/lsm_migration.go index 3ae343afd..e0796cbf6 100644 --- a/store/lsm_migration.go +++ b/store/lsm_migration.go @@ -10,6 +10,13 @@ import ( ) func (s *pebbleStore) ExportVersions(ctx context.Context, opts ExportVersionsOptions) (ExportVersionsResult, error) { + s.dbMu.RLock() + defer s.dbMu.RUnlock() + + return s.exportVersionsLocked(ctx, opts) +} + +func (s *pebbleStore) exportVersionsLocked(ctx context.Context, opts ExportVersionsOptions) (ExportVersionsResult, error) { pos, err := decodeExportCursor(opts.Cursor) if err != nil { return ExportVersionsResult{}, err @@ -18,9 +25,6 @@ func (s *pebbleStore) ExportVersions(ctx context.Context, opts ExportVersionsOpt return ExportVersionsResult{Done: true}, nil } - s.dbMu.RLock() - defer s.dbMu.RUnlock() - iter, err := s.db.NewIter(pebbleExportIterOptions(opts)) if err != nil { return ExportVersionsResult{}, errors.WithStack(err) diff --git a/store/migration_promote.go b/store/migration_promote.go new file mode 100644 index 000000000..cde302a1a --- /dev/null +++ b/store/migration_promote.go @@ -0,0 +1,397 @@ +package store + +import ( + "bytes" + "context" + "encoding/binary" + + "github.com/cockroachdb/errors" + "github.com/cockroachdb/pebble/v2" +) + +const migrationPromotionDoneFlag byte = 1 + +type promotedVersion struct { + staged MVCCVersion + target MVCCVersion +} + +func validatePromoteVersionsOptions(opts PromoteVersionsOptions) error { + if opts.TargetKey == nil { + return errors.New("migration promote target key mapper is required") + } + return nil +} + +func promotedVersionsFromStaged(opts PromoteVersionsOptions, versions []MVCCVersion) ([]promotedVersion, PromoteVersionsResult, error) { + out := make([]promotedVersion, 0, len(versions)) + result := PromoteVersionsResult{PromotedRows: uint64(len(versions))} //nolint:gosec // len is bounded by MaxVersions. + for _, staged := range versions { + targetKey, ok := opts.TargetKey(staged.Key) + if !ok { + return nil, PromoteVersionsResult{}, errors.WithStack(errors.Newf("migration promote target key rejected staged key %q", string(staged.Key))) + } + target := MVCCVersion{ + Key: bytes.Clone(targetKey), + CommitTS: staged.CommitTS, + Tombstone: staged.Tombstone, + Value: bytes.Clone(staged.Value), + KeyFamily: staged.KeyFamily, + ExpireAt: staged.ExpireAt, + } + if err := validateImportVersion(target); err != nil { + return nil, PromoteVersionsResult{}, err + } + result.PromotedBytes += versionExportSize(target.Key, len(target.Value)) + if target.CommitTS > result.MaxPromotedTS { + result.MaxPromotedTS = target.CommitTS + } + out = append(out, promotedVersion{ + staged: staged, + target: target, + }) + } + return out, result, nil +} + +func (s *mvccStore) PromoteVersions(ctx context.Context, opts PromoteVersionsOptions) (PromoteVersionsResult, error) { + if err := validatePromoteVersionsOptions(opts); err != nil { + return PromoteVersionsResult{}, err + } + if opts.MaxVersions <= 0 { + return PromoteVersionsResult{Done: true}, nil + } + + s.mtx.Lock() + defer s.mtx.Unlock() + + state, cursor := s.promotionStateAndCursorLocked(opts) + if opts.JobID != 0 && state.Done { + return PromoteVersionsResult{Done: true, TotalPromotedRows: state.PromotedRows}, nil + } + exported, toPromote, promoted, err := s.planMemoryPromotionLocked(ctx, opts, cursor) + if err != nil { + return PromoteVersionsResult{}, err + } + s.applyMemoryPromotionLocked(toPromote) + if promoted.MaxPromotedTS > s.lastCommitTS { + s.lastCommitTS = promoted.MaxPromotedTS + } + result, updatedState := finishPromotionResult(opts, state, exported, promoted) + if updatedState != nil { + s.migrationPromotions[opts.JobID] = clonePromotionState(*updatedState) + } + return result, nil +} + +func (s *mvccStore) planMemoryPromotionLocked( + ctx context.Context, + opts PromoteVersionsOptions, + cursor []byte, +) (ExportVersionsResult, []promotedVersion, PromoteVersionsResult, error) { + pos, err := decodeExportCursor(cursor) + if err != nil { + return ExportVersionsResult{}, nil, PromoteVersionsResult{}, err + } + opts.Cursor = cursor + exported, err := s.exportMemoryVersionsLocked(ctx, opts, pos) + if err != nil { + return ExportVersionsResult{}, nil, PromoteVersionsResult{}, err + } + toPromote, promoted, err := promotedVersionsFromStaged(opts, exported.Versions) + return exported, toPromote, promoted, err +} + +func (s *mvccStore) applyMemoryPromotionLocked(toPromote []promotedVersion) { + for _, version := range toPromote { + if version.target.Tombstone { + s.deleteVersionLocked(version.target.Key, version.target.CommitTS) + } else { + s.putVersionLocked(version.target.Key, version.target.Value, version.target.CommitTS, version.target.ExpireAt) + } + s.removeVersionLocked(version.staged.Key, version.staged.CommitTS) + } +} + +func (s *mvccStore) promotionStateAndCursorLocked(opts PromoteVersionsOptions) (PromotionState, []byte) { + if opts.JobID == 0 { + return PromotionState{}, opts.Cursor + } + state := clonePromotionState(s.migrationPromotions[opts.JobID]) + if len(state.Cursor) == 0 { + return state, opts.Cursor + } + return state, state.Cursor +} + +func (s *mvccStore) MigrationPromotionState(_ context.Context, jobID uint64) (PromotionState, bool, error) { + s.mtx.RLock() + defer s.mtx.RUnlock() + state, ok := s.migrationPromotions[jobID] + return clonePromotionState(state), ok, nil +} + +func (s *mvccStore) exportMemoryVersionsLocked(ctx context.Context, opts PromoteVersionsOptions, pos exportCursorPosition) (ExportVersionsResult, error) { + exportOpts := ExportVersionsOptions{ + StartKey: opts.StartKey, + EndKey: opts.EndKey, + Cursor: opts.Cursor, + MaxVersions: opts.MaxVersions, + MaxBytes: opts.MaxBytes, + MaxScannedBytes: opts.MaxScannedBytes, + } + result := newExportVersionsResult(opts.MaxVersions) + it := s.tree.Iterator() + if !s.seekMemoryExportStart(&it, opts.StartKey, pos.key) { + result.Done = true + return result, nil + } + for ok := true; ok; ok = it.Next() { + key, keyOK := it.Key().([]byte) + if err := checkExportKey(ctx, key, keyOK, opts.EndKey); err != nil { + if errors.Is(err, errExportReachedEnd) { + result.Done = true + result.NextCursor = nil + return result, nil + } + return ExportVersionsResult{}, err + } + if !keyOK { + continue + } + done, err := exportMemoryIteratorKey(ctx, exportOpts, pos, key, it.Value(), &result) + if err != nil || !done { + return result, err + } + } + result.Done = true + result.NextCursor = nil + return result, nil +} + +func (s *mvccStore) removeVersionLocked(key []byte, commitTS uint64) bool { + existing, ok := s.tree.Get(key) + if !ok { + return false + } + versions, _ := existing.([]VersionedValue) + idx := findVersionIndex(versions, commitTS) + if idx < 0 { + return false + } + next := make([]VersionedValue, len(versions)-1) + copy(next, versions[:idx]) + copy(next[idx:], versions[idx+1:]) + if len(next) == 0 { + s.tree.Remove(key) + return true + } + s.tree.Put(bytes.Clone(key), next) + return true +} + +func findVersionIndex(versions []VersionedValue, commitTS uint64) int { + for i := range versions { + if versions[i].TS == commitTS { + return i + } + } + return -1 +} + +func (s *pebbleStore) PromoteVersions(ctx context.Context, opts PromoteVersionsOptions) (PromoteVersionsResult, error) { + if err := validatePromoteVersionsOptions(opts); err != nil { + return PromoteVersionsResult{}, err + } + if opts.MaxVersions <= 0 { + return PromoteVersionsResult{Done: true}, nil + } + + s.dbMu.RLock() + defer s.dbMu.RUnlock() + + s.applyMu.Lock() + defer s.applyMu.Unlock() + + state, cursor, err := s.pebblePromotionStateAndCursor(opts) + if err != nil { + return PromoteVersionsResult{}, err + } + if opts.JobID != 0 && state.Done { + return PromoteVersionsResult{Done: true, TotalPromotedRows: state.PromotedRows}, nil + } + opts.Cursor = cursor + exported, toPromote, promoted, err := s.planPebblePromotionLocked(ctx, opts) + if err != nil { + return PromoteVersionsResult{}, err + } + result, stateToWrite := finishPromotionResult(opts, state, exported, promoted) + return s.finishPebblePromotion(toPromote, opts.JobID, stateToWrite, result) +} + +func (s *pebbleStore) finishPebblePromotion( + toPromote []promotedVersion, + jobID uint64, + stateToWrite *PromotionState, + result PromoteVersionsResult, +) (PromoteVersionsResult, error) { + if len(toPromote) == 0 && stateToWrite == nil { + return result, nil + } + if err := s.commitPebblePromoteVersions(toPromote, jobID, stateToWrite); err != nil { + return PromoteVersionsResult{}, err + } + return result, nil +} + +func (s *pebbleStore) planPebblePromotionLocked( + ctx context.Context, + opts PromoteVersionsOptions, +) (ExportVersionsResult, []promotedVersion, PromoteVersionsResult, error) { + exported, err := s.exportVersionsLocked(ctx, ExportVersionsOptions{ + StartKey: opts.StartKey, + EndKey: opts.EndKey, + Cursor: opts.Cursor, + MaxVersions: opts.MaxVersions, + MaxBytes: opts.MaxBytes, + MaxScannedBytes: opts.MaxScannedBytes, + }) + if err != nil { + return ExportVersionsResult{}, nil, PromoteVersionsResult{}, err + } + toPromote, promoted, err := promotedVersionsFromStaged(opts, exported.Versions) + return exported, toPromote, promoted, err +} + +func finishPromotionResult( + opts PromoteVersionsOptions, + state PromotionState, + exported ExportVersionsResult, + promoted PromoteVersionsResult, +) (PromoteVersionsResult, *PromotionState) { + promoted.NextCursor = exported.NextCursor + promoted.Done = exported.Done + promoted.ScannedBytes = exported.ScannedBytes + promoted.TotalPromotedRows = promoted.PromotedRows + if opts.JobID == 0 { + return promoted, nil + } + state.Cursor = bytes.Clone(exported.NextCursor) + state.Done = exported.Done + state.PromotedRows += promoted.PromotedRows + state.LastError = "" + promoted.TotalPromotedRows = state.PromotedRows + return promoted, &state +} + +func (s *pebbleStore) pebblePromotionStateAndCursor(opts PromoteVersionsOptions) (PromotionState, []byte, error) { + if opts.JobID == 0 { + return PromotionState{}, opts.Cursor, nil + } + state, ok, err := s.readPebblePromotionState(opts.JobID) + if err != nil { + return PromotionState{}, nil, err + } + if !ok || len(state.Cursor) == 0 { + return state, opts.Cursor, nil + } + return state, state.Cursor, nil +} + +func (s *pebbleStore) MigrationPromotionState(_ context.Context, jobID uint64) (PromotionState, bool, error) { + s.dbMu.RLock() + defer s.dbMu.RUnlock() + return s.readPebblePromotionState(jobID) +} + +func (s *pebbleStore) readPebblePromotionState(jobID uint64) (PromotionState, bool, error) { + val, closer, err := s.db.Get(migrationPromoteKey(jobID)) + if err != nil { + if errors.Is(err, pebble.ErrNotFound) { + return PromotionState{}, false, nil + } + return PromotionState{}, false, errors.WithStack(err) + } + defer func() { _ = closer.Close() }() + state, ok := decodePromotionState(val) + if !ok { + return PromotionState{}, false, errors.New("corrupt migration promotion state") + } + return state, true, nil +} + +func (s *pebbleStore) commitPebblePromoteVersions(versions []promotedVersion, jobID uint64, state *PromotionState) error { + batch := s.db.NewBatch() + defer batch.Close() + targets := make([]MVCCVersion, 0, len(versions)) + for _, version := range versions { + targets = append(targets, version.target) + } + if err := s.applyImportVersionsBatch(batch, targets); err != nil { + return err + } + for _, version := range versions { + if err := batch.Delete(encodeKey(version.staged.Key, version.staged.CommitTS), nil); err != nil { + return errors.WithStack(err) + } + } + if state != nil { + if err := batch.Set(migrationPromoteKey(jobID), encodePromotionState(*state), nil); err != nil { + return errors.WithStack(err) + } + } + if err := batch.Commit(s.directApplyWriteOpts()); err != nil { + return errors.WithStack(err) + } + return nil +} + +func encodePromotionState(state PromotionState) []byte { + buf := make([]byte, 0, 1+migrationUint64Bytes+binary.MaxVarintLen64*2+len(state.Cursor)+len(state.LastError)) + if state.Done { + buf = append(buf, migrationPromotionDoneFlag) + } else { + buf = append(buf, 0) + } + buf = binary.BigEndian.AppendUint64(buf, state.PromotedRows) + buf = binary.AppendUvarint(buf, lenAsUint64(len(state.Cursor))) + buf = append(buf, state.Cursor...) + buf = binary.AppendUvarint(buf, lenAsUint64(len(state.LastError))) + buf = append(buf, state.LastError...) + return buf +} + +func decodePromotionState(data []byte) (PromotionState, bool) { + if len(data) < 1+migrationUint64Bytes { + return PromotionState{}, false + } + state := PromotionState{ + Done: data[0]&migrationPromotionDoneFlag != 0, + PromotedRows: binary.BigEndian.Uint64(data[1 : 1+migrationUint64Bytes]), + } + rest := data[1+migrationUint64Bytes:] + cursorLen, n := binary.Uvarint(rest) + if n <= 0 || cursorLen > lenAsUint64(len(rest[n:])) { + return PromotionState{}, false + } + rest = rest[n:] + cursorEnd := int(cursorLen) //nolint:gosec // bounded by len(rest) above. + state.Cursor = bytes.Clone(rest[:cursorEnd]) + rest = rest[cursorEnd:] + errLen, n := binary.Uvarint(rest) + if n <= 0 || errLen != lenAsUint64(len(rest[n:])) { + return PromotionState{}, false + } + state.LastError = string(rest[n:]) + return state, true +} + +func clonePromotionState(state PromotionState) PromotionState { + state.Cursor = bytes.Clone(state.Cursor) + return state +} + +var _ MigrationPromoter = (*mvccStore)(nil) +var _ MigrationPromoter = (*pebbleStore)(nil) +var _ MigrationPromotionStateReader = (*mvccStore)(nil) +var _ MigrationPromotionStateReader = (*pebbleStore)(nil) diff --git a/store/migration_versions.go b/store/migration_versions.go index 54b7dd19b..28fd2f73d 100644 --- a/store/migration_versions.go +++ b/store/migration_versions.go @@ -15,6 +15,7 @@ const ( migrationAckPrefix = "!migstage|ack|" migrationHLCFloorPrefix = "!migstage|hlc_floor|" + migrationPromotePrefix = "!migstage|promote|" migrationUint64Bytes = 8 migrationAckKeyIDBytes = 2 * migrationUint64Bytes exportVersionSizeOverhead = 24 @@ -84,9 +85,17 @@ func migrationHLCFloorKey(jobID uint64) []byte { return key } +func migrationPromoteKey(jobID uint64) []byte { + key := make([]byte, len(migrationPromotePrefix)+migrationUint64Bytes) + copy(key, migrationPromotePrefix) + binary.BigEndian.PutUint64(key[len(migrationPromotePrefix):], jobID) + return key +} + func isMigrationMetadataKey(rawKey []byte) bool { return (len(rawKey) == len(migrationAckPrefix)+migrationAckKeyIDBytes && bytes.HasPrefix(rawKey, []byte(migrationAckPrefix))) || - (len(rawKey) == len(migrationHLCFloorPrefix)+migrationUint64Bytes && bytes.HasPrefix(rawKey, []byte(migrationHLCFloorPrefix))) + (len(rawKey) == len(migrationHLCFloorPrefix)+migrationUint64Bytes && bytes.HasPrefix(rawKey, []byte(migrationHLCFloorPrefix))) || + (len(rawKey) == len(migrationPromotePrefix)+migrationUint64Bytes && bytes.HasPrefix(rawKey, []byte(migrationPromotePrefix))) } func encodeMigrationImportAck(ack migrationImportAck) []byte { diff --git a/store/migration_versions_test.go b/store/migration_versions_test.go index 77badeedd..91548bb00 100644 --- a/store/migration_versions_test.go +++ b/store/migration_versions_test.go @@ -1,6 +1,7 @@ package store import ( + "bytes" "context" "os" "testing" @@ -181,6 +182,103 @@ func TestImportVersionsIdempotencyAndMetadata(t *testing.T) { }) } +func TestPromoteVersionsMovesStagedVersionsAndDeletesStagedRows(t *testing.T) { + runMigrationStoreSuite(t, func(t *testing.T, st MVCCStore) { + ctx := context.Background() + promoter, ok := st.(MigrationPromoter) + require.True(t, ok) + stateReader, ok := st.(MigrationPromotionStateReader) + require.True(t, ok) + + stage := func(raw string) []byte { + return append([]byte("stage|"), []byte(raw)...) + } + targetKey := func(staged []byte) ([]byte, bool) { + return bytes.TrimPrefix(staged, []byte("stage|")), bytes.HasPrefix(staged, []byte("stage|")) + } + prefix := []byte("stage|") + + require.NoError(t, st.PutAt(ctx, []byte("k"), []byte("old"), 5, 0)) + require.NoError(t, st.PutAt(ctx, stage("k"), []byte("v10"), 10, 0)) + require.NoError(t, st.PutWithTTLAt(ctx, stage("k"), []byte("v20"), 20, 55)) + require.NoError(t, st.DeleteAt(ctx, stage("k"), 30)) + require.NoError(t, st.PutAt(ctx, stage("z"), []byte("z15"), 15, 0)) + + first, err := promoter.PromoteVersions(ctx, PromoteVersionsOptions{ + JobID: 99, + StartKey: prefix, + EndKey: PrefixScanEnd(prefix), + MaxVersions: 2, + TargetKey: targetKey, + }) + require.NoError(t, err) + require.False(t, first.Done) + require.Equal(t, uint64(2), first.PromotedRows) + require.Equal(t, uint64(2), first.TotalPromotedRows) + require.Equal(t, uint64(30), first.MaxPromotedTS) + require.NotEmpty(t, first.NextCursor) + state, ok, err := stateReader.MigrationPromotionState(ctx, 99) + require.NoError(t, err) + require.True(t, ok) + require.False(t, state.Done) + require.Equal(t, first.NextCursor, state.Cursor) + require.Equal(t, uint64(2), state.PromotedRows) + + got, err := st.GetAt(ctx, []byte("k"), 25) + require.NoError(t, err) + require.Equal(t, []byte("v20"), got) + _, err = st.GetAt(ctx, []byte("k"), 35) + require.ErrorIs(t, err, ErrKeyNotFound) + + stagedLeft, err := st.ExportVersions(ctx, ExportVersionsOptions{ + StartKey: prefix, + EndKey: PrefixScanEnd(prefix), + MaxVersions: 10, + }) + require.NoError(t, err) + require.Equal(t, []MVCCVersion{ + {Key: stage("k"), CommitTS: 10, Value: []byte("v10")}, + {Key: stage("z"), CommitTS: 15, Value: []byte("z15")}, + }, stagedLeft.Versions) + + second, err := promoter.PromoteVersions(ctx, PromoteVersionsOptions{ + JobID: 99, + StartKey: prefix, + EndKey: PrefixScanEnd(prefix), + MaxVersions: 10, + TargetKey: targetKey, + }) + require.NoError(t, err) + require.True(t, second.Done) + require.Empty(t, second.NextCursor) + require.Equal(t, uint64(2), second.PromotedRows) + require.Equal(t, uint64(4), second.TotalPromotedRows) + require.Equal(t, uint64(15), second.MaxPromotedTS) + state, ok, err = stateReader.MigrationPromotionState(ctx, 99) + require.NoError(t, err) + require.True(t, ok) + require.True(t, state.Done) + require.Empty(t, state.Cursor) + require.Equal(t, uint64(4), state.PromotedRows) + + got, err = st.GetAt(ctx, []byte("k"), 10) + require.NoError(t, err) + require.Equal(t, []byte("v10"), got) + got, err = st.GetAt(ctx, []byte("z"), 15) + require.NoError(t, err) + require.Equal(t, []byte("z15"), got) + + stagedLeft, err = st.ExportVersions(ctx, ExportVersionsOptions{ + StartKey: prefix, + EndKey: PrefixScanEnd(prefix), + MaxVersions: 10, + }) + require.NoError(t, err) + require.True(t, stagedLeft.Done) + require.Empty(t, stagedLeft.Versions) + }) +} + func TestPebbleImportMetadataPersistsAcrossReopen(t *testing.T) { ctx := context.Background() dir, err := os.MkdirTemp("", "migration-import-persist-*") diff --git a/store/mvcc_store.go b/store/mvcc_store.go index 50be08ba3..b20c35254 100644 --- a/store/mvcc_store.go +++ b/store/mvcc_store.go @@ -60,13 +60,14 @@ func byteSliceComparator(a, b any) int { // mvccStore is an in-memory MVCC implementation backed by a treemap for // deterministic iteration order and range scans. type mvccStore struct { - tree *treemap.Map // key []byte -> []VersionedValue - mtx sync.RWMutex - log *slog.Logger - lastCommitTS uint64 - minRetainedTS uint64 - migrationAcks map[string]migrationImportAck - migrationHLCFloors map[uint64]uint64 + tree *treemap.Map // key []byte -> []VersionedValue + mtx sync.RWMutex + log *slog.Logger + lastCommitTS uint64 + minRetainedTS uint64 + migrationAcks map[string]migrationImportAck + migrationHLCFloors map[uint64]uint64 + migrationPromotions map[uint64]PromotionState // writeConflicts mirrors the per-(kind, key_prefix) counter from // the pebble-backed store so the in-memory implementation shows up // in the same Prometheus series (even if the counts are usually @@ -113,9 +114,10 @@ func NewMVCCStore(opts ...MVCCStoreOption) MVCCStore { log: slog.New(slog.NewJSONHandler(os.Stdout, &slog.HandlerOptions{ Level: slog.LevelWarn, })), - migrationAcks: make(map[string]migrationImportAck), - migrationHLCFloors: make(map[uint64]uint64), - writeConflicts: newWriteConflictCounter(), + migrationAcks: make(map[string]migrationImportAck), + migrationHLCFloors: make(map[uint64]uint64), + migrationPromotions: make(map[uint64]PromotionState), + writeConflicts: newWriteConflictCounter(), } for _, opt := range opts { opt(s) diff --git a/store/store.go b/store/store.go index 78cb59354..a645a95c5 100644 --- a/store/store.go +++ b/store/store.go @@ -116,6 +116,49 @@ type ImportVersionsResult struct { Duplicate bool } +// PromoteVersionsOptions atomically copies staged MVCC versions to their +// target keys and physically removes the staged versions. +type PromoteVersionsOptions struct { + JobID uint64 + StartKey []byte + EndKey []byte + Cursor []byte + MaxVersions int + MaxBytes uint64 + MaxScannedBytes uint64 + TargetKey func(stagedKey []byte) ([]byte, bool) +} + +// PromoteVersionsResult reports one resumable staged-version promotion chunk. +type PromoteVersionsResult struct { + NextCursor []byte + Done bool + PromotedRows uint64 + TotalPromotedRows uint64 + PromotedBytes uint64 + MaxPromotedTS uint64 + ScannedBytes uint64 +} + +// PromotionState is the target-local durable cursor for staged data promotion. +type PromotionState struct { + Cursor []byte + Done bool + PromotedRows uint64 + LastError string +} + +// MigrationPromoter is implemented by stores that can promote staged range +// migration data into the live keyspace. +type MigrationPromoter interface { + PromoteVersions(ctx context.Context, opts PromoteVersionsOptions) (PromoteVersionsResult, error) +} + +// MigrationPromotionStateReader reads target-local staged promotion state. +type MigrationPromotionStateReader interface { + MigrationPromotionState(ctx context.Context, jobID uint64) (PromotionState, bool, error) +} + // OpType describes a mutation kind. type OpType int From c1afe2706097644945bfefaf21e155dd7a6c1431 Mon Sep 17 00:00:00 2001 From: bootjp Date: Tue, 14 Jul 2026 02:08:26 +0900 Subject: [PATCH 02/10] store: persist promote applied index --- kv/fsm_migration_promote.go | 5 ++-- store/lsm_store_applied_index_test.go | 41 +++++++++++++++++++++++++++ store/migration_promote.go | 12 ++++++-- store/store.go | 5 +++- 4 files changed, 57 insertions(+), 6 deletions(-) diff --git a/kv/fsm_migration_promote.go b/kv/fsm_migration_promote.go index 0d9d1ed86..0551b9216 100644 --- a/kv/fsm_migration_promote.go +++ b/kv/fsm_migration_promote.go @@ -41,7 +41,7 @@ func (f *kvFSM) applyMigrationPromote(ctx context.Context, data []byte) any { if !ok { return errors.WithStack(store.ErrNotSupported) } - result, err := promoter.PromoteVersions(ctx, migrationPromoteOptionsFromProto(req)) + result, err := promoter.PromoteVersions(ctx, migrationPromoteOptionsFromProto(req, f.pendingApplyIdx)) if err != nil { return errors.WithStack(err) } @@ -51,7 +51,7 @@ func (f *kvFSM) applyMigrationPromote(ctx context.Context, data []byte) any { return result } -func migrationPromoteOptionsFromProto(req *pb.PromoteStagedVersionsRequest) store.PromoteVersionsOptions { +func migrationPromoteOptionsFromProto(req *pb.PromoteStagedVersionsRequest, appliedIndex uint64) store.PromoteVersionsOptions { maxVersions := int(req.GetMaxVersions()) if maxVersions <= 0 { maxVersions = defaultMigrationPromoteMaxVersions @@ -67,6 +67,7 @@ func migrationPromoteOptionsFromProto(req *pb.PromoteStagedVersionsRequest) stor prefix := distribution.MigrationStagedDataKeyPrefix(req.GetJobId()) return store.PromoteVersionsOptions{ JobID: req.GetJobId(), + AppliedIndex: appliedIndex, StartKey: prefix, EndKey: store.PrefixScanEnd(prefix), Cursor: req.GetCursor(), diff --git a/store/lsm_store_applied_index_test.go b/store/lsm_store_applied_index_test.go index ff00ea950..02c77b1d3 100644 --- a/store/lsm_store_applied_index_test.go +++ b/store/lsm_store_applied_index_test.go @@ -1,6 +1,7 @@ package store import ( + "bytes" "context" "os" "testing" @@ -96,6 +97,46 @@ func TestApplyMutationsRaftAt_BundlesMetaAppliedIndex(t *testing.T) { require.Equal(t, []byte("v1"), val) } +func TestPromoteVersions_BundlesMetaAppliedIndex(t *testing.T) { + ctx := context.Background() + st := newApplyIndexPebbleStore(t) + ps := pebbleStoreApplied(t, st) + + stage := func(raw string) []byte { + return append([]byte("stage|"), []byte(raw)...) + } + targetKey := func(staged []byte) ([]byte, bool) { + return bytes.TrimPrefix(staged, []byte("stage|")), bytes.HasPrefix(staged, []byte("stage|")) + } + prefix := []byte("stage|") + + require.NoError(t, ps.PutAt(ctx, stage("k"), []byte("v10"), 10, 0)) + + const entryIdx uint64 = 77 + result, err := ps.PromoteVersions(ctx, PromoteVersionsOptions{ + JobID: 9, + AppliedIndex: entryIdx, + StartKey: prefix, + EndKey: PrefixScanEnd(prefix), + MaxVersions: 10, + TargetKey: targetKey, + }) + require.NoError(t, err) + require.True(t, result.Done) + require.Equal(t, uint64(1), result.PromotedRows) + + got, present, err := ps.LastAppliedIndex() + require.NoError(t, err) + require.True(t, present, "PromoteVersions must persist metaAppliedIndex") + require.Equal(t, entryIdx, got) + + val, err := ps.GetAt(ctx, []byte("k"), 10) + require.NoError(t, err) + require.Equal(t, []byte("v10"), val) + _, err = ps.GetAt(ctx, stage("k"), 10) + require.ErrorIs(t, err, ErrKeyNotFound) +} + // TestApplyMutationsRaftAt_ZeroIndexLeavesMetaKey covers the // appliedIndex==0 escape hatch — callers without a raft entry index // (test fakes, legacy ApplyMutationsRaft) MUST NOT bump the meta key. diff --git a/store/migration_promote.go b/store/migration_promote.go index cde302a1a..91af77c7e 100644 --- a/store/migration_promote.go +++ b/store/migration_promote.go @@ -226,7 +226,7 @@ func (s *pebbleStore) PromoteVersions(ctx context.Context, opts PromoteVersionsO return PromoteVersionsResult{}, err } result, stateToWrite := finishPromotionResult(opts, state, exported, promoted) - return s.finishPebblePromotion(toPromote, opts.JobID, stateToWrite, result) + return s.finishPebblePromotion(toPromote, opts.JobID, stateToWrite, result, opts.AppliedIndex) } func (s *pebbleStore) finishPebblePromotion( @@ -234,11 +234,12 @@ func (s *pebbleStore) finishPebblePromotion( jobID uint64, stateToWrite *PromotionState, result PromoteVersionsResult, + appliedIndex uint64, ) (PromoteVersionsResult, error) { if len(toPromote) == 0 && stateToWrite == nil { return result, nil } - if err := s.commitPebblePromoteVersions(toPromote, jobID, stateToWrite); err != nil { + if err := s.commitPebblePromoteVersions(toPromote, jobID, stateToWrite, appliedIndex); err != nil { return PromoteVersionsResult{}, err } return result, nil @@ -320,7 +321,7 @@ func (s *pebbleStore) readPebblePromotionState(jobID uint64) (PromotionState, bo return state, true, nil } -func (s *pebbleStore) commitPebblePromoteVersions(versions []promotedVersion, jobID uint64, state *PromotionState) error { +func (s *pebbleStore) commitPebblePromoteVersions(versions []promotedVersion, jobID uint64, state *PromotionState, appliedIndex uint64) error { batch := s.db.NewBatch() defer batch.Close() targets := make([]MVCCVersion, 0, len(versions)) @@ -340,6 +341,11 @@ func (s *pebbleStore) commitPebblePromoteVersions(versions []promotedVersion, jo return errors.WithStack(err) } } + if appliedIndex > 0 { + if err := setPebbleUint64InBatch(batch, metaAppliedIndexBytes, appliedIndex); err != nil { + return err + } + } if err := batch.Commit(s.directApplyWriteOpts()); err != nil { return errors.WithStack(err) } diff --git a/store/store.go b/store/store.go index a645a95c5..b381ca789 100644 --- a/store/store.go +++ b/store/store.go @@ -119,7 +119,10 @@ type ImportVersionsResult struct { // PromoteVersionsOptions atomically copies staged MVCC versions to their // target keys and physically removes the staged versions. type PromoteVersionsOptions struct { - JobID uint64 + JobID uint64 + // AppliedIndex is the optional Raft entry index to bundle with Pebble + // promotion batches as metaAppliedIndex. Zero leaves the meta key unchanged. + AppliedIndex uint64 StartKey []byte EndKey []byte Cursor []byte From dd52db81988f75ecc1c7718935ac897dc2b52ac4 Mon Sep 17 00:00:00 2001 From: bootjp Date: Tue, 14 Jul 2026 03:09:48 +0900 Subject: [PATCH 03/10] Gate migration promotion replay safety --- adapter/internal.go | 61 ++++++++++++++++++++++----- adapter/internal_migration_test.go | 28 ++++++++++++ kv/fsm_migration_promote.go | 8 ++-- kv/fsm_migration_promote_test.go | 9 ++++ store/lsm_store_applied_index_test.go | 18 ++++++++ store/migration_promote.go | 5 ++- 6 files changed, 113 insertions(+), 16 deletions(-) diff --git a/adapter/internal.go b/adapter/internal.go index efacdab71..a3e5024d2 100644 --- a/adapter/internal.go +++ b/adapter/internal.go @@ -3,6 +3,8 @@ package adapter import ( "bytes" "context" + "os" + "strings" "github.com/bootjp/elastickv/distribution" "github.com/bootjp/elastickv/internal/raftengine" @@ -34,12 +36,19 @@ func WithInternalMigrationProposer(proposer raftengine.Proposer) InternalOption } } +func WithInternalMigrationPromoteGate(gate func(context.Context) error) InternalOption { + return func(i *Internal) { + i.migrationPromoteGate = gate + } +} + func NewInternalWithEngine(txm kv.Transactional, leader raftengine.LeaderView, clock *kv.HLC, relay *RedisPubSubRelay, opts ...InternalOption) *Internal { i := &Internal{ - leader: leader, - transactionManager: txm, - clock: clock, - relay: relay, + leader: leader, + transactionManager: txm, + clock: clock, + relay: relay, + migrationPromoteGate: defaultMigrationPromoteGate, } for _, opt := range opts { opt(i) @@ -48,13 +57,14 @@ func NewInternalWithEngine(txm kv.Transactional, leader raftengine.LeaderView, c } type Internal struct { - leader raftengine.LeaderView - transactionManager kv.Transactional - clock *kv.HLC - tsAllocator kv.TimestampAllocator - relay *RedisPubSubRelay - store store.MVCCStore - migrationProposer raftengine.Proposer + leader raftengine.LeaderView + transactionManager kv.Transactional + clock *kv.HLC + tsAllocator kv.TimestampAllocator + relay *RedisPubSubRelay + store store.MVCCStore + migrationProposer raftengine.Proposer + migrationPromoteGate func(context.Context) error pb.UnimplementedInternalServer } @@ -192,6 +202,9 @@ func (i *Internal) PromoteStagedVersions(ctx context.Context, req *pb.PromoteSta if err := i.verifyInternalLeader(ctx); err != nil { return nil, err } + if err := i.verifyMigrationPromoteEnabled(ctx); err != nil { + return nil, err + } result, err := i.proposeMigrationPromote(ctx, req) if err != nil { return nil, errors.WithStack(err) @@ -207,6 +220,32 @@ func (i *Internal) PromoteStagedVersions(ctx context.Context, req *pb.PromoteSta }, nil } +func (i *Internal) verifyMigrationPromoteEnabled(ctx context.Context) error { + if i.migrationPromoteGate == nil { + return nil + } + if err := i.migrationPromoteGate(ctx); err != nil { + return errors.WithStack(err) + } + return nil +} + +func defaultMigrationPromoteGate(context.Context) error { + if migrationPromoteOpcodeEnabledFromEnv() { + return nil + } + return errors.WithStack(status.Error(codes.FailedPrecondition, "migration promote opcode is disabled; enable after every voter is running a build that supports migration promotion")) +} + +func migrationPromoteOpcodeEnabledFromEnv() bool { + switch strings.ToLower(strings.TrimSpace(os.Getenv("ELASTICKV_ENABLE_MIGRATION_PROMOTE_OPCODE"))) { + case "1", "true", "yes", "on": + return true + default: + return false + } +} + func (i *Internal) verifyInternalLeader(ctx context.Context) error { if i.leader.State() != raftengine.StateLeader { return errors.WithStack(ErrNotLeader) diff --git a/adapter/internal_migration_test.go b/adapter/internal_migration_test.go index 54227d418..e4cde4997 100644 --- a/adapter/internal_migration_test.go +++ b/adapter/internal_migration_test.go @@ -153,6 +153,33 @@ func TestInternalImportRangeVersionsAppliesStoreBatch(t *testing.T) { require.GreaterOrEqual(t, clock.Current(), uint64(30)) } +func TestInternalPromoteStagedVersionsRejectsWhenOpcodeGateClosed(t *testing.T) { + t.Parallel() + + ctx := context.Background() + st := store.NewMVCCStore() + clock := kv.NewHLC() + proposer := &applyingMigrationProposer{ + fsm: kv.NewKvFSMWithHLC(st, clock), + } + internal := NewInternalWithEngine(nil, mockInternalLeader{}, clock, nil, + WithInternalStore(st), + WithInternalMigrationProposer(proposer), + WithInternalMigrationPromoteGate(func(context.Context) error { + return status.Error(codes.FailedPrecondition, "migration promote disabled for test") + }), + ) + + resp, err := internal.PromoteStagedVersions(ctx, &pb.PromoteStagedVersionsRequest{ + JobId: 7, + MaxVersions: 10, + }) + require.Nil(t, resp) + require.Error(t, err) + require.Equal(t, codes.FailedPrecondition, status.Code(err)) + require.Equal(t, uint64(0), proposer.calls) +} + func TestInternalPromoteStagedVersionsAppliesStoreBatch(t *testing.T) { t.Parallel() @@ -165,6 +192,7 @@ func TestInternalPromoteStagedVersionsAppliesStoreBatch(t *testing.T) { internal := NewInternalWithEngine(nil, mockInternalLeader{}, clock, nil, WithInternalStore(st), WithInternalMigrationProposer(proposer), + WithInternalMigrationPromoteGate(func(context.Context) error { return nil }), ) staged := distribution.MigrationStagedDataKey(7, []byte("k")) diff --git a/kv/fsm_migration_promote.go b/kv/fsm_migration_promote.go index 0551b9216..c4001b634 100644 --- a/kv/fsm_migration_promote.go +++ b/kv/fsm_migration_promote.go @@ -16,6 +16,8 @@ const ( defaultMigrationPromoteMaxScannedBytes = defaultMigrationPromoteMaxBytes * 4 ) +var ErrMigrationPromoteApply = errors.New("migration promote: FSM apply failed; halting apply") + // MarshalMigrationPromoteCommand encodes a target-group staged-data promotion // chunk as a Raft FSM command. func MarshalMigrationPromoteCommand(req *pb.PromoteStagedVersionsRequest) ([]byte, error) { @@ -35,15 +37,15 @@ func MarshalMigrationPromoteCommand(req *pb.PromoteStagedVersionsRequest) ([]byt func (f *kvFSM) applyMigrationPromote(ctx context.Context, data []byte) any { req := &pb.PromoteStagedVersionsRequest{} if err := proto.Unmarshal(data, req); err != nil { - return errors.WithStack(err) + return haltErr(errors.Wrap(errors.Mark(err, ErrMigrationPromoteApply), "kv/fsm: decode migration promote")) } promoter, ok := f.store.(store.MigrationPromoter) if !ok { - return errors.WithStack(store.ErrNotSupported) + return haltErr(errors.Wrap(errors.Mark(store.ErrNotSupported, ErrMigrationPromoteApply), "kv/fsm: migration promote store")) } result, err := promoter.PromoteVersions(ctx, migrationPromoteOptionsFromProto(req, f.pendingApplyIdx)) if err != nil { - return errors.WithStack(err) + return haltErr(errors.Wrap(errors.Mark(err, ErrMigrationPromoteApply), "kv/fsm: apply migration promote")) } if f.hlc != nil && result.MaxPromotedTS > 0 { f.hlc.Observe(result.MaxPromotedTS) diff --git a/kv/fsm_migration_promote_test.go b/kv/fsm_migration_promote_test.go index 2470c485d..2397243e2 100644 --- a/kv/fsm_migration_promote_test.go +++ b/kv/fsm_migration_promote_test.go @@ -7,6 +7,7 @@ import ( "github.com/bootjp/elastickv/distribution" pb "github.com/bootjp/elastickv/proto" "github.com/bootjp/elastickv/store" + "github.com/cockroachdb/errors" "github.com/stretchr/testify/require" ) @@ -67,3 +68,11 @@ func TestApplyMigrationPromoteMovesStagedVersions(t *testing.T) { require.NoError(t, err) require.Empty(t, left.Versions) } + +func TestApplyMigrationPromoteMalformedPayloadHalts(t *testing.T) { + t.Parallel() + + fsm := &kvFSM{store: store.NewMVCCStore()} + err := haltApplyOf(fsm.Apply([]byte{raftEncodeMigrationPromote, 0xff, 0xff})) + require.True(t, errors.Is(err, ErrMigrationPromoteApply), "got %v", err) +} diff --git a/store/lsm_store_applied_index_test.go b/store/lsm_store_applied_index_test.go index 02c77b1d3..4d4ecd4da 100644 --- a/store/lsm_store_applied_index_test.go +++ b/store/lsm_store_applied_index_test.go @@ -135,6 +135,24 @@ func TestPromoteVersions_BundlesMetaAppliedIndex(t *testing.T) { require.Equal(t, []byte("v10"), val) _, err = ps.GetAt(ctx, stage("k"), 10) require.ErrorIs(t, err, ErrKeyNotFound) + + const retryEntryIdx uint64 = 78 + retry, err := ps.PromoteVersions(ctx, PromoteVersionsOptions{ + JobID: 9, + AppliedIndex: retryEntryIdx, + StartKey: prefix, + EndKey: PrefixScanEnd(prefix), + MaxVersions: 10, + TargetKey: targetKey, + }) + require.NoError(t, err) + require.True(t, retry.Done) + require.Equal(t, uint64(1), retry.TotalPromotedRows) + + got, present, err = ps.LastAppliedIndex() + require.NoError(t, err) + require.True(t, present, "completed PromoteVersions retry must persist metaAppliedIndex") + require.Equal(t, retryEntryIdx, got) } // TestApplyMutationsRaftAt_ZeroIndexLeavesMetaKey covers the diff --git a/store/migration_promote.go b/store/migration_promote.go index 91af77c7e..a0b8174a5 100644 --- a/store/migration_promote.go +++ b/store/migration_promote.go @@ -218,7 +218,8 @@ func (s *pebbleStore) PromoteVersions(ctx context.Context, opts PromoteVersionsO return PromoteVersionsResult{}, err } if opts.JobID != 0 && state.Done { - return PromoteVersionsResult{Done: true, TotalPromotedRows: state.PromotedRows}, nil + result := PromoteVersionsResult{Done: true, TotalPromotedRows: state.PromotedRows} + return s.finishPebblePromotion(nil, opts.JobID, nil, result, opts.AppliedIndex) } opts.Cursor = cursor exported, toPromote, promoted, err := s.planPebblePromotionLocked(ctx, opts) @@ -236,7 +237,7 @@ func (s *pebbleStore) finishPebblePromotion( result PromoteVersionsResult, appliedIndex uint64, ) (PromoteVersionsResult, error) { - if len(toPromote) == 0 && stateToWrite == nil { + if len(toPromote) == 0 && stateToWrite == nil && appliedIndex == 0 { return result, nil } if err := s.commitPebblePromoteVersions(toPromote, jobID, stateToWrite, appliedIndex); err != nil { From b04d30fe3cfe3dab6bd7bfdadf07920d1a886681 Mon Sep 17 00:00:00 2001 From: bootjp Date: Tue, 14 Jul 2026 03:40:50 +0900 Subject: [PATCH 04/10] store: keep promote replay ungated --- store/lsm_migration.go | 8 +++-- store/lsm_store_registration_gate_test.go | 44 +++++++++++++++++++++++ store/migration_promote.go | 4 ++- 3 files changed, 52 insertions(+), 4 deletions(-) diff --git a/store/lsm_migration.go b/store/lsm_migration.go index e0796cbf6..d0c66b822 100644 --- a/store/lsm_migration.go +++ b/store/lsm_migration.go @@ -247,7 +247,9 @@ func (s *pebbleStore) validatePebbleImportBatch(opts ImportVersionsOptions) (boo func (s *pebbleStore) commitPebbleImportBatch(opts ImportVersionsOptions, batchMax uint64) error { batch := s.db.NewBatch() defer batch.Close() - if err := s.applyImportVersionsBatch(batch, opts.Versions); err != nil { + // Keep the existing import registration-gate behavior; promotion opts out + // at its call site because it replays committed Raft entries. + if err := s.applyImportVersionsBatch(batch, opts.Versions, true); err != nil { return err } if err := batch.Set(migrationAckKey(opts.JobID, opts.BracketID), encodeMigrationImportAck(migrationImportAck{ @@ -277,14 +279,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..11e6d14fc 100644 --- a/store/lsm_store_registration_gate_test.go +++ b/store/lsm_store_registration_gate_test.go @@ -1,6 +1,7 @@ package store import ( + "bytes" "context" "path/filepath" "testing" @@ -135,6 +136,49 @@ func TestRegistrationGate_FSMApplyPathNeverGated(t *testing.T) { mustGet(t, f.mvcc, []byte("raft"), 150, "applied") } +func TestRegistrationGate_PromoteVersionsNeverGated(t *testing.T) { + t.Parallel() + ctx := context.Background() + registered := true + f := newRegGateStore(t, ®istered) + stage := func(raw string) []byte { + return append([]byte("stage|"), []byte(raw)...) + } + targetKey := func(staged []byte) ([]byte, bool) { + return bytes.TrimPrefix(staged, []byte("stage|")), bytes.HasPrefix(staged, []byte("stage|")) + } + + if err := f.mvcc.PutAt(ctx, stage("promote"), []byte("value"), 100, 0); err != nil { + t.Fatalf("seed staged PutAt: %v", err) + } + registered = false + if err := f.mvcc.PutAt(ctx, []byte("direct"), []byte("blocked"), 110, 0); !errors.Is(err, ErrWriterNotRegistered) { + t.Fatalf("direct PutAt pre-registration: got %v, want ErrWriterNotRegistered", err) + } + promoter, ok := f.mvcc.(MigrationPromoter) + if !ok { + t.Fatalf("expected MigrationPromoter, got %T", f.mvcc) + } + + result, err := promoter.PromoteVersions(ctx, PromoteVersionsOptions{ + JobID: 11, + StartKey: []byte("stage|"), + EndKey: PrefixScanEnd([]byte("stage|")), + MaxVersions: 10, + TargetKey: targetKey, + }) + if err != nil { + t.Fatalf("PromoteVersions pre-registration: %v", err) + } + if !result.Done || result.PromotedRows != 1 { + t.Fatalf("PromoteVersions result = %+v, want done with one promoted row", result) + } + mustGet(t, f.mvcc, []byte("promote"), 150, "value") + if _, err := f.mvcc.GetAt(ctx, stage("promote"), 150); !errors.Is(err, ErrKeyNotFound) { + t.Fatalf("staged version after promotion: got %v, want ErrKeyNotFound", err) + } +} + // TestRegistrationGate_NotEncryptingIsUngated confirms the gate is only // consulted when the write WOULD encrypt. When the cutover gate is off // (envelope inactive) the direct path writes cleartext and never diff --git a/store/migration_promote.go b/store/migration_promote.go index a0b8174a5..12cc79389 100644 --- a/store/migration_promote.go +++ b/store/migration_promote.go @@ -329,7 +329,9 @@ func (s *pebbleStore) commitPebblePromoteVersions(versions []promotedVersion, jo for _, version := range versions { targets = append(targets, version.target) } - if err := s.applyImportVersionsBatch(batch, targets); err != nil { + // Promotion is replayed from the Raft FSM, so it must not fail closed on + // this node's local writer-registration state. + if err := s.applyImportVersionsBatch(batch, targets, false); err != nil { return err } for _, version := range versions { From c72eace4cd5007b515cda3d7fb961b4893e55e48 Mon Sep 17 00:00:00 2001 From: bootjp Date: Tue, 14 Jul 2026 06:18:35 +0900 Subject: [PATCH 05/10] Validate promote cursor before proposal --- adapter/internal.go | 14 ++++++++ adapter/internal_migration_test.go | 52 ++++++++++++++++++++++++++++++ store/migration_versions.go | 19 +++++++++++ 3 files changed, 85 insertions(+) diff --git a/adapter/internal.go b/adapter/internal.go index a909990f6..656b84380 100644 --- a/adapter/internal.go +++ b/adapter/internal.go @@ -216,6 +216,9 @@ func (i *Internal) PromoteStagedVersions(ctx context.Context, req *pb.PromoteSta if err := i.verifyMigrationPromoteEnabled(ctx); err != nil { return nil, err } + if err := validatePromoteStagedVersionsRequest(req); err != nil { + return nil, errors.WithStack(err) + } result, err := i.proposeMigrationPromote(ctx, req) if err != nil { return nil, errors.WithStack(err) @@ -241,6 +244,17 @@ func (i *Internal) verifyMigrationPromoteEnabled(ctx context.Context) error { return nil } +func validatePromoteStagedVersionsRequest(req *pb.PromoteStagedVersionsRequest) error { + prefix := distribution.MigrationStagedDataKeyPrefix(req.GetJobId()) + if err := store.ValidateExportCursorForRange(req.GetCursor(), prefix, store.PrefixScanEnd(prefix)); err != nil { + if errors.Is(err, store.ErrInvalidExportCursor) { + return errors.WithStack(status.Error(codes.InvalidArgument, store.ErrInvalidExportCursor.Error())) + } + return errors.WithStack(status.Errorf(codes.Internal, "validate promote cursor: %v", err)) + } + return nil +} + func defaultMigrationPromoteGate(context.Context) error { if migrationPromoteOpcodeEnabledFromEnv() { return nil diff --git a/adapter/internal_migration_test.go b/adapter/internal_migration_test.go index 9760a61b0..c1cf5df04 100644 --- a/adapter/internal_migration_test.go +++ b/adapter/internal_migration_test.go @@ -2,6 +2,7 @@ package adapter import ( "context" + "encoding/binary" "testing" "github.com/bootjp/elastickv/distribution" @@ -71,6 +72,15 @@ func (s *captureExportRangeVersionsStream) Send(resp *pb.ExportRangeVersionsResp return nil } +func encodeTestExportCursor(key []byte, commitTS uint64, tag byte) []byte { + var out []byte + out = binary.AppendUvarint(out, uint64(len(key))) + out = append(out, key...) + out = binary.AppendUvarint(out, commitTS) + out = append(out, tag) + return out +} + func testPrefixScanEnd(prefix []byte) []byte { out := append([]byte(nil), prefix...) for i := len(out) - 1; i >= 0; i-- { @@ -291,6 +301,48 @@ func TestInternalPromoteStagedVersionsRejectsWhenOpcodeGateClosed(t *testing.T) require.Equal(t, uint64(0), proposer.calls) } +func TestInternalPromoteStagedVersionsRejectsInvalidCursorBeforePropose(t *testing.T) { + t.Parallel() + + ctx := context.Background() + for _, tc := range []struct { + name string + cursor []byte + }{ + {name: "malformed cursor", cursor: []byte{0xff}}, + { + name: "cursor outside job staged prefix", + cursor: encodeTestExportCursor(distribution.MigrationStagedDataKey(8, []byte("k")), 30, 0), + }, + } { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + + st := store.NewMVCCStore() + clock := kv.NewHLC() + proposer := &applyingMigrationProposer{ + fsm: kv.NewKvFSMWithHLC(st, clock), + } + internal := NewInternalWithEngine(nil, mockInternalLeader{}, clock, nil, + WithInternalStore(st), + WithInternalMigrationProposer(proposer), + WithInternalMigrationPromoteGate(func(context.Context) error { return nil }), + ) + + resp, err := internal.PromoteStagedVersions(ctx, &pb.PromoteStagedVersionsRequest{ + JobId: 7, + Cursor: tc.cursor, + MaxVersions: 10, + }) + require.Nil(t, resp) + require.Error(t, err) + require.Equal(t, codes.InvalidArgument, status.Code(err)) + require.ErrorContains(t, err, store.ErrInvalidExportCursor.Error()) + require.Equal(t, uint64(0), proposer.calls) + }) + } +} + func TestInternalPromoteStagedVersionsAppliesStoreBatch(t *testing.T) { t.Parallel() diff --git a/store/migration_versions.go b/store/migration_versions.go index cd4c535fb..93ce86819 100644 --- a/store/migration_versions.go +++ b/store/migration_versions.go @@ -70,6 +70,25 @@ func decodeExportCursor(cursor []byte) (exportCursorPosition, error) { return exportCursorPosition{key: key, commitTS: commitTS, tag: tag}, nil } +// ValidateExportCursorForRange verifies that an export cursor decodes and +// resumes inside the supplied key interval. +func ValidateExportCursorForRange(cursor, startKey, endKey []byte) error { + pos, err := decodeExportCursor(cursor) + if err != nil { + return err + } + if len(pos.key) == 0 { + return nil + } + if startKey != nil && bytes.Compare(pos.key, startKey) < 0 { + return errors.WithStack(ErrInvalidExportCursor) + } + if endKey != nil && bytes.Compare(pos.key, endKey) >= 0 { + return errors.WithStack(ErrInvalidExportCursor) + } + return nil +} + func migrationAckKey(jobID, bracketID uint64) []byte { key := make([]byte, len(migrationAckPrefix)+migrationAckKeyIDBytes) copy(key, migrationAckPrefix) From 3cb3aad44078caf0a86233faf4dbac742fdc915d Mon Sep 17 00:00:00 2001 From: bootjp Date: Tue, 14 Jul 2026 10:52:15 +0900 Subject: [PATCH 06/10] store: advance promotion commit watermark --- store/migration_promote.go | 31 ++++++++++++++--- store/migration_versions_test.go | 58 ++++++++++++++++++++++++++++++++ 2 files changed, 84 insertions(+), 5 deletions(-) diff --git a/store/migration_promote.go b/store/migration_promote.go index 12cc79389..3d2d9de57 100644 --- a/store/migration_promote.go +++ b/store/migration_promote.go @@ -219,7 +219,7 @@ func (s *pebbleStore) PromoteVersions(ctx context.Context, opts PromoteVersionsO } if opts.JobID != 0 && state.Done { result := PromoteVersionsResult{Done: true, TotalPromotedRows: state.PromotedRows} - return s.finishPebblePromotion(nil, opts.JobID, nil, result, opts.AppliedIndex) + return s.finishPebblePromotion(nil, opts.JobID, nil, result, opts.AppliedIndex, 0) } opts.Cursor = cursor exported, toPromote, promoted, err := s.planPebblePromotionLocked(ctx, opts) @@ -227,7 +227,7 @@ func (s *pebbleStore) PromoteVersions(ctx context.Context, opts PromoteVersionsO return PromoteVersionsResult{}, err } result, stateToWrite := finishPromotionResult(opts, state, exported, promoted) - return s.finishPebblePromotion(toPromote, opts.JobID, stateToWrite, result, opts.AppliedIndex) + return s.finishPebblePromotion(toPromote, opts.JobID, stateToWrite, result, opts.AppliedIndex, promoted.MaxPromotedTS) } func (s *pebbleStore) finishPebblePromotion( @@ -236,11 +236,12 @@ func (s *pebbleStore) finishPebblePromotion( stateToWrite *PromotionState, result PromoteVersionsResult, appliedIndex uint64, + maxPromotedTS uint64, ) (PromoteVersionsResult, error) { - if len(toPromote) == 0 && stateToWrite == nil && appliedIndex == 0 { + if len(toPromote) == 0 && stateToWrite == nil && appliedIndex == 0 && maxPromotedTS == 0 { return result, nil } - if err := s.commitPebblePromoteVersions(toPromote, jobID, stateToWrite, appliedIndex); err != nil { + if err := s.commitPebblePromoteVersions(toPromote, jobID, stateToWrite, appliedIndex, maxPromotedTS); err != nil { return PromoteVersionsResult{}, err } return result, nil @@ -322,7 +323,7 @@ func (s *pebbleStore) readPebblePromotionState(jobID uint64) (PromotionState, bo return state, true, nil } -func (s *pebbleStore) commitPebblePromoteVersions(versions []promotedVersion, jobID uint64, state *PromotionState, appliedIndex uint64) error { +func (s *pebbleStore) commitPebblePromoteVersions(versions []promotedVersion, jobID uint64, state *PromotionState, appliedIndex, maxPromotedTS uint64) error { batch := s.db.NewBatch() defer batch.Close() targets := make([]MVCCVersion, 0, len(versions)) @@ -349,6 +350,26 @@ func (s *pebbleStore) commitPebblePromoteVersions(versions []promotedVersion, jo return err } } + return s.commitPebblePromotionBatch(batch, maxPromotedTS) +} + +func (s *pebbleStore) commitPebblePromotionBatch(batch *pebble.Batch, maxPromotedTS uint64) error { + if maxPromotedTS > 0 { + s.mtx.Lock() + defer s.mtx.Unlock() + newLastTS := s.lastCommitTS + if maxPromotedTS > newLastTS { + newLastTS = maxPromotedTS + } + if err := setPebbleUint64InBatch(batch, metaLastCommitTSBytes, newLastTS); err != nil { + return err + } + if err := batch.Commit(s.directApplyWriteOpts()); err != nil { + return errors.WithStack(err) + } + s.updateLastCommitTS(newLastTS) + return nil + } if err := batch.Commit(s.directApplyWriteOpts()); err != nil { return errors.WithStack(err) } diff --git a/store/migration_versions_test.go b/store/migration_versions_test.go index 91548bb00..d11fefaed 100644 --- a/store/migration_versions_test.go +++ b/store/migration_versions_test.go @@ -6,6 +6,7 @@ import ( "os" "testing" + "github.com/cockroachdb/pebble/v2" "github.com/stretchr/testify/require" ) @@ -279,6 +280,63 @@ func TestPromoteVersionsMovesStagedVersionsAndDeletesStagedRows(t *testing.T) { }) } +func TestPebblePromoteVersionsAdvancesLastCommitTS(t *testing.T) { + ctx := context.Background() + dir := t.TempDir() + st, err := NewPebbleStore(dir) + require.NoError(t, err) + closed := false + t.Cleanup(func() { + if !closed { + require.NoError(t, st.Close()) + } + }) + ps, ok := st.(*pebbleStore) + require.True(t, ok) + + stage := func(raw string) []byte { + return append([]byte("stage|"), []byte(raw)...) + } + targetKey := func(staged []byte) ([]byte, bool) { + return bytes.TrimPrefix(staged, []byte("stage|")), bytes.HasPrefix(staged, []byte("stage|")) + } + prefix := []byte("stage|") + + const promotedTS uint64 = 100 + require.NoError(t, ps.db.Set(encodeKey(stage("k"), promotedTS), encodeValue([]byte("v100"), false, 0, encStateCleartext), pebble.NoSync)) + require.Zero(t, ps.LastCommitTS()) + + result, err := ps.PromoteVersions(ctx, PromoteVersionsOptions{ + JobID: 101, + StartKey: prefix, + EndKey: PrefixScanEnd(prefix), + MaxVersions: 10, + TargetKey: targetKey, + }) + require.NoError(t, err) + require.True(t, result.Done) + require.Equal(t, uint64(1), result.PromotedRows) + require.Equal(t, promotedTS, result.MaxPromotedTS) + require.Equal(t, promotedTS, ps.LastCommitTS()) + + metaTS, err := readPebbleUint64(ps.db, metaLastCommitTSBytes) + require.NoError(t, err) + require.Equal(t, promotedTS, metaTS) + val, err := ps.GetAt(ctx, []byte("k"), ps.LastCommitTS()) + require.NoError(t, err) + require.Equal(t, []byte("v100"), val) + + require.NoError(t, st.Close()) + closed = true + reopened, err := NewPebbleStore(dir) + require.NoError(t, err) + defer func() { require.NoError(t, reopened.Close()) }() + require.Equal(t, promotedTS, reopened.LastCommitTS()) + val, err = reopened.GetAt(ctx, []byte("k"), reopened.LastCommitTS()) + require.NoError(t, err) + require.Equal(t, []byte("v100"), val) +} + func TestPebbleImportMetadataPersistsAcrossReopen(t *testing.T) { ctx := context.Background() dir, err := os.MkdirTemp("", "migration-import-persist-*") From 90fa6544b08d45dc6c963ddb584caf104047ae3c Mon Sep 17 00:00:00 2001 From: bootjp Date: Tue, 14 Jul 2026 11:45:49 +0900 Subject: [PATCH 07/10] store: honor raft promotion sync mode --- store/lsm_store_sync_mode_test.go | 60 +++++++++++++++++++++++++++++++ store/migration_promote.go | 51 +++++++++++++++++++++----- 2 files changed, 103 insertions(+), 8 deletions(-) diff --git a/store/lsm_store_sync_mode_test.go b/store/lsm_store_sync_mode_test.go index e5e065bfd..29349b997 100644 --- a/store/lsm_store_sync_mode_test.go +++ b/store/lsm_store_sync_mode_test.go @@ -1,6 +1,7 @@ package store import ( + "bytes" "context" "testing" @@ -116,6 +117,65 @@ func TestDirectApplyWriteOpts_AlwaysSync(t *testing.T) { }) } +func TestPromoteVersionsWriteOptsFollowApplyContext(t *testing.T) { + t.Run("raft-applied promotion observes nosync", func(t *testing.T) { + ps := newPebbleStoreWithFSMApplyWriteOptsForTest(t, t.TempDir(), pebble.NoSync, fsmSyncModeNoSync) + defer ps.Close() + + require.Same(t, pebble.NoSync, ps.promotionWriteOpts(123), + "promotion with an applied index must use raft apply write options") + }) + + t.Run("direct promotion stays sync", func(t *testing.T) { + ps := newPebbleStoreWithFSMApplyWriteOptsForTest(t, t.TempDir(), pebble.NoSync, fsmSyncModeNoSync) + defer ps.Close() + + require.Same(t, pebble.Sync, ps.promotionWriteOpts(0), + "promotion without an applied index has no raft durability backstop") + }) +} + +func TestPromoteVersionsRaftNoSyncFunctionalEquivalence(t *testing.T) { + dir := t.TempDir() + ps := newPebbleStoreWithFSMApplyWriteOptsForTest(t, dir, pebble.NoSync, fsmSyncModeNoSync) + defer ps.Close() + + ctx := context.Background() + stage := func(raw string) []byte { + return append([]byte("stage|"), []byte(raw)...) + } + targetKey := func(staged []byte) ([]byte, bool) { + return bytes.TrimPrefix(staged, []byte("stage|")), bytes.HasPrefix(staged, []byte("stage|")) + } + prefix := []byte("stage|") + + require.NoError(t, ps.PutAt(ctx, stage("k"), []byte("v10"), 10, 0)) + + const entryIdx uint64 = 88 + result, err := ps.PromoteVersions(ctx, PromoteVersionsOptions{ + JobID: 12, + AppliedIndex: entryIdx, + StartKey: prefix, + EndKey: PrefixScanEnd(prefix), + MaxVersions: 10, + TargetKey: targetKey, + }) + require.NoError(t, err) + require.True(t, result.Done) + require.Equal(t, uint64(1), result.PromotedRows) + + val, err := ps.GetAt(ctx, []byte("k"), 10) + require.NoError(t, err) + require.Equal(t, []byte("v10"), val) + _, err = ps.GetAt(ctx, stage("k"), 10) + require.ErrorIs(t, err, ErrKeyNotFound) + + got, present, err := ps.LastAppliedIndex() + require.NoError(t, err) + require.True(t, present) + require.Equal(t, entryIdx, got) +} + // TestDirectApplyMutations_NoSyncConfigured_StillWritesDurably is the // functional twin of TestDirectApplyWriteOpts_AlwaysSync: it exercises // the public ApplyMutations and DeletePrefixAt entry points with a diff --git a/store/migration_promote.go b/store/migration_promote.go index 3d2d9de57..b0e5b7fa5 100644 --- a/store/migration_promote.go +++ b/store/migration_promote.go @@ -217,9 +217,10 @@ func (s *pebbleStore) PromoteVersions(ctx context.Context, opts PromoteVersionsO if err != nil { return PromoteVersionsResult{}, err } + writeOpts := s.promotionWriteOpts(opts.AppliedIndex) if opts.JobID != 0 && state.Done { result := PromoteVersionsResult{Done: true, TotalPromotedRows: state.PromotedRows} - return s.finishPebblePromotion(nil, opts.JobID, nil, result, opts.AppliedIndex, 0) + return s.finishPebblePromotion(nil, opts.JobID, nil, result, opts.AppliedIndex, 0, writeOpts) } opts.Cursor = cursor exported, toPromote, promoted, err := s.planPebblePromotionLocked(ctx, opts) @@ -227,7 +228,22 @@ func (s *pebbleStore) PromoteVersions(ctx context.Context, opts PromoteVersionsO return PromoteVersionsResult{}, err } result, stateToWrite := finishPromotionResult(opts, state, exported, promoted) - return s.finishPebblePromotion(toPromote, opts.JobID, stateToWrite, result, opts.AppliedIndex, promoted.MaxPromotedTS) + return s.finishPebblePromotion( + toPromote, + opts.JobID, + stateToWrite, + result, + opts.AppliedIndex, + promoted.MaxPromotedTS, + writeOpts, + ) +} + +func (s *pebbleStore) promotionWriteOpts(appliedIndex uint64) *pebble.WriteOptions { + if appliedIndex > 0 { + return s.raftApplyWriteOpts() + } + return s.directApplyWriteOpts() } func (s *pebbleStore) finishPebblePromotion( @@ -237,11 +253,19 @@ func (s *pebbleStore) finishPebblePromotion( result PromoteVersionsResult, appliedIndex uint64, maxPromotedTS uint64, + writeOpts *pebble.WriteOptions, ) (PromoteVersionsResult, error) { if len(toPromote) == 0 && stateToWrite == nil && appliedIndex == 0 && maxPromotedTS == 0 { return result, nil } - if err := s.commitPebblePromoteVersions(toPromote, jobID, stateToWrite, appliedIndex, maxPromotedTS); err != nil { + if err := s.commitPebblePromoteVersions( + toPromote, + jobID, + stateToWrite, + appliedIndex, + maxPromotedTS, + writeOpts, + ); err != nil { return PromoteVersionsResult{}, err } return result, nil @@ -323,7 +347,14 @@ func (s *pebbleStore) readPebblePromotionState(jobID uint64) (PromotionState, bo return state, true, nil } -func (s *pebbleStore) commitPebblePromoteVersions(versions []promotedVersion, jobID uint64, state *PromotionState, appliedIndex, maxPromotedTS uint64) error { +func (s *pebbleStore) commitPebblePromoteVersions( + versions []promotedVersion, + jobID uint64, + state *PromotionState, + appliedIndex uint64, + maxPromotedTS uint64, + writeOpts *pebble.WriteOptions, +) error { batch := s.db.NewBatch() defer batch.Close() targets := make([]MVCCVersion, 0, len(versions)) @@ -350,10 +381,14 @@ func (s *pebbleStore) commitPebblePromoteVersions(versions []promotedVersion, jo return err } } - return s.commitPebblePromotionBatch(batch, maxPromotedTS) + return s.commitPebblePromotionBatch(batch, maxPromotedTS, writeOpts) } -func (s *pebbleStore) commitPebblePromotionBatch(batch *pebble.Batch, maxPromotedTS uint64) error { +func (s *pebbleStore) commitPebblePromotionBatch( + batch *pebble.Batch, + maxPromotedTS uint64, + writeOpts *pebble.WriteOptions, +) error { if maxPromotedTS > 0 { s.mtx.Lock() defer s.mtx.Unlock() @@ -364,13 +399,13 @@ func (s *pebbleStore) commitPebblePromotionBatch(batch *pebble.Batch, maxPromote if err := setPebbleUint64InBatch(batch, metaLastCommitTSBytes, newLastTS); err != nil { return err } - if err := batch.Commit(s.directApplyWriteOpts()); err != nil { + if err := batch.Commit(writeOpts); err != nil { return errors.WithStack(err) } s.updateLastCommitTS(newLastTS) return nil } - if err := batch.Commit(s.directApplyWriteOpts()); err != nil { + if err := batch.Commit(writeOpts); err != nil { return errors.WithStack(err) } return nil From a9c9afef3b080c6979440b1b703faea673eebd9b Mon Sep 17 00:00:00 2001 From: bootjp Date: Tue, 14 Jul 2026 12:00:25 +0900 Subject: [PATCH 08/10] Trim Lua negative cache bound test --- adapter/redis_lua_context.go | 26 ++++++++------ adapter/redis_lua_negative_type_cache_test.go | 35 ++++++------------- 2 files changed, 27 insertions(+), 34 deletions(-) diff --git a/adapter/redis_lua_context.go b/adapter/redis_lua_context.go index 4a57051b5..447423ee4 100644 --- a/adapter/redis_lua_context.go +++ b/adapter/redis_lua_context.go @@ -478,20 +478,26 @@ func (c *luaScriptContext) keyType(key []byte) (redisValueType, error) { if err != nil { return redisTypeNone, err } - if typ == redisTypeNone && len(c.negativeType) < maxNegativeTypeCacheEntries { - // Pin the absence result for the rest of this Eval so repeated - // BullMQ-style polling of a missing key (e.g. a "delayed" zset) - // does not re-run the ~8-seek rawKeyTypeAt probe on every - // redis.call. - // - // Bounded to keep adversarial scripts from growing the map - // unboundedly; once full, subsequent misses correctly fall - // through to the server probe without caching. - c.negativeType[string(key)] = true + if typ == redisTypeNone { + c.rememberNegativeType(key) } return typ, nil } +func (c *luaScriptContext) rememberNegativeType(key []byte) { + if len(c.negativeType) >= maxNegativeTypeCacheEntries { + return + } + // Pin the absence result for the rest of this Eval so repeated + // BullMQ-style polling of a missing key (e.g. a "delayed" zset) does + // not re-run the ~8-seek rawKeyTypeAt probe on every redis.call. + // + // Bounded to keep adversarial scripts from growing the map + // unboundedly; once full, subsequent misses correctly fall through to + // the server probe without caching. + c.negativeType[string(key)] = true +} + func (c *luaScriptContext) ensureKeyNotExpired(key []byte) error { ttl, err := c.loadTTL(key) if err != nil { diff --git a/adapter/redis_lua_negative_type_cache_test.go b/adapter/redis_lua_negative_type_cache_test.go index 6253f7e30..f44aec35a 100644 --- a/adapter/redis_lua_negative_type_cache_test.go +++ b/adapter/redis_lua_negative_type_cache_test.go @@ -149,44 +149,31 @@ func TestLuaNegativeTypeCache_SingleProbePerKey(t *testing.T) { // itself stays bounded. func TestLuaNegativeTypeCache_BoundedSize(t *testing.T) { t.Parallel() - nodes, _, _ := createNode(t, 3) - defer shutdown(nodes) - ctx := context.Background() - sc, err := newLuaScriptContext(ctx, nodes[0].redisServer) - require.NoError(t, err) - defer sc.Close() + sc := &luaScriptContext{negativeType: map[string]bool{}} // Probe cap+overflow unique missing keys. Each probe is a miss // (redisTypeNone); only the first `cap` should be memoized. const overflow = 50 total := maxNegativeTypeCacheEntries + overflow for i := 0; i < total; i++ { - typ, kerr := sc.keyType([]byte(fmt.Sprintf("lua:neg:cap:%d", i))) - require.NoError(t, kerr) - require.Equal(t, redisTypeNone, typ) + sc.rememberNegativeType([]byte(fmt.Sprintf("lua:neg:cap:%d", i))) } require.Equal(t, maxNegativeTypeCacheEntries, len(sc.negativeType), "negativeType map must be capped at maxNegativeTypeCacheEntries") - // Each unique key above required exactly one probe on first access. - require.Equal(t, total, sc.keyTypeProbeCount, - "each unique key must have triggered exactly one server probe") - - // Re-probing one of the first `cap` keys must hit the cache - // (no additional server probe). Re-probing an overflow key must - // miss the cache and issue another server probe. + // One of the first `cap` keys must be cached. An overflow key must + // remain uncached so keyType falls back to a server probe in real Eval + // execution while the map size stays bounded. cachedKey := []byte("lua:neg:cap:0") - _, kerr := sc.keyType(cachedKey) - require.NoError(t, kerr) - require.Equal(t, total, sc.keyTypeProbeCount, - "a key inserted before the cap must remain cached") + typ, ok := sc.cachedType(cachedKey) + require.True(t, ok, "a key inserted before the cap must remain cached") + require.Equal(t, redisTypeNone, typ) overflowKey := []byte(fmt.Sprintf("lua:neg:cap:%d", maxNegativeTypeCacheEntries+1)) - _, kerr = sc.keyType(overflowKey) - require.NoError(t, kerr) - require.Equal(t, total+1, sc.keyTypeProbeCount, - "a key probed after the cap was reached must fall back to the server probe") + _, ok = sc.cachedType(overflowKey) + require.False(t, ok, "a key probed after the cap must fall back to the server probe") + sc.rememberNegativeType(overflowKey) require.Equal(t, maxNegativeTypeCacheEntries, len(sc.negativeType), "fallback probe must NOT grow the bounded cache") } From b1381f3afb46d4a7c05f56427b539b94c35c0fa2 Mon Sep 17 00:00:00 2001 From: bootjp Date: Tue, 14 Jul 2026 19:13:57 +0900 Subject: [PATCH 09/10] migration: validate promotion resume state --- adapter/internal.go | 2 +- adapter/internal_migration_test.go | 21 ++++++++++++- store/migration_versions.go | 31 ++++++++++++++++++- store/migration_versions_test.go | 48 ++++++++++++++++++++++++++++++ store/mvcc_store.go | 1 + store/mvcc_store_snapshot_test.go | 46 ++++++++++++++++++++++++++++ 6 files changed, 146 insertions(+), 3 deletions(-) diff --git a/adapter/internal.go b/adapter/internal.go index 019abec82..7227296bf 100644 --- a/adapter/internal.go +++ b/adapter/internal.go @@ -253,7 +253,7 @@ func (i *Internal) verifyMigrationPromoteEnabled(ctx context.Context) error { func validatePromoteStagedVersionsRequest(req *pb.PromoteStagedVersionsRequest) error { prefix := distribution.MigrationStagedDataKeyPrefix(req.GetJobId()) - if err := store.ValidateExportCursorForRange(req.GetCursor(), prefix, store.PrefixScanEnd(prefix)); err != nil { + if err := store.ValidatePromotionCursorForRange(req.GetCursor(), prefix, store.PrefixScanEnd(prefix)); err != nil { if errors.Is(err, store.ErrInvalidExportCursor) { return errors.WithStack(status.Error(codes.InvalidArgument, store.ErrInvalidExportCursor.Error())) } diff --git a/adapter/internal_migration_test.go b/adapter/internal_migration_test.go index c1cf5df04..542a8b394 100644 --- a/adapter/internal_migration_test.go +++ b/adapter/internal_migration_test.go @@ -60,6 +60,13 @@ type captureExportRangeVersionsStream struct { responses []*pb.ExportRangeVersionsResponse } +const ( + testExportCursorTagEmitted byte = iota + testExportCursorTagScanned + testExportCursorTagPrunedKey + testExportCursorTagSkippedKey +) + func (s *captureExportRangeVersionsStream) Context() context.Context { if s.ctx != nil { return s.ctx @@ -312,7 +319,19 @@ func TestInternalPromoteStagedVersionsRejectsInvalidCursorBeforePropose(t *testi {name: "malformed cursor", cursor: []byte{0xff}}, { name: "cursor outside job staged prefix", - cursor: encodeTestExportCursor(distribution.MigrationStagedDataKey(8, []byte("k")), 30, 0), + cursor: encodeTestExportCursor(distribution.MigrationStagedDataKey(8, []byte("k")), 30, testExportCursorTagEmitted), + }, + { + name: "scanned cursor inside staged prefix", + cursor: encodeTestExportCursor(distribution.MigrationStagedDataKey(7, []byte("k")), 31, testExportCursorTagScanned), + }, + { + name: "pruned-key cursor inside staged prefix", + cursor: encodeTestExportCursor(distribution.MigrationStagedDataKey(7, []byte("k")), 32, testExportCursorTagPrunedKey), + }, + { + name: "skipped-key cursor inside staged prefix", + cursor: encodeTestExportCursor(distribution.MigrationStagedDataKey(7, []byte("k")), 33, testExportCursorTagSkippedKey), }, } { t.Run(tc.name, func(t *testing.T) { diff --git a/store/migration_versions.go b/store/migration_versions.go index df40b33fc..ef3ef8254 100644 --- a/store/migration_versions.go +++ b/store/migration_versions.go @@ -92,8 +92,20 @@ func decodeExportCursor(cursor []byte) (exportCursorPosition, error) { } // ValidateExportCursorForRange verifies that an export cursor decodes and -// resumes inside the supplied key interval. +// resumes inside the supplied key interval. Skipped-key cursors are accepted +// only when they describe a key outside the interval. func ValidateExportCursorForRange(cursor, startKey, endKey []byte) error { + pos, err := decodeExportCursor(cursor) + if err != nil { + return err + } + return validateExportCursorPositionForRange(pos, startKey, endKey) +} + +// ValidatePromotionCursorForRange verifies a promotion cursor before it is +// proposed to Raft. Promotion scans emit only accepted positions, so callers +// must not resume from sparse-scan-only cursor tags. +func ValidatePromotionCursorForRange(cursor, startKey, endKey []byte) error { pos, err := decodeExportCursor(cursor) if err != nil { return err @@ -101,6 +113,23 @@ func ValidateExportCursorForRange(cursor, startKey, endKey []byte) error { if !pos.hasKey { return nil } + if pos.tag != exportCursorTagEmitted { + return errors.WithStack(ErrInvalidExportCursor) + } + return validateExportCursorPositionForRange(pos, startKey, endKey) +} + +func validateExportCursorPositionForRange(pos exportCursorPosition, startKey, endKey []byte) error { + if !pos.hasKey { + return nil + } + if pos.tag == exportCursorTagSkippedKey { + opts := ExportVersionsOptions{StartKey: startKey, EndKey: endKey} + if !exportSkippedCursorOutsideRange(opts, pos.key) { + return errors.WithStack(ErrInvalidExportCursor) + } + return nil + } if startKey != nil && bytes.Compare(pos.key, startKey) < 0 { return errors.WithStack(ErrInvalidExportCursor) } diff --git a/store/migration_versions_test.go b/store/migration_versions_test.go index b4c0b964d..34cffdfeb 100644 --- a/store/migration_versions_test.go +++ b/store/migration_versions_test.go @@ -405,6 +405,54 @@ func TestExportVersionsRejectsCursorOutsideRequestedRange(t *testing.T) { }) } +func TestValidateExportCursorForRangeRejectsSkippedCursorInsideRange(t *testing.T) { + t.Parallel() + + err := ValidateExportCursorForRange( + encodeExportCursor([]byte("stage|k"), 10, exportCursorTagSkippedKey), + []byte("stage|"), + PrefixScanEnd([]byte("stage|")), + ) + require.ErrorIs(t, err, ErrInvalidExportCursor) + + err = ValidateExportCursorForRange( + encodeExportCursor([]byte("outside|k"), 10, exportCursorTagSkippedKey), + []byte("stage|"), + PrefixScanEnd([]byte("stage|")), + ) + require.NoError(t, err) +} + +func TestValidatePromotionCursorForRangeAcceptsOnlyEmittedPositions(t *testing.T) { + t.Parallel() + + prefix := []byte("stage|") + key := []byte("stage|k") + for _, tc := range []struct { + name string + cursor []byte + wantErr bool + }{ + {name: "empty cursor"}, + {name: "emitted cursor", cursor: encodeExportCursor(key, 10, exportCursorTagEmitted)}, + {name: "scanned cursor", cursor: encodeExportCursor(key, 10, exportCursorTagScanned), wantErr: true}, + {name: "pruned-key cursor", cursor: encodeExportCursor(key, 10, exportCursorTagPrunedKey), wantErr: true}, + {name: "skipped-key cursor", cursor: encodeExportCursor(key, 10, exportCursorTagSkippedKey), wantErr: true}, + {name: "emitted cursor outside range", cursor: encodeExportCursor([]byte("other|k"), 10, exportCursorTagEmitted), wantErr: true}, + } { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + + err := ValidatePromotionCursorForRange(tc.cursor, prefix, PrefixScanEnd(prefix)) + if tc.wantErr { + require.ErrorIs(t, err, ErrInvalidExportCursor) + return + } + require.NoError(t, err) + }) + } +} + func TestExportVersionsDoesNotTreatMigrationPrefixUserKeyAsMetadata(t *testing.T) { runMigrationStoreSuite(t, func(t *testing.T, st MVCCStore) { ctx := context.Background() diff --git a/store/mvcc_store.go b/store/mvcc_store.go index 9590a1c44..88817d94d 100644 --- a/store/mvcc_store.go +++ b/store/mvcc_store.go @@ -938,6 +938,7 @@ func (s *mvccStore) restoreStreamingSnapshot(r io.Reader) error { s.minRetainedTS = minRetainedTS s.migrationAcks = make(map[migrationAckID]migrationImportAck) s.migrationHLCFloors = make(map[uint64]uint64) + s.migrationPromotions = make(map[uint64]PromotionState) return nil } diff --git a/store/mvcc_store_snapshot_test.go b/store/mvcc_store_snapshot_test.go index b16cbe68b..3032a1cc3 100644 --- a/store/mvcc_store_snapshot_test.go +++ b/store/mvcc_store_snapshot_test.go @@ -61,6 +61,19 @@ func TestMVCCStore_RestoreClearsMigrationMetadata(t *testing.T) { ctx := context.Background() st := newTestMVCCStore(t) + promoter, ok := any(st).(MigrationPromoter) + require.True(t, ok) + stateReader, ok := any(st).(MigrationPromotionStateReader) + require.True(t, ok) + + prefix := []byte("stage|") + stage := func(raw string) []byte { + return append([]byte("stage|"), []byte(raw)...) + } + targetKey := func(staged []byte) ([]byte, bool) { + return bytes.TrimPrefix(staged, prefix), bytes.HasPrefix(staged, prefix) + } + require.NoError(t, st.PutAt(ctx, []byte("base"), []byte("v1"), 10, 0)) snap, err := st.Snapshot() @@ -80,13 +93,46 @@ func TestMVCCStore_RestoreClearsMigrationMetadata(t *testing.T) { require.NoError(t, err) require.Equal(t, uint64(50), floor) + require.NoError(t, st.PutAt(ctx, stage("stale"), []byte("old"), 70, 0)) + promoted, err := promoter.PromoteVersions(ctx, PromoteVersionsOptions{ + JobID: 7, + StartKey: prefix, + EndKey: PrefixScanEnd(prefix), + MaxVersions: 10, + TargetKey: targetKey, + }) + require.NoError(t, err) + require.True(t, promoted.Done) + state, ok, err := stateReader.MigrationPromotionState(ctx, 7) + require.NoError(t, err) + require.True(t, ok) + require.True(t, state.Done) + require.NoError(t, st.Restore(bytes.NewReader(raw))) floor, err = st.MigrationHLCFloor(ctx, 7) require.NoError(t, err) require.Zero(t, floor) + _, ok, err = stateReader.MigrationPromotionState(ctx, 7) + require.NoError(t, err) + require.False(t, ok) _, err = st.GetAt(ctx, []byte("imported"), 50) require.ErrorIs(t, err, ErrKeyNotFound) + require.NoError(t, st.PutAt(ctx, stage("fresh"), []byte("new"), 80, 0)) + promoted, err = promoter.PromoteVersions(ctx, PromoteVersionsOptions{ + JobID: 7, + StartKey: prefix, + EndKey: PrefixScanEnd(prefix), + MaxVersions: 10, + TargetKey: targetKey, + }) + require.NoError(t, err) + require.True(t, promoted.Done) + require.Equal(t, uint64(1), promoted.PromotedRows) + got, err := st.GetAt(ctx, []byte("fresh"), 80) + require.NoError(t, err) + require.Equal(t, []byte("new"), got) + res, err := st.ImportVersions(ctx, ImportVersionsOptions{ JobID: 7, BracketID: 3, From 40c2f527b3fd6d10d22e633112d703de373c6f6a Mon Sep 17 00:00:00 2001 From: bootjp Date: Tue, 14 Jul 2026 20:04:06 +0900 Subject: [PATCH 10/10] store: ignore stale promotion cursors --- store/migration_promote.go | 11 ++++--- store/migration_versions_test.go | 56 ++++++++++++++++++++++++++++++++ 2 files changed, 62 insertions(+), 5 deletions(-) diff --git a/store/migration_promote.go b/store/migration_promote.go index 8ee56e862..bfcd1af44 100644 --- a/store/migration_promote.go +++ b/store/migration_promote.go @@ -124,10 +124,11 @@ func (s *mvccStore) promotionStateAndCursorLocked(opts PromoteVersionsOptions) ( if opts.JobID == 0 { return PromotionState{}, opts.Cursor } - state := clonePromotionState(s.migrationPromotions[opts.JobID]) - if len(state.Cursor) == 0 { - return state, opts.Cursor + state, ok := s.migrationPromotions[opts.JobID] + if !ok { + return PromotionState{}, nil } + state = clonePromotionState(state) return state, state.Cursor } @@ -318,8 +319,8 @@ func (s *pebbleStore) pebblePromotionStateAndCursor(opts PromoteVersionsOptions) if err != nil { return PromotionState{}, nil, err } - if !ok || len(state.Cursor) == 0 { - return state, opts.Cursor, nil + if !ok { + return PromotionState{}, nil, nil } return state, state.Cursor, nil } diff --git a/store/migration_versions_test.go b/store/migration_versions_test.go index 34cffdfeb..c83ab6ca9 100644 --- a/store/migration_versions_test.go +++ b/store/migration_versions_test.go @@ -934,6 +934,62 @@ func TestPromoteVersionsMovesStagedVersionsAndDeletesStagedRows(t *testing.T) { }) } +func TestPromoteVersionsIgnoresClientCursorWhenStateMissing(t *testing.T) { + runMigrationStoreSuite(t, func(t *testing.T, st MVCCStore) { + ctx := context.Background() + promoter, ok := st.(MigrationPromoter) + require.True(t, ok) + stateReader, ok := st.(MigrationPromotionStateReader) + require.True(t, ok) + + stage := func(raw string) []byte { + return append([]byte("stage|"), []byte(raw)...) + } + targetKey := func(staged []byte) ([]byte, bool) { + return bytes.TrimPrefix(staged, []byte("stage|")), bytes.HasPrefix(staged, []byte("stage|")) + } + prefix := []byte("stage|") + + require.NoError(t, st.PutAt(ctx, stage("a"), []byte("a10"), 10, 0)) + require.NoError(t, st.PutAt(ctx, stage("z"), []byte("z20"), 20, 0)) + staleCursor := encodeExportCursor(stage("m"), 1, exportCursorTagEmitted) + + result, err := promoter.PromoteVersions(ctx, PromoteVersionsOptions{ + JobID: 202, + StartKey: prefix, + EndKey: PrefixScanEnd(prefix), + Cursor: staleCursor, + MaxVersions: 10, + TargetKey: targetKey, + }) + require.NoError(t, err) + require.True(t, result.Done) + require.Equal(t, uint64(2), result.PromotedRows) + require.Equal(t, uint64(2), result.TotalPromotedRows) + + state, ok, err := stateReader.MigrationPromotionState(ctx, 202) + require.NoError(t, err) + require.True(t, ok) + require.True(t, state.Done) + require.Equal(t, uint64(2), state.PromotedRows) + + got, err := st.GetAt(ctx, []byte("a"), 10) + require.NoError(t, err) + require.Equal(t, []byte("a10"), got) + got, err = st.GetAt(ctx, []byte("z"), 20) + require.NoError(t, err) + require.Equal(t, []byte("z20"), got) + stagedLeft, err := st.ExportVersions(ctx, ExportVersionsOptions{ + StartKey: prefix, + EndKey: PrefixScanEnd(prefix), + MaxVersions: 10, + }) + require.NoError(t, err) + require.True(t, stagedLeft.Done) + require.Empty(t, stagedLeft.Versions) + }) +} + func TestPebblePromoteVersionsAdvancesLastCommitTS(t *testing.T) { ctx := context.Background() dir := t.TempDir()