migration: promote staged versions

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 {

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))
+}

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 {