diff --git a/adapter/internal.go b/adapter/internal.go index 85b7caf29..30029da94 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" @@ -35,6 +37,12 @@ func WithInternalMigrationProposer(proposer raftengine.Proposer) InternalOption } } +func WithInternalMigrationPromoteGate(gate func(context.Context) error) InternalOption { + return func(i *Internal) { + i.migrationPromoteGate = gate + } +} + func WithInternalRouteEngine(engine *distribution.Engine) InternalOption { return func(i *Internal) { i.routeEngine = engine @@ -43,10 +51,11 @@ func WithInternalRouteEngine(engine *distribution.Engine) InternalOption { 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) @@ -55,14 +64,15 @@ func NewInternalWithEngine(txm kv.Transactional, leader raftengine.LeaderView, c } type Internal struct { - leader raftengine.LeaderView - transactionManager kv.Transactional - clock *kv.HLC - tsAllocator kv.TimestampAllocator - relay *RedisPubSubRelay - store store.MVCCStore - migrationProposer raftengine.Proposer - routeEngine *distribution.Engine + leader raftengine.LeaderView + transactionManager kv.Transactional + clock *kv.HLC + tsAllocator kv.TimestampAllocator + relay *RedisPubSubRelay + store store.MVCCStore + migrationProposer raftengine.Proposer + migrationPromoteGate func(context.Context) error + routeEngine *distribution.Engine pb.UnimplementedInternalServer } @@ -210,6 +220,77 @@ func validateImportRangeVersionsRequest(req *pb.ImportRangeVersionsRequest) erro return nil } +func (i *Internal) PromoteStagedVersions(ctx context.Context, req *pb.PromoteStagedVersionsRequest) (*pb.PromoteStagedVersionsResponse, error) { + if req == nil { + return nil, errors.WithStack(status.Error(codes.InvalidArgument, "promote staged versions request is nil")) + } + if req.GetJobId() == 0 { + return nil, errors.WithStack(status.Error(codes.InvalidArgument, "promote staged versions job_id is required")) + } + if i.migrationProposer == nil { + return nil, errors.WithStack(status.Error(codes.FailedPrecondition, "migration promote proposer is not configured")) + } + if err := i.verifyInternalLeader(ctx); err != nil { + return nil, err + } + if err := i.verifyMigrationPromoteEnabled(ctx); err != nil { + return nil, err + } + if err := validatePromoteStagedVersionsRequest(req); err != nil { + return nil, errors.WithStack(err) + } + result, err := i.proposeMigrationPromote(ctx, req) + if err != nil { + return nil, errors.WithStack(err) + } + if i.clock != nil && result.MaxPromotedTS > 0 { + i.clock.Observe(result.MaxPromotedTS) + } + return &pb.PromoteStagedVersionsResponse{ + NextCursor: result.NextCursor, + Done: result.Done, + PromotedRows: result.PromotedRows, + MaxPromotedTs: result.MaxPromotedTS, + }, nil +} + +func (i *Internal) verifyMigrationPromoteEnabled(ctx context.Context) error { + if i.migrationPromoteGate == nil { + return nil + } + if err := i.migrationPromoteGate(ctx); err != nil { + return errors.WithStack(err) + } + return nil +} + +func validatePromoteStagedVersionsRequest(req *pb.PromoteStagedVersionsRequest) error { + prefix := distribution.MigrationStagedDataKeyPrefix(req.GetJobId()) + if err := store.ValidatePromotionCursorForRange(req.GetCursor(), prefix, store.PrefixScanEnd(prefix)); err != nil { + if errors.Is(err, store.ErrInvalidExportCursor) { + return errors.WithStack(status.Error(codes.InvalidArgument, store.ErrInvalidExportCursor.Error())) + } + return errors.WithStack(status.Errorf(codes.Internal, "validate promote cursor: %v", err)) + } + return nil +} + +func defaultMigrationPromoteGate(context.Context) error { + if migrationPromoteOpcodeEnabledFromEnv() { + return nil + } + return errors.WithStack(status.Error(codes.FailedPrecondition, "migration promote opcode is disabled; enable after every voter is running a build that supports migration promotion")) +} + +func migrationPromoteOpcodeEnabledFromEnv() bool { + switch strings.ToLower(strings.TrimSpace(os.Getenv("ELASTICKV_ENABLE_MIGRATION_PROMOTE_OPCODE"))) { + case "1", "true", "yes", "on": + return true + default: + return false + } +} + func (i *Internal) verifyInternalLeader(ctx context.Context) error { if i.leader.State() != raftengine.StateLeader { return errors.WithStack(ErrNotLeader) @@ -225,14 +306,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: @@ -247,6 +325,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 05e74e76a..519b8647b 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" @@ -59,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 @@ -71,6 +79,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-- { @@ -364,3 +381,120 @@ func TestInternalImportRangeVersionsRejectsMissingIdentifiers(t *testing.T) { require.Error(t, err) require.Equal(t, codes.InvalidArgument, status.Code(err)) } + +func TestInternalPromoteStagedVersionsRejectsWhenOpcodeGateClosed(t *testing.T) { + t.Parallel() + + ctx := context.Background() + st := store.NewMVCCStore() + clock := kv.NewHLC() + proposer := &applyingMigrationProposer{ + fsm: kv.NewKvFSMWithHLC(st, clock), + } + internal := NewInternalWithEngine(nil, mockInternalLeader{}, clock, nil, + WithInternalStore(st), + WithInternalMigrationProposer(proposer), + WithInternalMigrationPromoteGate(func(context.Context) error { + return status.Error(codes.FailedPrecondition, "migration promote disabled for test") + }), + ) + + resp, err := internal.PromoteStagedVersions(ctx, &pb.PromoteStagedVersionsRequest{ + JobId: 7, + MaxVersions: 10, + }) + require.Nil(t, resp) + require.Error(t, err) + require.Equal(t, codes.FailedPrecondition, status.Code(err)) + require.Equal(t, uint64(0), proposer.calls) +} + +func TestInternalPromoteStagedVersionsRejectsInvalidCursorBeforePropose(t *testing.T) { + t.Parallel() + + ctx := context.Background() + for _, tc := range []struct { + name string + cursor []byte + }{ + {name: "malformed cursor", cursor: []byte{0xff}}, + { + name: "cursor outside job staged prefix", + cursor: encodeTestExportCursor(distribution.MigrationStagedDataKey(8, []byte("k")), 30, testExportCursorTagEmitted), + }, + { + name: "scanned cursor inside staged prefix", + cursor: encodeTestExportCursor(distribution.MigrationStagedDataKey(7, []byte("k")), 31, testExportCursorTagScanned), + }, + { + name: "pruned-key cursor inside staged prefix", + cursor: encodeTestExportCursor(distribution.MigrationStagedDataKey(7, []byte("k")), 32, testExportCursorTagPrunedKey), + }, + { + name: "skipped-key cursor inside staged prefix", + cursor: encodeTestExportCursor(distribution.MigrationStagedDataKey(7, []byte("k")), 33, testExportCursorTagSkippedKey), + }, + } { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + + st := store.NewMVCCStore() + clock := kv.NewHLC() + proposer := &applyingMigrationProposer{ + fsm: kv.NewKvFSMWithHLC(st, clock), + } + internal := NewInternalWithEngine(nil, mockInternalLeader{}, clock, nil, + WithInternalStore(st), + WithInternalMigrationProposer(proposer), + WithInternalMigrationPromoteGate(func(context.Context) error { return nil }), + ) + + resp, err := internal.PromoteStagedVersions(ctx, &pb.PromoteStagedVersionsRequest{ + JobId: 7, + Cursor: tc.cursor, + MaxVersions: 10, + }) + require.Nil(t, resp) + require.Error(t, err) + require.Equal(t, codes.InvalidArgument, status.Code(err)) + require.ErrorContains(t, err, store.ErrInvalidExportCursor.Error()) + require.Equal(t, uint64(0), proposer.calls) + }) + } +} + +func TestInternalPromoteStagedVersionsAppliesStoreBatch(t *testing.T) { + t.Parallel() + + ctx := context.Background() + st := store.NewMVCCStore() + clock := kv.NewHLC() + proposer := &applyingMigrationProposer{ + fsm: kv.NewKvFSMWithHLC(st, clock), + } + internal := NewInternalWithEngine(nil, mockInternalLeader{}, clock, nil, + WithInternalStore(st), + WithInternalMigrationProposer(proposer), + WithInternalMigrationPromoteGate(func(context.Context) error { return nil }), + ) + + staged := distribution.MigrationStagedDataKey(7, []byte("k")) + require.NoError(t, st.PutAt(ctx, staged, []byte("v"), 30, 0)) + + resp, err := internal.PromoteStagedVersions(ctx, &pb.PromoteStagedVersionsRequest{ + JobId: 7, + MaxVersions: 10, + }) + require.NoError(t, err) + require.True(t, resp.GetDone()) + require.Equal(t, uint64(1), resp.GetPromotedRows()) + require.Equal(t, uint64(30), resp.GetMaxPromotedTs()) + require.Equal(t, uint64(1), proposer.calls) + + got, err := st.GetAt(ctx, []byte("k"), 30) + require.NoError(t, err) + require.Equal(t, []byte("v"), got) + _, err = st.GetAt(ctx, staged, 30) + require.ErrorIs(t, err, store.ErrKeyNotFound) + require.GreaterOrEqual(t, clock.Current(), uint64(30)) +} diff --git a/adapter/redis_lua_context.go b/adapter/redis_lua_context.go index 4a57051b5..447423ee4 100644 --- a/adapter/redis_lua_context.go +++ b/adapter/redis_lua_context.go @@ -478,20 +478,26 @@ func (c *luaScriptContext) keyType(key []byte) (redisValueType, error) { if err != nil { return redisTypeNone, err } - if typ == redisTypeNone && len(c.negativeType) < maxNegativeTypeCacheEntries { - // Pin the absence result for the rest of this Eval so repeated - // BullMQ-style polling of a missing key (e.g. a "delayed" zset) - // does not re-run the ~8-seek rawKeyTypeAt probe on every - // redis.call. - // - // Bounded to keep adversarial scripts from growing the map - // unboundedly; once full, subsequent misses correctly fall - // through to the server probe without caching. - c.negativeType[string(key)] = true + if typ == redisTypeNone { + c.rememberNegativeType(key) } return typ, nil } +func (c *luaScriptContext) rememberNegativeType(key []byte) { + if len(c.negativeType) >= maxNegativeTypeCacheEntries { + return + } + // Pin the absence result for the rest of this Eval so repeated + // BullMQ-style polling of a missing key (e.g. a "delayed" zset) does + // not re-run the ~8-seek rawKeyTypeAt probe on every redis.call. + // + // Bounded to keep adversarial scripts from growing the map + // unboundedly; once full, subsequent misses correctly fall through to + // the server probe without caching. + c.negativeType[string(key)] = true +} + func (c *luaScriptContext) ensureKeyNotExpired(key []byte) error { ttl, err := c.loadTTL(key) if err != nil { diff --git a/adapter/redis_lua_negative_type_cache_test.go b/adapter/redis_lua_negative_type_cache_test.go index 6253f7e30..f44aec35a 100644 --- a/adapter/redis_lua_negative_type_cache_test.go +++ b/adapter/redis_lua_negative_type_cache_test.go @@ -149,44 +149,31 @@ func TestLuaNegativeTypeCache_SingleProbePerKey(t *testing.T) { // itself stays bounded. func TestLuaNegativeTypeCache_BoundedSize(t *testing.T) { t.Parallel() - nodes, _, _ := createNode(t, 3) - defer shutdown(nodes) - ctx := context.Background() - sc, err := newLuaScriptContext(ctx, nodes[0].redisServer) - require.NoError(t, err) - defer sc.Close() + sc := &luaScriptContext{negativeType: map[string]bool{}} // Probe cap+overflow unique missing keys. Each probe is a miss // (redisTypeNone); only the first `cap` should be memoized. const overflow = 50 total := maxNegativeTypeCacheEntries + overflow for i := 0; i < total; i++ { - typ, kerr := sc.keyType([]byte(fmt.Sprintf("lua:neg:cap:%d", i))) - require.NoError(t, kerr) - require.Equal(t, redisTypeNone, typ) + sc.rememberNegativeType([]byte(fmt.Sprintf("lua:neg:cap:%d", i))) } require.Equal(t, maxNegativeTypeCacheEntries, len(sc.negativeType), "negativeType map must be capped at maxNegativeTypeCacheEntries") - // Each unique key above required exactly one probe on first access. - require.Equal(t, total, sc.keyTypeProbeCount, - "each unique key must have triggered exactly one server probe") - - // Re-probing one of the first `cap` keys must hit the cache - // (no additional server probe). Re-probing an overflow key must - // miss the cache and issue another server probe. + // One of the first `cap` keys must be cached. An overflow key must + // remain uncached so keyType falls back to a server probe in real Eval + // execution while the map size stays bounded. cachedKey := []byte("lua:neg:cap:0") - _, kerr := sc.keyType(cachedKey) - require.NoError(t, kerr) - require.Equal(t, total, sc.keyTypeProbeCount, - "a key inserted before the cap must remain cached") + typ, ok := sc.cachedType(cachedKey) + require.True(t, ok, "a key inserted before the cap must remain cached") + require.Equal(t, redisTypeNone, typ) overflowKey := []byte(fmt.Sprintf("lua:neg:cap:%d", maxNegativeTypeCacheEntries+1)) - _, kerr = sc.keyType(overflowKey) - require.NoError(t, kerr) - require.Equal(t, total+1, sc.keyTypeProbeCount, - "a key probed after the cap was reached must fall back to the server probe") + _, ok = sc.cachedType(overflowKey) + require.False(t, ok, "a key probed after the cap must fall back to the server probe") + sc.rememberNegativeType(overflowKey) require.Equal(t, maxNegativeTypeCacheEntries, len(sc.negativeType), "fallback probe must NOT grow the bounded cache") } diff --git a/kv/fsm.go b/kv/fsm.go index 771f359c5..bd2edd712 100644 --- a/kv/fsm.go +++ b/kv/fsm.go @@ -371,6 +371,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: @@ -406,6 +408,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..383426ae2 --- /dev/null +++ b/kv/fsm_migration_promote.go @@ -0,0 +1,98 @@ +package kv + +import ( + "context" + + "github.com/bootjp/elastickv/distribution" + pb "github.com/bootjp/elastickv/proto" + "github.com/bootjp/elastickv/store" + "github.com/cockroachdb/errors" + "google.golang.org/protobuf/proto" +) + +const ( + defaultMigrationPromoteMaxVersions = 1024 + defaultMigrationPromoteMaxBytes = 4 << 20 + defaultMigrationPromoteMaxScannedBytes = defaultMigrationPromoteMaxBytes * 4 +) + +var ErrMigrationPromoteApply = errors.New("migration promote: FSM apply failed; halting apply") + +// MarshalMigrationPromoteCommand encodes a target-group staged-data promotion +// chunk as a Raft FSM command. +func MarshalMigrationPromoteCommand(req *pb.PromoteStagedVersionsRequest) ([]byte, error) { + if req == nil { + return nil, errors.WithStack(ErrInvalidRequest) + } + b, err := proto.Marshal(req) + if err != nil { + return nil, errors.WithStack(err) + } + if len(b) >= maxMarshaledCommandSize { + return nil, errors.New("marshaled migration promote request too large") + } + return prependByte(raftEncodeMigrationPromote, b), nil +} + +func (f *kvFSM) applyMigrationPromote(ctx context.Context, data []byte) any { + req := &pb.PromoteStagedVersionsRequest{} + if err := proto.Unmarshal(data, req); err != nil { + return haltErr(errors.Wrap(errors.Mark(err, ErrMigrationPromoteApply), "kv/fsm: decode migration promote")) + } + promoter, ok := f.store.(store.MigrationPromoter) + if !ok { + return haltErr(errors.Wrap(errors.Mark(store.ErrNotSupported, ErrMigrationPromoteApply), "kv/fsm: migration promote store")) + } + result, err := promoter.PromoteVersions(ctx, migrationPromoteOptionsFromProto(req, f.pendingApplyIdx)) + if err != nil { + if isMigrationPromoteOrdinaryApplyError(err) { + return errors.Wrap(err, "kv/fsm: apply migration promote") + } + return haltErr(errors.Wrap(errors.Mark(err, ErrMigrationPromoteApply), "kv/fsm: apply migration promote")) + } + if f.hlc != nil && result.MaxPromotedTS > 0 { + f.hlc.Observe(result.MaxPromotedTS) + } + return result +} + +func migrationPromoteOptionsFromProto(req *pb.PromoteStagedVersionsRequest, appliedIndex uint64) store.PromoteVersionsOptions { + maxVersions := int(req.GetMaxVersions()) + if maxVersions <= 0 { + maxVersions = defaultMigrationPromoteMaxVersions + } + maxBytes := req.GetMaxBytes() + if maxBytes == 0 { + maxBytes = defaultMigrationPromoteMaxBytes + } + maxScannedBytes := req.GetMaxScannedBytes() + if maxScannedBytes == 0 { + maxScannedBytes = defaultMigrationPromoteMaxScannedBytes + } + prefix := distribution.MigrationStagedDataKeyPrefix(req.GetJobId()) + return store.PromoteVersionsOptions{ + JobID: req.GetJobId(), + AppliedIndex: appliedIndex, + StartKey: prefix, + EndKey: store.PrefixScanEnd(prefix), + Cursor: req.GetCursor(), + MaxVersions: maxVersions, + MaxBytes: maxBytes, + MaxScannedBytes: maxScannedBytes, + TargetKey: migrationPromoteTargetKey(req.GetJobId()), + } +} + +func isMigrationPromoteOrdinaryApplyError(err error) bool { + return errors.Is(err, store.ErrInvalidExportCursor) +} + +func migrationPromoteTargetKey(jobID uint64) func([]byte) ([]byte, bool) { + return func(stagedKey []byte) ([]byte, bool) { + gotJobID, rawKey, ok := distribution.MigrationStagedDataKeyParts(stagedKey) + if !ok || gotJobID != jobID { + return nil, false + } + return rawKey, true + } +} diff --git a/kv/fsm_migration_promote_test.go b/kv/fsm_migration_promote_test.go new file mode 100644 index 000000000..192b4af1a --- /dev/null +++ b/kv/fsm_migration_promote_test.go @@ -0,0 +1,96 @@ +package kv + +import ( + "context" + "testing" + + "github.com/bootjp/elastickv/distribution" + pb "github.com/bootjp/elastickv/proto" + "github.com/bootjp/elastickv/store" + "github.com/cockroachdb/errors" + "github.com/stretchr/testify/require" +) + +func TestMigrationPromoteTargetKeyRestoresRawKey(t *testing.T) { + t.Parallel() + + targetKey := migrationPromoteTargetKey(9) + raw, ok := targetKey(distribution.MigrationStagedDataKey(9, []byte("user|k"))) + require.True(t, ok) + require.Equal(t, []byte("user|k"), raw) + + _, ok = targetKey(distribution.MigrationStagedDataKey(10, []byte("user|k"))) + require.False(t, ok) +} + +func TestApplyMigrationPromoteMovesStagedVersions(t *testing.T) { + t.Parallel() + + ctx := context.Background() + st := store.NewMVCCStore() + hlc := NewHLC() + fsm := &kvFSM{store: st, hlc: hlc} + staged := distribution.MigrationStagedDataKey(9, []byte("user|k")) + require.NoError(t, st.PutAt(ctx, staged, []byte("v10"), 10, 0)) + require.NoError(t, st.DeleteAt(ctx, staged, 20)) + + cmd, err := MarshalMigrationPromoteCommand(&pb.PromoteStagedVersionsRequest{ + JobId: 9, + MaxVersions: 10, + }) + require.NoError(t, err) + applied := fsm.Apply(cmd) + result, ok := applied.(store.PromoteVersionsResult) + require.True(t, ok, "got %T: %v", applied, applied) + require.True(t, result.Done) + require.Equal(t, uint64(2), result.PromotedRows) + require.Equal(t, uint64(2), result.TotalPromotedRows) + require.Equal(t, uint64(20), result.MaxPromotedTS) + require.GreaterOrEqual(t, hlc.Current(), uint64(20)) + stateReader, ok := st.(store.MigrationPromotionStateReader) + require.True(t, ok) + state, ok, err := stateReader.MigrationPromotionState(ctx, 9) + require.NoError(t, err) + require.True(t, ok) + require.True(t, state.Done) + require.Equal(t, uint64(2), state.PromotedRows) + require.Equal(t, uint64(20), state.MaxPromotedTS) + + got, err := st.GetAt(ctx, []byte("user|k"), 10) + require.NoError(t, err) + require.Equal(t, []byte("v10"), got) + _, err = st.GetAt(ctx, []byte("user|k"), 20) + require.ErrorIs(t, err, store.ErrKeyNotFound) + left, err := st.ExportVersions(ctx, store.ExportVersionsOptions{ + StartKey: distribution.MigrationStagedDataKeyPrefix(9), + EndKey: store.PrefixScanEnd(distribution.MigrationStagedDataKeyPrefix(9)), + MaxVersions: 10, + }) + require.NoError(t, err) + require.Empty(t, left.Versions) +} + +func TestApplyMigrationPromoteMalformedPayloadHalts(t *testing.T) { + t.Parallel() + + fsm := &kvFSM{store: store.NewMVCCStore()} + err := haltApplyOf(fsm.Apply([]byte{raftEncodeMigrationPromote, 0xff, 0xff})) + require.True(t, errors.Is(err, ErrMigrationPromoteApply), "got %v", err) +} + +func TestApplyMigrationPromoteInvalidCursorReturnsOrdinaryError(t *testing.T) { + t.Parallel() + + fsm := &kvFSM{store: store.NewMVCCStore()} + cmd, err := MarshalMigrationPromoteCommand(&pb.PromoteStagedVersionsRequest{ + Cursor: []byte{0xff}, + MaxVersions: 10, + }) + require.NoError(t, err) + resp := fsm.Apply(cmd) + require.Nil(t, haltApplyOf(resp)) + err, ok := resp.(error) + require.True(t, ok, "got %T: %v", resp, resp) + require.ErrorIs(t, err, store.ErrInvalidExportCursor) + require.False(t, errors.Is(err, ErrMigrationPromoteApply)) +} diff --git a/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 daa5cf68e..c114b24ce 100644 --- a/store/lsm_migration.go +++ b/store/lsm_migration.go @@ -10,6 +10,13 @@ import ( ) func (s *pebbleStore) ExportVersions(ctx context.Context, opts ExportVersionsOptions) (ExportVersionsResult, error) { + s.dbMu.RLock() + defer s.dbMu.RUnlock() + + return s.exportVersionsLocked(ctx, opts) +} + +func (s *pebbleStore) exportVersionsLocked(ctx context.Context, opts ExportVersionsOptions) (ExportVersionsResult, error) { opts = normalizeExportVersionsOptions(opts) pos, err := decodeExportCursorForOptions(opts) if err != nil { @@ -19,9 +26,6 @@ func (s *pebbleStore) ExportVersions(ctx context.Context, opts ExportVersionsOpt return ExportVersionsResult{Done: true}, nil } - s.dbMu.RLock() - defer s.dbMu.RUnlock() - iter, err := s.db.NewIter(pebbleExportIterOptions(opts)) if err != nil { return ExportVersionsResult{}, errors.WithStack(err) diff --git a/store/lsm_store_applied_index_test.go b/store/lsm_store_applied_index_test.go index ff00ea950..4d4ecd4da 100644 --- a/store/lsm_store_applied_index_test.go +++ b/store/lsm_store_applied_index_test.go @@ -1,6 +1,7 @@ package store import ( + "bytes" "context" "os" "testing" @@ -96,6 +97,64 @@ func TestApplyMutationsRaftAt_BundlesMetaAppliedIndex(t *testing.T) { require.Equal(t, []byte("v1"), val) } +func TestPromoteVersions_BundlesMetaAppliedIndex(t *testing.T) { + ctx := context.Background() + st := newApplyIndexPebbleStore(t) + ps := pebbleStoreApplied(t, st) + + stage := func(raw string) []byte { + return append([]byte("stage|"), []byte(raw)...) + } + targetKey := func(staged []byte) ([]byte, bool) { + return bytes.TrimPrefix(staged, []byte("stage|")), bytes.HasPrefix(staged, []byte("stage|")) + } + prefix := []byte("stage|") + + require.NoError(t, ps.PutAt(ctx, stage("k"), []byte("v10"), 10, 0)) + + const entryIdx uint64 = 77 + result, err := ps.PromoteVersions(ctx, PromoteVersionsOptions{ + JobID: 9, + AppliedIndex: entryIdx, + StartKey: prefix, + EndKey: PrefixScanEnd(prefix), + MaxVersions: 10, + TargetKey: targetKey, + }) + require.NoError(t, err) + require.True(t, result.Done) + require.Equal(t, uint64(1), result.PromotedRows) + + got, present, err := ps.LastAppliedIndex() + require.NoError(t, err) + require.True(t, present, "PromoteVersions must persist metaAppliedIndex") + require.Equal(t, entryIdx, got) + + val, err := ps.GetAt(ctx, []byte("k"), 10) + require.NoError(t, err) + require.Equal(t, []byte("v10"), val) + _, err = ps.GetAt(ctx, stage("k"), 10) + require.ErrorIs(t, err, ErrKeyNotFound) + + const retryEntryIdx uint64 = 78 + retry, err := ps.PromoteVersions(ctx, PromoteVersionsOptions{ + JobID: 9, + AppliedIndex: retryEntryIdx, + StartKey: prefix, + EndKey: PrefixScanEnd(prefix), + MaxVersions: 10, + TargetKey: targetKey, + }) + require.NoError(t, err) + require.True(t, retry.Done) + require.Equal(t, uint64(1), retry.TotalPromotedRows) + + got, present, err = ps.LastAppliedIndex() + require.NoError(t, err) + require.True(t, present, "completed PromoteVersions retry must persist metaAppliedIndex") + require.Equal(t, retryEntryIdx, got) +} + // TestApplyMutationsRaftAt_ZeroIndexLeavesMetaKey covers the // appliedIndex==0 escape hatch — callers without a raft entry index // (test fakes, legacy ApplyMutationsRaft) MUST NOT bump the meta key. diff --git a/store/lsm_store_registration_gate_test.go b/store/lsm_store_registration_gate_test.go index 8039a9afb..f8b6dd713 100644 --- a/store/lsm_store_registration_gate_test.go +++ b/store/lsm_store_registration_gate_test.go @@ -1,6 +1,7 @@ package store import ( + "bytes" "context" "path/filepath" "testing" @@ -173,6 +174,49 @@ func TestRegistrationGate_FSMApplyPathNeverGated(t *testing.T) { mustGet(t, f.mvcc, []byte("raft-import"), 250, "imported") } +func TestRegistrationGate_PromoteVersionsNeverGated(t *testing.T) { + t.Parallel() + ctx := context.Background() + registered := true + f := newRegGateStore(t, ®istered) + stage := func(raw string) []byte { + return append([]byte("stage|"), []byte(raw)...) + } + targetKey := func(staged []byte) ([]byte, bool) { + return bytes.TrimPrefix(staged, []byte("stage|")), bytes.HasPrefix(staged, []byte("stage|")) + } + + if err := f.mvcc.PutAt(ctx, stage("promote"), []byte("value"), 100, 0); err != nil { + t.Fatalf("seed staged PutAt: %v", err) + } + registered = false + if err := f.mvcc.PutAt(ctx, []byte("direct"), []byte("blocked"), 110, 0); !errors.Is(err, ErrWriterNotRegistered) { + t.Fatalf("direct PutAt pre-registration: got %v, want ErrWriterNotRegistered", err) + } + promoter, ok := f.mvcc.(MigrationPromoter) + if !ok { + t.Fatalf("expected MigrationPromoter, got %T", f.mvcc) + } + + result, err := promoter.PromoteVersions(ctx, PromoteVersionsOptions{ + JobID: 11, + StartKey: []byte("stage|"), + EndKey: PrefixScanEnd([]byte("stage|")), + MaxVersions: 10, + TargetKey: targetKey, + }) + if err != nil { + t.Fatalf("PromoteVersions pre-registration: %v", err) + } + if !result.Done || result.PromotedRows != 1 { + t.Fatalf("PromoteVersions result = %+v, want done with one promoted row", result) + } + mustGet(t, f.mvcc, []byte("promote"), 150, "value") + if _, err := f.mvcc.GetAt(ctx, stage("promote"), 150); !errors.Is(err, ErrKeyNotFound) { + t.Fatalf("staged version after promotion: got %v, want ErrKeyNotFound", err) + } +} + // TestRegistrationGate_NotEncryptingIsUngated confirms the gate is only // 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/lsm_store_sync_mode_test.go b/store/lsm_store_sync_mode_test.go index e5e065bfd..29349b997 100644 --- a/store/lsm_store_sync_mode_test.go +++ b/store/lsm_store_sync_mode_test.go @@ -1,6 +1,7 @@ package store import ( + "bytes" "context" "testing" @@ -116,6 +117,65 @@ func TestDirectApplyWriteOpts_AlwaysSync(t *testing.T) { }) } +func TestPromoteVersionsWriteOptsFollowApplyContext(t *testing.T) { + t.Run("raft-applied promotion observes nosync", func(t *testing.T) { + ps := newPebbleStoreWithFSMApplyWriteOptsForTest(t, t.TempDir(), pebble.NoSync, fsmSyncModeNoSync) + defer ps.Close() + + require.Same(t, pebble.NoSync, ps.promotionWriteOpts(123), + "promotion with an applied index must use raft apply write options") + }) + + t.Run("direct promotion stays sync", func(t *testing.T) { + ps := newPebbleStoreWithFSMApplyWriteOptsForTest(t, t.TempDir(), pebble.NoSync, fsmSyncModeNoSync) + defer ps.Close() + + require.Same(t, pebble.Sync, ps.promotionWriteOpts(0), + "promotion without an applied index has no raft durability backstop") + }) +} + +func TestPromoteVersionsRaftNoSyncFunctionalEquivalence(t *testing.T) { + dir := t.TempDir() + ps := newPebbleStoreWithFSMApplyWriteOptsForTest(t, dir, pebble.NoSync, fsmSyncModeNoSync) + defer ps.Close() + + ctx := context.Background() + stage := func(raw string) []byte { + return append([]byte("stage|"), []byte(raw)...) + } + targetKey := func(staged []byte) ([]byte, bool) { + return bytes.TrimPrefix(staged, []byte("stage|")), bytes.HasPrefix(staged, []byte("stage|")) + } + prefix := []byte("stage|") + + require.NoError(t, ps.PutAt(ctx, stage("k"), []byte("v10"), 10, 0)) + + const entryIdx uint64 = 88 + result, err := ps.PromoteVersions(ctx, PromoteVersionsOptions{ + JobID: 12, + AppliedIndex: entryIdx, + StartKey: prefix, + EndKey: PrefixScanEnd(prefix), + MaxVersions: 10, + TargetKey: targetKey, + }) + require.NoError(t, err) + require.True(t, result.Done) + require.Equal(t, uint64(1), result.PromotedRows) + + val, err := ps.GetAt(ctx, []byte("k"), 10) + require.NoError(t, err) + require.Equal(t, []byte("v10"), val) + _, err = ps.GetAt(ctx, stage("k"), 10) + require.ErrorIs(t, err, ErrKeyNotFound) + + got, present, err := ps.LastAppliedIndex() + require.NoError(t, err) + require.True(t, present) + require.Equal(t, entryIdx, got) +} + // TestDirectApplyMutations_NoSyncConfigured_StillWritesDurably is the // functional twin of TestDirectApplyWriteOpts_AlwaysSync: it exercises // the public ApplyMutations and DeletePrefixAt entry points with a diff --git a/store/migration_promote.go b/store/migration_promote.go new file mode 100644 index 000000000..d8a8bc27b --- /dev/null +++ b/store/migration_promote.go @@ -0,0 +1,496 @@ +package store + +import ( + "bytes" + "context" + "encoding/binary" + + "github.com/cockroachdb/errors" + "github.com/cockroachdb/pebble/v2" +) + +const ( + migrationPromotionDoneFlag byte = 1 + migrationPromotionStateVersion2Flag byte = 1 << 7 +) + +type promotedVersion struct { + staged MVCCVersion + target MVCCVersion +} + +func validatePromoteVersionsOptions(opts PromoteVersionsOptions) error { + if opts.TargetKey == nil { + return errors.New("migration promote target key mapper is required") + } + return nil +} + +func promotedVersionsFromStaged(opts PromoteVersionsOptions, versions []MVCCVersion) ([]promotedVersion, PromoteVersionsResult, error) { + out := make([]promotedVersion, 0, len(versions)) + result := PromoteVersionsResult{PromotedRows: uint64(len(versions))} //nolint:gosec // len is bounded by MaxVersions. + for _, staged := range versions { + targetKey, ok := opts.TargetKey(staged.Key) + if !ok { + return nil, PromoteVersionsResult{}, errors.WithStack(errors.Newf("migration promote target key rejected staged key %q", string(staged.Key))) + } + target := MVCCVersion{ + Key: bytes.Clone(targetKey), + CommitTS: staged.CommitTS, + Tombstone: staged.Tombstone, + Value: bytes.Clone(staged.Value), + KeyFamily: staged.KeyFamily, + ExpireAt: staged.ExpireAt, + } + if err := validateImportVersion(target); err != nil { + return nil, PromoteVersionsResult{}, err + } + result.PromotedBytes += versionExportSize(target.Key, len(target.Value)) + if target.CommitTS > result.MaxPromotedTS { + result.MaxPromotedTS = target.CommitTS + } + out = append(out, promotedVersion{ + staged: staged, + target: target, + }) + } + return out, result, nil +} + +func (s *mvccStore) PromoteVersions(ctx context.Context, opts PromoteVersionsOptions) (PromoteVersionsResult, error) { + if err := validatePromoteVersionsOptions(opts); err != nil { + return PromoteVersionsResult{}, err + } + if opts.MaxVersions <= 0 { + return PromoteVersionsResult{Done: true}, nil + } + + s.mtx.Lock() + defer s.mtx.Unlock() + + state, cursor := s.promotionStateAndCursorLocked(opts) + if opts.JobID != 0 && state.Done { + return PromoteVersionsResult{Done: true, TotalPromotedRows: state.PromotedRows, MaxPromotedTS: state.MaxPromotedTS}, nil + } + exported, toPromote, promoted, err := s.planMemoryPromotionLocked(ctx, opts, cursor) + if err != nil { + return PromoteVersionsResult{}, err + } + s.applyMemoryPromotionLocked(toPromote) + result, updatedState := finishPromotionResult(opts, state, exported, promoted) + if result.MaxPromotedTS > s.lastCommitTS { + s.lastCommitTS = result.MaxPromotedTS + } + if updatedState != nil { + s.migrationPromotions[opts.JobID] = clonePromotionState(*updatedState) + } + return result, nil +} + +func (s *mvccStore) planMemoryPromotionLocked( + ctx context.Context, + opts PromoteVersionsOptions, + cursor []byte, +) (ExportVersionsResult, []promotedVersion, PromoteVersionsResult, error) { + exportOpts := normalizeExportVersionsOptions(ExportVersionsOptions{ + StartKey: opts.StartKey, + EndKey: opts.EndKey, + Cursor: cursor, + MaxVersions: opts.MaxVersions, + MaxBytes: opts.MaxBytes, + MaxScannedBytes: opts.MaxScannedBytes, + }) + pos, err := decodeExportCursorForOptions(exportOpts) + if err != nil { + return ExportVersionsResult{}, nil, PromoteVersionsResult{}, err + } + exported, err := s.exportMemoryVersionsLocked(ctx, exportOpts, pos) + if err != nil { + return ExportVersionsResult{}, nil, PromoteVersionsResult{}, err + } + toPromote, promoted, err := promotedVersionsFromStaged(opts, exported.Versions) + return exported, toPromote, promoted, err +} + +func (s *mvccStore) applyMemoryPromotionLocked(toPromote []promotedVersion) { + for _, version := range toPromote { + if version.target.Tombstone { + s.deleteVersionLocked(version.target.Key, version.target.CommitTS) + } else { + s.putVersionLocked(version.target.Key, version.target.Value, version.target.CommitTS, version.target.ExpireAt) + } + s.removeVersionLocked(version.staged.Key, version.staged.CommitTS) + } +} + +func (s *mvccStore) promotionStateAndCursorLocked(opts PromoteVersionsOptions) (PromotionState, []byte) { + if opts.JobID == 0 { + return PromotionState{}, opts.Cursor + } + state, ok := s.migrationPromotions[opts.JobID] + if !ok { + return PromotionState{}, nil + } + state = clonePromotionState(state) + return state, state.Cursor +} + +func (s *mvccStore) MigrationPromotionState(_ context.Context, jobID uint64) (PromotionState, bool, error) { + s.mtx.RLock() + defer s.mtx.RUnlock() + state, ok := s.migrationPromotions[jobID] + return clonePromotionState(state), ok, nil +} + +func (s *mvccStore) exportMemoryVersionsLocked(ctx context.Context, opts ExportVersionsOptions, pos exportCursorPosition) (ExportVersionsResult, error) { + result := newExportVersionsResult(opts.MaxVersions) + it := s.tree.Iterator() + if !s.seekMemoryExportStart(&it, opts.StartKey, pos) { + result.Done = true + return result, nil + } + for ok := true; ok; ok = it.Next() { + key, keyOK := it.Key().([]byte) + if err := checkExportKey(ctx, key, keyOK, opts.EndKey); err != nil { + if errors.Is(err, errExportReachedEnd) { + result.Done = true + result.NextCursor = nil + return result, nil + } + return ExportVersionsResult{}, err + } + if !keyOK { + continue + } + done, err := exportMemoryIteratorKey(ctx, opts, pos, key, it.Value(), &result) + if err != nil || !done { + return result, err + } + } + result.Done = true + result.NextCursor = nil + return result, nil +} + +func (s *mvccStore) removeVersionLocked(key []byte, commitTS uint64) bool { + existing, ok := s.tree.Get(key) + if !ok { + return false + } + versions, _ := existing.([]VersionedValue) + idx := findVersionIndex(versions, commitTS) + if idx < 0 { + return false + } + next := make([]VersionedValue, len(versions)-1) + copy(next, versions[:idx]) + copy(next[idx:], versions[idx+1:]) + if len(next) == 0 { + s.tree.Remove(key) + return true + } + s.tree.Put(bytes.Clone(key), next) + return true +} + +func findVersionIndex(versions []VersionedValue, commitTS uint64) int { + for i := range versions { + if versions[i].TS == commitTS { + return i + } + } + return -1 +} + +func (s *pebbleStore) PromoteVersions(ctx context.Context, opts PromoteVersionsOptions) (PromoteVersionsResult, error) { + if err := validatePromoteVersionsOptions(opts); err != nil { + return PromoteVersionsResult{}, err + } + if opts.MaxVersions <= 0 { + return PromoteVersionsResult{Done: true}, nil + } + + s.dbMu.RLock() + defer s.dbMu.RUnlock() + + s.applyMu.Lock() + defer s.applyMu.Unlock() + + state, cursor, err := s.pebblePromotionStateAndCursor(opts) + if err != nil { + return PromoteVersionsResult{}, err + } + writeOpts := s.promotionWriteOpts(opts.AppliedIndex) + if opts.JobID != 0 && state.Done { + result := PromoteVersionsResult{Done: true, TotalPromotedRows: state.PromotedRows, MaxPromotedTS: state.MaxPromotedTS} + return s.finishPebblePromotion(nil, opts.JobID, nil, result, opts.AppliedIndex, state.MaxPromotedTS, writeOpts) + } + opts.Cursor = cursor + exported, toPromote, promoted, err := s.planPebblePromotionLocked(ctx, opts) + if err != nil { + return PromoteVersionsResult{}, err + } + result, stateToWrite := finishPromotionResult(opts, state, exported, promoted) + return s.finishPebblePromotion( + toPromote, + opts.JobID, + stateToWrite, + result, + opts.AppliedIndex, + result.MaxPromotedTS, + writeOpts, + ) +} + +func (s *pebbleStore) promotionWriteOpts(appliedIndex uint64) *pebble.WriteOptions { + if appliedIndex > 0 { + return s.raftApplyWriteOpts() + } + return s.directApplyWriteOpts() +} + +func (s *pebbleStore) finishPebblePromotion( + toPromote []promotedVersion, + jobID uint64, + stateToWrite *PromotionState, + result PromoteVersionsResult, + appliedIndex uint64, + maxPromotedTS uint64, + writeOpts *pebble.WriteOptions, +) (PromoteVersionsResult, error) { + if len(toPromote) == 0 && stateToWrite == nil && appliedIndex == 0 && maxPromotedTS == 0 { + return result, nil + } + if err := s.commitPebblePromoteVersions( + toPromote, + jobID, + stateToWrite, + appliedIndex, + maxPromotedTS, + writeOpts, + ); err != nil { + return PromoteVersionsResult{}, err + } + return result, nil +} + +func (s *pebbleStore) planPebblePromotionLocked( + ctx context.Context, + opts PromoteVersionsOptions, +) (ExportVersionsResult, []promotedVersion, PromoteVersionsResult, error) { + exported, err := s.exportVersionsLocked(ctx, ExportVersionsOptions{ + StartKey: opts.StartKey, + EndKey: opts.EndKey, + Cursor: opts.Cursor, + MaxVersions: opts.MaxVersions, + MaxBytes: opts.MaxBytes, + MaxScannedBytes: opts.MaxScannedBytes, + }) + if err != nil { + return ExportVersionsResult{}, nil, PromoteVersionsResult{}, err + } + toPromote, promoted, err := promotedVersionsFromStaged(opts, exported.Versions) + return exported, toPromote, promoted, err +} + +func finishPromotionResult( + opts PromoteVersionsOptions, + state PromotionState, + exported ExportVersionsResult, + promoted PromoteVersionsResult, +) (PromoteVersionsResult, *PromotionState) { + promoted.NextCursor = exported.NextCursor + promoted.Done = exported.Done + promoted.ScannedBytes = exported.ScannedBytes + promoted.TotalPromotedRows = promoted.PromotedRows + if opts.JobID == 0 { + return promoted, nil + } + state.Cursor = bytes.Clone(exported.NextCursor) + state.Done = exported.Done + state.PromotedRows += promoted.PromotedRows + if promoted.MaxPromotedTS > state.MaxPromotedTS { + state.MaxPromotedTS = promoted.MaxPromotedTS + } + state.LastError = "" + promoted.TotalPromotedRows = state.PromotedRows + promoted.MaxPromotedTS = state.MaxPromotedTS + return promoted, &state +} + +func (s *pebbleStore) pebblePromotionStateAndCursor(opts PromoteVersionsOptions) (PromotionState, []byte, error) { + if opts.JobID == 0 { + return PromotionState{}, opts.Cursor, nil + } + state, ok, err := s.readPebblePromotionState(opts.JobID) + if err != nil { + return PromotionState{}, nil, err + } + if !ok { + return PromotionState{}, nil, nil + } + return state, state.Cursor, nil +} + +func (s *pebbleStore) MigrationPromotionState(_ context.Context, jobID uint64) (PromotionState, bool, error) { + s.dbMu.RLock() + defer s.dbMu.RUnlock() + return s.readPebblePromotionState(jobID) +} + +func (s *pebbleStore) readPebblePromotionState(jobID uint64) (PromotionState, bool, error) { + states, err := s.readPebblePromotionStates() + if err != nil { + return PromotionState{}, false, err + } + state, ok := states[jobID] + return clonePromotionState(state), ok, nil +} + +func (s *pebbleStore) readPebblePromotionStates() (map[uint64]PromotionState, error) { + val, closer, err := s.db.Get(migrationPromoteMetaKeyBytes) + if err != nil { + if errors.Is(err, pebble.ErrNotFound) { + return make(map[uint64]PromotionState), nil + } + return nil, errors.WithStack(err) + } + defer func() { _ = closer.Close() }() + states, ok := decodeMigrationPromotionStates(val) + if !ok { + return nil, errors.New("corrupt migration promotion state metadata") + } + return states, nil +} + +func (s *pebbleStore) stagePebblePromotionState(batch *pebble.Batch, jobID uint64, state PromotionState) error { + states, err := s.readPebblePromotionStates() + if err != nil { + return err + } + states[jobID] = clonePromotionState(state) + return errors.WithStack(batch.Set(migrationPromoteMetaKeyBytes, encodeMigrationPromotionStates(states), nil)) +} + +func (s *pebbleStore) commitPebblePromoteVersions( + versions []promotedVersion, + jobID uint64, + state *PromotionState, + appliedIndex uint64, + maxPromotedTS uint64, + writeOpts *pebble.WriteOptions, +) error { + batch := s.db.NewBatch() + defer batch.Close() + targets := make([]MVCCVersion, 0, len(versions)) + for _, version := range versions { + targets = append(targets, version.target) + } + // Promotion is replayed from the Raft FSM, so it must not fail closed on + // this node's local writer-registration state. + if err := s.applyImportVersionsBatch(batch, targets, false); err != nil { + return err + } + for _, version := range versions { + if err := batch.Delete(encodeKey(version.staged.Key, version.staged.CommitTS), nil); err != nil { + return errors.WithStack(err) + } + } + if state != nil { + if err := s.stagePebblePromotionState(batch, jobID, *state); err != nil { + return err + } + } + if appliedIndex > 0 { + if err := setPebbleUint64InBatch(batch, metaAppliedIndexBytes, appliedIndex); err != nil { + return err + } + } + return s.commitPebblePromotionBatch(batch, maxPromotedTS, writeOpts) +} + +func (s *pebbleStore) commitPebblePromotionBatch( + batch *pebble.Batch, + maxPromotedTS uint64, + writeOpts *pebble.WriteOptions, +) error { + if maxPromotedTS > 0 { + s.mtx.Lock() + defer s.mtx.Unlock() + newLastTS := s.lastCommitTS + if maxPromotedTS > newLastTS { + newLastTS = maxPromotedTS + } + if err := setPebbleUint64InBatch(batch, metaLastCommitTSBytes, newLastTS); err != nil { + return err + } + if err := batch.Commit(writeOpts); err != nil { + return errors.WithStack(err) + } + s.updateLastCommitTS(newLastTS) + return nil + } + if err := batch.Commit(writeOpts); err != nil { + return errors.WithStack(err) + } + return nil +} + +func encodePromotionState(state PromotionState) []byte { + buf := make([]byte, 0, 1+2*migrationUint64Bytes+binary.MaxVarintLen64*2+len(state.Cursor)+len(state.LastError)) + flags := migrationPromotionStateVersion2Flag + if state.Done { + flags |= migrationPromotionDoneFlag + } + buf = append(buf, flags) + buf = binary.BigEndian.AppendUint64(buf, state.PromotedRows) + buf = binary.BigEndian.AppendUint64(buf, state.MaxPromotedTS) + buf = binary.AppendUvarint(buf, lenAsUint64(len(state.Cursor))) + buf = append(buf, state.Cursor...) + buf = binary.AppendUvarint(buf, lenAsUint64(len(state.LastError))) + buf = append(buf, state.LastError...) + return buf +} + +func decodePromotionState(data []byte) (PromotionState, bool) { + if len(data) < 1+migrationUint64Bytes { + return PromotionState{}, false + } + flags := data[0] + state := PromotionState{ + Done: flags&migrationPromotionDoneFlag != 0, + PromotedRows: binary.BigEndian.Uint64(data[1 : 1+migrationUint64Bytes]), + } + rest := data[1+migrationUint64Bytes:] + if flags&migrationPromotionStateVersion2Flag != 0 { + if len(rest) < migrationUint64Bytes { + return PromotionState{}, false + } + state.MaxPromotedTS = binary.BigEndian.Uint64(rest[:migrationUint64Bytes]) + rest = rest[migrationUint64Bytes:] + } + cursorLen, n := binary.Uvarint(rest) + if n <= 0 || cursorLen > lenAsUint64(len(rest[n:])) { + return PromotionState{}, false + } + rest = rest[n:] + cursorEnd := int(cursorLen) //nolint:gosec // bounded by len(rest) above. + state.Cursor = bytes.Clone(rest[:cursorEnd]) + rest = rest[cursorEnd:] + errLen, n := binary.Uvarint(rest) + if n <= 0 || errLen != lenAsUint64(len(rest[n:])) { + return PromotionState{}, false + } + state.LastError = string(rest[n:]) + return state, true +} + +func clonePromotionState(state PromotionState) PromotionState { + state.Cursor = bytes.Clone(state.Cursor) + return state +} + +var _ MigrationPromoter = (*mvccStore)(nil) +var _ MigrationPromoter = (*pebbleStore)(nil) +var _ MigrationPromotionStateReader = (*mvccStore)(nil) +var _ MigrationPromotionStateReader = (*pebbleStore)(nil) diff --git a/store/migration_versions.go b/store/migration_versions.go index 22d9572c7..be764e360 100644 --- a/store/migration_versions.go +++ b/store/migration_versions.go @@ -18,6 +18,7 @@ const ( migrationAckMetaKey = "_migack" migrationHLCFloorMetaKey = "_mighlc" + migrationPromoteMetaKey = "_migpromote" migrationMetadataVersion = 1 migrationAckPrefix = "!migstage|ack|" @@ -29,6 +30,7 @@ const ( var ( migrationAckMetaKeyBytes = []byte(migrationAckMetaKey) migrationHLCFloorMetaKeyBytes = []byte(migrationHLCFloorMetaKey) + migrationPromoteMetaKeyBytes = []byte(migrationPromoteMetaKey) ) type exportCursorPosition struct { @@ -89,6 +91,54 @@ func decodeExportCursor(cursor []byte) (exportCursorPosition, error) { return exportCursorPosition{key: key, commitTS: commitTS, tag: tag, hasKey: true}, nil } +// ValidateExportCursorForRange verifies that an export cursor decodes and +// resumes inside the supplied key interval. Skipped-key cursors are accepted +// only when they describe a key outside the interval. +func ValidateExportCursorForRange(cursor, startKey, endKey []byte) error { + pos, err := decodeExportCursor(cursor) + if err != nil { + return err + } + return validateExportCursorPositionForRange(pos, startKey, endKey) +} + +// ValidatePromotionCursorForRange verifies a promotion cursor before it is +// proposed to Raft. Promotion scans emit only accepted positions, so callers +// must not resume from sparse-scan-only cursor tags. +func ValidatePromotionCursorForRange(cursor, startKey, endKey []byte) error { + pos, err := decodeExportCursor(cursor) + if err != nil { + return err + } + if !pos.hasKey { + return nil + } + if pos.tag != exportCursorTagEmitted { + return errors.WithStack(ErrInvalidExportCursor) + } + return validateExportCursorPositionForRange(pos, startKey, endKey) +} + +func validateExportCursorPositionForRange(pos exportCursorPosition, startKey, endKey []byte) error { + if !pos.hasKey { + return nil + } + if pos.tag == exportCursorTagSkippedKey { + opts := ExportVersionsOptions{StartKey: startKey, EndKey: endKey} + if !exportSkippedCursorOutsideRange(opts, pos.key) { + return errors.WithStack(ErrInvalidExportCursor) + } + return nil + } + if startKey != nil && bytes.Compare(pos.key, startKey) < 0 { + return errors.WithStack(ErrInvalidExportCursor) + } + if endKey != nil && bytes.Compare(pos.key, endKey) >= 0 { + return errors.WithStack(ErrInvalidExportCursor) + } + return nil +} + func decodeExportCursorForOptions(opts ExportVersionsOptions) (exportCursorPosition, error) { pos, err := decodeExportCursor(opts.Cursor) if err != nil { @@ -97,7 +147,7 @@ func decodeExportCursorForOptions(opts ExportVersionsOptions) (exportCursorPosit if err := validateExportCursorRange(opts, pos); err != nil { return exportCursorPosition{}, err } - return pos, nil + return normalizeExportCursorPositionForRange(opts, pos), nil } func validateExportCursorRange(opts ExportVersionsOptions, pos exportCursorPosition) error { @@ -124,6 +174,16 @@ func exportSkippedCursorOutsideRange(opts ExportVersionsOptions, key []byte) boo (opts.EndKey != nil && bytes.Compare(key, opts.EndKey) >= 0) } +func normalizeExportCursorPositionForRange(opts ExportVersionsOptions, pos exportCursorPosition) exportCursorPosition { + if !pos.hasKey || pos.tag != exportCursorTagSkippedKey || opts.StartKey == nil { + return pos + } + if bytes.Compare(pos.key, opts.StartKey) >= 0 { + return pos + } + return exportCursorPosition{} +} + func normalizeExportVersionsOptions(opts ExportVersionsOptions) ExportVersionsOptions { if opts.EndKey != nil && len(opts.EndKey) == 0 { opts.EndKey = nil @@ -144,7 +204,8 @@ func exportUsesSparseScanBudget(opts ExportVersionsOptions) bool { func isMigrationMetadataKey(rawKey []byte) bool { return bytes.Equal(rawKey, migrationAckMetaKeyBytes) || - bytes.Equal(rawKey, migrationHLCFloorMetaKeyBytes) + bytes.Equal(rawKey, migrationHLCFloorMetaKeyBytes) || + bytes.Equal(rawKey, migrationPromoteMetaKeyBytes) } func encodeMigrationImportAcks(acks map[migrationAckID]migrationImportAck) []byte { @@ -249,6 +310,61 @@ func decodeMigrationHLCFloors(data []byte) (map[uint64]uint64, bool) { return floors, len(rest) == 0 } +func encodeMigrationPromotionStates(states map[uint64]PromotionState) []byte { + jobIDs := make([]uint64, 0, len(states)) + for jobID := range states { + jobIDs = append(jobIDs, jobID) + } + sort.Slice(jobIDs, func(i, j int) bool { return jobIDs[i] < jobIDs[j] }) + + buf := make([]byte, 0, 1+binary.MaxVarintLen64+len(jobIDs)*(migrationUint64Bytes+binary.MaxVarintLen64)) + buf = append(buf, migrationMetadataVersion) + buf = binary.AppendUvarint(buf, lenAsUint64(len(jobIDs))) + for _, jobID := range jobIDs { + encoded := encodePromotionState(states[jobID]) + buf = binary.BigEndian.AppendUint64(buf, jobID) + buf = binary.AppendUvarint(buf, lenAsUint64(len(encoded))) + buf = append(buf, encoded...) + } + return buf +} + +func decodeMigrationPromotionStates(data []byte) (map[uint64]PromotionState, bool) { + if len(data) == 0 || data[0] != migrationMetadataVersion { + return nil, false + } + rest := data[1:] + count, n := binary.Uvarint(rest) + if n <= 0 { + return nil, false + } + rest = rest[n:] + states := make(map[uint64]PromotionState) + for i := uint64(0); i < count; i++ { + if len(rest) < migrationUint64Bytes { + return nil, false + } + jobID := binary.BigEndian.Uint64(rest[:migrationUint64Bytes]) + rest = rest[migrationUint64Bytes:] + stateLen, n := binary.Uvarint(rest) + if n <= 0 { + return nil, false + } + rest = rest[n:] + if stateLen > lenAsUint64(len(rest)) { + return nil, false + } + stateEnd := int(stateLen) //nolint:gosec // bounded by len(rest) above. + state, ok := decodePromotionState(rest[:stateEnd]) + if !ok { + return nil, false + } + states[jobID] = state + rest = rest[stateEnd:] + } + return states, len(rest) == 0 +} + func validateImportVersion(version MVCCVersion) error { if version.CommitTS == 0 { return errors.New("migration import version has zero commit_ts") diff --git a/store/migration_versions_test.go b/store/migration_versions_test.go index d76a10ee3..0fb58e4bd 100644 --- a/store/migration_versions_test.go +++ b/store/migration_versions_test.go @@ -9,6 +9,7 @@ import ( "testing" "github.com/bootjp/elastickv/internal/encryption" + "github.com/cockroachdb/pebble/v2" "github.com/stretchr/testify/require" ) @@ -404,6 +405,72 @@ func TestExportVersionsRejectsCursorOutsideRequestedRange(t *testing.T) { }) } +func TestExportVersionsSkippedCursorBeforeStartResumesAtStartKey(t *testing.T) { + runMigrationStoreSuite(t, func(t *testing.T, st MVCCStore) { + ctx := context.Background() + require.NoError(t, st.PutAt(ctx, []byte("a"), []byte("a10"), 10, 0)) + require.NoError(t, st.PutAt(ctx, []byte("m"), []byte("m20"), 20, 0)) + + res, err := st.ExportVersions(ctx, ExportVersionsOptions{ + StartKey: []byte("m"), + EndKey: []byte("z"), + Cursor: encodeExportCursor([]byte("a"), 10, exportCursorTagSkippedKey), + MaxVersions: 10, + }) + require.NoError(t, err) + require.True(t, res.Done) + require.Equal(t, []MVCCVersion{{Key: []byte("m"), CommitTS: 20, Value: []byte("m20")}}, res.Versions) + }) +} + +func TestValidateExportCursorForRangeRejectsSkippedCursorInsideRange(t *testing.T) { + t.Parallel() + + err := ValidateExportCursorForRange( + encodeExportCursor([]byte("stage|k"), 10, exportCursorTagSkippedKey), + []byte("stage|"), + PrefixScanEnd([]byte("stage|")), + ) + require.ErrorIs(t, err, ErrInvalidExportCursor) + + err = ValidateExportCursorForRange( + encodeExportCursor([]byte("outside|k"), 10, exportCursorTagSkippedKey), + []byte("stage|"), + PrefixScanEnd([]byte("stage|")), + ) + require.NoError(t, err) +} + +func TestValidatePromotionCursorForRangeAcceptsOnlyEmittedPositions(t *testing.T) { + t.Parallel() + + prefix := []byte("stage|") + key := []byte("stage|k") + for _, tc := range []struct { + name string + cursor []byte + wantErr bool + }{ + {name: "empty cursor"}, + {name: "emitted cursor", cursor: encodeExportCursor(key, 10, exportCursorTagEmitted)}, + {name: "scanned cursor", cursor: encodeExportCursor(key, 10, exportCursorTagScanned), wantErr: true}, + {name: "pruned-key cursor", cursor: encodeExportCursor(key, 10, exportCursorTagPrunedKey), wantErr: true}, + {name: "skipped-key cursor", cursor: encodeExportCursor(key, 10, exportCursorTagSkippedKey), wantErr: true}, + {name: "emitted cursor outside range", cursor: encodeExportCursor([]byte("other|k"), 10, exportCursorTagEmitted), wantErr: true}, + } { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + + err := ValidatePromotionCursorForRange(tc.cursor, prefix, PrefixScanEnd(prefix)) + if tc.wantErr { + require.ErrorIs(t, err, ErrInvalidExportCursor) + return + } + require.NoError(t, err) + }) + } +} + func TestExportVersionsDoesNotTreatMigrationPrefixUserKeyAsMetadata(t *testing.T) { runMigrationStoreSuite(t, func(t *testing.T, st MVCCStore) { ctx := context.Background() @@ -788,6 +855,279 @@ func TestImportVersionsIdempotencyAndMetadata(t *testing.T) { }) } +func TestPromoteVersionsMovesStagedVersionsAndDeletesStagedRows(t *testing.T) { + runMigrationStoreSuite(t, func(t *testing.T, st MVCCStore) { + ctx := context.Background() + promoter, ok := st.(MigrationPromoter) + require.True(t, ok) + stateReader, ok := st.(MigrationPromotionStateReader) + require.True(t, ok) + + stage := func(raw string) []byte { + return append([]byte("stage|"), []byte(raw)...) + } + targetKey := func(staged []byte) ([]byte, bool) { + return bytes.TrimPrefix(staged, []byte("stage|")), bytes.HasPrefix(staged, []byte("stage|")) + } + prefix := []byte("stage|") + + require.NoError(t, st.PutAt(ctx, []byte("k"), []byte("old"), 5, 0)) + require.NoError(t, st.PutAt(ctx, stage("k"), []byte("v10"), 10, 0)) + require.NoError(t, st.PutWithTTLAt(ctx, stage("k"), []byte("v20"), 20, 55)) + require.NoError(t, st.DeleteAt(ctx, stage("k"), 30)) + require.NoError(t, st.PutAt(ctx, stage("z"), []byte("z15"), 15, 0)) + + first, err := promoter.PromoteVersions(ctx, PromoteVersionsOptions{ + JobID: 99, + StartKey: prefix, + EndKey: PrefixScanEnd(prefix), + MaxVersions: 2, + TargetKey: targetKey, + }) + require.NoError(t, err) + require.False(t, first.Done) + require.Equal(t, uint64(2), first.PromotedRows) + require.Equal(t, uint64(2), first.TotalPromotedRows) + require.Equal(t, uint64(30), first.MaxPromotedTS) + require.NotEmpty(t, first.NextCursor) + state, ok, err := stateReader.MigrationPromotionState(ctx, 99) + require.NoError(t, err) + require.True(t, ok) + require.False(t, state.Done) + require.Equal(t, first.NextCursor, state.Cursor) + require.Equal(t, uint64(2), state.PromotedRows) + require.Equal(t, uint64(30), state.MaxPromotedTS) + + got, err := st.GetAt(ctx, []byte("k"), 25) + require.NoError(t, err) + require.Equal(t, []byte("v20"), got) + _, err = st.GetAt(ctx, []byte("k"), 35) + require.ErrorIs(t, err, ErrKeyNotFound) + + stagedLeft, err := st.ExportVersions(ctx, ExportVersionsOptions{ + StartKey: prefix, + EndKey: PrefixScanEnd(prefix), + MaxVersions: 10, + }) + require.NoError(t, err) + require.Equal(t, []MVCCVersion{ + {Key: stage("k"), CommitTS: 10, Value: []byte("v10")}, + {Key: stage("z"), CommitTS: 15, Value: []byte("z15")}, + }, stagedLeft.Versions) + + second, err := promoter.PromoteVersions(ctx, PromoteVersionsOptions{ + JobID: 99, + StartKey: prefix, + EndKey: PrefixScanEnd(prefix), + MaxVersions: 10, + TargetKey: targetKey, + }) + require.NoError(t, err) + require.True(t, second.Done) + require.Empty(t, second.NextCursor) + require.Equal(t, uint64(2), second.PromotedRows) + require.Equal(t, uint64(4), second.TotalPromotedRows) + require.Equal(t, uint64(30), second.MaxPromotedTS) + state, ok, err = stateReader.MigrationPromotionState(ctx, 99) + require.NoError(t, err) + require.True(t, ok) + require.True(t, state.Done) + require.Empty(t, state.Cursor) + require.Equal(t, uint64(4), state.PromotedRows) + require.Equal(t, uint64(30), state.MaxPromotedTS) + + retry, err := promoter.PromoteVersions(ctx, PromoteVersionsOptions{ + JobID: 99, + StartKey: prefix, + EndKey: PrefixScanEnd(prefix), + MaxVersions: 10, + TargetKey: targetKey, + }) + require.NoError(t, err) + require.True(t, retry.Done) + require.Zero(t, retry.PromotedRows) + require.Equal(t, uint64(4), retry.TotalPromotedRows) + require.Equal(t, uint64(30), retry.MaxPromotedTS) + + got, err = st.GetAt(ctx, []byte("k"), 10) + require.NoError(t, err) + require.Equal(t, []byte("v10"), got) + got, err = st.GetAt(ctx, []byte("z"), 15) + require.NoError(t, err) + require.Equal(t, []byte("z15"), got) + + stagedLeft, err = st.ExportVersions(ctx, ExportVersionsOptions{ + StartKey: prefix, + EndKey: PrefixScanEnd(prefix), + MaxVersions: 10, + }) + require.NoError(t, err) + require.True(t, stagedLeft.Done) + require.Empty(t, stagedLeft.Versions) + }) +} + +func TestPromoteVersionsIgnoresClientCursorWhenStateMissing(t *testing.T) { + runMigrationStoreSuite(t, func(t *testing.T, st MVCCStore) { + ctx := context.Background() + promoter, ok := st.(MigrationPromoter) + require.True(t, ok) + stateReader, ok := st.(MigrationPromotionStateReader) + require.True(t, ok) + + stage := func(raw string) []byte { + return append([]byte("stage|"), []byte(raw)...) + } + targetKey := func(staged []byte) ([]byte, bool) { + return bytes.TrimPrefix(staged, []byte("stage|")), bytes.HasPrefix(staged, []byte("stage|")) + } + prefix := []byte("stage|") + + require.NoError(t, st.PutAt(ctx, stage("a"), []byte("a10"), 10, 0)) + require.NoError(t, st.PutAt(ctx, stage("z"), []byte("z20"), 20, 0)) + staleCursor := encodeExportCursor(stage("m"), 1, exportCursorTagEmitted) + + result, err := promoter.PromoteVersions(ctx, PromoteVersionsOptions{ + JobID: 202, + StartKey: prefix, + EndKey: PrefixScanEnd(prefix), + Cursor: staleCursor, + MaxVersions: 10, + TargetKey: targetKey, + }) + require.NoError(t, err) + require.True(t, result.Done) + require.Equal(t, uint64(2), result.PromotedRows) + require.Equal(t, uint64(2), result.TotalPromotedRows) + + state, ok, err := stateReader.MigrationPromotionState(ctx, 202) + require.NoError(t, err) + require.True(t, ok) + require.True(t, state.Done) + require.Equal(t, uint64(2), state.PromotedRows) + require.Equal(t, uint64(20), state.MaxPromotedTS) + + got, err := st.GetAt(ctx, []byte("a"), 10) + require.NoError(t, err) + require.Equal(t, []byte("a10"), got) + got, err = st.GetAt(ctx, []byte("z"), 20) + require.NoError(t, err) + require.Equal(t, []byte("z20"), got) + stagedLeft, err := st.ExportVersions(ctx, ExportVersionsOptions{ + StartKey: prefix, + EndKey: PrefixScanEnd(prefix), + MaxVersions: 10, + }) + require.NoError(t, err) + require.True(t, stagedLeft.Done) + require.Empty(t, stagedLeft.Versions) + }) +} + +func TestPebblePromoteVersionsAdvancesLastCommitTS(t *testing.T) { + ctx := context.Background() + dir := t.TempDir() + st, err := NewPebbleStore(dir) + require.NoError(t, err) + closed := false + t.Cleanup(func() { + if !closed { + require.NoError(t, st.Close()) + } + }) + ps, ok := st.(*pebbleStore) + require.True(t, ok) + + stage := func(raw string) []byte { + return append([]byte("stage|"), []byte(raw)...) + } + targetKey := func(staged []byte) ([]byte, bool) { + return bytes.TrimPrefix(staged, []byte("stage|")), bytes.HasPrefix(staged, []byte("stage|")) + } + prefix := []byte("stage|") + + const promotedTS uint64 = 100 + require.NoError(t, ps.db.Set(encodeKey(stage("k"), promotedTS), encodeValue([]byte("v100"), false, 0, encStateCleartext), pebble.NoSync)) + require.Zero(t, ps.LastCommitTS()) + + result, err := ps.PromoteVersions(ctx, PromoteVersionsOptions{ + JobID: 101, + StartKey: prefix, + EndKey: PrefixScanEnd(prefix), + MaxVersions: 10, + TargetKey: targetKey, + }) + require.NoError(t, err) + require.True(t, result.Done) + require.Equal(t, uint64(1), result.PromotedRows) + require.Equal(t, promotedTS, result.MaxPromotedTS) + require.Equal(t, promotedTS, ps.LastCommitTS()) + state, ok, err := ps.MigrationPromotionState(ctx, 101) + require.NoError(t, err) + require.True(t, ok) + require.Equal(t, promotedTS, state.MaxPromotedTS) + + metaTS, err := readPebbleUint64(ps.db, metaLastCommitTSBytes) + require.NoError(t, err) + require.Equal(t, promotedTS, metaTS) + val, err := ps.GetAt(ctx, []byte("k"), ps.LastCommitTS()) + require.NoError(t, err) + require.Equal(t, []byte("v100"), val) + + require.NoError(t, st.Close()) + closed = true + reopened, err := NewPebbleStore(dir) + require.NoError(t, err) + defer func() { require.NoError(t, reopened.Close()) }() + require.Equal(t, promotedTS, reopened.LastCommitTS()) + val, err = reopened.GetAt(ctx, []byte("k"), reopened.LastCommitTS()) + require.NoError(t, err) + require.Equal(t, []byte("v100"), val) + reopenedPromoter, ok := reopened.(MigrationPromoter) + require.True(t, ok) + retry, err := reopenedPromoter.PromoteVersions(ctx, PromoteVersionsOptions{ + JobID: 101, + StartKey: prefix, + EndKey: PrefixScanEnd(prefix), + MaxVersions: 10, + TargetKey: targetKey, + }) + require.NoError(t, err) + require.True(t, retry.Done) + require.Zero(t, retry.PromotedRows) + require.Equal(t, uint64(1), retry.TotalPromotedRows) + require.Equal(t, promotedTS, retry.MaxPromotedTS) +} + +func TestPromotionStateCodecPreservesMaxPromotedTS(t *testing.T) { + t.Parallel() + + state := PromotionState{ + Cursor: []byte("cursor"), + Done: true, + PromotedRows: 7, + MaxPromotedTS: 42, + LastError: "boom", + } + decoded, ok := decodePromotionState(encodePromotionState(state)) + require.True(t, ok) + require.Equal(t, state, decoded) + + old := []byte{migrationPromotionDoneFlag} + old = binary.BigEndian.AppendUint64(old, 3) + old = binary.AppendUvarint(old, lenAsUint64(len("old-cursor"))) + old = append(old, "old-cursor"...) + old = binary.AppendUvarint(old, lenAsUint64(len("old-error"))) + old = append(old, "old-error"...) + decoded, ok = decodePromotionState(old) + require.True(t, ok) + require.True(t, decoded.Done) + require.Equal(t, uint64(3), decoded.PromotedRows) + require.Zero(t, decoded.MaxPromotedTS) + require.Equal(t, []byte("old-cursor"), decoded.Cursor) + require.Equal(t, "old-error", decoded.LastError) +} + func TestPebbleImportMetadataPersistsAcrossReopen(t *testing.T) { ctx := context.Background() dir, err := os.MkdirTemp("", "migration-import-persist-*") diff --git a/store/mvcc_store.go b/store/mvcc_store.go index 8c8712eb9..88817d94d 100644 --- a/store/mvcc_store.go +++ b/store/mvcc_store.go @@ -60,13 +60,14 @@ func byteSliceComparator(a, b any) int { // mvccStore is an in-memory MVCC implementation backed by a treemap for // deterministic iteration order and range scans. type mvccStore struct { - tree *treemap.Map // key []byte -> []VersionedValue - mtx sync.RWMutex - log *slog.Logger - lastCommitTS uint64 - minRetainedTS uint64 - migrationAcks map[migrationAckID]migrationImportAck - migrationHLCFloors map[uint64]uint64 + tree *treemap.Map // key []byte -> []VersionedValue + mtx sync.RWMutex + log *slog.Logger + lastCommitTS uint64 + minRetainedTS uint64 + migrationAcks map[migrationAckID]migrationImportAck + migrationHLCFloors map[uint64]uint64 + migrationPromotions map[uint64]PromotionState // writeConflicts mirrors the per-(kind, key_prefix) counter from // the pebble-backed store so the in-memory implementation shows up // in the same Prometheus series (even if the counts are usually @@ -113,9 +114,10 @@ func NewMVCCStore(opts ...MVCCStoreOption) MVCCStore { log: slog.New(slog.NewJSONHandler(os.Stdout, &slog.HandlerOptions{ Level: slog.LevelWarn, })), - migrationAcks: make(map[migrationAckID]migrationImportAck), - migrationHLCFloors: make(map[uint64]uint64), - writeConflicts: newWriteConflictCounter(), + migrationAcks: make(map[migrationAckID]migrationImportAck), + migrationHLCFloors: make(map[uint64]uint64), + migrationPromotions: make(map[uint64]PromotionState), + writeConflicts: newWriteConflictCounter(), } for _, opt := range opts { opt(s) @@ -936,6 +938,7 @@ func (s *mvccStore) restoreStreamingSnapshot(r io.Reader) error { s.minRetainedTS = minRetainedTS s.migrationAcks = make(map[migrationAckID]migrationImportAck) s.migrationHLCFloors = make(map[uint64]uint64) + s.migrationPromotions = make(map[uint64]PromotionState) return nil } diff --git a/store/mvcc_store_snapshot_test.go b/store/mvcc_store_snapshot_test.go index b16cbe68b..3032a1cc3 100644 --- a/store/mvcc_store_snapshot_test.go +++ b/store/mvcc_store_snapshot_test.go @@ -61,6 +61,19 @@ func TestMVCCStore_RestoreClearsMigrationMetadata(t *testing.T) { ctx := context.Background() st := newTestMVCCStore(t) + promoter, ok := any(st).(MigrationPromoter) + require.True(t, ok) + stateReader, ok := any(st).(MigrationPromotionStateReader) + require.True(t, ok) + + prefix := []byte("stage|") + stage := func(raw string) []byte { + return append([]byte("stage|"), []byte(raw)...) + } + targetKey := func(staged []byte) ([]byte, bool) { + return bytes.TrimPrefix(staged, prefix), bytes.HasPrefix(staged, prefix) + } + require.NoError(t, st.PutAt(ctx, []byte("base"), []byte("v1"), 10, 0)) snap, err := st.Snapshot() @@ -80,13 +93,46 @@ func TestMVCCStore_RestoreClearsMigrationMetadata(t *testing.T) { require.NoError(t, err) require.Equal(t, uint64(50), floor) + require.NoError(t, st.PutAt(ctx, stage("stale"), []byte("old"), 70, 0)) + promoted, err := promoter.PromoteVersions(ctx, PromoteVersionsOptions{ + JobID: 7, + StartKey: prefix, + EndKey: PrefixScanEnd(prefix), + MaxVersions: 10, + TargetKey: targetKey, + }) + require.NoError(t, err) + require.True(t, promoted.Done) + state, ok, err := stateReader.MigrationPromotionState(ctx, 7) + require.NoError(t, err) + require.True(t, ok) + require.True(t, state.Done) + require.NoError(t, st.Restore(bytes.NewReader(raw))) floor, err = st.MigrationHLCFloor(ctx, 7) require.NoError(t, err) require.Zero(t, floor) + _, ok, err = stateReader.MigrationPromotionState(ctx, 7) + require.NoError(t, err) + require.False(t, ok) _, err = st.GetAt(ctx, []byte("imported"), 50) require.ErrorIs(t, err, ErrKeyNotFound) + require.NoError(t, st.PutAt(ctx, stage("fresh"), []byte("new"), 80, 0)) + promoted, err = promoter.PromoteVersions(ctx, PromoteVersionsOptions{ + JobID: 7, + StartKey: prefix, + EndKey: PrefixScanEnd(prefix), + MaxVersions: 10, + TargetKey: targetKey, + }) + require.NoError(t, err) + require.True(t, promoted.Done) + require.Equal(t, uint64(1), promoted.PromotedRows) + got, err := st.GetAt(ctx, []byte("fresh"), 80) + require.NoError(t, err) + require.Equal(t, []byte("new"), got) + res, err := st.ImportVersions(ctx, ImportVersionsOptions{ JobID: 7, BracketID: 3, diff --git a/store/store.go b/store/store.go index 9733626ba..1dcd96df8 100644 --- a/store/store.go +++ b/store/store.go @@ -120,6 +120,53 @@ type ImportVersionsResult struct { Duplicate bool } +// PromoteVersionsOptions atomically copies staged MVCC versions to their +// target keys and physically removes the staged versions. +type PromoteVersionsOptions struct { + JobID uint64 + // AppliedIndex is the optional Raft entry index to bundle with Pebble + // promotion batches as metaAppliedIndex. Zero leaves the meta key unchanged. + AppliedIndex uint64 + StartKey []byte + EndKey []byte + Cursor []byte + MaxVersions int + MaxBytes uint64 + MaxScannedBytes uint64 + TargetKey func(stagedKey []byte) ([]byte, bool) +} + +// PromoteVersionsResult reports one resumable staged-version promotion chunk. +type PromoteVersionsResult struct { + NextCursor []byte + Done bool + PromotedRows uint64 + TotalPromotedRows uint64 + PromotedBytes uint64 + MaxPromotedTS uint64 + ScannedBytes uint64 +} + +// PromotionState is the target-local durable cursor for staged data promotion. +type PromotionState struct { + Cursor []byte + Done bool + PromotedRows uint64 + MaxPromotedTS uint64 + LastError string +} + +// MigrationPromoter is implemented by stores that can promote staged range +// migration data into the live keyspace. +type MigrationPromoter interface { + PromoteVersions(ctx context.Context, opts PromoteVersionsOptions) (PromoteVersionsResult, error) +} + +// MigrationPromotionStateReader reads target-local staged promotion state. +type MigrationPromotionStateReader interface { + MigrationPromotionState(ctx context.Context, jobID uint64) (PromotionState, bool, error) +} + // OpType describes a mutation kind. type OpType int