migration: add target readiness guard

adapter/internal.go

@@ -254,6 +254,28 @@ func (i *Internal) PromoteStagedVersions(ctx context.Context, req *pb.PromoteSta
 	}, nil
 }
 
+func (i *Internal) ApplyTargetStagedReadiness(ctx context.Context, req *pb.TargetStagedReadinessRequest) (*pb.TargetStagedReadinessResponse, error) {
+	if req == nil {
+		return nil, errors.WithStack(status.Error(codes.InvalidArgument, "target staged readiness request is nil"))
+	}
+	if req.GetJobId() == 0 {
+		return nil, errors.WithStack(status.Error(codes.InvalidArgument, "target staged readiness job_id is required"))
+	}
+	if req.GetMigrationJobId() == 0 {
+		return nil, errors.WithStack(status.Error(codes.InvalidArgument, "target staged readiness migration_job_id is required"))
+	}
+	if i.migrationProposer == nil {
+		return nil, errors.WithStack(status.Error(codes.FailedPrecondition, "target staged readiness proposer is not configured"))
+	}
+	if err := i.verifyInternalLeader(ctx); err != nil {
+		return nil, err
+	}
+	if err := i.proposeTargetStagedReadiness(ctx, req); err != nil {
+		return nil, errors.WithStack(err)
+	}
+	return &pb.TargetStagedReadinessResponse{}, nil
+}
+
 func (i *Internal) verifyMigrationPromoteEnabled(ctx context.Context) error {
 	if i.migrationPromoteGate == nil {
 		return nil
@@ -349,6 +371,24 @@ func (i *Internal) proposeMigrationPromote(ctx context.Context, req *pb.PromoteS
 	}
 }
 
+func (i *Internal) proposeTargetStagedReadiness(ctx context.Context, req *pb.TargetStagedReadinessRequest) error {
+	cmd, err := kv.MarshalTargetStagedReadinessCommand(req)
+	if err != nil {
+		return errors.WithStack(err)
+	}
+	resp, err := i.proposeMigrationCommand(ctx, cmd, "target staged readiness")
+	if err != nil {
+		return errors.WithStack(err)
+	}
+	if resp == nil {
+		return nil
+	}
+	if err, ok := resp.(error); ok {
+		return errors.WithStack(err)
+	}
+	return errors.WithStack(errors.Newf("unexpected target readiness 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 {

adapter/internal_migration_test.go

@@ -498,3 +498,43 @@ func TestInternalPromoteStagedVersionsAppliesStoreBatch(t *testing.T) {
 	require.ErrorIs(t, err, store.ErrKeyNotFound)
 	require.GreaterOrEqual(t, clock.Current(), uint64(30))
 }
+
+func TestInternalApplyTargetStagedReadinessProposesThroughRaft(t *testing.T) {
+	t.Parallel()
+
+	ctx := context.Background()
+	st := store.NewMVCCStore()
+	proposer := &applyingMigrationProposer{
+		fsm: kv.NewKvFSMWithHLC(st, nil),
+	}
+	internal := NewInternalWithEngine(nil, mockInternalLeader{}, nil, nil,
+		WithInternalStore(st),
+		WithInternalMigrationProposer(proposer),
+	)
+
+	_, err := internal.ApplyTargetStagedReadiness(ctx, &pb.TargetStagedReadinessRequest{
+		JobId:                  9,
+		RouteStart:             []byte("a"),
+		RouteEnd:               []byte("z"),
+		ExpectedCutoverVersion: 3,
+		MigrationJobId:         7,
+		MinWriteTsExclusive:    100,
+		Armed:                  true,
+	})
+	require.NoError(t, err)
+	require.Equal(t, uint64(1), proposer.calls)
+
+	reader, ok := st.(store.MigrationTargetReadinessReader)
+	require.True(t, ok)
+	states, err := reader.MigrationTargetReadinessStates(ctx)
+	require.NoError(t, err)
+	require.Equal(t, []store.TargetStagedReadinessState{{
+		JobID:                  9,
+		RouteStart:             []byte("a"),
+		RouteEnd:               []byte("z"),
+		ExpectedCutoverVersion: 3,
+		MigrationJobID:         7,
+		MinWriteTSExclusive:    100,
+		Armed:                  true,
+	}}, states)
+}

adapter/redis_compat_commands_stream_test.go

@@ -279,7 +279,7 @@ func TestRedis_StreamXReadLatencyIsConstant(t *testing.T) {
 
 	rdb := redis.NewClient(&redis.Options{Addr: nodes[0].redisAddress})
 	defer func() { _ = rdb.Close() }()
-	ctx := context.Background()
+	ctx := t.Context()
 
 	const (
 		total  = 10_000
@@ -291,13 +291,21 @@ func TestRedis_StreamXReadLatencyIsConstant(t *testing.T) {
 	}
 
 	measure := func() time.Duration {
-		start := time.Now()
-		streams, err := rdb.XRead(ctx, &redis.XReadArgs{
-			Streams: []string{"stream-lat", lastID},
-			Count:   10,
-			Block:   10 * time.Millisecond,
-		}).Result()
-		elapsed := time.Since(start)
+		var (
+			streams []redis.XStream
+			elapsed time.Duration
+		)
+		err := retryNotLeader(ctx, func() error {
+			start := time.Now()
+			var xerr error
+			streams, xerr = rdb.XRead(ctx, &redis.XReadArgs{
+				Streams: []string{"stream-lat", lastID},
+				Count:   10,
+				Block:   10 * time.Millisecond,
+			}).Result()
+			elapsed = time.Since(start)
+			return xerr
+		})
 		require.True(t, errors.Is(err, redis.Nil) || err == nil)
 		require.Empty(t, streams)
 		return elapsed